| """Copy mounted inference files sequentially and verify Hub content hashes.""" |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| PROTOCOL = 'sequential-local-copy-hub-hash-v1' |
|
|
|
|
| def stage_checkpoint(mount, destination, manifest_path): |
| mount, destination, manifest_path = map(Path, (mount, destination, manifest_path)) |
| manifest_bytes = manifest_path.read_bytes() |
| manifest = json.loads(manifest_bytes) |
| destination.mkdir(parents=True, exist_ok=True) |
| total = 0 |
| for item in manifest['files']: |
| name = item['name'] |
| assert not Path(name).is_absolute() and '..' not in Path(name).parts, 'Unsafe checkpoint relative path' |
| (destination/name).parent.mkdir(parents=True,exist_ok=True) |
| source = mount/name |
| assert source.stat().st_size == item['size'], 'Mounted file size differs: '+name |
| print('SERVER_MODEL_STAGE_BEGIN', name, item['size'], flush=True) |
| if item['hash_algorithm'] == 'sha256': |
| digest = hashlib.sha256() |
| else: |
| assert item['hash_algorithm'] == 'git-sha1' |
| digest = hashlib.sha1() |
| digest.update(('blob '+str(item['size'])+'\0').encode()) |
| temporary = destination/(name+'.partial') |
| count = 0; next_progress = 512*1024*1024 |
| try: |
| with source.open('rb') as reader, temporary.open('wb') as writer: |
| while block := reader.read(8*1024*1024): |
| writer.write(block); digest.update(block); count += len(block) |
| if count >= next_progress: |
| print('SERVER_MODEL_STAGE_PROGRESS', name, count, item['size'], flush=True) |
| next_progress += 512*1024*1024 |
| assert count == item['size'], 'Incomplete file copy: '+name |
| assert digest.hexdigest() == item['hash'], 'Content hash differs: '+name |
| temporary.replace(destination/name) |
| except Exception: |
| temporary.unlink(missing_ok=True) |
| raise |
| total += count |
| print('SERVER_MODEL_FILE_VERIFIED', name, count, flush=True) |
| print('SERVER_MODEL_STAGE_COMPLETE', len(manifest['files']), total, flush=True) |
| return dict(protocol=PROTOCOL, manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(), |
| source_files_revision=manifest['source_files_revision'], verified_files=len(manifest['files']), |
| verified_bytes=total) |
|
|
