#!/usr/bin/env python3 """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): # Both sources must pass before downloading the first artifact. The pinned # commit, not a later branch head, defines every expected file and fingerprint. 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) # local_dir avoids a second full HF cache copy. Keep one largest-file margin # for an interrupted transfer plus 2GiB for the runtime and small metadata. return sum(missing) + max(missing, default=0) + RESERVE def copy_evidence(source, destination, expected): if checked_existing(destination, expected): return # A fresh exclusive temporary name preserves interrupted copies and permits # another attempt without overwriting either one or the original evidence. 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 # Commit completion markers only after all files verify. 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: # noqa: BLE001 - network errors can contain credential details. 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')): # Use Hub's sequential HTTP stream instead of optional parallel # native transfer backends; only this helper process is affected. 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: # noqa: BLE001 - do not expose authentication exception payloads. raise RuntimeError('Private export recovery stopped (' + type(error).__name__ + '); no credential details were logged.') from None finally: token = None if __name__ == '__main__': main()