source/scripts/publish_wrapup_evidence.py

"""Prepare and publish an explicit, reviewed wrap-up evidence file map.

prepare --plan PLAN.json --output NEW_EXPORT creates an immutable local stage.
publish --stage NEW_EXPORT --repo-id andyshu/opensysone --output NEW_STATUS uploads
profiles/<evidence_id>/payload and updates only PROFILE_RESULTS.json after verification.
It never changes the final-model pointer, training snapshot pointer, card or visibility.

Plan format: {"format":"opensysone-wrapup-evidence-plan-v1", "repo_id":"...",
 "evidence_id":"20260917T...", "files":{"reports/report.md":{"source":"/absolute/file"}},
 "completion_checks":[{"path":"profiling/summary.json","json_equals":{"status":"complete"}}],
 "checkpoints":[{"path":"artifacts/gx10/checkpoint.pt","sha256":"...","step":159,
   "source_commit":"40 hex characters","role":"resumable","proof":"proofs/gx10.json"}],
 "provenance":{...}}
Every checkpoint needs its expected hash and a mapped CPU provenance proof. Source
archives include committed HEAD plus all checkpoint source revisions. No model is loaded.
"""
import argparse
import json
import logging
import os
from pathlib import Path
import re
import shutil
import signal
import sys
import tempfile

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from scripts.publish_hf_final import digest, read_json, write_json, utc, deadline_alarm, stop_requested
from scripts.publish_hf_snapshot import relative_name, git_bytes, archive_source

ALLOWED_ROOTS = (ROOT, Path.home() / 'ai/opensysone')
CATEGORIES = {'artifacts', 'validation', 'profiling', 'proofs', 'reports'}
WEIGHTS = {'best.pt', 'checkpoint.pt', 'latest.pt', 'latest.evaluated.pt'}
SECRET_NAMES = {'.env', 'token', 'access_token', 'credentials', 'credentials.json', 'credentials.ini',
                'auth.json', '.netrc', 'id_rsa', 'id_ed25519'}
SECRET_DIRECTORIES = {'.ssh', '.aws', '.gnupg', '.cache', 'secrets'}
SECRET_PATTERN = re.compile(r'(^|[_.-])(credentials?|secrets?|access[_-]?tokens?|api[_-]?keys?|private[_-]?keys?)([_.-]|$)', re.I)
SAFE_SUFFIXES = {'.json', '.jsonl', '.md', '.txt', '.csv', '.log', '.png', '.py', '.sh', '.cjs', '.pt', '.sha256'}


def safe_source(source, roots):
    path = Path(source)
    if not path.is_absolute() or '..' in path.parts or path.resolve() != path or not path.is_file():
        raise ValueError('Evidence source must be an absolute regular file without symlinks')
    if not any(path.is_relative_to(Path(root).resolve()) for root in roots):
        raise ValueError('Evidence source is outside the reviewed project/runtime roots')
    if any(part in SECRET_DIRECTORIES or part == 'models' for part in path.parts):
        raise ValueError('Secret directories and base-model directories are prohibited')
    if (path.name.lower() in SECRET_NAMES or SECRET_PATTERN.search(path.name) or path.name.startswith('.env.') or
            (path.suffix.lower() not in SAFE_SUFFIXES and path.name != 'exit_code') or
            (path.suffix.lower() == '.pt' and path.name not in WEIGHTS)):
        raise ValueError('Secret, base-weight or unsupported evidence file is prohibited')
    return path


def validate_plan(plan, roots):
    if plan.get('format') != 'opensysone-wrapup-evidence-plan-v1':
        raise ValueError('Unknown evidence plan format')
    if not re.fullmatch(r'[a-zA-Z0-9][a-zA-Z0-9_.-]*', plan.get('evidence_id', '')):
        raise ValueError('Invalid immutable evidence identifier')
    if not re.fullmatch(r'[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+', plan.get('repo_id', '')):
        raise ValueError('Invalid repository identifier')
    files = plan.get('files', {})
    if not files or not plan.get('completion_checks'):
        raise ValueError('Explicit file map and completion gates are required')
    for name, record in files.items():
        path = Path(relative_name(name))
        if path.parts[0] not in CATEGORIES or len(path.parts) < 2:
            raise ValueError('Evidence destination is outside the reviewed categories')
        source = safe_source(record['source'], roots)
        if (path.suffix.lower() != source.suffix.lower() or path.name.lower() in SECRET_NAMES or
                (not path.suffix and (path.name != 'exit_code' or source.name != 'exit_code'))):
            raise ValueError('Evidence cannot disguise a prohibited source/file type')
    checkpoints = {record['path']: record for record in plan.get('checkpoints', [])}
    weights = {name for name in files if name.endswith('.pt')}
    if len(checkpoints) != len(plan.get('checkpoints', [])) or set(checkpoints) != weights:
        raise ValueError('Every weight requires exactly one checkpoint provenance record')
    for name, record in checkpoints.items():
        if (Path(name).name not in WEIGHTS or not re.fullmatch('[0-9a-f]{64}', record.get('sha256', '')) or
                not re.fullmatch('[0-9a-f]{40}', record.get('source_commit', '')) or
                type(record.get('step')) is not int or record['step'] < 0 or
                record.get('role') not in ('selected', 'resumable', 'latest_evaluated') or
                record.get('proof') not in files or not record['proof'].endswith('.json')):
            raise ValueError('Checkpoint provenance is incomplete')
    for check in plan['completion_checks']:
        if check.get('path') not in files or (('json_equals' in check) == ('text_equals' in check)):
            raise ValueError('Completion gate must name a mapped file and one explicit expectation')
        path = Path(files[check['path']]['source'])
        if 'json_equals' in check:
            expected = check['json_equals']
            if not expected or any(read_json(path).get(key) != value for key, value in expected.items()):
                raise ValueError('Required evidence is not complete')
        elif path.read_text().strip() != check['text_equals']:
            raise ValueError('Required completion exit/status does not match')
    return checkpoints


def prepare(plan_path, output, source_root=ROOT, roots=ALLOWED_ROOTS):
    plan = read_json(plan_path)
    checkpoints = validate_plan(plan, roots)
    output = Path(output).resolve()
    source_root = Path(source_root).resolve()
    if output.exists() or output.is_relative_to(source_root):
        raise ValueError('Export must be a new directory outside the source checkout')
    head = git_bytes(source_root, 'rev-parse', 'HEAD').decode().strip()
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = Path(tempfile.mkdtemp(prefix=output.name + '.tmp-', dir=output.parent))
    try:
        copied = {}
        for name, record in plan['files'].items():
            source = Path(record['source'])
            before = digest(source)
            expected = checkpoints[name]['sha256'] if name in checkpoints else record.get('sha256')
            if expected is not None and before != expected:
                raise ValueError('Evidence does not match its expected checksum')
            target = temporary / 'payload' / name
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(source, target)
            if digest(source) != before or digest(target) != before:
                raise ValueError('Evidence changed while staging')
            copied['payload/' + name] = {'sha256': before, 'size': target.stat().st_size}
        # Recheck completion against the frozen copied files, closing the copy race.
        frozen_plan = {**plan, 'files': {name: {**record, 'source': str(temporary / 'payload' / name)}
                                      for name, record in plan['files'].items()}}
        validate_plan(frozen_plan, (temporary,))
        revisions = sorted({head} | {record['source_commit'] for record in checkpoints.values()})
        for revision in revisions:
            archive_source(source_root, revision, temporary / 'payload/sources' / revision)
        for path in (temporary / 'payload/sources').rglob('*'):
            if path.is_file():
                copied[str(path.relative_to(temporary))] = {'sha256': digest(path), 'size': path.stat().st_size}
        manifest = {'format': 'opensysone-wrapup-evidence-v1', 'created_utc': utc(), 'repo_id': plan['repo_id'],
            'evidence_id': plan['evidence_id'], 'source_commit': head, 'source_revisions': revisions,
            'checkpoints': plan.get('checkpoints', []), 'completion_checks': plan['completion_checks'],
            'provenance': plan.get('provenance', {}), 'files': copied, 'source_plan_sha256': digest(plan_path),
            'file_count': len(copied), 'payload_bytes': sum(record['size'] for record in copied.values()),
            'model_loaded': False, 'gpu_initialized': False,
            'scope': 'Supplemental stopped-training, validation, profiling and report evidence; final release is separate.'}
        write_json(temporary / 'manifest.json', manifest)
        verify_stage(temporary)
        for path in temporary.rglob('*'):
            if path.is_file():
                path.chmod(0o444)
        temporary.rename(output)
        return manifest
    except BaseException:
        shutil.rmtree(temporary)
        raise


def verify_stage(directory):
    directory = Path(directory).resolve()
    manifest = read_json(directory / 'manifest.json')
    if manifest.get('format') != 'opensysone-wrapup-evidence-v1':
        raise ValueError('Unknown evidence stage format')
    if not re.fullmatch(r'[a-zA-Z0-9][a-zA-Z0-9_.-]*', manifest.get('evidence_id', '')):
        raise ValueError('Invalid immutable evidence identifier')
    actual = {str(path.relative_to(directory)) for path in directory.rglob('*') if path.is_file()}
    if actual != set(manifest['files']) | {'manifest.json'}:
        raise ValueError('Evidence stage inventory changed')
    for name, expected in manifest['files'].items():
        path = directory / relative_name(name)
        if (not name.startswith('payload/') or path.is_symlink() or not path.resolve().is_relative_to(directory) or
                path.stat().st_size != expected['size'] or digest(path) != expected['sha256']):
            raise ValueError('Evidence stage checksum/path mismatch')
    for checkpoint in manifest['checkpoints']:
        if manifest['files']['payload/' + checkpoint['path']]['sha256'] != checkpoint['sha256']:
            raise ValueError('Staged checkpoint disagrees with provenance')
    return manifest


def publish(api, repo_id, directory, progress, operation_factory=None):
    directory = Path(directory)
    manifest = verify_stage(directory)
    if manifest['repo_id'] != repo_id:
        raise ValueError('Evidence target repository mismatch')
    initial = api.repo_info(repo_id, repo_type='model', timeout=30)
    if initial.private is not True:
        raise ValueError('Expected the existing private repository')
    prefix = 'profiles/' + manifest['evidence_id']
    existing = api.get_paths_info(repo_id, [prefix + '/manifest.json'], revision=initial.sha, repo_type='model')
    if existing:
        previous = api.hf_hub_download(repo_id, prefix + '/manifest.json', revision=initial.sha, repo_type='model', etag_timeout=30)
        if Path(previous).read_bytes() != (directory / 'manifest.json').read_bytes():
            raise ValueError('Immutable evidence identifier already belongs to different content')
    progress('uploading_payload', export_directory=str(directory), evidence_id=manifest['evidence_id'])
    uploaded = api.upload_folder(repo_id=repo_id, repo_type='model', folder_path=str(directory), path_in_repo=prefix,
        parent_commit=initial.sha, commit_message='Back up verified OpenSysOne wrap-up and profiling evidence')
    progress('verifying_payload', payload_commit=uploaded.oid)
    files = {**manifest['files'], 'manifest.json': {'size': (directory / 'manifest.json').stat().st_size,
                                                  'sha256': digest(directory / 'manifest.json')}}
    names = [prefix + '/' + name for name in files]
    remote = []
    for start in range(0, len(names), 250):
        remote.extend(api.get_paths_info(repo_id, names[start:start+250], revision=uploaded.oid, repo_type='model'))
    by_name = {record.path: record for record in remote}
    for name, expected in files.items():
        remote_file = by_name.get(prefix + '/' + name)
        if remote_file is None or remote_file.size != expected['size']:
            raise ValueError('Remote evidence size mismatch')
        lfs = getattr(remote_file, 'lfs', None)
        if name.endswith('.pt') and lfs is None:
            raise ValueError('Remote checkpoint lacks LFS hash metadata')
        remote_hash = (lfs.get('sha256') if isinstance(lfs, dict) else lfs.sha256) if lfs else remote_file.blob_id
        local_hash = expected['sha256'] if lfs else digest(directory / name, 'sha1', git_blob=True)
        if remote_hash != local_hash:
            raise ValueError('Remote evidence checksum mismatch')
    downloaded = api.hf_hub_download(repo_id, prefix + '/manifest.json', revision=uploaded.oid, repo_type='model', etag_timeout=30)
    if Path(downloaded).read_bytes() != (directory / 'manifest.json').read_bytes():
        raise ValueError('Remote evidence manifest bytes mismatch')
    current = api.repo_info(repo_id, repo_type='model', timeout=30)
    if current.private is not True:
        raise ValueError('Repository visibility changed')
    pointer = {'format': 'opensysone-profile-results-pointer-v1', 'repo_id': repo_id,
        'evidence_id': manifest['evidence_id'], 'path': prefix, 'manifest_path': prefix + '/manifest.json',
        'manifest_sha256': files['manifest.json']['sha256'], 'payload_commit': uploaded.oid,
        'source_commit': manifest['source_commit'], 'verified_utc': utc()}
    if operation_factory is None:
        from huggingface_hub import CommitOperationAdd
        operation_factory = CommitOperationAdd
    content = (json.dumps(pointer, indent=2) + '\n').encode()
    progress('publishing_evidence_pointer')
    committed = api.create_commit(repo_id, repo_type='model', parent_commit=current.sha,
        operations=[operation_factory(path_in_repo='PROFILE_RESULTS.json', path_or_fileobj=content)],
        commit_message='Point to verified OpenSysOne profiling and wrap-up evidence')
    downloaded = api.hf_hub_download(repo_id, 'PROFILE_RESULTS.json', revision=committed.oid, repo_type='model', etag_timeout=30)
    if Path(downloaded).read_bytes() != content:
        raise ValueError('Published evidence pointer bytes mismatch')
    return {**pointer, 'pointer_commit': committed.oid, 'repository_private': True}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest='command', required=True)
    preparation = commands.add_parser('prepare')
    preparation.add_argument('--plan', required=True)
    preparation.add_argument('--output', required=True)
    publication = commands.add_parser('publish')
    publication.add_argument('--stage', required=True)
    publication.add_argument('--repo-id', required=True)
    publication.add_argument('--output', required=True, help='Fresh publication-only status directory')
    args = parser.parse_args()
    Path('/proc/self/oom_score_adj').write_text('0')
    signal.signal(signal.SIGALRM, deadline_alarm)
    signal.signal(signal.SIGTERM, stop_requested)
    signal.alarm(1800)
    if args.command == 'prepare':
        try:
            result = prepare(args.plan, args.output)
            print(json.dumps({'status': 'prepared', 'output': args.output, 'files': result['file_count'],
                              'payload_bytes': result['payload_bytes'], 'uploaded': False}))
            return 0
        except BaseException as error:
            print(json.dumps({'status': 'failed', 'error_type': type(error).__name__}), file=sys.stderr)
            return 1
        finally:
            signal.alarm(0)
    output = Path(args.output).resolve()
    output.mkdir(parents=True, exist_ok=False)
    state = {'status': 'running', 'pid': os.getpid(), 'command': [sys.executable, *sys.argv],
             'started_utc': utc(), 'repo_id': args.repo_id}
    def progress(stage, **values):
        state.update(stage=stage, heartbeat_utc=utc(), **values)
        write_json(output / 'state.json', state)
    code = 1
    try:
        os.environ.update(HF_HUB_DISABLE_PROGRESS_BARS='1', HF_HUB_DISABLE_XET='1',
                          HF_HUB_DOWNLOAD_TIMEOUT='60', HF_HUB_ETAG_TIMEOUT='30')
        from huggingface_hub import HfApi, set_client_factory
        import httpx
        for name in ('huggingface_hub', 'httpx', 'httpcore'):
            logging.getLogger(name).setLevel(logging.CRITICAL)
        set_client_factory(lambda: httpx.Client(timeout=httpx.Timeout(60, connect=15), follow_redirects=True))
        result = publish(HfApi(), args.repo_id, args.stage, progress)
        progress('complete', status='complete', finished_utc=utc(), **result)
        code = 0
    except BaseException as error:
        progress('failed', status='failed', error_type=type(error).__name__, finished_utc=utc())
    finally:
        signal.alarm(0)
        (output / 'exit_code').write_text(str(code) + '\n')
    return code


if __name__ == '__main__':
    raise SystemExit(main())