| |
| """Portable, lossless OpenWAM checkpoint archive. Python standard library only. |
| |
| Weights are copied, never linked. Import is append-only, content-addressed shared |
| components are immutable, and publication uses a filesystem lock and atomic JSON. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import base64 |
| import concurrent.futures |
| import contextlib |
| import datetime |
| import fcntl |
| import hashlib |
| import json |
| import math |
| import os |
| from pathlib import Path, PurePosixPath |
| import re |
| import shutil |
| import struct |
| import subprocess |
| import sys |
| import tempfile |
| import uuid |
|
|
| VERSION = 1 |
| BLOCK = 8 * 1024 * 1024 |
| PROCESS_HASH_CACHE = {} |
| DTYPE_BYTES = {'BOOL': 1, 'U8': 1, 'I8': 1, 'F8_E4M3': 1, 'F8_E5M2': 1, |
| 'I16': 2, 'U16': 2, 'F16': 2, 'BF16': 2, 'I32': 4, 'U32': 4, |
| 'F32': 4, 'I64': 8, 'U64': 8, 'F64': 8} |
|
|
|
|
| def now(): |
| return datetime.datetime.now(datetime.timezone.utc).isoformat() |
|
|
|
|
| def json_bytes(value): |
| return json.dumps(value, sort_keys=True, separators=(',', ':'), ensure_ascii=False).encode() |
|
|
|
|
| def atomic_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temp = path.with_name(path.name + '.tmp-' + uuid.uuid4().hex) |
| try: |
| with temp.open('xb') as f: |
| f.write(json.dumps(value, indent=2, ensure_ascii=False).encode() + b'\n') |
| f.flush() |
| os.fsync(f.fileno()) |
| os.replace(temp, path) |
| finally: |
| temp.unlink(missing_ok=True) |
|
|
|
|
| def sha256_file(path): |
| h = hashlib.sha256() |
| with Path(path).open('rb') as f: |
| while chunk := f.read(BLOCK): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def stable_stat(path): |
| s = Path(path).stat() |
| return {'size': s.st_size, 'mtime_ns': s.st_mtime_ns, 'inode': s.st_ino, 'device': s.st_dev} |
|
|
|
|
| def checked_path(root, relative): |
| """No absolute paths, traversal, symlinks, or references outside the archive.""" |
| rel = PurePosixPath(relative) |
| if rel.is_absolute() or not rel.parts or any(p in ('..', '.') for p in rel.parts): |
| raise ValueError(f'Invalid archive-relative path: {relative!r}') |
| root = Path(root).resolve() |
| path = root |
| for part in rel.parts: |
| path = path / part |
| if path.is_symlink(): |
| raise ValueError(f'Archive contains a symlink: {path}') |
| if not path.resolve().is_relative_to(root): |
| raise ValueError(f'Archive path escaped root: {relative}') |
| return path |
|
|
|
|
| @contextlib.contextmanager |
| def locked(path): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open('a') as f: |
| fcntl.flock(f, fcntl.LOCK_EX) |
| yield |
|
|
|
|
| def canonical_run(run): |
| if not re.fullmatch(r'run\d+', run): |
| raise ValueError('run id must be run followed by digits, e.g. run41') |
| return f'run{int(run[3:]):02d}' |
|
|
|
|
| def read_header(path): |
| with Path(path).open('rb') as f: |
| raw_n = f.read(8) |
| if len(raw_n) != 8: |
| raise ValueError(f'Truncated safetensors: {path}') |
| n = struct.unpack('<Q', raw_n)[0] |
| if n > 100_000_000 or n < 2: |
| raise ValueError(f'Invalid safetensors header: {path}') |
| raw = f.read(n) |
| if len(raw) != n: |
| raise ValueError(f'Truncated safetensors header: {path}') |
| def pairs_hook(pairs): |
| result = {} |
| for key, value in pairs: |
| if key in result: |
| raise ValueError(f'Duplicate JSON key: {key}') |
| result[key] = value |
| return result |
| header = json.loads(raw, object_pairs_hook=pairs_hook) |
| tensors = {k: v for k, v in header.items() if k != '__metadata__'} |
| end = 0 |
| for key, value in sorted(tensors.items(), key=lambda kv: kv[1]['data_offsets']): |
| a, b = value['data_offsets'] |
| dtype = value['dtype'] |
| shape = value['shape'] |
| if dtype not in DTYPE_BYTES or any(not isinstance(x, int) or x < 0 for x in shape): |
| raise ValueError(f'Unsupported tensor layout: {path}: {key}') |
| if a != end or b - a != math.prod(shape) * DTYPE_BYTES[dtype]: |
| raise ValueError(f'Invalid tensor offsets: {path}: {key}') |
| end = b |
| if len(raw_n) + len(raw) + end != Path(path).stat().st_size: |
| raise ValueError(f'Checkpoint length mismatch: {path}') |
| return raw_n + raw, tensors |
|
|
|
|
| def shard_layout(keys, tensors): |
| offsets, header, end = {}, {}, 0 |
| for key in keys: |
| src = tensors[key] |
| n = src['data_offsets'][1] - src['data_offsets'][0] |
| offsets[key] = (end, end + n) |
| header[key] = {'dtype': src['dtype'], 'shape': src['shape'], 'data_offsets': [end, end + n]} |
| end += n |
| raw = json_bytes(header) |
| raw += b' ' * (-len(raw) % 8) |
| return struct.pack('<Q', len(raw)) + raw, offsets |
|
|
|
|
| def component(key, share_decoder): |
| for prefix, name in [ |
| ('video_backbone.reason1.', 'reason1'), |
| ('video_backbone.vae.', 'wan-vae'), |
| ('video_backbone.video_encoder._m.', 'visual-encoder'), |
| ('video_backbone.video_encoder._svae.', 'svae'), |
| ('video_backbone.video_encoder.teacher_model.', 'videorae-teacher'), |
| ('video_backbone.video_encoder.latent_compressor.', 'videorae-compressor'), |
| ('vlm_expert.vlm.visual.', 'qwen-vision'), |
| ('vlm_expert.vlm.language_model.embed_tokens.', 'qwen-embeddings'), |
| ]: |
| if key.startswith(prefix): |
| return name |
| if share_decoder and key.startswith('vlm_expert.vlm.'): |
| return 'qwen-frozen-decoder' |
| return 'private' |
|
|
|
|
| class Archive: |
| def __init__(self, root): |
| self.root = Path(root).resolve() |
| self.hash_cache = PROCESS_HASH_CACHE |
|
|
| def init(self): |
| self.root.mkdir(parents=True, exist_ok=True) |
| with locked(self.root / '.archive.lock'): |
| path = self.root / 'catalog.json' |
| if not path.exists(): |
| atomic_json(path, {'format': 'openwam-checkpoint-archive', 'schema_version': VERSION, |
| 'created_at': now(), 'updated_at': now(), 'entries': {}}) |
| self.catalog() |
|
|
| def catalog(self): |
| c = json.loads((self.root / 'catalog.json').read_text()) |
| if c.get('schema_version') != VERSION or c.get('format') != 'openwam-checkpoint-archive': |
| raise ValueError('Unsupported catalog version/format') |
| return c |
|
|
| def check_file(self, record, full=True): |
| path = checked_path(self.root, record['path']) |
| if not path.is_file() or path.stat().st_size != record['bytes']: |
| raise ValueError(f'Missing or truncated archive file: {path}') |
| if full: |
| stat = path.stat() |
| cache_key = (str(path), stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns) |
| digest = self.hash_cache.get(cache_key) |
| if digest is None: |
| digest = sha256_file(path) |
| self.hash_cache[cache_key] = digest |
| if digest != record['sha256']: |
| raise ValueError(f'Checksum mismatch: {path}') |
| return path |
|
|
| def resolve(self, run, stage='sft'): |
| entries = self.catalog()['entries'] |
| key = run if '/' in run else f'{stage}/{canonical_run(run)}' |
| if key not in entries: |
| raise KeyError(f'No archive entry {key}; use list to see available runs/stages') |
| return entries[key] |
|
|
| def index(self, entry): |
| path = self.check_file(entry['index']) |
| index = json.loads(path.read_text()) |
| if index.get('schema_version') != VERSION or index['id'] != entry['id']: |
| raise ValueError('Index/catalog identity mismatch') |
| return index |
|
|
| def _publish(self, stage_dir, destination, index, entry): |
| final = checked_path(self.root, destination) |
| with locked(self.root / '.archive.lock'): |
| catalog = self.catalog() |
| for dep in entry.get('dependencies', {}).values(): |
| if dep not in catalog['entries']: |
| raise ValueError(f'Missing catalog dependency: {dep}') |
| if entry['id'] in catalog['entries'] or final.exists(): |
| raise FileExistsError(f'Archive entry already exists: {entry["id"]}; immutable entries cannot be overwritten') |
| atomic_json(stage_dir / 'checkpoint.index.json', index) |
| entry['index'] = {'path': destination + '/checkpoint.index.json', |
| 'bytes': (stage_dir / 'checkpoint.index.json').stat().st_size, |
| 'sha256': sha256_file(stage_dir / 'checkpoint.index.json')} |
| final.parent.mkdir(parents=True, exist_ok=True) |
| os.rename(stage_dir, final) |
| catalog['entries'][entry['id']] = entry |
| catalog['updated_at'] = now() |
| try: |
| atomic_json(self.root / 'catalog.json', catalog) |
| except BaseException: |
| os.rename(final, stage_dir) |
| raise |
| return entry |
|
|
| def _copy_file(self, source, dest, relative): |
| source, dest = Path(source), Path(dest) |
| before = stable_stat(source) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| h = hashlib.sha256() |
| with source.open('rb') as src, dest.open('xb') as out: |
| while chunk := src.read(BLOCK): |
| out.write(chunk) |
| h.update(chunk) |
| out.flush() |
| os.fsync(out.fileno()) |
| if stable_stat(source) != before or sha256_file(dest) != h.hexdigest(): |
| raise ValueError(f'Copy verification failed: {source}') |
| if dest.stat().st_ino == source.stat().st_ino and dest.stat().st_dev == source.stat().st_dev: |
| raise ValueError('Copy unexpectedly shares an inode with its source') |
| return {'path': relative, 'bytes': before['size'], 'sha256': h.hexdigest(), |
| 'source': str(source), 'source_stat': before} |
|
|
| def _asset_sources(self, source_dir): |
| for root, dirs, files in os.walk(source_dir, followlinks=False): |
| dirs[:] = [d for d in dirs if not d.startswith(('accel_state_', '.')) and d not in ('wandb', 'logs', 'eval', 'features')] |
| for d in dirs: |
| if (Path(root) / d).is_symlink(): |
| raise ValueError(f'Unexpected symlinked metadata directory: {Path(root) / d}') |
| for name in sorted(files): |
| if name.startswith(('checkpoint_step_', '.')) or name.endswith(('.log', '.mp4')): |
| continue |
| src = Path(root) / name |
| rel = src.relative_to(source_dir).as_posix() |
| yield src, rel |
|
|
| def _assets(self, source_dir, stage_dir, destination): |
| assets = [] |
| for src, rel in self._asset_sources(source_dir): |
| record = self._copy_file(src, stage_dir / 'private/assets' / rel, |
| destination + '/private/assets/' + rel) |
| record['restore_path'] = rel |
| assets.append(record) |
| return assets |
|
|
| def _install_shared(self, source, prefix, tensors, keys, hashes, label, stage_dir): |
| identity = hashlib.sha256(json_bytes([ |
| [k, tensors[k]['dtype'], tensors[k]['shape'], hashes[k]] for k in sorted(keys) |
| ])).hexdigest() |
| base = ('svae/_components' if label == 'svae' else 'shared/' + label) + '/' + identity |
| relative = base + '/model.safetensors' |
| path = checked_path(self.root, relative) |
| meta_path = checked_path(self.root, base + '/component.json') |
| with locked(self.root / '.archive.lock'): |
| if path.exists() and meta_path.exists(): |
| record = json.loads(meta_path.read_text()) |
| if record['identity'] != identity or record['path'] != relative: |
| raise ValueError('Invalid shared component metadata') |
| self.check_file(record) |
| return record |
| header, _ = shard_layout(keys, tensors) |
| temp = stage_dir / (label + '-' + uuid.uuid4().hex + '.safetensors') |
| h = hashlib.sha256(header) |
| with source.open('rb') as src, temp.open('xb') as out: |
| out.write(header) |
| for key in keys: |
| a, b = tensors[key]['data_offsets'] |
| src.seek(len(prefix) + a) |
| remaining = b - a |
| tensor_hash = hashlib.sha256() |
| while remaining: |
| data = src.read(min(BLOCK, remaining)) |
| if not data: |
| raise ValueError('Source truncated during shared-component copy') |
| out.write(data) |
| h.update(data) |
| tensor_hash.update(data) |
| remaining -= len(data) |
| if tensor_hash.hexdigest() != hashes[key]: |
| raise ValueError(f'Source changed while importing: {source}: {key}') |
| out.flush() |
| os.fsync(out.fileno()) |
| record = {'path': relative, 'bytes': temp.stat().st_size, 'sha256': h.hexdigest(), |
| 'identity': identity, 'component': label, 'tensor_count': len(keys)} |
| if sha256_file(temp) != record['sha256']: |
| raise ValueError('Shared component readback verification failed') |
| with locked(self.root / '.archive.lock'): |
| if path.exists(): |
| if not meta_path.exists(): |
| raise ValueError(f'Incomplete component publication: {path}') |
| self.check_file(record) |
| temp.unlink() |
| else: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| os.rename(temp, path) |
| atomic_json(meta_path, record) |
| return record |
|
|
| def add_policy(self, *, run, name, stage, source, share_decoder=False, dependencies=None): |
| self.init() |
| run = canonical_run(run) |
| if stage not in ('sft', 'pretrain') or not re.fullmatch(r'[A-Za-z0-9_.-]+', name): |
| raise ValueError('Invalid stage/name') |
| if not re.match(r'^run0*' + str(int(run[3:])) + r'(?:_|-|$)', name): |
| raise ValueError('Run name must start with the requested run number') |
| source = Path(source).resolve() |
| match = re.fullmatch(r'checkpoint_step_(\d+)\.safetensors', source.name) |
| if not match or not source.is_file(): |
| raise ValueError('Pass the explicit completed checkpoint_step_N.safetensors file') |
| for required in ['config.yaml'] + (['normalization_stats.npy'] if stage == 'sft' else []): |
| if not (source.parent / required).is_file(): |
| raise ValueError(f'Missing deploy asset: {source.parent / required}') |
| key = stage + '/' + run |
| existing = self.catalog()['entries'].get(key) |
| if existing: |
| if existing['source']['sha256'] != sha256_file(source) or existing.get('dependencies', {}) != (dependencies or {}): |
| raise FileExistsError(f'Different checkpoint/dependencies already registered for {key}') |
| prior_assets = self.index(existing)['assets'] |
| if {rel for _, rel in self._asset_sources(source.parent)} != {a['restore_path'] for a in prior_assets}: |
| raise FileExistsError(f'Deploy metadata file set changed for existing entry {key}') |
| for asset in prior_assets: |
| current = checked_path(source.parent, asset['restore_path']) |
| if not current.is_file() or sha256_file(current) != asset['sha256']: |
| raise FileExistsError(f'Deploy metadata changed for existing entry {key}: {current}') |
| self.verify_entry(existing) |
| return existing |
| name = re.sub(r'^run\d+', run, name) |
| destination = stage + '/' + name |
| (self.root / '.staging').mkdir(exist_ok=True) |
| temp = Path(tempfile.mkdtemp(prefix=run + '-', dir=self.root / '.staging')) |
| try: |
| before = stable_stat(source) |
| prefix, tensors = read_header(source) |
| keys = sorted(tensors, key=lambda k: tensors[k]['data_offsets']) |
| groups = {} |
| for k in keys: |
| groups.setdefault(component(k, share_decoder), []).append(k) |
| private_keys = groups.pop('private', []) |
| private_header, _ = shard_layout(private_keys, tensors) |
| private = temp / 'private/weights.safetensors' |
| private.parent.mkdir() |
| private_hash = hashlib.sha256(private_header) |
| original_hash = hashlib.sha256(prefix) |
| hashes = {} |
| with source.open('rb') as src, private.open('xb') as out: |
| out.write(private_header) |
| src.seek(len(prefix)) |
| for k in keys: |
| a, b = tensors[k]['data_offsets'] |
| left = b - a |
| h = hashlib.sha256() |
| own = component(k, share_decoder) == 'private' |
| while left: |
| chunk = src.read(min(BLOCK, left)) |
| if not chunk: |
| raise ValueError(f'Truncated source: {source}') |
| original_hash.update(chunk) |
| h.update(chunk) |
| if own: |
| out.write(chunk) |
| private_hash.update(chunk) |
| left -= len(chunk) |
| hashes[k] = h.hexdigest() |
| out.flush() |
| os.fsync(out.fileno()) |
| if stable_stat(source) != before: |
| raise ValueError(f'Source changed during import: {source}') |
| private_record = {'path': destination + '/private/weights.safetensors', |
| 'bytes': private.stat().st_size, 'sha256': private_hash.hexdigest()} |
| if sha256_file(private) != private_record['sha256']: |
| raise ValueError('Private weights readback verification failed') |
| records = [private_record] |
| mapping = {k: private_record['path'] for k in private_keys} |
| for label, group_keys in groups.items(): |
| record = self._install_shared(source, prefix, tensors, group_keys, hashes, label, temp) |
| records.append(record) |
| mapping.update({k: record['path'] for k in group_keys}) |
| assets = self._assets(source.parent, temp, destination) |
| if stable_stat(source) != before: |
| raise ValueError('Source changed while metadata was copied') |
| provenance = {'path': str(source), 'bytes': before['size'], 'sha256': original_hash.hexdigest(), |
| 'stat': before, 'filename': source.name} |
| index = {'schema_version': VERSION, 'kind': 'policy', 'id': key, |
| 'original': provenance, 'original_header_base64': base64.b64encode(prefix).decode(), |
| 'weight_map': mapping, 'files': records, 'assets': assets, 'tensor_sha256': hashes} |
| entry = {'id': key, 'run_id': run, 'run_name': name, 'stage': stage, 'step': int(match[1]), |
| 'kind': 'policy', 'directory': destination, 'source': provenance, |
| 'dependencies': dependencies or {}, 'created_at': now(), 'tensor_count': len(keys)} |
| self._publish(temp, destination, index, entry) |
| print(json.dumps({'imported': key, 'step': entry['step'], 'original_bytes': before['size'], |
| 'private_bytes': private_record['bytes']}), flush=True, file=sys.stderr) |
| return entry |
| finally: |
| if temp.exists(): |
| shutil.rmtree(temp) |
|
|
| def add_model(self, *, key, name, source, files, metadata=None, extra_assets=None): |
| self.init() |
| if not re.fullmatch(r'(?:svae|pretrain)/[A-Za-z0-9_.-]+', key) or not re.fullmatch(r'[A-Za-z0-9_.-]+', name): |
| raise ValueError('Invalid model key/name') |
| stage, run = key.split('/') |
| source = Path(source).resolve() |
| destination = stage + '/' + name |
| self.root.joinpath('.staging').mkdir(exist_ok=True) |
| temp = Path(tempfile.mkdtemp(prefix='model-', dir=self.root / '.staging')) |
| try: |
| assets = [] |
| sources = [(checked_path(source, rel), rel) for rel in sorted(files)] |
| sources.extend((Path(src), rel) for rel, src in (extra_assets or {}).items()) |
| if len({rel for _, rel in sources}) != len(sources): |
| raise ValueError('Duplicate destination in model assets') |
| for src, rel in sources: |
| if not src.is_file(): |
| raise FileNotFoundError(src) |
| r = self._copy_file(src, checked_path(temp / 'private', rel), destination + '/private/' + rel) |
| r['restore_path'] = rel |
| assets.append(r) |
| existing = self.catalog()['entries'].get(key) |
| if existing: |
| prior = self.index(existing) |
| if {r['restore_path']: r['sha256'] for r in prior['assets']} != {r['restore_path']: r['sha256'] for r in assets}: |
| raise FileExistsError(f'Different model already registered: {key}') |
| self.verify_entry(existing) |
| return existing |
| index = {'schema_version': VERSION, 'kind': 'model', 'id': key, 'assets': assets, 'files': [], 'metadata': metadata or {}} |
| entry = {'id': key, 'run_id': run, 'run_name': name, 'stage': stage, 'kind': 'model', |
| 'directory': destination, 'created_at': now(), 'dependencies': {}, 'metadata': metadata or {}} |
| return self._publish(temp, destination, index, entry) |
| finally: |
| if temp.exists(): |
| shutil.rmtree(temp) |
|
|
| def verify_entry(self, entry, reconstruct=False): |
| index = self.index(entry) |
| for record in index['files'] + index['assets']: |
| self.check_file(record) |
| for dependency in entry.get('dependencies', {}).values(): |
| self.resolve(dependency) |
| if index['kind'] == 'policy': |
| self._validate_mapping(index) |
| if reconstruct: |
| self._reconstruct(index) |
| return {'id': entry['id'], 'files': len(index['files']) + len(index['assets']), |
| 'original_byte_equivalence': bool(reconstruct and index['kind'] == 'policy')} |
|
|
| def _validate_mapping(self, index): |
| prefix = base64.b64decode(index['original_header_base64'], validate=True) |
| n = struct.unpack('<Q', prefix[:8])[0] |
| if len(prefix) != n + 8: |
| raise ValueError('Invalid original checkpoint header length') |
| original = {k: v for k, v in json.loads(prefix[8:]).items() if k != '__metadata__'} |
| if set(original) != set(index['weight_map']) or set(original) != set(index['tensor_sha256']): |
| raise ValueError('Index must cover exactly every original tensor') |
| files = {f['path']: f for f in index['files']} |
| if len(files) != len(index['files']) or set(index['weight_map'].values()) - files.keys(): |
| raise ValueError('Duplicate or undeclared shard reference') |
| shards = {} |
| for relative in files: |
| p = self.check_file(files[relative], full=False) |
| h, ts = read_header(p) |
| assigned = {k for k, path in index['weight_map'].items() if path == relative} |
| if set(ts) != assigned: |
| raise ValueError('Shard/index keys differ; cannot silently load extra or missing parameters') |
| for k in assigned: |
| if any(ts[k][f] != original[k][f] for f in ['dtype', 'shape']): |
| raise ValueError(f'Tensor schema mismatch: {k}') |
| shards[relative] = (p, len(h), ts) |
| return prefix, original, shards |
|
|
| def _reconstruct(self, index, output=None): |
| prefix, tensors, shards = self._validate_mapping(index) |
| h = hashlib.sha256(prefix) |
| if output: |
| output.write(prefix) |
| handles = {} |
| try: |
| for relative, (path, start, _) in shards.items(): |
| handles[relative] = path.open('rb') |
| written = len(prefix) |
| for key in sorted(tensors, key=lambda k: tensors[k]['data_offsets']): |
| rel = index['weight_map'][key] |
| _, offset, table = shards[rel] |
| a, b = table[key]['data_offsets'] |
| f = handles[rel] |
| f.seek(offset + a) |
| left = b - a |
| tensor_hash = hashlib.sha256() |
| while left: |
| chunk = f.read(min(BLOCK, left)) |
| if not chunk: |
| raise ValueError('Truncated archive shard') |
| tensor_hash.update(chunk) |
| h.update(chunk) |
| if output: |
| output.write(chunk) |
| written += len(chunk) |
| left -= len(chunk) |
| if tensor_hash.hexdigest() != index['tensor_sha256'][key]: |
| raise ValueError(f'Tensor byte mismatch: {key}') |
| if h.hexdigest() != index['original']['sha256'] or written != index['original']['bytes']: |
| raise ValueError('Reconstruction does not match the original checkpoint byte-for-byte') |
| finally: |
| for f in handles.values(): |
| f.close() |
|
|
| def materialize(self, run, output, stage='sft'): |
| entry = self.resolve(run, stage) |
| index = self.index(entry) |
| output = Path(output).resolve() |
| if output.is_relative_to(self.root) or self.root.is_relative_to(output): |
| raise ValueError('Materialization must be outside the archive, not an ancestor of it') |
| output.parent.mkdir(parents=True, exist_ok=True) |
| with locked(output.parent / ('.' + output.name + '.lock')): |
| marker = output / '.archive-materialized.json' |
| if output.exists(): |
| if not marker.is_file() or json.loads(marker.read_text()).get('index_sha256') != entry['index']['sha256']: |
| raise FileExistsError(f'Output exists and is not a matching verified archive export: {output}') |
| proof = json.loads(marker.read_text()) |
| expected = [{'path': a['restore_path'], 'bytes': a['bytes'], 'sha256': a['sha256']} for a in index['assets']] |
| if index['kind'] == 'policy': |
| expected.append({'path': index['original']['filename'], 'bytes': index['original']['bytes'], 'sha256': index['original']['sha256']}) |
| if sorted(proof.get('files', []), key=lambda r: r['path']) != sorted(expected, key=lambda r: r['path']): |
| raise ValueError(f'Materialization marker does not match the archive index: {marker}') |
| for record in expected: |
| p = checked_path(output, record['path']) |
| if not p.is_file() or p.stat().st_size != record['bytes'] or sha256_file(p) != record['sha256']: |
| raise ValueError(f'Modified materialized file: {p}; remove the cache explicitly and recreate it') |
| return output |
| |
| temp = Path(tempfile.mkdtemp(prefix='.' + output.name + '-', dir=output.parent)) |
| try: |
| proof = [] |
| for r in index['files'] + index['assets']: |
| self.check_file(r) |
| for record in index['assets']: |
| src = checked_path(self.root, record['path']) |
| rel = record['restore_path'] |
| dest = checked_path(temp, rel) |
| copied = self._copy_file(src, dest, rel) |
| proof.append({k: copied[k] for k in ['path', 'bytes', 'sha256']}) |
| if index['kind'] == 'policy': |
| dest = checked_path(temp, index['original']['filename']) |
| with dest.open('xb') as f: |
| self._reconstruct(index, f) |
| f.flush() |
| os.fsync(f.fileno()) |
| if sha256_file(dest) != index['original']['sha256']: |
| raise ValueError('Materialized checkpoint readback verification failed') |
| proof.append({'path': dest.name, 'bytes': dest.stat().st_size, 'sha256': index['original']['sha256']}) |
| atomic_json(temp / '.archive-materialized.json', {'id': entry['id'], 'index_sha256': entry['index']['sha256'], 'files': proof}) |
| os.rename(temp, output) |
| finally: |
| if temp.exists(): |
| shutil.rmtree(temp) |
| return output |
|
|
| def verify(self, reconstruct=False, report=None, run=None, stage='sft', workers=1): |
| initial_catalog_hash = sha256_file(self.root / 'catalog.json') |
| catalog = self.catalog() |
| results = [] |
| selected = set(catalog['entries']) |
| if run: |
| selected = {self.resolve(run, stage)['id']} |
| pending = list(selected) |
| while pending: |
| entry = catalog['entries'][pending.pop()] |
| for dep in entry.get('dependencies', {}).values(): |
| if dep not in selected: |
| self.resolve(dep) |
| selected.add(dep) |
| pending.append(dep) |
| work = [] |
| for key, entry in sorted(catalog['entries'].items()): |
| if key != entry['id']: |
| raise ValueError('Catalog key/entry identity mismatch') |
| if key not in selected: |
| continue |
| work.append(entry) |
| def verify_one(entry): |
| result = self.verify_entry(entry, reconstruct=reconstruct) |
| print(json.dumps(result), flush=True, file=sys.stderr) |
| return result |
| if workers < 1: |
| raise ValueError('workers must be positive') |
| with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool: |
| results = list(pool.map(verify_one, work)) |
| for root, dirs, files in os.walk(self.root): |
| dirs[:] = [d for d in dirs if d != '.staging'] |
| for name in dirs + files: |
| p = Path(root) / name |
| if p.is_symlink() or (p.is_file() and p.stat().st_nlink != 1): |
| raise ValueError(f'Archive files must be independent real files: {p}') |
| if name == 'checkpoint.index.json' and p.relative_to(self.root).as_posix() not in { |
| e['index']['path'] for e in catalog['entries'].values() |
| }: |
| raise ValueError(f'Unregistered checkpoint index (interrupted publication): {p}') |
| if sha256_file(self.root / 'catalog.json') != initial_catalog_hash: |
| raise ValueError('Catalog changed during verification; verify the new snapshot again') |
| result = {'status': 'verified', 'schema_version': VERSION, 'time': now(), |
| 'catalog_sha256': initial_catalog_hash, 'reconstructed': reconstruct, |
| 'scope': 'all' if run is None else self.resolve(run, stage)['id'], |
| 'entries': results} |
| if report: |
| atomic_json(report, result) |
| return result |
|
|
|
|
| def cli(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--root', type=Path, default=Path(__file__).resolve().parent.parent) |
| sub = parser.add_subparsers(dest='cmd', required=True) |
| sub.add_parser('init') |
| sub.add_parser('list') |
| p = sub.add_parser('resolve'); p.add_argument('run'); p.add_argument('--stage', default='sft') |
| p = sub.add_parser('add-policy') |
| p.add_argument('--run', required=True); p.add_argument('--name', required=True) |
| p.add_argument('--stage', choices=['sft', 'pretrain'], required=True) |
| p.add_argument('--checkpoint', type=Path, required=True) |
| p.add_argument('--share-frozen-decoder', action='store_true') |
| p.add_argument('--svae'); p.add_argument('--pretrain') |
| p = sub.add_parser('add-model') |
| p.add_argument('--id', required=True); p.add_argument('--name', required=True) |
| p.add_argument('--source-dir', required=True); p.add_argument('--files', nargs='+', required=True) |
| p.add_argument('--extra-file', action='append', default=[], metavar='RELATIVE=SOURCE') |
| p = sub.add_parser('verify'); p.add_argument('--reconstruct', action='store_true'); p.add_argument('--report', type=Path) |
| p.add_argument('--run'); p.add_argument('--stage', default='sft') |
| p.add_argument('--workers', type=int, default=1) |
| p = sub.add_parser('materialize'); p.add_argument('run'); p.add_argument('--stage', default='sft'); p.add_argument('--output', required=True) |
| p = sub.add_parser('eval'); p.add_argument('run'); p.add_argument('--repo', default=os.environ.get('OPENWAM_REPO', '/mnt/data/limingleyang/OpenWAM-3sys')) |
| p.add_argument('--cache-root', default=os.environ.get('OPENWAM_CHECKPOINT_EVAL_CACHE')) |
| p.add_argument('--dry-run', action='store_true') |
| args = parser.parse_args() |
| archive = Archive(args.root) |
| if args.cmd == 'init': |
| archive.init(); return |
| if args.cmd == 'list': |
| for key, entry in sorted(archive.catalog()['entries'].items()): |
| print(key, entry.get('step', '-'), entry['directory'], sep='\t') |
| elif args.cmd == 'resolve': |
| print(json.dumps(archive.resolve(args.run, args.stage), indent=2)) |
| elif args.cmd == 'add-policy': |
| deps = {k: v for k, v in [('svae', args.svae), ('pretrain', args.pretrain)] if v} |
| result = archive.add_policy(run=args.run, name=args.name, stage=args.stage, source=args.checkpoint, |
| share_decoder=args.share_frozen_decoder, dependencies=deps) |
| print(json.dumps(result, indent=2)) |
| elif args.cmd == 'add-model': |
| extras = dict(item.split('=', 1) for item in args.extra_file) |
| print(json.dumps(archive.add_model(key=args.id, name=args.name, source=args.source_dir, files=args.files, extra_assets=extras), indent=2)) |
| elif args.cmd == 'verify': |
| result = archive.verify(args.reconstruct, args.report, args.run, args.stage, args.workers) |
| print(json.dumps({'status': result['status'], 'entries': len(result['entries']), 'reconstructed': args.reconstruct})) |
| elif args.cmd == 'materialize': |
| print(archive.materialize(args.run, args.output, args.stage)) |
| elif args.cmd == 'eval': |
| entry = archive.resolve(args.run) |
| if entry['kind'] != 'policy' or entry['stage'] != 'sft': |
| raise ValueError('LIBERO evaluation requires an SFT policy') |
| repo = Path(args.repo).resolve() |
| script = repo / 'scripts/libero/eval.sh' |
| if not script.is_file(): |
| raise FileNotFoundError(script) |
| cache = Path(args.cache_root) if args.cache_root else archive.root.parent / 'checkpoint_eval_cache' |
| output = cache / entry['run_id'] / entry['index']['sha256'][:16] |
| materialized = archive.materialize(args.run, output) |
| env = os.environ.copy() |
| env.setdefault('EVAL_RUN_NAME', entry['run_name'] + '_step' + str(entry['step'])) |
| if args.dry_run: |
| env['EVAL_DRY_RUN'] = '1' |
| raise SystemExit(subprocess.call(['bash', str(script), str(materialized), entry['source']['filename']], cwd=repo, env=env)) |
|
|
|
|
| if __name__ == '__main__': |
| cli() |
|
|
