| |
| """Restore an archived training source tree into a new directory.""" |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--list', action='store_true') |
| parser.add_argument('--snapshot') |
| parser.add_argument('--output', type=Path) |
| args = parser.parse_args() |
| index = json.loads((ROOT/'training/SNAPSHOTS.json').read_text()) |
| if args.list: |
| print(json.dumps(index, indent=2)); return |
| if args.snapshot not in index or args.output is None: |
| parser.error('Choose --snapshot from --list and supply a new --output directory.') |
| output = args.output.resolve() |
| if output.exists(): |
| parser.error('Output already exists; choose a new directory. No files were changed.') |
| manifest = json.loads((ROOT/index[args.snapshot]['snapshot']).read_text()) |
| checked = [] |
| for name, record in manifest['files'].items(): |
| destination = (output/name).resolve() |
| source = (ROOT/record['file']).resolve() |
| if not destination.is_relative_to(output) or not source.is_relative_to(ROOT): |
| raise ValueError('Invalid snapshot path') |
| data = source.read_bytes() |
| if hashlib.sha256(data).hexdigest() != record['packaged_sha256']: |
| raise ValueError(f'Snapshot integrity mismatch: {name}') |
| checked.append((destination, data)) |
| output.mkdir(parents=True, exist_ok=False) |
| for destination, data in checked: |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| destination.write_bytes(data) |
| print(json.dumps({'snapshot': args.snapshot, 'files': len(checked), 'output': str(output)}, indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
