File size: 12,220 Bytes
aa73ff9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """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 # Never log prompts, credentials or commands.
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()
|