Add owned OpenAI endpoint and coordinated native model switching

test_endpoint.py

@@ -0,0 +1,119 @@
+import contextlib
+import hashlib
+import http.client
+import json
+import socket
+import tempfile
+import threading
+import time
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
+from unittest.mock import Mock
+from endpoint import Endpoint
+from inference import Scheduler,InferenceError
+
+class FakeWorker:
+ def __init__(self):self.name=None;self.started=[];self.stopped=[];self.requests=[];self.block=threading.Event();self.block.set()
+ def stop(self):
+  if self.name:self.stopped.append(self.name)
+  self.name=None
+ def ensure(self,p):self.name=p['name'];self.started.append(self.name)
+ def status(self):return dict(state='ready' if self.name else 'stopped',profile_name=self.name,profile_id=self.name,port=0,error=None,gpus=[])
+ def connect(self):return http.client.HTTPConnection('127.0.0.1',self.http.server_port,timeout=3),'internal-test-only'
+
+class Upstream(BaseHTTPRequestHandler):
+ def log_message(self,*args):pass
+ def do_POST(self):
+  w=self.server.worker;data=json.loads(self.rfile.read(int(self.headers['Content-Length'])));w.requests.append(data)
+  if data.get('stream'):
+   self.send_response(200);self.send_header('Content-Type','text/event-stream');self.end_headers();self.wfile.write(b'data: {"choices":[]}\n\n');self.wfile.flush();w.block.wait(3);self.wfile.write(b'data: [DONE]\n\n')
+  else:
+   body=json.dumps(dict(id='test',model=data['model'],choices=[{'message':{'role':'assistant','content':'synthetic'}}])).encode();self.send_response(200);self.send_header('Content-Length',str(len(body)));self.end_headers();self.wfile.write(body)
+
+class EndpointTests(unittest.TestCase):
+ def setUp(self):
+  self.tmp=tempfile.TemporaryDirectory();self.worker=FakeWorker();self.worker.http=ThreadingHTTPServer(('127.0.0.1',0),Upstream);self.worker.http.worker=self.worker;threading.Thread(target=self.worker.http.serve_forever,daemon=True).start()
+  self.rows=[dict(id=n,name=n,kind='chat',revision=1,updated_at=1,runnable=True,blockers=[],parameters=dict(slots=2,temperature=.2,top_p=.8,top_k=20)) for n in ('alpha','beta')]
+  self.rows.append(dict(id='image',name='image',kind='image',revision=1,updated_at=1,runnable=True,blockers=[],parameters=dict(width=512,height=512)))
+  self.rows.append(dict(id='audio',name='audio',kind='audio',revision=1,updated_at=1,runnable=False,blockers=['No worker'],parameters={}))
+  self.profiles=SimpleNamespace(status=lambda:dict(profiles=self.rows));self.record={'api_token_hash':hashlib.sha256(b'A'*32).hexdigest()}
+  self.images=Mock();self.images.status.return_value=dict(job=None)
+  self.ep=Endpoint(self.tmp.name,self.profiles,self.worker,Scheduler(self.worker),self.images,SimpleNamespace(read=lambda:self.record),1)
+  with socket.socket() as s:s.bind(('127.0.0.1',0));self.port=s.getsockname()[1]
+  self.ep.configure({'port':self.port});self.ep.start()
+ def tearDown(self):
+  self.worker.block.set();self.ep.close();self.worker.http.shutdown();self.worker.http.server_close();self.tmp.cleanup()
+ def request(self,path='/v1/models',data=None,token='A'*32):
+  conn=http.client.HTTPConnection('127.0.0.1',self.port,timeout=5);headers={'Content-Type':'application/json','Authorization':'Bearer '+token}
+  conn.request('POST' if data is not None else 'GET',path,body=json.dumps(data) if data is not None else None,headers=headers);r=conn.getresponse();body=r.read();conn.close();return r.status,(body if data and data.get('stream') else json.loads(body))
+ def enable(self,name):self.ep.enable(dict(id=name,enabled=True))
+ def test_explicit_publication_auth_rotation_and_missing_audio(self):
+  self.assertEqual(self.request()[1]['data'],[]);self.assertEqual(self.request(token='bad')[0],401)
+  self.enable('alpha');self.assertEqual([p['id'] for p in self.request()[1]['data']],['alpha']);self.assertFalse(self.worker.started)
+  with self.assertRaises(ValueError):self.enable('audio')
+  self.assertEqual(self.request('/v1/audio/speech',{})[0],501)
+  self.record['api_token_hash']=hashlib.sha256(b'B'*32).hexdigest();self.assertEqual(self.request()[0],401);self.assertEqual(self.request(token='B'*32)[0],200)
+ def test_profile_switch_and_sampling_defaults(self):
+  for name in ('alpha','beta'):self.enable(name)
+  for name in ('alpha','beta'):
+   status,data=self.request('/v1/chat/completions',dict(model=name,messages=[dict(role='user',content='test')]))
+   self.assertEqual(status,200);self.assertEqual(data['model'],name)
+  self.assertEqual(self.worker.stopped,['alpha']);self.assertEqual(self.worker.requests[-1]['temperature'],.2)
+  self.assertEqual(self.ep.status()['counts']['llm']['enabled'],2)
+ def test_unknown_or_unsupported_request_does_not_load(self):
+  self.enable('alpha')
+  for req in [dict(model='unknown',messages=[{}]),dict(model='alpha',messages=[dict(content=[dict(type='image_url')])]),dict(model='alpha',messages=[{}],cache_file='/tmp/foo')]:
+   self.assertGreaterEqual(self.request('/v1/chat/completions',req)[0],400)
+  self.assertFalse(self.worker.started)
+ def test_stream_lease_blocks_switch_until_stream_finishes(self):
+  self.enable('alpha');self.enable('beta');self.worker.block.clear()
+  conn=http.client.HTTPConnection('127.0.0.1',self.port,timeout=5)
+  conn.request('POST','/v1/chat/completions',json.dumps(dict(model='alpha',messages=[{}],stream=True)),headers={'Content-Type':'application/json','Authorization':'Bearer '+'A'*32});r=conn.getresponse();self.assertEqual(r.status,200)
+  result=[];thread=threading.Thread(target=lambda:result.append(self.request('/v1/chat/completions',dict(model='beta',messages=[{}]))));thread.start()
+  deadline=time.monotonic()+2
+  while not self.ep.scheduler.status()['waiting_requests'] and time.monotonic()<deadline:time.sleep(.01)
+  self.assertEqual(self.worker.name,'alpha');self.assertEqual(self.worker.stopped,[])
+  self.worker.block.set();self.assertIn(b'[DONE]',r.read());conn.close();thread.join(3);self.assertEqual(result[0][0],200);self.assertEqual(self.worker.name,'beta')
+ def test_port_conflict_and_graceful_stop(self):
+  with self.assertRaises(ValueError):self.ep.configure(dict(port=self.port+1))
+  self.ep.stop();deadline=time.monotonic()+3
+  while self.ep.state!='stopped' and time.monotonic()<deadline:time.sleep(.02)
+  self.assertEqual(self.ep.state,'stopped')
+  with socket.socket() as sock:
+   sock.bind(('127.0.0.1',self.port))
+   with self.assertRaises(ValueError):self.ep.start()
+  self.ep.start();self.assertTrue(self.ep.status()['reachable'])
+ def test_image_unloads_chat_and_returns_openai_shape(self):
+  self.enable('alpha');self.enable('image');self.request('/v1/chat/completions',dict(model='alpha',messages=[{}]))
+  self.images.start.return_value={'id':'new'};self.images.status.return_value={'job':{'id':'new','state':'complete'}};self.images.image.return_value=b'synthetic-png'
+  status,data=self.request('/v1/images/generations',dict(model='image',prompt='synthetic',response_format='b64_json'))
+  self.assertEqual(status,200);self.assertIn('b64_json',data['data'][0]);self.assertIsNone(self.worker.name);self.images.start.assert_called_once_with('image','synthetic',reserved=True)
+ def test_disabled_and_edited_queued_profiles_are_rejected(self):
+  self.enable('alpha');self.enable('beta');self.worker.block.clear()
+  first=threading.Thread(target=lambda:self.request('/v1/chat/completions',dict(model='alpha',messages=[{}],stream=True)));first.start()
+  deadline=time.monotonic()+2
+  while not self.worker.requests and time.monotonic()<deadline:time.sleep(.01)
+  result=[];second=threading.Thread(target=lambda:result.append(self.request('/v1/chat/completions',dict(model='beta',messages=[{}]))));second.start()
+  deadline=time.monotonic()+2
+  while not self.ep.scheduler.status()['waiting_requests'] and time.monotonic()<deadline:time.sleep(.01)
+  self.ep.enable(dict(id='beta',enabled=False));self.worker.block.set();first.join(4);second.join(4);self.assertEqual(result[0][0],503);self.assertNotIn('beta',self.worker.started)
+
+class SchedulerTests(unittest.TestCase):
+ def test_parallel_slots_and_fifo_switch(self):
+  w=FakeWorker();s=Scheduler(w)
+  with s.lease('a',2):
+   with s.lease('a',2):self.assertEqual(s.status()['active_requests'],2)
+   with self.assertRaises(InferenceError):
+    with s.lease('b',timeout=0):pass
+  with s.lease('b'):self.assertEqual(s.status()['active_requests'],1)
+  self.assertEqual(s.status()['active_requests'],0)
+ def test_failed_prepare_releases_queue_and_can_retry(self):
+  s=Scheduler(FakeWorker())
+  with self.assertRaises(ValueError):
+   with s.lease('bad',prepare=lambda:(_ for _ in ()).throw(ValueError('bad'))):pass
+  with s.lease('good'):self.assertFalse(s.status()['switching'])
+  self.assertEqual(s.status()['active_requests'],0)
+
+if __name__=='__main__':unittest.main()