| """One-time migration driver. Only copies/imports/verifies; NEVER deletes sources.""" |
| import concurrent.futures |
| import json |
| from pathlib import Path |
| import re |
| import sys |
|
|
| from checkpoint_archive import Archive, atomic_json |
|
|
| BASE = Path('/mnt/data/limingleyang') |
| ROOT = BASE / 'checkpoints' |
| PLAN = json.loads((ROOT / 'migration/plan.json').read_text()) |
|
|
|
|
| def add_models(): |
| archive = Archive(ROOT) |
| archive.init() |
| dependencies = {} |
| for row in PLAN['svae']: |
| source = Path(row['source']) |
| origin = min(int(x['id'][3:]) for x in row['consumers']) |
| key = f'svae/run{origin:02d}' |
| name = Path(row['destination']).name |
| files = [p.name for p in source.parent.iterdir() if p.is_file() and p.name in ['svae.pt', 'config.yaml', 'complete.json']] |
| extras = {} |
| roots = [source.parent.parent] |
| if source.parent.parent.name == 'svae_train': |
| roots.append(source.parent.parent.parent) |
| for root in list(roots): |
| roots.extend(root / name for name in ['features', 'frames'] if (root / name).is_dir()) |
| for root in roots: |
| for f in root.iterdir(): |
| if f.is_file() and (f.suffix in ['.json', '.yaml', '.yml', '.txt'] or 'stats' in f.name and f.suffix in ['.pt', '.npy']): |
| extras['collection/' + f.relative_to(BASE).as_posix()] = str(f) |
| archive.add_model(key=key, name=name, source=source.parent, files=files, extra_assets=extras, |
| metadata={'consumers': row['consumers'], 'original_checkpoint': str(source), |
| 'note': 'Original independent model precision retained; policy bf16 tensors are separate exact copies.'}) |
| dependencies[str(source)] = key |
| print(json.dumps({'model_imported': key, 'metadata_files': len(extras)}), flush=True) |
| row = PLAN['extra_models'][0] |
| source = Path(row['source']) |
| files = [p.name for p in source.iterdir() if p.is_file() and |
| (p.name in ['last.pt', 'encoder.pt', 'latent_stats_libero.pt'] or p.suffix in ['.json', '.yaml', '.yml'])] |
| archive.add_model(key='pretrain/run19-encoder', name=Path(row['destination']).name, |
| source=source, files=files, metadata={'final_training_step': 1000, 'kind': 'VideoRAE encoder fine-tuning'}) |
| return dependencies |
|
|
|
|
| def import_row(pair): |
| row, svae_keys = pair |
| p = Path(row['source']) |
| actual = p.stat() |
| if actual.st_size != row['source_size'] or abs(actual.st_mtime - row['source_mtime']) > 0.001: |
| raise ValueError(f'Source changed since approved inventory: {p}') |
| latest = max(int(re.fullmatch(r'checkpoint_step_(\d+)\.safetensors', f.name)[1]) for f in p.parent.glob('checkpoint_step_*.safetensors')) |
| if latest != row['step']: |
| raise ValueError(f'Final checkpoint selection changed: {p.parent}') |
| deps = {} |
| for svae in PLAN['svae']: |
| if any(x['id'] == row['id'] and x['stage'] == row['stage'] for x in svae['consumers']): |
| deps['svae'] = svae_keys[svae['source']] |
| if row['stage'] == 'sft' and row['run_number'] in (23, 39, 40): |
| deps['pretrain'] = 'pretrain/run39' if row['run_number'] == 39 else 'pretrain/run23' |
| if row['run_number'] == 19: |
| deps['encoder'] = 'pretrain/run19-encoder' |
| return Archive(ROOT).add_policy(run=row['id'], name=row['run_name'], stage=row['stage'], source=p, |
| share_decoder=row['share_vlm'], dependencies=deps) |
|
|
|
|
| def main(): |
| ROOT.joinpath('migration').mkdir(parents=True, exist_ok=True) |
| atomic_json(ROOT / 'migration/status.json', {'status': 'importing', 'source_cleanup': 'not_started'}) |
| try: |
| dependencies = add_models() |
| for stage in ['pretrain', 'sft']: |
| rows = [r for r in PLAN['policies'] if r['stage'] == stage] |
| with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: |
| for entry in pool.map(import_row, [(row, dependencies) for row in rows]): |
| print(json.dumps({'published': entry['id']}), flush=True) |
| atomic_json(ROOT / 'migration/status.json', {'status': 'copied_verifying', 'source_cleanup': 'not_started'}) |
| archive = Archive(ROOT) |
| report = archive.verify(reconstruct=True, report=ROOT / 'migration/verification.json') |
| assert len(report['entries']) == 53, len(report['entries']) |
| atomic_json(ROOT / 'migration/status.json', {'status': 'verified', 'entries': 53, 'source_cleanup': 'not_started'}) |
| print('IMPORT_AND_BYTE_EQUIVALENCE_VERIFICATION_COMPLETE', flush=True) |
| except BaseException as exc: |
| atomic_json(ROOT / 'migration/status.json', {'status': 'failed', 'source_cleanup': 'not_started', 'error': str(exc)}) |
| raise |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
