source/scripts/verify_artifact.py

"""Verify a trusted trained checkpoint, longest training inputs and real HTTP.

No optimizer step is taken and the source checkpoint is never changed. Optimizer
state is restored so the stress pass measures the memory needed during training.
"""
import argparse
from datetime import datetime, timezone
import fcntl
import json
import os
from pathlib import Path
import sys
import threading
import time

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch
from experiment import (data_for, decision_backward, guard_memory, load_artifact,
                        predict, sha256, write_json)
from jev_harness import LocalBackend, RemoteBackend, compile_request, create_server


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--checkpoint', required=True)
    parser.add_argument('--reference', required=True)
    parser.add_argument('--output', required=True)
    parser.add_argument('--inference-max-tokens',type=int,default=1024)
    args = parser.parse_args()
    lock_path = Path.home()/'ai/opensysone/runs/.smoke.lock'
    lock = lock_path.open('w')
    try:
        fcntl.flock(lock,fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        raise RuntimeError('Another project run holds the model-load lock') from None
    out = Path(args.output).resolve()
    out.mkdir(parents=True, exist_ok=False)
    before = sha256(args.checkpoint)
    write_json(out/'manifest.json',{'pid':os.getpid(),'checkpoint_sha256':before,
                                   'started_utc':datetime.now(timezone.utc).isoformat(),
                                   'verification_source_sha256':sha256(__file__)})
    guard_memory()
    scorer, artifact = load_artifact(args.checkpoint)
    if not 256 <= args.inference_max_tokens <= scorer.lm.config.max_position_embeddings:
        raise ValueError('Verification inference limit must be 256 to the base context limit')
    config = artifact['config']
    data, signature = data_for(scorer, config['dataset'], out)
    if signature != artifact['data_signature']:
        raise ValueError('Frozen data/model signature mismatch')
    reference = json.loads(Path(args.reference).read_text())
    by_id = {row['id']: row for row in reference}
    selected = []
    for family in sorted({r['family'] for r in data['validation']}):
        selected.extend([r for r in data['validation'] if r['family'] == family][:4])
    fresh = predict(scorer, selected)
    max_difference = max(abs(a-b) for row in fresh
                         for a,b in zip(row['probabilities'], by_id[row['id']]['probabilities']))
    if max_difference > 1e-4:
        raise RuntimeError(f'Fresh reconstruction differs: {max_difference}')
    write_json(out / 'reload_predictions.json', fresh)
    parameters = [p for p in scorer.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW([
        {'params': [p for p in scorer.lm.parameters() if p.requires_grad], 'lr': config['lr']},
        {'params': scorer.head.parameters(), 'lr': config['head_lr']}], weight_decay=0.01)
    optimizer.load_state_dict(artifact['optimizer'])
    # Longest complete decision in each family, plus each candidate-count class.
    stress_rows = {}
    for key in sorted({r['family'] for r in data['train']}):
        row = max((r for r in data['train'] if r['family']==key),
                  key=lambda r:max(map(len,r['_sequences'])))
        stress_rows[row['id']] = row
    for count in sorted({len(r['choices']) for r in data['train']}):
        row = max((r for r in data['train'] if len(r['choices'])==count),
                  key=lambda r:max(map(len,r['_sequences'])))
        stress_rows[row['id']] = row
    stress = []
    for row in stress_rows.values():
        scorer.train()
        optimizer.zero_grad(set_to_none=True)
        loss = decision_backward(scorer,row,1,config.get('two_pass',False))
        norm = torch.nn.utils.clip_grad_norm_(parameters,1.0,error_if_nonfinite=True)
        stress.append({'id':row['id'],'family':row['family'],'branches':len(row['choices']),
                       'max_branch_tokens':max(map(len,row['_sequences'])),
                       'loss':loss,'gradient_norm':norm.item(),
                       'peak_cuda_allocated_bytes':torch.cuda.max_memory_allocated()})
        write_json(out / 'stress.json', stress)
    optimizer.zero_grad(set_to_none=True)
    training_peak = torch.cuda.max_memory_allocated()
    training_reserved = torch.cuda.max_memory_reserved()
    del optimizer
    torch.cuda.reset_peak_memory_stats()
    scorer.eval()
    backend = LocalBackend.__new__(LocalBackend)
    backend.scorer, backend.temperature = scorer, artifact.get('temperature',1.0)
    backend.checkpoint, backend.calibrated = str(Path(args.checkpoint).resolve()), 'temperature' in artifact
    backend.max_tokens = scorer.max_tokens
    backend.model_name = 'opensysone-'+artifact['model_provenance']['model_id'].split('/')[-1].lower()
    payload = json.loads((Path(__file__).resolve().parents[1]/'examples/jev_request.json').read_text())
    direct = backend(payload)
    server = create_server(backend,port=0,api_key='local-verification-only')
    worker = threading.Thread(target=server.serve_forever,daemon=True)
    worker.start()
    try:
        remote = RemoteBackend(f'http://127.0.0.1:{server.server_address[1]}',
                               api_key='local-verification-only',attempts=1)
        tick = time.perf_counter()
        response = remote(payload)
        elapsed = time.perf_counter()-tick
        if response != direct:
            raise RuntimeError('HTTP response differs from direct trained-model inference')
        try:
            RemoteBackend(remote.base_url,api_key='wrong-key',attempts=1)(payload)
        except RuntimeError as error:
            if str(error) != 'Jev HTTP 401; request failed':
                raise
        else:
            raise RuntimeError('HTTP authentication did not reject an invalid key')
        try:
            remote({**payload,'state':'oversized '* (config['max_tokens']*2)})
        except RuntimeError as error:
            if str(error) != 'Jev HTTP 422; request failed':
                raise
        else:
            raise RuntimeError('HTTP did not reject an oversized complete candidate')
        # Inference has no gradient graphs or optimizer-step temporaries. Verify
        # a useful larger API context separately from the training token limit.
        backend.scorer.max_tokens = backend.max_tokens = args.inference_max_tokens
        for words in range(args.inference_max_tokens,32,-8):
            long_payload = {'model':'opensysone','state':'background '*words+'Please help today.',
                            'questions':{'urgent':{'type':'noul','instructions':'Does the customer express urgency?'}}}
            try:
                long_sequences = scorer.sequences(compile_request(long_payload)[0][4])
                break
            except ValueError as error:
                if 'no truncation' not in str(error):
                    raise
        else:
            raise RuntimeError('Could not construct a bounded long-context request')
        if max(map(len,long_sequences)) < args.inference_max_tokens-24:
            raise RuntimeError('Long-context probe did not reach its intended token length')
        long_direct = backend(long_payload)
        long_response = remote(long_payload)
        if long_response != long_direct:
            raise RuntimeError('Long-context HTTP output differs from direct inference')
        write_json(out/'http_long_context_response.json',long_response)
        large_payload = {'model':'opensysone','state':'A customer needs help with a failed payment.',
                         'questions':{'routing':{'type':'choice',
                         'instructions':'Which candidate routing category best fits the request?',
                         'criteria':{f'category_{i}':f'Routing category {i}' for i in range(255)}}}}
        remote.timeout = 300
        tick = time.perf_counter()
        large_response = remote(large_payload)
        large_elapsed = time.perf_counter()-tick
        if len(large_response['answers']['routing']['probabilities']) != 255:
            raise RuntimeError('HTTP maximum-choice request lost alternatives')
        write_json(out/'http_255_choices_response.json',large_response)
    finally:
        server.shutdown()
        server.server_close()
        worker.join(timeout=5)
    write_json(out / 'http_response.json', response)
    if sha256(args.checkpoint) != before:
        raise RuntimeError('Verification changed the source checkpoint')
    result = {'status':'passed','checkpoint':backend.checkpoint,'checkpoint_sha256':before,
              'step':artifact['step'],'fresh_reload_decisions':len(fresh),
              'fresh_reload_probability_max_abs':max_difference,
              'optimizer_state_restored':True,'optimizer_steps_taken':0,
              'longest_input_stress':stress,'http_model':backend.model_name,
              'http_matches_direct':True,'http_invalid_key_status':401,
              'http_oversized_input_status':422,'http_255_choices':True,
              'inference_max_tokens':args.inference_max_tokens,'http_long_context_branch_tokens':list(map(len,long_sequences)),
              'http_long_context_matches_direct':True,
              'http_255_choices_seconds':large_elapsed,
              'http_end_to_end_seconds':elapsed,'temperature_fitted':backend.calibrated,
              'training_peak_cuda_allocated_bytes':training_peak,
              'inference_peak_cuda_allocated_bytes':torch.cuda.max_memory_allocated(),
              'peak_cuda_allocated_bytes':max(training_peak,torch.cuda.max_memory_allocated()),
              'peak_cuda_reserved_bytes':max(training_reserved,torch.cuda.max_memory_reserved())}
    write_json(out / 'verification.json', result)
    print(json.dumps(result),flush=True)


if __name__ == '__main__':
    main()