diff --git a/README.md b/README.md index b83aed0..500614f 100644 --- a/README.md +++ b/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 diff --git a/c/coli b/c/coli index 56f25e6..88f407f 100755 --- a/c/coli +++ b/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") diff --git a/c/openai_server.py b/c/openai_server.py index d97bd64..48ad34a 100644 --- a/c/openai_server.py +++ b/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,77 +422,91 @@ 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()) - if not stream: - output = [] - stats = self.server.engine.generate(prompt, maximum, temperature, top_p, output.append) - text = "".join(output) - finish = "length" if stats["length_limited"] else "stop" - choice = ({"index": 0, "message": {"role": "assistant", "content": text, - "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) - 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) - self.send_cors_headers() - self.end_headers() - connected = True - - def event(choices, usage_marker=False): - nonlocal connected - if not connected: + 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) + text = "".join(output) + finish = "length" if stats["length_limited"] else "stop" + choice = ({"index": 0, "message": {"role": "assistant", "content": text, + "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, queue_headers) return - event_body = {"id": completion_id, "object": stream_object, "created": created, - "model": self.server.model_id, "choices": choices} + + 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 + + def event(choices, usage_marker=False): + nonlocal connected + if not connected: + return + event_body = {"id": completion_id, "object": stream_object, "created": created, + "model": self.server.model_id, "choices": choices} + if include_usage: + event_body["usage"] = None if not usage_marker else usage_marker + try: + data = json.dumps(event_body, ensure_ascii=False, separators=(",", ":")) + self.wfile.write(f"data: {data}\n\n".encode()) + self.wfile.flush() + except OSError: + connected = False + + if chat: + event([{"index": 0, "delta": {"role": "assistant", "content": ""}, + "logprobs": None, "finish_reason": None}]) + + def emit(text): + choice = ({"index": 0, "delta": {"content": text}, "logprobs": None, + "finish_reason": None} if chat else + {"index": 0, "text": text, "logprobs": None, "finish_reason": None}) + event([choice]) + + stats = self.server.engine.generate(prompt, maximum, temperature, top_p, emit) + finish = "length" if stats["length_limited"] else "stop" + final_choice = ({"index": 0, "delta": {}, "logprobs": None, "finish_reason": finish} + if chat else {"index": 0, "text": "", "logprobs": None, + "finish_reason": finish}) + event([final_choice]) if include_usage: - event_body["usage"] = None if not usage_marker else usage_marker - try: - data = json.dumps(event_body, ensure_ascii=False, separators=(",", ":")) - self.wfile.write(f"data: {data}\n\n".encode()) - self.wfile.flush() - except OSError: - connected = False + event([], self.usage(stats)) + if connected: + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except OSError: + pass + self.close_connection = True - if chat: - event([{"index": 0, "delta": {"role": "assistant", "content": ""}, - "logprobs": None, "finish_reason": None}]) - - def emit(text): - choice = ({"index": 0, "delta": {"content": text}, "logprobs": None, - "finish_reason": None} if chat else - {"index": 0, "text": text, "logprobs": None, "finish_reason": None}) - event([choice]) - - stats = self.server.engine.generate(prompt, maximum, temperature, top_p, emit) - finish = "length" if stats["length_limited"] else "stop" - final_choice = ({"index": 0, "delta": {}, "logprobs": None, "finish_reason": finish} - if chat else {"index": 0, "text": "", "logprobs": None, - "finish_reason": finish}) - event([final_choice]) - if include_usage: - event([], self.usage(stats)) - if connected: - try: - self.wfile.write(b"data: [DONE]\n\n") - self.wfile.flush() - except OSError: - 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): @@ -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__": diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py index 06382ea..2ddc241 100644 --- a/c/tests/test_openai_server.py +++ b/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|>", 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()