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

README.md

@@ -151,6 +151,14 @@ Browser access from the Vite development server and Tauri local origins is enabl
 by default. Repeat `--cors-origin https://your-ui.example` to allow another exact
 origin, or use `--cors-origin '*'` only on a trusted local network.
 
+The engine owns one mutable KV context, so HTTP generation uses a bounded FIFO
+admission queue instead of pretending to run unsafe parallel sequences. Configure it
+with `--max-queue N` (default 8) and `--queue-timeout SECONDS` (default 300), or the
+`COLI_MAX_QUEUE` / `COLI_QUEUE_TIMEOUT` environment variables. Saturated and timed-out
+requests receive OpenAI-shaped HTTP 429 errors before streaming headers are sent.
+`GET /health` exposes active/queued/completed/rejected counters, and successful
+generation responses include `x-colibri-queue-wait-ms`.
+
 ### Experimental resident CUDA backend
 
 colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming

c/coli

@@ -440,7 +440,7 @@ def cmd_serve(a):
     need_model(a.model)
     from openai_server import serve
     serve(a.model, a.host, a.port, a.model_id, a.api_key,
-          a.cap, a.ngen, GLM, env_for(a), a.cors_origin)
+          a.cap,a.ngen,GLM,env_for(a),a.cors_origin,a.max_queue,a.queue_timeout)
 
 def cmd_bench(a):
     need_model(a.model)
@@ -504,6 +504,8 @@ def main():
     ps.add_argument("--model-id",default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))
     ps.add_argument("--api-key",default=os.environ.get("COLI_API_KEY"))
     ps.add_argument("--cors-origin",action="append",default=None)
+    ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))
+    ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))
     pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
     pb.add_argument("--limit",type=int,default=40); pb.add_argument("--data",default=os.path.join(HERE,"bench"))
     pc=sub.add_parser("convert", parents=[common]); pc.add_argument("--repo",default="zai-org/GLM-5.2-FP8")

c/openai_server.py

@@ -3,9 +3,13 @@
 
 import argparse
 import codecs
+import collections
+import contextlib
 import json
 import os
+import select
 import signal
+import socket
 import subprocess
 import sys
 import threading
@@ -29,13 +33,19 @@ DEFAULT_CORS_ORIGINS = (
 
 
 class APIError(Exception):
-    def __init__(self, status, message, param=None, code=None, error_type="invalid_request_error"):
+    def __init__(self, status, message, param=None, code=None, error_type="invalid_request_error",
+                 headers=None):
         super().__init__(message)
         self.status = status
         self.message = message
         self.param = param
         self.code = code
         self.error_type = error_type
+        self.headers = headers or {}
+
+
+class ClientCancelled(Exception):
+    pass
 
 
 def error_object(error):
@@ -43,6 +53,87 @@ def error_object(error):
                       "param": error.param, "code": error.code}}
 
 
+class GenerationScheduler:
+    """Bounded FIFO admission for the engine's single mutable KV context."""
+
+    def __init__(self, max_queue=8, queue_timeout=300):
+        if max_queue < 0:
+            raise ValueError("max_queue cannot be negative")
+        if queue_timeout <= 0:
+            raise ValueError("queue_timeout must be positive")
+        self.max_queue = max_queue
+        self.queue_timeout = queue_timeout
+        self.condition = threading.Condition()
+        self.queue = collections.deque()
+        self.active = False
+        self.closed = False
+        self.admitted = 0
+        self.completed = 0
+        self.rejected = 0
+        self.timed_out = 0
+        self.cancelled = 0
+
+    @contextlib.contextmanager
+    def admit(self, cancelled=None):
+        ticket = object()
+        queued_at = time.monotonic()
+        with self.condition:
+            if self.closed:
+                raise APIError(503, "The inference scheduler is shutting down.", None,
+                               "scheduler_closed", "server_error")
+            if (self.active or self.queue) and len(self.queue) >= self.max_queue:
+                self.rejected += 1
+                raise APIError(429, "The inference queue is full.", None, "queue_full",
+                               "rate_limit_error", {"Retry-After": "1"})
+            self.queue.append(ticket)
+            deadline = queued_at + self.queue_timeout
+            while True:
+                if self.closed:
+                    self.queue.remove(ticket)
+                    self.condition.notify_all()
+                    raise APIError(503, "The inference scheduler is shutting down.", None,
+                                   "scheduler_closed", "server_error")
+                if not self.active and self.queue[0] is ticket:
+                    break
+                if cancelled and cancelled():
+                    self.queue.remove(ticket)
+                    self.cancelled += 1
+                    self.condition.notify_all()
+                    raise ClientCancelled()
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    self.queue.remove(ticket)
+                    self.timed_out += 1
+                    self.condition.notify_all()
+                    raise APIError(429, "Timed out waiting for the inference engine.", None,
+                                   "queue_timeout", "rate_limit_error", {"Retry-After": "1"})
+                self.condition.wait(min(remaining, 0.25))
+            self.queue.popleft()
+            self.active = True
+            self.admitted += 1
+            wait_seconds = time.monotonic() - queued_at
+        try:
+            yield wait_seconds
+        finally:
+            with self.condition:
+                self.active = False
+                self.completed += 1
+                self.condition.notify_all()
+
+    def snapshot(self):
+        with self.condition:
+            return {"active": self.active, "queued": len(self.queue),
+                    "max_queue": self.max_queue, "queue_timeout_seconds": self.queue_timeout,
+                    "admitted": self.admitted, "completed": self.completed,
+                    "rejected": self.rejected, "timed_out": self.timed_out,
+                    "cancelled": self.cancelled}
+
+    def close(self):
+        with self.condition:
+            self.closed = True
+            self.condition.notify_all()
+
+
 def content_text(content, param):
     if isinstance(content, str):
         return content
@@ -205,12 +296,13 @@ class APIServer(ThreadingHTTPServer):
     daemon_threads = True
 
     def __init__(self, address, engine, model_id, api_key=None, max_tokens=1024,
-                 cors_origins=DEFAULT_CORS_ORIGINS):
+                 cors_origins=DEFAULT_CORS_ORIGINS, max_queue=8, queue_timeout=300):
         super().__init__(address, APIHandler)
         self.engine = engine
         self.model_id = model_id
         self.api_key = api_key
         self.max_tokens = max_tokens
+        self.scheduler = GenerationScheduler(max_queue, queue_timeout)
         self.cors_origins = tuple(cors_origins)
         self.created = int(time.time())
 
@@ -222,13 +314,15 @@ class APIHandler(BaseHTTPRequestHandler):
     def log_message(self, fmt, *args):
         sys.stderr.write("[api] %s - %s\n" % (self.address_string(), fmt % args))
 
-    def send_json(self, status, body, request_id=None):
+    def send_json(self, status, body, request_id=None, headers=None):
         data = json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode()
         self.send_response(status)
         self.send_header("Content-Type", "application/json")
         self.send_header("Content-Length", str(len(data)))
         if request_id:
             self.send_header("x-request-id", request_id)
+        for name, value in (headers or {}).items():
+            self.send_header(name, value)
         self.send_cors_headers()
         self.end_headers()
         self.wfile.write(data)
@@ -240,7 +334,8 @@ class APIHandler(BaseHTTPRequestHandler):
         self.send_header("Access-Control-Allow-Origin", "*" if "*" in self.server.cors_origins else origin)
         self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
         self.send_header("Access-Control-Allow-Headers", "Authorization, Content-Type")
-        self.send_header("Access-Control-Expose-Headers", "x-request-id")
+        self.send_header("Access-Control-Expose-Headers",
+                         "x-request-id, x-colibri-queue-wait-ms, Retry-After")
         self.send_header("Access-Control-Max-Age", "600")
         if "*" not in self.server.cors_origins:
             self.send_header("Vary", "Origin")
@@ -275,7 +370,7 @@ class APIHandler(BaseHTTPRequestHandler):
         try:
             path = urlsplit(self.path).path
             if path == "/health":
-                self.send_json(200, {"status": "ok"}, request_id)
+                self.send_json(200, {"status": "ok", "scheduler": self.server.scheduler.snapshot()}, request_id)
                 return
             self.require_auth()
             if path == "/v1/models":
@@ -286,7 +381,7 @@ class APIHandler(BaseHTTPRequestHandler):
             else:
                 raise APIError(404, "Not found.", None, "not_found")
         except APIError as error:
-            self.send_json(error.status, error_object(error), request_id)
+            self.send_json(error.status, error_object(error), request_id, error.headers)
 
     def do_OPTIONS(self):
         self.send_response(204)
@@ -308,7 +403,9 @@ class APIHandler(BaseHTTPRequestHandler):
             else:
                 raise APIError(404, "Not found.", None, "not_found")
         except APIError as error:
-            self.send_json(error.status, error_object(error), request_id)
+            self.send_json(error.status, error_object(error), request_id, error.headers)
+        except ClientCancelled:
+            pass
         except (BrokenPipeError, ConnectionResetError):
             pass
         except Exception as error:
@@ -325,11 +422,17 @@ class APIHandler(BaseHTTPRequestHandler):
         stream = body.get("stream", False)
         if not isinstance(stream, bool):
             raise APIError(400, "`stream` must be a boolean.", "stream")
+        stream_options = body.get("stream_options") if stream else None
+        if stream and stream_options is not None and not isinstance(stream_options, dict):
+            raise APIError(400, "`stream_options` must be an object.", "stream_options")
+        include_usage = bool((stream_options or {}).get("include_usage"))
         object_name = "chat.completion" if chat else "text_completion"
         id_prefix = "chatcmpl-" if chat else "cmpl-"
         completion_id = id_prefix + uuid.uuid4().hex
         created = int(time.time())
 
+        with self.server.scheduler.admit(self.client_disconnected) as queue_wait:
+            queue_headers = {"x-colibri-queue-wait-ms": str(round(queue_wait * 1000))}
             if not stream:
                 output = []
                 stats = self.server.engine.generate(prompt, maximum, temperature, top_p, output.append)
@@ -339,19 +442,17 @@ class APIHandler(BaseHTTPRequestHandler):
                            "refusal": None}, "logprobs": None, "finish_reason": finish} if chat else
                           {"index": 0, "text": text, "logprobs": None, "finish_reason": finish})
                 self.send_json(200, {"id": completion_id, "object": object_name, "created": created,
-                "model": self.server.model_id, "choices": [choice], "usage": self.usage(stats)}, request_id)
+                    "model": self.server.model_id, "choices": [choice], "usage": self.usage(stats)},
+                    request_id, queue_headers)
                 return
 
-        stream_options = body.get("stream_options")
-        if stream_options is not None and not isinstance(stream_options, dict):
-            raise APIError(400, "`stream_options` must be an object.", "stream_options")
-        include_usage = bool((stream_options or {}).get("include_usage"))
             stream_object = "chat.completion.chunk" if chat else object_name
             self.send_response(200)
             self.send_header("Content-Type", "text/event-stream")
             self.send_header("Cache-Control", "no-cache")
             self.send_header("X-Accel-Buffering", "no")
             self.send_header("x-request-id", request_id)
+            for name, value in queue_headers.items(): self.send_header(name, value)
             self.send_cors_headers()
             self.end_headers()
             connected = True
@@ -397,6 +498,16 @@ class APIHandler(BaseHTTPRequestHandler):
                     pass
             self.close_connection = True
 
+    def client_disconnected(self):
+        try:
+            readable, _, _ = select.select([self.connection], [], [], 0)
+            if not readable:
+                return False
+            flags = socket.MSG_PEEK | getattr(socket, "MSG_DONTWAIT", 0)
+            return self.connection.recv(1, flags) == b""
+        except (OSError, ValueError):
+            return True
+
     @staticmethod
     def usage(stats):
         prompt = stats["prompt_tokens"]
@@ -424,16 +535,21 @@ class APIHandler(BaseHTTPRequestHandler):
 
 
 def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_key=None,
-          cap=8, max_tokens=1024, engine=HERE / "glm", env=None, cors_origins=None):
+          cap=8, max_tokens=1024, engine=HERE / "glm", env=None, cors_origins=None,
+          max_queue=8, queue_timeout=300):
     if not 1 <= max_tokens:
         raise ValueError("max_tokens must be positive")
     if not 1 <= port <= 65535:
         raise ValueError("port must be between 1 and 65535")
+    if max_queue < 0:
+        raise ValueError("max_queue cannot be negative")
+    if queue_timeout <= 0:
+        raise ValueError("queue_timeout must be positive")
     if host not in ("127.0.0.1", "localhost", "::1") and not api_key:
         print("WARNING: API is listening beyond localhost without COLI_API_KEY", file=sys.stderr)
     runtime = Engine(engine, model, cap, max_tokens, env)
     origins = DEFAULT_CORS_ORIGINS if cors_origins is None else tuple(cors_origins)
-    server = APIServer((host, port), runtime, model_id, api_key, max_tokens, origins)
+    server = APIServer((host, port),runtime,model_id,api_key,max_tokens,origins,max_queue,queue_timeout)
     print(f"OpenAI-compatible API listening on http://{host}:{port}/v1", file=sys.stderr)
     previous_sigterm = signal.getsignal(signal.SIGTERM)
     signal.signal(signal.SIGTERM, lambda *_: threading.Thread(target=server.shutdown, daemon=True).start())
@@ -441,6 +557,7 @@ def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_ke
         server.serve_forever()
     finally:
         signal.signal(signal.SIGTERM, previous_sigterm)
+        server.scheduler.close()
         server.server_close()
         runtime.close()
 
@@ -457,9 +574,13 @@ def main():
                         help="allowed browser origin; repeat as needed (use '*' for any origin)")
     parser.add_argument("--cap", type=int, default=8)
     parser.add_argument("--max-tokens", type=int, default=1024)
+    parser.add_argument("--max-queue", type=int, default=int(os.environ.get("COLI_MAX_QUEUE", "8")))
+    parser.add_argument("--queue-timeout", type=float,
+                        default=float(os.environ.get("COLI_QUEUE_TIMEOUT", "300")))
     args = parser.parse_args()
     serve(args.model, args.host, args.port, args.model_id, args.api_key,
-          args.cap, args.max_tokens, args.engine, cors_origins=args.cors_origin)
+          args.cap,args.max_tokens,args.engine,cors_origins=args.cors_origin,
+          max_queue=args.max_queue,queue_timeout=args.queue_timeout)
 
 
 if __name__ == "__main__":

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()