Bounded FIFO inference scheduling: admission queue, OpenAI-shaped 429/503, queued-disconnect detection (#28)

c/tests/test_openai_server.py

@@ -5,7 +5,8 @@ import unittest
 from urllib.error import HTTPError
 from urllib.request import Request, urlopen
 
-from openai_server import APIError, APIServer, END, generation_options, read_engine_turn, render_chat
+from openai_server import (APIError, APIServer, ClientCancelled, END, GenerationScheduler,
+                           generation_options, read_engine_turn, render_chat)
 
 
 class FakeEngine:
@@ -19,6 +20,18 @@ class FakeEngine:
         return {"prompt_tokens": 7, "completion_tokens": 2, "length_limited": False}
 
 
+class BlockingEngine(FakeEngine):
+    def __init__(self):
+        super().__init__()
+        self.entered = threading.Event()
+        self.release = threading.Event()
+
+    def generate(self, prompt, maximum, temperature, top_p, on_text):
+        self.entered.set()
+        self.release.wait(2)
+        return super().generate(prompt, maximum, temperature, top_p, on_text)
+
+
 class TemplateTest(unittest.TestCase):
     def test_renders_text_subset_of_official_template(self):
         prompt = render_chat([
@@ -66,6 +79,83 @@ class ProtocolTest(unittest.TestCase):
         self.assertTrue(stats["length_limited"])
 
 
+class SchedulerTest(unittest.TestCase):
+    def test_rejects_when_waiting_queue_is_full(self):
+        scheduler = GenerationScheduler(max_queue=0, queue_timeout=1)
+        with scheduler.admit():
+            with self.assertRaises(APIError) as caught:
+                with scheduler.admit():
+                    pass
+        self.assertEqual(caught.exception.status, 429)
+        self.assertEqual(caught.exception.code, "queue_full")
+        self.assertEqual(scheduler.snapshot()["rejected"], 1)
+
+    def test_times_out_and_cancels_queued_requests(self):
+        scheduler = GenerationScheduler(max_queue=2, queue_timeout=0.02)
+        with scheduler.admit():
+            with self.assertRaises(APIError) as timed_out:
+                with scheduler.admit():
+                    pass
+            with self.assertRaises(ClientCancelled):
+                with scheduler.admit(lambda: True):
+                    pass
+        stats = scheduler.snapshot()
+        self.assertEqual(timed_out.exception.code, "queue_timeout")
+        self.assertEqual(stats["timed_out"], 1)
+        self.assertEqual(stats["cancelled"], 1)
+
+    def test_admits_waiters_in_fifo_order(self):
+        scheduler = GenerationScheduler(max_queue=2, queue_timeout=1)
+        entered = threading.Event()
+        release = threading.Event()
+        order = []
+
+        def run(name, block=False):
+            with scheduler.admit():
+                order.append(name)
+                if block:
+                    entered.set()
+                    release.wait(1)
+
+        first = threading.Thread(target=run, args=("first", True))
+        second = threading.Thread(target=run, args=("second",))
+        third = threading.Thread(target=run, args=("third",))
+        first.start(); entered.wait(1)
+        second.start()
+        for _ in range(100):
+            if scheduler.snapshot()["queued"] == 1: break
+            threading.Event().wait(0.005)
+        third.start()
+        for _ in range(100):
+            if scheduler.snapshot()["queued"] == 2: break
+            threading.Event().wait(0.005)
+        release.set()
+        first.join(1); second.join(1); third.join(1)
+        self.assertEqual(order, ["first", "second", "third"])
+        self.assertEqual(scheduler.snapshot()["completed"], 3)
+
+    def test_close_rejects_waiters(self):
+        scheduler = GenerationScheduler(max_queue=1, queue_timeout=1)
+        entered = threading.Event()
+        release = threading.Event()
+        errors = []
+
+        def active():
+            with scheduler.admit():
+                entered.set(); release.wait(1)
+
+        def waiting():
+            try:
+                with scheduler.admit(): pass
+            except APIError as error:
+                errors.append(error.code)
+
+        first = threading.Thread(target=active); first.start(); entered.wait(1)
+        second = threading.Thread(target=waiting); second.start()
+        scheduler.close(); release.set(); first.join(1); second.join(1)
+        self.assertEqual(errors, ["scheduler_closed"])
+
+
 class HTTPTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -77,6 +167,7 @@ class HTTPTest(unittest.TestCase):
 
     @classmethod
     def tearDownClass(cls):
+        cls.server.scheduler.close()
         cls.server.shutdown()
         cls.server.server_close()
         cls.thread.join(timeout=2)
@@ -96,6 +187,12 @@ class HTTPTest(unittest.TestCase):
             self.request("/v1/models", key="wrong")
         self.assertEqual(caught.exception.code, 401)
 
+    def test_health_reports_scheduler(self):
+        with self.request("/health") as response:
+            scheduler = json.load(response)["scheduler"]
+        self.assertEqual(scheduler["max_queue"], 8)
+        self.assertIn("queued", scheduler)
+
     def test_browser_preflight(self):
         request = Request(self.base + "/v1/chat/completions", method="OPTIONS", headers={
             "Origin": "http://localhost:5173",
@@ -113,9 +210,11 @@ class HTTPTest(unittest.TestCase):
             "max_tokens": 4,
         }) as response:
             body = json.load(response)
+            queue_wait = response.headers.get("x-colibri-queue-wait-ms")
         self.assertEqual(body["object"], "chat.completion")
         self.assertEqual(body["choices"][0]["message"]["content"], "Héllo")
         self.assertEqual(body["usage"], {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9})
+        self.assertIsNotNone(queue_wait)
         self.assertIn("<|user|>Hi<|assistant|><think></think>", self.engine.calls[-1][0])
 
     def test_streaming_chat_completion(self):
@@ -148,5 +247,45 @@ class HTTPTest(unittest.TestCase):
         self.assertEqual(caught.exception.code, 400)
 
 
+class SchedulerHTTPTest(unittest.TestCase):
+    def setUp(self):
+        self.engine = BlockingEngine()
+        self.server = APIServer(("127.0.0.1", 0), self.engine, "test-model",
+                                max_tokens=16, max_queue=0)
+        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
+        self.thread.start()
+        self.url = f"http://127.0.0.1:{self.server.server_port}/v1/chat/completions"
+
+    def tearDown(self):
+        self.engine.release.set()
+        self.server.scheduler.close()
+        self.server.shutdown(); self.server.server_close(); self.thread.join(timeout=2)
+
+    def request(self):
+        body = json.dumps({"model": "test-model", "messages": [
+            {"role": "user", "content": "Hi"}]}).encode()
+        return urlopen(Request(self.url, data=body, headers={"Content-Type": "application/json"}), timeout=2)
+
+    def test_queue_full_returns_429_before_generation(self):
+        first_errors = []
+
+        def first_request():
+            try:
+                with self.request() as response: response.read()
+            except Exception as error:
+                first_errors.append(error)
+
+        first = threading.Thread(target=first_request); first.start()
+        self.assertTrue(self.engine.entered.wait(1))
+        with self.assertRaises(HTTPError) as caught:
+            self.request()
+        error = json.loads(caught.exception.read())["error"]
+        self.assertEqual(caught.exception.code, 429)
+        self.assertEqual(caught.exception.headers["Retry-After"], "1")
+        self.assertEqual(error["code"], "queue_full")
+        self.engine.release.set(); first.join(2)
+        self.assertEqual(first_errors, [])
+
+
 if __name__ == "__main__":
     unittest.main()