code/scripts/evaluate_pro.py

#!/usr/bin/env python3
"""Opt in to a separate, budget-checked A100 evaluation; never alter the T4 run."""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import signal
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'src'))
sys.path.insert(0, str(ROOT / 'scripts'))
from generate import directory_fingerprint, file_hash

from vibegame import BASE_MODEL, SYSTEM_PROMPT

NAME = 'evaluation-a100-v1'
BATCH = 8
MAX_TOKENS = 3072
COUNT = 500
WORKERS = 4
GENERATION_MARGIN = 1.5
VALIDATION_SECONDS_PER_PROMPT = 90  # 60s gameplay plus 30s import/startup/I/O allowance.
SETUP_SECONDS = 1800  # Reload, sandbox, warmup, and scheduling allowance.
UNIT_RESERVE_FRACTION = 0.1
RATIONALE = ('Predeclared separate full A100 pass to reduce completion time after the free T4 quota was exhausted. '
             'The completed adapter and decoding are unchanged. The A100 500-prompt pass is the primary final '
             'protocol; the partial T4 run remains separate and will be disclosed. No prompt-level results '
             'are pooled or selected between runs. Capacity decisions use only admitted training prompts.')


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


def write_new(path, value):
    with path.open('x') as handle:
        json.dump(value, handle, sort_keys=True, indent=2)
        handle.write('\n')
        handle.flush()
        os.fsync(handle.fileno())


def next_path(output, prefix):
    for number in range(10000):
        path = output / f'{prefix}-{number:04d}.json'
        if not path.exists():
            return path
    raise RuntimeError('Too many preserved attempts; inspect this evaluation directory')


def runtime_snapshot():
    # A separate process leaves no model or CUDA context in this supervisor.
    probe = '''import torch,importlib.metadata as m,json
names=['unsloth','transformers','peft','bitsandbytes']
print(json.dumps({'cuda_available':torch.cuda.is_available(),'device_count':torch.cuda.device_count(),
'gpu':torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
'torch_version':torch.__version__,'cuda_version':torch.version.cuda,
'packages':{name:m.version(name) for name in names}}))'''
    return json.loads(subprocess.check_output([sys.executable, '-c', probe], text=True))


def validate_sources(run, data, output, godot, runtime):
    if output.is_symlink() or output.resolve() != run.resolve() / NAME:
        raise ValueError('A100 evaluation must use the separate RUN/evaluation-a100-v1 directory')
    if not runtime.get('cuda_available') or runtime.get('device_count') != 1 or 'A100' not in str(runtime.get('gpu')):
        raise ValueError('This entry point requires one already allocated A100 CUDA GPU')
    adapter = run / 'adapter'
    completed = read_json(adapter / 'TRAINING_COMPLETE.json')
    if (completed.get('status') != 'complete' or completed.get('target_steps', 0) <= 0
            or completed.get('global_step') != completed.get('target_steps')):
        raise ValueError('A100 evaluation requires the completed training adapter')
    saved = read_json(adapter / 'run_manifest.json')['signature']
    if read_json(run / 'run_manifest.json')['signature'] != saved:
        raise ValueError('Adapter and original training manifest disagree')
    for key, filename in [('train_sha256', 'train.jsonl'), ('dev_sha256', 'dev.jsonl'),
                          ('blind_sha256', 'eval_prompts.jsonl'), ('validation_manifest_sha256', 'validation_manifest.jsonl')]:
        if file_hash(data / filename) != saved.get(key):
            raise ValueError('Original frozen data changed: ' + filename)
    system_hash = hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()
    if saved.get('base_model') != BASE_MODEL or saved.get('system_prompt_sha256') != system_hash:
        raise ValueError('Model or system prompt differs from the completed training run')
    prior_path = run / 'evaluation/manifest.json'
    prior = read_json(prior_path)
    generation = prior['generation']
    if (prior.get('expected_count') != COUNT or prior.get('max_new_tokens') != MAX_TOKENS
            or prior.get('run_seconds') != 60 or prior.get('sample_indices') != [0, 200, 400]
            or prior.get('smoke_test') is not False or prior.get('batch_size') != BATCH
            or prior.get('workers') != WORKERS or prior.get('decoding') != {'temperature': 0, 'seed': 42}
            or prior.get('latency_version') != 2 or 'T4' not in generation.get('gpu', '')
            or prior.get('prompts_sha256') != saved['blind_sha256']
            or prior.get('runtime_sha256') != file_hash(godot)):
        raise ValueError('The original T4 evaluation does not match the fixed protocol')
    if (generation.get('backend') != 'unsloth' or generation.get('base_model') != saved['base_model']
            or generation.get('base_revision') != saved['base_revision']
            or generation.get('system_prompt_sha256') != system_hash
            or generation.get('adapter') != directory_fingerprint(adapter)):
        raise ValueError('Completed adapter/tokenizer identity differs from the original T4 run')
    for key in ('torch_version', 'cuda_version', 'packages'):
        if runtime.get(key) != generation.get(key):
            raise ValueError('Restore the original evaluation runtime before A100 evaluation: ' + key)
    return {'schema_version': 1, 'experiment': NAME, 'rationale': RATIONALE,
            'hardware_tier': 'Colab Pro A100; training previously completed on a free T4',
            'source_t4_manifest_sha256': file_hash(prior_path), 'source_t4_directory': str(prior_path.parent.resolve()),
            'adapter': generation['adapter'], 'base_model': saved['base_model'], 'base_revision': saved['base_revision'],
            'system_prompt_sha256': system_hash, 'train_sha256': saved['train_sha256'],
            'prompts_sha256': saved['blind_sha256'], 'godot_sha256': prior['runtime_sha256'],
            'gpu_runtime': runtime, 'expected_count': COUNT, 'batch_size': BATCH, 'workers': WORKERS,
            'max_new_tokens': MAX_TOKENS, 'run_seconds': 60, 'sample_indices': [0, 200, 400],
            'decoding': prior['decoding'], 'dtype': 'float16', 'base_quantization': 'unchanged pinned 4-bit base',
            't4_evidence_policy': 'Read only; preserve every result and interrupted attempt; never pool scores',
            'source_sha256': {name: file_hash(ROOT / 'scripts' / name) for name in
                              ('generate.py', 'evaluate.py', 'benchmark_batch.py', 'evaluate_pro.py')},
            'budget_policy': {'generation_margin': GENERATION_MARGIN, 'validation_seconds_per_prompt': VALIDATION_SECONDS_PER_PROMPT,
                              'setup_seconds': SETUP_SECONDS, 'balance_reserve_fraction': UNIT_RESERVE_FRACTION,
                              'overlap_credit': False,
                              'limitations': 'Conservative estimate, not a guaranteed upper bound or a billing cap; throughput, import time and Colab rates can vary.'}}


def declare_protocol(output, protocol):
    path = output / 'pro-protocol.json'
    if path.exists():
        if read_json(path).get('protocol') != protocol:
            raise ValueError('A100 protocol changed; original predeclaration and evidence are preserved')
    else:
        if output.exists() and any(output.iterdir()):
            raise ValueError('Refusing to adopt existing results without the A100 predeclaration')
        output.mkdir(parents=True, exist_ok=True)
        write_new(path, {'declared_utc': datetime.now(timezone.utc).isoformat(), 'protocol': protocol})


def capacity_valid(report, protocol, prior_generation, training_prompts):
    expected_prompt_hash = hashlib.sha256(json.dumps(training_prompts, ensure_ascii=False).encode()).hexdigest()
    generation = report.get('generation', {})
    if (report.get('status') != 'complete' or report.get('evaluation_profile') != 'pro-a100-opt-in'
            or report.get('blind_prompts_read') is not False or report.get('batch_sizes') != [BATCH]
            or report.get('recommended_batch_size') != BATCH or report.get('full_max_new_tokens') != MAX_TOKENS
            or report.get('memory_estimate_policy') != 'kv_full_context_v1'
            or report.get('training_jsonl_sha256') != protocol['train_sha256']
            or report.get('training_prompts_sha256') != expected_prompt_hash
            or report.get('gpu') != protocol['gpu_runtime']['gpu']):
        return False
    for key in ('backend', 'adapter', 'base_model', 'base_revision', 'torch_version', 'cuda_version', 'packages',
                'system_prompt_sha256', 'base_artifacts'):
        if key not in prior_generation or generation.get(key) != prior_generation[key]:
            return False
    if generation.get('gpu') != protocol['gpu_runtime']['gpu']:
        return False
    full = report.get('full_check') or {}
    short = next((row for row in report.get('measurements', []) if row.get('batch_size') == BATCH), {})
    for row in (short, full):
        estimate = row.get('full_context_memory_estimate') or {}
        if (row.get('status') != 'success' or row.get('batch_size') != BATCH
                or estimate.get('within_limit') is not True or estimate.get('assessed_batch_size') != BATCH
                or estimate.get('target_context_tokens') != 4096 or estimate.get('cache_growth_allowance_tokens', 0) < 512
                or estimate.get('reserve_bytes', 0) < 2_000_000_000
                or not 0 < estimate.get('estimated_peak_bytes', 0) < estimate.get('effective_limit_bytes', 0)
                or estimate.get('effective_limit_bytes', 0) > min(12_500_000_000, report.get('gpu_memory_bytes', 0) - 2_000_000_000)
                or not 0 < row.get('peak_allocated_bytes', 0) < 12_500_000_000
                or not 0 < row.get('peak_reserved_bytes', 0) < 12_500_000_000):
            return False
    return full.get('max_new_tokens') == MAX_TOKENS and full.get('within_memory_limit') is True


def budget_projection(full, rate, available, elapsed, fresh_prompts=COUNT, pending_results=COUNT):
    if any(not math.isfinite(x) or x <= 0 for x in (rate, available)) or not math.isfinite(elapsed) or elapsed < 0:
        raise ValueError('Observed compute-unit rate and available balance must be finite and positive')
    tokens = full.get('per_prompt_tokens', [])
    seconds = full.get('end_to_end_call_seconds')
    if (len(tokens) != BATCH or any(type(x) is not int or not 0 < x <= MAX_TOKENS for x in tokens)
            or not isinstance(seconds, (int, float)) or not math.isfinite(seconds) or seconds <= 0):
        raise ValueError('Budget requires a valid measured full-cap batch')
    if not 0 <= fresh_prompts <= COUNT or not 0 <= pending_results <= COUNT:
        raise ValueError('Invalid remaining evaluation count')
    normalized = seconds * MAX_TOKENS / max(tokens)
    generation = math.ceil(fresh_prompts / BATCH) * normalized * GENERATION_MARGIN
    validation = math.ceil(pending_results / WORKERS) * VALIDATION_SECONDS_PER_PROMPT
    projected = generation + validation + SETUP_SECONDS
    spent = elapsed * rate / 3600
    reserve = available * UNIT_RESERVE_FRACTION
    usable = available - spent - reserve
    units = projected * rate / 3600
    return {'observed_compute_units_per_hour': rate, 'observed_available_units_at_start': available,
            'elapsed_session_seconds': elapsed, 'estimated_units_used_since_start': spent, 'reserved_units': reserve,
            'fresh_prompts': fresh_prompts, 'pending_results': pending_results,
            'measured_batch_seconds': seconds, 'normalized_3072_batch_seconds': normalized,
            'generation_seconds_with_margin': generation, 'validation_seconds_without_overlap_credit': validation,
            'setup_seconds': SETUP_SECONDS, 'projected_remaining_seconds': projected,
            'projected_remaining_units': units, 'usable_remaining_units': usable,
            'fits_observed_balance': units <= usable,
            'budget_deadline_remaining_seconds': max(0, usable * 3600 / rate)}


def run_budgeted(command, seconds):
    if seconds <= 0:
        raise RuntimeError('No observed compute-unit budget remains')
    process = subprocess.Popen([str(x) for x in command], cwd=ROOT, start_new_session=True)
    try:
        code = process.wait(timeout=seconds)
    except subprocess.TimeoutExpired:
        os.killpg(process.pid, signal.SIGTERM)
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            os.killpg(process.pid, signal.SIGKILL)
            process.wait()
        raise RuntimeError('Stopped at the observed compute-unit time budget; immutable evidence was retained') from None
    if code:
        raise subprocess.CalledProcessError(code, command)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    for name in ('run', 'data', 'godot'):
        parser.add_argument('--' + name, type=Path, required=True)
    parser.add_argument('--output', type=Path)
    parser.add_argument('--compute-units-per-hour', type=float, required=True, help='Current A100 rate observed in Colab')
    parser.add_argument('--available-compute-units', type=float, required=True, help='Current remaining balance observed immediately before this invocation')
    args = parser.parse_args()
    started = time.monotonic()
    rate, balance = args.compute_units_per_hour, args.available_compute_units
    if any(not math.isfinite(x) or x <= 0 for x in (rate, balance)):
        parser.error('Both observed compute-unit values must be finite and positive')
    args.run, args.data, args.godot = (path.expanduser().resolve(strict=True) for path in (args.run, args.data, args.godot))
    drive = Path('/content/drive/MyDrive')
    if not drive.is_dir() or not args.run.is_relative_to(drive.resolve()) or not args.data.is_relative_to(drive.resolve()):
        parser.error('Run, original data and A100 evidence must persist on mounted Drive')
    output = args.output or args.run / NAME
    runtime = runtime_snapshot()
    protocol = validate_sources(args.run, args.data, output, args.godot, runtime)
    import fcntl
    lock_name = hashlib.sha256(str(output.resolve()).encode()).hexdigest()
    with (Path('/content') / f'.vibegame-pro-{lock_name}.lock').open('a+') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        declare_protocol(output, protocol)
        prior_generation = read_json(args.run / 'evaluation/manifest.json')['generation']
        from benchmark_batch import capacity_prompts
        training_prompts = [row['instruction'] for row in capacity_prompts(args.run / 'adapter', args.data / 'train.jsonl')]
        capacity_path, capacity = None, None
        for path in sorted(output.glob('capacity-*.json')):
            try:
                candidate = read_json(path)
            except (ValueError, OSError):
                continue
            if capacity_valid(candidate, protocol, prior_generation, training_prompts):
                capacity_path, capacity = path, candidate
                break
        if capacity is None:
            capacity_path = next_path(output, 'capacity')
            remaining = balance * (1 - UNIT_RESERVE_FRACTION) * 3600 / rate - (time.monotonic() - started)
            run_budgeted([sys.executable, ROOT / 'scripts/benchmark_batch.py', capacity_path,
                          '--adapter', args.run / 'adapter', '--batch-sizes', BATCH, '--pro-a100',
                          '--training-jsonl', args.data / 'train.jsonl'], remaining)
            capacity = read_json(capacity_path)
            if not capacity_valid(capacity, protocol, prior_generation, training_prompts):
                raise RuntimeError('A100 batch8 capacity or exact model/runtime identity check failed; no blind attempts started')
        items = [output / 'items' / f'{index:04d}' for index in range(COUNT)]
        fresh = sum(not (item / 'attempt.json').exists() for item in items)
        pending = sum(not (item / 'result.json').exists() for item in items)
        projection = budget_projection(capacity['full_check'], rate, balance, time.monotonic() - started, fresh, pending)
        budget_path = next_path(output, 'budget')
        write_new(budget_path, {'observed_utc': datetime.now(timezone.utc).isoformat(),
                               'capacity_report_sha256': file_hash(capacity_path), **projection})
        if not projection['fits_observed_balance']:
            raise RuntimeError('Conservative A100 evaluation projection exceeds the observed balance; no blind attempts started. See ' + str(budget_path))
        command = [sys.executable, ROOT / 'scripts/evaluate.py', '--prompts', args.data / 'eval_prompts.jsonl',
                   '--output', output, '--backend', 'unsloth', '--adapter', args.run / 'adapter',
                   '--godot', args.godot, '--run-seconds', 60, '--expected-count', COUNT,
                   '--max-new-tokens', MAX_TOKENS, '--batch-size', BATCH, '--workers', WORKERS]
        if prior_generation.get('requested_base_model') is not None:
            command += ['--base-model', prior_generation['requested_base_model']]
        run_budgeted(command, projection['budget_deadline_remaining_seconds'])


if __name__ == '__main__':
    main()