| """Check a downloaded release without training or evaluating a model.""" |
| from pathlib import Path |
| import argparse |
| import ast |
| import hashlib |
| import importlib |
| import json |
| import sys |
| import numpy as np |
|
|
|
|
| def sha256(path): |
| h = hashlib.sha256() |
| with path.open('rb') as f: |
| for block in iter(lambda: f.read(4194304), b''): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--package', type=Path, default=Path(__file__).resolve().parent) |
| p.add_argument('--models', action='store_true', help='Also load every checkpoint on CPU and check its interface and tensors') |
| p.add_argument('--output', type=Path, help='Optional external JSON receipt') |
| a = p.parse_args() |
| root = a.package.resolve() |
| manifest = json.loads((root / 'MANIFEST.json').read_text()) |
| errors = [] |
| for row in manifest['files']: |
| path = root / row['path'] |
| if root not in path.resolve().parents: |
| errors.append('Path escapes release: ' + row['path']); continue |
| if not path.is_file(): |
| errors.append('Missing: ' + row['path']); continue |
| if path.stat().st_size != row['bytes'] or sha256(path) != row['sha256']: |
| errors.append('Release content differs: ' + row['path']) |
| if path.suffix == '.py': |
| try: |
| ast.parse(path.read_text()) |
| except SyntaxError as error: |
| errors.append(f'{row["path"]}: {error}') |
| data = json.loads((root / 'data/manifest.json').read_text()) |
| statuses = {} |
| for row in data['datasets']: |
| statuses[row['status']] = statuses.get(row['status'], 0) + 1 |
| with np.load(root / row['path'], allow_pickle=False) as arrays: |
| schema = {key: {'shape': list(arrays[key].shape), 'dtype': str(arrays[key].dtype)} for key in arrays.files if key != 'metadata'} |
| if schema != row['arrays']: |
| errors.append('Dataset schema differs: ' + row['path']) |
| for key in arrays.files: |
| value = arrays[key] |
| if np.issubdtype(value.dtype, np.number) and not np.isfinite(value).all(): |
| errors.append('Non-finite dataset array: ' + row['path'] + ':' + key) |
| model_count = 0 |
| if a.models: |
| import torch |
| torch.set_num_threads(1) |
| sys.path.insert(0, str(root)) |
| for row in manifest['files']: |
| path = root / row['path'] |
| if path.suffix != '.pt': |
| continue |
| method = Path(row['path']).parts[1] |
| config = importlib.import_module(f'experiments.{method}.src.config').default_config() |
| model = importlib.import_module(f'experiments.{method}.src.model').build_model(config) |
| state = torch.load(path, map_location='cpu', weights_only=True) |
| try: |
| model.load_state_dict(state, strict=True) |
| except RuntimeError as error: |
| errors.append(f'Checkpoint interface: {row["path"]}: {error}') |
| if not all(torch.isfinite(t).all() for t in state.values()): |
| errors.append('Non-finite checkpoint: ' + row['path']) |
| model_count += 1 |
| result = {'passed': not errors, 'files_checked': len(manifest['files']), |
| 'datasets_checked': len(data['datasets']), 'dataset_provenance_counts': statuses, |
| 'checkpoints_loaded_on_cpu': model_count, 'models_executed': False, |
| 'release_integrity_only': True, 'errors': errors} |
| if a.output: |
| output = a.output.resolve() |
| supplied = {str((root / r['path']).resolve()) for r in manifest['files']} |
| if str(output) in supplied or output == root / 'MANIFEST.json': |
| raise ValueError('Receipt must not overwrite a supplied file') |
| output.parent.mkdir(parents=True, exist_ok=True) |
| output.write_text(json.dumps(result, indent=2) + '\n') |
| print(json.dumps(result, indent=2)) |
| if errors: |
| raise SystemExit(1) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
