code/scripts/restore_exports.py

#!/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()