source/scripts/verify_expanded_startup.py

"""CPU-only audit of trusted expanded-data pilot and optional campaign startup.

Run after the eight-step pilot exits. Repeat with --campaign after at least four
new updates. Capture the campaign's initial checkpoint before its first periodic
save if a durable exact optimizer/RNG comparison is required. No model is
instantiated and CUDA is never initialized.
"""
import argparse
from collections import Counter
from datetime import datetime, timezone
import hashlib
import importlib
import json
import math
import os
from pathlib import Path
import random
import subprocess
import sys

sys.dont_write_bytecode = True
CAP = 16 * 2**30
IDENTITY_KEYS = ('group', 'family', 'target', 'choices')
STATE_KEYS = ('trainable_state', 'optimizer', 'random_state', 'torch_rng', 'cuda_rng')
RESUME_CONFIG = ('rank', 'alpha', 'adapters', 'max_tokens', 'branch_batch_size',
                 'seed', 'effective_batch', 'lr', 'head_lr', 'schedule_steps',
                 'epochs', 'validation_per_family', 'two_pass', 'selection_metric')


def require(condition, message):
    if not condition:
        raise ValueError(message)


def sha256(path):
    digest = hashlib.sha256()
    with Path(path).open('rb') as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b''):
            digest.update(chunk)
    return digest.hexdigest()


def read_json(path):
    return json.loads(Path(path).read_text())


def prediction_parity(reference, actual):
    """Require every frozen identity and numeric output, independent of row order."""
    require(len(reference) == len(actual) == 512, 'Expected all 512 validation rows')
    left, right = ({row['id']: row for row in rows} for rows in (reference, actual))
    require(len(left) == len(right) == 512 and left.keys() == right.keys(),
            'Validation IDs differ or contain duplicates')
    expected = {family: 128 for family in ('arc', 'banking', 'boolq', 'snli')}
    require(Counter(row['family'] for row in actual) == expected,
            'Validation family counts differ')
    maxima = {key: 0.0 for key in ('logits', 'probabilities', 'log_probabilities')}
    for identifier, a in left.items():
        b = right[identifier]
        require(all(a[key] == b[key] for key in IDENTITY_KEYS),
                'Validation group/family/target/choices differ')
        for key in maxima:
            require(len(a[key]) == len(b[key]) == len(a['choices']),
                    'Prediction vector length differs')
            require(all(isinstance(x, (float, int)) and math.isfinite(x)
                        for x in [*a[key], *b[key]]), 'Non-finite prediction')
            maxima[key] = max(maxima[key], max(abs(x-y) for x, y in zip(a[key], b[key])))
    require(all(value == 0 for value in maxima.values()), 'Prediction replay is not exact')
    return {'rows': 512, 'identity_exact': True, 'max_abs_difference': maxima}


def exact_state(left, right, torch, path='state'):
    if torch.is_tensor(left) or torch.is_tensor(right):
        require(torch.is_tensor(left) and torch.is_tensor(right)
                and left.dtype == right.dtype and left.shape == right.shape
                and torch.equal(left, right), 'State mismatch: ' + path)
    elif isinstance(left, dict):
        require(isinstance(right, dict) and left.keys() == right.keys(),
                'State mapping mismatch: ' + path)
        for key in left:
            exact_state(left[key], right[key], torch, path + '.' + str(key))
    elif isinstance(left, (list, tuple)):
        require(type(left) is type(right) and len(left) == len(right),
                'State sequence mismatch: ' + path)
        for index, (a, b) in enumerate(zip(left, right)):
            exact_state(a, b, torch, path + '.' + str(index))
    else:
        require(type(left) is type(right) and left == right, 'State mismatch: ' + path)


def optimizer_steps(artifact, expected, torch):
    optimizer = artifact['optimizer']
    parameter_ids = [p for group in optimizer['param_groups'] for p in group['params']]
    require(bool(parameter_ids) and len(parameter_ids) == len(set(parameter_ids))
            and set(parameter_ids) == set(optimizer['state']), 'Incomplete Adam state')
    for state in optimizer['state'].values():
        require(float(state['step']) == expected, 'Adam step differs from pilot step')
        for key in ('exp_avg', 'exp_avg_sq'):
            require(torch.is_tensor(state[key]) and bool(torch.isfinite(state[key]).all()),
                    'Non-finite or missing Adam tensor')
    return {'parameter_states': len(parameter_ids), 'all_steps': expected}


def load_checkpoint(path, torch):
    before = sha256(path)
    artifact = torch.load(path, map_location='cpu', weights_only=False)
    require(sha256(path) == before, 'Checkpoint changed during CPU inspection; retry')
    require(artifact['format'] == 'opensysone-adapter-v1', 'Wrong checkpoint format')
    require(bool(artifact['trainable_state']) and all(
        value.device.type == 'cpu' and bool(torch.isfinite(value).all())
        for value in artifact['trainable_state'].values()), 'Invalid trainable tensor')
    return artifact, before


def verify_initial_state(initial, pilot, parent, torch):
    require(initial['step'] == 0 and not initial['optimizer']['state'],
            'Initial warm checkpoint must have empty Adam at step zero')
    for key in ('config', 'initialization', 'data_signature', 'source_commit', 'model_provenance'):
        require(initial[key] == pilot[key], 'Initial warm provenance differs: ' + key)
    exact_state(parent['trainable_state'], initial['trainable_state'], torch)
    require(initial['random_state'] == random.Random(pilot['config']['seed']).getstate(),
            'Initial Python RNG does not match fresh seed')
    require(initial['torch_rng'].numel() > 0 and len(initial['cuda_rng']) == 1,
            'Initial Torch/CUDA RNG state missing')
    return {'step': 0, 'all_parent_weights_exact': True, 'adam_state_empty': True,
            'python_rng_matches_fresh_seed': True, 'torch_cuda_rng_present': True,
            'torch_rng_scope': 'Seed/no-restore source verified; GPU initialization sequence not rerun.'}


def verify_source(manifest, source):
    require(not manifest['git_status'].strip(), 'Execution source was not clean')
    revision = subprocess.check_output(['git', '-C', str(source), 'rev-parse', 'HEAD'],
                                       text=True, timeout=10).strip()
    require(manifest['git_commit'] == revision, 'Use the frozen execution source checkout')
    for filename, expected in manifest['source_sha256'].items():
        require(sha256(source/filename) == expected, 'Execution source file changed: ' + filename)
    require(manifest['cuda_cap_bytes'] == CAP and str(manifest['oom_score_adj']) == '0',
            'Training memory/OOM gate differs')
    return {'commit': revision, 'verified_files': len(manifest['source_sha256'])}


def verify_correctness(directory, final=False):
    names = ['correctness_initial.json'] + (['correctness_final.json'] if final else [])
    for name in names:
        checks = read_json(directory/name)
        values = [v for key, v in checks.items() if key.endswith('max_abs')]
        require(len(values) >= 6 and all(math.isfinite(v) and 0 <= v <= 1e-4 for v in values),
                'Correctness gate failed: ' + name)


def finite_updates(path, start, minimum):
    rows = []
    lines = Path(path).read_text().splitlines()
    for index, line in enumerate(lines):
        try:
            row = json.loads(line)
        except json.JSONDecodeError:
            require(index == len(lines)-1, 'Malformed training log')
            break  # A concurrent writer may be appending the final record.
        if row['step'] > start:
            rows.append(row)
    require(len(rows) >= minimum, 'Waiting for sufficient subsequent optimizer updates')
    require([r['step'] for r in rows] == list(range(start+1, start+1+len(rows))),
            'Non-contiguous optimizer updates')
    require(all(math.isfinite(row[key]) for row in rows for key in ('loss', 'gradient_norm')),
            'Non-finite optimizer update')
    require(all(0 < row['peak_cuda_allocated_bytes'] <= CAP for row in rows),
            'CUDA allocation exceeded 16 GiB')
    return {'count': len(rows), 'first_step': rows[0]['step'], 'last_step': rows[-1]['step'],
            'all_finite': True, 'peak_cuda_allocated_bytes': max(r['peak_cuda_allocated_bytes'] for r in rows)}


def verify(args, torch):
    source = args.source_root.resolve()
    sys.path.insert(0, str(source))
    transition_module = importlib.import_module('data_transition')
    selection_module = importlib.import_module('selection')
    parent, parent_sha = load_checkpoint(args.parent, torch)
    pilot_dir = args.pilot.resolve()
    pilot, pilot_sha = load_checkpoint(pilot_dir/'checkpoint.pt', torch)
    manifest = read_json(pilot_dir/'manifest.json')
    summary = read_json(pilot_dir/'summary.json')
    config = manifest['config']
    require(pilot['config'] == config and pilot['source_commit'] == manifest['git_commit']
            and pilot['data_signature'] == manifest['data_signature'], 'Pilot provenance differs')
    require(pilot['model_provenance'] == manifest['model_provenance'] == parent['model_provenance'],
            'Pinned base model differs')
    for key in ('prompt_version', 'adapter_version'):
        require(pilot[key] == parent[key], 'Parent implementation version differs')
    expected = {'seed': args.seed, 'lr': args.lr, 'head_lr': args.lr,
                'schedule_steps': args.schedule_steps, 'allow_train_data_change': True,
                'rank': 8, 'alpha': 16.0, 'max_tokens': 512, 'effective_batch': 4,
                'branch_batch_size': 1, 'two_pass': True, 'validation_per_family': 128,
                'selection_metric': 'crossfit_temperature_nll_v1'}
    require(all(config.get(k) == v for k, v in expected.items()), 'Unexpected pilot configuration')
    require(pilot['step'] == args.pilot_steps == summary['completed_steps'], 'Pilot step differs')
    require(summary['checkpoint_sha256'] == pilot_sha, 'Pilot final checksum differs')
    require(summary['final_correctness_status'] == 'passed', 'Pilot final correctness incomplete')
    require(0 < summary['peak_cuda_allocated_bytes'] <= CAP, 'Pilot exceeded CUDA allocation cap')
    if (pilot_dir.parent/'exit_code').exists():
        require((pilot_dir.parent/'exit_code').read_text().strip() == '0', 'Pilot exited unsuccessfully')
    initialization = manifest['initialization']
    require(pilot['initialization'] == initialization, 'Pilot initialization metadata differs')
    require(initialization['kind'] == 'warm_start' and initialization['restores_optimizer'] is False
            and initialization['restores_rng'] is False, 'Warm start restored optimizer/RNG')
    require(initialization['parent_checkpoint_sha256'] == parent_sha
            and initialization['parent_step'] == parent['step']
            and initialization['parent_source_commit'] == parent['source_commit'], 'Wrong warm parent')
    transition = transition_module.verify_train_data_transition(parent, config, pilot['data_signature'])
    require(initialization['data_transition'] == transition, 'Recorded data transition differs')
    parent_rows = read_json(args.parent.parent/f"validation_step_{parent['step']:06d}_predictions.json")
    initial_rows = read_json(pilot_dir/'initial_validation_predictions.json')
    initial_parity = prediction_parity(parent_rows, initial_rows)
    selected = selection_module.validation_selection(initial_rows)
    require(selected == read_json(pilot_dir/'initial_validation_selection.json'), 'Initial selector differs')
    require(abs(selected['score']-parent['best_validation_selection_score']) < 1e-12,
            'Parent selection score differs')
    require(pilot['selection_metric'] == selected['metric'], 'Pilot selection policy changed')
    verify_correctness(pilot_dir, final=True)
    report = {'parent': {'path': str(args.parent), 'sha256': parent_sha, 'step': parent['step']},
              'pilot': {'directory': str(pilot_dir), 'sha256': pilot_sha, 'step': pilot['step'],
                        'source': verify_source(manifest, source), 'initial_parity': initial_parity,
                        'fresh_adam': optimizer_steps(pilot, args.pilot_steps, torch),
                        'updates': finite_updates(pilot_dir/'training.jsonl', 0, args.pilot_steps)},
              'data_transition': transition,
              'fresh_optimizer_evidence': 'Initialization flags plus all final Adam counters equal pilot steps; initial resumable checkpoint was overwritten.'}
    if args.pilot_initial_state:
        initial_state, initial_sha = load_checkpoint(args.pilot_initial_state, torch)
        report['pilot']['initial_checkpoint'] = {
            **verify_initial_state(initial_state, pilot, parent, torch),
            'path': str(args.pilot_initial_state), 'sha256': initial_sha}
        report['fresh_optimizer_evidence'] = 'Captured step-zero checkpoint has empty Adam; final Adam counters equal pilot steps.'
    best, _ = load_checkpoint(pilot_dir/'best.pt', torch)
    if best['step'] == 0:
        exact_state(parent['trainable_state'], best['trainable_state'], torch)
        report['pilot']['retained_step0_weights_exact'] = True
    if args.campaign:
        directory = args.campaign.resolve()/'training'
        resumed_manifest = read_json(directory/'manifest.json')
        resumed_init = resumed_manifest['initialization']
        require(resumed_init['kind'] == 'resume' and resumed_init['restores_optimizer'] is True
                and resumed_init['restores_rng'] is True
                and resumed_init['parent_checkpoint_sha256'] == pilot_sha
                and resumed_init['parent_step'] == pilot['step']
                and resumed_init['parent_initialization'] == initialization, 'Wrong resume lineage')
        require(resumed_manifest['model_provenance'] == pilot['model_provenance']
                and resumed_manifest['data_signature'] == pilot['data_signature'], 'Resume provenance differs')
        require(all(resumed_manifest['config'].get(k) == config.get(k) for k in RESUME_CONFIG),
                'Resume configuration differs')
        reference = read_json(pilot_dir/f"validation_step_{pilot['step']:06d}_predictions.json")
        replay = read_json(directory/'resumed_initial_predictions.json')
        campaign = {'directory': str(args.campaign), 'source': verify_source(resumed_manifest, source),
                    'prediction_parity': prediction_parity(reference, replay),
                    'updates': finite_updates(directory/'training.jsonl', pilot['step'], args.minimum_updates)}
        verify_correctness(directory)
        state_path = args.resume_state or directory/'checkpoint.pt'
        resumed_state, state_sha = load_checkpoint(state_path, torch)
        require(resumed_state['data_signature'] == pilot['data_signature']
                and resumed_state['model_provenance'] == pilot['model_provenance']
                and resumed_state['source_commit'] == resumed_manifest['git_commit']
                and resumed_state['config'] == resumed_manifest['config']
                and resumed_state['initialization'] == resumed_init, 'Resume checkpoint provenance differs')
        if resumed_state['step'] == pilot['step']:
            for key in STATE_KEYS:
                exact_state(pilot[key], resumed_state[key], torch, key)
            campaign['initial_state_comparison'] = {'status': 'exact', 'keys': STATE_KEYS,
                                                    'path': str(state_path), 'sha256': state_sha}
        else:
            require(not args.resume_state, 'Explicit resume-state is not the initial pilot step')
            campaign['initial_state_comparison'] = {'status': 'unavailable_overwritten',
                                                    'current_step': resumed_state['step']}
        report['campaign'] = campaign
    require(not torch.cuda.is_initialized(), 'Verifier unexpectedly initialized CUDA')
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    for name in ('parent', 'pilot', 'output'):
        parser.add_argument('--'+name, type=Path, required=True)
    parser.add_argument('--source-root', type=Path, default=Path(__file__).resolve().parents[1],
                        help='Frozen execution checkout; defaults to this repository')
    parser.add_argument('--campaign', type=Path)
    parser.add_argument('--pilot-initial-state', type=Path, help='Captured warm step-zero checkpoint')
    parser.add_argument('--resume-state', type=Path, help='Captured initial campaign checkpoint; requires exact state')
    parser.add_argument('--pilot-steps', type=int, default=8)
    parser.add_argument('--seed', type=int, default=433)
    parser.add_argument('--lr', type=float, default=2e-5)
    parser.add_argument('--schedule-steps', type=int, default=3500)
    parser.add_argument('--minimum-updates', type=int, default=4)
    args = parser.parse_args()
    if args.resume_state and not args.campaign:
        parser.error('--resume-state requires --campaign')
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    Path('/proc/self/oom_score_adj').write_text('0')
    import torch
    torch.set_num_threads(1)
    report = verify(args, torch)
    report.update(status='passed', verified_utc=datetime.now(timezone.utc).isoformat(),
                  verifier_sha256=sha256(__file__), cuda_initialized=False,
                  reserved_predictions_accessed=False, model_instantiated=False)
    with args.output.open('x') as handle:
        json.dump(report, handle, indent=2, allow_nan=False)
        handle.write('\n')
    print(json.dumps({'status': 'passed', 'output': str(args.output)}))


if __name__ == '__main__':
    main()