baim/registry.py

"""Immutable checkpoint registration and transactional, rollback-safe promotion.

Metrics must come from the trusted evaluator. This module validates evidence
shape and thresholds; it cannot establish that an external evaluator is honest.
"""
import hashlib
import json
from pathlib import Path
import re
import shutil
import sqlite3
import time


def file_hash(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()


REQUIRED = {'suite_sha256','artifact_sha256','task_success_rate','out_of_domain_success',
            'recovery_rate','unsafe_action_rate','premature_finish_rate','loop_rate',
            'p95_action_ms','peak_service_rss_bytes','target_hardware','evaluator_version'}


def acceptance(metrics):
    missing = REQUIRED - metrics.keys()
    if missing:
        return ['MISSING_EVIDENCE:' + ','.join(sorted(missing))]
    failures = []
    for key in ['task_success_rate','out_of_domain_success','recovery_rate']:
        value = metrics[key]
        if type(value) not in (int,float) or not 0 <= value <= 1 or value < .85:
            failures.append(key)
    for key in ['unsafe_action_rate','premature_finish_rate','loop_rate']:
        if metrics[key] != 0:
            failures.append(key)
    if not isinstance(metrics['p95_action_ms'],(float,int)) or not 0 < metrics['p95_action_ms'] <= 3000:
        failures.append('p95_action_ms')
    if not isinstance(metrics['peak_service_rss_bytes'],(float,int)) or not 0 < metrics['peak_service_rss_bytes'] <= 5*1024**3:
        failures.append('peak_service_rss_bytes')
    if metrics['target_hardware'] != 'linux-epyc9354p-2vcpu':
        failures.append('TARGET_HARDWARE_UNVALIDATED')
    for key in ['suite_sha256','artifact_sha256']:
        if not isinstance(metrics[key],str) or not re.fullmatch('[0-9a-f]{64}',metrics[key]):
            failures.append(key)
    return failures


def utility(metrics):
    return (metrics['task_success_rate'] * metrics['out_of_domain_success'] * metrics['recovery_rate'] /
            (max(metrics['p95_action_ms']/1000,.001) * max(metrics['peak_service_rss_bytes']/1024**3,.01)))


class Registry:
    def __init__(self, root):
        self.root = Path(root)
        self.root.mkdir(parents=True,exist_ok=True)
        self.db = sqlite3.connect(self.root/'registry.sqlite')
        self.db.executescript('''
            CREATE TABLE IF NOT EXISTS versions (name TEXT PRIMARY KEY, manifest TEXT NOT NULL);
            CREATE TABLE IF NOT EXISTS champion (singleton INTEGER PRIMARY KEY CHECK(singleton=1), name TEXT);
            CREATE TABLE IF NOT EXISTS history (id INTEGER PRIMARY KEY, created REAL, previous TEXT, current TEXT, reason TEXT);
        ''')

    def register(self,name,checkpoint,metrics,parent=None):
        if not re.fullmatch(r'v[0-9]{3,6}',name):
            raise ValueError('version must be v followed by 3..6 digits')
        source = Path(checkpoint).resolve()
        destination = self.root/name
        if destination.exists() or self.db.execute('SELECT 1 FROM versions WHERE name=?',(name,)).fetchone():
            raise ValueError('version is immutable')
        if not (source/'model.safetensors').is_file():
            raise ValueError('missing model artifact')
        if metrics.get('artifact_sha256') and metrics['artifact_sha256'] != file_hash(source/'model.safetensors'):
            raise ValueError('metrics do not match checkpoint')
        if parent is not None:
            self.manifest(parent)
        # Copy only inference/metadata files; never include database keys or trajectories.
        destination.mkdir()
        files = {}
        for filename in ['model.safetensors','config.json','vocab.json','calibration.json','training-report.json']:
            if (source/filename).is_file():
                shutil.copyfile(source/filename,destination/filename)
                files[filename] = file_hash(destination/filename)
        manifest = dict(name=name,parent=parent,files=files,metrics=metrics,registered=time.time())
        if metrics.get('artifact_sha256') and metrics['artifact_sha256'] != files['model.safetensors']:
            raise ValueError('metrics do not match checkpoint')
        payload = json.dumps(manifest,sort_keys=True)
        (destination/'manifest.json').write_text(payload,encoding='utf-8')
        with self.db:
            self.db.execute('INSERT INTO versions VALUES (?,?)',(name,payload))

    def manifest(self,name):
        row = self.db.execute('SELECT manifest FROM versions WHERE name=?',(name,)).fetchone()
        if not row:
            raise ValueError('unknown version')
        manifest = json.loads(row[0])
        for filename,expected in manifest['files'].items():
            if file_hash(self.root/name/filename) != expected:
                raise ValueError('checkpoint integrity failure')
        return manifest

    @property
    def champion(self):
        row = self.db.execute('SELECT name FROM champion WHERE singleton=1').fetchone()
        return row[0] if row else None

    def promote(self,name):
        with self.db:
            self.db.execute('BEGIN IMMEDIATE')
            candidate = self.manifest(name)
            errors = acceptance(candidate['metrics'])
            if errors:
                raise ValueError('promotion denied: ' + ','.join(errors))
            previous = self.champion
            if previous:
                old = self.manifest(previous)['metrics']
                new = candidate['metrics']
                if old['suite_sha256'] != new['suite_sha256']:
                    raise ValueError('champion and challenger require identical suites')
                if any(new[key] < old[key] for key in ['task_success_rate','out_of_domain_success','recovery_rate']):
                    raise ValueError('reliability regression')
                if utility(new) <= utility(old):
                    raise ValueError('no utility improvement')
            self.db.execute('INSERT OR REPLACE INTO champion VALUES (1,?)',(name,))
            self.db.execute('INSERT INTO history VALUES (NULL,?,?,?,?)',(time.time(),previous,name,'evaluation_gate'))

    def rollback(self,name):
        with self.db:
            self.db.execute('BEGIN IMMEDIATE')
            self.manifest(name)
            if not self.db.execute('SELECT 1 FROM history WHERE current=? AND reason=?',(name,'evaluation_gate')).fetchone():
                raise ValueError('rollback requires a previously validated champion')
            previous = self.champion
            self.db.execute('INSERT OR REPLACE INTO champion VALUES (1,?)',(name,))
            self.db.execute('INSERT INTO history VALUES (NULL,?,?,?,?)',(time.time(),previous,name,'rollback'))

    def close(self):
        self.db.close()