| |
| """Restore completed private HF export backups without changing remote repositories.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import getpass |
| import json |
| import os |
| import re |
| import shutil |
| import sys |
| import tempfile |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from continue_run import exclusive_export_lock |
| from stage_exports import ( |
| atomic_json, |
| encoded, |
| field, |
| fingerprint_bytes, |
| fingerprint_file, |
| require_no_providers, |
| staging_files, |
| verify_inventory, |
| ) |
|
|
| KINDS = {'merged': 'merged-fp16', 'gguf': 'gguf'} |
| COMPLETE = {'merged': 'MERGE_COMPLETE.json', 'gguf': 'GGUF_COMPLETE.json'} |
| RESERVE = 2 * 1024**3 |
|
|
|
|
| def read_json(path): |
| return json.loads(Path(path).read_text()) |
|
|
|
|
| def member(root, name): |
| if (not isinstance(name, str) or not name or name in {'.', '..'} |
| or '/' in name or '\\' in name or '\0' in name): |
| raise ValueError('Unsafe export filename') |
| path = root / name |
| if path.is_symlink() or path.resolve().parent != root.resolve(): |
| raise ValueError('Export member leaves its target directory') |
| return path |
|
|
|
|
| def validate_fingerprint(value): |
| if (not isinstance(value, dict) or type(value.get('size')) is not int or value['size'] < 0 |
| or not re.fullmatch(r'[0-9a-f]{64}', str(value.get('sha256', ''))) |
| or not re.fullmatch(r'[0-9a-f]{40}', str(value.get('git_blob_sha1', '')))): |
| raise ValueError('Invalid recorded export fingerprint') |
|
|
|
|
| def plan_restore(receipt_path, evidence_root, adapter): |
| """Validate persistent evidence before authentication or large downloads.""" |
| receipt = read_json(receipt_path) |
| if (receipt.get('schema_version') != 1 or receipt.get('status') != 'complete_private' |
| or receipt.get('visibility') != 'private' or set(receipt.get('repos', {})) != set(KINDS)): |
| raise ValueError('A completed private staging receipt for both exports is required') |
| adapter_hash = fingerprint_file(adapter / 'adapter_model.safetensors')['sha256'] |
| if receipt.get('adapter_sha256') != adapter_hash: |
| raise ValueError('Staged exports belong to a different adapter') |
| training = read_json(adapter / 'TRAINING_COMPLETE.json') |
| if (training.get('status') != 'complete' or training.get('target_steps', 0) <= 0 |
| or training.get('global_step') != training.get('target_steps')): |
| raise ValueError('Adapter training is incomplete') |
| source_signature = read_json(adapter / 'run_manifest.json')['signature'] |
| manifests, commits, local_markers = {}, {}, {} |
| for kind, directory in KINDS.items(): |
| row = receipt['repos'][kind] |
| if row.get('status') != 'complete_private' or not re.fullmatch(r'[0-9a-f]{40}', str(row.get('commit', ''))): |
| raise ValueError('Staging receipt lacks a completed immutable commit') |
| repo = row.get('repo_id', '') |
| expected_suffix = '/VibeGame-8B' + ('-GGUF' if kind == 'gguf' else '') |
| if not re.fullmatch(r'[A-Za-z0-9_.-]+/VibeGame-8B(?:-GGUF)?', repo) or not repo.endswith(expected_suffix): |
| raise ValueError('Unexpected staged repository identity') |
| files = row.get('files') |
| if not isinstance(files, dict) or not files or 'merge_progress.json' in files: |
| raise ValueError('Invalid staged file manifest') |
| if kind == 'merged' and 'model.safetensors.index.json' not in files: |
| raise ValueError('Staged merged export lacks its tensor index') |
| for name, value in files.items(): |
| member(evidence_root, name) |
| validate_fingerprint(value) |
| manifest = {'schema_version': 1, 'kind': 'vibegame-private-export-staging', |
| 'visibility_when_staged': 'private', 'evaluation_at_staging': 'pending', |
| 'repo_id': repo, 'artifact_kind': kind, 'adapter_sha256': adapter_hash, 'files': files} |
| if fingerprint_bytes(encoded(manifest))['sha256'] != row.get('manifest_sha256'): |
| raise ValueError('Staging manifest fingerprint does not match its receipt') |
| marker = evidence_root / directory / COMPLETE[kind] |
| if not marker.is_file() or fingerprint_file(marker) != files.get(COMPLETE[kind]): |
| raise ValueError('Drive completion evidence differs from staged export') |
| local_markers[kind] = marker |
| manifests[kind], commits[kind] = manifest, row['commit'] |
| if manifests['merged']['repo_id'].split('/')[0] != manifests['gguf']['repo_id'].split('/')[0]: |
| raise ValueError('Export repositories have different owners') |
| merged = read_json(local_markers['merged']) |
| progress_path = evidence_root / 'merged-fp16/merge_progress.json' |
| progress = read_json(progress_path) |
| signature = progress.get('signature', {}) |
| if (merged.get('adapter_sha256') != adapter_hash or signature.get('adapter_sha256') != adapter_hash |
| or any(signature.get(key) != merged.get(key) for key in ('base', 'revision', 'merge')) |
| or merged.get('base') != source_signature.get('original_base') |
| or merged.get('revision') != source_signature.get('original_revision')): |
| raise ValueError('Drive merge evidence belongs to a different adapter or base') |
| tensors = progress.get('tensors', {}) |
| if not tensors or len(tensors) != merged.get('tensor_count'): |
| raise ValueError('Drive merge progress is incomplete') |
| for row in tensors.values(): |
| recorded = manifests['merged']['files'].get(row.get('file'), {}) |
| if (row.get('sha256') != recorded.get('sha256') or type(row.get('tensor_bytes')) is not int |
| or not 0 < row['tensor_bytes'] <= recorded.get('size', 0)): |
| raise ValueError('Drive tensor evidence differs from staged shard fingerprints') |
| if sum(row['tensor_bytes'] for row in tensors.values()) != merged.get('total_size'): |
| raise ValueError('Drive tensor byte count differs from completed merge') |
| gguf = read_json(local_markers['gguf']) |
| quantized = manifests['gguf']['files'].get('VibeGame-8B-Q4_K_M.gguf', {}) |
| if (gguf.get('source_merge') != merged or gguf.get('quantization') != 'Q4_K_M' |
| or gguf.get('sha256') != quantized.get('sha256') or gguf.get('bytes') != quantized.get('size') |
| or 'Modelfile' not in manifests['gguf']['files']): |
| raise ValueError('GGUF receipt does not match the same completed merge') |
| return {'manifests': manifests, 'commits': commits, 'markers': local_markers, |
| 'progress': progress_path, 'adapter_sha256': adapter_hash, |
| 'receipt_sha256': fingerprint_file(receipt_path)['sha256'], |
| 'receipt_path': Path(receipt_path), 'adapter_path': Path(adapter), |
| 'evidence_root': Path(evidence_root)} |
|
|
|
|
| def check_remote(api, plan): |
| |
| |
| for kind, manifest in plan['manifests'].items(): |
| commit = plan['commits'][kind] |
| info = api.repo_info(repo_id=manifest['repo_id'], repo_type='model', revision=commit, |
| files_metadata=True, timeout=30) |
| if field(info, 'sha') != commit: |
| raise ValueError('Hugging Face returned a different export commit') |
| verify_inventory(info, staging_files(manifest)) |
| require_no_providers(api, manifest['repo_id']) |
|
|
|
|
| def checked_existing(path, expected): |
| if not path.exists(): |
| return False |
| if not path.is_file() or fingerprint_file(path) != expected: |
| raise ValueError('Existing local export differs from the staged artifact: ' + path.name) |
| return True |
|
|
|
|
| def remaining_space(plan, output): |
| missing = [] |
| for kind, manifest in plan['manifests'].items(): |
| root = output / KINDS[kind] |
| for name, value in staging_files(manifest).items(): |
| if not checked_existing(member(root, name), value): |
| missing.append(value['size']) |
| progress = output / 'merged-fp16/merge_progress.json' |
| if not checked_existing(member(progress.parent, progress.name), fingerprint_file(plan['progress'])): |
| missing.append(plan['progress'].stat().st_size) |
| |
| |
| return sum(missing) + max(missing, default=0) + RESERVE |
|
|
|
|
| def copy_evidence(source, destination, expected): |
| if checked_existing(destination, expected): |
| return |
| |
| |
| with source.open('rb') as incoming, tempfile.NamedTemporaryFile( |
| dir=destination.parent, prefix=destination.name + '.restore-partial-', delete=False) as outgoing: |
| temporary = Path(outgoing.name) |
| shutil.copyfileobj(incoming, outgoing, length=8 * 1024 * 1024) |
| outgoing.flush() |
| os.fsync(outgoing.fileno()) |
| if fingerprint_file(temporary) != expected: |
| raise ValueError('Copied Drive evidence failed verification') |
| temporary.replace(destination) |
|
|
|
|
| def restore(api, download, plan, output, report_path, *, resources=None): |
| """Read-only Hub calls, one streamed download at a time, bounded host memory.""" |
| if resources is None: |
| import psutil |
| resources = lambda path: (shutil.disk_usage(path).free, psutil.virtual_memory().available) |
| resolved_report = report_path.resolve() |
| if (resolved_report == plan['receipt_path'].resolve() |
| or resolved_report.is_relative_to(plan['adapter_path'].resolve()) |
| or resolved_report.is_relative_to(plan['evidence_root'].resolve()) |
| or resolved_report.is_relative_to(output.resolve())): |
| raise ValueError('Recovery report must not replace receipt, evidence, adapter, or export artifacts') |
| if output.is_symlink() or report_path.is_symlink(): |
| raise ValueError('Restore destination or report is a symlink') |
| output.mkdir(parents=True, exist_ok=True) |
| for directory in KINDS.values(): |
| root = output / directory |
| if root.is_symlink() or (root / '.cache').is_symlink(): |
| raise ValueError('Export destination directory is a symlink') |
| root.mkdir(exist_ok=True) |
| report = {'schema_version': 1, 'status': 'preflight', 'receipt_sha256': plan['receipt_sha256'], |
| 'adapter_sha256': plan['adapter_sha256'], 'repos': { |
| kind: {'repo_id': manifest['repo_id'], 'commit': plan['commits'][kind]} |
| for kind, manifest in plan['manifests'].items()}, |
| 'verified_files': {}, 'remote_mutations': 0, 'inference_calls': 0} |
| atomic_json(report_path, report) |
| try: |
| need = remaining_space(plan, output) |
| free_disk, free_ram = resources(output) |
| if free_disk < need or free_ram < RESERVE: |
| raise RuntimeError('Insufficient disk or host memory to restore exports') |
| check_remote(api, plan) |
| for kind, manifest in plan['manifests'].items(): |
| root = output / KINDS[kind] |
| for name, expected in sorted(staging_files(manifest).items()): |
| destination = member(root, name) |
| if name == COMPLETE[kind]: |
| continue |
| if not checked_existing(destination, expected): |
| free_disk, free_ram = resources(output) |
| if free_disk < expected['size'] + RESERVE or free_ram < RESERVE: |
| raise RuntimeError('Insufficient resources before the next export download') |
| fetched = Path(download(repo_id=manifest['repo_id'], repo_type='model', filename=name, |
| revision=plan['commits'][kind], local_dir=str(root))) |
| if fetched.resolve() != destination.resolve() or destination.is_symlink(): |
| raise ValueError('Download did not land in the expected artifact directory') |
| if not checked_existing(destination, expected): |
| raise ValueError('Download is missing') |
| report['verified_files'][kind + '/' + name] = expected['sha256'] |
| report['status'] = 'restoring' |
| atomic_json(report_path, report) |
| merged_root = output / 'merged-fp16' |
| index = read_json(merged_root / 'model.safetensors.index.json') |
| progress = read_json(plan['progress']) |
| expected_map = {name: row['file'] for name, row in progress['tensors'].items()} |
| if (index.get('weight_map') != expected_map |
| or index.get('metadata', {}).get('total_size') != read_json(plan['markers']['merged'])['total_size']): |
| raise ValueError('Restored tensor index differs from Drive merge evidence') |
| with (output / 'gguf/VibeGame-8B-Q4_K_M.gguf').open('rb') as handle: |
| if handle.read(4) != b'GGUF': |
| raise ValueError('Restored quantization has no GGUF header') |
| check_remote(api, plan) |
| copy_evidence(plan['progress'], member(merged_root, 'merge_progress.json'), fingerprint_file(plan['progress'])) |
| for kind, directory in KINDS.items(): |
| name = COMPLETE[kind] |
| copy_evidence(plan['markers'][kind], member(output / directory, name), |
| plan['manifests'][kind]['files'][name]) |
| report['verified_files'][kind + '/' + name] = plan['manifests'][kind]['files'][name]['sha256'] |
| report['status'] = 'complete' |
| atomic_json(report_path, report) |
| return report |
| except Exception as error: |
| report.update(status='incomplete', failure_type=type(error).__name__) |
| atomic_json(report_path, report) |
| raise RuntimeError('Export restore stopped; verified files and a credential-free report were preserved.') from None |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| for name in ('receipt', 'evidence', 'adapter'): |
| parser.add_argument('--' + name, type=Path, required=True) |
| parser.add_argument('--output', type=Path, default=Path('/content/vibegame-exports')) |
| parser.add_argument('--report', type=Path) |
| args = parser.parse_args() |
| args.report = args.report or args.receipt.with_name('export-restore.json') |
| drive = Path('/content/drive/MyDrive') |
| if not drive.is_dir() or any(not path.resolve().is_relative_to(drive.resolve()) |
| for path in (args.receipt, args.evidence, args.adapter, args.report)): |
| parser.error('Receipt, original evidence, adapter, and recovery report must stay on mounted Drive') |
| if args.output.resolve().is_relative_to(drive.resolve()): |
| parser.error('Restore the large exports to VM disk, not the free 15GB Drive') |
| token = None |
| try: |
| plan = plan_restore(args.receipt, args.evidence, args.adapter) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| with exclusive_export_lock(Path('/content/.vibegame-export.lock')): |
| |
| |
| os.environ['HF_HUB_DISABLE_XET'] = '1' |
| os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '0' |
| from huggingface_hub import HfApi, hf_hub_download |
| token = os.environ.get('HF_TOKEN') or getpass.getpass('Hugging Face token for private backup access (not saved): ') |
| if not token: |
| raise ValueError('Private backup access needs a token') |
| api = HfApi(endpoint='https://huggingface.co', token=token) |
| def download(**kwargs): |
| return hf_hub_download(**kwargs, endpoint='https://huggingface.co', token=token) |
| report = restore(api, download, plan, args.output, args.report) |
| print(json.dumps({'status': report['status'], 'output': str(args.output), 'remote_mutations': 0})) |
| except Exception as error: |
| raise RuntimeError('Private export recovery stopped (' + type(error).__name__ + '); no credential details were logged.') from None |
| finally: |
| token = None |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
