File size: 4,815 Bytes
9020600 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """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()
|