| """Small authenticated classifier API; one lazy GPU worker, unloaded when idle.""" |
| import argparse |
| import contextlib |
| import hmac |
| import json |
| import os |
| import resource |
| import selectors |
| import signal |
| import subprocess |
| import sys |
| import threading |
| import time |
| import uuid |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| from pathlib import Path |
|
|
| ROOT=Path(__file__).resolve().parent |
| MODEL_NAME='laya-devops-permissions' |
|
|
|
|
| def decision_state(body,systemone=False): |
| """Expose the trained permission head, never pretend to answer other questions.""" |
| if body.get('model')!=MODEL_NAME:raise ValueError('Unknown model') |
| if body.get('stream') or body.get('images') or body.get('assets'):raise ValueError('Only non-streaming text decisions are supported') |
| canonical=json.loads((ROOT/'questions.json').read_text())['permission'] |
| questions=body.get('questions') |
| if systemone: |
| if not isinstance(questions,dict) or set(questions)!={'permission'}:raise ValueError('Only the trained permission question is supported; use questions.json') |
| question=questions['permission'];state=body.get('state') |
| else: |
| if not isinstance(questions,list) or len(questions)!=1:raise ValueError('Provide exactly one trained permission question') |
| q=questions[0] |
| if q.get('name')!='permission':raise ValueError('Question name must be permission') |
| choices=q.get('choices',[]) |
| if not isinstance(choices,list) or len(choices)!=3:raise ValueError('Exactly three permission choices are required') |
| question={'type':q.get('type'),'instructions':q.get('instructions'),'criteria':{c['value']:c.get('description','') for c in choices}} |
| state=body.get('input') |
| if isinstance(state,list): |
| if len(state)!=1 or state[0].get('role')!='user':raise ValueError('Provide one user input containing the JSON state') |
| state=state[0].get('content') |
| if isinstance(state,list): |
| if len(state)!=1 or state[0].get('type')!='input_text':raise ValueError('Provide one input_text containing the JSON state') |
| state=state[0].get('text') |
| if not isinstance(question,dict) or any(question.get(k)!=canonical[k] for k in canonical):raise ValueError('This fixed classifier requires the exact trained question from questions.json') |
| if isinstance(state,str):state=json.loads(state) |
| return validate_state(state) |
|
|
|
|
| def decision_response(result,systemone=False): |
| p=result['probabilities'] |
| answer={'type':'choice','choice':result['decision'],'confidence':p[result['decision']], |
| 'probabilities':p if systemone else [{'value':k,'probability':v} for k,v in p.items()]} |
| if not systemone:answer['name']='permission' |
| return {'model':MODEL_NAME,'answers':{'permission':answer} if systemone else [answer], |
| 'usage':{'input_tokens':result.get('input_tokens',0),'output_tokens':0,'total_tokens':result.get('input_tokens',0)}, |
| 'permission_policy':{'decision':result['decision'],'raw_choice':result['raw_decision'], |
| 'allow_threshold':result['allow_threshold'],'probabilities_are_raw':True}, |
| 'executed':False,'benchmark_qualified':False} |
|
|
|
|
| def validate_state(state): |
| if not isinstance(state,dict):raise ValueError('State must be a JSON object') |
| for key in ('command','trusted_user_request','environment'): |
| if not isinstance(state.get(key),str) or not state[key].strip():raise ValueError('Required nonempty string: '+key) |
| return state |
|
|
|
|
| def chat_state(body): |
| if body.get('model')!=MODEL_NAME:raise ValueError('Unknown model') |
| if body.get('stream'):raise ValueError('Streaming is not supported by this classifier') |
| if body.get('tools') or body.get('functions'):raise ValueError('This endpoint classifies JSON states; it does not execute tools') |
| messages=body.get('messages') |
| if not isinstance(messages,list) or len(messages)!=1 or messages[0].get('role')!='user':raise ValueError('Provide exactly one user message containing the JSON state') |
| content=messages[0].get('content') |
| if not isinstance(content,str):raise ValueError('Message content must be a JSON string') |
| return validate_state(json.loads(content)) |
|
|
|
|
| def worker(model): |
| os.environ['HF_HUB_OFFLINE']='1' |
| os.environ.setdefault('HF_HOME',str(ROOT/'.hf-cache')) |
| with contextlib.redirect_stdout(sys.stderr): |
| import laya_mlx |
| import mlx.core as mx |
| from layastudio import engine |
| from evaluate import decision |
| from input_preprocessing import load_preprocessor |
| from operating_policy import load_allow_threshold |
| engine.limit_mlx_cache() |
| agent=laya_mlx.load(str(model));q=json.loads((model/'questions.json').read_text()) |
| normalizer=load_preprocessor(json.loads((model/'input-preprocessor.json').read_text())) |
| threshold=load_allow_threshold(model) |
| prefix,_=engine.build_prefix(agent.tok,engine.internal_question(q['permission']),agent.cfg['head_max_len']) |
| room=agent.cfg['max_len']-len(prefix)-1 |
| for line in sys.stdin: |
| try: |
| state=normalizer(validate_state(json.loads(line))) |
| tokens=len(agent.tok(engine.serialize_state(state).replace(agent.tok.mask_token,' '))['input_ids']) |
| if tokens>room:raise ValueError('State exceeds token budget; manual review required') |
| with contextlib.redirect_stdout(sys.stderr):p=agent.predict(state,q)['answers']['permission']['probabilities'] |
| result=dict(decision=decision(p,threshold),raw_decision=decision(p),probabilities=p,allow_threshold=threshold,input_tokens=tokens+len(prefix)+1,executed=False,benchmark_qualified=False) |
| except Exception as e:result=dict(decision='ASK',error=str(e),executed=False,benchmark_qualified=False) |
| peak=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
| result['memory']={'process_peak_rss_bytes':peak if sys.platform=='darwin' else peak*1024,'mlx_active_bytes':mx.get_active_memory(),'mlx_peak_bytes':mx.get_peak_memory(),'mlx_cache_bytes':mx.get_cache_memory()} |
| print(json.dumps(result),flush=True) |
|
|
|
|
| class LazyWorker: |
| def __init__(self,model,idle_seconds,timeout=60): |
| self.model=model;self.idle_seconds=idle_seconds;self.timeout=timeout |
| self.proc=None;self.lock=threading.Lock();self.last_used=time.monotonic();self.last_memory=None |
| self.closed=threading.Event() |
| threading.Thread(target=self.reaper,daemon=True).start() |
| def stop_locked(self): |
| if self.proc: |
| self.proc.terminate() |
| try:self.proc.wait(timeout=5) |
| except subprocess.TimeoutExpired:self.proc.kill();self.proc.wait() |
| self.proc=None |
| def reaper(self): |
| while not self.closed.wait(2): |
| with self.lock: |
| if self.proc and time.monotonic()-self.last_used>=self.idle_seconds:self.stop_locked() |
| def classify(self,state): |
| if not self.lock.acquire(timeout=2):return dict(decision='ASK',error='Classifier busy; retry or request manual review',executed=False,benchmark_qualified=False) |
| try: |
| if self.proc is None or self.proc.poll() is not None: |
| self.stop_locked() |
| self.proc=subprocess.Popen([sys.executable,str(Path(__file__).resolve()),'--worker','--model',str(self.model)],stdin=subprocess.PIPE,stdout=subprocess.PIPE,text=True,bufsize=1) |
| self.proc.stdin.write(json.dumps(state)+'\n');self.proc.stdin.flush() |
| with selectors.DefaultSelector() as sel: |
| sel.register(self.proc.stdout,selectors.EVENT_READ) |
| if not sel.select(self.timeout):raise TimeoutError('Model timeout; manual review required') |
| line=self.proc.stdout.readline() |
| if not line:raise RuntimeError('Model worker exited; manual review required') |
| result=json.loads(line);self.last_memory=result.get('memory');return result |
| except Exception as e: |
| self.stop_locked();return dict(decision='ASK',error=str(e),executed=False,benchmark_qualified=False) |
| finally:self.last_used=time.monotonic();self.lock.release() |
| def close(self): |
| self.closed.set() |
| with self.lock:self.stop_locked() |
|
|
|
|
| def handler_for(runtime,api_key): |
| class Handler(BaseHTTPRequestHandler): |
| def log_message(self,*args):pass |
| def reply(self,status,data): |
| b=json.dumps(data).encode();self.send_response(status);self.send_header('Content-Type','application/json');self.send_header('Content-Length',str(len(b)));self.end_headers();self.wfile.write(b) |
| def authorized(self): |
| if hmac.compare_digest(self.headers.get('Authorization',''),'Bearer '+api_key):return True |
| self.reply(401,{'error':{'message':'Invalid API key','type':'authentication_error'}});return False |
| def do_GET(self): |
| if not self.authorized():return |
| if self.path=='/v1/models':self.reply(200,{'object':'list','data':[{'id':MODEL_NAME,'object':'model','created':0,'owned_by':'local'}]}) |
| elif self.path=='/health': |
| with runtime.lock:loaded=runtime.proc is not None and runtime.proc.poll() is None |
| self.reply(200,{'status':'ok','model_loaded':loaded,'idle_unload_seconds':runtime.idle_seconds,'last_loaded_memory':runtime.last_memory,'benchmark_qualified':False,'model_path':str(runtime.model)}) |
| else:self.reply(404,{'error':{'message':'Not found'}}) |
| def do_POST(self): |
| if not self.authorized():return |
| if self.path not in ('/v1/chat/completions','/classify','/v1/decisions','/v1/systemone'):self.reply(404,{'error':{'message':'Not found'}});return |
| try: |
| size=int(self.headers.get('Content-Length','0')) |
| if not 0<size<=65536:raise ValueError('Body must be 1–65536 bytes') |
| body=json.loads(self.rfile.read(size)) |
| if self.path in ('/v1/decisions','/v1/systemone'):state=decision_state(body,self.path=='/v1/systemone') |
| else:state=chat_state(body) if self.path=='/v1/chat/completions' else validate_state(body) |
| except (ValueError,TypeError,AttributeError,KeyError) as e:self.reply(400,{'error':{'message':str(e),'type':'invalid_request_error'}});return |
| result=runtime.classify(state) |
| if self.path=='/classify':self.reply(200,result);return |
| if self.path in ('/v1/decisions','/v1/systemone'): |
| if 'error' in result:self.reply(503,{'error':{'message':result['error'],'type':'classifier_unavailable'},'decision':'ASK','executed':False});return |
| self.reply(200,decision_response(result,self.path=='/v1/systemone'));return |
| self.reply(200,{'id':'chatcmpl-'+uuid.uuid4().hex,'object':'chat.completion','created':int(time.time()),'model':MODEL_NAME,'choices':[{'index':0,'message':{'role':'assistant','content':json.dumps(result)},'finish_reason':'stop'}]}) |
| return Handler |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser(description=__doc__);p.add_argument('--worker',action='store_true');p.add_argument('--model',type=Path,default=ROOT/'workspace/runs/devops-v23-policy090/model');p.add_argument('--host',default='127.0.0.1');p.add_argument('--port',type=int,default=8765);p.add_argument('--idle-seconds',type=int,default=120);p.add_argument('--key-file',type=Path,default=ROOT/'workspace/local-api.key');a=p.parse_args() |
| if a.worker:worker(a.model.resolve());return |
| if a.idle_seconds<5:p.error('idle-seconds must be >=5') |
| key=a.key_file.read_text().strip() |
| if len(key)<24:p.error('API key must have at least 24 characters') |
| runtime=LazyWorker(a.model.resolve(),a.idle_seconds);server=ThreadingHTTPServer((a.host,a.port),handler_for(runtime,key));server.daemon_threads=True |
| def shutdown_signal(signum,frame):raise KeyboardInterrupt |
| signal.signal(signal.SIGTERM,shutdown_signal);signal.signal(signal.SIGHUP,shutdown_signal) |
| try:server.serve_forever() |
| except KeyboardInterrupt:pass |
| finally:server.server_close();runtime.close() |
|
|
|
|
| if __name__=='__main__':main() |
|
|
