Roll back a model with feature-contract checks

"""Local routing around a recorded Rice predictor; no cloud deployment."""
import copy
import hashlib
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from predictor import decode, encode, digest, validate, predict, CONTRACT
HERE=Path(__file__).resolve().parent
SALT='rice-canary-fixture-v1'

def selected(data, expected, contract=CONTRACT):
    if digest(data)!=expected:raise ValueError('selected model digest mismatch')
    bundle=validate(decode(data))
    if bundle['contract']!=contract:raise ValueError('serving feature contract incompatible')
    return {'bundle':bundle,'version':expected}

def bundles():
    raw=(HERE/'baseline-model.json').read_bytes();live=selected(raw,digest(raw))
    altered=copy.deepcopy(live['bundle']);altered['bias']+=1.0
    # Deliberate parameter fixture, not an independently trained candidate.
    altered['fixture_variant']='recorded baseline bias plus 1.0; no training'
    raw=encode(altered)
    return live,selected(raw,digest(raw))

def identity(value):
    if not isinstance(value,str) or not 1<=len(value)<=80:raise ValueError('bounded nonempty identity required')
    return value

def cohort(user, basis_points=2000):
    identity(user)
    if type(basis_points) is not int or not 0<=basis_points<=10000:raise ValueError('invalid canary allocation')
    bucket=int.from_bytes(hashlib.sha256((SALT+':'+user).encode()).digest()[:8],'big')%10000
    return 'canary' if bucket<basis_points else 'control'

class Router:
    def __init__(self,live,candidate):
        self.live,self.candidate=live,candidate
        self.active=live
        self.audit=[]
        self.lock=threading.Lock()
    def infer(self,request,mode='shadow',basis_points=2000):
        if not isinstance(request,dict) or set(request)!={'id','user','input'}:raise ValueError('routing envelope mismatch')
        identity(request['id']);identity(request['user'])
        if mode not in ('shadow','canary','active'):raise ValueError('unknown route mode')
        arm=cohort(request['user'],basis_points) if mode=='canary' else 'control'
        with self.lock:chosen=self.active if mode=='active' else self.candidate if arm=='canary' else self.live
        answer=predict(chosen['bundle'],request['input'],chosen['version'])
        entry={'id':request['id'],'user':request['user'],'arm':arm,'served':answer}
        if mode=='shadow':
            try:entry['shadow']=predict(self.candidate['bundle'],request['input'],self.candidate['version'])
            except (ValueError,TypeError,OverflowError,KeyError) as error:entry['shadow_error']=type(error).__name__
        with self.lock:self.audit.append(entry)
        return answer
    def rollback(self,data,expected,contract=CONTRACT):
        candidate=selected(data,expected,contract)
        # Swap the complete local selection after every check has succeeded.
        with self.lock:
            previous=self.active['version'];self.active=candidate
        return {'previous':previous,'selected':expected,'contract':contract}

def server(router,mode='shadow',basis_points=2000):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self,*_):pass
        def reply(self,status,value):
            data=encode(value);self.send_response(status);self.send_header('Content-Type','application/json')
            self.send_header('Content-Length',str(len(data)));self.end_headers();self.wfile.write(data)
        def do_POST(self):
            if self.path!='/predict':self.reply(404,{'error':'unknown route'});return
            try:
                size=int(self.headers.get('Content-Length','0'))
                if not 0<size<=4096:self.reply(413,{'error':'body size'});return
                if self.headers.get_content_type()!='application/json':self.reply(415,{'error':'content type'});return
                self.reply(200,router.infer(decode(self.rfile.read(size)),mode,basis_points))
            except (ValueError,TypeError,OverflowError,KeyError):self.reply(422,{'error':'invalid request'})
    return HTTPServer(('127.0.0.1',0),Handler)