70 lines
3.9 KiB
Python
70 lines
3.9 KiB
Python
|
"""Bounded, cancellable persistent JSONL model process; never executes a shell."""
|
||
|
|
import json, os, selectors, signal, subprocess, threading, time
|
||
|
|
from robot_bt_coordinator.plan import canonical, strict_json
|
||
|
|
|
||
|
|
class InferenceError(RuntimeError):
|
||
|
|
def __init__(self,code,message=''):self.code=code;super().__init__(message or code)
|
||
|
|
|
||
|
|
class FixtureBackend:
|
||
|
|
model_version='fixture-only'
|
||
|
|
def __init__(self,raw):self.raw=raw
|
||
|
|
def infer(self,request,timeout,cancel=None):
|
||
|
|
if cancel is not None and cancel.is_set():raise InferenceError('CANCELED')
|
||
|
|
return self.raw(request) if callable(self.raw) else self.raw
|
||
|
|
def close(self):pass
|
||
|
|
|
||
|
|
class ProcessBackend:
|
||
|
|
"""One in-flight call; timeout/cancel kills the entire inference process group.
|
||
|
|
Restart reloads weights. The local GPU worker has no robot-control authority.
|
||
|
|
"""
|
||
|
|
def __init__(self,argv,model_version):
|
||
|
|
if not isinstance(argv,list) or not argv or any(not isinstance(x,str) or not x for x in argv):raise ValueError('explicit argv required')
|
||
|
|
if not model_version:raise ValueError('pinned model version required')
|
||
|
|
self.argv=argv;self.model_version=model_version;self.lock=threading.Lock();self.process=None
|
||
|
|
def close(self):
|
||
|
|
p=self.process;self.process=None
|
||
|
|
if p:
|
||
|
|
if p.poll() is None:
|
||
|
|
os.killpg(p.pid,signal.SIGKILL)
|
||
|
|
p.wait(timeout=5)
|
||
|
|
p.stdin.close();p.stdout.close()
|
||
|
|
def infer(self,request,timeout,cancel=None):
|
||
|
|
if not 0<timeout<=3600:raise InferenceError('INVALID_TIMEOUT')
|
||
|
|
if not self.lock.acquire(blocking=False):raise InferenceError('BUSY')
|
||
|
|
deadline=time.monotonic()+timeout
|
||
|
|
try:
|
||
|
|
if cancel is not None and cancel.is_set():raise InferenceError('CANCELED')
|
||
|
|
if self.process is None:
|
||
|
|
self.process=subprocess.Popen(self.argv,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=None,start_new_session=True,bufsize=0)
|
||
|
|
p=self.process
|
||
|
|
payload=(canonical(request)+'\n').encode()
|
||
|
|
if len(payload)>262144:raise InferenceError('INPUT_TOO_LARGE')
|
||
|
|
# Nonblocking write/read includes process startup in the same deadline.
|
||
|
|
os.set_blocking(p.stdin.fileno(),False);os.set_blocking(p.stdout.fileno(),False)
|
||
|
|
sent=0;data=b''
|
||
|
|
with selectors.DefaultSelector() as sel:
|
||
|
|
sel.register(p.stdout,selectors.EVENT_READ);sel.register(p.stdin,selectors.EVENT_WRITE)
|
||
|
|
while True:
|
||
|
|
if cancel is not None and cancel.is_set():raise InferenceError('CANCELED')
|
||
|
|
if time.monotonic()>=deadline:raise InferenceError('TIMEOUT')
|
||
|
|
for key,_ in sel.select(min(.05,max(0,deadline-time.monotonic()))):
|
||
|
|
if key.fileobj is p.stdin:
|
||
|
|
sent+=os.write(p.stdin.fileno(),payload[sent:])
|
||
|
|
if sent==len(payload):sel.unregister(p.stdin)
|
||
|
|
else:
|
||
|
|
chunk=os.read(p.stdout.fileno(),65536)
|
||
|
|
if not chunk:raise InferenceError('WORKER_EXITED')
|
||
|
|
data+=chunk
|
||
|
|
if len(data)>262144:raise InferenceError('OUTPUT_TOO_LARGE')
|
||
|
|
if b'\n' in data:
|
||
|
|
raw,extra=data.split(b'\n',1)
|
||
|
|
if extra.strip():raise InferenceError('WORKER_PROTOCOL')
|
||
|
|
msg=strict_json(raw.decode())
|
||
|
|
if not isinstance(msg,dict) or set(msg)!={'raw'} or not isinstance(msg['raw'],str):raise InferenceError('WORKER_PROTOCOL')
|
||
|
|
return msg['raw']
|
||
|
|
except InferenceError:
|
||
|
|
self.close();raise
|
||
|
|
except Exception as ex:
|
||
|
|
self.close();raise InferenceError('INFERENCE_FAILED',str(ex)) from ex
|
||
|
|
finally:self.lock.release()
|