From 4574c724de8a93fdf89f8bda638e98ddfc0a87a9 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 02:08:55 +0530 Subject: [PATCH 001/135] serve: /tokenize and /detokenize, as SGLang and vLLM serve them, counting a chat request as /v1/chat/completions sends it --- trnrt/serve/api.py | 55 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/trnrt/serve/api.py b/trnrt/serve/api.py index 4b4ab1b..341a026 100644 --- a/trnrt/serve/api.py +++ b/trnrt/serve/api.py @@ -1019,4 +1019,59 @@ async def chat(http: HTTPRequest): parse=parse, ) + # SGLang's and vLLM's /tokenize and /detokenize (also under /v1). A chat request is rendered exactly as + # /v1/chat/completions renders it (chat_prompt), so `count` is what the engine would be sent; a prompt is encoded as + # /v1/completions encodes it. Both run on the prompt thread with its tokenizer, as a chat prompt does: never on the + # event loop, and never in a second thread inside one tokenizer. + def tokenize(body: dict) -> dict: + if (body.get("prompt") is None) == (body.get("messages") is None): + raise _bad("exactly one of prompt or messages is required") + if body.get("messages") is not None: + # The fields chat refuses unless neutral (_NEUTRAL) and that change the rendered prompt. + for key in ("add_generation_prompt", "continue_final_message"): + if body.get(key) not in _NEUTRAL[key]: + raise _bad(f"{key} is not supported; send the request without it") + ids, _ = chat_prompt(prompt_tok, body) + seqs, batch = [ids], False + else: + prompt = body["prompt"] + batch = isinstance(prompt, list) + texts = prompt if batch else [prompt] + if not texts or not all(isinstance(t, str) for t in texts): + raise _bad("prompt must be a string or a non-empty list of strings") + special = _flag(body, "add_special_tokens", True) + seqs = [list(prompt_tok(t, add_special_tokens=special)["input_ids"]) for t in texts] + out: dict = { + "tokens": seqs if batch else seqs[0], + "count": [len(s) for s in seqs] if batch else len(seqs[0]), + "max_model_len": engine.cfg.max_model_len, + } + if _flag(body, "return_token_strs"): # vLLM's + strs = [prompt_tok.convert_ids_to_tokens(s) for s in seqs] + out["token_strs"] = strs if batch else strs[0] + return out + + def detokenize(body: dict) -> dict: + tokens = body.get("tokens") + batch = isinstance(tokens, list) and bool(tokens) and all(isinstance(s, list) for s in tokens) + seqs: list = cast("list", tokens) if batch else [tokens] + if not all(isinstance(s, list) and all(type(t) is int and 0 <= t < loop_.vocab for t in s) for s in seqs): + raise _bad(f"tokens must be a list of token ids in [0, {loop_.vocab}), or a list of such lists") + skip = _flag(body, "skip_special_tokens", True) # SGLang's default + texts = [prompt_tok.decode(s, skip_special_tokens=skip) for s in seqs] + text = texts if batch else texts[0] + return {"text": text, "prompt": text} # SGLang answers `text`, vLLM `prompt` + + @app.post("/tokenize") + @app.post("/v1/tokenize") + async def tokenize_route(http: HTTPRequest): + body = await _body(http) + return await asyncio.get_running_loop().run_in_executor(prompt_pool, tokenize, body) + + @app.post("/detokenize") + @app.post("/v1/detokenize") + async def detokenize_route(http: HTTPRequest): + body = await _body(http) + return await asyncio.get_running_loop().run_in_executor(prompt_pool, detokenize, body) + return app From c291db0bae998b1586595cb1afc7168f908ea4a4 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 02:08:55 +0530 Subject: [PATCH 002/135] router: send /tokenize and /detokenize to any live engine, outside the prefix trees and sessions --- trnrt/serve/router.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/trnrt/serve/router.py b/trnrt/serve/router.py index 3085e1b..03ab1ec 100644 --- a/trnrt/serve/router.py +++ b/trnrt/serve/router.py @@ -87,6 +87,9 @@ LOG_SESSIONS_ENV = "TRNRT_ROUTER_LOG_SESSIONS" SESSION_KEY_CHARS = 4096 ROUTES = ("/v1/completions", "/v1/chat/completions") +# Every worker answers these the same (one model, one tokenizer), so they go to the least loaded live worker and touch +# no prefix tree and no session: a /tokenize of an agent's context must not steer where its next turn is routed. +STATELESS_ROUTES = ("/tokenize", "/v1/tokenize", "/detokenize", "/v1/detokenize") class _Node: @@ -691,8 +694,34 @@ async def handler(http: Request): # annotated so FastAPI passes the request return handler + async def stateless(path: str, http: Request): + """STATELESS_ROUTES: the first live worker that answers, least loaded first. A refused connection or a 503 (a + failed engine) marks the worker down, as in a request, and the next one is asked.""" + body = await _json_object(http) + for w in sorted(router._live(), key=lambda i: (router.load[i], i)): + try: + r = await client.post(router.workers[w] + path, json=body) + except httpx.ConnectError: + router.note_failure(w, hard=True) + continue + except httpx.HTTPError: + continue + if r.status_code == 503: + router.note_failure(w, hard=True) + continue + return Response(r.content, status_code=r.status_code, media_type=r.headers.get("content-type")) + raise HTTPException(status_code=503, detail=f"no engine answered {path}") + + def stateless_route(path: str): + async def handler(http: Request): + return await stateless(path, http) + + return handler + for path in ROUTES: app.add_api_route(path, route(path), methods=["POST"]) + for path in STATELESS_ROUTES: + app.add_api_route(path, stateless_route(path), methods=["POST"]) return app From a5cccf3099d3076995651c6c2675f90a369f52d8 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 02:08:55 +0530 Subject: [PATCH 003/135] tests: /tokenize counts what the engine would be sent, /detokenize answers SGLang and vLLM clients, both refuse with 400 --- tests/test_serve_api.py | 80 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/tests/test_serve_api.py b/tests/test_serve_api.py index ebce5ad..66abcd4 100644 --- a/tests/test_serve_api.py +++ b/tests/test_serve_api.py @@ -628,3 +628,83 @@ def params_with(gen: dict | None, body: dict, path: str = "/v1/chat/completions" assert params_with({"temperature": 0.6, "top_k": 20, "min_p": 0.05}, msgs) == (0.6, 1.0, 20, 0.05) assert params_with({"top_p": 0, "temperature": -1, "top_k": "x"}, msgs) == (1.0, 1.0, 0, 0.0) assert params_with(None, msgs) == (1.0, 1.0, 0, 0.0) + + +@pytest.mark.parametrize("prefix", ["", "/v1"]) +def test_tokenize_counts_what_the_engine_would_be_sent(serve, prefix): + """A chat request's tokens are the prompt /v1/chat/completions sends the engine for the same body (tools and + template arguments included); a prompt's are /v1/completions' encoding. No request reaches the engine.""" + client, eng = serve(PLAIN) + chat = { + "messages": [{"role": "user", "content": "weather?"}], + "tools": TOOLS, + "chat_template_kwargs": {"enable_thinking": False}, + } + r = client.post(prefix + "/tokenize", json=chat) + assert r.status_code == 200, r.text + assert eng.added == [] + assert client.post("/v1/chat/completions", json={**chat, "max_tokens": 2}).status_code == 200 + sent = eng.added[-1].prompt_ids + assert r.json() == {"tokens": sent, "count": len(sent), "max_model_len": 64} + + r = client.post(prefix + "/tokenize", json={"prompt": "Hello !", "return_token_strs": True}) + assert r.json() == {"tokens": [ID["Hello"], ID["!"]], "count": 2, "max_model_len": 64, "token_strs": ["Hello", "!"]} + r = client.post(prefix + "/tokenize", json={"prompt": ["Hello", "Hello !"], "add_special_tokens": False}) + assert r.json() == {"tokens": [[ID["Hello"]], [ID["Hello"], ID["!"]]], "count": [1, 2], "max_model_len": 64} + + +@pytest.mark.parametrize( + "body", + [ + {}, + {"prompt": "Hello", "messages": [{"role": "user", "content": "hi"}]}, + {"prompt": [ID["Hello"]]}, # token ids: nothing to tokenize + {"prompt": []}, + {"prompt": ["Hello", 3]}, + {"prompt": "Hello", "add_special_tokens": "yes"}, + {"messages": []}, + {"messages": [{"role": "user", "content": "hi"}], "continue_final_message": True}, + {"messages": [{"role": "user", "content": "hi"}], "add_generation_prompt": False}, + {"messages": [{"role": "user", "content": "hi"}], "chat_template_kwargs": {"unread": 1}}, + ], +) +def test_tokenize_refuses_what_it_cannot_answer_with_a_400(serve, body): + client, eng = serve(PLAIN) + r = client.post("/tokenize", json=body) + assert r.status_code == 400, r.text + assert eng.added == [] + + +@pytest.mark.parametrize("prefix", ["", "/v1"]) +def test_detokenize_answers_sglang_and_vllm_clients(prefix): + from fastapi.testclient import TestClient + + from trnrt.serve import api + + eng = ScriptedEngine(PLAIN) + # registered as a special token (its id stays 0), so skip_special_tokens has one to drop + eng.tokenizer.add_special_tokens({"unk_token": ""}) + with TestClient(api.build_app(eng, "glm")) as client: + _detokenize_checks(client, prefix) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def _detokenize_checks(client, prefix: str) -> None: + r = client.post(prefix + "/detokenize", json={"tokens": [ID[""], ID["Hello"], ID[" there"]]}) + assert r.status_code == 200, r.text + assert r.json() == {"text": "Hello there", "prompt": "Hello there"} # is special: skipped by default + r = client.post(prefix + "/detokenize", json={"tokens": [ID[""], ID["Hello"]], "skip_special_tokens": False}) + assert r.json()["text"] == "Hello" + r = client.post(prefix + "/detokenize", json={"tokens": [[ID["Hello"]], [ID["Hello"], ID["!"]]]}) + assert r.json() == {"text": ["Hello", "Hello!"], "prompt": ["Hello", "Hello!"]} + assert client.post(prefix + "/detokenize", json={"tokens": []}).json()["text"] == "" + + +@pytest.mark.parametrize( + "tokens", [None, "Hello", [-1], [len(WORDS)], [True], [1.0], [[1], 2], [[1], [len(WORDS)]]], ids=repr +) +def test_detokenize_refuses_ids_outside_the_vocabulary(serve, tokens): + client, _ = serve(PLAIN) + r = client.post("/detokenize", json={} if tokens is None else {"tokens": tokens}) + assert r.status_code == 400, r.text From 6ae262751cd39e1826344f3021111d66e83b802c Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 02:08:55 +0530 Subject: [PATCH 004/135] tests: the router's tokenize routes fail over and change no routing state --- tests/test_router.py | 56 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 50 insertions(+), 6 deletions(-) diff --git a/tests/test_router.py b/tests/test_router.py index 1e55155..9060b71 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -10,7 +10,7 @@ for _m in ("httpx", "fastapi", "uvicorn"): # the router's stack pytest.importorskip(_m) -from trnrt.serve.router import PrefixTree, Router # noqa: E402 +from trnrt.serve.router import STATELESS_ROUTES, PrefixTree, Router # noqa: E402 def test_prefix_tree_match_insert_and_prune(): @@ -167,13 +167,13 @@ async def _stop(servers, tasks, i=None): def _stub_worker(i: int, chat_status: int = 200): - """A worker that answers /health, /v1/models and /v1/chat/completions (with chat_status: 503 is a failed engine, - whose serve/api.py refuses every request before it queues it), streaming or not.""" + """A worker that answers /health, /v1/models, /v1/chat/completions and STATELESS_ROUTES (with chat_status: 503 is a + failed engine, whose serve/api.py refuses every request before it queues it), streaming or not.""" from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse app = FastAPI() - app.state.chats, app.state.status = 0, chat_status + app.state.chats, app.state.toks, app.state.status = 0, 0, chat_status @app.get("/health") async def health(): @@ -200,6 +200,14 @@ async def stream(): return StreamingResponse(stream(), media_type="text/event-stream") return {"choices": [{"message": {"role": "assistant", "content": f"w{i}"}}]} + async def stateless(req: Request): + app.state.toks += 1 + if app.state.status != 200: + raise HTTPException(status_code=app.state.status, detail="engine failed: rank 3 exited") + return {"worker": i, "body": await req.json()} + + for path in STATELESS_ROUTES: + app.add_api_route(path, stateless, methods=["POST"]) return app @@ -255,7 +263,7 @@ async def main() -> None: asyncio.run(main()) -@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/completions"]) +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/completions", "/tokenize", "/v1/detokenize"]) @pytest.mark.parametrize("body", [b"{not json", b"[1, 2]", b'"text"', b"null", b"true", b"3", b"", b"\xff\xfe"]) def test_a_body_that_is_not_a_json_object_is_answered_400_and_reaches_no_worker(path, body): """The router refuses a malformed body, or JSON that is not an object, with 400, as each engine does, instead of @@ -277,7 +285,43 @@ async def main() -> None: ) assert r.status_code == 400, r.text assert r.json()["detail"] == "the request body must be a JSON object" - assert router.load == [0] and worker.state.chats == 0 + assert router.load == [0] and worker.state.chats == 0 and worker.state.toks == 0 + finally: + await _stop(servers, tasks) + + asyncio.run(main()) + + +@pytest.mark.parametrize("path", STATELESS_ROUTES) +def test_tokenize_and_detokenize_go_to_a_live_worker_and_change_no_routing_state(path): + """STATELESS_ROUTES go to the least loaded live worker and touch no prefix tree, no session and no load count: a + /tokenize of an agent's context must not steer where its next turn goes. A refused connection or a 503 marks the + worker down and the next one is asked; 503 once none answers.""" + import asyncio + + import httpx + + from trnrt.serve.router import build_router_app + + async def main() -> None: + ports = [_free_port() for _ in range(4)] # ports[0] serves nothing: a dead worker + router = Router([f"http://127.0.0.1:{p}" for p in ports[:3]]) + failed, ok = _stub_worker(1, chat_status=503), _stub_worker(2) + servers, tasks = await _start( + [(failed, ports[1]), (ok, ports[2]), (build_router_app(router, poll_s=0), ports[3])] + ) + body = {"messages": [{"role": "user", "content": "x" * 64}], "session_id": "s1"} + try: + async with httpx.AsyncClient(timeout=30) as c: + r = await c.post(f"http://127.0.0.1:{ports[3]}{path}", json=body) + assert r.status_code == 200, r.text + assert r.json() == {"worker": 2, "body": body} + assert router.up == [False, False, True] # refused, then 503: both marked down + assert failed.state.toks == 1 and ok.state.toks == 1 + assert router.sessions == {} and all(t.size == 0 for t in router.trees) and router.load == [0, 0, 0] + await _stop(servers, tasks, 1) + r = await c.post(f"http://127.0.0.1:{ports[3]}{path}", json=body) + assert r.status_code == 503 and r.json()["detail"] == f"no engine answered {path}", r.text finally: await _stop(servers, tasks) From 6d9624c5b16b25c7a64df96f3af345eb037e5a3b Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 02:08:55 +0530 Subject: [PATCH 005/135] docs: the tokenize and detokenize endpoints --- docs/API.md | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/docs/API.md b/docs/API.md index 3b401a0..7cd2d4b 100644 --- a/docs/API.md +++ b/docs/API.md @@ -1,6 +1,6 @@ # API -The router on `--port` (default 8000) speaks the OpenAI API: `/v1/chat/completions`, `/v1/completions`, `/v1/models` and `/health`. Any OpenAI client works with `base_url="http://:8000/v1"`. +The router on `--port` (default 8000) speaks the OpenAI API: `/v1/chat/completions`, `/v1/completions`, `/v1/models` and `/health`, plus `/tokenize` and `/detokenize` ([below](#tokenize-and-detokenize)). Any OpenAI client works with `base_url="http://:8000/v1"`. ## Example @@ -30,7 +30,8 @@ curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' - - A response cut off by `max_tokens` has `finish_reason` `"length"`, even when it ends in a tool call: that call may be incomplete. - A client that disconnects aborts its request, streamed or not. -- Endpoints on `--port` (the router): `/v1/chat/completions`, `/v1/completions`, `/v1/models` and `/health`. Each +- Endpoints on `--port` (the router): `/v1/chat/completions`, `/v1/completions`, `/v1/models`, `/health`, `/tokenize` + and `/detokenize`. Each engine also serves Prometheus metrics at `/metrics` on its own port (8001 and 8002 with the default `--port`), on 127.0.0.1 only, among them the sparse-attention selection's counters (`trnrt:dsa_local_k_*`, [KNOWN_ISSUES.md](KNOWN_ISSUES.md)). @@ -45,3 +46,23 @@ curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' - - Status codes: 400 for a request that cannot be served as sent (a field of the wrong type or out of range, a prompt longer than the context, a chat template variable the template never reads); 500 for an error inside the engine; 503 once the engine has failed, with the reason on `/health`. + +## Tokenize and detokenize + +`/tokenize` and `/detokenize` (also under `/v1`) answer as SGLang and vLLM do. Neither runs the model, and the router +sends them to any live engine. + +| Endpoint | Request | Response | +|---|---|---| +| `/tokenize` | exactly one of `prompt` (a string, or a list of strings for a batch) or `messages` (with the chat fields `/v1/chat/completions` takes: `tools`, `tool_choice`, `chat_template_kwargs`, `reasoning_effort`); `add_special_tokens` (prompt only, default `true`); `return_token_strs` | `tokens`, `count`, `max_model_len`, and `token_strs` when asked | +| `/detokenize` | `tokens`: a list of token ids, or a list of such lists for a batch; `skip_special_tokens` (default `true`) | the text as `text` (SGLang's name) and as `prompt` (vLLM's) | + +A chat request is counted exactly as `/v1/chat/completions` would send it to the engine, rendered by the model's chat +template, so `count` is the prompt length that request would use against `--max-model-len`. Fields the chat endpoint +refuses that would change the rendering (`add_generation_prompt` false, `continue_final_message` true) are refused here +too. + +```bash +curl -s localhost:8000/tokenize -H 'Content-Type: application/json' \ + -d '{"messages": [{"role": "user", "content": "What is the weather in Paris?"}]}' +``` From 6c4742b83ba1744e93fa5fca1a8273bd1c9ecb5c Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 02:08:55 +0530 Subject: [PATCH 006/135] changelog: /tokenize and /detokenize --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c1aaa3..0748e1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,12 @@ and versions follow [Semantic Versioning](https://semver.org/). Rebuild the comp ## [Unreleased] +### Added + +- `/tokenize` and `/detokenize` (also under `/v1`), as SGLang and vLLM serve them. A chat request is counted exactly as + `/v1/chat/completions` would send it to the engine. The router sends them to any live engine without touching its + routing state ([docs/API.md](docs/API.md)). + ## [0.1.0] - 2026-11-08 The first public release. From 78a66575c556ec04f0a060ded6dda39681da30e5 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:04 +0530 Subject: [PATCH 007/135] radix_cache: evict by SGLang's lfu, slru and priority policies, and tlru --- trnrt/engine/radix_cache.py | 92 ++++++++++++++++++++++++++++++++++--- 1 file changed, 86 insertions(+), 6 deletions(-) diff --git a/trnrt/engine/radix_cache.py b/trnrt/engine/radix_cache.py index dff8adb..2e3c14b 100644 --- a/trnrt/engine/radix_cache.py +++ b/trnrt/engine/radix_cache.py @@ -9,7 +9,8 @@ - A page referenced by a node belongs to the tree. It is freed only by eviction. - A running request holds a lock on the deepest node of its cached prefix, which pins every page on the path from the root to that node. -- Eviction frees whole unlocked leaves in least-recently-used order. +- Eviction frees whole unlocked leaves in the eviction policy's order (EVICTION_POLICIES), least recently used first by + default. - Session references (SGLang v0.5.17 "session-reference-aware radix cache", sgl-project/sglang #29173): a finished request that carries a session_id registers its prefix under that session; eviction takes unreferenced prefixes first, then fewer-referenced ones, each least @@ -33,6 +34,7 @@ import heapq import itertools +import math from collections import OrderedDict from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -59,6 +61,12 @@ class RadixNode: # The checkpoint ends a request's prompt (the prompt-end checkpoint, engine/scheduler.py): a returning turn whose # replayed reply differs from the cached one resumes there, so free_dominated keeps the nearest such one on a path. anchor: bool = False + # The eviction policies' bookkeeping (RadixCache.touch), read only by the policy that uses it: requests admitted on a + # path through this node (lfu, slru; a node starts at 1, for the request that inserted it), the highest importance of + # those requests (priority), and the tokens already evicted below it along its deepest evicted branch (tlru). + hit_count: int = 1 + importance: float = -math.inf + tail_evicted: int = 0 def is_leaf(self) -> bool: return not self.children @@ -74,8 +82,56 @@ def num_pages(self) -> int: return len(self.pages) +# SGLang's --radix-eviction-policy names (server_args.py RADIX_EVICTION_POLICY_CHOICES, mem_cache/evict_policy.py), plus +# tlru. Each policy's own settings (--radix-eviction-policy-config) and their defaults; tlru's have none. +EVICTION_POLICIES = ("lru", "lfu", "slru", "priority", "tlru") +POLICY_PARAMS: dict[str, dict[str, int | None]] = { + "slru": {"protected_threshold": 2}, # SGLang SLRUStrategy's default + "tlru": {"threshold": None, "next_prompt_estimate": None}, +} + + +def policy_params(policy: str, config: dict | None = None) -> dict: + """An eviction policy's settings: its defaults (POLICY_PARAMS) updated from `config`. ValueError for an unknown + policy, a setting it does not take, a value that is not a non-negative integer, or tlru without its two settings.""" + if policy not in EVICTION_POLICIES: + raise ValueError(f"the eviction policy must be one of {', '.join(EVICTION_POLICIES)}") + params = dict(POLICY_PARAMS.get(policy, {})) + for key, val in (config or {}).items(): + if key not in params: + raise ValueError(f"{policy} eviction has no setting {key!r} (its settings: {', '.join(params) or 'none'})") + if type(val) is not int or val < 0: + raise ValueError(f"{policy} eviction's {key} must be a non-negative integer") + params[key] = val + if policy == "tlru": + th, est = params["threshold"], params["next_prompt_estimate"] + if th is None or est is None or not 0 <= est < th: + raise ValueError("tlru eviction needs threshold > next_prompt_estimate >= 0 (tokens)") + return params + + class RadixCache: - def __init__(self, pool: PagePool, page_size: int): + """`eviction_policy` orders what eviction (evict, evict_checkpoints) takes first, within each session-reference + class (unreferenced first, then fewer references): + - lru (the default): the prefix unused for longest; + - lfu: the fewest hits (requests admitted on a path through it), then the least recent (SGLang LFUStrategy); + - slru: every probationary prefix (fewer than protected_threshold hits) before any protected one, the least recent + first within each (SGLang SLRUStrategy); + - priority: the lowest importance of the requests that used it, then the least recent (SGLang PriorityStrategy); + - tlru (Tail-Optimized LRU, arXiv 2510.15152): first the tail of each cached sequence that its next turn can + recompute and still prefill at most `threshold` uncached tokens (threshold - next_prompt_estimate tokens from its + end), then plain recency. + `policy_config` holds the policy's own settings (POLICY_PARAMS); a key the policy does not take is an error.""" + + def __init__(self, pool: PagePool, page_size: int, eviction_policy: str = "lru", policy_config: dict | None = None): + params = policy_params(eviction_policy, policy_config) + self.tel_budget = 0 + if eviction_policy == "tlru": + self.tel_budget = int(params["threshold"] or 0) - int(params["next_prompt_estimate"] or 0) + self.eviction_policy = eviction_policy + self.policy_params = params + # Whether admissions and inserts record hits and importance (touch): only for the policies that read them. + self.tracks_use = eviction_policy in ("lfu", "slru", "priority") self.pool = pool self.page_size = page_size self.root = RadixNode() @@ -108,11 +164,31 @@ def __init__(self, pool: PagePool, page_size: int): def _page_key(self, ids: list[int], i: int) -> tuple[int, ...]: return tuple(ids[i : i + self.page_size]) - @staticmethod - def _key(n: RadixNode): + def _key(self, n: RadixNode): """Eviction order, smallest first: unreferenced before session-referenced, fewer references - first, then least recently used.""" - return (n.session_refs > 0, n.session_refs, n.last_access) + first, then the eviction policy's order (lru: least recently used).""" + policy = self.eviction_policy + if policy == "lru": + return (n.session_refs > 0, n.session_refs, n.last_access) + refs = (n.session_refs > 0, n.session_refs) + if policy == "lfu": + return (*refs, n.hit_count, n.last_access) + if policy == "slru": + return (*refs, n.hit_count >= self.policy_params["protected_threshold"], n.last_access) + if policy == "priority": + return (*refs, n.importance, n.last_access) + safe = n.tail_evicted + len(n.key) <= self.tel_budget # tlru: the next turn can recompute it within budget + return (*refs, not safe, n.last_access) + + def touch(self, node: RadixNode, importance: float, hit: bool = True) -> None: + """A request of `importance` used the path to `node`: a hit on every node of it (hit), and each node's + importance raised to the request's. Called only when tracks_use.""" + while node is not self.root: + if hit: + node.hit_count += 1 + if importance > node.importance: + node.importance = importance + node = node.parent # type: ignore[assignment] def match_prefix(self, ids: list[int], max_pages: int | None = None) -> MatchResult: """Longest page-aligned cached prefix of `ids`, at most `max_pages` pages. @@ -156,6 +232,8 @@ def _split(self, child: RadixNode, n_pages: int) -> RadixNode: lock_ref=child.lock_ref, last_access=child.last_access, session_refs=child.session_refs, # registrations stay on the deeper node + hit_count=child.hit_count, + importance=child.importance, ) parent = child.parent if parent is None: @@ -191,6 +269,7 @@ def insert(self, ids: list[int], pages: list[int]) -> MatchResult: last_access=next(self._clock), ) node.children[tuple(leaf.key[:ps])] = leaf + node.tail_evicted = 0 # tlru: the sequence grew again, so its tail evicted earlier is no longer its tail self.evictable_pages += len(leaf.pages) tree_pages.extend(leaf.pages) node = leaf @@ -325,6 +404,7 @@ def evict(self, num_pages: int) -> int: if parent is not self.root: self._register(sid, parent) del parent.children[tuple(leaf.key[: self.page_size])] + parent.tail_evicted = max(parent.tail_evicted, leaf.tail_evicted + len(leaf.key)) leaf.parent = None if parent is not self.root and parent.is_leaf() and parent.lock_ref == 0: heapq.heappush(heap, (self._key(parent), id(parent), parent)) From 6499230c07779047f080c108569cb9d263449867 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:04 +0530 Subject: [PATCH 008/135] tests: each eviction policy decides against recency, and lru is unchanged --- tests/test_radix_eviction.py | 147 +++++++++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 tests/test_radix_eviction.py diff --git a/tests/test_radix_eviction.py b/tests/test_radix_eviction.py new file mode 100644 index 0000000..79c6880 --- /dev/null +++ b/tests/test_radix_eviction.py @@ -0,0 +1,147 @@ +"""The prefix cache's eviction policies (engine/radix_cache.py EVICTION_POLICIES, SGLang's --radix-eviction-policy). + +Each test builds prefixes whose recency alone orders them one way and whose policy key orders them the other, so the +policy is what decides; lru on the same tree is the control.""" + +from __future__ import annotations + +import pytest + +from trnrt.engine.kv_pool import PagePool +from trnrt.engine.radix_cache import RadixCache, policy_params + +PS = 4 + + +def tree(policy="lru", config=None, n=3): + """n two-page prefixes [k]*8 for k = 1 .. n, inserted (so last used) in that order: lru evicts 1 first.""" + pool = PagePool(64) + rc = RadixCache(pool, PS, policy, config) + nodes = [rc.insert([k] * 8, pool.alloc(2)).node for k in range(1, n + 1)] + return rc, nodes + + +def evicted_order(rc, n) -> list[int]: + """Evict one prefix at a time: which ones went, in order (by their token).""" + out = [] + for _ in range(n): + left = {k for k in range(1, n + 1) if rc.match_prefix([k] * 8, max_pages=2).num_pages} + rc.evict(1) + gone = left - {k for k in left if rc.match_prefix([k] * 8, max_pages=2).num_pages} + out.extend(sorted(gone)) + return out + + +def test_lru_control(): + rc, _ = tree() + assert evicted_order(rc, 3) == [1, 2, 3] + + +def test_lfu_evicts_the_fewest_hits_first(): + rc, (a, _b, c) = tree("lfu") + for _ in range(3): + rc.touch(a, 0) + rc.touch(c, 0) + assert evicted_order(rc, 3) == [2, 3, 1] + + +def test_slru_evicts_every_probationary_prefix_before_a_protected_one(): + rc, (a, _, _) = tree("slru") # protected from 2 hits: a node starts at 1 (its insert), one touch protects it + rc.touch(a, 0) + assert evicted_order(rc, 3) == [2, 3, 1] + rc, (a, _b, _c) = tree("slru", {"protected_threshold": 3}) + rc.touch(a, 0) # 2 hits: still probationary under a threshold of 3 + assert evicted_order(rc, 3) == [1, 2, 3] + + +def test_priority_evicts_the_least_important_first(): + rc, (a, b, _c) = tree("priority") + rc.touch(a, 5, hit=False) + rc.touch(b, -1, hit=False) # c: never touched, the lowest importance of all + assert evicted_order(rc, 3) == [3, 2, 1] + + +def test_tlru_evicts_the_tail_a_next_turn_can_recompute_first(): + """Budget 8 tokens (threshold 10, next prompt 2): a 2-page leaf (8 tokens) is safe, a 3-page one (12) is not. The + older 3-page prefix outlives the newer 2-page one, where lru takes it first (the control).""" + for policy, config, kept in (("tlru", {"threshold": 10, "next_prompt_estimate": 2}, 1), ("lru", None, 2)): + pool = PagePool(64) + rc = RadixCache(pool, PS, policy, config) + rc.insert([1] * 12, pool.alloc(3)) # the older + rc.insert([2] * 8, pool.alloc(2)) + rc.evict(1) + left = [k for k, n in ((1, 12), (2, 8)) if rc.match_prefix([k] * n).num_pages] + assert left == [kept], (policy, left) + + +def test_tlru_chain_counts_the_evicted_tail(): + pool = PagePool(64) + rc = RadixCache(pool, PS, "tlru", {"threshold": 10, "next_prompt_estimate": 2}) # budget 8 tokens + chain = [3] * 4 + [4] * 4 + [5] * 4 + pages = pool.alloc(3) + for n in (1, 2, 3): # three nodes of one page each, the deepest inserted last + rc.insert(chain[: 4 * n], pages[:n]) + rc.insert([1] * 12, pool.alloc(3)) # the newest, 12 tokens: unsafe from the start + rc.evict(1) # the chain's leaf (4 tokens, safe) goes first, though it is older than the long one + assert rc.total_pages() == 5 + rc.evict(1) # its parent: 4 + the 4 evicted below = 8 tokens, still safe + assert rc.total_pages() == 4 + rc.evict(1) # the chain's first page: 4 + 8 = 12, unsafe as the long one is; by recency the chain is older + assert rc.total_pages() == 3 and rc.match_prefix([1] * 12).num_pages == 3 + + +def test_a_sequence_that_grows_again_resets_its_evicted_tail(): + pool = PagePool(64) + rc = RadixCache(pool, PS, "tlru", {"threshold": 10, "next_prompt_estimate": 2}) + rc.insert([3] * 4, pool.alloc(1)) + rc.insert([3] * 4 + [4] * 4, [rc.match_prefix([3] * 4).pages[0], *pool.alloc(1)]) + rc.evict(1) + parent = rc.match_prefix([3] * 4).node + assert parent.tail_evicted == 4 + rc.insert([3] * 4 + [6] * 4, [parent.pages[0], *pool.alloc(1)]) + assert parent.tail_evicted == 0 + + +def test_session_references_still_come_first_under_every_policy(): + for policy, config in ( + ("lfu", None), + ("slru", None), + ("priority", None), + ("tlru", {"threshold": 9, "next_prompt_estimate": 0}), + ): + rc, (a, _b, c) = tree(policy, config) + gen = rc.session_open("s") + assert rc.session_register("s", gen, c) + if rc.tracks_use: + rc.touch(a, 9) + order = evicted_order(rc, 3) + assert order[-1] == 3, (policy, order) + + +def test_policy_settings_are_checked(): + assert policy_params("lru") == {} and policy_params("slru") == {"protected_threshold": 2} + assert policy_params("slru", {"protected_threshold": 5}) == {"protected_threshold": 5} + with pytest.raises(ValueError, match="one of"): + policy_params("fifo") + with pytest.raises(ValueError, match="no setting 'threshold'"): + policy_params("slru", {"threshold": 3}) + with pytest.raises(ValueError, match="no setting"): + policy_params("lru", {"protected_threshold": 3}) + with pytest.raises(ValueError, match="non-negative integer"): + policy_params("slru", {"protected_threshold": -1}) + with pytest.raises(ValueError, match="non-negative integer"): + policy_params("slru", {"protected_threshold": 2.5}) + with pytest.raises(ValueError, match="tlru eviction needs"): + policy_params("tlru") + with pytest.raises(ValueError, match="tlru eviction needs"): + policy_params("tlru", {"threshold": 4, "next_prompt_estimate": 4}) + with pytest.raises(ValueError, match="one of"): + RadixCache(PagePool(8), PS, "mru") + + +def test_lru_keeps_its_order_key_and_records_nothing(): + """The default's eviction key is the one it had before the policies existed (session references, then recency), and + it never tracks hits or importance: an engine without the flags evicts exactly as before.""" + rc, (a, _b, _c) = tree() + assert not rc.tracks_use + assert rc._key(a) == (a.session_refs > 0, a.session_refs, a.last_access) From 90f3489429972ba57254c116622880edf9726022 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:04 +0530 Subject: [PATCH 009/135] request: carry a priority, and the error of a request ended unserved --- trnrt/engine/request.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/trnrt/engine/request.py b/trnrt/engine/request.py index 79ca99b..d39ae67 100644 --- a/trnrt/engine/request.py +++ b/trnrt/engine/request.py @@ -59,6 +59,10 @@ class Request: # session (soft eviction protection) until the session is closed. session_id: str | None = None session_gen: int = -1 + # SGLang's per-request priority (engine/scheduler.py: priority scheduling, the priority eviction); None: none given. + priority: int | None = None + # (HTTP status, message) when the scheduler ended the request without serving it (a full queue's displacement). + error: tuple[int, str] | None = None # Prompt followed by generated tokens. After a preemption the whole list is the # "prompt" of the recompute, so it is kept as one list. From 9ec5c03ddc8f1405c447aecb8c5273c452807d2e Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:04 +0530 Subject: [PATCH 010/135] scheduler: admission policies, priority scheduling and queue limits --- trnrt/engine/scheduler.py | 218 +++++++++++++++++++++++++++++++++++++- 1 file changed, 213 insertions(+), 5 deletions(-) diff --git a/trnrt/engine/scheduler.py b/trnrt/engine/scheduler.py index 60d3765..4bd565d 100644 --- a/trnrt/engine/scheduler.py +++ b/trnrt/engine/scheduler.py @@ -34,11 +34,21 @@ running request becomes a checkpoint target of the one computing it first (at its admission, or, when the other arrives while it is still prefilling, added to its targets then), and the other waits until it can resume from it. + +Scheduling policies (SchedulerConfig.policy and the fields after it, all host only; the defaults schedule exactly as +before they existed). The waiting queue is ordered before each admission round: fcfs keeps arrival order; lpm puts the +longest cached prefix first; lof the longest max_new_tokens first; spf the shortest prompt first. With priority +scheduling (SGLang --enable-priority-scheduling) a request's `priority` comes first, higher values first unless +low_priority_first; preemption under KV pressure takes the worst priority first; and a waiting request better than the +worst running one by more than priority_preemption_threshold takes that one's seat (SGLang preempt_to_schedule). A +queue limit (max_queued_requests, max_queued_tokens) refuses a request with QueueFull; under priority scheduling a +request better than the worst waiting one displaces it instead (SGLang _abort_on_queued_limit). """ from __future__ import annotations import bisect +import math import time from collections import deque from dataclasses import dataclass, field @@ -62,6 +72,22 @@ class NeedSync(Exception): recomputed from REAL tokens, so the caller commits the in-flight step and retries.""" +class QueueFull(Exception): + """The waiting queue is at its limit (SchedulerConfig.max_queued_requests / max_queued_tokens): the request is + refused before it is queued. SGLang answers this with 503 "The request queue is full." (scheduler.py + _abort_on_queued_limit).""" + + +# SGLang's --schedule-policy names (managers/schedule_policy.py), the ones served here, plus spf; "priority" is fcfs with +# priority scheduling on. +POLICIES = ("fcfs", "lpm", "lof", "spf", "priority") +# SGLang turns lpm's prefix matching off above this many waiting requests (SchedulePolicy._determine_active_policy): +# each match walks the radix tree, once per waiting request per step. +LPM_MAX_QUEUE = 128 +# Displaced from a full queue by a request of higher priority (SGLang's message). +DISPLACED = "The request is aborted by a higher priority request." + + @dataclass class SchedulerConfig: page_size: int @@ -78,6 +104,70 @@ class SchedulerConfig: admission_decode_tokens: int = 256 # EngineConfig.max_running: running requests at most (None: max_num_seqs). Admission only, never the graphs' seats. max_running: int | None = None + # Scheduling policy (module docstring, POLICIES). Host only; the defaults are the scheduling without them. + policy: str = "fcfs" + priority_scheduling: bool = False + low_priority_first: bool = False # SGLang --schedule-low-priority-values-first + priority_preemption: bool = True # SGLang's default; --disable-priority-preemption turns it off + priority_preemption_threshold: int = 10 # SGLang --priority-scheduling-preemption-threshold's default + max_queued_requests: int | None = None # SGLang --max-queued-requests + max_queued_tokens: int | None = None # prompt tokens waiting, at most + + def __post_init__(self) -> None: + if self.policy not in POLICIES: + raise ValueError(f"schedule policy must be one of {', '.join(POLICIES)}") + if self.policy == "priority": + self.policy, self.priority_scheduling = "fcfs", True + if self.priority_scheduling and self.policy == "lpm": + # SGLang's own check (server_args.py): an order by priority and one by cached prefix cannot both lead + raise ValueError("priority scheduling takes schedule policy fcfs, lof or spf, not lpm") + for name in ("max_queued_requests", "max_queued_tokens"): + v = getattr(self, name) + if v is not None and v < 1: + raise ValueError(f"{name} must be at least 1") + + +@dataclass(frozen=True) +class Scheduling: + """The host-only scheduling settings (the `serve` and `engine` scheduling flags, trnrt/__main__.py), kept out of + EngineConfig, whose fields the captured graphs are keyed on. The defaults schedule and evict exactly as an engine + built without them.""" + + policy: str = "fcfs" + priority_scheduling: bool = False + low_priority_first: bool = False + priority_preemption: bool = True + priority_preemption_threshold: int = 10 + max_queued_requests: int | None = None + max_queued_tokens: int | None = None + eviction_policy: str = "lru" # radix_cache.EVICTION_POLICIES + eviction_config: dict | None = None # the eviction policy's own settings (radix_cache.POLICY_PARAMS) + + @property + def takes_priority(self) -> bool: + """Whether a request's `priority` changes anything: priority scheduling (by the policy's name or the flag) or + the priority eviction. A request that sends one otherwise is refused, not silently ignored.""" + return self.priority_scheduling or self.policy == "priority" or self.eviction_policy == "priority" + + def check_priority(self, priority: int | None) -> None: + """ValueError (the API's 400) for a request's priority that these settings would ignore.""" + if priority is not None and not self.takes_priority: + raise ValueError( + "priority needs priority scheduling (--enable-priority-scheduling) or the priority eviction " + "(--radix-eviction-policy priority); send the request without it" + ) + + def scheduler_fields(self) -> dict: + """The SchedulerConfig fields these settings give.""" + return { + "policy": self.policy, + "priority_scheduling": self.priority_scheduling, + "low_priority_first": self.low_priority_first, + "priority_preemption": self.priority_preemption, + "priority_preemption_threshold": self.priority_preemption_threshold, + "max_queued_requests": self.max_queued_requests, + "max_queued_tokens": self.max_queued_tokens, + } def shared_pages(a: list[int], b: list[int], ps: int, limit: int) -> int: @@ -211,6 +301,7 @@ def __init__(self, cfg: SchedulerConfig, pool: PagePool, radix: RadixCache, draf self.waiting: deque[Request] = deque() self.running: list[Request] = [] # admission order, oldest first self.num_preemptions = 0 + self.num_priority_preemptions = 0 # of which for a waiting request of higher priority (_priority_preempt) # In-chunk checkpoints placed, and skipped for want of a free row (the chunk is not split instead). self.ckpt_mid_placed = 0 self.ckpt_mid_skipped = 0 @@ -231,13 +322,52 @@ def states(self, pool: StatePool) -> None: def close_session(self, sid: str) -> bool: return self.radix.session_close(sid) - def add(self, req: Request) -> None: + def rank(self, req: Request) -> float: + """req's priority as a sort key, smaller first: higher priority values first (SGLang's default), lower ones with + low_priority_first; a request without one comes last (SGLang gives it the worst value, sys.maxsize).""" + p = req.priority + if p is None: + return math.inf + return float(p) if self.cfg.low_priority_first else -float(p) + + def add(self, req: Request) -> Request | None: + """Queue `req`. Returns the waiting request it displaced from a full queue (already finished, finish_reason + "abort", Request.error set), or None; raises QueueFull when the queue is at its limit and nothing is displaced.""" if req.num_tokens >= self.cfg.max_model_len: raise ValueError(f"prompt of {req.num_tokens} tokens does not fit max_model_len={self.cfg.max_model_len}") + displaced = self._queue_room(req) if req.session_id is not None: req.session_gen = self.radix.session_open(req.session_id) self._mark_running_junctions(req) self.waiting.append(req) + return displaced + + def _importance(self, req: Request) -> float: + """req's priority for the priority eviction, larger kept longer: SGLang's request priority, 0 when it has none + (radix_cache.py _insert_helper), negated with low_priority_first so that eviction keeps what scheduling favours.""" + p = req.priority or 0 + return -float(p) if self.cfg.low_priority_first else float(p) + + def _queue_room(self, req: Request) -> Request | None: + """Make room for `req` under the queue limits, or raise QueueFull. Over max_queued_requests with priority + scheduling, the worst waiting request (the newest among equals) is displaced if `req` ranks strictly better; + one that has already generated tokens (a preempted request, partly streamed) is never displaced.""" + c = self.cfg + if c.max_queued_tokens is not None: + queued = sum(r.num_tokens for r in self.waiting) + if queued + req.num_tokens > c.max_queued_tokens: + raise QueueFull(f"The request queue is full: {queued} prompt tokens are waiting.") + if c.max_queued_requests is None or len(self.waiting) < c.max_queued_requests: + return None + if c.priority_scheduling: + fresh = [r for r in self.waiting if not r.output_ids] + if fresh: + worst = max(fresh, key=lambda r: (self.rank(r), r.arrival_time)) + if self.rank(req) < self.rank(worst): + self.abort(worst) + worst.error = (503, DISPLACED) + return worst + raise QueueFull("The request queue is full.") def has_work(self) -> bool: return bool(self.waiting or self.running) @@ -393,14 +523,78 @@ def _seats(self) -> int: cap = self.cfg.max_running return self.cfg.max_num_seqs if cap is None else min(self.cfg.max_num_seqs, cap) + def _order_waiting(self) -> None: + """Sort the waiting queue by the scheduling policy (stable: arrival order among equals). fcfs without priority + scheduling leaves it as it is.""" + c = self.cfg + if (c.policy == "fcfs" and not c.priority_scheduling) or len(self.waiting) < 2: + return + if c.policy == "lpm": + if len(self.waiting) > LPM_MAX_QUEUE: + return + ps = c.page_size + hit = { + id(r): self.radix.match_checkpoint(r.token_ids, self._prompt_pages(r, ps))[0].num_pages + for r in self.waiting + } + self.waiting = deque(sorted(self.waiting, key=lambda r: -hit[id(r)])) + return + self.waiting = deque(sorted(self.waiting, key=self._wait_key)) + + def _wait_key(self, req: Request) -> tuple[float, int]: + """The admission order of the fcfs, lof and spf policies, smallest first: the priority (priority scheduling), + then the longest max_new_tokens (lof) or the shortest prompt (spf).""" + c = self.cfg + rank = self.rank(req) if c.priority_scheduling else 0.0 + if c.policy == "lof": + return (rank, -req.params.max_new_tokens) + if c.policy == "spf": + return (rank, req.num_tokens) + return (rank, 0) + + def _victim_key(self, req: Request) -> tuple[float, float]: + """Preemption order under priority scheduling, the largest first: the worst priority, then the youngest.""" + return (self.rank(req), req.arrival_time) + + def _priority_preempt(self) -> None: + """Priority scheduling: when every seat is taken and the best waiting request (that is not waiting for a prefix + another request is computing) ranks better than the worst running one by more than + priority_preemption_threshold, that one is preempted for its seat (SGLang PrefillAdder.preempt_to_schedule, + which also requires the difference to exceed the threshold). Before anything is planned, so the victim is in no + entry of this step; with a step in flight, NeedSync first (a preempted request recomputes from real tokens).""" + c = self.cfg + if not (c.priority_scheduling and c.priority_preemption) or not self.waiting: + return + if len(self.running) < self._seats(): + return + best = next((r for r in self.waiting if not self._prefix_in_progress(r)), None) + if best is None: + return + victim = self._youngest_running() + if not self.rank(victim) - self.rank(best) > c.priority_preemption_threshold: + return + if self.in_flight: + raise NeedSync() + self._preempt(victim) + self.num_priority_preemptions += 1 + def _schedule(self) -> StepPlan: plan = StepPlan() ps = self.cfg.page_size + preempted = self.num_preemptions + self._order_waiting() + self._priority_preempt() # 1. Every decoding sequence gets its one token, preempting the youngest if KV # runs out. Iterating oldest-first while preempting youngest-first means a - # victim is never a sequence already scheduled this step. - for req in list(self.running): + # victim is never a sequence already scheduled this step. Under priority + # scheduling the victim is the worst priority: the loop then visits the best + # first (the same invariant) and the entries go back to running order. + order = list(self.running) + if self.cfg.priority_scheduling: + pos = {id(r): i for i, r in enumerate(order)} + order.sort(key=lambda r: (self._victim_key(r), pos[id(r)])) + for req in order: if req.status is not Status.RUNNING or not req.is_decoding: continue draft = self.draft_fn(req) if self.draft_fn is not None else [] @@ -423,6 +617,9 @@ def _schedule(self) -> StepPlan: if req.status is Status.RUNNING: save = None if draft else self.track(req, req.num_computed, 1) plan.decodes.append(ScheduledSeq(req, start, 1 + len(draft), True, draft, save=save)) + if self.cfg.priority_scheduling: + at = {id(r): i for i, r in enumerate(self.running)} + plan.decodes.sort(key=lambda s: at[id(s.req)]) # 2. Continue partial prefills, oldest first. budget = self.cfg.max_prefill_tokens @@ -437,7 +634,10 @@ def _schedule(self) -> StepPlan: plan.prefills.append(self._chunk(req, req.num_computed, n)) budget -= n - # 3. Admit waiting requests, first come first served, while budget and capacity remain. + # 3. Admit waiting requests, first come first served (or in the policy's order), while budget and capacity + # remain. A request preempted above went to the queue's head; the policy puts it where it ranks. + if self.num_preemptions != preempted: + self._order_waiting() owed = sum(self._still_needs(r) for r in self.running) deferred = [] # waiting for a running request to checkpoint the prefix they share while self.waiting and budget > 0 and len(self.running) < self._seats(): @@ -470,6 +670,8 @@ def _schedule(self) -> StepPlan: self.waiting.popleft() req.pages = list(match.pages) + new_pages req.node = match.node + if self.radix.tracks_use: # a hit on the prefix it resumes from (the lfu, slru and priority evictions) + self.radix.touch(match.node, self._importance(req)) req.num_computed = cached if req.num_cached_tokens < 0: req.num_cached_tokens = cached @@ -533,7 +735,11 @@ def _reserve(self, req: Request, total_tokens: int) -> bool: return True def _youngest_running(self) -> Request: - """Preemption victim: the youngest running request.""" + """Preemption victim: the youngest running request, or under priority scheduling the worst priority (then the + youngest; _schedule visits the decodes in the order that keeps a victim out of the plan).""" + if self.cfg.priority_scheduling: + pos = {id(r): i for i, r in enumerate(self.running)} + return max(self.running, key=lambda r: (self._victim_key(r), pos[id(r)])) return max(self.running, key=lambda r: r.arrival_time) # -- recurrent-state checkpoints ------------------------------------------------------ @@ -712,6 +918,8 @@ def _cache(self, req: Request) -> None: if n_full == 0: return res = self.radix.insert(req.token_ids[: n_full * ps], req.pages[:n_full]) + if self.radix.tracks_use: + self.radix.touch(res.node, self._importance(req), hit=False) for i, page in enumerate(res.pages): if req.pages[i] != page: self.pool.free([req.pages[i]]) From 4d31cc428c67dcbc4152a49e4b61bfcb404c2424 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:04 +0530 Subject: [PATCH 011/135] tests: the admission policies, priority preemption and queue limits keep outputs exact --- tests/test_scheduler_policy.py | 345 +++++++++++++++++++++++++++++++++ 1 file changed, 345 insertions(+) create mode 100644 tests/test_scheduler_policy.py diff --git a/tests/test_scheduler_policy.py b/tests/test_scheduler_policy.py new file mode 100644 index 0000000..7179ad2 --- /dev/null +++ b/tests/test_scheduler_policy.py @@ -0,0 +1,345 @@ +"""The scheduling policies (engine/scheduler.py: SchedulerConfig.policy, priority scheduling, the queue limits). + +The "model" is test_scheduler's: the next token is a pure function of the history, so every request's output is known +exactly however it was ordered, preempted or displaced. Each policy is checked by the order it admits requests in, and +every run also checks that the outputs are the reference and the pages and checkpoint rows balance.""" + +from __future__ import annotations + +import dataclasses +import random +import zlib +from collections import deque + +import pytest + +from tests.test_scheduler import ( + VOCAB, + check_invariants, + make, + next_token, + reference_output, + reference_output_no_eos, + run, +) +from tests.test_scheduler_spec import K, emitted, speculating +from trnrt.engine.request import Request, SamplingParams, Status +from trnrt.engine.scheduler import DISPLACED, NeedSync, QueueFull, SchedulerConfig, Scheduling + + +def policy(sched, **kw): + """`sched` (tests.test_scheduler.make) with these SchedulerConfig fields.""" + sched.cfg = dataclasses.replace(sched.cfg, **kw) + return sched + + +def req(rid, n, max_new=2, priority=None): + """A prompt of n copies of one token of its own (no two requests share a page: no junction reorders them).""" + r = Request(rid, [8 + zlib.crc32(rid.encode()) % (VOCAB - 8)] * n, SamplingParams(max_new_tokens=max_new)) + r.priority = priority + return r + + +def step(sched): + plan = sched.schedule() + sched.update(plan, [next_token(s.req.token_ids[: s.end]) if s.sample else -1 for s in plan.seqs()]) + check_invariants(sched) + return plan + + +def admitted(sched, reqs) -> list[str]: + """Queue every request, then run to the end: the order of their first prefill entries.""" + for r in reqs: + sched.add(r) + order: list[str] = [] + while sched.has_work(): + for s in step(sched).prefills: + if s.req.rid not in order: + order.append(s.req.rid) + for r in reqs: + assert r.output_ids == reference_output(r.prompt_ids, r.params.max_new_tokens), r.rid + return order + + +def drive(sched, arrivals: dict[int, list[Request]], rng, depth=2, max_steps=20_000) -> list[Request]: + """The engine's overlapped loop (tests/test_scheduler_spec.py run_overlapped), with arrivals: arrivals[t] are added + before step t is scheduled, as the server's engine thread adds what its inbox holds before each step. Up to `depth` + steps in flight; a NeedSync reads them all back and retries. Blind speculative entries emit 1 .. 1 + K tokens.""" + queue: deque = deque() + added: list[Request] = [] + + def commit_one(): + plan = queue.popleft() + tokens = [] + for s in plan.seqs(): + if s.draft: + tokens.append(emitted(s.req.token_ids, 1 + rng.randrange(K + 1))) + else: + tokens.append(next_token(s.req.token_ids[: s.end]) if s.sample else -1) + sched.commit(plan, tokens) + sched.in_flight = bool(queue) + + t = 0 + last = max(arrivals, default=0) + while sched.has_work() or queue or t <= last: + for r in arrivals.get(t, ()): + sched.add(r) + added.append(r) + t += 1 + try: + plan = sched.schedule() + except NeedSync: + while queue: + commit_one() + continue + if not plan: + if queue: + commit_one() + continue + for s in plan.seqs(): + assert s.req.status is Status.RUNNING, "a planned entry belongs to a request that is not running" + sched.advance(plan) + queue.append(plan) + sched.in_flight = True + if len(queue) >= depth: + commit_one() + check_invariants(sched) + assert t < max_steps + return added + + +def _arrivals(rng, n, prios, spread=6) -> dict[int, list[Request]]: + out: dict[int, list[Request]] = {} + for i in range(n): + r = Request( + f"q{i}", + [rng.randrange(8, VOCAB) for _ in range(rng.randrange(4, 30))], + SamplingParams(max_new_tokens=rng.randrange(5, 40)), + ) + r.priority = rng.choice(prios) + out.setdefault(i * spread // 2 + rng.randrange(spread), []).append(r) + return out + + +def test_config_validation(): + assert SchedulerConfig(1, 1, 1, 2, policy="priority").priority_scheduling # fcfs with priority scheduling + assert SchedulerConfig(1, 1, 1, 2, policy="priority").policy == "fcfs" + with pytest.raises(ValueError, match="one of"): + SchedulerConfig(1, 1, 1, 2, policy="random") + with pytest.raises(ValueError, match="not lpm"): + SchedulerConfig(1, 1, 1, 2, policy="lpm", priority_scheduling=True) + with pytest.raises(ValueError, match="at least 1"): + SchedulerConfig(1, 1, 1, 2, max_queued_requests=0) + + +def test_fcfs_spf_and_lof_admission_orders(): + mk = lambda: [req("a", 30, 9), req("bb", 5, 2), req("ccc", 12, 4)] # noqa: E731 + assert admitted(make(max_seqs=1, budget=64), mk()) == ["a", "bb", "ccc"] + assert admitted(policy(make(max_seqs=1, budget=64), policy="spf"), mk()) == ["bb", "ccc", "a"] + assert admitted(policy(make(max_seqs=1, budget=64), policy="lof"), mk()) == ["a", "ccc", "bb"] + + +def test_lpm_admits_the_longest_cached_prefix_first(): + system = list(range(100, 132)) # 8 pages + + def order(sched) -> tuple[list[str], Request]: + run(sched, [Request("warm", [*system, 9], SamplingParams(max_new_tokens=2))]) # caches it, with a checkpoint + blocker = Request("blk", [700, 701, 702], SamplingParams(max_new_tokens=2)) + cold = Request("cold", [500 + i for i in range(20)], SamplingParams(max_new_tokens=2)) + hot = Request("hot", [*system, 11, 12], SamplingParams(max_new_tokens=2)) + return admitted(sched, [blocker, cold, hot]), hot + + got, hot = order(policy(make(max_seqs=1, budget=64), policy="lpm")) + assert got == ["hot", "blk", "cold"] and hot.num_cached_tokens == len(system) + assert order(make(max_seqs=1, budget=64))[0] == ["blk", "cold", "hot"] # the control: arrival order + + +def test_priority_scheduling_orders_by_priority_both_ways_and_none_last(): + mk = lambda: [req("x", 3, priority=0), req("a", 4, priority=1), req("b", 4), req("c", 4, priority=5)] # noqa: E731 + hi = policy(make(max_seqs=1, budget=64), priority_scheduling=True) + assert admitted(hi, mk()) == ["c", "a", "x", "b"] # higher values first (SGLang's default), no priority last + lo = policy(make(max_seqs=1, budget=64), priority_scheduling=True, low_priority_first=True) + assert admitted(lo, mk()) == ["x", "a", "c", "b"] + assert admitted(make(max_seqs=1, budget=64), mk()) == ["x", "a", "b", "c"] # off: arrival order + # with lof the priority still leads, and max_new_tokens orders within one priority + mk2 = lambda: [req("x", 3, 2), req("a", 4, 3, 1), req("bb", 4, 9, 1), req("c", 4, 2, 2)] # noqa: E731 + lof = policy(make(max_seqs=1, budget=64), priority_scheduling=True, policy="lof") + assert admitted(lof, mk2()) == ["c", "bb", "a", "x"] + + +def _seats_full(**kw): + """Two seats taken by long requests of priority 0.""" + sched = policy(make(max_seqs=2, budget=64), priority_scheduling=True, **kw) + sp = SamplingParams(max_new_tokens=30, ignore_eos=True) + low = [req(f"low{i}", 6, 30, priority=0) for i in range(2)] + for r in low: + r.params = sp + sched.add(r) + for _ in range(3): + step(sched) + assert len(sched.running) == 2 + return sched, low + + +@pytest.mark.parametrize(("prio", "preempts"), [(20, True), (11, True), (10, False), (5, False)]) +def test_a_waiting_request_better_by_more_than_the_threshold_takes_a_seat(prio, preempts): + sched, low = _seats_full() + urgent = req("urgent", 5, 30, priority=prio) + urgent.params = low[0].params + sched.add(urgent) + plan = step(sched) + assert (urgent in [s.req for s in plan.prefills]) == preempts + assert sched.num_priority_preemptions == int(preempts) + if preempts: # the younger of the two equally bad ones goes, and waits behind the urgent one + assert low[1].status is Status.WAITING and low[1].num_preemptions == 1 and low[0] in sched.running + run(sched, []) + for r in [*low, urgent]: + assert r.output_ids == reference_output_no_eos(r.prompt_ids, 30), r.rid + + +def test_priority_preemption_can_be_disabled_and_needs_priority_scheduling(): + sched, _ = _seats_full(priority_preemption=False) + urgent = req("urgent", 5, 30, priority=100) + sched.add(urgent) + assert urgent not in [s.req for s in step(sched).prefills] and sched.num_priority_preemptions == 0 + off = make(max_seqs=2, budget=64) # no priority scheduling: the priority field orders nothing + for r in [req("l0", 6, 30), req("l1", 6, 30)]: + off.add(r) + step(off) + late = req("late", 5, 30, priority=100) + off.add(late) + assert late not in [s.req for s in step(off).prefills] and off.num_priority_preemptions == 0 + + +@pytest.mark.parametrize("depth", [1, 2, 3]) +def test_priority_preemption_with_steps_in_flight_asks_for_a_sync_first(depth): + """The overlapped loop with arrivals over time: a preemption for a better request never runs while a step is in + flight (NeedSync first; drive() checks every plan's entries), and the outputs stay exact.""" + rng = random.Random(10 + depth) + sched = policy(make(num_pages=200, max_seqs=3, budget=32), priority_scheduling=True) + reqs = drive(sched, _arrivals(rng, 16, [None, 0, 5, 50, 100]), rng, depth) + assert sched.num_priority_preemptions > 0 + for r in reqs: + assert r.status is Status.FINISHED + assert r.output_ids == reference_output(r.prompt_ids, r.params.max_new_tokens), r.rid + + +def test_kv_pressure_preempts_the_worst_priority_first(): + """The KV-short pool of test_scheduler's preemption test, with priorities: the victims are the worst priority, so the + best request (a young one) is never preempted, where youngest-first preempts it (the control); every step's entries + are of running requests and the outputs exact.""" + + def victims(prio: bool) -> tuple[list[Request], list[Request]]: + rng = random.Random(1) + sched = make(num_pages=41, max_seqs=8, budget=32, decode_tokens=8) + if prio: + policy(sched, priority_scheduling=True) + out: list[Request] = [] + real = sched._preempt + sched._preempt = lambda r: (out.append(r), real(r)) + reqs = [] + for i in range(8): + sp = SamplingParams(max_new_tokens=100, ignore_eos=True) + r = Request(f"p{i}", [rng.randrange(8, VOCAB) for _ in range(20)], sp) + r.priority = 1000 if i == 6 else i + reqs.append(r) + drive(sched, {0: reqs}, rng, depth=1) + for r in reqs: + assert r.output_ids == reference_output_no_eos(r.prompt_ids, 100), r.rid + return out, reqs + + out, reqs = victims(True) + assert out and reqs[6] not in out + out, reqs = victims(False) + assert reqs[6] in out, "the control: youngest-first preempts the young best request" + + +@pytest.mark.parametrize("depth", [1, 2, 3]) +@pytest.mark.parametrize("seed", [0, 1]) +def test_priority_scheduling_with_blind_speculative_steps_in_flight(depth, seed): + """Priority scheduling under the blind speculative steps of tests/test_scheduler_spec.py (up to 3 in flight, any + acceptance), KV pressure and arrivals over time: exact outputs, balanced pages and rows, preemptions of both + kinds.""" + rng = random.Random(seed) + sched = speculating(policy(make(num_pages=32, max_seqs=4, budget=24, decode_tokens=0), priority_scheduling=True)) + reqs = drive(sched, _arrivals(rng, 20, [None, 0, 3, 20, 40]), rng, depth) + for r in reqs: + assert r.status is Status.FINISHED and r.spec_inflight == 0 + assert r.output_ids == reference_output(r.prompt_ids, r.params.max_new_tokens), r.rid + assert sched.num_preemptions > sched.num_priority_preemptions > 0 + + +def test_max_queued_requests_refuses_past_the_limit_and_frees_room_as_requests_start(): + sched = policy(make(max_seqs=1, budget=64), max_queued_requests=2) + a, b, c = req("a", 4), req("bb", 4), req("ccc", 4) + sched.add(a) + sched.add(b) + with pytest.raises(QueueFull, match="The request queue is full"): + sched.add(c) + assert [r.rid for r in sched.waiting] == ["a", "bb"] and c.status is Status.WAITING + step(sched) # a starts: one waiting + sched.add(c) + assert [r.rid for r in sched.waiting] == ["bb", "ccc"] + + +def test_max_queued_tokens_counts_the_waiting_prompts(): + sched = policy(make(max_seqs=1, budget=64), max_queued_tokens=10) + sched.add(req("a", 6)) + with pytest.raises(QueueFull, match="6 prompt tokens"): + sched.add(req("bb", 5)) + sched.add(req("ccc", 4)) # 10 exactly + + +def test_a_full_queue_displaces_its_worst_request_for_a_better_one(): + sched = policy(make(max_seqs=1, budget=64), max_queued_requests=3, priority_scheduling=True) + a, bb, cc = req("a", 4, priority=5), req("bb", 4, priority=1), req("cc", 4, priority=1) + for r in (a, bb, cc): + sched.add(r) + with pytest.raises(QueueFull): # no priority: never better than anyone + sched.add(req("none", 4)) + with pytest.raises(QueueFull): # equal to the worst: refused, as SGLang (strictly better only) + sched.add(req("tie", 4, priority=1)) + best = req("best", 4, priority=9) + assert sched.add(best) is cc # of the two worst (1, 1) the newer + assert cc.status is Status.FINISHED and cc.finish_reason == "abort" and cc.error == (503, DISPLACED) + assert [r.rid for r in sched.waiting] == ["a", "bb", "best"] + assert admitted(sched, []) == ["best", "a", "bb"] + + +def test_a_full_queue_never_displaces_a_request_that_has_generated(): + sched = policy(make(max_seqs=1, budget=64), max_queued_requests=1, priority_scheduling=True) + w = req("w", 4, priority=0) + sched.add(w) + w.token_ids.append(55) # as a preempted request that already streamed a token + with pytest.raises(QueueFull): + sched.add(req("x", 4, priority=99)) + + +def test_scheduling_settings_and_the_priority_field(): + assert not Scheduling().takes_priority + Scheduling().check_priority(None) + with pytest.raises(ValueError, match="enable-priority-scheduling"): + Scheduling().check_priority(3) + for s in ( + Scheduling(priority_scheduling=True), + Scheduling(policy="priority"), + Scheduling(eviction_policy="priority"), + ): + assert s.takes_priority + s.check_priority(-4) + base = SchedulerConfig(4, 8, 16, 64) + assert dataclasses.replace(base, **Scheduling().scheduler_fields()) == base # the defaults are SchedulerConfig's + + +def test_the_eviction_bookkeeping_is_recorded_only_under_a_policy_that_reads_it(): + """lru (the default) records no hits and no importance; lfu and priority do, at admission and insert.""" + for ev, tracked in (("lru", False), ("lfu", True), ("priority", True)): + sched = make() + sched.radix.eviction_policy, sched.radix.tracks_use = ev, tracked + system = list(range(100, 132)) + run(sched, [Request("a", [*system, 9], SamplingParams(max_new_tokens=2))]) + b = Request("b", [*system, 10], SamplingParams(max_new_tokens=2)) + b.priority = 7 + run(sched, [b]) + node = sched.radix.match_prefix(system).node + assert (node.hit_count > 1) == tracked and (node.importance == 7) == tracked, ev From add1f5a703cabedc8256d92d04531a559cbe4887 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:04 +0530 Subject: [PATCH 012/135] engine: take host-only scheduling settings and a request's priority --- trnrt/engine/engine.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/trnrt/engine/engine.py b/trnrt/engine/engine.py index d2d0fa8..931028e 100644 --- a/trnrt/engine/engine.py +++ b/trnrt/engine/engine.py @@ -19,7 +19,7 @@ from .radix_cache import RadixCache from .request import Request, SamplingParams, Status, unsupported from .sampler import NUM_TOP_LOGPROBS, unpack -from .scheduler import BLIND, NeedSync, Scheduler, SchedulerConfig, StepPlan +from .scheduler import BLIND, NeedSync, Scheduler, SchedulerConfig, Scheduling, StepPlan if TYPE_CHECKING: from .host_cache import HostTier @@ -219,8 +219,12 @@ def build_shard(cfg: EngineConfig, path: str, num_pages: int, rank: int = 0): class LLMEngine: - def __init__(self, cfg: EngineConfig): + scheduling = Scheduling() # the defaults, for an engine assembled without __init__ (the tests' stand-ins) + + def __init__(self, cfg: EngineConfig, scheduling: Scheduling | None = None): self.cfg = cfg + # Host-only scheduling settings (policy, priorities, queue limits, eviction): outside cfg, which keys the graphs. + self.scheduling = scheduling or Scheduling() self.t_start = time.perf_counter() # startup_seconds counts from here to the end of warmup() self.startup_seconds: float | None = None path = resolve_model_path(cfg.model_path, cfg.model_revision) @@ -278,7 +282,8 @@ def __init__(self, cfg: EngineConfig): self.runner.tp_send = tp.send _OPEN.add(self) self.pool = PagePool(num_pages) - self.radix = RadixCache(self.pool, cfg.page_size) + sch = self.scheduling + self.radix = RadixCache(self.pool, cfg.page_size, sch.eviction_policy, sch.eviction_config) scfg = SchedulerConfig( page_size=cfg.page_size, max_num_seqs=cfg.max_num_seqs, @@ -287,6 +292,7 @@ def __init__(self, cfg: EngineConfig): eos_token_ids=self.mcfg.eos_token_ids, admission_decode_tokens=cfg.admission_decode_tokens, max_running=cfg.max_running, + **sch.scheduler_fields(), ) # The scheduler (engine/scheduler.py) runs the served prefix-cache rules for a recurrent model: a hit resumes from # the deepest state checkpoint at or below it, checkpoints are written inside prefill chunks and at every page @@ -332,13 +338,24 @@ def __init__(self, cfg: EngineConfig): self._inflight: tuple[StepPlan, list[Any]] | None = None def add_request( - self, prompt_ids: list[int], params: SamplingParams, rid: str | None = None, session_id: str | None = None + self, + prompt_ids: list[int], + params: SamplingParams, + rid: str | None = None, + session_id: str | None = None, + priority: int | None = None, ) -> Request: + """Queue a request. ValueError for one the engine does not serve; scheduler.QueueFull when the waiting queue is + at its limit. A waiting request a higher-priority one displaced from a full queue ends with Request.error set.""" why = unsupported(params) if why: # before anything is built or queued: the request never reaches a step raise ValueError(f"{why} is not supported; send the request without it") + self.scheduling.check_priority(priority) req = Request(rid or f"req-{next(self._ids)}", list(prompt_ids), params, session_id=session_id) - self.scheduler.add(req) + req.priority = priority + displaced = self.scheduler.add(req) + if displaced is not None: + self.runner.forget(displaced) return req def flush_cache(self) -> None: From 25f0638182a5881529a2a722bfe0c1a7ee8750e6 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:05 +0530 Subject: [PATCH 013/135] tests: priority scheduling and a full queue through the engine's speculative step --- tests/test_engine_priority.py | 79 +++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 tests/test_engine_priority.py diff --git a/tests/test_engine_priority.py b/tests/test_engine_priority.py new file mode 100644 index 0000000..3ff9d1f --- /dev/null +++ b/tests/test_engine_priority.py @@ -0,0 +1,79 @@ +"""Priority scheduling and the queue limits through the engine's own step (LLMEngine.add_request and step), with blind +speculative steps, dynamic speculation and the state pool, against tests/test_spec_engine_fake.py's fake device: the +device checks every call's state rows and pages, so a preemption or a displacement the host books wrongly fails a call. +The outputs must be the reference whatever was preempted, displaced or reordered.""" + +from __future__ import annotations + +import dataclasses +import random + +import pytest + +from tests.test_spec_engine_fake import check_done, make_engine, prompts +from trnrt.engine.request import SamplingParams, Status +from trnrt.engine.scheduler import DISPLACED, QueueFull, Scheduling + + +def with_scheduling(eng, **kw): + """make_engine's engine with these host-only settings, as LLMEngine.__init__ applies them.""" + eng.scheduling = Scheduling(**kw) + eng.scheduler.cfg = dataclasses.replace(eng.scheduler.cfg, **eng.scheduling.scheduler_fields()) + return eng + + +@pytest.mark.parametrize("dyn", [False, True]) +@pytest.mark.parametrize("seed", [0, 1]) +def test_priority_scheduling_through_the_engine_generates_the_reference(dyn, seed): + if dyn: + eng = make_engine(num_pages=48, seats=8, ckpts=10, spec_seats=3, spec_max_batch=3, max_len=200, seed=seed) + else: + eng = make_engine(num_pages=40, seats=4, ckpts=6, max_len=120, seed=seed) + with_scheduling(eng, priority_scheduling=True, priority_preemption_threshold=5) + eng.scheduler.cfg = dataclasses.replace(eng.scheduler.cfg, max_running=3) # host only: seats fill sooner + rng = random.Random(seed) + reqs = [] + pending = prompts(20, seed=seed + 3) + while pending: # arrivals in small bursts over the run + n = rng.randrange(1, 4) + burst, pending = pending[:n], pending[n:] + for p in burst: + prio = rng.choice([0, 3, 30]) + reqs.append(eng.add_request(p, SamplingParams(max_new_tokens=rng.randrange(6, 45)), priority=prio)) + for _ in range(rng.randrange(0, 4)): + eng.step() + steps = 0 + while eng.has_work(): + eng.step() + steps += 1 + assert steps < 8000 + check_done(eng, reqs) + assert eng.scheduler.num_priority_preemptions > 0, "no request ever took a seat by priority" + + +def test_a_full_queue_refuses_or_displaces_through_the_engine_and_frees_what_it_held(): + eng = with_scheduling(make_engine(seats=2), priority_scheduling=True, max_queued_requests=2) + sp = SamplingParams(max_new_tokens=8) + running = [eng.add_request(p, sp, priority=5) for p in prompts(2, seed=5, shared=False)] + eng.step() # both admitted: the queue is empty again + low = [eng.add_request(p, sp, priority=0) for p in prompts(2, seed=6, shared=False)] + with pytest.raises(QueueFull): + eng.add_request([10, 11, 12], sp, priority=0) # no better than the worst waiting + high = eng.add_request([10, 11, 12], sp, priority=9) + gone = low[1] # the newer of the two worst + assert gone.status is Status.FINISHED and gone.error == (503, DISPLACED) and gone.finish_reason == "abort" + while eng.has_work(): + eng.step() + check_done(eng, [*running, low[0], high]) + assert gone.output_ids == [] + + +def test_the_engine_refuses_a_priority_it_would_ignore(): + eng = make_engine() + with pytest.raises(ValueError, match="enable-priority-scheduling"): + eng.add_request([10, 11, 12], SamplingParams(max_new_tokens=4), priority=1) + assert not eng.scheduler.waiting + r = eng.add_request([10, 11, 12], SamplingParams(max_new_tokens=4)) # no priority: the default engine as before + while eng.has_work(): + eng.step() + check_done(eng, [r]) From ed062bdddd1052b386caf1d5cb927c2da1cc12fe Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:05 +0530 Subject: [PATCH 014/135] cli: SGLang's scheduling and eviction flags, host only --- trnrt/__main__.py | 113 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 111 insertions(+), 2 deletions(-) diff --git a/trnrt/__main__.py b/trnrt/__main__.py index a789b49..87d8354 100644 --- a/trnrt/__main__.py +++ b/trnrt/__main__.py @@ -10,7 +10,8 @@ prefix cache (`--preset batch12` names it). `--preset batch8` is the other supported configuration: decode batches of up to 8 per engine (PRESETS below). The bucket ladders and the decode kernel routes are graph structure: a server finds its graphs in the compile cache only if `compile capture` ran with the same arguments (the same --preset included). ---max-running (fewer requests running than the compiled seats) is host only. An engine compiles nothing on the serving +--max-running (fewer requests running than the compiled seats) and the scheduling flags (--schedule-policy, priorities, +queue limits, --radix-eviction-policy) are host only. An engine compiles nothing on the serving host: a graph missing from the cache ends its start-up with the likely cause (NEURON_LIBTORCH_ASSERT_CACHE_HIT=1; --allow-compile compiles it there instead, minutes to hours a graph). Before it starts, an engine on Neuron checks its host (engine/start_checks.py): its devices free of other processes, and enough contiguous free host memory for the @@ -372,6 +373,106 @@ def add_engine_args(ap: argparse.ArgumentParser) -> None: default=(2,), help="the decode row counts at which a dense MLP layer runs as the XLA graph instead of the dense half's kernel", ) + add_scheduling_args(ap) + + +def _json_object(s: str) -> dict: + import json + + try: + v = json.loads(s) + except ValueError as e: + raise argparse.ArgumentTypeError(f"not JSON: {e}") from None + if not isinstance(v, dict): + raise argparse.ArgumentTypeError("must be a JSON object") + return v + + +def add_scheduling_args(ap: argparse.ArgumentParser) -> None: + """Each engine's scheduling, under SGLang's flag names wherever SGLang has the setting (engine/scheduler.py, + engine/radix_cache.py). Host only: none of them changes a graph, and the defaults schedule as an engine without + them.""" + g = ap.add_argument_group( + "scheduling", "host only, per engine (no graph changes); the defaults are first come, first served and LRU" + ) + g.add_argument( + "--schedule-policy", + default="fcfs", + choices=["fcfs", "lpm", "lof", "spf", "priority"], + help="the order waiting requests are admitted in: fcfs arrival, lpm longest cached prefix, lof longest " + "max_tokens, spf shortest prompt; priority is fcfs with --enable-priority-scheduling", + ) + g.add_argument( + "--enable-priority-scheduling", + action="store_true", + help="admit by each request's `priority` first (higher values first), preempt the worst priority first", + ) + g.add_argument( + "--schedule-low-priority-values-first", + action="store_true", + help="with priority scheduling: lower priority values first", + ) + g.add_argument( + "--priority-scheduling-preemption-threshold", + type=_at_least(0), + default=10, + help="with priority scheduling: a waiting request takes a running one's seat when its priority is better by " + "more than this", + ) + g.add_argument( + "--disable-priority-preemption", + action="store_true", + help="with priority scheduling: never preempt a running request for a waiting one of higher priority", + ) + g.add_argument( + "--max-queued-requests", + type=_at_least(1), + default=None, + help="waiting requests per engine at most; past it a request is answered 503 (default: no limit)", + ) + g.add_argument( + "--max-queued-tokens", + type=_at_least(1), + default=None, + help="prompt tokens waiting per engine at most; past it a request is answered 503 (default: no limit)", + ) + g.add_argument( + "--radix-eviction-policy", + default="lru", + choices=["lru", "lfu", "slru", "priority", "tlru"], + help="which cached prefix the prefix cache evicts first", + ) + g.add_argument( + "--radix-eviction-policy-config", + type=_json_object, + default=None, + help='the eviction policy\'s settings as JSON: slru {"protected_threshold": 2}; tlru {"threshold": N, ' + '"next_prompt_estimate": M} (both required)', + ) + + +def scheduling(argv: list[str]): + """The engine's host-only scheduling settings (engine/scheduler.Scheduling) from the serve or engine arguments + `argv`. ValueError for a combination the scheduler or the prefix cache refuses.""" + from .engine.radix_cache import policy_params + from .engine.scheduler import SchedulerConfig, Scheduling + + a = parse_engine_args(argv) + s = Scheduling( + policy=a.schedule_policy, + priority_scheduling=a.enable_priority_scheduling, + low_priority_first=a.schedule_low_priority_values_first, + priority_preemption=not a.disable_priority_preemption, + priority_preemption_threshold=a.priority_scheduling_preemption_threshold, + max_queued_requests=a.max_queued_requests, + max_queued_tokens=a.max_queued_tokens, + eviction_policy=a.radix_eviction_policy, + eviction_config=a.radix_eviction_policy_config, + ) + # the checks the engine would make at start-up, made here so that serve refuses before it starts any engine + SchedulerConfig(page_size=1, max_num_seqs=1, max_prefill_tokens=1, max_model_len=2, **s.scheduler_fields()) + policy_params(s.eviction_policy, s.eviction_config) + return s def _env_synth_al() -> float | None: @@ -504,6 +605,10 @@ def cmd_engine(argv: list[str]) -> int: add_compile_arg(ap) a = ap.parse_args(argv) cfg = engine_config(argv, a.core_base) + try: + sched = scheduling(argv) + except ValueError as e: + ap.error(str(e)) apply_preset_env(argv) # before anything imports trnrt.kernels; the ranks the engine spawns inherit it from .engine import start_checks @@ -517,7 +622,7 @@ def cmd_engine(argv: list[str]) -> int: t = time.perf_counter() try: - eng = LLMEngine(cfg) + eng = LLMEngine(cfg, sched) print( f"engine up in {time.perf_counter() - t:.1f} s (cores {a.core_base}-{a.core_base + cfg.tp - 1})", flush=True ) @@ -623,6 +728,10 @@ def cmd_serve(argv: list[str]) -> int: ) add_compile_arg(ap) a = ap.parse_args(argv) + try: + scheduling(argv) # refused here, before any engine starts, rather than by every engine + except ValueError as e: + ap.error(str(e)) # The engines start one at a time, each healthy (its weights loaded, every graph warmed) before the next starts. The # host-wide start lock stays on too (engine/start_lock.py: each graph's first call holds it): two engines whose # ranks load collective graphs at the same moment collide in the runtime's collectives bootstrap. From c26685753dc76adbd8346a855d428606be8fee10 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:05 +0530 Subject: [PATCH 015/135] tests: the stand-in engine takes the scheduling settings --- tests/test_serve_cli.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_serve_cli.py b/tests/test_serve_cli.py index e373aeb..ce086a6 100644 --- a/tests/test_serve_cli.py +++ b/tests/test_serve_cli.py @@ -216,13 +216,15 @@ def get(url, timeout=None): class _FakeLLMEngine: seen_env: ClassVar[list] = [] cfgs: ClassVar[list] = [] + schedulings: ClassVar[list] = [] fail = None - def __init__(self, cfg): + def __init__(self, cfg, scheduling=None): import os type(self).seen_env.append(os.environ.get(cli.CACHE_HIT_ENV)) type(self).cfgs.append(cfg) + type(self).schedulings.append(scheduling) def warmup(self): if type(self).fail is not None: @@ -247,6 +249,7 @@ def engine(monkeypatch): import trnrt.serve.api as api _FakeLLMEngine.seen_env, _FakeLLMEngine.cfgs, _FakeLLMEngine.fail = [], [], None + _FakeLLMEngine.schedulings = [] monkeypatch.setattr(eng, "LLMEngine", _FakeLLMEngine) monkeypatch.setattr(api, "serve", lambda *a, **k: None) monkeypatch.setattr(api, "build_app", lambda *a, **k: None) From 4dd3c23a0af1a66382ee6d80d94b6734eb527301 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:05 +0530 Subject: [PATCH 016/135] tests: the scheduling flags reach every engine and change no engine config --- tests/test_sched_cli.py | 106 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 tests/test_sched_cli.py diff --git a/tests/test_sched_cli.py b/tests/test_sched_cli.py new file mode 100644 index 0000000..0238ea5 --- /dev/null +++ b/tests/test_sched_cli.py @@ -0,0 +1,106 @@ +"""The scheduling flags (trnrt/__main__.py add_scheduling_args, scheduling): SGLang's names, host only, refused before +any engine starts when the scheduler or the prefix cache would refuse them, and handed by serve to every engine.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("httpx") +pytest.importorskip("uvicorn") + +import trnrt.__main__ as cli +from tests import test_serve_cli as base +from trnrt.engine.scheduler import Scheduling + +host, engine, stand_in = base.host, base.engine, base.stand_in # its fixtures: fake engine processes, no device + +BASE = ["--model", "m", "--device", "cpu"] + + +def test_no_flag_is_the_default_scheduling(): + assert cli.scheduling(BASE) == Scheduling() + + +def test_the_flags_give_the_scheduling_settings(): + s = cli.scheduling( + [ + *BASE, + "--schedule-policy", + "lof", + "--enable-priority-scheduling", + "--schedule-low-priority-values-first", + "--priority-scheduling-preemption-threshold", + "3", + "--disable-priority-preemption", + "--max-queued-requests", + "64", + "--max-queued-tokens", + "100000", + "--radix-eviction-policy", + "slru", + "--radix-eviction-policy-config", + '{"protected_threshold": 4}', + ] + ) + assert s == Scheduling( + policy="lof", + priority_scheduling=True, + low_priority_first=True, + priority_preemption=False, + priority_preemption_threshold=3, + max_queued_requests=64, + max_queued_tokens=100000, + eviction_policy="slru", + eviction_config={"protected_threshold": 4}, + ) + + +def test_the_scheduling_flags_change_no_engine_config_field(): + """Host only: the EngineConfig (whose fields key the captured graphs) is the same with every flag set.""" + flags = [ + "--schedule-policy", + "priority", + "--max-queued-requests", + "8", + "--radix-eviction-policy", + "tlru", + "--radix-eviction-policy-config", + '{"threshold": 4096, "next_prompt_estimate": 512}', + ] + assert cli.engine_config([*BASE, *flags]) == cli.engine_config(BASE) + + +@pytest.mark.parametrize( + ("flags", "why"), + [ + (["--schedule-policy", "lpm", "--enable-priority-scheduling"], "not lpm"), + (["--radix-eviction-policy", "tlru"], "tlru eviction needs"), + (["--radix-eviction-policy", "lru", "--radix-eviction-policy-config", '{"threshold": 3}'], "no setting"), + ( + ["--radix-eviction-policy", "slru", "--radix-eviction-policy-config", '{"protected_threshold": -2}'], + "non-neg", + ), + ], +) +def test_a_combination_the_engine_would_refuse_is_refused_by_serve_before_any_engine_starts(host, flags, why, capsys): + with pytest.raises(SystemExit) as e: + cli.cmd_serve([*BASE, *flags]) + assert e.value.code == 2 and why in capsys.readouterr().err + assert host.engines == [] + + +@pytest.mark.parametrize("flags", [["--radix-eviction-policy-config", "[1]"], ["--radix-eviction-policy-config", "{x"]]) +def test_the_eviction_config_must_be_a_json_object(host, flags, capsys): + with pytest.raises(SystemExit) as e: + cli.cmd_serve([*BASE, *flags]) + assert e.value.code == 2 and "--radix-eviction-policy-config" in capsys.readouterr().err + + +def test_serve_hands_the_scheduling_flags_to_every_engine(stand_in): + flags = ["--enable-priority-scheduling", "--max-queued-requests", "16", "--radix-eviction-policy", "priority"] + assert cli.cmd_serve([*BASE, "--port", "9100", *flags]) == 0 + want = Scheduling(priority_scheduling=True, max_queued_requests=16, eviction_policy="priority") + assert stand_in.schedulings == [want, want] + stand_in.schedulings.clear() + assert cli.cmd_serve([*BASE, "--port", "9100"]) == 0 + assert stand_in.schedulings == [Scheduling(), Scheduling()] From 1d1a2ca18b19ac3ea1ec3142ba66d165be84c018 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:05 +0530 Subject: [PATCH 017/135] api: a request's priority, and 503 for a full queue --- trnrt/serve/api.py | 49 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 11 deletions(-) diff --git a/trnrt/serve/api.py b/trnrt/serve/api.py index 341a026..fdccb0b 100644 --- a/trnrt/serve/api.py +++ b/trnrt/serve/api.py @@ -46,6 +46,7 @@ from fastapi.responses import JSONResponse, PlainTextResponse, Response, StreamingResponse from ..engine.request import Request, SamplingParams, Status +from ..engine.scheduler import QueueFull from ..metrics import Metrics, kv_pages, spec_stats from .parsers import THINK_OPEN, StreamParser, chat_finish_reason, parse_full, template_messages, tool_choice_error @@ -68,18 +69,26 @@ class _Stream: @dataclass class _Submit: """A request on its way to the engine thread. `out` receives, in order: ("accepted", Request), or ("error", why, - HTTP status) if it was not accepted; then ("tokens", new ids, finish reason or None, their logprobs or None) until - the finish reason is set, or ("error", why, HTTP status) if the engine fails. The status is 400 for a request the - engine refuses, 500 for an internal error that leaves the engine serving and 503 once the engine has failed.""" + HTTP status[, queue full]) if it was not accepted; then ("tokens", new ids, finish reason or None, their logprobs or + None) until the finish reason is set, or ("error", why, HTTP status[, queue full]) if the engine fails or displaces + it. The status is 400 for a request the engine refuses, 500 for an internal error that leaves the engine serving, and + 503 once the engine has failed or for a request a full queue refused or displaced (queue full: True, and the response + carries QUEUE_FULL_HEADER so the router does not take the engine for failed).""" rid: str prompt_ids: list[int] params: SamplingParams loop: asyncio.AbstractEventLoop session_id: str | None = None + priority: int | None = None out: asyncio.Queue = field(default_factory=asyncio.Queue) +# Set on a 503 that answers a full waiting queue (scheduler.QueueFull) or a request displaced from it: the engine is +# serving, unlike every other 503 it sends, so the router (serve/router.py) passes the answer on and keeps routing to it. +QUEUE_FULL_HEADER = "x-trnrt-queue-full" + + # Engine loops of this process, held weakly: at exit each loop's thread stops after its step, then its engine closes # (LLMEngine.close is idempotent), which stops its tensor-parallel workers. _LOOPS: weakref.WeakSet = weakref.WeakSet() @@ -164,7 +173,9 @@ def _put(self, kind: str, payload) -> None: raise HTTPException(status_code=503, detail=f"engine failed: {err}") self.inbox.put((kind, payload)) - def submit(self, prompt_ids: list[int], params: SamplingParams, session_id: str | None = None) -> _Submit: + def submit( + self, prompt_ids: list[int], params: SamplingParams, session_id: str | None = None, priority: int | None = None + ) -> _Submit: # Checked here, before the engine sees the request: an id outside the vocabulary fails the step, and at tp > 1 # one failed step fails the engine for every request (_fail). if not prompt_ids or not all(type(t) is int and 0 <= t < self.vocab for t in prompt_ids): @@ -173,7 +184,7 @@ def submit(self, prompt_ids: list[int], params: SamplingParams, session_id: str ) if session_id is not None and not isinstance(session_id, str): raise HTTPException(status_code=400, detail="session_id must be a string") - sub = _Submit(uuid.uuid4().hex, prompt_ids, params, asyncio.get_running_loop(), session_id) + sub = _Submit(uuid.uuid4().hex, prompt_ids, params, asyncio.get_running_loop(), session_id, priority) self._put("add", sub) return sub @@ -189,10 +200,14 @@ def _drain(self, block: bool) -> None: kind, payload = item if kind == "add": post = payload.loop.call_soon_threadsafe + # priority only when sent: an engine without priority scheduling refuses any value (LLMEngine.add_request) + extra = {} if payload.priority is None else {"priority": payload.priority} try: req = self.engine.add_request( - payload.prompt_ids, payload.params, rid=payload.rid, session_id=payload.session_id + payload.prompt_ids, payload.params, rid=payload.rid, session_id=payload.session_id, **extra ) + except QueueFull as e: + post(payload.out.put_nowait, ("error", str(e), 503, True)) except ValueError as e: post(payload.out.put_nowait, ("error", str(e), 400)) except Exception as e: # noqa: BLE001 add_request raised before the request reached the scheduler @@ -302,6 +317,10 @@ def _loop(self) -> None: with_lp = req.params.logprobs is not None # With logprobs on, only the tokens whose logprobs are recorded go out, so every token carries its own. ready = len(out) if finished or not with_lp else min(len(out), len(req.logprobs)) + if finished and req.error is not None: # ended unserved: displaced from a full queue + batch.setdefault(st.loop, []).append((st.out, ("error", req.error[1], req.error[0], True))) + del self.streams[rid] + continue if ready > st.sent or finished: new = out[st.sent : ready] lps = req.logprobs[st.sent : ready] if with_lp else None @@ -461,6 +480,12 @@ def _bad(detail: str) -> HTTPException: return HTTPException(status_code=400, detail=detail) +def _error(detail: str, status: int, queue_full: bool = False) -> HTTPException: + """The HTTP error of an engine's ("error", detail, status[, queue_full]) (_Submit): a full queue's 503 carries + QUEUE_FULL_HEADER.""" + return HTTPException(status_code=status, detail=detail, headers={QUEUE_FULL_HEADER: "1"} if queue_full else None) + + def _int(body: dict, key: str, default, lo: int, hi: int = _INT64_MAX): """body[key] as an int in [lo, hi] (a float with no fraction counts), or `default` when it is absent or null. 400 otherwise: a value the engine would only fail on in a step must not get there (numpy refuses a negative or @@ -820,12 +845,14 @@ async def _generate( stream = _flag(body, "stream") opts = body.get("stream_options") or {} include_usage = isinstance(opts, dict) and _flag(opts, "include_usage") - sub = loop_.submit(prompt_ids, params, body.get("session_id")) - # Admission errors (a prompt that does not fit, an unsupported parameter) arrive before any token: surface - # them as HTTP status codes rather than inside a 200 stream. + # SGLang's per-request priority (engine/scheduler.py); an engine without priority scheduling refuses it (400) + priority = _int(body, "priority", None, -_INT64_MAX - 1) + sub = loop_.submit(prompt_ids, params, body.get("session_id"), priority) + # Admission errors (a prompt that does not fit, an unsupported parameter, a full queue) arrive before any token: + # surface them as HTTP status codes rather than inside a 200 stream. first = await sub.out.get() if first[0] == "error": - raise HTTPException(status_code=first[2], detail=first[1]) + raise _error(*first[1:]) req = first[1] created = int(time.time()) oid = ("chatcmpl-" if chat else "cmpl-") + sub.rid @@ -950,7 +977,7 @@ async def collect(): while finish is None: kind, *rest = await sub.out.get() if kind == "error": - raise HTTPException(status_code=rest[1], detail=rest[0]) + raise _error(*rest) ids, finish, lps = rest n_out += len(ids) if lps is not None: From 19007a6da82d6ad995bd987d41069f649208228c Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:05 +0530 Subject: [PATCH 018/135] router: pass a full queue's 503 on without marking the engine down --- trnrt/serve/router.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/trnrt/serve/router.py b/trnrt/serve/router.py index 03ab1ec..e1741a0 100644 --- a/trnrt/serve/router.py +++ b/trnrt/serve/router.py @@ -90,6 +90,10 @@ # Every worker answers these the same (one model, one tokenizer), so they go to the least loaded live worker and touch # no prefix tree and no session: a /tokenize of an agent's context must not steer where its next turn is routed. STATELESS_ROUTES = ("/tokenize", "/v1/tokenize", "/detokenize", "/v1/detokenize") +# An engine's 503 for a full waiting queue carries this header (serve/api.py QUEUE_FULL_HEADER): the engine is serving, +# so the answer goes to the client as it is (SGLang's 503 "The request queue is full.") and the engine stays up. Another +# engine is not tried: a session belongs on its own, and one without a session was routed where its prefix is cached. +QUEUE_FULL_HEADER = "x-trnrt-queue-full" class _Node: @@ -657,7 +661,16 @@ async def pick() -> int: except BaseException: router.load[w] -= 1 raise - if (resp if body.get("stream") else r).status_code == 503: + got = resp if body.get("stream") else r # a streamed 503 was read whole above + if got.status_code == 503 and got.headers.get(QUEUE_FULL_HEADER): + router.load[w] -= 1 + return Response( + got.content, + status_code=503, + media_type=got.headers.get("content-type"), + headers={QUEUE_FULL_HEADER: "1"}, + ) + if got.status_code == 503: router.load[w] -= 1 router.note_failure(w, hard=True) if attempt == 1 or all(not ok for ok in router.up): From e31210d9d6471701218051e11cbbe9e3441d4a6a Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:03:05 +0530 Subject: [PATCH 019/135] tests: priority and a full queue over HTTP and through the router --- tests/test_serve_sched.py | 210 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 210 insertions(+) create mode 100644 tests/test_serve_sched.py diff --git a/tests/test_serve_sched.py b/tests/test_serve_sched.py new file mode 100644 index 0000000..abea55a --- /dev/null +++ b/tests/test_serve_sched.py @@ -0,0 +1,210 @@ +"""The HTTP side of scheduling (serve/api.py, serve/router.py): a request's `priority`, and the 503 of a full queue. + +The engine is tests/test_serve_api.py's scripted stand-in, extended with what LLMEngine.add_request does for these: the +priority check of its Scheduling settings, scheduler.QueueFull for a full queue, and a waiting request displaced by a +better one (finished, Request.error set).""" + +from __future__ import annotations + +import json + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from tests.test_serve_api import PLAIN, ScriptedEngine +from trnrt.engine.request import Status +from trnrt.engine.scheduler import DISPLACED, QueueFull, Scheduling + +MSG = {"messages": [{"role": "user", "content": "Hello"}]} + + +class QueueEngine(ScriptedEngine): + """ScriptedEngine with LLMEngine.add_request's scheduling: `scheduling` checks a request's priority; with `limit` + set, a request of priority below 5 waits (it never steps) and one that finds `limit` waiting is refused with + QueueFull, unless its priority is 5 or more, when it displaces the oldest waiting request and runs.""" + + def __init__(self, script, scheduling=None, limit=None): + super().__init__(script) + self.scheduling = scheduling or Scheduling() + self.limit = limit + self.parked: list = [] + self.priorities: list = [] + + def add_request(self, prompt_ids, params, rid=None, session_id=None, priority=None): + self.scheduling.check_priority(priority) + self.priorities.append(priority) + if self.limit is None: + return super().add_request(prompt_ids, params, rid=rid, session_id=session_id) + urgent = priority is not None and priority >= 5 + if len(self.parked) >= self.limit and not urgent: + raise QueueFull("The request queue is full.") + req = super().add_request(prompt_ids, params, rid=rid, session_id=session_id) + if not urgent: + self.scheduler.running.remove(req) + req.status = Status.WAITING + self.parked.append(req) + elif len(self.parked) >= self.limit: + gone = self.parked.pop(0) + gone.status, gone.finish_reason, gone.error = Status.FINISHED, "abort", (503, DISPLACED) + return req + + +@pytest.fixture +def app(): + from fastapi.testclient import TestClient + + from trnrt.serve import api + + started = [] + + def start(eng): + client = TestClient(api.build_app(eng, "glm")) + client.__enter__() + started.append((client, eng)) + return client + + yield start + for client, eng in started: + client.__exit__(None, None, None) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def test_a_priority_the_engine_would_ignore_is_refused_with_400(app): + eng = QueueEngine(PLAIN) + client = app(eng) + r = client.post("/v1/chat/completions", json={**MSG, "priority": 3}) + assert r.status_code == 400 and "--enable-priority-scheduling" in r.json()["detail"], r.text + assert client.post("/v1/chat/completions", json={**MSG, "priority": None}).status_code == 200 # null is absent + assert eng.priorities == [None] # the refused one never reached the queue + + +@pytest.mark.parametrize("value", ["high", 1.5, True, [1], 2**63]) +def test_a_priority_that_is_not_an_integer_is_refused_before_the_engine(app, value): + eng = QueueEngine(PLAIN, Scheduling(priority_scheduling=True)) + client = app(eng) + r = client.post("/v1/completions", json={"prompt": [1, 2], "priority": value}) + assert r.status_code == 400 and "priority must be an integer" in r.text, r.text + assert eng.priorities == [] + + +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/completions"]) +def test_the_priority_reaches_the_engine_on_both_endpoints(app, path): + eng = QueueEngine(PLAIN, Scheduling(priority_scheduling=True)) + client = app(eng) + body = {**MSG} if "chat" in path else {"prompt": [1, 2]} + for p in (7, -3, 2.0): + assert client.post(path, json={**body, "priority": p}).status_code == 200 + assert eng.priorities == [7, -3, 2] + + +@pytest.mark.parametrize("stream", [False, True]) +def test_a_full_queue_answers_503_with_the_queue_full_header(app, stream): + from trnrt.serve.api import QUEUE_FULL_HEADER + + eng = QueueEngine(PLAIN, limit=0) + client = app(eng) + r = client.post("/v1/chat/completions", json={**MSG, "stream": stream}) + assert r.status_code == 503 and r.json()["detail"] == "The request queue is full.", r.text + assert r.headers[QUEUE_FULL_HEADER] == "1" + h = client.get("/health") # the engine is serving: only that request was refused + assert h.status_code == 200 + eng.limit = None + assert client.post("/v1/chat/completions", json=MSG).status_code == 200 + + +def test_an_engine_failure_503_has_no_queue_full_header(app): + from trnrt.serve.api import QUEUE_FULL_HEADER + + eng = QueueEngine(PLAIN) + eng.cfg.tp = 2 + client = app(eng) + eng.dead = [1] + r = client.post("/v1/chat/completions", json=MSG) + assert r.status_code == 503 and QUEUE_FULL_HEADER not in r.headers + + +@pytest.mark.parametrize("stream", [False, True]) +def test_a_displaced_request_gets_a_503_and_the_better_one_runs(app, stream): + import threading + + from trnrt.serve.api import QUEUE_FULL_HEADER + + eng = QueueEngine(PLAIN, Scheduling(priority_scheduling=True), limit=1) + client = app(eng) + out: dict = {} + + def waiting_request(): + if stream: + with client.stream("POST", "/v1/chat/completions", json={**MSG, "stream": True, "priority": 0}) as r: + out["status"] = r.status_code + out["lines"] = [ln for ln in r.iter_lines() if ln.startswith("data: ")] + else: + r = client.post("/v1/chat/completions", json={**MSG, "priority": 0}) + out["status"], out["body"], out["headers"] = r.status_code, r.json(), r.headers + + t = threading.Thread(target=waiting_request) + t.start() + for _ in range(500): # until the engine holds it waiting + if eng.parked: + break + t.join(0.01) + assert eng.parked + r = client.post("/v1/chat/completions", json={**MSG, "priority": 9}) + assert r.status_code == 200 and r.json()["choices"][0]["finish_reason"] == "stop", r.text + t.join(10) + assert not t.is_alive() + if stream: # accepted with a 200 before it was displaced: the stream ends with the error + assert out["status"] == 200 + err = json.loads(out["lines"][-1][6:])["error"] + assert err == {"message": DISPLACED, "code": 503} + else: + assert out["status"] == 503 and out["body"]["detail"] == DISPLACED and out["headers"][QUEUE_FULL_HEADER] == "1" + + +def test_the_router_passes_a_queue_full_503_on_and_keeps_the_engine_up(): + """The real engine app behind the real router: a full queue's 503 reaches the client with its header, and the router + neither marks the engine down nor retries elsewhere; an engine failure's 503 still marks it down (the control).""" + import asyncio + + import httpx + + from tests.test_router import _free_port, _start, _stop + from trnrt.serve import api + from trnrt.serve.router import QUEUE_FULL_HEADER, Router, build_router_app + + assert api.QUEUE_FULL_HEADER == QUEUE_FULL_HEADER + + async def main() -> None: + ports = [_free_port() for _ in range(3)] + full, ok = QueueEngine(PLAIN, limit=0), QueueEngine(PLAIN) + full.cfg.tp = 2 # its rank watch: the control below kills a rank + router = Router([f"http://127.0.0.1:{p}" for p in ports[:2]]) + servers, tasks = await _start( + [ + (api.build_app(full, "glm"), ports[0]), + (api.build_app(ok, "glm"), ports[1]), + (build_router_app(router, poll_s=0), ports[2]), + ] + ) + base = f"http://127.0.0.1:{ports[2]}" + try: + async with httpx.AsyncClient(timeout=30) as c: + for stream in (False, True): + body = {**MSG, "session_id": "s1", "stream": stream} + router.sessions["s1"] = 0 # pinned to the engine whose queue is full + r = await c.post(base + "/v1/chat/completions", json=body) + assert r.status_code == 503 and r.headers.get(QUEUE_FULL_HEADER) == "1", r.text + assert r.json()["detail"] == "The request queue is full." + assert router.up == [True, True] and router.load == [0, 0] and ok.added == [] + full.dead = [1] # the control: the same engine failing answers 503 without the header + r = await c.post(base + "/v1/chat/completions", json={**MSG, "session_id": "s1"}) + assert r.status_code == 200 and router.up == [False, True] and len(ok.added) == 1 + finally: + await _stop(servers, tasks) + for lp in [lp for lp in api._LOOPS if lp.engine in (full, ok)]: + lp.stop() + + asyncio.run(main()) From 940906db59063eb7c0d47c2d40c277a099384a64 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 020/135] serve: SGLang's native /generate and its info, load, health, abort and flush endpoints --- trnrt/serve/native.py | 616 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 616 insertions(+) create mode 100644 trnrt/serve/native.py diff --git a/trnrt/serve/native.py b/trnrt/serve/native.py new file mode 100644 index 0000000..747a5e6 --- /dev/null +++ b/trnrt/serve/native.py @@ -0,0 +1,616 @@ +"""SGLang's native HTTP API on one engine, beside the OpenAI API of serve/api.py (whose build_app adds these routes). + +| Route | What it does | +|---|---| +| POST, PUT /generate | SGLang's GenerateReqInput: `text` or `input_ids`, one prompt or a batch, SGLang's `sampling_params`, `stream`, output logprobs; SGLang's response (`text`, `output_ids`, `meta_info`) | +| GET /model_info, /get_model_info | the model: its path as the operator passed it, served name, parsers, sampling defaults | +| GET /server_info, /get_server_info | the engine's configuration, its version and preset, and its load | +| GET /get_load | requests running and waiting, KV tokens in use (SGLang's legacy /get_load shape) | +| GET /health_generate | a one-token generation through the engine (SGLang's semantics: any step ending meanwhile counts) | +| POST /abort_request | abort a request by its `rid`, or every request (`abort_all`) | +| GET, POST /flush_cache | drop the prefix cache, only while the engine is idle (`?timeout=` seconds to wait for that) | + +Shapes follow SGLang's main branch as of September 2026 (after v0.5.18): python/sglang/srt/managers/io_struct.py +(GenerateReqInput, AbortReq), python/sglang/srt/entrypoints/http_server.py, python/sglang/srt/managers/ +tokenizer_manager.py (meta_info and its logprob tuples), python/sglang/srt/managers/schedule_batch.py (finish reasons) and +python/sglang/srt/sampling/sampling_params.py. Every route reaches the engine through EngineLoop's inbox (AGENTS.md: +the engine thread owns engine state). What the served graphs cannot compute (prompt logprobs, guided decoding, several +samples per prompt, multimodal input, LoRA) is refused with a 400 that names the field, as on the OpenAI routes. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import logging +import math +import os +import time +from typing import TYPE_CHECKING, Any + +from fastapi import HTTPException +from fastapi import Request as HTTPRequest +from fastapi.responses import JSONResponse, PlainTextResponse, Response, StreamingResponse + +from .. import __version__ +from ..engine.request import SamplingParams +from ..metrics import kv_pages +from .api import _NEUTRAL, _bad, _body, _Detok, _flag, _int, _params, _unless_disconnected + +if TYPE_CHECKING: + from concurrent.futures import ThreadPoolExecutor + + from fastapi import FastAPI + + from .api import EngineLoop, _Submit + +logger = logging.getLogger(__name__) + +# SGLang's SamplingParams.max_new_tokens default. +GENERATE_MAX_NEW_TOKENS = 128 +# SGLang's HEALTH_CHECK_TIMEOUT default (http_server.py): how long /health_generate waits for a step to end. +HEALTH_GENERATE_TIMEOUT_S = 20.0 +# SGLang's names for parsers equivalent to serve/parsers.py's (function_call_parser.py, reasoning_parser.py). +REASONING_PARSER, TOOL_CALL_PARSER = "glm45", "glm47" +# SGLang's /flush_cache answer on success (http_server.py). +FLUSHED = ( + "Cache flushed.\nPlease check backend logs for more details. (When there are running or waiting requests, the " + "operation will not be performed.)\n" +) + +# SGLang sampling_params -> the field api._params reads it from. The mapped fields _NEUTRAL lists are refused unless +# neutral, named here by their SGLang names. +_SAMPLING: dict[str, str] = { + "max_new_tokens": "max_tokens", + "stop": "stop", + "stop_token_ids": "stop_token_ids", + "temperature": "temperature", + "top_p": "top_p", + "top_k": "top_k", + "min_p": "min_p", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", + "repetition_penalty": "repetition_penalty", + "ignore_eos": "ignore_eos", + "sampling_seed": "seed", + "n": "n", + "logit_bias": "logit_bias", + "min_new_tokens": "min_tokens", + "no_stop_trim": "include_stop_str_in_output", + "skip_special_tokens": "skip_special_tokens", + "json_schema": "guided_json", + "regex": "guided_regex", + "ebnf": "guided_grammar", +} +# SGLang sampling parameters with no counterpart here, accepted only at these values. +_SAMPLING_NEUTRAL: dict[str, tuple[object, ...]] = { + "stop_regex": (None, [], ""), + "structural_tag": (None,), + "spaces_between_special_tokens": (None, True), + "stream_interval": (None, 1), + "beam_width": (None, 1), + "custom_params": (None, {}), +} +# GenerateReqInput fields this server does not serve (multimodal input, LoRA, hidden states, disaggregation, priority +# scheduling, cache salts, ...), accepted only absent or at a neutral value. Fields that change nothing a client sees +# (conversation_id, a trace header) are accepted as unknown fields are on every route. +_GENERATE_REFUSED = ( + "input_embeds", + "image_data", + "video_data", + "audio_data", + "mm_hashes", + "mm_content_hashes", + "use_audio_in_video", + "video_config", + "modalities", + "token_ids_logprob", + "return_sampling_mask", + "return_flat_raw_top_logprobs", + "return_flat_raw_top_logprobs_b64", + "return_hidden_states", + "return_routed_experts", + "routed_experts_start_len", + "return_indexer_topk", + "return_entropy", + "return_bytes", + "session_params", + "lora_path", + "lora_id", + "custom_logit_processor", + "positional_embed_overrides", + "bootstrap_host", + "bootstrap_port", + "bootstrap_room", + "bootstrap_pair_key", + "decode_tp_size", + "routed_dp_rank", + "data_parallel_rank", + "disagg_prefill_dp_rank", + "routing_key", + "background", + "require_reasoning", + "max_thinking_tokens", + "priority", + "extra_key", + "cache_salt", + "custom_labels", + "multi_item_delimiter_indices", + "max_dynamic_patch", + "min_dynamic_patch", + "image_max_dynamic_patch", + "video_max_dynamic_patch", + "images_config", + "need_wait_for_mm_inputs", + "num_items_assigned", + "encoder_urls", +) +_UNSET: tuple[object, ...] = (None, False, 0, "", [], {}) # 0 == False: a dp rank or priority of 0 is the default + + +@dataclasses.dataclass +class _Item: + """One prompt of a /generate request: its ids, sampling, rid (None: the server's) and logprob options.""" + + ids: list[int] + params: SamplingParams + rid: str | None + top: int # top logprobs per output token (with return_logprob) + logprob: bool + + +def _sampling(sp: Any, logprobs: int | None, vocab: int, defaults: dict[str, float]) -> SamplingParams: + """SGLang sampling_params as SamplingParams, checked by api._params under SGLang's names. 400 for an unknown key + (SGLang's SamplingParams(**sp) raises on one), and for a parameter not served at a value that is not neutral.""" + if sp is None: + sp = {} + if not isinstance(sp, dict): + raise _bad("sampling_params must be an object") + unknown = sorted(k for k in sp if k not in _SAMPLING and k not in _SAMPLING_NEUTRAL) + if unknown: + raise _bad(f"unknown sampling parameters: {', '.join(unknown)}") + refused = [k for k, neutral in _SAMPLING_NEUTRAL.items() if sp.get(k) not in neutral] + refused += [k for k, v in _SAMPLING.items() if v in _NEUTRAL and sp.get(k) not in _NEUTRAL[v]] + if refused: + raise _bad(f"not supported: {', '.join(refused)}") + body: dict = {v: sp[k] for k, v in _SAMPLING.items() if k in sp} + body["logprobs"] = logprobs + try: + return _params(body, GENERATE_MAX_NEW_TOKENS, chat=False, vocab=vocab, defaults=defaults) + except HTTPException as e: # its message names the field as _params reads it: name SGLang's + detail = str(e.detail) + for sg, ours in _SAMPLING.items(): + if sg != ours and detail.startswith(ours + " "): + detail = sg + detail[len(ours) :] + break + raise _bad(detail) from None + + +def _per_prompt(body: dict, key: str, n: int, single: bool) -> list: + """body[key] for each of n prompts: one value for all, or (a batch only) a list of n.""" + v = body.get(key) + if isinstance(v, list): + if single or len(v) != n: + raise _bad(f"{key} must be one value, or a list of one per prompt of a batch ({n})") + return v + return [v] * n + + +def _prompts(body: dict, tok) -> tuple[list[list[int]], bool]: + """The request's prompts as token ids, and whether it is one prompt (SGLang's is_single).""" + text, ids = body.get("text"), body.get("input_ids") + if (text is None) == (ids is None): + raise _bad("exactly one of text or input_ids is required") + if text is not None: + if isinstance(text, str): + return [list(tok(text)["input_ids"])], True + if isinstance(text, list) and text and all(isinstance(t, str) for t in text): + return [list(tok(t)["input_ids"]) for t in text], False + raise _bad("text must be a string or a non-empty list of strings") + if isinstance(ids, list) and ids and all(type(t) is int for t in ids): + return [ids], True + if isinstance(ids, list) and ids and all(isinstance(s, list) and s for s in ids): + return ids, False # each one's ids are checked by EngineLoop.submit + raise _bad("input_ids must be a non-empty list of token ids, or a list of such lists") + + +def parse_generate(body: dict, tok, vocab: int, defaults: dict[str, float]) -> tuple[list[_Item], bool]: + """A /generate request (SGLang's GenerateReqInput) as its prompts, and whether it is one prompt. 400 for anything + malformed or not served (_GENERATE_REFUSED, _sampling, prompt logprobs). Tokenizes `text`: run it off the event + loop.""" + refused = [k for k in _GENERATE_REFUSED if body.get(k) not in _UNSET] + if refused: + raise _bad(f"not supported: {', '.join(refused)}") + seqs, single = _prompts(body, tok) + n = len(seqs) + rid = body.get("rid") + if rid is None: + rids: list[str | None] = [None] * n + elif isinstance(rid, str) and rid: + rids = [rid] if single else [f"{rid}_{i}" for i in range(n)] # SGLang's batch rids (_normalize_rid) + elif not single and isinstance(rid, list) and len(rid) == n and all(isinstance(r, str) and r for r in rid): + rids = rid + else: + raise _bad("rid must be a non-empty string, or for a batch a list of one per prompt") + if len({r for r in rids if r is not None}) < sum(r is not None for r in rids): + raise _bad("the rids of a batch must differ") + for key in ("stream", "return_text_in_logprobs", "return_prompt_token_ids", "log_metrics", "no_logs"): + _flag(body, key) + if body.get("session_id") is not None and not isinstance(body.get("session_id"), str): + raise _bad("session_id must be a string") + items = [] + for i, (ids, sp, want_v, start_v, top_v, r) in enumerate( + zip( + seqs, + _per_prompt(body, "sampling_params", n, single), + _per_prompt(body, "return_logprob", n, single), + _per_prompt(body, "logprob_start_len", n, single), + _per_prompt(body, "top_logprobs_num", n, single), + rids, + ) + ): + one = {"return_logprob": want_v, "logprob_start_len": start_v, "top_logprobs_num": top_v} + want = _flag(one, "return_logprob") + start = _int(one, "logprob_start_len", -1, -1) + top = _int(one, "top_logprobs_num", 0, 0, 20) # the device returns at most 20 alternatives per token + if want and start >= 0: + raise _bad( + f"logprob_start_len {start} asks for prompt logprobs (prompt {i}), which no served graph computes; " + "send -1 for output logprobs only" + ) + params = _sampling(sp, top if want else None, vocab, defaults) + items.append(_Item(list(ids), params, r, top, want)) + return items, single + + +class _Run: + """One prompt of a /generate request in the engine: its output as it arrives, and SGLang's view of it.""" + + def __init__(self, item: _Item, sub: _Submit, tok, text_in_logprobs: bool, prompt_ids: bool, t0: float): + self.item, self.sub, self.tok = item, sub, tok + self.detok = _Detok(tok, item.params.stop, incremental=True, rid=sub.rid) + self.text_in, self.prompt_ids, self.t0 = text_in_logprobs, prompt_ids, t0 + self.ids: list[int] = [] + # SGLang's logprob entries of the output so far, built once per token (a stream's events repeat them all) + self.token_lps: list[list] = [] + self.top_lps: list[list[list]] = [] + self.text = "" + self.finish: str | None = None + self.req: Any = None # the engine's Request, once accepted + + @property + def rid(self) -> str: + return self.sub.rid + + def push(self, ids: list[int], finish: str | None, lps) -> bool: + """The engine's next tokens; whether anything changed.""" + if lps is not None: # one per id (the last delivery of a finished request may have fewer) + self.token_lps += [self._entry(lp, t) for t, (lp, _, _) in zip(ids, lps)] + if self.item.top: + self.top_lps += [[self._entry(v, i) for i, v in zip(ti, tl)] for _, ti, tl in lps] + self.ids += ids + delta = self.detok.push(ids, finish is not None) + self.text += delta + self.finish = finish + return bool(ids or delta or finish) + + async def run(self) -> None: + """Read the engine's tokens until the request finishes; HTTPException on an engine error.""" + while self.finish is None: + kind, *rest = await self.sub.out.get() + if kind == "error": + raise HTTPException(status_code=rest[1], detail=rest[0]) + self.push(*rest) + + def _finish_reason(self) -> dict | None: + """SGLang's finish reason (schedule_batch.py FINISH_*.to_json): a stop names what matched, the stop string or the + token id (EOS or a stop_token_ids entry); a length, the tokens generated.""" + f = self.finish + if f is None: + return None + if f == "length": + return {"type": "length", "length": len(self.ids)} + if f == "abort": + return {"type": "abort", "message": "Aborted"} + if f == "stop" and self.detok.stopped: + full = self.tok.decode(self.ids, skip_special_tokens=True) + hits = [(full.find(x), x) for x in self.item.params.stop if x in full] + if hits: + return {"type": "stop", "matched": min(hits)[1]} + return {"type": f, "matched": self.ids[-1] if self.ids else None} + + def _entry(self, lp: float, t: int) -> list: + return [lp, t, self.tok.decode([t]) if self.text_in else None] + + def result(self) -> dict: + """SGLang's output so far (tokenizer_manager.py _handle_batch_output; a stream's events are cumulative, as with + SGLang's default incremental_streaming_output off).""" + req = self.req + meta: dict = { + "id": self.rid, + "finish_reason": self._finish_reason(), + "prompt_tokens": len(self.item.ids), + "completion_tokens": len(self.ids), + "cached_tokens": max(req.num_cached_tokens, 0), + "num_retractions": req.num_preemptions, + } + if self.finish is not None: + meta["e2e_latency"] = time.perf_counter() - self.t0 + if self.item.logprob: + meta["input_token_logprobs"] = [] # output logprobs only: SGLang's logprob_start_len -1 + meta["output_token_logprobs"] = list(self.token_lps) + if self.item.top: + meta["input_top_logprobs"] = [] + meta["output_top_logprobs"] = list(self.top_lps) + out = {"text": self.text, "output_ids": list(self.ids), "meta_info": meta} + if self.prompt_ids: + out["prompt_token_ids"] = self.item.ids + return out + + +def _jsonable(v): + if v is None or isinstance(v, (bool, int, float, str)): + return v + if isinstance(v, (list, tuple)): + return [_jsonable(x) for x in v] + return str(v) # a torch dtype + + +def _checkpoint_ids(path: str | None) -> tuple[str | None, list | None]: + """model_type and architectures from the checkpoint's config.json (None when unreadable).""" + if not path: + return None, None + try: + with open(os.path.join(path, "config.json")) as f: + cfg = json.load(f) + except (OSError, ValueError): + return None, None + if not isinstance(cfg, dict): + return None, None + return cfg.get("model_type"), cfg.get("architectures") + + +def add_routes( + app: FastAPI, + loop_: EngineLoop, + *, + model_name: str, + preset: str | None, + prompt_pool: ThreadPoolExecutor, + prompt_tok, + gen_defaults: dict[str, float], +) -> None: + """SGLang's native routes on the app of one engine (serve/api.py build_app).""" + engine = loop_.engine + tok = engine.tokenizer + model_type, architectures = _checkpoint_ids(getattr(engine, "model_path", None)) + + def load() -> dict: + """SGLang's legacy /get_load entry for this engine (http_server.py get_load: num_tokens is the KV tokens in use + plus the waiting requests' tokens, num_pending_tokens the latter). Read off the engine thread, as /v1/stats + is: list() copies each queue whole while holding the GIL, so the engine thread's changes cannot interleave.""" + sched = engine.scheduler + running, waiting = list(sched.running), list(sched.waiting) + used, total = kv_pages(engine) + pending = sum(r.num_tokens for r in waiting) + return { + "dp_rank": 0, + "num_reqs": len(running) + len(waiting), + "num_running_reqs": len(running), + "num_waiting_reqs": len(waiting), + "num_tokens": used * engine.cfg.page_size + pending, + "num_pending_tokens": pending, + "token_usage": round(used / total, 4) if total else 0.0, + "ts_tic": time.perf_counter(), + } + + @app.get("/get_load") + async def get_load(): + return [load()] + + @app.get("/model_info") + @app.get("/get_model_info") + async def model_info(): + """SGLang's /model_info. model_path is the checkpoint as the operator passed it (a directory or a Hugging Face + repo id), never the download cache it resolved to.""" + return { + "model_path": engine.cfg.model_path, + "tokenizer_path": engine.cfg.model_path, + "served_model_name": model_name, + "is_generation": True, + "preferred_sampling_params": gen_defaults or None, + "reasoning_parser": REASONING_PARSER, + "tool_call_parser": TOOL_CALL_PARSER, + "has_image_understanding": False, + "has_audio_understanding": False, + "model_type": model_type, + "architectures": architectures, + } + + @app.get("/server_info") + @app.get("/get_server_info") + async def server_info(): + """SGLang's /server_info: the engine's configuration (EngineConfig: shapes, pools and routes, all from the + command line; it holds no credential), SGLang's names for the ones its clients read, the version and preset, + and the load (internal_states).""" + cfg = engine.cfg + fields = ( + {f.name: getattr(cfg, f.name) for f in dataclasses.fields(cfg)} + if dataclasses.is_dataclass(cfg) + else dict(vars(cfg)) + ) + return { + **{k: _jsonable(v) for k, v in fields.items()}, + "served_model_name": model_name, + "tp_size": cfg.tp, + "dp_size": 1, + "context_length": cfg.max_model_len, + "preset": preset, + "version": __version__, + "startup_seconds": getattr(engine, "startup_seconds", None), + "internal_states": [load()], + } + + @app.get("/health_generate") + async def health_generate(): + """SGLang's /health_generate: one token generated from prompt [0] (greedy, not counted in /metrics). A step that + ends after the probe is sent shows the engine serving, so a busy engine answers as soon as its current step ends + (and the probe is aborted); 503 if no step ends within HEALTH_GENERATE_TIMEOUT_S, or once the engine failed.""" + tic = time.monotonic() + err = loop_.health_error() + if err is None: + try: + sub = loop_.submit([0], SamplingParams(max_new_tokens=1, temperature=0.0), metrics=False) + except HTTPException as e: # the engine failed meanwhile + err = str(e.detail) + if err is not None: + return JSONResponse(status_code=503, content={"status": "failed", "error": err}) + + async def finished() -> str | None: + while True: + kind, *rest = await sub.out.get() + if kind == "error": + return str(rest[0]) + if kind == "tokens" and rest[1] is not None: + return None + + done = asyncio.ensure_future(finished()) + try: + while not done.done() and loop_.last_step <= tic and time.monotonic() < tic + HEALTH_GENERATE_TIMEOUT_S: + await asyncio.wait({done}, timeout=0.05) + if done.done(): + err = done.result() + elif loop_.last_step <= tic: + err = f"no engine step ended within {HEALTH_GENERATE_TIMEOUT_S:.0f} s" + finally: + if not done.done(): + done.cancel() + loop_.abort(sub.rid) + err = loop_.health_error() or err + if err is not None: + return JSONResponse(status_code=503, content={"status": "failed", "error": err}) + return {"status": "ok"} + + @app.post("/abort_request") + async def abort_request(http: HTTPRequest): + """SGLang's AbortReq: `rid` (exact: a batch's rids are rid_0, rid_1, ...) or `abort_all`. An aborted request + answers with what it generated and finish_reason {"type": "abort"}. 200 whether or not the rid is running, as + SGLang.""" + body = await _body(http) + for key in ("finished_reason", "abort_message"): + if body.get(key) is not None: + raise _bad(f"not supported: {key}") + if _flag(body, "abort_all"): + loop_.abort_all() + return Response(status_code=200) + rid = body.get("rid") + if not (isinstance(rid, str) and rid): + raise _bad("rid (a non-empty string) is required unless abort_all is true") + loop_.abort(rid) + return Response(status_code=200) + + @app.api_route("/flush_cache", methods=["GET", "POST"]) + async def flush_cache(http: HTTPRequest): + """SGLang's /flush_cache: the prefix cache (device pages, state checkpoints and the host tier) dropped, only + while no request runs or waits; `timeout` (seconds, default 0) waits that long for the engine to be idle. 200 + with SGLang's text when flushed, 400 with the reason when not.""" + raw = http.query_params.get("timeout", "0") + try: + timeout = float(raw) + except ValueError: + timeout = math.nan + if not (math.isfinite(timeout) and timeout >= 0): + raise _bad("timeout must be a number of seconds >= 0") + f = loop_.flush(timeout) + flushed, message = await f.fut + if flushed: + return PlainTextResponse(FLUSHED) + return PlainTextResponse(message + "\n", status_code=400) + + @app.api_route("/generate", methods=["POST", "PUT"]) + async def generate(http: HTTPRequest): + body = await _body(http) + items, single = await asyncio.get_running_loop().run_in_executor( + prompt_pool, parse_generate, body, prompt_tok, loop_.vocab, gen_defaults + ) + t0 = time.perf_counter() + subs: list[_Submit] = [] + log_metrics = _flag(body, "log_metrics", True) + try: # extend() keeps the requests queued before one raised: they are aborted + subs.extend( + loop_.submit(it.ids, it.params, body.get("session_id"), rid=it.rid, metrics=log_metrics) for it in items + ) + except HTTPException: + for s in subs: + loop_.abort(s.rid) + raise + # Admission errors (a prompt that does not fit, a rid in use) arrive before any token: an HTTP status, not a + # 200 stream, and the batch's other prompts are aborted. + firsts = [await s.out.get() for s in subs] + bad = next((f for f in firsts if f[0] == "error"), None) + if bad is not None: + for s, f in zip(subs, firsts): + if f[0] == "accepted": + loop_.abort(s.rid) + raise HTTPException(status_code=bad[2], detail=bad[1]) + text_in, prompt_ids = _flag(body, "return_text_in_logprobs"), _flag(body, "return_prompt_token_ids") + runs = [_Run(it, s, tok, text_in, prompt_ids, t0) for it, s in zip(items, subs)] + for r, f in zip(runs, firsts): + r.req = f[1] + + def abort_unfinished() -> None: + for r in runs: + if r.finish is None: + loop_.abort(r.rid) + + if _flag(body, "stream"): + return StreamingResponse(_events(runs, single, abort_unfinished), media_type="text/event-stream") + + async def collect() -> bool: + await asyncio.gather(*(r.run() for r in runs)) + return True + + try: # a client that goes away aborts its requests (api._generate) + done = await _unless_disconnected(http, collect()) + finally: + abort_unfinished() + if done is None: + return Response(status_code=499) # nginx's "client closed request"; nobody reads it + return runs[0].result() if single else [r.result() for r in runs] + + +async def _events(runs: list[_Run], single: bool, abort_unfinished): + """A /generate stream: SGLang's server-sent events, one per prompt that has new output after a step (cumulative + text, output_ids and meta_info; a batch's carry its prompt's `index`), then `data: [DONE]`. An engine error ends the + stream with SGLang's error event; a client that goes away closes this generator, which aborts what is left.""" + merged: asyncio.Queue = asyncio.Queue() + + async def pump(i: int, r: _Run) -> None: + while True: + m = await r.sub.out.get() + merged.put_nowait((i, m)) + if m[0] == "error" or m[2] is not None: + return + + pumps = [asyncio.ensure_future(pump(i, r)) for i, r in enumerate(runs)] + left = len(runs) + try: + while left: + i, (kind, *rest) = await merged.get() + if kind == "error": + err = {"message": rest[0], "type": "server_error", "code": rest[1], "retryable": False} + yield f"data: {json.dumps({'error': err})}\n\n" + break + r = runs[i] + if r.push(*rest): + out = r.result() + if not single: + out["index"] = i + yield f"data: {json.dumps(out)}\n\n" + left -= r.finish is not None + yield "data: [DONE]\n\n" + finally: + for p in pumps: + p.cancel() + abort_unfinished() From b032df6eb69cc68253fa9bad75e9cddbb7398436 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 021/135] serve: client request ids, abort with delivery, abort_all and an idle-only cache flush in the engine loop --- trnrt/serve/api.py | 167 +++++++++++++++++++++++++++++++++++++++------ 1 file changed, 146 insertions(+), 21 deletions(-) diff --git a/trnrt/serve/api.py b/trnrt/serve/api.py index 341a026..913c761 100644 --- a/trnrt/serve/api.py +++ b/trnrt/serve/api.py @@ -7,7 +7,9 @@ - GET /v1/models; - GET /health: 200 while the engine can serve, 503 with the reason once it cannot; - GET /v1/stats: the KV-cache gauges the router (serve/router.py) uses to break ties between engines; -- GET /metrics: Prometheus counters and histograms (trnrt/metrics.py). +- GET /metrics: Prometheus counters and histograms (trnrt/metrics.py); +- SGLang's native API (serve/native.py): POST /generate, the info endpoints /model_info and /server_info, /get_load, + /health_generate, /abort_request and /flush_cache. A request may carry a `session_id` (a string): the engine keeps that session's prefix referenced in its prefix cache, and the router pins the session to one engine. @@ -63,6 +65,7 @@ class _Stream: loop: asyncio.AbstractEventLoop out: asyncio.Queue sent: int = 0 + metrics: bool = True # counted in /metrics when it finishes (a /health_generate probe is not) @dataclass @@ -70,7 +73,8 @@ class _Submit: """A request on its way to the engine thread. `out` receives, in order: ("accepted", Request), or ("error", why, HTTP status) if it was not accepted; then ("tokens", new ids, finish reason or None, their logprobs or None) until the finish reason is set, or ("error", why, HTTP status) if the engine fails. The status is 400 for a request the - engine refuses, 500 for an internal error that leaves the engine serving and 503 once the engine has failed.""" + engine refuses, 500 for an internal error that leaves the engine serving and 503 once the engine has failed. An + aborted request (EngineLoop.abort, abort_all) ends with ("tokens", its undelivered ids, "abort", their logprobs).""" rid: str prompt_ids: list[int] @@ -78,6 +82,22 @@ class _Submit: loop: asyncio.AbstractEventLoop session_id: str | None = None out: asyncio.Queue = field(default_factory=asyncio.Queue) + metrics: bool = True + + +@dataclass +class _Flush: + """A /flush_cache on its way to the engine thread, which answers `fut` with (flushed, message): at once when the + engine is idle, else once it is idle, or with a refusal past `deadline` (time.monotonic()).""" + + deadline: float + loop: asyncio.AbstractEventLoop + fut: asyncio.Future + + +def _resolve(fut: asyncio.Future, value) -> None: + if not fut.done(): # its handler may have been cancelled (the client went away) + fut.set_result(value) # Engine loops of this process, held weakly: at exit each loop's thread stops after its step, then its engine closes @@ -129,6 +149,10 @@ def __init__(self, engine: LLMEngine): # the engine has them at all (an engine without context-parallel row groups has none). self._local_k_stale = True self._local_k_none = False + # time.monotonic() at the end of the last step that ran (/health_generate: a step that ends after a probe was + # sent shows the engine is serving), and the /flush_cache waiting for the engine to be idle, if any. + self.last_step = 0.0 + self._flush: _Flush | None = None self.thread = threading.Thread(target=self._run, name="trnrt-engine", daemon=True) self.thread.start() if engine.cfg.tp > 1: @@ -164,7 +188,16 @@ def _put(self, kind: str, payload) -> None: raise HTTPException(status_code=503, detail=f"engine failed: {err}") self.inbox.put((kind, payload)) - def submit(self, prompt_ids: list[int], params: SamplingParams, session_id: str | None = None) -> _Submit: + def submit( + self, + prompt_ids: list[int], + params: SamplingParams, + session_id: str | None = None, + rid: str | None = None, + metrics: bool = True, + ) -> _Submit: + """Queue a request for the engine thread. `rid`: the client's request id (SGLang's), which /abort_request + names; a new one when None. `metrics`: whether /metrics counts it.""" # Checked here, before the engine sees the request: an id outside the vocabulary fails the step, and at tp > 1 # one failed step fails the engine for every request (_fail). if not prompt_ids or not all(type(t) is int and 0 <= t < self.vocab for t in prompt_ids): @@ -173,13 +206,29 @@ def submit(self, prompt_ids: list[int], params: SamplingParams, session_id: str ) if session_id is not None and not isinstance(session_id, str): raise HTTPException(status_code=400, detail="session_id must be a string") - sub = _Submit(uuid.uuid4().hex, prompt_ids, params, asyncio.get_running_loop(), session_id) + if rid is not None and not (isinstance(rid, str) and rid): + raise HTTPException(status_code=400, detail="rid must be a non-empty string") + sub = _Submit( + rid or uuid.uuid4().hex, prompt_ids, params, asyncio.get_running_loop(), session_id, metrics=metrics + ) self._put("add", sub) return sub def abort(self, rid: str) -> None: self.inbox.put(("abort", rid)) + def abort_all(self) -> None: + """Abort every request the engine holds, and those queued for it before this call.""" + self.inbox.put(("abort_all", None)) + + def flush(self, timeout_s: float) -> _Flush: + """Ask the engine thread to drop its prefix cache (LLMEngine.flush_cache) once it is idle, waiting up to + `timeout_s` for that (SGLang's /flush_cache: 0 flushes only an engine idle when the request arrives).""" + loop = asyncio.get_running_loop() + f = _Flush(time.monotonic() + timeout_s, loop, loop.create_future()) + self._put("flush", f) + return f + def _drain(self, block: bool) -> None: try: item = self.inbox.get(timeout=0.05) if block else self.inbox.get_nowait() @@ -188,29 +237,84 @@ def _drain(self, block: bool) -> None: while True: kind, payload = item if kind == "add": - post = payload.loop.call_soon_threadsafe - try: - req = self.engine.add_request( - payload.prompt_ids, payload.params, rid=payload.rid, session_id=payload.session_id - ) - except ValueError as e: - post(payload.out.put_nowait, ("error", str(e), 400)) - except Exception as e: # noqa: BLE001 add_request raised before the request reached the scheduler - traceback.print_exc() # (its last step): that request fails, the engine keeps serving - post(payload.out.put_nowait, ("error", f"request refused: {e!r}", 500)) - else: - self.streams[payload.rid] = _Stream(req, payload.loop, payload.out) - post(payload.out.put_nowait, ("accepted", req)) + self._add(payload) elif kind == "abort": st = self.streams.pop(payload, None) - if st is not None and st.req.status is not Status.FINISHED: - self.engine.abort(st.req) + if st is not None: + self._abort(st) + elif kind == "abort_all": + streams, self.streams = self.streams, {} + for st in streams.values(): + self._abort(st) + elif kind == "flush": + self._start_flush(payload) # "wake": nothing to do, the loop re-checks whether it is stopping try: item = self.inbox.get_nowait() except queue.Empty: return + def _add(self, sub: _Submit) -> None: + post = sub.loop.call_soon_threadsafe + if sub.rid in self.streams: # a client's rid (SGLang's /generate) names one request at a time + post(sub.out.put_nowait, ("error", f"rid {sub.rid!r} is already in use", 400)) + return + try: + req = self.engine.add_request(sub.prompt_ids, sub.params, rid=sub.rid, session_id=sub.session_id) + except ValueError as e: + post(sub.out.put_nowait, ("error", str(e), 400)) + except Exception as e: # noqa: BLE001 add_request raised before the request reached the scheduler + traceback.print_exc() # (its last step): that request fails, the engine keeps serving + post(sub.out.put_nowait, ("error", f"request refused: {e!r}", 500)) + else: + self.streams[sub.rid] = _Stream(req, sub.loop, sub.out, metrics=sub.metrics) + post(sub.out.put_nowait, ("accepted", req)) + + def _abort(self, st: _Stream) -> None: + """Abort a request taken out of `streams`. Its handler, if still listening (an /abort_request; a disconnected + client's is gone), gets the tokens not yet delivered and the finish reason "abort".""" + req = st.req + if req.status is not Status.FINISHED: + self.engine.abort(req) + out = req.output_ids + with_lp = req.params.logprobs is not None + ready = min(len(out), len(req.logprobs)) if with_lp else len(out) + new = out[st.sent : ready] + lps = req.logprobs[st.sent : ready] if with_lp else None + st.loop.call_soon_threadsafe(st.out.put_nowait, ("tokens", new, "abort", lps)) + + def _start_flush(self, f: _Flush) -> None: + """A /flush_cache, in inbox order: done now if the engine is idle, refused if another is waiting or the request + set no time to wait, else kept until the engine is idle or its deadline passes (_check_flush; SGLang's + scheduler_components/flush_wrapper.py).""" + if self._flush is not None: + f.loop.call_soon_threadsafe(_resolve, f.fut, (False, "Another flush_cache is already in progress.")) + return + self._flush = f + self._check_flush() + + def _check_flush(self) -> None: + f = self._flush + if f is None: + return + if not self.engine.has_work(): + self._flush = None + try: + self.engine.flush_cache() + except Exception as e: # noqa: BLE001 a failed flush is the request's answer, never the engine's end + traceback.print_exc() + result = (False, f"flush_cache failed: {e!r}") + else: + logger.info("trnrt: prefix cache flushed") + result = (True, "") + f.loop.call_soon_threadsafe(_resolve, f.fut, result) + elif time.monotonic() >= f.deadline: + self._flush = None + busy = len(self.engine.scheduler.running) + len(self.engine.scheduler.waiting) + f.loop.call_soon_threadsafe( + _resolve, f.fut, (False, f"Cache not flushed: {busy} requests are running or waiting.") + ) + def _run(self) -> None: try: self._loop() @@ -260,6 +364,8 @@ def _fail(self, reason: str) -> None: streams, self.streams = self.streams, {} for st in streams.values(): st.loop.call_soon_threadsafe(st.out.put_nowait, ("error", f"engine failed: {reason}", 503)) + flushes = [self._flush] if self._flush is not None else [] + self._flush = None while True: try: kind, payload = self.inbox.get_nowait() @@ -267,6 +373,10 @@ def _fail(self, reason: str) -> None: break if kind == "add": payload.loop.call_soon_threadsafe(payload.out.put_nowait, ("error", f"engine failed: {reason}", 503)) + elif kind == "flush": + flushes.append(payload) + for f in flushes: + f.loop.call_soon_threadsafe(_resolve, f.fut, (False, f"engine failed: {reason}")) try: self.engine.close_failed() except Exception: # noqa: BLE001 the requests are answered; a teardown error must not hide the reason @@ -275,6 +385,7 @@ def _fail(self, reason: str) -> None: def _loop(self) -> None: while not self._stopping.is_set() and self.failed is None: self._drain(block=not self.engine.has_work()) + self._check_flush() if not self.engine.has_work(): self._read_local_k() continue @@ -294,6 +405,7 @@ def _loop(self) -> None: self.engine.abort(st.req) self.streams.clear() continue + self.last_step = time.monotonic() batch: dict = {} # one cross-thread call per event loop and step, not one per stream for rid, st in list(self.streams.items()): req = st.req @@ -311,7 +423,8 @@ def _loop(self) -> None: ) if finished: del self.streams[rid] - self.metrics.on_finish(req) + if st.metrics: + self.metrics.on_finish(req) for loop, items in batch.items(): loop.call_soon_threadsafe(_deliver, items) @@ -770,7 +883,8 @@ async def _unless_disconnected(http: HTTPRequest, aw): return task.result() -def build_app(engine: LLMEngine, model_name: str) -> FastAPI: +def build_app(engine: LLMEngine, model_name: str, preset: str | None = None) -> FastAPI: + """The server of one engine. `preset`: the --preset it was started with, for /server_info.""" app = FastAPI(title="TrnRT") loop_ = EngineLoop(engine) tok = engine.tokenizer @@ -1074,4 +1188,15 @@ async def detokenize_route(http: HTTPRequest): body = await _body(http) return await asyncio.get_running_loop().run_in_executor(prompt_pool, detokenize, body) + from . import native # here, not at the top: native imports this module + + native.add_routes( + app, + loop_, + model_name=model_name, + preset=preset, + prompt_pool=prompt_pool, + prompt_tok=prompt_tok, + gen_defaults=gen_defaults, + ) return app From ca95b73e35986108597e1c407fa55836eee85d7a Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 022/135] cli: hand the engine's preset to its server for /server_info --- trnrt/__main__.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/trnrt/__main__.py b/trnrt/__main__.py index a789b49..e96bb88 100644 --- a/trnrt/__main__.py +++ b/trnrt/__main__.py @@ -531,7 +531,13 @@ def cmd_engine(argv: list[str]) -> int: print(f"bucket warmup {w['seconds']:.1f} s over {w['graphs']} graphs", flush=True) # The router keeps its connections to the engines alive and drops idle ones after 30 s: uvicorn closing them first # (its default keep-alive is 5 s) raced a request onto a closing connection. - serve(build_app(eng, a.served_model_name), host=a.host, port=a.port, log_level="warning", timeout_keep_alive=600) + serve( + build_app(eng, a.served_model_name, preset=a.preset), + host=a.host, + port=a.port, + log_level="warning", + timeout_keep_alive=600, + ) return 0 From fe8ec679fceba734ded0029e68e15b4cd0b8eac2 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 023/135] router: route /generate by its prompt, and fan SGLang's info, load, health, abort and flush routes out to the engines --- trnrt/serve/router.py | 153 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 141 insertions(+), 12 deletions(-) diff --git a/trnrt/serve/router.py b/trnrt/serve/router.py index 03ab1ec..561baea 100644 --- a/trnrt/serve/router.py +++ b/trnrt/serve/router.py @@ -55,7 +55,18 @@ sessions go to the live ones, so a load balancer in front keeps sending); 503 {"status": "down"} once none is. Each body names the workers down ("down": their URLs) and carries the routing counters. Up and down are the router's view (above): the /health polls and the requests' refused connections and 503s; -- GET /v1/models: the answer of the first live worker that gives one (least loaded first), 503 if none does. +- GET /v1/models: the answer of the first live worker that gives one (least loaded first), 503 if none does; +- SGLang's native API (serve/native.py). POST /generate is a request like /v1/completions: routed by its prompt (`text` + or `input_ids`, a batch as one request) or its session. The rest go to every worker or to one: + + | Route | Workers | Answer | + |---|---|---| + | /model_info, /get_model_info | the first live one | its answer (every worker serves the same model) | + | /server_info, /get_server_info | all | the first answer, with dp_size the number of workers and internal_states every worker's load | + | /get_load | all | every worker's entry, dp_rank its index | + | /health_generate | all | 200 "ok" when each generated, "degraded" when some did, 503 "down" when none did | + | /abort_request | all | 200 once any took it: only the worker holding the rid has it, so no rid map is kept | + | /flush_cache | all | 200 when every one flushed (its prefix tree here is cleared too), else 400 naming the others | python -m trnrt.serve.router --worker-urls http://127.0.0.1:8111,http://127.0.0.1:8112 --port 8000 """ @@ -73,7 +84,7 @@ import httpx from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse, Response, StreamingResponse +from fastapi.responses import JSONResponse, PlainTextResponse, Response, StreamingResponse from .. import log @@ -86,7 +97,7 @@ # without it. LOG_SESSIONS_ENV = "TRNRT_ROUTER_LOG_SESSIONS" SESSION_KEY_CHARS = 4096 -ROUTES = ("/v1/completions", "/v1/chat/completions") +ROUTES = ("/v1/completions", "/v1/chat/completions", "/generate") # Every worker answers these the same (one model, one tokenizer), so they go to the least loaded live worker and touch # no prefix tree and no session: a /tokenize of an agent's context must not steer where its next turn is routed. STATELESS_ROUTES = ("/tokenize", "/v1/tokenize", "/detokenize", "/v1/detokenize") @@ -347,6 +358,11 @@ def _forget(self, sid: str) -> None: self.keys.pop(sid, None) self.tentative.discard(sid) + def clear_tree(self, w: int) -> None: + """Forget what was sent to worker w (its prefix cache was flushed): no request is routed there for it.""" + with self._trees_lock: + self.trees[w] = PrefixTree(self.trees[w].max_chars) + def note_ok(self, w: int) -> None: if not self.up[w]: logger.info("trnrt.router: worker %s (%s) is up", w, self.workers[w]) @@ -473,6 +489,15 @@ def _request_head(path: str, body: dict) -> int: return len(json.dumps(msgs[:k], ensure_ascii=False)) - 1 if k else 0 # "[A, B" is a prefix of "[A, B, C]" +def _prompt(path: str, body: dict): + """A prompt request's prompt: /v1/completions' `prompt`, SGLang /generate's `text` or else `input_ids` (a string, a + list of ids, or a batch's list of either).""" + if path == "/generate": + text = body.get("text") + return body.get("input_ids", "") if text is None else text + return body.get("prompt", "") + + def _session_key(path: str, body: dict) -> tuple[str, str]: """A session request's history key, (sys_hash, first_hash): hashes of its leading system messages and of the first SESSION_KEY_CHARS characters of the first message after them (for a completion: "" and its prompt's first @@ -481,7 +506,7 @@ def _session_key(path: str, body: dict) -> tuple[str, str]: compaction). No dump of the whole request (~1 MB at 200K tokens).""" h = lambda b: hashlib.sha1(b.encode("utf-8", "replace"), usedforsecurity=False).hexdigest() # noqa: E731 if not path.endswith("chat/completions"): - p = body.get("prompt", "") + p = _prompt(path, body) return h(""), h((p if isinstance(p, str) else json.dumps(p))[:SESSION_KEY_CHARS]) msgs = body.get("messages") or [] k = _leading_system(msgs) @@ -494,7 +519,7 @@ def _session_key(path: str, body: dict) -> tuple[str, str]: def _request_text(path: str, body: dict) -> str: if path.endswith("chat/completions"): return json.dumps(body.get("messages", []), ensure_ascii=False) - p = body.get("prompt", "") + p = _prompt(path, body) return p if isinstance(p, str) else json.dumps(p) @@ -599,13 +624,12 @@ async def health(): {"status": status, "down": down, **router.health()}, status_code=503 if status == "down" else 200 ) - @app.get("/v1/models") - async def models(): - """Every worker serves the same model: the first live one that answers, least loaded first. A refused - connection marks a worker down, as in a request.""" + async def first_live(path: str) -> JSONResponse: + """What every worker answers the same (one model): the first live one that answers, least loaded first. A + refused connection marks a worker down, as in a request.""" for w in sorted(router._live(), key=lambda i: (router.load[i], i)): try: - r = await client.get(router.workers[w] + "/v1/models", timeout=health_timeout_s) + r = await client.get(router.workers[w] + path, timeout=health_timeout_s) except httpx.ConnectError: router.note_failure(w, hard=True) continue @@ -614,7 +638,112 @@ async def models(): if r.status_code == 200: with contextlib.suppress(ValueError): return JSONResponse(r.json()) - raise HTTPException(status_code=503, detail="no engine answered /v1/models") + raise HTTPException(status_code=503, detail=f"no engine answered {path}") + + @app.get("/v1/models") + async def models(): + return await first_live("/v1/models") + + @app.get("/model_info") + @app.get("/get_model_info") + async def model_info(): + return await first_live("/model_info") + + async def each(method: str, path: str, timeout: float | None, **kw) -> list[httpx.Response | None]: + """Every worker's answer to one request, sent to all at once; None for one that did not answer. A refused + connection marks it down, as in a request; a 503 too (a failed engine).""" + + async def one(w: int) -> httpx.Response | None: + try: + r = await client.request(method, router.workers[w] + path, timeout=timeout, **kw) + except httpx.ConnectError: + router.note_failure(w, hard=True) + return None + except httpx.HTTPError: + return None + if r.status_code == 503: + router.note_failure(w, hard=True) + return r + + return list(await asyncio.gather(*(one(w) for w in range(len(router.workers))))) + + def json_of(r: httpx.Response | None): + if r is None or r.status_code != 200: + return None + try: + return r.json() + except ValueError: + return None + + @app.get("/get_load") + async def get_load(): + """Every engine's SGLang load entry, dp_rank its worker index (an engine is the data-parallel rank here).""" + out = [] + for w, r in enumerate(await each("GET", "/get_load", health_timeout_s)): + got = json_of(r) + if isinstance(got, list): + out += [dict(e, dp_rank=w) for e in got if isinstance(e, dict)] + if not out: + raise HTTPException(status_code=503, detail="no engine answered /get_load") + return out + + @app.get("/server_info") + @app.get("/get_server_info") + async def server_info(): + """The first engine's configuration (the engines share it), as SGLang's data-parallel /server_info: dp_size + the number of engines, internal_states each answering engine's load.""" + answers = [(w, json_of(r)) for w, r in enumerate(await each("GET", "/server_info", health_timeout_s))] + answers = [(w, a) for w, a in answers if isinstance(a, dict)] + if not answers: + raise HTTPException(status_code=503, detail="no engine answered /server_info") + states = [dict(e, dp_rank=w) for w, a in answers for e in a.get("internal_states") or () if isinstance(e, dict)] + return {**answers[0][1], "dp_size": len(router.workers), "internal_states": states} + + @app.get("/health_generate") + async def health_generate(): + """Each engine's /health_generate (a token through its engine): 200 "ok" when every engine generated, + "degraded" when some did, 503 "down" when none did. Each body names the engines that did not ("down").""" + from .native import HEALTH_GENERATE_TIMEOUT_S + + rs = await each("GET", "/health_generate", HEALTH_GENERATE_TIMEOUT_S + 10) + ok = [r is not None and r.status_code == 200 for r in rs] + for w, good in enumerate(ok): + if good: + router.note_ok(w) + status = "ok" if all(ok) else "degraded" if any(ok) else "down" + down = [u for u, good in zip(router.workers, ok) if not good] + return JSONResponse({"status": status, "down": down}, status_code=503 if status == "down" else 200) + + @app.post("/abort_request") + async def abort_request(http: Request): + """Sent to every engine: only the one holding the rid has it (abort_all: every engine aborts all). 200 once + any engine took it; else the first engine's refusal (a malformed body), or 503 if none answered.""" + body = await _json_object(http) + rs = await each("POST", "/abort_request", health_timeout_s, json=body) + if any(r is not None and r.status_code == 200 for r in rs): + return Response(status_code=200) + r = next((r for r in rs if r is not None), None) + if r is None: + raise HTTPException(status_code=503, detail="no engine answered /abort_request") + return Response(r.content, status_code=r.status_code, media_type=r.headers.get("content-type")) + + @app.api_route("/flush_cache", methods=["GET", "POST"]) + async def flush_cache(http: Request): + """Every engine's prefix cache flushed (each only while idle; `timeout` passed on). 200 when every engine + flushed, else 400 naming the engines that did not and why. A flushed engine's prefix tree here is cleared: + the tree stands for what its cache holds.""" + query = http.url.query + rs = await each(http.method, "/flush_cache" + (f"?{query}" if query else ""), None) + failed = [] + for w, r in enumerate(rs): + if r is not None and r.status_code == 200: + router.clear_tree(w) + else: + why = "no answer" if r is None else r.text.strip() or f"status {r.status_code}" + failed.append(f"engine {w}: {why}") + if failed: + return PlainTextResponse("Cache not flushed on every engine.\n" + "\n".join(failed) + "\n", status_code=400) + return PlainTextResponse(rs[0].text if rs[0] is not None else "") async def proxy(path: str, http: Request): body = await _json_object(http) @@ -719,7 +848,7 @@ async def handler(http: Request): return handler for path in ROUTES: - app.add_api_route(path, route(path), methods=["POST"]) + app.add_api_route(path, route(path), methods=["POST", "PUT"] if path == "/generate" else ["POST"]) for path in STATELESS_ROUTES: app.add_api_route(path, stateless_route(path), methods=["POST"]) return app From 5b143f34869aaf809f0ea352fbfe699eef2e8c84 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 024/135] tests: the scripted engine has flush_cache, a page size and a model path --- tests/test_serve_api.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_serve_api.py b/tests/test_serve_api.py index 66abcd4..1384cc2 100644 --- a/tests/test_serve_api.py +++ b/tests/test_serve_api.py @@ -1,8 +1,8 @@ """The HTTP layer (trnrt/serve/api.py) over a scripted stand-in for the engine, so it runs anywhere in seconds. -The stand-in implements exactly the engine API the server uses (add_request, has_work, step, abort, dead_ranks, -close_failed, close, tokenizer, cfg.tp, cfg.max_model_len, mcfg.vocab_size, scheduler, pool, radix) with real Request -objects, and generates a fixed GLM-format output one token per step. Its tokenizer is a word-level one whose words are +The stand-in implements exactly the engine API the server uses (add_request, has_work, step, abort, flush_cache, +dead_ranks, close_failed, close, tokenizer, cfg.tp, cfg.max_model_len, cfg.page_size, cfg.model_path, mcfg.vocab_size, +scheduler, pool, radix) with real Request objects, and generates a fixed GLM-format output one token per step. Its tokenizer is a word-level one whose words are GLM's tags and a few phrases, decoded by plain concatenation, and whose chat template opens `` in the prompt unless enable_thinking is false. That switch is the stand-in's: GLM-5.3-Flash's own template has no enable_thinking and always opens `` (tests/test_glm53_template.py runs the server on it). The real engine behind the same server: @@ -94,7 +94,7 @@ class ScriptedEngine: def __init__(self, script: list[int]): self.tokenizer = _tokenizer() - self.cfg = types.SimpleNamespace(tp=1, max_model_len=64) + self.cfg = types.SimpleNamespace(tp=1, max_model_len=64, page_size=4, model_path="org/glm-test") self.mcfg = types.SimpleNamespace(vocab_size=len(self.tokenizer)) self.scheduler = types.SimpleNamespace(running=[], waiting=[], num_preemptions=0) self.pool = types.SimpleNamespace(num_usable=100, num_free=90) @@ -105,6 +105,7 @@ def __init__(self, script: list[int]): self.added: list[Request] = [] self.closed = False self.step_delay = 0.0 + self.flushes: list[int] = [] # at each flush_cache: the requests the engine held (LLMEngine.flush_cache) def add_request(self, prompt_ids, params, rid=None, session_id=None): why = unsupported(params) # the engine's own refusal (LLMEngine.add_request) @@ -143,6 +144,9 @@ def abort(self, req) -> None: self.scheduler.running.remove(req) req.status, req.finish_reason = Status.FINISHED, "abort" + def flush_cache(self) -> None: + self.flushes.append(len(self.scheduler.running) + len(self.scheduler.waiting)) + def dead_ranks(self) -> list[int]: return list(self.dead) From 0d3196d00376de6bb7fc47b0453f7db265dcadcf Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 025/135] tests: SGLang's native endpoints on the scripted engine --- tests/test_serve_native.py | 405 +++++++++++++++++++++++++++++++++++++ 1 file changed, 405 insertions(+) create mode 100644 tests/test_serve_native.py diff --git a/tests/test_serve_native.py b/tests/test_serve_native.py new file mode 100644 index 0000000..eb8a4fe --- /dev/null +++ b/tests/test_serve_native.py @@ -0,0 +1,405 @@ +"""SGLang's native API on one engine (trnrt/serve/native.py), over the scripted engine of tests/test_serve_api.py: +/generate in SGLang's request and response shapes, the info and load endpoints, /health_generate, /abort_request and +/flush_cache. The real engine behind the same routes: tests/test_api_native.py.""" + +from __future__ import annotations + +import contextlib +import json +import socket +import threading +import time + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from tests.test_serve_api import ID, PLAIN, WORDS, ScriptedEngine +from trnrt.engine.request import Request, SamplingParams + + +@pytest.fixture +def serve(): + from fastapi.testclient import TestClient + + from trnrt.serve import api + + started = [] + + def start(script, tp=1, preset=None): + eng = ScriptedEngine(script) + eng.cfg.tp = tp + client = TestClient(api.build_app(eng, "glm", preset=preset)) + client.__enter__() + started.append((client, eng)) + return client, eng + + yield start + for client, eng in started: + client.__exit__(None, None, None) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +@contextlib.contextmanager +def _uvicorn(eng, preset=None): + """The app on a real server in a thread (the TestClient runs one request at a time): its base URL.""" + import uvicorn + + from trnrt.serve import api + + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + srv = uvicorn.Server( + uvicorn.Config(api.build_app(eng, "glm", preset=preset), host="127.0.0.1", port=port, log_level="warning") + ) + th = threading.Thread(target=srv.run, daemon=True) + th.start() + try: + deadline = time.monotonic() + 30 + while not srv.started and time.monotonic() < deadline: + time.sleep(0.02) + assert srv.started + yield f"http://127.0.0.1:{port}" + finally: + srv.should_exit = True + th.join(timeout=30) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def _background(fn): + """fn() in a thread; returns a getter that joins it and gives its result.""" + out: dict = {} + th = threading.Thread(target=lambda: out.setdefault("r", fn()), daemon=True) + th.start() + + def get(): + th.join(30) + assert not th.is_alive() + return out["r"] + + return get + + +def _wait_for(cond, s: float = 10.0) -> None: + deadline = time.monotonic() + s + while not cond() and time.monotonic() < deadline: + time.sleep(0.01) + assert cond() + + +@pytest.mark.parametrize("method", ["POST", "PUT"]) +def test_generate_answers_in_sglangs_shape(serve, method): + client, eng = serve(PLAIN) + r = client.request(method, "/generate", json={"text": "Hello", "rid": "a1"}) + assert r.status_code == 200, r.text + out = r.json() + assert out["text"] == "Hello there again!" and out["output_ids"] == PLAIN + meta = out["meta_info"] + assert meta["id"] == "a1" and meta["finish_reason"] == {"type": "stop", "matched": PLAIN[-1]} + assert (meta["prompt_tokens"], meta["completion_tokens"], meta["cached_tokens"]) == (1, 4, 2) + assert meta["e2e_latency"] >= 0 and "output_token_logprobs" not in meta and "prompt_token_ids" not in out + assert eng.added[-1].params.max_new_tokens == 128 # SGLang's default + assert eng.added[-1].prompt_ids == [ID["Hello"]] + + r = client.post("/generate", json={"input_ids": [1, 2], "sampling_params": {"max_new_tokens": 2}}) + assert r.json()["meta_info"]["finish_reason"] == {"type": "length", "length": 2} + assert r.json()["text"] == "Hello there" + r = client.post("/generate", json={"input_ids": [1], "sampling_params": {"stop": [" again", "!"]}}) + assert r.json()["text"] == "Hello there" and r.json()["meta_info"]["finish_reason"] == { + "type": "stop", + "matched": " again", + } + p = eng.added[-1].params + assert p.stop == (" again", "!") + + +def test_generate_maps_sglangs_sampling_params(serve): + client, eng = serve(PLAIN) + sp = { + "max_new_tokens": 3, + "temperature": 0.5, + "top_p": 0.9, + "top_k": 1 << 30, # SGLang's "all" + "min_p": 0.05, + "stop_token_ids": [ID["!"]], + "ignore_eos": True, + "sampling_seed": 7, + "n": 1, + "skip_special_tokens": True, + "spaces_between_special_tokens": True, + "no_stop_trim": False, + "min_new_tokens": 0, + "json_schema": None, + "stop_regex": None, + } + r = client.post("/generate", json={"input_ids": [1], "sampling_params": sp, "priority": 0, "image_data": None}) + assert r.status_code == 200, r.text + p = eng.added[-1].params + assert (p.max_new_tokens, p.temperature, p.top_p, p.top_k, p.min_p) == (3, 0.5, 0.9, 1 << 30, 0.05) + assert (p.stop_token_ids, p.ignore_eos, p.seed) == ((ID["!"],), True, 7) + + +def test_generate_batch_rids_and_prompt_ids(serve): + client, _ = serve(PLAIN) + r = client.post( + "/generate", + json={ + "input_ids": [[1], [2, 3]], + "rid": "b", + "sampling_params": [{"max_new_tokens": 1}, {"max_new_tokens": 2}], + "return_prompt_token_ids": True, + }, + ) + assert r.status_code == 200, r.text + a, b = r.json() + assert (a["meta_info"]["id"], b["meta_info"]["id"]) == ("b_0", "b_1") # SGLang's batch rids + assert (a["text"], b["text"]) == ("Hello", "Hello there") + assert (a["prompt_token_ids"], b["prompt_token_ids"]) == ([1], [2, 3]) + r = client.post("/generate", json={"text": ["Hello", "Hello !"], "rid": ["x", "y"]}) + assert [o["meta_info"]["id"] for o in r.json()] == ["x", "y"] + assert [o["meta_info"]["prompt_tokens"] for o in r.json()] == [1, 2] + + +def test_generate_output_logprobs_in_sglangs_tuples(serve): + client, _ = serve(PLAIN) + body = {"input_ids": [1], "return_logprob": True, "top_logprobs_num": 1, "sampling_params": {"max_new_tokens": 2}} + meta = client.post("/generate", json=body).json()["meta_info"] + assert meta["input_token_logprobs"] == [] and meta["input_top_logprobs"] == [] + assert meta["output_token_logprobs"] == [[-0.25, ID["Hello"], None], [-0.25, ID[" there"], None]] + assert meta["output_top_logprobs"][0] == [[-0.25, ID["Hello"], None], [-2.0, 0, None]] + meta = client.post("/generate", json={**body, "return_text_in_logprobs": True, "logprob_start_len": -1}).json()[ + "meta_info" + ] + assert meta["output_token_logprobs"][1] == [-0.25, ID[" there"], " there"] + assert meta["output_top_logprobs"][0][1] == [-2.0, 0, ""] + meta = client.post("/generate", json={**body, "top_logprobs_num": 0}).json()["meta_info"] + assert "output_top_logprobs" not in meta and len(meta["output_token_logprobs"]) == 2 + + +@pytest.mark.parametrize( + ("body", "why"), + [ + ({"input_ids": [1], "return_logprob": True, "logprob_start_len": 0}, "prompt logprobs"), + ({"input_ids": [1], "return_logprob": True, "top_logprobs_num": 21}, "top_logprobs_num"), + ({"input_ids": [1], "image_data": "cat.png"}, "not supported: image_data"), + ({"input_ids": [1], "lora_path": "adapter"}, "not supported: lora_path"), + ({"input_ids": [1], "priority": 2}, "not supported: priority"), + ({"input_ids": [1], "sampling_params": {"n": 2}}, "not supported: n"), + ({"input_ids": [1], "sampling_params": {"json_schema": "{}"}}, "not supported: json_schema"), + ({"input_ids": [1], "sampling_params": {"no_stop_trim": True}}, "not supported: no_stop_trim"), + ({"input_ids": [1], "sampling_params": {"max_new_tokens": 0}}, "max_new_tokens must be"), + ({"input_ids": [1], "sampling_params": {"sampling_seed": -1}}, "sampling_seed must be"), + ({"input_ids": [1], "sampling_params": {"bogus": 1}}, "unknown sampling parameters: bogus"), + ({"input_ids": [1], "sampling_params": [{}]}, "sampling_params"), + ({"input_ids": [1], "return_logprob": [True]}, "return_logprob"), + ({"input_ids": [1], "rid": ["a"]}, "rid"), + ({"input_ids": [[1], [2]], "rid": ["a"]}, "rid"), + ({"input_ids": [[1], [2]], "rid": ["a", "a"]}, "must differ"), + ({"input_ids": [[1], [2]], "sampling_params": [{}]}, "sampling_params"), + ({"text": "Hello", "input_ids": [1]}, "exactly one of text or input_ids"), + ({}, "exactly one of text or input_ids"), + ({"input_ids": []}, "input_ids"), + ({"text": [1]}, "text"), + ({"input_ids": [len(WORDS)]}, "token ids"), + ({"input_ids": [1], "stream": "yes"}, "stream"), + ({"input_ids": [1], "session_id": 3}, "session_id"), + ({"input_ids": [1], "sampling_params": {"frequency_penalty": 0.5}}, "penalties"), + ], + ids=lambda v: v if isinstance(v, str) else None, +) +def test_generate_refuses_what_it_cannot_serve_with_a_400(serve, body, why): + client, eng = serve(PLAIN) + r = client.post("/generate", json=body) + assert r.status_code == 400 and why in r.json()["detail"], r.text + assert eng.scheduler.running == [] and all(q.status.value == "finished" for q in eng.added) + + +def _sse(r) -> list: + lines = [ln for ln in r.iter_lines() if ln.startswith("data: ")] + assert lines[-1] == "data: [DONE]" + return [json.loads(ln[6:]) for ln in lines[:-1]] + + +def test_generate_streams_cumulative_events(serve): + client, _ = serve(PLAIN) + with client.stream("POST", "/generate", json={"input_ids": [1], "stream": True}) as r: + assert r.status_code == 200 + events = _sse(r) + assert [e["text"] for e in events] == ["Hello", "Hello there", "Hello there again", "Hello there again!"] + assert [e["output_ids"] for e in events][-1] == PLAIN and "index" not in events[0] + assert [e["meta_info"]["finish_reason"] for e in events[:-1]] == [None] * 3 + assert events[-1]["meta_info"]["finish_reason"] == {"type": "stop", "matched": PLAIN[-1]} + assert events[-1]["meta_info"]["completion_tokens"] == 4 + + body = {"input_ids": [[1], [2]], "stream": True, "sampling_params": [{"max_new_tokens": 1}, {}]} + with client.stream("POST", "/generate", json=body) as r: + events = _sse(r) + last = {e["index"]: e for e in events} + assert last[0]["text"] == "Hello" and last[0]["meta_info"]["finish_reason"]["type"] == "length" + assert last[1]["text"] == "Hello there again!" and last[1]["meta_info"]["finish_reason"]["type"] == "stop" + + +def test_abort_request_ends_a_request_by_rid_with_what_it_generated(): + """The client still waiting gets its output so far and finish_reason abort (a disconnect-only abort would leave it + hanging); abort_all ends every request, an OpenAI one included; a rid in use is refused while it runs.""" + import httpx + + eng = ScriptedEngine(PLAIN * 40) + eng.step_delay = 0.02 # 160 tokens: 3 s on its own + with _uvicorn(eng) as base, httpx.Client(base_url=base, timeout=30) as c: + body = {"input_ids": [1], "rid": "r1", "sampling_params": {"max_new_tokens": 160}} + got = _background(lambda: c.post("/generate", json=body)) + _wait_for(lambda: any(len(q.output_ids) >= 3 for q in eng.scheduler.running)) + r = c.post("/generate", json=body) + assert r.status_code == 400 and "already in use" in r.text + assert c.post("/abort_request", json={"rid": "nope"}).status_code == 200 # no such rid: nothing happens + assert eng.aborted == [] + assert c.post("/abort_request", json={"rid": "r1"}).status_code == 200 + r = got() + assert r.status_code == 200, r.text + meta = r.json()["meta_info"] + assert meta["finish_reason"]["type"] == "abort" and 3 <= meta["completion_tokens"] < 160 + assert len(r.json()["output_ids"]) == meta["completion_tokens"] and eng.aborted == ["r1"] + + gen = _background(lambda: c.post("/generate", json={**body, "rid": "r2"})) + cmpl = _background(lambda: c.post("/v1/completions", json={"prompt": [1], "max_tokens": 160})) + _wait_for(lambda: len(eng.scheduler.running) == 2) + assert c.post("/abort_request", json={"abort_all": True}).status_code == 200 + assert gen().json()["meta_info"]["finish_reason"]["type"] == "abort" + assert cmpl().json()["choices"][0]["finish_reason"] == "abort" + assert len(eng.aborted) == 3 + for bad in ({}, {"rid": ""}, {"rid": 3}, {"rid": "r1", "abort_message": "x"}, {"abort_all": "yes"}): + assert c.post("/abort_request", json=bad).status_code == 400, bad + + +def test_flush_cache_flushes_only_an_idle_engine(): + """Idle: flushed at once. Busy with timeout 0 (SGLang's default): refused with the reason, nothing flushed. Busy + with a timeout: flushed once the running request finishes, never while it runs; a second flush meanwhile is + refused.""" + import httpx + + eng = ScriptedEngine(PLAIN * 40) + eng.step_delay = 0.02 + with _uvicorn(eng) as base, httpx.Client(base_url=base, timeout=30) as c: + r = c.get("/flush_cache") + assert r.status_code == 200 and r.text.startswith("Cache flushed.") and eng.flushes == [0] + assert c.post("/flush_cache").status_code == 200 and eng.flushes == [0, 0] + + busy = _background(lambda: c.post("/v1/completions", json={"prompt": [1], "max_tokens": 100})) # 2 s + _wait_for(lambda: eng.scheduler.running) + r = c.post("/flush_cache") + assert r.status_code == 400 and "1 requests are running or waiting" in r.text and eng.flushes == [0, 0] + waiting = _background(lambda: c.post("/flush_cache?timeout=20")) + time.sleep(0.1) + r = c.get("/flush_cache?timeout=5") + assert r.status_code == 400 and "Another flush_cache" in r.text + assert busy().status_code == 200 + assert waiting().status_code == 200 and eng.flushes == [0, 0, 0] # flushed after the request, while idle + for t in ("-1", "x", "nan", "inf"): + assert c.post(f"/flush_cache?timeout={t}").status_code == 400, t + + +def test_health_generate_runs_a_token_through_the_engine(monkeypatch): + """Idle: the probe (prompt [0], one greedy token) runs, and /metrics does not count it. Stuck in a step past the + timeout: 503, and the probe is aborted.""" + import httpx + + from trnrt.serve import native + + eng = ScriptedEngine(PLAIN) + with _uvicorn(eng) as base, httpx.Client(base_url=base, timeout=30) as c: + r = c.get("/health_generate") + assert r.status_code == 200 and r.json() == {"status": "ok"} + (probe,) = eng.added + assert probe.prompt_ids == [0] and probe.params.max_new_tokens == 1 and probe.params.temperature == 0.0 + assert len(probe.output_ids) == 1 and eng.aborted == [] + assert "trnrt:generation_tokens_total 0" in c.get("/metrics").text # a probe is not a request + c.post("/v1/completions", json={"prompt": [1]}) + assert "trnrt:generation_tokens_total 4" in c.get("/metrics").text + + monkeypatch.setattr(native, "HEALTH_GENERATE_TIMEOUT_S", 0.3) + eng.step_delay = 2.0 # no step ends within the timeout: the engine is stuck + started = _background(lambda: c.post("/v1/completions", json={"prompt": [1], "max_tokens": 1})) + _wait_for(lambda: eng.scheduler.running) + time.sleep(0.1) + r = c.get("/health_generate") + assert r.status_code == 503 and "no engine step ended" in r.json()["error"], r.text + assert started().status_code == 200 + _wait_for(lambda: len(eng.added) == 4 and eng.aborted == [eng.added[3].rid]) + + +def test_health_generate_is_503_once_a_rank_died(serve): + client, eng = serve(PLAIN, tp=2) + eng.dead = [1] + r = client.get("/health_generate") + assert r.status_code == 503 and "rank(s) [1] exited" in r.json()["error"] + + +def test_load_and_info_endpoints(serve): + client, eng = serve(PLAIN, preset="batch12") + eng.scheduler.waiting.append(Request("w", [1, 2, 3], SamplingParams())) + (load,) = client.get("/get_load").json() + assert load["dp_rank"] == 0 and (load["num_reqs"], load["num_running_reqs"], load["num_waiting_reqs"]) == (1, 0, 1) + # 6 KV pages in use (the stand-in's pool: 100 usable, 90 free, 4 evictable) of 4 tokens each, plus 3 waiting + assert (load["num_tokens"], load["num_pending_tokens"], load["token_usage"]) == (27, 3, 0.06) + + for path in ("/model_info", "/get_model_info"): + info = client.get(path).json() + assert info["model_path"] == "org/glm-test" and info["served_model_name"] == "glm" + assert info["is_generation"] is True and (info["reasoning_parser"], info["tool_call_parser"]) == ( + "glm45", + "glm47", + ) + for path in ("/server_info", "/get_server_info"): + info = client.get(path).json() + assert (info["preset"], info["tp_size"], info["dp_size"], info["context_length"]) == ("batch12", 1, 1, 64) + assert info["max_model_len"] == 64 and info["page_size"] == 4 and info["model_path"] == "org/glm-test" + assert info["internal_states"][0]["num_waiting_reqs"] == 1 + from trnrt import __version__ + + assert info["version"] == __version__ + + +def test_model_info_reads_the_checkpoint_but_never_names_its_download_path(tmp_path): + """The engine's resolved checkpoint (a download cache under the server's home) is read for model_type and + architectures; the path the clients see is the one the operator passed.""" + from fastapi.testclient import TestClient + + from trnrt.serve import api + + (tmp_path / "config.json").write_text(json.dumps({"model_type": "glm5_next", "architectures": ["Glm5NextForCLM"]})) + eng = ScriptedEngine(PLAIN) + eng.model_path = str(tmp_path) + with TestClient(api.build_app(eng, "glm")) as c: + info = c.get("/model_info").json() + assert (info["model_type"], info["architectures"]) == ("glm5_next", ["Glm5NextForCLM"]) + for path in ("/model_info", "/server_info"): + assert str(tmp_path) not in c.get(path).text + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +@pytest.mark.parametrize("stream", [False, True]) +def test_a_generate_client_that_disconnects_aborts_its_requests(stream): + """Every prompt of a batch is aborted when its client goes away, streamed or not, as on the OpenAI routes.""" + eng = ScriptedEngine(PLAIN * 40) + eng.step_delay = 0.05 # 160 tokens: 8 s to finish on its own + with _uvicorn(eng) as base: + port = int(base.rsplit(":", 1)[1]) + body = json.dumps( + {"input_ids": [[1], [2]], "sampling_params": {"max_new_tokens": 160}, "stream": stream} + ).encode() + c = socket.create_connection(("127.0.0.1", port)) + c.sendall( + b"POST /generate HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\n" + b"Content-Length: %d\r\n\r\n%s" % (len(body), body) + ) + _wait_for(lambda: len(eng.scheduler.running) == 2) + rids = sorted(q.rid for q in eng.scheduler.running) + c.close() + _wait_for(lambda: len(eng.aborted) == 2, 3) + assert sorted(eng.aborted) == rids and all(len(q.output_ids) < 160 for q in eng.added) From 3a4e66b61d706419d00f5a76969e10addd62785e Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 026/135] tests: the router's SGLang native routes --- tests/test_router_native.py | 212 ++++++++++++++++++++++++++++++++++++ 1 file changed, 212 insertions(+) create mode 100644 tests/test_router_native.py diff --git a/tests/test_router_native.py b/tests/test_router_native.py new file mode 100644 index 0000000..f463ba3 --- /dev/null +++ b/tests/test_router_native.py @@ -0,0 +1,212 @@ +"""The router's SGLang native routes (trnrt/serve/router.py): /generate routed like a completion by its prompt or +session, and the info, load, health, abort and flush routes sent to one worker or to all and their answers combined.""" + +import asyncio +import json + +import pytest + +for _m in ("httpx", "fastapi", "uvicorn"): + pytest.importorskip(_m) + +from tests.test_router import _free_port, _start, _stop # noqa: E402 after the importorskip of its stack +from trnrt.serve.router import Router, _request_text, _session_key, build_router_app # noqa: E402 the same + + +def test_generate_prompts_are_the_routing_text(): + """/generate's `text` or `input_ids` (a batch's list too) is what the prefix trees match, as a completion's + `prompt`; without this every /generate looked like an empty prompt.""" + assert _request_text("/generate", {"text": "You are helpful. Q: 2+2?"}) == "You are helpful. Q: 2+2?" + assert _request_text("/generate", {"input_ids": [5, 6, 7]}) == "[5, 6, 7]" + assert _request_text("/generate", {"text": ["a b", "c"]}) == json.dumps(["a b", "c"]) + assert _request_text("/generate", {"input_ids": [[1, 2], [3]]}) == "[[1, 2], [3]]" + assert _request_text("/v1/completions", {"prompt": "x", "text": "y"}) == "x" + assert _session_key("/generate", {"text": "abc"}) == _session_key("/v1/completions", {"prompt": "abc"}) + assert _session_key("/generate", {"text": "abc"}) != _session_key("/generate", {"text": "abd"}) + + +def _native_worker(i: int, flush_status: int = 200, health_status: int = 200): + """A worker with SGLang's native routes: it answers each with its index, and records what it was sent. (No + `from __future__ import annotations` in this module: FastAPI resolves a handler's `req: Request` in the module's + globals, where this local import is not, and would take `req` for a query parameter.)""" + from fastapi import FastAPI, Request + from fastapi.responses import JSONResponse, PlainTextResponse, Response + + app = FastAPI() + app.state.seen = [] + + @app.get("/health") + async def health(): + return {"status": "ok"} + + @app.post("/generate") + async def generate(req: Request): + body = await req.json() + app.state.seen.append(("generate", body)) + return {"text": f"w{i}", "output_ids": [i], "meta_info": {"id": body.get("rid", "x")}} + + @app.get("/get_load") + async def get_load(): + return [{"dp_rank": 0, "num_reqs": i + 1, "num_waiting_reqs": 0, "num_tokens": 10 * i}] + + @app.get("/model_info") + async def model_info(): + app.state.seen.append(("model_info", None)) + return {"model_path": "org/glm", "worker": i} + + @app.get("/server_info") + async def server_info(): + return {"tp_size": 32, "dp_size": 1, "worker": i, "internal_states": [{"dp_rank": 0, "num_reqs": i}]} + + @app.get("/health_generate") + async def health_generate(): + if health_status != 200: + return JSONResponse(status_code=health_status, content={"status": "failed", "error": "stuck"}) + return {"status": "ok"} + + @app.post("/abort_request") + async def abort_request(req: Request): + body = await req.json() + app.state.seen.append(("abort", body)) + if not body.get("rid") and not body.get("abort_all"): + return JSONResponse(status_code=400, content={"detail": "rid is required"}) + return Response(status_code=200) + + @app.api_route("/flush_cache", methods=["GET", "POST"]) + async def flush_cache(req: Request): + app.state.seen.append(("flush", req.method, req.url.query)) + if flush_status != 200: + return PlainTextResponse("Cache not flushed: 2 requests are running or waiting.\n", status_code=400) + return PlainTextResponse("Cache flushed.\n") + + return app + + +def _run(workers, check, router_kw=None) -> None: + async def main() -> None: + ports = [_free_port() for _ in range(len(workers) + 1)] + router = Router([f"http://127.0.0.1:{p}" for p in ports[:-1]], **(router_kw or {})) + apps = [(w, p) for w, p in zip(workers, ports) if w is not None] + servers, tasks = await _start([*apps, (build_router_app(router, poll_s=0), ports[-1])]) + try: + import httpx + + async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{ports[-1]}", timeout=60) as c: + await check(c, router) + finally: + await _stop(servers, tasks) + + asyncio.run(main()) + + +def test_generate_goes_through_the_prefix_routed_proxy(): + """A /generate is routed like a completion: a repeated long prompt (one prompt or a batch) stays on the worker + that has its prefix, a session's requests on its pin; the body reaches the worker whole.""" + w0, w1 = _native_worker(0), _native_worker(1) + + async def check(c, router) -> None: + long = "You are a careful assistant. " * 40 + r = await c.post("/generate", json={"text": long + "Q1", "sampling_params": {"max_new_tokens": 4}}) + assert r.status_code == 200, r.text + first = r.json()["text"] + for q in range(2, 6): + r = await c.put("/generate", json={"text": long + f"Q{q}"}) + assert r.json()["text"] == first + assert router.routed["cache"] == 4 and router.load == [0, 0] + batch = {"input_ids": [[7] * 300 + [1], [7] * 300 + [2]], "rid": "b"} + w = (await c.post("/generate", json=batch)).json()["text"] + assert (await c.post("/generate", json={**batch, "rid": "c"})).json()["text"] == w + assert ("generate", batch) in (w0 if w == "w0" else w1).state.seen + r = await c.post("/generate", json={"text": "hi", "session_id": "s9"}) + assert router.sessions["s9"] == int(r.json()["text"][1]) + + _run([w0, w1], check) + + +def test_info_and_load_routes_combine_the_workers(): + """model_info: one live worker's answer. get_load: every worker's entry with dp_rank its index. server_info: the + first answer, dp_size the workers, internal_states every worker's. A dead worker is skipped (and marked down).""" + w1, w2 = _native_worker(1), _native_worker(2) + + async def check(c, router) -> None: + r = await c.get("/get_load") + assert r.status_code == 200, r.text + assert [(e["dp_rank"], e["num_reqs"]) for e in r.json()] == [(1, 2), (2, 3)] + assert router.up == [False, True, True] + for path in ("/server_info", "/get_server_info"): + info = (await c.get(path)).json() + assert info["dp_size"] == 3 and info["tp_size"] == 32 and info["worker"] == 1 + assert [(s["dp_rank"], s["num_reqs"]) for s in info["internal_states"]] == [(1, 1), (2, 2)] + for path in ("/model_info", "/get_model_info"): + assert (await c.get(path)).json()["model_path"] == "org/glm" + assert len([s for s in w1.state.seen + w2.state.seen if s[0] == "model_info"]) == 2 # one worker each time + + _run([None, w1, w2], check) + + +def test_health_generate_follows_every_worker(): + async def check(c, router) -> None: + r = await c.get("/health_generate") + assert r.status_code == 200 and r.json()["status"] == "degraded", r.text + assert r.json()["down"] == [router.workers[1]] and router.up == [True, False] + + _run([_native_worker(0), _native_worker(1, health_status=503)], check) + + async def none(c, router) -> None: + r = await c.get("/health_generate") + assert r.status_code == 503 and r.json()["status"] == "down" + + _run([None, _native_worker(1, health_status=503)], none) + + async def both(c, router) -> None: + router.up = [False, False] + r = await c.get("/health_generate") + assert r.status_code == 200 and r.json() == {"status": "ok", "down": []} + assert router.up == [True, True] # each one generated: up + + _run([_native_worker(0), _native_worker(1)], both) + + +def test_abort_request_reaches_every_worker(): + """A rid's request is on one worker, which the router need not know: every worker gets the abort. A malformed abort + is refused as the workers refuse it.""" + w0, w1 = _native_worker(0), _native_worker(1) + + async def check(c, router) -> None: + for body in ({"rid": "r1"}, {"abort_all": True}): + r = await c.post("/abort_request", json=body) + assert r.status_code == 200, r.text + assert ("abort", body) in w0.state.seen and ("abort", body) in w1.state.seen + r = await c.post("/abort_request", json={}) + assert r.status_code == 400 and "rid" in r.text + r = await c.post("/abort_request", content=b"[1]", headers={"content-type": "application/json"}) + assert r.status_code == 400 + + _run([w0, w1], check) + + +def test_flush_cache_flushes_every_worker_and_their_prefix_trees(): + """200 only when every worker flushed; the router forgets a flushed worker's prefixes (its tree) but keeps the + others', and the query (timeout) reaches the workers.""" + w0, w1 = _native_worker(0), _native_worker(1, flush_status=400) + + async def check(c, router) -> None: + for _ in range(2): + await c.post("/generate", json={"text": "x" * 500}) + await c.post("/generate", json={"text": "y" * 500}) + router.trees[0].insert("z" * 100) + router.trees[1].insert("z" * 100) + r = await c.get("/flush_cache?timeout=3") + assert r.status_code == 400 and "engine 1: Cache not flushed: 2 requests" in r.text, r.text + assert router.trees[0].size == 0 and router.trees[1].size > 0 + assert ("flush", "GET", "timeout=3") in w0.state.seen and ("flush", "GET", "timeout=3") in w1.state.seen + + _run([w0, w1], check) + + async def ok(c, router) -> None: + router.trees[0].insert("z" * 100) + r = await c.post("/flush_cache") + assert r.status_code == 200 and r.text == "Cache flushed.\n" + assert all(t.size == 0 for t in router.trees) + + _run([_native_worker(0), _native_worker(1)], ok) From 89bf2d020ab2b82778c294e89563c02afaa2cb95 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 027/135] tests: /generate matches /v1/completions and /flush_cache empties both cache tiers on a real engine --- tests/test_api_native.py | 100 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 tests/test_api_native.py diff --git a/tests/test_api_native.py b/tests/test_api_native.py new file mode 100644 index 0000000..897730c --- /dev/null +++ b/tests/test_api_native.py @@ -0,0 +1,100 @@ +"""SGLang's native API (trnrt/serve/native.py) on a GLM-5.3-Flash engine on CPU (a tiny random model, +tests/test_glm5_next.py), with the host-memory tier on: /generate gives the tokens and logprobs /v1/completions gives, +/flush_cache empties the prefix cache of both tiers so the next turn is cold and the one after it hits again, and +/health_generate runs its token through the engine. +""" + +from __future__ import annotations + +import json + +import pytest +import torch + +HOST = 0.05 # GiB per rank of host-memory tier: ample for the tiny model's pages and rows + + +@pytest.fixture(scope="module") +def client(tmp_path_factory): + from fastapi.testclient import TestClient + + from tests.test_api_cached_tokens import _save_tokenizer + from trnrt.serve.api import build_app + + glm = pytest.importorskip("tests.test_glm5_next") + d = tmp_path_factory.mktemp("api_native") + glm.build(str(d), seed=1, index_topk=16) + cfg = json.loads((d / "config.json").read_text()) + vocab = cfg.get("vocab_size") or cfg["text_config"]["vocab_size"] + _save_tokenizer(d, vocab) + eng = glm.engine(str(d), max_prefill_tokens=64, host_cache_gb=HOST) + with TestClient(build_app(eng, "test-model")) as c: + c.engine, c.vocab = eng, vocab + yield c + + +def _random(c, g, n): + return torch.randint(1, c.vocab, (n,), generator=g).tolist() + + +def _completion(c, prompt, n): + r = c.post( + "/v1/completions", json={"prompt": prompt, "max_tokens": n, "temperature": 0, "logprobs": 0, "ignore_eos": True} + ) + assert r.status_code == 200, r.text + choice = r.json()["choices"][0] + ids = c.engine.tokenizer.convert_tokens_to_ids(choice["logprobs"]["tokens"]) + return ids, choice["logprobs"]["token_logprobs"], r.json()["usage"]["prompt_tokens_details"]["cached_tokens"] + + +def _generate(c, prompt, n): + body = { + "input_ids": prompt, + "sampling_params": {"max_new_tokens": n, "temperature": 0, "ignore_eos": True}, + "return_logprob": True, + } + r = c.post("/generate", json=body) + assert r.status_code == 200, r.text + return r.json() + + +def test_generate_matches_completions(client): + c = client + prompt = _random(c, torch.Generator().manual_seed(11), 21) + ids, lps, _ = _completion(c, prompt, 5) + out = _generate(c, prompt, 5) + meta = out["meta_info"] + assert out["output_ids"] == ids and meta["completion_tokens"] == 5 and meta["prompt_tokens"] == 21 + assert [t for _, t, _ in meta["output_token_logprobs"]] == ids + assert [lp for lp, _, _ in meta["output_token_logprobs"]] == pytest.approx(lps, abs=1e-5) + assert meta["finish_reason"] == {"type": "length", "length": 5} + assert out["text"] == c.engine.tokenizer.decode(ids, skip_special_tokens=True) + + +def test_flush_cache_empties_both_tiers_and_the_cache_works_after(client): + c = client + eng = c.engine + g = torch.Generator().manual_seed(12) + first = _random(c, g, 25) + reply, _, cold = _completion(c, first, 4) + second = first + reply + _random(c, g, 5) # the next turn: resumes from the first's checkpoint + _, _, hit = _completion(c, second, 4) + assert cold == 0 and hit > 0 + assert eng.radix.total_pages() > 0 and (eng.host_tier.entries or eng.radix.num_ckpts) + + r = c.post("/flush_cache?timeout=30") # the engine may still be reading back its last step: wait for idle + assert r.status_code == 200 and r.text.startswith("Cache flushed."), r.text + assert eng.radix.total_pages() == 0 and eng.radix.num_ckpts == 0 + assert eng.pool.num_free == eng.pool.num_usable + assert not eng.host_tier.entries and not eng.host_tier.ckpts + + again, _, cached = _completion(c, second, 4) + assert cached == 0 # cold: nothing of the first turn is left + third = second + again + _random(c, g, 5) + _, _, hit = _completion(c, third, 4) + assert hit > 0 # the cache serves again after the flush + + +def test_health_generate_on_the_engine(client): + r = client.get("/health_generate") + assert r.status_code == 200 and r.json() == {"status": "ok"} From cbea0e761111d9c22ce328922744a9b5ea071f48 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 028/135] docs: SGLang's /generate and the info and admin endpoints --- docs/API.md | 66 ++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 63 insertions(+), 3 deletions(-) diff --git a/docs/API.md b/docs/API.md index 7cd2d4b..efcc566 100644 --- a/docs/API.md +++ b/docs/API.md @@ -1,6 +1,6 @@ # API -The router on `--port` (default 8000) speaks the OpenAI API: `/v1/chat/completions`, `/v1/completions`, `/v1/models` and `/health`, plus `/tokenize` and `/detokenize` ([below](#tokenize-and-detokenize)). Any OpenAI client works with `base_url="http://:8000/v1"`. +The router on `--port` (default 8000) speaks the OpenAI API: `/v1/chat/completions`, `/v1/completions`, `/v1/models` and `/health`, plus `/tokenize` and `/detokenize` ([below](#tokenize-and-detokenize)). Any OpenAI client works with `base_url="http://:8000/v1"`. It also speaks SGLang's native API: [`/generate`](#generate) and the [info and admin endpoints](#info-and-admin-endpoints). ## Example @@ -30,8 +30,9 @@ curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' - - A response cut off by `max_tokens` has `finish_reason` `"length"`, even when it ends in a tool call: that call may be incomplete. - A client that disconnects aborts its request, streamed or not. -- Endpoints on `--port` (the router): `/v1/chat/completions`, `/v1/completions`, `/v1/models`, `/health`, `/tokenize` - and `/detokenize`. Each +- Endpoints on `--port` (the router): `/v1/chat/completions`, `/v1/completions`, `/v1/models`, `/health`, `/tokenize`, + `/detokenize`, and SGLang's `/generate`, `/model_info`, `/server_info`, `/get_load`, `/health_generate`, + `/abort_request` and `/flush_cache`. Each engine also serves Prometheus metrics at `/metrics` on its own port (8001 and 8002 with the default `--port`), on 127.0.0.1 only, among them the sparse-attention selection's counters (`trnrt:dsa_local_k_*`, [KNOWN_ISSUES.md](KNOWN_ISSUES.md)). @@ -66,3 +67,62 @@ too. curl -s localhost:8000/tokenize -H 'Content-Type: application/json' \ -d '{"messages": [{"role": "user", "content": "What is the weather in Paris?"}]}' ``` + +## Generate + +`/generate` (POST or PUT) is SGLang's native generation endpoint, with SGLang's request and response shapes, so SGLang +clients work unchanged. It runs through the same engine path as `/v1/completions`: the same checks, stop strings, and a +client that disconnects aborts its request. The router routes it like a completion, by its prompt or its `session_id`. + +| Request field | Meaning | +|---|---| +| `text` or `input_ids` | One prompt (a string, or a list of token ids), or a batch (a list of either). Exactly one of the two. | +| `sampling_params` | SGLang's names: `max_new_tokens` (default 128), `temperature`, `top_p`, `top_k`, `min_p`, `stop`, `stop_token_ids`, `ignore_eos`, `sampling_seed`. A batch may send one object per prompt. Defaults come from the checkpoint's `generation_config.json`, as on the OpenAI routes. | +| `stream` | Server-sent events, then `data: [DONE]`. Each event carries the output so far (SGLang's default): `text`, `output_ids` and `meta_info`. A batch's events also carry their prompt's `index`. | +| `return_logprob`, `top_logprobs_num` (0 to 20), `return_text_in_logprobs` | Output-token logprobs in `meta_info`. | +| `logprob_start_len` | `-1` (or absent) only, when `return_logprob` is set. | +| `rid` | The request's id, which `/abort_request` names. A batch's ids are `rid_0`, `rid_1`, ..., or a list of one per prompt. A rid already running is refused. | +| `session_id` | As on the OpenAI routes: the router pins the session to one engine. | +| `return_prompt_token_ids` | Adds `prompt_token_ids` to the response. | +| `log_metrics` | `false` leaves the request out of `/metrics`. | + +The response is one object for one prompt and a list for a batch: + +| Field | Value | +|---|---| +| `text` | The output, cut before a stop string. | +| `output_ids` | The output token ids. | +| `meta_info.id` | The rid. | +| `meta_info.finish_reason` | `{"type": "stop", "matched": }`, `{"type": "length", "length": n}`, `{"type": "abort", "message": "Aborted"}`, or `null` while a stream runs. | +| `meta_info.prompt_tokens`, `completion_tokens`, `cached_tokens` | Token counts; `cached_tokens` is the prompt tokens resumed from the prefix cache. | +| `meta_info.num_retractions`, `e2e_latency` | Times the request was preempted; seconds from arrival to finish. | +| `meta_info.output_token_logprobs` | With `return_logprob`: `[logprob, token_id, text or null]` per output token. | +| `meta_info.output_top_logprobs` | With `top_logprobs_num` > 0: a list of `[logprob, token_id, text or null]` per output token. | +| `meta_info.input_token_logprobs`, `input_top_logprobs` | Always `[]`: prompt logprobs are not computed. | + +Answered with 400, naming the field: prompt logprobs (`logprob_start_len` >= 0 with `return_logprob`), `n` > 1, +guided decoding (`json_schema`, `regex`, `ebnf`, `structural_tag`), `min_new_tokens`, `no_stop_trim`, `logit_bias`, +`skip_special_tokens` false, penalties, multimodal input, LoRA, hidden states, `priority`, `session_params` and the +other GenerateReqInput fields this server does not serve. An unknown `sampling_params` key is refused too, as SGLang +does. Each is accepted at its neutral value. + +```bash +curl -s localhost:8000/generate -H 'Content-Type: application/json' \ + -d '{"text": "The capital of France is", "sampling_params": {"max_new_tokens": 16, "temperature": 0}}' +``` + +## Info and admin endpoints + +SGLang's names (the `get_` forms are SGLang's older ones, served too). The router answers each from one engine or from +all of them: + +| Endpoint | Engine answer | Router | +|---|---|---| +| `GET /model_info`, `/get_model_info` | `model_path` (as passed to `--model`), `served_model_name`, `is_generation`, `preferred_sampling_params` (the checkpoint's defaults), `reasoning_parser` and `tool_call_parser` (SGLang's names for the equivalent parsers: `glm45`, `glm47`), `model_type`, `architectures` | The first live engine | +| `GET /server_info`, `/get_server_info` | The engine's configuration (every `EngineConfig` field), `tp_size`, `dp_size`, `context_length`, `preset`, `version`, `startup_seconds`, and `internal_states` (its `/get_load` entry) | The first answer, with `dp_size` the number of engines and `internal_states` one entry per engine | +| `GET /get_load` | `[{dp_rank, num_reqs, num_running_reqs, num_waiting_reqs, num_tokens, num_pending_tokens, token_usage, ts_tic}]`: `num_tokens` is the KV tokens in use plus the waiting requests' tokens | Every engine's entry, `dp_rank` its index | +| `GET /health_generate` | 200 once a step ends after a one-token probe is sent (the probe itself on an idle engine); 503 if none ends within 20 s, or once the engine failed. The probe is not counted in `/metrics` | Every engine: 200 `"ok"` or `"degraded"` while any generated, 503 `"down"` when none did, naming the engines down | +| `POST /abort_request` | `{"rid": ...}` aborts that request (an exact match); `{"abort_all": true}` aborts every request. The aborted request answers with its output so far and finish reason `abort`. 200 whether or not the rid runs | Every engine: only the one holding the rid has it. 200 once any engine took it | +| `GET` or `POST /flush_cache` | Drops the prefix cache (device pages, state checkpoints and the host-memory tier), only while no request runs or waits. `?timeout=` (default 0) waits up to that long for the engine to be idle. 200 when flushed, 400 with the reason when not | Every engine; 200 only when all flushed. The router also forgets a flushed engine's prefixes | + +None of these returns a secret or a host path the operator did not pass. From 49ac78a0613839244e468c241c898b23c656ff09 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 029/135] docs: serve/native.py in the architecture map --- docs/ARCHITECTURE.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 9b2f063..78b9893 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -17,6 +17,7 @@ | `trnrt/platform.py` | Trainium2 facts, the neuronx-cc argument set and the runtime environment. | | `trnrt/metrics.py` | Prometheus text for each engine's `/metrics`. | | `trnrt/serve/api.py` | The OpenAI-compatible server for one engine: request validation, chat template, streaming. | +| `trnrt/serve/native.py` | SGLang's native API on one engine: `/generate`, the info and load endpoints, `/health_generate`, `/abort_request`, `/flush_cache`. | | `trnrt/serve/router.py` | The session-affine, prefix-aware router in front of the engines, with failover. | | `trnrt/serve/parsers.py` | GLM reasoning and tool-call parsing, streamed and whole. | | `trnrt/engine/engine.py` | `LLMEngine`: model build, the step loop, start and close. | From 7b3d7323a222d8fba9dbc4972ac39a40db487cd4 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:04:17 +0530 Subject: [PATCH 030/135] changelog: SGLang's /generate and the info and admin endpoints --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0748e1b..dc25d60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,12 @@ and versions follow [Semantic Versioning](https://semver.org/). Rebuild the comp - `/tokenize` and `/detokenize` (also under `/v1`), as SGLang and vLLM serve them. A chat request is counted exactly as `/v1/chat/completions` would send it to the engine. The router sends them to any live engine without touching its routing state ([docs/API.md](docs/API.md)). +- SGLang's native `/generate` (also PUT): `text` or `input_ids`, one prompt or a batch, SGLang's `sampling_params`, + streaming, output logprobs and request ids. The router routes it by its prompt or session + ([docs/API.md](docs/API.md#generate)). +- SGLang's `/model_info`, `/server_info`, `/get_load`, `/health_generate`, `/abort_request` and `/flush_cache` (and + their `get_` names), on each engine and on the router, which asks one engine or all of them + ([docs/API.md](docs/API.md#info-and-admin-endpoints)). ## [0.1.0] - 2026-11-08 From a6b706e0e4b706d75404d8292c5a07400b707c10 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 031/135] api: n choices, SGLang's stop_regex, kept stop text, skip_special_tokens and matched_stop --- trnrt/serve/api.py | 619 ++++++++++++++++++++++++++++++++++----------- 1 file changed, 465 insertions(+), 154 deletions(-) diff --git a/trnrt/serve/api.py b/trnrt/serve/api.py index 341a026..0b2a641 100644 --- a/trnrt/serve/api.py +++ b/trnrt/serve/api.py @@ -4,11 +4,14 @@ output's reasoning (`reasoning_content`) and tool calls (`tool_calls`) parsed out of it (serve/parsers.py), as one response or as server-sent events; - POST /v1/completions: a prompt as text or token ids, through the same generation path; +- POST /tokenize and /detokenize (also under /v1); - GET /v1/models; - GET /health: 200 while the engine can serve, 503 with the reason once it cannot; - GET /v1/stats: the KV-cache gauges the router (serve/router.py) uses to break ties between engines; - GET /metrics: Prometheus counters and histograms (trnrt/metrics.py). +A request may ask for `n` choices: each is its own engine request (Run), sharing the prompt through the prefix cache. + A request may carry a `session_id` (a string): the engine keeps that session's prefix referenced in its prefix cache, and the router pins the session to one engine. @@ -23,7 +26,9 @@ import atexit import contextlib import copy +import dataclasses import functools +import importlib import json import logging import math @@ -47,7 +52,15 @@ from ..engine.request import Request, SamplingParams, Status from ..metrics import Metrics, kv_pages, spec_stats -from .parsers import THINK_OPEN, StreamParser, chat_finish_reason, parse_full, template_messages, tool_choice_error +from .parsers import ( + THINK_CLOSE, + THINK_OPEN, + StreamParser, + chat_finish_reason, + parse_full, + template_messages, + tool_choice_error, +) if TYPE_CHECKING: # only annotations: importing this module does not load the engine and its model code from ..engine.engine import LLMEngine @@ -61,7 +74,7 @@ class _Stream: req: Request loop: asyncio.AbstractEventLoop - out: asyncio.Queue + out: asyncio.Queue | _Tagged sent: int = 0 @@ -77,7 +90,19 @@ class _Submit: params: SamplingParams loop: asyncio.AbstractEventLoop session_id: str | None = None - out: asyncio.Queue = field(default_factory=asyncio.Queue) + out: asyncio.Queue | _Tagged = field(default_factory=asyncio.Queue) + + +class _Tagged: + """One choice's view of a shared queue: each item put here arrives on `q` as (index, item). The engine thread only + ever calls put_nowait on a request's queue (on its event loop), so the n requests of one HTTP request (`n` > 1) can + share one queue and be read in arrival order.""" + + def __init__(self, q: asyncio.Queue, index: int): + self.q, self.index = q, index + + def put_nowait(self, item) -> None: + self.q.put_nowait((self.index, item)) # Engine loops of this process, held weakly: at exit each loop's thread stops after its step, then its engine closes @@ -164,7 +189,9 @@ def _put(self, kind: str, payload) -> None: raise HTTPException(status_code=503, detail=f"engine failed: {err}") self.inbox.put((kind, payload)) - def submit(self, prompt_ids: list[int], params: SamplingParams, session_id: str | None = None) -> _Submit: + def submit( + self, prompt_ids: list[int], params: SamplingParams, session_id: str | None = None, out: _Tagged | None = None + ) -> _Submit: # Checked here, before the engine sees the request: an id outside the vocabulary fails the step, and at tp > 1 # one failed step fails the engine for every request (_fail). if not prompt_ids or not all(type(t) is int and 0 <= t < self.vocab for t in prompt_ids): @@ -174,6 +201,8 @@ def submit(self, prompt_ids: list[int], params: SamplingParams, session_id: str if session_id is not None and not isinstance(session_id, str): raise HTTPException(status_code=400, detail="session_id must be a string") sub = _Submit(uuid.uuid4().hex, prompt_ids, params, asyncio.get_running_loop(), session_id) + if out is not None: + sub.out = out self._put("add", sub) return sub @@ -334,14 +363,35 @@ class _Detok: Never emits half of a multi-byte character, and with stop strings it holds back the last `max(len(stop)) - 1` characters until they cannot start a stop string, so a stop sequence is never partially streamed. Output is cut before the first stop string, as - the OpenAI API specifies. + the OpenAI API specifies, or after it with `keep` (vLLM's include_stop_str_in_output, + SGLang's no_stop_trim). + + `regexes` are SGLang's stop_regex patterns, as (compiled, the longest text a match can span or None): the output + ends at the end of the first match, the matched text included, as SGLang's does (its trim looks for the pattern + itself in the text, which a pattern rarely is). A match never takes back text already streamed. `matched`: the stop + string or pattern that ended the output, None until one does. """ - def __init__(self, tokenizer, stops: tuple[str, ...] = (), incremental: bool = False, rid: str = ""): + def __init__( + self, + tokenizer, + stops: tuple[str, ...] = (), + incremental: bool = False, + rid: str = "", + regexes: tuple[tuple[re.Pattern, int | None], ...] = (), + keep: bool = False, + skip_special: bool = True, + ): self.tok = tokenizer self.rid = rid self.stops = stops - self.hold = max((len(x) for x in stops), default=1) - 1 + self.keep = keep + # With keep the text up to a stop string's end is emitted either way, so nothing needs holding back. + self.hold = 0 if keep else max((len(x) for x in stops), default=1) - 1 + self.regexes = regexes + self.skip_special = skip_special + self.matched: str | None = None + self._searched = 0 # how much of the text the regexes have been searched over self.ids: list[int] = [] self.emitted = 0 self.stopped = False @@ -360,7 +410,7 @@ def _check(self, full: str | None) -> None: stream ended without a stop string. `full`: that decode when the caller has it (the final push), else decoded here.""" if full is None: - full = self.tok.decode(self.ids, skip_special_tokens=True) + full = self.tok.decode(self.ids, skip_special_tokens=self.skip_special) got = "".join(self._streamed or ()) # set for an incremental stream, the only caller (push) DETOK_CHECK["checked"] += 1 if full.startswith(got) and (self.stopped or got == full): @@ -379,10 +429,11 @@ def _check(self, full: str | None) -> None: ) def _decoded(self, final: bool) -> str: + skip = self.skip_special if final or not self.incremental: - return self.tok.decode(self.ids, skip_special_tokens=True) - prefix = self.tok.decode(self.ids[self._po : self._ro], skip_special_tokens=True) - new = self.tok.decode(self.ids[self._po :], skip_special_tokens=True) + return self.tok.decode(self.ids, skip_special_tokens=skip) + prefix = self.tok.decode(self.ids[self._po : self._ro], skip_special_tokens=skip) + new = self.tok.decode(self.ids[self._po :], skip_special_tokens=skip) if new.startswith(prefix) and len(new) > len(prefix) and not new.endswith("\ufffd"): self._text += new[len(prefix) :] self._po, self._ro = self._ro, len(self.ids) @@ -395,13 +446,41 @@ def _decoded(self, final: bool) -> str: return self._text + new[len(prefix) :] return self._text + "\ufffd" # an incomplete character: held back below, as the full decode's would be - def push(self, ids: list[int], final: bool) -> str: + def _stop_at(self, text: str, final: bool) -> int | None: + """Where the output ends, if a stop string or a stop regex is in `text`: the earliest to start wins (a stop + string on a tie). Sets `matched`.""" + best: tuple[int, int, str] | None = None # (start, end of the output, what matched) + for x in self.stops: + i = text.find(x) + if i >= 0 and (best is None or i < best[0]): + best = (i, i + len(x) if self.keep else i, x) + if self.regexes: + # An incomplete character's U+FFFD is not text yet: a pattern must not match it. + body = text if final else text.rstrip("\ufffd") + for rx, width in self.regexes: + # Only a match that ends in the new text is new, and a bounded one starts at most `width` before it. + # Searching from a position (not a slice) keeps anchors and lookbehinds reading the whole text. + start = 0 if width is None else max(0, self._searched - width - 8) + m = rx.search(body, start) + if m is not None and (best is None or m.start() < best[0]): + best = (m.start(), m.end(), rx.pattern) + self._searched = len(body) + if best is None: + return None + self.matched = best[2] + return max(best[1], self.emitted) + + def push(self, ids: list[int], final: bool, drop_last: bool = False) -> str: + """The text these ids add. drop_last: the last id is the stop token that ended the request (EOS or a + stop_token_ids entry), which is not part of the text (vLLM and SGLang leave it out).""" if self.stopped: return "" self.ids.extend(ids) + if drop_last and ids: + self.ids.pop() text = self._decoded(final) - cut = min((i for i in (text.find(x) for x in self.stops) if i >= 0), default=-1) - if cut >= 0: + cut = self._stop_at(text, final) + if cut is not None: self.stopped = True end = cut elif final: @@ -420,10 +499,13 @@ def push(self, ids: list[int], final: bool) -> str: # Request fields that would change the output in a way no served graph computes (guided decoding, prompt logprobs, -# several choices, logit biases, an echoed or suffixed prompt): refused with a 400, never ignored. Each is accepted at +# logit biases, an echoed or suffixed prompt, a minimum length): refused with a 400, never ignored. Each is accepted at # its neutral values, absent or listed here, which OpenAI SDKs and agent harnesses send by default. A plain-text # response_format is neutral too; penalties and a thinking budget are the engine's to refuse (engine.unsupported), -# also only when they would change the output. +# also only when they would change the output. min_tokens (vLLM) and min_new_tokens (SGLang) need EOS and the stop +# tokens masked out of the distribution a token is drawn from, and every served graph draws on the device from the +# logits alone, with no per-row token mask (engine/sampler.py, sample_candidates): ignoring EOS on the host instead +# would put the EOS in the output and let the model write past it. _NEUTRAL: dict[str, tuple[object, ...]] = { "prompt_logprobs": (None,), "guided_json": (None,), @@ -431,21 +513,26 @@ def push(self, ids: list[int], final: bool) -> str: "guided_choice": (None,), "guided_grammar": (None,), "structured_outputs": (None,), - "n": (None, 1), - "best_of": (None, 1), "logit_bias": (None, {}), "echo": (None, False), "suffix": (None, ""), "min_tokens": (None, 0), - "include_stop_str_in_output": (None, False), + "min_new_tokens": (None, 0), "add_generation_prompt": (None, True), "continue_final_message": (None, False), - "skip_special_tokens": (None, True), + # A slow tokenizer's decode option; the fast tokenizer the server decodes with has no such switch. + "spaces_between_special_tokens": (None, True), "functions": (None, []), "function_call": (None, "none", "auto"), } +def _stream_flags(body: dict) -> tuple[bool, bool]: + """stream, and stream_options.include_usage.""" + opts = body.get("stream_options") or {} + return _flag(body, "stream"), isinstance(opts, dict) and _flag(opts, "include_usage") + + def _refused(body: dict) -> list[str]: out = [k for k, neutral in _NEUTRAL.items() if body.get(k) not in neutral] rf = body.get("response_format") @@ -581,6 +668,72 @@ def _params( ) +# `n` choices per request at most: each is its own engine request, holding a decode seat. +MAX_N = 128 +# stop_regex: patterns per request, and characters per pattern, at most. Each pattern is searched on the event loop at +# every streamed piece of text (over the new text and the longest span a match can have, or the whole text for a +# pattern with an unbounded repeat). +MAX_STOP_REGEX = 16 +MAX_STOP_REGEX_CHARS = 1024 + + +@dataclass(frozen=True) +class StopRules: + """A request's host-side stop rules (the engine matches stop strings and stop tokens itself): SGLang's stop_regex + as (compiled, the longest text a match spans or None); keep: the stop text stays in the output (vLLM's + include_stop_str_in_output, SGLang's no_stop_trim); skip_special: skip_special_tokens.""" + + regexes: tuple[tuple[re.Pattern, int | None], ...] = () + keep: bool = False + skip_special: bool = True + + +def _match_width(pattern: str) -> int | None: + """The longest text a match of `pattern` can span, None if unbounded (or unknown): the window a new match can start + in. Read from the regex parser's own width (SGLang computes the same bound, sampling_params.get_max_seq_length).""" + try: + parser = importlib.import_module("re._parser") + hi = parser.parse(pattern).getwidth()[1] + except Exception: # noqa: BLE001 a private module: without it every pattern is searched over the whole text + return None + return int(hi) if hi < 1 << 20 else None + + +def _stop_rules(body: dict) -> StopRules: + """The request's StopRules, checked: 400 for a pattern that does not compile, and for keep or skip_special_tokens + that is not a boolean.""" + rx = body.get("stop_regex") + rx = [] if rx is None else [rx] if isinstance(rx, str) else rx + if not (isinstance(rx, list) and all(isinstance(x, str) and x for x in rx)): + raise _bad("stop_regex must be a non-empty string or a list of them") + if len(rx) > MAX_STOP_REGEX or any(len(x) > MAX_STOP_REGEX_CHARS for x in rx): + raise _bad(f"stop_regex takes at most {MAX_STOP_REGEX} patterns of at most {MAX_STOP_REGEX_CHARS} characters") + compiled = [] + for x in rx: + try: + compiled.append((re.compile(x), _match_width(x))) + except (re.error, RecursionError, OverflowError) as e: + raise _bad(f"stop_regex {x!r} is not a valid regular expression: {e}") from None + keep = _flag(body, "include_stop_str_in_output") or _flag(body, "no_stop_trim") + return StopRules(tuple(compiled), keep, _flag(body, "skip_special_tokens", True)) + + +def _fanout(body: dict, params: SamplingParams) -> list[SamplingParams]: + """One SamplingParams per choice (`n`, default 1). best_of is accepted only when it asks for no more than n (it would + rank extra candidates by their logprob, which no request here returns). A seed gives choice i the seed + i, as vLLM + does (v1/engine/parallel_sampling.py), so the choices differ and each is reproducible; without one each choice draws + its own noise, keyed by its own request id.""" + n = _int(body, "n", 1, 1, MAX_N) + best_of = _int(body, "best_of", n, 1) + if best_of != n: + raise _bad(f"not supported: best_of {best_of} with n {n}; send best_of equal to n, or leave it out") + if params.seed is not None and params.seed + n - 1 > _INT64_MAX: + raise _bad(f"seed must be at most {_INT64_MAX - n + 1} with n {n}: choice i draws with seed + i") + if n == 1: + return [params] + return [params if params.seed is None else dataclasses.replace(params, seed=params.seed + i) for i in range(n)] + + # A role test naming "developer" in a Jinja template: `m.role == 'developer'`, `message['role'] != "developer"`, # `role in ['system', 'developer']`. The word in prose or in an emitted string is not one. _DEVELOPER_ROLE = re.compile(r"""role['"\]]*\s*(?:==|!=|in)\s*[\[(]?[^\])\n%}]*?['"]developer['"]""") @@ -770,6 +923,213 @@ async def _unless_disconnected(http: HTTPRequest, aw): return task.result() +class EngineError(Exception): + """An engine request failed after it was accepted: `detail` and the HTTP `status` (_Submit).""" + + def __init__(self, detail: str, status: int): + super().__init__(detail) + self.detail, self.status = detail, status + + +class Run: + """The engine requests of one HTTP request: `n` choices of one prompt (OpenAI's `n`), each its own engine request, + sharing the prompt. The prefix cache makes the copies cheap: the scheduler sees the requests queued together, makes + their shared prompt pages a checkpoint target of the first one it admits, and holds the others until they can resume + from it (engine/scheduler.py: _junctions_ahead, _prefix_in_progress), so the prompt is computed once and each other + choice computes only its last partial page. Their tokens arrive tagged with the choice's index on one queue. + + Each choice has its own detokenizer (stop strings, stop_regex, skip_special_tokens) and, for a chat request, its own + GLM parser. A choice the host stops (a stop_regex match) is aborted in the engine. Whatever ends the run early (an + error, a client gone, a consumer that stops reading) aborts every choice still running (close).""" + + def __init__( + self, + loop_: EngineLoop, + tok, + prompt_ids: list[int], + params: list[SamplingParams], + session_id, + rules: StopRules, + eos: tuple[int, ...] = (), + parse: dict | None = None, + ): + self.loop_, self.tok, self.prompt_ids, self.params = loop_, tok, prompt_ids, params + self.session_id, self.rules, self.parse = session_id, rules, parse + n = len(params) + self.q: asyncio.Queue = asyncio.Queue() + self.subs: list[_Submit] = [] + self.reqs: list[Request | None] = [None] * n + self.detoks: list[_Detok] = [] + self.parsers = [StreamParser(**parse) for _ in range(n)] if parse is not None else [] + # The ids whose token ends a request without being part of its text (EOS and the request's stop ids). + self.stop_ids = [frozenset(eos) | frozenset(p.stop_token_ids) for p in params] + self.ids: list[list[int]] = [[] for _ in range(n)] + self.text = [""] * n + self.lps: list[list] = [[] for _ in range(n)] + self.finish: list[str | None] = [None] * n + self.created = int(time.time()) + self._early: list = [] # tokens that arrived while start() still waited for another choice's acceptance + self._open = n + self.completed = False + + @property + def n(self) -> int: + return len(self.params) + + @property + def rid(self) -> str: + return self.subs[0].rid + + async def start(self) -> None: + """Hand every choice to the engine and wait until each is accepted. HTTPException (the engine's status) if one is + not: the others are aborted, so no choice runs for a request that is answered with an error.""" + try: + for i, p in enumerate(self.params): + self.subs.append(self.loop_.submit(self.prompt_ids, p, self.session_id, out=_Tagged(self.q, i))) + waiting = set(range(self.n)) + while waiting: + i, item = await self.q.get() + if item[0] == "error": + raise HTTPException(status_code=item[2], detail=item[1]) + if item[0] == "accepted": + self.reqs[i] = item[1] + waiting.discard(i) + else: + self._early.append((i, item)) + except BaseException: + self.close() + raise + for sub in self.subs: + self.detoks.append( + _Detok( + self.tok, + sub.params.stop, + incremental=True, + rid=sub.rid, + regexes=self.rules.regexes, + keep=self.rules.keep, + skip_special=self.rules.skip_special, + ) + ) + + def close(self) -> None: + """Abort every choice that has not finished (idempotent; nothing once the run completed).""" + if self.completed: + return + for i, sub in enumerate(self.subs): + if self.finish[i] is None: + self.loop_.abort(sub.rid) + + async def deliveries(self): + """(choice index, new text, finish reason or None, (id, logprob entry) pairs or None) per delivery, until every + choice has finished. EngineError if the engine fails a choice; close() runs on any early exit.""" + try: + while self._open: + i, item = self._early.pop(0) if self._early else await self.q.get() + if item[0] == "error": + raise EngineError(item[1], item[2]) + if self.finish[i] is not None: # a choice the host stopped: the tokens it made before its abort + continue + _, ids, finish, lps = item + self.ids[i] += ids + if lps is not None: + self.lps[i] += zip(ids, lps) + d = self.detoks[i] + drop = finish == "stop" and not self.rules.keep and bool(ids) and ids[-1] in self.stop_ids[i] + text = d.push(ids, finish is not None, drop_last=drop) + if d.matched is not None: + if finish is None: # a stop_regex match: the engine knows nothing of it + self.loop_.abort(self.subs[i].rid) + finish = "stop" + self.text[i] += text + if finish is not None: + self.finish[i] = finish + self._open -= 1 + yield i, text, finish, (list(zip(ids, lps)) if lps is not None else None) + self.completed = True + finally: + self.close() + + async def parsed(self): + """A chat run's deliveries through each choice's GLM parser: (index, delta dicts (StreamParser.push), finish + reason as the chat API gives it (chat_finish_reason), logprob pairs or None).""" + async for i, text, finish, pairs in self.deliveries(): + p = self.parsers[i] + deltas = p.push(text, finish is not None) + yield i, deltas, chat_finish_reason(finish, p.saw_tool), pairs + + async def wait(self, http: HTTPRequest) -> bool: + """Every choice to its end (non-streaming); False if the client disconnected first (every choice is then + aborted). EngineError if the engine fails one.""" + + async def drain() -> bool: + async for _ in self.deliveries(): + pass + return True + + try: + done = await _unless_disconnected(http, drain()) + except BaseException: + self.close() + raise + return done is not None + + def message(self, i: int) -> tuple[dict, str | None]: + """A finished chat choice's message (parse_full of its whole text, as the non-streaming chat API returns it) and + its finish reason.""" + msg = parse_full(self.text[i], **(self.parse or {})) + return msg, chat_finish_reason(self.finish[i], bool(msg.get("tool_calls"))) + + def matched_stop(self, i: int) -> str | int | None: + """SGLang's matched_stop for a finished choice: the stop string or stop_regex pattern that ended it, else the + token id (EOS or a stop id) that did, else None (a length or an abort).""" + d = self.detoks[i] if i < len(self.detoks) else None + if d is not None and d.matched is not None: + return d.matched + ids = self.ids[i] + if self.finish[i] == "stop" and ids and ids[-1] in self.stop_ids[i]: + return ids[-1] + return None + + def cached_tokens(self) -> int: + """The prompt tokens resumed from the prefix cache: the fewest of any choice's, since the prompt is counted once + (the choice that computed it).""" + return min((max(r.num_cached_tokens, 0) for r in self.reqs if r is not None), default=0) + + def usage(self) -> dict: + """OpenAI usage: the prompt counted once (as OpenAI counts it for n choices), every choice's output tokens.""" + out = sum(len(x) for x in self.ids) + return { + "prompt_tokens": len(self.prompt_ids), + "completion_tokens": out, + "total_tokens": len(self.prompt_ids) + out, + "prompt_tokens_details": {"cached_tokens": self.cached_tokens()}, + } + + def reasoning_tokens(self, i: int) -> int | None: + """The output tokens choice i spent reasoning: those before its first `` (between `` and it + when the prompt did not open the block), counted on its ids; None if the tokenizer has no single id for the + tags.""" + close = _tag_id(self.tok, THINK_CLOSE) + ids = self.ids[i] + if close is None: + return None + start = 0 + if not (self.parse or {}).get("thinking_open"): + opening = _tag_id(self.tok, THINK_OPEN) + if opening is None or opening not in ids: + return 0 + start = ids.index(opening) + 1 + end = ids.index(close, start) if close in ids[start:] else len(ids) + return end - start + + +def _tag_id(tok, tag: str) -> int | None: + """The id of a tag the tokenizer holds as one token (GLM's `` and `` are added tokens), else None.""" + ids = tok.encode(tag, add_special_tokens=False) + return ids[0] if len(ids) == 1 else None + + def build_app(engine: LLMEngine, model_name: str) -> FastAPI: app = FastAPI(title="TrnRT") loop_ = EngineLoop(engine) @@ -808,39 +1168,33 @@ async def metrics(): async def models(): return {"object": "list", "data": [{"id": model_name, "object": "model", "owned_by": "trnrt"}]} - async def _generate( - http: HTTPRequest, - body: dict, - prompt_ids: list[int], - params: SamplingParams, - chat: bool, - parse: dict | None = None, - ): - parse_args = parse or {} - stream = _flag(body, "stream") - opts = body.get("stream_options") or {} - include_usage = isinstance(opts, dict) and _flag(opts, "include_usage") - sub = loop_.submit(prompt_ids, params, body.get("session_id")) - # Admission errors (a prompt that does not fit, an unsupported parameter) arrive before any token: surface - # them as HTTP status codes rather than inside a 200 stream. - first = await sub.out.get() - if first[0] == "error": - raise HTTPException(status_code=first[2], detail=first[1]) - req = first[1] - created = int(time.time()) - oid = ("chatcmpl-" if chat else "cmpl-") + sub.rid - obj = "chat.completion" if chat else "text_completion" - detok = _Detok(tok, params.stop, incremental=True, rid=sub.rid) + eos = tuple(getattr(engine.mcfg, "eos_token_ids", ()) or ()) - def usage_of(n_out: int) -> dict: - """OpenAI usage, with prompt_tokens_details.cached_tokens: the prompt tokens this request resumed from the - prefix cache (0 cold). Final once the request has finished.""" - return { - "prompt_tokens": len(prompt_ids), - "completion_tokens": n_out, - "total_tokens": len(prompt_ids) + n_out, - "prompt_tokens_details": {"cached_tokens": max(req.num_cached_tokens, 0)}, - } + async def start_run(body: dict, prompt_ids: list[int], params: SamplingParams, parse: dict | None = None) -> Run: + """A request's choices handed to the engine and accepted (Run.start). parse: a chat request's parser arguments + (chat_prompt), None for a completion. Every request field is checked before anything is queued.""" + rules = _stop_rules(body) + choices = _fanout(body, params) + run = Run(loop_, tok, prompt_ids, choices, body.get("session_id"), rules, eos, parse) + await run.start() + return run + + async def start_chat(body: dict) -> Run: + """A /v1/chat/completions body's run: its prompt rendered on the prompt thread, its sampling checked, its choices + accepted. The Responses and Messages APIs (serve/responses.py, serve/anthropic.py) translate into this.""" + ids, parse = await asyncio.get_running_loop().run_in_executor(prompt_pool, chat_prompt, prompt_tok, body) + # max_tokens defaults to the rest of the context (vLLM's default): reasoning at the template's default effort + # (Max) can run long, and a smaller hidden default would cut the answer off + room = max(engine.cfg.max_model_len - len(ids), 1) + return await start_run( + body, ids, _params(body, room, chat=True, vocab=loop_.vocab, defaults=gen_defaults), parse + ) + + async def _generate(http: HTTPRequest, body: dict, run: Run, chat: bool): + stream, include_usage = _stream_flags(body) # checked before the run started: raises nothing here + created = run.created + oid = ("chatcmpl-" if chat else "cmpl-") + run.rid + obj = "chat.completion" if chat else "text_completion" def lp_block(lps): """OpenAI logprobs from (token id, (logprob, top ids, top logprobs)) pairs.""" @@ -864,13 +1218,15 @@ def entry(t, lp): "top_logprobs": [{tok.decode([i]): v for i, v in zip(ti, tl)} for _, (_, ti, tl) in lps], } - def chunk(text: str, finish: str | None, lps=None) -> dict: + def chunk(i: int, text: str, finish: str | None, lps=None) -> dict: if chat: - choice = {"index": 0, "delta": {"content": text} if text else {}, "finish_reason": finish} + choice = {"index": i, "delta": {"content": text} if text else {}, "finish_reason": finish} else: - choice = {"index": 0, "text": text, "finish_reason": finish} + choice = {"index": i, "text": text, "finish_reason": finish} if lps is not None: choice["logprobs"] = lp_block(lps) + if finish is not None: + choice["matched_stop"] = run.matched_stop(i) # SGLang's field return { "id": oid, "object": obj + (".chunk" if chat else ""), @@ -879,112 +1235,77 @@ def chunk(text: str, finish: str | None, lps=None) -> dict: "choices": [choice], } - sparser = StreamParser(**parse_args) if chat else None - - def chat_chunks(text: str, finish: str | None, pairs): - if sparser is None: # only a chat request builds chunks through the parser - raise ValueError("no stream parser for a completions request") - deltas = sparser.push(text, finish is not None) - finish = chat_finish_reason(finish, sparser.saw_tool) + def chat_chunks(i: int, deltas: list[dict], finish: str | None, pairs) -> list[dict]: if not deltas: - return [chunk("", finish, pairs)] if (finish or pairs) else [] + return [chunk(i, "", finish, pairs)] if (finish or pairs) else [] out = [] - for i, d in enumerate(deltas): - last = i == len(deltas) - 1 - c = chunk("", finish if last else None, pairs if last else None) + for k, d in enumerate(deltas): + last = k == len(deltas) - 1 + c = chunk(i, "", finish if last else None, pairs if last else None) c["choices"][0]["delta"] = d out.append(c) return out async def events(): - completed = False - n_streamed = 0 try: - if chat: # OpenAI's first chunk names the role - c = chunk("", None) - c["choices"][0]["delta"] = {"role": "assistant", "content": ""} - yield f"data: {json.dumps(c)}\n\n" - while True: - kind, *rest = await sub.out.get() - if kind == "error": - completed = True - yield f"data: {json.dumps({'error': {'message': rest[0], 'code': rest[1]}})}\n\n" - return - ids, finish, lps = rest - n_streamed += len(ids) - text = detok.push(ids, finish is not None) - pairs = list(zip(ids, lps)) if lps is not None else None - if text or finish or pairs: - for c in chat_chunks(text, finish, pairs) if chat else [chunk(text, finish, pairs)]: + if chat: # OpenAI's first chunk of each choice names the role + for i in range(run.n): + c = chunk(i, "", None) + c["choices"][0]["delta"] = {"role": "assistant", "content": ""} + yield f"data: {json.dumps(c)}\n\n" + if chat: + async for i, deltas, finish, pairs in run.parsed(): + for c in chat_chunks(i, deltas, finish, pairs): yield f"data: {json.dumps(c)}\n\n" - if finish is not None: - completed = True - if include_usage: # OpenAI stream_options.include_usage: one last chunk, no choices - yield ( - "data: " - + json.dumps( - { - "id": oid, - "object": obj + (".chunk" if chat else ""), - "created": created, - "model": model_name, - "choices": [], - "usage": usage_of(n_streamed), - } - ) - + "\n\n" - ) - yield "data: [DONE]\n\n" - return + else: + async for i, text, finish, pairs in run.deliveries(): + if text or finish or pairs: + yield f"data: {json.dumps(chunk(i, text, finish, pairs))}\n\n" + if include_usage: # OpenAI stream_options.include_usage: one last chunk, no choices + last = { + "id": oid, + "object": obj + (".chunk" if chat else ""), + "created": created, + "model": model_name, + "choices": [], + "usage": run.usage(), + } + yield f"data: {json.dumps(last)}\n\n" + yield "data: [DONE]\n\n" + except EngineError as e: + yield f"data: {json.dumps({'error': {'message': e.detail, 'code': e.status}})}\n\n" finally: - # The generator is closed early when the client goes away. - if not completed: - loop_.abort(sub.rid) + run.close() # the generator is closed early when the client goes away: abort what still runs if stream: # StreamingResponse closes events() when the client goes away, and its finally aborts return StreamingResponse(events(), media_type="text/event-stream") - async def collect(): - text, finish, n_out = "", None, 0 - all_lps: list[tuple[int, tuple[float, list[int], list[float]]]] = [] - while finish is None: - kind, *rest = await sub.out.get() - if kind == "error": - raise HTTPException(status_code=rest[1], detail=rest[0]) - ids, finish, lps = rest - n_out += len(ids) - if lps is not None: - all_lps.extend(zip(ids, lps)) - text += detok.push(ids, finish is not None) - return text, finish, n_out, all_lps - # Nothing cancels a plain response's handler when its client disconnects, so the request would hold its seat - # until max_tokens: watch for the disconnect, and abort. + # until max_tokens: Run.wait watches for the disconnect, and aborts. try: - done = await _unless_disconnected(http, collect()) - except asyncio.CancelledError: - loop_.abort(sub.rid) - raise - if done is None: - loop_.abort(sub.rid) - return Response(status_code=499) # nginx's "client closed request"; nobody reads it - text, finish, n_out, all_lps = done - if chat: - msg = parse_full(text, **parse_args) - finish = chat_finish_reason(finish, bool(msg.get("tool_calls"))) - choice = {"index": 0, "message": msg, "finish_reason": finish} - else: - choice = {"index": 0, "text": text, "finish_reason": finish} - if params.logprobs is not None: - choice["logprobs"] = lp_block(all_lps) + if not await run.wait(http): + return Response(status_code=499) # nginx's "client closed request"; nobody reads it + except EngineError as e: + raise HTTPException(status_code=e.status, detail=e.detail) from None + choices = [] + for i in range(run.n): + if chat: + msg, finish = run.message(i) + choice = {"index": i, "message": msg, "finish_reason": finish} + else: + choice = {"index": i, "text": run.text[i], "finish_reason": run.finish[i]} + if run.params[i].logprobs is not None: + choice["logprobs"] = lp_block(run.lps[i]) + choice["matched_stop"] = run.matched_stop(i) + choices.append(choice) return JSONResponse( { "id": oid, "object": obj, "created": created, "model": model_name, - "choices": [choice], - "usage": usage_of(n_out), + "choices": choices, + "usage": run.usage(), } ) @@ -999,25 +1320,15 @@ async def completions(http: HTTPRequest): else: raise HTTPException(status_code=400, detail="prompt must be a string or a list of token ids") # max_tokens defaults to 16, the OpenAI completions API's default - return await _generate( - http, body, ids, _params(body, 16, chat=False, vocab=loop_.vocab, defaults=gen_defaults), chat=False - ) + params = _params(body, 16, chat=False, vocab=loop_.vocab, defaults=gen_defaults) + _stream_flags(body) # checked before anything is queued + return await _generate(http, body, await start_run(body, ids, params), chat=False) @app.post("/v1/chat/completions") async def chat(http: HTTPRequest): body = await _body(http) - ids, parse = await asyncio.get_running_loop().run_in_executor(prompt_pool, chat_prompt, prompt_tok, body) - # max_tokens defaults to the rest of the context (vLLM's default): reasoning at the template's default effort - # (Max) can run long, and a smaller hidden default would cut the answer off - room = max(engine.cfg.max_model_len - len(ids), 1) - return await _generate( - http, - body, - ids, - _params(body, room, chat=True, vocab=loop_.vocab, defaults=gen_defaults), - chat=True, - parse=parse, - ) + _stream_flags(body) # checked before anything is queued + return await _generate(http, body, await start_chat(body), chat=True) # SGLang's and vLLM's /tokenize and /detokenize (also under /v1). A chat request is rendered exactly as # /v1/chat/completions renders it (chat_prompt), so `count` is what the engine would be sent; a prompt is encoded as From 5d9b04a15bca01ebea6801848560e4e66b96f627 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 032/135] tests: n and include_stop_str_in_output are served, and refused only when malformed --- tests/test_serve_api.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_serve_api.py b/tests/test_serve_api.py index 66abcd4..97d94a7 100644 --- a/tests/test_serve_api.py +++ b/tests/test_serve_api.py @@ -349,7 +349,6 @@ def test_an_sdk_request_with_every_neutral_field_is_served(serve): ("frequency_penalty", -0.2), ("repetition_penalty", 1.1), ("thinking_token_budget", 128), - ("n", 2), ("logit_bias", {"13": 5}), ("response_format", {"type": "json_object"}), ("prompt_logprobs", 0), @@ -421,7 +420,7 @@ def test_request_fields_are_checked_before_the_engine(serve): ("prompt", [len(WORDS)]), ("prompt", [1, True]), ("min_tokens", 4), - ("include_stop_str_in_output", True), + ("include_stop_str_in_output", "yes"), ("thinking_token_budget", "lots"), ("temperature", 10**400), ] From b48a20a21bb4bc8c326b5f2e53fd2e7d848a0d77 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 033/135] tests: n choices, stop_regex, kept stop text, the stop token and the refused minimum length --- tests/test_serve_sampling.py | 276 +++++++++++++++++++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 tests/test_serve_sampling.py diff --git a/tests/test_serve_sampling.py b/tests/test_serve_sampling.py new file mode 100644 index 0000000..6d92c34 --- /dev/null +++ b/tests/test_serve_sampling.py @@ -0,0 +1,276 @@ +"""The sampling extras of /v1/chat/completions and /v1/completions (trnrt/serve/api.py): `n` choices, SGLang's +stop_regex, keeping the stop text (include_stop_str_in_output / no_stop_trim), skip_special_tokens, the stop token left +out of the text, and the refusal of a minimum length. Over the scripted engine of tests/test_serve_api.py.""" + +from __future__ import annotations + +import json +import time + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from tests.test_serve_api import ID, PLAIN, TOOL_TURN, TOOLS, ScriptedEngine, _events + +MSG = {"messages": [{"role": "user", "content": "hi"}]} + + +@pytest.fixture +def app(): + from fastapi.testclient import TestClient + + from trnrt.serve import api + + started = [] + + def start(script, setup=None): + eng = ScriptedEngine(script) + if setup is not None: + setup(eng) + client = TestClient(api.build_app(eng, "glm")) + client.__enter__() + started.append((client, eng)) + return client, eng + + yield start + for client, eng in started: + client.__exit__(None, None, None) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def test_n_choices_share_the_prompt_and_count_it_once(app): + """n engine requests for one HTTP request, one prompt; each choice its index; usage counts the prompt once and + every choice's output. A seed gives choice i seed + i (vLLM), so seeded choices differ and reproduce.""" + client, eng = app(TOOL_TURN) + r = client.post("/v1/chat/completions", json={**MSG, "tools": TOOLS, "n": 3, "seed": 7}) + assert r.status_code == 200, r.text + body = r.json() + assert [c["index"] for c in body["choices"]] == [0, 1, 2] + for c in body["choices"]: + assert c["finish_reason"] == "tool_calls" and c["message"]["tool_calls"][0]["function"]["name"] == "get_weather" + reqs = eng.added[-3:] + assert len({tuple(q.prompt_ids) for q in reqs}) == 1 and len({q.rid for q in reqs}) == 3 + assert [q.params.seed for q in reqs] == [7, 8, 9] + n_prompt = len(reqs[0].prompt_ids) + assert body["usage"] == { + "prompt_tokens": n_prompt, + "completion_tokens": 3 * len(TOOL_TURN), + "total_tokens": n_prompt + 3 * len(TOOL_TURN), + "prompt_tokens_details": {"cached_tokens": 2}, + } + assert client.post("/v1/chat/completions", json={**MSG, "n": 2}).status_code == 200 + assert [q.params.seed for q in eng.added[-2:]] == [None, None] # unseeded: each its own noise (its request id) + + +def test_n_choices_stream_per_choice_chunks(app): + """Streamed, each choice opens with its role chunk, its deltas carry its index, each ends with its finish reason, + and the usage chunk sums the choices.""" + client, _ = app(PLAIN) + events = _events( + client, {**MSG, "n": 2, "stream": True, "stream_options": {"include_usage": True}, "max_tokens": 3} + ) + usage = events.pop() + assert usage["usage"]["completion_tokens"] == 6 and usage["choices"] == [] + by = {0: [], 1: []} + for e in events: + (c,) = e["choices"] + by[c["index"]].append(c) + for i in (0, 1): + assert by[i][0]["delta"] == {"role": "assistant", "content": ""} + assert "".join(c["delta"].get("reasoning_content", "") for c in by[i]) == "Hello there again" + assert [c["finish_reason"] for c in by[i]][-1] == "length" + assert sum(c["finish_reason"] is not None for c in by[i]) == 1 + + +def test_n_completions_with_logprobs(app): + client, _ = app(PLAIN) + r = client.post("/v1/completions", json={"prompt": [1], "n": 2, "max_tokens": 2, "logprobs": 1}) + choices = r.json()["choices"] + assert [c["index"] for c in choices] == [0, 1] + for c in choices: + assert c["text"] == "Hello there" and c["logprobs"]["tokens"] == ["Hello", " there"] + assert c["finish_reason"] == "length" and c["matched_stop"] is None + + +@pytest.mark.parametrize( + "extra", + [ + {"n": 0}, + {"n": 129}, + {"n": "2"}, + {"n": 2, "best_of": 3}, + {"best_of": 2}, + {"n": 2, "seed": 2**63 - 1}, + {"min_tokens": 4}, + {"min_new_tokens": 4}, + {"spaces_between_special_tokens": False}, + {"stop_regex": ""}, + {"stop_regex": [3]}, + {"stop_regex": "("}, + {"stop_regex": ["a"] * 17}, + {"no_stop_trim": "yes"}, + {"skip_special_tokens": "no"}, + ], + ids=repr, +) +def test_what_cannot_be_served_is_refused_before_the_engine(app, extra): + client, eng = app(PLAIN) + for path, body in (("/v1/chat/completions", MSG), ("/v1/completions", {"prompt": [1]})): + r = client.post(path, json={**body, **extra}) + assert r.status_code == 400, (path, r.text) + assert eng.added == [] + ok = {"n": 2, "best_of": 2, "seed": 2**63 - 2, "min_tokens": 0, "spaces_between_special_tokens": True} + assert client.post("/v1/completions", json={"prompt": [1], **ok}).status_code == 200 + + +def test_a_choice_the_engine_refuses_aborts_the_others(): + """The second choice is refused after the first was accepted: the client gets the refusal, and the first is aborted + rather than left running for nobody.""" + from fastapi.testclient import TestClient + + from trnrt.serve import api + + eng = ScriptedEngine(PLAIN * 40) + eng.step_delay = 0.01 + real, calls = eng.add_request, [] + + def second_refused(*a, **k): + calls.append(1) + if len(calls) == 2: + raise ValueError("no room") + return real(*a, **k) + + eng.add_request = second_refused + with TestClient(api.build_app(eng, "glm")) as client: + r = client.post("/v1/completions", json={"prompt": [1], "n": 2, "max_tokens": 100}) + assert r.status_code == 400 and "no room" in r.text + deadline = time.monotonic() + 5 + while eng.scheduler.running and time.monotonic() < deadline: + time.sleep(0.01) + assert eng.scheduler.running == [] and eng.aborted == [eng.added[0].rid] + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def test_a_client_that_disconnects_aborts_every_choice(): + """Real uvicorn (the TestClient never disconnects): every one of the n engine requests is aborted.""" + import socket + import threading + + import uvicorn + + from trnrt.serve import api + + eng = ScriptedEngine(PLAIN * 40) + eng.step_delay = 0.05 + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + srv = uvicorn.Server(uvicorn.Config(api.build_app(eng, "glm"), host="127.0.0.1", port=port, log_level="warning")) + th = threading.Thread(target=srv.run, daemon=True) + th.start() + try: + deadline = time.monotonic() + 30 + while not srv.started and time.monotonic() < deadline: + time.sleep(0.02) + body = json.dumps({"prompt": [1], "max_tokens": 160, "stream": True, "n": 3}).encode() + c = socket.create_connection(("127.0.0.1", port)) + c.sendall( + b"POST /v1/completions HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\n" + b"Content-Length: %d\r\n\r\n%s" % (len(body), body) + ) + deadline = time.monotonic() + 10 + while len(eng.scheduler.running) < 3 and time.monotonic() < deadline: + time.sleep(0.02) + rids = sorted(q.rid for q in eng.scheduler.running) + assert len(rids) == 3 + c.close() + deadline = time.monotonic() + 3 + while len(eng.aborted) < 3 and time.monotonic() < deadline: + time.sleep(0.02) + assert sorted(eng.aborted) == rids + finally: + srv.should_exit = True + th.join(timeout=30) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +@pytest.mark.parametrize("stream", [False, True]) +def test_stop_regex_ends_the_output_after_its_match_and_aborts_the_request(app, stream): + """SGLang's stop_regex: the output ends at the end of the first match, the matched text kept (SGLang's trim looks + for the pattern itself), finish "stop", matched_stop the pattern. The engine knows nothing of it, so the request is + aborted there. The match spans two tokens (" there", " again"), so it is only found by searching back into text + already pushed: a bounded pattern's search starts its longest match before the new text.""" + client, eng = app(PLAIN * 10) + body = {"prompt": [1], "max_tokens": 40, "stop_regex": [r"x{3}", r"re\s?ag"]} + if stream: + events = _events(client, {**body, "stream": True}) + text = "".join(e["choices"][0]["text"] for e in events) + last = events[-1]["choices"][0] + else: + last = client.post("/v1/completions", json=body).json()["choices"][0] + text = last["text"] + assert text == "Hello there ag" + assert last["finish_reason"] == "stop" and last["matched_stop"] == r"re\s?ag" + deadline = time.monotonic() + 3 + while not eng.aborted and time.monotonic() < deadline: + time.sleep(0.01) + assert eng.aborted == [eng.added[-1].rid] + + +def test_stop_regex_unbounded_and_anchored(app): + """A pattern with an unbounded repeat is searched over the whole text; an anchor keeps its meaning (the search + starts at a position, never at a slice).""" + client, _ = app(PLAIN * 4) + r = client.post("/v1/completions", json={"prompt": [1], "max_tokens": 16, "stop_regex": r"!.*again"}) + assert r.json()["choices"][0]["text"] == "Hello there again!Hello there again" + r = client.post("/v1/completions", json={"prompt": [1], "max_tokens": 16, "stop_regex": r"^there"}) + assert r.json()["choices"][0]["text"] == "Hello there again!" * 4 # "there" never starts the text + + +@pytest.mark.parametrize("key", ["include_stop_str_in_output", "no_stop_trim"]) +def test_the_stop_text_is_kept_when_asked(app, key): + client, _ = app(PLAIN) + base = {"prompt": [1], "stop": " again"} + assert client.post("/v1/completions", json=base).json()["choices"][0]["text"] == "Hello there" + r = client.post("/v1/completions", json={**base, key: True}) + assert r.json()["choices"][0]["text"] == "Hello there again" + assert r.json()["choices"][0]["matched_stop"] == " again" + events = _events(client, {**base, key: True, "stream": True}) + assert "".join(e["choices"][0]["text"] for e in events) == "Hello there again" + + +def test_the_stop_token_is_not_in_the_text(app): + """The EOS (or a stop_token_ids entry) that ends a request is left out of its text even when the tokenizer does not + mark it special, as vLLM and SGLang leave it out; matched_stop names it. Kept with no_stop_trim.""" + + def eos(eng): + eng.mcfg.eos_token_ids = (ID["!"],) + + client, _ = app(PLAIN, eos) + choice = client.post("/v1/completions", json={"prompt": [1]}).json()["choices"][0] + assert choice["text"] == "Hello there again" and choice["matched_stop"] == ID["!"] + assert client.post("/v1/completions", json={"prompt": [1], "no_stop_trim": True}).json()["choices"][0]["text"] == ( + "Hello there again!" + ) + client2, _ = app(PLAIN) + choice = client2.post("/v1/completions", json={"prompt": [1], "stop_token_ids": [ID["!"]]}).json()["choices"][0] + assert choice["text"] == "Hello there again" and choice["matched_stop"] == ID["!"] + events = _events(client2, {"prompt": [1], "stop_token_ids": [ID["!"]], "stream": True}) + assert "".join(e["choices"][0]["text"] for e in events) == "Hello there again" + + +def test_skip_special_tokens_false_keeps_special_tokens(app): + def special(eng): + eng.tokenizer.add_special_tokens({"unk_token": ""}) # id 0 becomes special + + client, _ = app([ID["Hello"], ID[""], ID["!"]], special) + assert client.post("/v1/completions", json={"prompt": [1]}).json()["choices"][0]["text"] == "Hello!" + r = client.post("/v1/completions", json={"prompt": [1], "skip_special_tokens": False}) + assert r.json()["choices"][0]["text"] == "Hello!" + events = _events(client, {"prompt": [1], "skip_special_tokens": False, "stream": True}) + assert "".join(e["choices"][0]["text"] for e in events) == "Hello!" From 56162518bbb5ec90df68cab135f509bf0ffc545a Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 034/135] responses: the OpenAI Responses API, stateless, over the chat path --- trnrt/serve/responses.py | 539 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 539 insertions(+) create mode 100644 trnrt/serve/responses.py diff --git a/trnrt/serve/responses.py b/trnrt/serve/responses.py new file mode 100644 index 0000000..4c6f3d4 --- /dev/null +++ b/trnrt/serve/responses.py @@ -0,0 +1,539 @@ +"""The OpenAI Responses API, `POST /v1/responses`, stateless, served by the chat path (serve/api.py). + +A request is translated into a /v1/chat/completions body: `instructions` and `input` become messages, function tools +become chat tools, `max_output_tokens` becomes max_tokens and `reasoning.effort` becomes reasoning_effort. Its output, +parsed by the same GLM parsers as a chat response, is translated back into Responses output items: + +| Chat output | Responses item | +|---|---| +| `reasoning_content` | a `reasoning` item, its text as `reasoning_text` content (and as a `summary_text` when `reasoning.summary` is set, as SGLang mirrors it) | +| `content` | a `message` item with one `output_text` part | +| each of `tool_calls` | a `function_call` item; its `call_id` is the chat call's id | + +Every other field goes to the chat path unchanged, which serves or refuses it as /v1/chat/completions does (sampling +extras, `session_id`, `chat_template_kwargs`). A streamed response sends the Responses events: `response.created`, +`response.in_progress`, then per item `response.output_item.added`, its deltas and its `.done` events, then +`response.completed` (`response.incomplete` when max_output_tokens cut it off, `response.failed` on an engine error). + +Nothing is stored: `store` true, `previous_response_id` and `background` are refused with a 400, and so are the fields +whose meaning needs what no served graph computes (structured output in `text.format`, `top_logprobs`) or a server-side +tool. The fields that matter follow SGLang's /v1/responses (python/sglang/srt/entrypoints/openai/serving_responses.py and +protocol.py ResponsesRequest), and the event shapes the OpenAI API's. +""" + +from __future__ import annotations + +import json +import uuid +from typing import TYPE_CHECKING, Any + +from fastapi import HTTPException +from fastapi import Request as HTTPRequest +from fastapi.responses import JSONResponse, Response, StreamingResponse + +from .api import EngineError, Run + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + +# The Responses fields translated here. Any other field goes to the chat path as it is. +_OWN = frozenset( + { + "input", + "instructions", + "max_output_tokens", + "tools", + "tool_choice", + "reasoning", + "text", + "store", + "previous_response_id", + "background", + "include", + "truncation", + "max_tool_calls", + "top_logprobs", + "metadata", + "user", + "service_tier", + "stream", + "stream_options", + "prompt_cache_key", + "safety_identifier", + } +) +# Chat fields with no meaning here, or another name in this API: refused unless absent or neutral. +_CHAT_ONLY: dict[str, tuple[object, ...]] = { + "messages": (None,), + "n": (None, 1), + "best_of": (None, 1), + "logprobs": (None, False), + "max_tokens": (None,), + "max_completion_tokens": (None,), + "prompt": (None,), + "conversation": (None,), +} +# include values served: the reasoning text, returned as the reasoning item's encrypted_content (not encrypted: this +# server keeps no key) so a stateless client can send its reasoning back with the next turn. +_INCLUDE = frozenset({"reasoning.encrypted_content"}) +_TEXT_PARTS = ("input_text", "output_text", "text") + + +def _bad(detail: str) -> HTTPException: + return HTTPException(status_code=400, detail=detail) + + +def _text_of(content, where: str) -> str: + """A message item's content as text: a string, or a list of text parts (input_text, output_text), joined as GLM's + template joins a list of text parts. 400 for any other part (an image or a file).""" + if content is None: + return "" + if isinstance(content, str): + return content + if not isinstance(content, list): + raise _bad(f"{where}: content must be a string or a list of parts") + out = [] + for part in content: + if not (isinstance(part, dict) and part.get("type") in _TEXT_PARTS and isinstance(part.get("text"), str)): + kind = part.get("type") if isinstance(part, dict) else type(part).__name__ + raise _bad(f"{where}: content part {kind!r} is not supported; this model reads text only") + out.append(part["text"]) + return "".join(out) + + +def _reasoning_text(item: dict) -> str: + """A reasoning item's text: its reasoning_text content, else its summary, else its encrypted_content (which this + server fills with the reasoning text itself).""" + for key, kind in (("content", "reasoning_text"), ("summary", "summary_text")): + parts = item.get(key) + if isinstance(parts, list): + texts = [p.get("text") for p in parts if isinstance(p, dict) and p.get("type") == kind] + if any(isinstance(t, str) and t for t in texts): + return "\n".join(t for t in texts if isinstance(t, str)) + enc = item.get("encrypted_content") + return enc if isinstance(enc, str) else "" + + +def input_messages(body: dict) -> list[dict]: + """The chat messages of a request's instructions and input. An assistant turn's items (reasoning, an assistant + message, function calls) join into one assistant message, as the chat API carries a turn; a function_call_output + becomes a tool message.""" + msgs: list[dict] = [] + ins = body.get("instructions") + if ins is not None: + if not isinstance(ins, str): + raise _bad("instructions must be a string") + msgs.append({"role": "system", "content": ins}) + inp = body.get("input") + if isinstance(inp, str): + msgs.append({"role": "user", "content": inp}) + return msgs + if not isinstance(inp, list) or not inp: + raise _bad("input must be a string or a non-empty list of items") + + def assistant() -> dict: + """The assistant message the next assistant-side item joins: the last message if it is one, else a new one.""" + if msgs and msgs[-1].get("role") == "assistant": + return msgs[-1] + msgs.append({"role": "assistant", "content": ""}) + return msgs[-1] + + for k, item in enumerate(inp): + where = f"input[{k}]" + if not isinstance(item, dict): + raise _bad(f"{where} must be an object") + kind = item.get("type", "message") + if kind == "message": + role = item.get("role") + if role not in ("user", "system", "developer", "assistant"): + raise _bad(f"{where}: role must be user, system, developer or assistant") + text = _text_of(item.get("content"), where) + if role == "assistant": + m = assistant() + m["content"] = (m.get("content") or "") + text + else: + msgs.append({"role": role, "content": text}) + elif kind == "reasoning": + m = assistant() + text = _reasoning_text(item) + m["reasoning_content"] = (m.get("reasoning_content") or "") + text + elif kind == "function_call": + name, args = item.get("name"), item.get("arguments", "") + if not isinstance(name, str) or not name or not isinstance(args, str): + raise _bad(f"{where}: a function_call needs a name and its arguments as a JSON string") + call_id = item.get("call_id") or item.get("id") or "call_" + uuid.uuid4().hex[:24] + m = assistant() + m.setdefault("tool_calls", []).append( + {"id": call_id, "type": "function", "function": {"name": name, "arguments": args}} + ) + elif kind == "function_call_output": + out = item.get("output") + msgs.append( + { + "role": "tool", + "tool_call_id": item.get("call_id"), + "content": out if isinstance(out, str) else _text_of(out, where), + } + ) + elif kind == "item_reference": + raise _bad(f"{where}: item_reference is not supported; this server stores no items") + else: + raise _bad(f"{where}: input item type {kind!r} is not supported") + return msgs + + +def _chat_tools(tools) -> list[dict] | None: + if tools is None: + return None + if not isinstance(tools, list): + raise _bad("tools must be a list") + out = [] + for t in tools: + if not (isinstance(t, dict) and t.get("type") == "function" and isinstance(t.get("name"), str)): + kind = t.get("type") if isinstance(t, dict) else None + raise _bad(f"tool type {kind!r} is not supported: only function tools with a name") + fn = {k: t[k] for k in ("name", "description", "parameters", "strict") if k in t} + out.append({"type": "function", "function": fn}) + return out + + +def _tool_choice(tc): + """The chat API's tool_choice: "auto" and "none" are served; "required" and a named function are passed on for the + chat path to refuse with its reason (they need guided decoding).""" + if isinstance(tc, dict) and tc.get("type") == "function" and "function" not in tc: + return {"type": "function", "function": {"name": tc.get("name")}} + return tc + + +def to_chat(body: dict) -> tuple[dict, dict]: + """The /v1/chat/completions body for a Responses request, and what translating its output back needs: wants_summary, + encrypted (include reasoning.encrypted_content). 400 for what cannot be served.""" + if body.get("store") is True: + raise _bad("store is not supported: this server keeps no responses; send store false") + if body.get("previous_response_id") is not None: + raise _bad("previous_response_id is not supported: this server keeps no responses; send the whole input") + if body.get("background") is True: + raise _bad("background is not supported") + if body.get("truncation") not in (None, "disabled"): + raise _bad('truncation is not supported; send "disabled" or leave it out') + if body.get("max_tool_calls") is not None: + raise _bad("max_tool_calls is not supported") + if body.get("top_logprobs") not in (None, 0): + raise _bad("top_logprobs is not supported on /v1/responses") + include = body.get("include") or [] + if not isinstance(include, list) or any(x not in _INCLUDE for x in include): + raise _bad(f"include supports only {', '.join(sorted(_INCLUDE))}") + text = body.get("text") + if text is not None: + fmt = text.get("format") if isinstance(text, dict) else None + if not isinstance(text, dict) or (fmt is not None and (not isinstance(fmt, dict) or fmt.get("type") != "text")): + raise _bad('text.format other than {"type": "text"} is not supported: it needs guided decoding') + if text.get("verbosity") is not None: + raise _bad("text.verbosity is not supported") + for key, neutral in _CHAT_ONLY.items(): + if body.get(key) not in neutral: + raise _bad(f"{key} is not a /v1/responses field; see the Responses API") + reasoning = body.get("reasoning") + if reasoning is not None and not isinstance(reasoning, dict): + raise _bad("reasoning must be an object") + reasoning = reasoning or {} + summary = reasoning.get("summary", reasoning.get("generate_summary")) + if summary not in (None, "auto", "concise", "detailed"): + raise _bad('reasoning.summary must be "auto", "concise" or "detailed"') + + chat = {k: v for k, v in body.items() if k not in _OWN} + chat["messages"] = input_messages(body) + tools = _chat_tools(body.get("tools")) + if tools: + chat["tools"] = tools + if body.get("tool_choice") is not None: + chat["tool_choice"] = _tool_choice(body["tool_choice"]) + if body.get("max_output_tokens") is not None: + chat["max_tokens"] = body["max_output_tokens"] + if reasoning.get("effort") is not None: + chat["reasoning_effort"] = reasoning["effort"] + return chat, {"wants_summary": summary is not None, "encrypted": bool(include)} + + +def _id(prefix: str) -> str: + return prefix + uuid.uuid4().hex + + +def _usage(run: Run) -> dict: + u = run.usage() + return { + "input_tokens": u["prompt_tokens"], + "input_tokens_details": {"cached_tokens": u["prompt_tokens_details"]["cached_tokens"]}, + "output_tokens": u["completion_tokens"], + "output_tokens_details": {"reasoning_tokens": run.reasoning_tokens(0) or 0}, + "total_tokens": u["total_tokens"], + } + + +def _reasoning_item(item_id: str, text: str, info: dict, status: str = "completed") -> dict: + item: dict = { + "id": item_id, + "type": "reasoning", + "summary": [{"type": "summary_text", "text": text}] if info["wants_summary"] and text else [], + "content": [{"type": "reasoning_text", "text": text}] if text else [], + "status": status, + } + if info["encrypted"]: + item["encrypted_content"] = text + return item + + +def _message_item(item_id: str, text: str, status: str = "completed") -> dict: + content = ( + [{"type": "output_text", "text": text, "annotations": [], "logprobs": []}] if status == "completed" else [] + ) + return {"id": item_id, "type": "message", "role": "assistant", "status": status, "content": content} + + +def _call_item(item_id: str, call_id: str, name: str, arguments: str, status: str = "completed") -> dict: + return { + "id": item_id, + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": arguments, + "status": status, + } + + +def output_items(msg: dict, info: dict) -> list[dict]: + """A parsed chat message (parsers.parse_full) as Responses output items: reasoning, the answer, then the calls.""" + out = [] + if msg.get("reasoning_content"): + out.append(_reasoning_item(_id("rs_"), msg["reasoning_content"], info)) + if msg.get("content"): + out.append(_message_item(_id("msg_"), msg["content"])) + for tc in msg.get("tool_calls") or (): + fn = tc["function"] + out.append(_call_item(_id("fc_"), tc["id"], fn["name"], fn["arguments"])) + return out + + +def _response(body: dict, model: str, rid: str, created: int, status: str, output: list, usage) -> dict: + """The Response object, echoing the request's settings as the OpenAI API does.""" + r = body.get("reasoning") + reasoning: dict = r if isinstance(r, dict) else {} + return { + "id": "resp_" + rid, + "object": "response", + "created_at": created, + "status": status, + "background": False, + "error": None, + "incomplete_details": {"reason": "max_output_tokens"} if status == "incomplete" else None, + "instructions": body.get("instructions"), + "max_output_tokens": body.get("max_output_tokens"), + "max_tool_calls": None, + "model": model, + "output": output, + "parallel_tool_calls": body.get("parallel_tool_calls") is not False, + "previous_response_id": None, + "reasoning": {"effort": reasoning.get("effort"), "summary": reasoning.get("summary")}, + "store": False, + "temperature": body.get("temperature"), + "text": {"format": {"type": "text"}}, + "tool_choice": body.get("tool_choice") if body.get("tool_choice") is not None else "auto", + "tools": body.get("tools") or [], + "top_p": body.get("top_p"), + "truncation": "disabled", + "usage": usage, + "user": body.get("user"), + "metadata": body.get("metadata") or {}, + } + + +class _Stream: + """The Responses events of one streamed run, from its parsed chat deltas. One output item is open at a time (GLM + writes reasoning, then the answer, then its calls): opening the next closes it. Whitespace-only content between two + calls opens no message item (it is held, and leads the next content if any comes).""" + + def __init__(self, body: dict, info: dict, model: str, run: Run): + self.body, self.info, self.model, self.run = body, info, model, run + self.seq = 0 + self.output: list[dict] = [] # the items closed so far + self.cur: dict[str, Any] | None = None # the open item: kind, id, index, text, and a call's call_id, name + self.space = "" # whitespace-only content not yet in a message item + + def ev(self, kind: str, **fields) -> str: + data = {"type": kind, "sequence_number": self.seq, **fields} + self.seq += 1 + return f"event: {kind}\ndata: {json.dumps(data)}\n\n" + + def response(self, status: str, usage=None) -> dict: + return _response(self.body, self.model, self.run.rid, self.run.created, status, list(self.output), usage) + + def _open(self, kind: str, **extra) -> list[str]: + out = self.close() + index = len(self.output) + cur: dict[str, Any] = { + "kind": kind, + "id": _id({"reasoning": "rs_", "message": "msg_", "call": "fc_"}[kind]), + "index": index, + } + cur.update(text="", **extra) + self.cur = cur + if kind == "reasoning": + item = _reasoning_item(cur["id"], "", self.info, "in_progress") + elif kind == "message": + item = _message_item(cur["id"], "", "in_progress") + else: + item = _call_item(cur["id"], cur["call_id"], cur["name"], "", "in_progress") + out.append(self.ev("response.output_item.added", output_index=index, item=item)) + if kind == "message": + part = {"type": "output_text", "text": "", "annotations": [], "logprobs": []} + out.append( + self.ev( + "response.content_part.added", item_id=cur["id"], output_index=index, content_index=0, part=part + ) + ) + elif kind == "reasoning" and self.info["wants_summary"]: + part = {"type": "summary_text", "text": ""} + out.append( + self.ev( + "response.reasoning_summary_part.added", + item_id=cur["id"], + output_index=index, + summary_index=0, + part=part, + ) + ) + return out + + def close(self) -> list[str]: + """The .done events of the open item, if any.""" + cur, self.cur = self.cur, None + if cur is None: + return [] + at = {"item_id": cur["id"], "output_index": cur["index"]} + text = cur["text"] + out = [] + if cur["kind"] == "reasoning": + item = _reasoning_item(cur["id"], text, self.info) + if self.info["wants_summary"]: + out.append(self.ev("response.reasoning_summary_text.done", **at, summary_index=0, text=text)) + part = {"type": "summary_text", "text": text} + out.append(self.ev("response.reasoning_summary_part.done", **at, summary_index=0, part=part)) + else: + out.append(self.ev("response.reasoning_text.done", **at, content_index=0, text=text)) + elif cur["kind"] == "message": + item = _message_item(cur["id"], text) + out.append(self.ev("response.output_text.done", **at, content_index=0, text=text, logprobs=[])) + out.append(self.ev("response.content_part.done", **at, content_index=0, part=item["content"][0])) + else: + item = _call_item(cur["id"], cur["call_id"], cur["name"], text) + out.append(self.ev("response.function_call_arguments.done", **at, arguments=text, name=cur["name"])) + out.append(self.ev("response.output_item.done", output_index=cur["index"], item=item)) + self.output.append(item) + return out + + def _current(self) -> dict[str, Any]: + if self.cur is None: + raise RuntimeError("no output item is open") + return self.cur + + def delta(self, d: dict) -> list[str]: + """The events of one parsed chat delta (StreamParser: reasoning_content, content or tool_calls).""" + out: list[str] = [] + if "reasoning_content" in d: + if self.cur is None or self.cur["kind"] != "reasoning": + out += self._open("reasoning") + cur = self._current() + cur["text"] += d["reasoning_content"] + at = {"item_id": cur["id"], "output_index": cur["index"]} + if self.info["wants_summary"]: + out.append( + self.ev( + "response.reasoning_summary_text.delta", **at, summary_index=0, delta=d["reasoning_content"] + ) + ) + else: + out.append( + self.ev("response.reasoning_text.delta", **at, content_index=0, delta=d["reasoning_content"]) + ) + elif "content" in d: + text = d["content"] + if self.cur is None or self.cur["kind"] != "message": + if not text.strip(): + self.space += text + return out + text, self.space = self.space + text, "" + out += self._open("message") + cur = self._current() + cur["text"] += text + at = {"item_id": cur["id"], "output_index": cur["index"]} + out.append(self.ev("response.output_text.delta", **at, content_index=0, delta=text, logprobs=[])) + else: + for call in d.get("tool_calls", ()): + fn = call["function"] + if "id" in call: # a new call: its name is complete + self.space = "" + out += self._open("call", call_id=call["id"], name=fn["name"], call_index=call["index"]) + call_item = self.cur # the call's item, or None (a muted or skipped call) + if call_item is None or call_item.get("call_index") != call["index"] or not fn["arguments"]: + continue + call_item["text"] += fn["arguments"] + out.append( + self.ev( + "response.function_call_arguments.delta", + item_id=call_item["id"], + output_index=call_item["index"], + delta=fn["arguments"], + ) + ) + return out + + +async def respond(http: HTTPRequest, body: dict, start_chat: Callable[[dict], Awaitable[Run]], model: str) -> Response: + """POST /v1/responses: the chat path's run of the translated request, answered as a Response or its events.""" + chat, info = to_chat(body) + stream = body.get("stream") + if stream is not None and not isinstance(stream, bool): + raise _bad("stream must be true or false") + run = await start_chat(chat) + if stream: + return StreamingResponse(_events(run, body, info, model), media_type="text/event-stream") + try: + if not await run.wait(http): + return Response(status_code=499) # the client went away; nobody reads it + except EngineError as e: + raise HTTPException(status_code=e.status, detail=e.detail) from None + msg, finish = run.message(0) + status = "incomplete" if finish == "length" else "completed" + return JSONResponse(_response(body, model, run.rid, run.created, status, output_items(msg, info), _usage(run))) + + +async def _events(run: Run, body: dict, info: dict, model: str): + s = _Stream(body, info, model, run) + try: + yield s.ev("response.created", response=s.response("in_progress")) + yield s.ev("response.in_progress", response=s.response("in_progress")) + finish = None + async for _, deltas, fin, _ in run.parsed(): + for d in deltas: + for e in s.delta(d): + yield e + finish = fin or finish + for e in s.close(): + yield e + status = "incomplete" if finish == "length" else "completed" + yield s.ev(f"response.{status}", response=s.response(status, _usage(run))) + except EngineError as e: + for ev in s.close(): + yield ev + failed = s.response("failed") + failed["error"] = {"code": "server_error" if e.status >= 500 else "invalid_request", "message": e.detail} + yield s.ev("response.failed", response=failed) + finally: + run.close() # the client went away (the generator is closed early): abort the request + + +def not_stored(response_id: str) -> HTTPException: + """GET, DELETE or cancel of a stored response: 404, since nothing is stored.""" + return HTTPException( + status_code=404, detail=f"response {response_id!r} not found: this server stores no responses (store is false)" + ) From e03c388c69e4e0dc3cbc558f2e62388d3fce8959 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 035/135] anthropic: the Anthropic Messages API and count_tokens over the chat path --- trnrt/serve/anthropic.py | 396 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 396 insertions(+) create mode 100644 trnrt/serve/anthropic.py diff --git a/trnrt/serve/anthropic.py b/trnrt/serve/anthropic.py new file mode 100644 index 0000000..5d79dd1 --- /dev/null +++ b/trnrt/serve/anthropic.py @@ -0,0 +1,396 @@ +"""The Anthropic Messages API, `POST /v1/messages` and `POST /v1/messages/count_tokens`, served by the chat path +(serve/api.py). + +A request is translated into a /v1/chat/completions body, and its output, parsed by the same GLM parsers as a chat +response, back into Anthropic content blocks: + +| Messages API | Chat path | +|---|---| +| `system` (a string or text blocks) | a leading system message | +| `text` blocks | message content | +| an assistant's `thinking` blocks | its `reasoning_content` (GLM's template replays it) | +| an assistant's `tool_use` blocks | its `tool_calls` | +| a user's `tool_result` blocks | `tool` messages, in their place among the user's text | +| `tools` (`name`, `description`, `input_schema`) | function tools | +| `tool_choice` `auto` / `none`; `disable_parallel_tool_use` | `tool_choice` `auto` / `none`; `parallel_tool_calls` false | +| `max_tokens` (required), `temperature`, `top_p`, `top_k`, `stop_sequences` | the same sampling fields; `stop` | +| `thinking` `{"type": "disabled"}` / `{"type": "enabled", "budget_tokens": N}` | `enable_thinking` false / `thinking_token_budget` N, which the chat path refuses as it always does | +| `output_config.effort` | `reasoning_effort` | +| output `reasoning_content`, `content`, `tool_calls` | `thinking`, `text` and `tool_use` blocks, in that order | + +GLM-5.3-Flash always reasons (its template has no switch), so its reasoning comes back as a `thinking` block whether or +not the request set `thinking`, with an empty `signature` (nothing here signs it). `tool_choice` `any` and `tool` need +guided decoding and are refused, as the chat path refuses `required` and a named tool. Usage follows Anthropic's split: +`input_tokens` are the prompt tokens not read from the prefix cache, `cache_read_input_tokens` those that were. + +Errors use Anthropic's shape, `{"type": "error", "error": {"type": ..., "message": ...}}`. The translation follows +SGLang's Anthropic adapter (python/sglang/srt/entrypoints/anthropic/serving.py), except that SGLang accepts and ignores a +thinking budget, which this server refuses (AGENTS.md: unsupported fields are refused, never ignored). +""" + +from __future__ import annotations + +import json +import uuid +from typing import TYPE_CHECKING, Any + +from fastapi import HTTPException +from fastapi import Request as HTTPRequest +from fastapi.responses import JSONResponse, Response, StreamingResponse + +from .api import EngineError, Run + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + +# Anthropic's error types by HTTP status (https://docs.anthropic.com/en/api/errors). +ERROR_TYPES = { + 400: "invalid_request_error", + 401: "authentication_error", + 403: "permission_error", + 404: "not_found_error", + 413: "request_too_large", + 429: "rate_limit_error", + 500: "api_error", + 503: "overloaded_error", + 529: "overloaded_error", +} +# Request fields that change the output in a way this server cannot serve: refused unless absent or null. +_UNSUPPORTED = ("mcp_servers", "container", "context_management") + + +def _bad(detail: str) -> HTTPException: + return HTTPException(status_code=400, detail=detail) + + +def error_body(status: int, detail) -> dict: + return { + "type": "error", + "error": {"type": ERROR_TYPES.get(status, "api_error"), "message": str(detail)}, + } + + +def error_response(e: HTTPException) -> JSONResponse: + """An HTTPException in Anthropic's error shape, at its status.""" + return JSONResponse(error_body(e.status_code, e.detail), status_code=e.status_code) + + +def _text_of(content, where: str) -> str: + """A string, or text blocks joined (as Anthropic concatenates a turn's text blocks); 400 for any other block.""" + if isinstance(content, str): + return content + if not isinstance(content, list): + raise _bad(f"{where} must be a string or a list of text blocks") + out = [] + for b in content: + if not (isinstance(b, dict) and b.get("type") == "text" and isinstance(b.get("text"), str)): + kind = b.get("type") if isinstance(b, dict) else type(b).__name__ + raise _bad(f"{where}: block type {kind!r} is not supported here; this model reads text only") + out.append(b["text"]) + return "".join(out) + + +def _system(system) -> str | None: + if system is None: + return None + return _text_of(system, "system") + + +def messages(body: dict) -> list[dict]: + """The chat messages of a request's system and messages.""" + out: list[dict] = [] + sys = _system(body.get("system")) + if sys: + out.append({"role": "system", "content": sys}) + msgs = body.get("messages") + if not isinstance(msgs, list) or not msgs: + raise _bad("messages must be a non-empty list") + for k, m in enumerate(msgs): + where = f"messages[{k}]" + if not isinstance(m, dict) or m.get("role") not in ("user", "assistant"): + raise _bad(f"{where}: role must be user or assistant") + role, content = m["role"], m.get("content") + if isinstance(content, str): + out.append({"role": role, "content": content}) + continue + if not isinstance(content, list): + raise _bad(f"{where}: content must be a string or a list of blocks") + text: list[str] = [] + thinking: list[str] = [] + calls: list[dict] = [] + for j, b in enumerate(content): + at = f"{where}.content[{j}]" + kind = b.get("type") if isinstance(b, dict) else None + if kind == "text" and isinstance(b.get("text"), str): + text.append(b["text"]) + elif kind == "thinking" and role == "assistant": + if isinstance(b.get("thinking"), str): + thinking.append(b["thinking"]) + elif kind == "tool_use" and role == "assistant": + name, args = b.get("name"), b.get("input", {}) + if not isinstance(name, str) or not name or not isinstance(args, dict): + raise _bad(f"{at}: tool_use needs a name and an input object") + calls.append( + { + "id": b.get("id") or "call_" + uuid.uuid4().hex[:24], + "type": "function", + "function": {"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, + } + ) + elif kind == "tool_result" and role == "user": + if text: # the user's text before the result stays before it + out.append({"role": "user", "content": "".join(text)}) + text = [] + result = b.get("content", "") + out.append( + { + "role": "tool", + "tool_call_id": b.get("tool_use_id"), + "content": result if isinstance(result, str) else _text_of(result, at), + } + ) + else: + raise _bad(f"{at}: block type {kind!r} is not supported in a {role} message") + if role == "assistant": + msg: dict = {"role": "assistant", "content": "".join(text)} + if thinking: + msg["reasoning_content"] = "\n".join(thinking) + if calls: + msg["tool_calls"] = calls + out.append(msg) + elif text: + out.append({"role": "user", "content": "".join(text)}) + return out + + +def _tools(tools) -> list[dict] | None: + if tools is None: + return None + if not isinstance(tools, list): + raise _bad("tools must be a list") + out = [] + for t in tools: + ok = isinstance(t, dict) and isinstance(t.get("name"), str) and isinstance(t.get("input_schema"), dict) + if not ok or t.get("type", "custom") != "custom": + kind = t.get("type") if isinstance(t, dict) else None + raise _bad(f"tool {kind!r} is not supported: only custom tools with a name and an input_schema") + fn = {"name": t["name"], "parameters": t["input_schema"]} + if t.get("description") is not None: + fn["description"] = t["description"] + out.append({"type": "function", "function": fn}) + return out + + +def to_chat(body: dict) -> dict: + """The /v1/chat/completions body for a Messages request (count_tokens takes the same, without max_tokens). 400 for + what cannot be served.""" + for key in _UNSUPPORTED: + if body.get(key) not in (None, []): + raise _bad(f"{key} is not supported") + chat: dict[str, Any] = {"messages": messages(body)} + tools = _tools(body.get("tools")) + if tools: + chat["tools"] = tools + tc = body.get("tool_choice") + if tc is not None: + kind = tc.get("type") if isinstance(tc, dict) else None + if kind == "auto" or kind == "none": + chat["tool_choice"] = kind + elif kind == "any": + chat["tool_choice"] = "required" # refused by the chat path: it needs guided decoding + elif kind == "tool": + chat["tool_choice"] = {"type": "function", "function": {"name": tc.get("name")}} # refused, as above + else: + raise _bad('tool_choice type must be "auto", "any", "tool" or "none"') + if tc.get("disable_parallel_tool_use") is True: + chat["parallel_tool_calls"] = False + thinking = body.get("thinking") + if thinking is not None: + kind = thinking.get("type") if isinstance(thinking, dict) else None + if kind == "disabled": + chat["chat_template_kwargs"] = {"enable_thinking": False} + elif kind == "enabled": + budget = thinking.get("budget_tokens") + if type(budget) is not int: + raise _bad("thinking.budget_tokens is required with thinking enabled") + chat["thinking_token_budget"] = budget # the engine refuses any budget, with its reason + elif kind != "adaptive": + raise _bad('thinking type must be "enabled", "disabled" or "adaptive"') + oc = body.get("output_config") + if oc is not None: + if not isinstance(oc, dict) or set(oc) - {"effort"}: + raise _bad("output_config supports only effort") + if oc.get("effort") is not None: + chat["reasoning_effort"] = oc["effort"] + for src, dst in ( + ("temperature", "temperature"), + ("top_p", "top_p"), + ("top_k", "top_k"), + ("stop_sequences", "stop"), + ): + if body.get(src) is not None: + chat[dst] = body[src] + if body.get("session_id") is not None: # the router's and the engine's session affinity + chat["session_id"] = body["session_id"] + return chat + + +def _usage(run: Run) -> dict: + u = run.usage() + cached = u["prompt_tokens_details"]["cached_tokens"] + return { + "input_tokens": u["prompt_tokens"] - cached, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": cached, + "output_tokens": u["completion_tokens"], + } + + +def _stop(run: Run, finish: str | None, stops: list[str]) -> tuple[str | None, str | None]: + """Anthropic's stop_reason and stop_sequence for a finished chat choice.""" + if finish == "length": + return "max_tokens", None + if finish == "tool_calls": + return "tool_use", None + matched = run.matched_stop(0) + if isinstance(matched, str) and matched in stops: + return "stop_sequence", matched + return "end_turn", None + + +def content_blocks(msg: dict) -> list[dict]: + """A parsed chat message (parsers.parse_full) as Anthropic content blocks: thinking, text, then tool_use.""" + out: list[dict] = [] + if msg.get("reasoning_content"): + out.append({"type": "thinking", "thinking": msg["reasoning_content"], "signature": ""}) + if msg.get("content"): + out.append({"type": "text", "text": msg["content"]}) + for tc in msg.get("tool_calls") or (): + fn = tc["function"] + try: + args = json.loads(fn["arguments"]) if fn["arguments"] else {} + except ValueError: # a call cut off by max_tokens: its arguments as written, under one key + args = {"_raw": fn["arguments"]} + out.append({"type": "tool_use", "id": tc["id"], "name": fn["name"], "input": args}) + return out + + +def _check(body: dict) -> list[str]: + """The checks of a Messages request beyond the translation: max_tokens is required; stream is a boolean. Returns + stop_sequences.""" + mt = body.get("max_tokens") + if type(mt) is not int or mt < 1: + raise _bad("max_tokens is required: an integer of at least 1") + stream = body.get("stream") + if stream is not None and not isinstance(stream, bool): + raise _bad("stream must be true or false") + stops = body.get("stop_sequences") or [] + return stops if isinstance(stops, list) else [] + + +async def respond(http: HTTPRequest, body: dict, start_chat: Callable[[dict], Awaitable[Run]], model: str) -> Response: + """POST /v1/messages: the chat path's run of the translated request, answered as a Message or its events.""" + stops = _check(body) + chat = to_chat(body) + chat["max_tokens"] = body["max_tokens"] + run = await start_chat(chat) + if body.get("stream"): + return StreamingResponse(_events(run, stops, model), media_type="text/event-stream") + try: + if not await run.wait(http): + return Response(status_code=499) # the client went away; nobody reads it + except EngineError as e: + raise HTTPException(status_code=e.status, detail=e.detail) from None + msg, finish = run.message(0) + reason, seq = _stop(run, finish, stops) + return JSONResponse( + { + "id": "msg_" + run.rid, + "type": "message", + "role": "assistant", + "model": model, + "content": content_blocks(msg), + "stop_reason": reason, + "stop_sequence": seq, + "usage": _usage(run), + } + ) + + +def _ev(kind: str, **fields) -> str: + return f"event: {kind}\ndata: {json.dumps({'type': kind, **fields})}\n\n" + + +async def _events(run: Run, stops: list[str], model: str): + """Anthropic's stream: message_start, then per block content_block_start, its deltas and content_block_stop, then + message_delta (stop_reason and the final usage) and message_stop. message_start waits for the first delivery: by + then the engine has scheduled the request, so its prefix-cache hit (cache_read_input_tokens) is known.""" + index = -1 + open_kind: str | None = None # the open block's kind: thinking, text or tool_use + call_index: int | None = None # the open tool_use block's call (the parser's index) + started = False + finish = None + space = "" # whitespace-only text not yet in a text block + try: + async for _, deltas, fin, _ in run.parsed(): + if not started: + started = True + start: dict[str, Any] = { + "id": "msg_" + run.rid, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": dict(_usage(run), output_tokens=0), + } + yield _ev("message_start", message=start) + for d in deltas: + if "tool_calls" in d: + for call in d["tool_calls"]: + fn = call["function"] + if "id" in call: + space = "" + if open_kind is not None: + yield _ev("content_block_stop", index=index) + index, open_kind, call_index = index + 1, "tool_use", call["index"] + use = {"type": "tool_use", "id": call["id"], "name": fn["name"], "input": {}} + yield _ev("content_block_start", index=index, content_block=use) + if open_kind == "tool_use" and call_index == call["index"] and fn["arguments"]: + part = {"type": "input_json_delta", "partial_json": fn["arguments"]} + yield _ev("content_block_delta", index=index, delta=part) + continue + kind, key = ("thinking", "reasoning_content") if "reasoning_content" in d else ("text", "content") + piece = d[key] + if open_kind != kind: + if kind == "text": + if not piece.strip(): # whitespace between two calls opens no text block: held + space += piece + continue + piece, space = space + piece, "" + if open_kind is not None: + yield _ev("content_block_stop", index=index) + index, open_kind = index + 1, kind + block = ( + {"type": "thinking", "thinking": "", "signature": ""} + if kind == "thinking" + else {"type": "text", "text": ""} + ) + yield _ev("content_block_start", index=index, content_block=block) + delta = ( + {"type": "thinking_delta", "thinking": piece} + if kind == "thinking" + else {"type": "text_delta", "text": piece} + ) + yield _ev("content_block_delta", index=index, delta=delta) + finish = fin or finish + if open_kind is not None: + yield _ev("content_block_stop", index=index) + reason, seq = _stop(run, finish, stops) + yield _ev("message_delta", delta={"stop_reason": reason, "stop_sequence": seq}, usage=_usage(run)) + yield _ev("message_stop") + except EngineError as e: + yield _ev("error", error=error_body(e.status, e.detail)["error"]) + finally: + run.close() # the client went away (the generator is closed early): abort the request From 2738810f4d31aa48fe7cca19d79c3ac1bff563a0 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 036/135] api: serve /v1/responses and /v1/messages --- trnrt/serve/api.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/trnrt/serve/api.py b/trnrt/serve/api.py index 0b2a641..c62b724 100644 --- a/trnrt/serve/api.py +++ b/trnrt/serve/api.py @@ -4,6 +4,8 @@ output's reasoning (`reasoning_content`) and tool calls (`tool_calls`) parsed out of it (serve/parsers.py), as one response or as server-sent events; - POST /v1/completions: a prompt as text or token ids, through the same generation path; +- POST /v1/responses (serve/responses.py) and POST /v1/messages, /v1/messages/count_tokens (serve/anthropic.py): the + OpenAI Responses and Anthropic Messages APIs, translated into the chat path; - POST /tokenize and /detokenize (also under /v1); - GET /v1/models; - GET /health: 200 while the engine can serve, 503 with the reason once it cannot; @@ -1385,4 +1387,40 @@ async def detokenize_route(http: HTTPRequest): body = await _body(http) return await asyncio.get_running_loop().run_in_executor(prompt_pool, detokenize, body) + # The Responses and Messages APIs, translated into the chat path (start_chat). Imported here: both import this module. + from . import anthropic, responses + + @app.post("/v1/responses") + async def responses_route(http: HTTPRequest): + body = await _body(http) + return await responses.respond(http, body, start_chat, model_name) + + @app.get("/v1/responses/{response_id}") + @app.delete("/v1/responses/{response_id}") + @app.post("/v1/responses/{response_id}/cancel") + async def stored_response(response_id: str): + raise responses.not_stored(response_id) + + @app.post("/v1/messages") + async def messages_route(http: HTTPRequest): + try: + body = await _body(http) + return await anthropic.respond(http, body, start_chat, model_name) + except HTTPException as e: # Anthropic's error shape + return anthropic.error_response(e) + + def count_tokens(body: dict) -> dict: + """Anthropic's count_tokens: the prompt the same request to /v1/messages would send the engine, rendered on the + prompt thread as a chat prompt is (/tokenize counts the same way).""" + ids, _ = chat_prompt(prompt_tok, anthropic.to_chat(body)) + return {"input_tokens": len(ids)} + + @app.post("/v1/messages/count_tokens") + async def count_tokens_route(http: HTTPRequest): + try: + body = await _body(http) + return await asyncio.get_running_loop().run_in_executor(prompt_pool, count_tokens, body) + except HTTPException as e: + return anthropic.error_response(e) + return app From be623638272208328b841fdb19cae779a9390672 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 037/135] tests: /v1/responses gives the chat API's output as Responses items and events --- tests/test_serve_responses.py | 280 ++++++++++++++++++++++++++++++++++ 1 file changed, 280 insertions(+) create mode 100644 tests/test_serve_responses.py diff --git a/tests/test_serve_responses.py b/tests/test_serve_responses.py new file mode 100644 index 0000000..2e303d6 --- /dev/null +++ b/tests/test_serve_responses.py @@ -0,0 +1,280 @@ +"""The OpenAI Responses API, POST /v1/responses (trnrt/serve/responses.py), over the scripted engine of +tests/test_serve_api.py: the same reasoning, answer and tool calls /v1/chat/completions gives, as Responses items and +events; what a stateless server cannot do is refused with 400.""" + +from __future__ import annotations + +import json + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from tests.test_serve_api import PLAIN, TOOL_TURN, TOOLS, ScriptedEngine + +# The chat TOOLS as Responses function tools. +R_TOOLS = [{"type": "function", **t["function"]} for t in TOOLS] + + +@pytest.fixture +def client(): + from fastapi.testclient import TestClient + + from trnrt.serve import api + + started = [] + + def start(script, setup=None): + eng = ScriptedEngine(script) + if setup is not None: + setup(eng) + c = TestClient(api.build_app(eng, "glm")) + c.__enter__() + started.append((c, eng)) + return c, eng + + yield start + for c, eng in started: + c.__exit__(None, None, None) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def _sse(c, body) -> list[dict]: + with c.stream("POST", "/v1/responses", json={**body, "stream": True}) as r: + assert r.status_code == 200, r.read() + lines = list(r.iter_lines()) + names = [ln[7:] for ln in lines if ln.startswith("event: ")] + data = [json.loads(ln[6:]) for ln in lines if ln.startswith("data: ")] + assert names == [d["type"] for d in data] + assert [d["sequence_number"] for d in data] == list(range(len(data))) + return data + + +def _strip_ids(items: list[dict]) -> list[dict]: + return [{k: v for k, v in it.items() if k not in ("id", "call_id")} for it in items] + + +def test_a_tool_turn_matches_the_chat_api(client): + """Reasoning, the answer and the call come back as reasoning, message and function_call items carrying what + /v1/chat/completions returns for the same request; usage counts the reasoning tokens (those before ).""" + c, eng = client(TOOL_TURN) + r = c.post("/v1/responses", json={"input": "weather?", "tools": R_TOOLS, "instructions": "Be brief."}) + assert r.status_code == 200, r.text + resp = r.json() + chat = c.post( + "/v1/chat/completions", + json={ + "messages": [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "weather?"}], + "tools": TOOLS, + }, + ) + msg = chat.json()["choices"][0]["message"] + assert eng.added[0].prompt_ids == eng.added[1].prompt_ids # the same prompt + reasoning, message, call = resp["output"] + assert reasoning["type"] == "reasoning" and reasoning["content"] == [ + {"type": "reasoning_text", "text": msg["reasoning_content"]} + ] + assert reasoning["summary"] == [] and "encrypted_content" not in reasoning + assert message["type"] == "message" and message["role"] == "assistant" and message["status"] == "completed" + assert message["content"] == [{"type": "output_text", "text": msg["content"], "annotations": [], "logprobs": []}] + (tc,) = msg["tool_calls"] + assert call["type"] == "function_call" and call["name"] == tc["function"]["name"] == "get_weather" + assert json.loads(call["arguments"]) == json.loads(tc["function"]["arguments"]) == {"city": "Paris"} + assert call["call_id"].startswith("call_") and call["id"].startswith("fc_") + assert resp["status"] == "completed" and resp["object"] == "response" and resp["id"].startswith("resp_") + assert resp["store"] is False and resp["instructions"] == "Be brief." and resp["incomplete_details"] is None + n_prompt = len(eng.added[0].prompt_ids) + assert resp["usage"] == { + "input_tokens": n_prompt, + "input_tokens_details": {"cached_tokens": 2}, + "output_tokens": len(TOOL_TURN), + "output_tokens_details": {"reasoning_tokens": 1}, + "total_tokens": n_prompt + len(TOOL_TURN), + } + + +def test_the_stream_sends_the_responses_events_and_the_same_items(client): + c, _ = client(TOOL_TURN) + body = {"input": "weather?", "tools": R_TOOLS} + whole = c.post("/v1/responses", json=body).json() + ev = _sse(c, body) + kinds = [e["type"] for e in ev] + assert kinds[:2] == ["response.created", "response.in_progress"] and kinds[-1] == "response.completed" + assert ev[0]["response"]["status"] == "in_progress" and ev[0]["response"]["output"] == [] + expected = [ + "response.output_item.added", + "response.reasoning_text.delta", + "response.reasoning_text.done", + "response.output_item.done", + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.output_item.added", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + "response.output_item.done", + ] + assert [k for k in kinds[2:-1] if not k.endswith(".delta")] == [k for k in expected if not k.endswith(".delta")] + final = ev[-1]["response"] + assert _strip_ids(final["output"]) == _strip_ids(whole["output"]) # each parse makes its own call id + args = "".join(e["delta"] for e in ev if e["type"] == "response.function_call_arguments.delta") + assert json.loads(args) == {"city": "Paris"} + added = [e["item"] for e in ev if e["type"] == "response.output_item.added"] + assert [a["status"] for a in added] == ["in_progress"] * 3 + assert final["usage"] == whole["usage"] + + +def test_length_is_incomplete(client): + c, _ = client(PLAIN) + r = c.post("/v1/responses", json={"input": "hi", "max_output_tokens": 2}).json() + assert r["status"] == "incomplete" and r["incomplete_details"] == {"reason": "max_output_tokens"} + assert r["usage"]["output_tokens"] == 2 + ev = _sse(c, {"input": "hi", "max_output_tokens": 2}) + assert ev[-1]["type"] == "response.incomplete" and ev[-1]["response"]["status"] == "incomplete" + + +def test_a_reasoning_summary_and_encrypted_content_mirror_the_reasoning(client): + """reasoning.summary mirrors the reasoning into a summary part (as SGLang does: there is no summarizer), streamed + as reasoning_summary events; include reasoning.encrypted_content returns the reasoning text there, for a stateless + client to send back.""" + c, _ = client(TOOL_TURN) + body = { + "input": "weather?", + "tools": R_TOOLS, + "reasoning": {"summary": "auto"}, + "include": ["reasoning.encrypted_content"], + } + item = c.post("/v1/responses", json=body).json()["output"][0] + assert item["summary"] == [{"type": "summary_text", "text": "The user wants weather."}] + assert item["encrypted_content"] == "The user wants weather." + kinds = [e["type"] for e in _sse(c, body)] + assert "response.reasoning_summary_text.delta" in kinds and "response.reasoning_text.delta" not in kinds + assert "response.reasoning_summary_part.added" in kinds and "response.reasoning_summary_part.done" in kinds + + +@pytest.mark.parametrize( + "extra", + [ + {"store": True}, + {"previous_response_id": "resp_1"}, + {"background": True}, + {"truncation": "auto"}, + {"include": ["file_search_call.results"]}, + {"text": {"format": {"type": "json_schema", "name": "x", "schema": {}}}}, + {"tools": [{"type": "web_search"}]}, + {"tool_choice": "required"}, + {"tool_choice": {"type": "function", "name": "get_weather"}}, + {"top_logprobs": 2}, + {"max_tool_calls": 1}, + {"n": 2}, + {"max_tokens": 5}, + {"messages": []}, + {"reasoning": {"effort": "medium"}}, + {"reasoning": {"summary": "long"}}, + {"input": [{"role": "user", "content": [{"type": "input_image", "image_url": "x"}]}]}, + {"input": [{"type": "item_reference", "id": "x"}]}, + {"input": []}, + {"input": 3}, + {"stream": "yes"}, + ], + ids=repr, +) +def test_what_cannot_be_served_is_refused(client, extra): + c, eng = client(PLAIN) + r = c.post("/v1/responses", json={"input": "hi", "tools": R_TOOLS, **extra}) + assert r.status_code == 400, r.text + assert eng.added == [] + + +def test_neutral_fields_are_served(client): + """What an OpenAI SDK may send by default, or a stateless agent harness always sends.""" + c, _ = client(PLAIN) + body = { + "input": "hi", + "store": False, + "truncation": "disabled", + "text": {"format": {"type": "text"}}, + "include": [], + "tool_choice": "auto", + "parallel_tool_calls": True, + "metadata": {"k": "v"}, + "user": "u-1", + "service_tier": "auto", + "top_logprobs": 0, + "prompt_cache_key": "conversation-1", + } + r = c.post("/v1/responses", json=body) + assert r.status_code == 200, r.text + assert r.json()["metadata"] == {"k": "v"} + + +def test_stored_responses_are_not_found(client): + c, _ = client(PLAIN) + for method, path in (("GET", "/v1/responses/resp_1"), ("DELETE", "/v1/responses/resp_1")): + assert c.request(method, path).status_code == 404 + assert c.post("/v1/responses/resp_1/cancel").status_code == 404 + + +def test_input_items_become_the_chat_messages(): + """An assistant turn's items (reasoning, its message, its calls) join into one assistant message; a call's output + becomes a tool message; text parts join.""" + from trnrt.serve.responses import input_messages + + body = { + "instructions": "sys", + "input": [ + {"role": "developer", "content": "dev"}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "a"}, {"text": "b"}]}, + {"type": "reasoning", "summary": [], "content": [{"type": "reasoning_text", "text": "think"}]}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}, + {"type": "function_call", "call_id": "call_1", "name": "f", "arguments": '{"x": 1}'}, + {"type": "function_call", "call_id": "call_2", "name": "g", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "r1"}, + {"type": "function_call_output", "call_id": "call_2", "output": [{"type": "input_text", "text": "r2"}]}, + {"type": "reasoning", "encrypted_content": "again"}, + {"role": "assistant", "content": "done"}, + ], + } + with pytest.raises(Exception, match="not supported"): # a part without a type is not text + input_messages(body) + body["input"][1]["content"][1]["type"] = "input_text" + assert input_messages(body) == [ + {"role": "system", "content": "sys"}, + {"role": "developer", "content": "dev"}, + {"role": "user", "content": "ab"}, + { + "role": "assistant", + "content": "ok", + "reasoning_content": "think", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "f", "arguments": '{"x": 1}'}}, + {"id": "call_2", "type": "function", "function": {"name": "g", "arguments": "{}"}}, + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "r1"}, + {"role": "tool", "tool_call_id": "call_2", "content": "r2"}, + {"role": "assistant", "content": "done", "reasoning_content": "again"}, + ] + + +def _two_calls(eng) -> None: + """The turn calls the tool twice, a newline between the calls (an added token of the stand-in tokenizer).""" + eng.tokenizer.add_tokens(["\n"]) + nl = eng.tokenizer.convert_tokens_to_ids("\n") + eng.script = [*TOOL_TURN, nl, *TOOL_TURN[3:]] + + +def test_whitespace_between_two_calls_opens_no_message_item(client): + """The newline between two calls is content to the parser; streamed, it opens no message item, so the streamed + items are the non-streamed ones (where parse_full strips it).""" + c, _ = client(TOOL_TURN, _two_calls) + body = {"input": "weather?", "tools": R_TOOLS} + whole = c.post("/v1/responses", json=body).json()["output"] + assert [it["type"] for it in whole] == ["reasoning", "message", "function_call", "function_call"] + final = _sse(c, body)[-1]["response"]["output"] + assert _strip_ids(final) == _strip_ids(whole) From 343d3822963dc5e3badc94bed9358e83e5d8c5e0 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 038/135] tests: /v1/messages gives the chat API's output as Anthropic blocks and events --- tests/test_serve_anthropic.py | 273 ++++++++++++++++++++++++++++++++++ 1 file changed, 273 insertions(+) create mode 100644 tests/test_serve_anthropic.py diff --git a/tests/test_serve_anthropic.py b/tests/test_serve_anthropic.py new file mode 100644 index 0000000..23bd61c --- /dev/null +++ b/tests/test_serve_anthropic.py @@ -0,0 +1,273 @@ +"""The Anthropic Messages API, POST /v1/messages and /v1/messages/count_tokens (trnrt/serve/anthropic.py), over the +scripted engine of tests/test_serve_api.py: the same reasoning, answer and tool calls /v1/chat/completions gives, as +thinking, text and tool_use blocks and Anthropic's events; errors in Anthropic's shape.""" + +from __future__ import annotations + +import json + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from tests.test_serve_api import PLAIN, TOOL_TURN, TOOLS, ScriptedEngine + +A_TOOLS = [ + {"name": t["function"]["name"], "description": "Weather.", "input_schema": t["function"]["parameters"]} + for t in TOOLS +] +ASK = {"messages": [{"role": "user", "content": "weather?"}], "max_tokens": 64} + + +@pytest.fixture +def client(): + from fastapi.testclient import TestClient + + from trnrt.serve import api + + started = [] + + def start(script, setup=None): + eng = ScriptedEngine(script) + if setup is not None: + setup(eng) + c = TestClient(api.build_app(eng, "glm")) + c.__enter__() + started.append((c, eng)) + return c, eng + + yield start + for c, eng in started: + c.__exit__(None, None, None) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def _sse(c, body) -> list[dict]: + with c.stream("POST", "/v1/messages", json={**body, "stream": True}) as r: + assert r.status_code == 200, r.read() + lines = list(r.iter_lines()) + names = [ln[7:] for ln in lines if ln.startswith("event: ")] + data = [json.loads(ln[6:]) for ln in lines if ln.startswith("data: ")] + assert names == [d["type"] for d in data] + return data + + +def _blocks(events: list[dict]) -> list[dict]: + """The content blocks a client rebuilds from the events, as the Anthropic SDK accumulates them.""" + blocks: list[dict] = [] + for e in events: + if e["type"] == "content_block_start": + assert e["index"] == len(blocks) + blocks.append(dict(e["content_block"])) + elif e["type"] == "content_block_delta": + b, d = blocks[e["index"]], e["delta"] + if d["type"] == "thinking_delta": + b["thinking"] += d["thinking"] + elif d["type"] == "text_delta": + b["text"] += d["text"] + else: + b["_json"] = b.get("_json", "") + d["partial_json"] + for b in blocks: + if b["type"] == "tool_use": + b["input"] = json.loads(b.pop("_json", "{}")) + return blocks + + +def test_a_tool_turn_matches_the_chat_api(client): + """Thinking, text and tool_use blocks carry what /v1/chat/completions returns for the same request; a call ends the + turn with stop_reason tool_use; usage splits the prompt into uncached input and cache reads, as Anthropic does.""" + c, eng = client(TOOL_TURN) + r = c.post("/v1/messages", json={**ASK, "tools": A_TOOLS, "system": [{"type": "text", "text": "Be brief."}]}) + assert r.status_code == 200, r.text + m = r.json() + chat = c.post( + "/v1/chat/completions", + json={ + "messages": [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "weather?"}], + "tools": [dict(t, function=dict(t["function"], description="Weather.")) for t in TOOLS], + "max_tokens": 64, + }, + ) + msg = chat.json()["choices"][0]["message"] + assert eng.added[0].prompt_ids == eng.added[1].prompt_ids + thinking, text, use = m["content"] + assert thinking == {"type": "thinking", "thinking": msg["reasoning_content"], "signature": ""} + assert text == {"type": "text", "text": msg["content"]} + (tc,) = msg["tool_calls"] + assert use["type"] == "tool_use" and use["name"] == "get_weather" and use["id"].startswith("call_") + assert use["input"] == json.loads(tc["function"]["arguments"]) == {"city": "Paris"} + assert m["stop_reason"] == "tool_use" and m["stop_sequence"] is None + assert m["type"] == "message" and m["role"] == "assistant" and m["id"].startswith("msg_") and m["model"] == "glm" + n_prompt = len(eng.added[0].prompt_ids) + assert m["usage"] == { + "input_tokens": n_prompt - 2, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 2, + "output_tokens": len(TOOL_TURN), + } + + +def test_the_stream_rebuilds_the_same_message(client): + c, _ = client(TOOL_TURN) + body = {**ASK, "tools": A_TOOLS} + whole = c.post("/v1/messages", json=body).json() + ev = _sse(c, body) + kinds = [e["type"] for e in ev] + assert kinds[0] == "message_start" and kinds[-2:] == ["message_delta", "message_stop"] + start = ev[0]["message"] + assert start["content"] == [] and start["stop_reason"] is None and start["usage"]["output_tokens"] == 0 + assert start["usage"]["cache_read_input_tokens"] == 2 + got = _blocks(ev) + for b in got: + b.pop("id", None) + want = [{k: v for k, v in b.items() if k != "id"} for b in whole["content"]] + assert got == want + assert kinds.count("content_block_start") == kinds.count("content_block_stop") == 3 + assert ev[-2]["delta"] == {"stop_reason": "tool_use", "stop_sequence": None} + assert ev[-2]["usage"] == whole["usage"] + + +def test_stop_sequences_and_max_tokens(client): + c, _ = client(PLAIN) + # thinking disabled: the stand-in template has a thinking switch (GLM-5.3-Flash's has none, and refuses it) + body = {**ASK, "stop_sequences": [" again"], "thinking": {"type": "disabled"}} + m = c.post("/v1/messages", json=body).json() + assert m["content"] == [{"type": "text", "text": "Hello there"}] + assert m["stop_reason"] == "stop_sequence" and m["stop_sequence"] == " again" + ev = _sse(c, body) + assert _blocks(ev) == m["content"] and ev[-2]["delta"]["stop_sequence"] == " again" + m = c.post("/v1/messages", json={**ASK, "max_tokens": 2, "thinking": {"type": "disabled"}}).json() + assert m["stop_reason"] == "max_tokens" and m["content"] == [{"type": "text", "text": "Hello there"}] + m = c.post("/v1/messages", json={**ASK, "thinking": {"type": "disabled"}}).json() + assert m["stop_reason"] == "end_turn" and m["content"] == [{"type": "text", "text": "Hello there again!"}] + m = c.post("/v1/messages", json={**ASK, "thinking": {"type": "adaptive"}}).json() + assert m["content"] == [{"type": "thinking", "thinking": "Hello there again!", "signature": ""}] + + +@pytest.mark.parametrize( + "body", + [ + {"messages": [{"role": "user", "content": "x"}]}, # no max_tokens + {**ASK, "max_tokens": 0}, + {**ASK, "messages": []}, + {**ASK, "messages": [{"role": "system", "content": "x"}]}, + {**ASK, "messages": [{"role": "user", "content": [{"type": "image", "source": {}}]}]}, + {**ASK, "messages": [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "x"}]}]}, + {**ASK, "tools": A_TOOLS, "tool_choice": {"type": "any"}}, + {**ASK, "tools": A_TOOLS, "tool_choice": {"type": "tool", "name": "get_weather"}}, + {**ASK, "tools": [{"type": "web_search_20250305", "name": "web_search"}]}, + {**ASK, "thinking": {"type": "enabled", "budget_tokens": 2048}}, + {**ASK, "thinking": {"type": "enabled"}}, + {**ASK, "thinking": {"type": "sometimes"}}, + {**ASK, "output_config": {"effort": "medium"}}, + {**ASK, "mcp_servers": [{"type": "url", "url": "x"}]}, + {**ASK, "stream": "yes"}, + {**ASK, "temperature": -1}, + ], + ids=repr, +) +def test_refusals_are_400_in_anthropics_error_shape(client, body): + c, eng = client(PLAIN) + r = c.post("/v1/messages", json=body) + assert r.status_code == 400, r.text + err = r.json() + assert err["type"] == "error" and err["error"]["type"] == "invalid_request_error" and err["error"]["message"] + assert all(q.status.name == "FINISHED" for q in eng.added) # a refusal leaves nothing running + assert eng.scheduler.running == [] + + +def test_a_malformed_body_and_an_engine_failure_answer_in_anthropics_shape(client): + c, eng = client(PLAIN) + r = c.post("/v1/messages", content=b"{oops", headers={"content-type": "application/json"}) + assert r.status_code == 400 and r.json()["error"]["type"] == "invalid_request_error" + + def failing_step(): + raise RuntimeError("device fault") + + eng.step = failing_step + r = c.post("/v1/messages", json=ASK) + assert r.status_code == 500 and r.json() == { + "type": "error", + "error": {"type": "api_error", "message": "engine step failed: device fault"}, + } + + +def test_count_tokens_counts_what_messages_would_send(client): + """count_tokens renders the request exactly as /v1/messages does, so its count is that request's prompt length + (and /tokenize's for the same chat messages). No request reaches the engine.""" + c, eng = client(PLAIN) + body = {"messages": ASK["messages"], "system": "Be brief.", "tools": A_TOOLS} + r = c.post("/v1/messages/count_tokens", json=body) + assert r.status_code == 200, r.text + assert eng.added == [] + assert c.post("/v1/messages", json={**body, "max_tokens": 4}).status_code == 200 + assert r.json() == {"input_tokens": len(eng.added[-1].prompt_ids)} + tok = c.post( + "/tokenize", + json={ + "messages": [{"role": "system", "content": "Be brief."}, *ASK["messages"]], + "tools": [dict(t, function=dict(t["function"], description="Weather.")) for t in TOOLS], + }, + ) + assert tok.json()["count"] == r.json()["input_tokens"] + bad = c.post("/v1/messages/count_tokens", json={"messages": []}) + assert bad.status_code == 400 and bad.json()["type"] == "error" + + +def test_the_messages_become_the_chat_messages(): + """A user's tool results become tool messages in their place among its text; an assistant's thinking becomes its + reasoning_content and its tool_use blocks its tool_calls.""" + from trnrt.serve.anthropic import messages + + body = { + "system": "sys", + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}}]}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "hmm", "signature": "s"}, + {"type": "text", "text": "ok"}, + {"type": "tool_use", "id": "toolu_1", "name": "f", "input": {"x": 1}}, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "before"}, + {"type": "tool_result", "tool_use_id": "toolu_1", "content": [{"type": "text", "text": "r"}]}, + {"type": "text", "text": "after"}, + ], + }, + ], + } + assert messages(body) == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "a"}, + { + "role": "assistant", + "content": "ok", + "reasoning_content": "hmm", + "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": '{"x": 1}'}}], + }, + {"role": "user", "content": "before"}, + {"role": "tool", "tool_call_id": "toolu_1", "content": "r"}, + {"role": "user", "content": "after"}, + ] + + +def test_whitespace_between_two_calls_opens_no_text_block(client): + """The newline between two calls is content to the parser; streamed, it opens no text block, so the streamed + blocks are the non-streamed ones (where parse_full strips it).""" + from tests.test_serve_responses import _two_calls + + c, _ = client(TOOL_TURN, _two_calls) + body = {**ASK, "tools": A_TOOLS} + whole = c.post("/v1/messages", json=body).json()["content"] + assert [b["type"] for b in whole] == ["thinking", "text", "tool_use", "tool_use"] + got = _blocks(_sse(c, body)) + assert [{k: v for k, v in b.items() if k != "id"} for b in got] == [ + {k: v for k, v in b.items() if k != "id"} for b in whole + ] From 432984d4e5edb95f26bf5eeb61e78d23b3c69f95 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 039/135] router: route /v1/responses and /v1/messages by their conversation, count_tokens to any engine --- trnrt/serve/router.py | 43 ++++++++++++++++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/trnrt/serve/router.py b/trnrt/serve/router.py index 03ab1ec..e8483b9 100644 --- a/trnrt/serve/router.py +++ b/trnrt/serve/router.py @@ -50,6 +50,10 @@ evicting least-recently-used leaves (128M characters: above 32 sessions of a 1M-token context's ~4 MB of text, so long sessions' prefixes are not evicted under load). +The Responses API's and the Messages API's requests (/v1/responses, /v1/messages) carry a conversation as a chat's do: +its instructions or system part leads it as its system messages, so they are routed by session and prefix the same way +(_conversation). /v1/messages/count_tokens, like /tokenize, goes to any live worker and touches no routing state. + The router's own endpoints: - GET /health: 200 {"status": "ok"} while every worker is up; 200 {"status": "degraded"} while some are (requests and sessions go to the live ones, so a load balancer in front keeps sending); 503 {"status": "down"} once none is. Each @@ -86,10 +90,12 @@ # without it. LOG_SESSIONS_ENV = "TRNRT_ROUTER_LOG_SESSIONS" SESSION_KEY_CHARS = 4096 -ROUTES = ("/v1/completions", "/v1/chat/completions") +# The generation routes: each carries a conversation (a chat's messages, a Responses request's instructions and input, a +# Messages request's system and messages; _conversation) or a prompt, routed by its session or its prefix. +ROUTES = ("/v1/completions", "/v1/chat/completions", "/v1/responses", "/v1/messages") # Every worker answers these the same (one model, one tokenizer), so they go to the least loaded live worker and touch # no prefix tree and no session: a /tokenize of an agent's context must not steer where its next turn is routed. -STATELESS_ROUTES = ("/tokenize", "/v1/tokenize", "/detokenize", "/v1/detokenize") +STATELESS_ROUTES = ("/tokenize", "/v1/tokenize", "/detokenize", "/v1/detokenize", "/v1/messages/count_tokens") class _Node: @@ -463,12 +469,29 @@ def _leading_system(msgs: list) -> int: return k +def _conversation(path: str, body: dict) -> list | None: + """A request's conversation as a list of chat-like messages (the system part first), None for a prompt. The + Responses API's instructions and Anthropic's system lead as a system message, so each API's conversation grows at + its end like a chat's and its system part is the shared head.""" + if path.endswith("chat/completions"): + return body.get("messages", []) + if path.endswith("/responses"): + ins, inp = body.get("instructions"), body.get("input") + head = [{"role": "system", "content": ins}] if ins else [] + return head + (inp if isinstance(inp, list) else [{"role": "user", "content": inp}] if inp else []) + if path.endswith("/messages"): + system, msgs = body.get("system"), body.get("messages") + head = [{"role": "system", "content": system}] if system else [] + return head + (msgs if isinstance(msgs, list) else []) + return None + + def _request_head(path: str, body: dict) -> int: - """The length of _request_text's prefix through a chat request's leading system messages (0 for a prompt): the part + """The length of _request_text's prefix through a conversation's leading system messages (0 for a prompt): the part every session of one client shares (its instructions, environment and file tree).""" - if not path.endswith("chat/completions"): + msgs = _conversation(path, body) + if not msgs or not isinstance(msgs, list): return 0 - msgs = body.get("messages") or [] k = _leading_system(msgs) return len(json.dumps(msgs[:k], ensure_ascii=False)) - 1 if k else 0 # "[A, B" is a prefix of "[A, B, C]" @@ -480,10 +503,11 @@ def _session_key(path: str, body: dict) -> tuple[str, str]: under another system prompt (a title or summary request), first_hash under another conversation start (a compaction). No dump of the whole request (~1 MB at 200K tokens).""" h = lambda b: hashlib.sha1(b.encode("utf-8", "replace"), usedforsecurity=False).hexdigest() # noqa: E731 - if not path.endswith("chat/completions"): + msgs = _conversation(path, body) + if msgs is None: p = body.get("prompt", "") return h(""), h((p if isinstance(p, str) else json.dumps(p))[:SESSION_KEY_CHARS]) - msgs = body.get("messages") or [] + msgs = msgs if isinstance(msgs, list) else [] k = _leading_system(msgs) first = msgs[k] if k < len(msgs) and isinstance(msgs[k], dict) else {} c = first.get("content") @@ -492,8 +516,9 @@ def _session_key(path: str, body: dict) -> tuple[str, str]: def _request_text(path: str, body: dict) -> str: - if path.endswith("chat/completions"): - return json.dumps(body.get("messages", []), ensure_ascii=False) + msgs = _conversation(path, body) + if msgs is not None: + return json.dumps(msgs, ensure_ascii=False) p = body.get("prompt", "") return p if isinstance(p, str) else json.dumps(p) From 6d5fbd45c87a95bb06f52fe417c18e9b4f351832 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 040/135] tests: the router routes Responses and Messages conversations as it routes chat --- tests/test_router.py | 93 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 92 insertions(+), 1 deletion(-) diff --git a/tests/test_router.py b/tests/test_router.py index 9060b71..f671ed1 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -263,7 +263,9 @@ async def main() -> None: asyncio.run(main()) -@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/completions", "/tokenize", "/v1/detokenize"]) +@pytest.mark.parametrize( + "path", ["/v1/chat/completions", "/v1/completions", "/v1/responses", "/v1/messages", "/tokenize", "/v1/detokenize"] +) @pytest.mark.parametrize("body", [b"{not json", b"[1, 2]", b'"text"', b"null", b"true", b"3", b"", b"\xff\xfe"]) def test_a_body_that_is_not_a_json_object_is_answered_400_and_reaches_no_worker(path, body): """The router refuses a malformed body, or JSON that is not an object, with 400, as each engine does, instead of @@ -537,3 +539,92 @@ async def main() -> None: await _stop(servers, tasks) asyncio.run(main()) + + +# One conversation in each API that carries one: a long shared system part, then the turns. +SYSTEM = "You are a coding agent. Environment and file tree follow. " + "src/x.py " * 4500 +TURNS = [ + {"role": "user", "content": "task 0: fix the bug"}, + {"role": "assistant", "content": "reading files " * 3000}, +] + + +def _bodies(turns: list) -> dict[str, dict]: + return { + "/v1/chat/completions": {"messages": [{"role": "system", "content": SYSTEM}, *turns]}, + "/v1/responses": {"instructions": SYSTEM, "input": turns}, + "/v1/messages": {"system": SYSTEM, "messages": turns, "max_tokens": 8}, + } + + +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/responses", "/v1/messages"]) +def test_each_conversation_api_routes_by_its_conversation(path): + """A Responses request's instructions and input, and a Messages request's system and messages, are its + conversation as a chat's messages are: its routing text grows at its end as the conversation does (so its prefix + sticks to the worker holding it), its head is the system part (so a shared system prompt does not pin new sessions), + and its session key is stable while it grows and differs by system prompt.""" + from trnrt.serve.router import _request_head, _request_text, _session_key + + first, later = _bodies(TURNS[:1])[path], _bodies([*TURNS, {"role": "user", "content": "next"}])[path] + t1, t2 = _request_text(path, first), _request_text(path, later) + head = _request_head(path, first) + assert head > 40_000 and SYSTEM[:1000] in t1[:head] and "task 0" not in t1[:head] + assert t2.startswith(t1[:-1]) # "[A, B" is a prefix of "[A, B, C]" + assert _request_head(path, later) == head + assert _session_key(path, first) == _session_key(path, later) + other = dict(first, **({"instructions": "Write a title."} if path == "/v1/responses" else {})) + if path == "/v1/messages": + other["system"] = "Write a title." + elif path == "/v1/chat/completions": + other = {"messages": [{"role": "system", "content": "Write a title."}, *TURNS[:1]]} + assert _session_key(path, other)[0] != _session_key(path, first)[0] + assert _session_key(path, other)[1] == _session_key(path, first)[1] + r = Router(["a", "b"], balance_abs=1, balance_rel=1.0, sticky_chars=32_768) + home = r.pick(_request_text(path, _bodies(TURNS)[path]), head=head) + r.load[home] += 20 # far past balance_abs: only the sticky rule keeps the conversation home + assert r.pick(t2, head=head) == home and r.routed["sticky"] == 1 + + +def test_the_router_proxies_responses_and_messages_with_their_sessions(): + """/v1/responses and /v1/messages go through the proxy: a session_id pins them as it pins a chat.""" + import asyncio + + import httpx + from fastapi import FastAPI, Request + + from trnrt.serve.router import build_router_app + + def worker(i: int) -> FastAPI: + app = FastAPI() + + @app.get("/health") + async def health(): + return {"status": "ok"} + + async def answer(req: Request): + return {"worker": i, "body": await req.json()} + + for path in ("/v1/responses", "/v1/messages"): + app.add_api_route(path, answer, methods=["POST"]) + return app + + async def main() -> None: + ports = [_free_port() for _ in range(3)] + router = Router([f"http://127.0.0.1:{p}" for p in ports[:2]]) + servers, tasks = await _start( + [(worker(0), ports[0]), (worker(1), ports[1]), (build_router_app(router, poll_s=0), ports[2])] + ) + base = f"http://127.0.0.1:{ports[2]}" + try: + async with httpx.AsyncClient(timeout=30) as c: + for path, conversation in _bodies(TURNS[:1]).items(): + if path == "/v1/chat/completions": + continue + body = {**conversation, "session_id": "s-" + path} + got = {(await c.post(base + path, json=body)).json()["worker"] for _ in range(3)} + assert got == {router.sessions["s-" + path]}, (path, got) + assert router.load == [0, 0] and all(t.size == 0 for t in router.trees) # session requests: no tree + finally: + await _stop(servers, tasks) + + asyncio.run(main()) From f5cebf022874bda1ebf12e1186dea84cf1e81ba3 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 041/135] docs: sampling extras, the Responses API and the Anthropic Messages API --- docs/API.md | 84 +++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 79 insertions(+), 5 deletions(-) diff --git a/docs/API.md b/docs/API.md index 7cd2d4b..ee8795f 100644 --- a/docs/API.md +++ b/docs/API.md @@ -1,6 +1,6 @@ # API -The router on `--port` (default 8000) speaks the OpenAI API: `/v1/chat/completions`, `/v1/completions`, `/v1/models` and `/health`, plus `/tokenize` and `/detokenize` ([below](#tokenize-and-detokenize)). Any OpenAI client works with `base_url="http://:8000/v1"`. +The router on `--port` (default 8000) speaks the OpenAI API: `/v1/chat/completions`, `/v1/completions`, `/v1/models` and `/health`, plus `/tokenize` and `/detokenize` ([below](#tokenize-and-detokenize)). Any OpenAI client works with `base_url="http://:8000/v1"`. It also serves the OpenAI Responses API ([below](#responses-api)) and the Anthropic Messages API ([below](#anthropic-messages-api)). ## Example @@ -30,8 +30,8 @@ curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' - - A response cut off by `max_tokens` has `finish_reason` `"length"`, even when it ends in a tool call: that call may be incomplete. - A client that disconnects aborts its request, streamed or not. -- Endpoints on `--port` (the router): `/v1/chat/completions`, `/v1/completions`, `/v1/models`, `/health`, `/tokenize` - and `/detokenize`. Each +- Endpoints on `--port` (the router): `/v1/chat/completions`, `/v1/completions`, `/v1/responses`, `/v1/messages`, + `/v1/messages/count_tokens`, `/v1/models`, `/health`, `/tokenize` and `/detokenize`. Each engine also serves Prometheus metrics at `/metrics` on its own port (8001 and 8002 with the default `--port`), on 127.0.0.1 only, among them the sparse-attention selection's counters (`trnrt:dsa_local_k_*`, [KNOWN_ISSUES.md](KNOWN_ISSUES.md)). @@ -41,12 +41,86 @@ curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' - in front of it should drop idle connections sooner than that; otherwise a request can go out on a connection the router is closing. - Not supported, answered with 400: repetition, frequency and presence penalties, guided decoding (`response_format` - other than text, `tool_choice` `"required"` or a named tool), prompt logprobs and thinking budgets. Every graph - samples on the device, and these need host inputs no served graph takes. + other than text, `tool_choice` `"required"` or a named tool), prompt logprobs, thinking budgets, and a minimum length + (`min_tokens`, SGLang's `min_new_tokens`). Every graph samples on the device, and these need host inputs no served + graph takes. A minimum length needs EOS masked out of the distribution a token is drawn from; ignoring EOS on the + host instead would put it in the output and let the model write past it. - Status codes: 400 for a request that cannot be served as sent (a field of the wrong type or out of range, a prompt longer than the context, a chat template variable the template never reads); 500 for an error inside the engine; 503 once the engine has failed, with the reason on `/health`. +## Sampling extras + +These fields work on `/v1/chat/completions` and `/v1/completions`, as vLLM and SGLang take them. + +| Field | What it does | +|---|---| +| `n` | Up to 128 choices of one prompt, each its own engine request. The prefix cache computes the shared prompt once: the scheduler holds the other choices until they can resume from the first one's prompt checkpoint. Each choice has its `index`, streamed chunks included. `usage` counts the prompt once and every choice's output; `cached_tokens` is the fewest any choice read from the cache. A client that disconnects aborts every choice. | +| `best_of` | Accepted only equal to `n` (ranking extra candidates needs their logprobs, which no request returns). | +| `seed` with `n` | Choice i samples with `seed + i`, as vLLM does, so the choices differ and each is reproducible. | +| `stop_regex` | SGLang's: a pattern or a list of up to 16, matched on the output text. The output ends at the end of the first match, the matched text included (as in SGLang), with `finish_reason` `"stop"`. | +| `include_stop_str_in_output`, `no_stop_trim` | vLLM's and SGLang's names for one switch: keep the stop string, or the stop token's text, in the output. | +| `skip_special_tokens` | `false` keeps special tokens in the output text. | +| `stop_token_ids` | Tokens that end the request. As in vLLM and SGLang, the stop token (or EOS) is not part of the text. | + +Every finished choice carries SGLang's `matched_stop`: the stop string or `stop_regex` pattern that ended it, the token id +that did (EOS or a stop token), or `null` (a length). + +## Responses API + +`POST /v1/responses` serves the OpenAI Responses API, stateless. It is translated into a chat request, so the same +prompt, sampling and GLM parsers apply, and its output comes back as Responses items. + +| Request field | Served as | +|---|---| +| `input` (a string, or message, `reasoning`, `function_call` and `function_call_output` items), `instructions` | the chat messages; an assistant turn's items join into one assistant message | +| `tools` (function tools), `tool_choice` | chat tools; `"auto"` and `"none"` (`"required"` and a named function are refused, as in chat) | +| `max_output_tokens`, `temperature`, `top_p`, `stream`, `parallel_tool_calls`, `session_id`, the chat sampling extras | the same chat fields | +| `reasoning.effort` | `reasoning_effort`: `"low"`, `"high"` or `"max"`, as in chat | +| `reasoning.summary` | a `summary_text` that mirrors the reasoning (there is no summarizer), as SGLang does | +| `include: ["reasoning.encrypted_content"]` | the reasoning text in the reasoning item's `encrypted_content`, not encrypted, so a client can send it back | + +| Output | Item | +|---|---| +| reasoning | a `reasoning` item with `reasoning_text` content | +| the answer | a `message` item with one `output_text` part | +| each tool call | a `function_call` item; send its result back as `function_call_output` with its `call_id` | + +A response cut off by `max_output_tokens` has status `"incomplete"`. `usage.output_tokens_details.reasoning_tokens` +counts the output tokens before ``. Streamed, the events are `response.created`, `response.in_progress`, each +item's `response.output_item.added`, deltas and `.done` events, then `response.completed` (`response.incomplete`, or +`response.failed` on an engine error). + +Nothing is stored. Refused with 400: `store: true`, `previous_response_id`, `background`, `truncation: "auto"`, +`max_tool_calls`, `top_logprobs`, `text.format` other than text, other `include` values, item references, non-text +input parts and tools other than functions. `GET`, `DELETE` and `cancel` of a response answer 404. + +## Anthropic Messages API + +`POST /v1/messages` serves the Anthropic Messages API, and `POST /v1/messages/count_tokens` counts a request's input +tokens exactly as `/v1/messages` would render it (the same count `/tokenize` gives for the chat messages). Point an +Anthropic client at `base_url="http://:8000"`. + +| Request field | Served as | +|---|---| +| `system` (a string or text blocks), `messages` with `text`, `tool_use`, `tool_result` and `thinking` blocks | the chat messages: a user's tool results become tool messages in their place, an assistant's thinking its `reasoning_content` | +| `tools` (custom tools: `name`, `description`, `input_schema`), `tool_choice` `auto` / `none`, `disable_parallel_tool_use` | chat tools, `tool_choice`, `parallel_tool_calls: false` | +| `max_tokens` (required), `temperature`, `top_p`, `top_k`, `stop_sequences`, `stream`, `session_id` | the same chat fields; `stop` | +| `thinking` `{"type": "adaptive"}` or absent | GLM-5.3-Flash reasons as always | +| `output_config.effort` | `reasoning_effort`: `"low"`, `"high"` or `"max"` | + +The reply's content is a `thinking` block (with an empty `signature`), a `text` block and `tool_use` blocks. +GLM-5.3-Flash always reasons, so the `thinking` block comes back whether or not the request set `thinking`. +`stop_reason` is `end_turn`, `max_tokens`, `stop_sequence` (with `stop_sequence`) or `tool_use`. `usage.input_tokens` +counts the prompt tokens not read from the prefix cache and `cache_read_input_tokens` those that were, as Anthropic +splits them. Streamed, the events are Anthropic's: `message_start`, each block's `content_block_start`, deltas +(`thinking_delta`, `text_delta`, `input_json_delta`) and `content_block_stop`, then `message_delta` and `message_stop`. + +Errors use Anthropic's shape, `{"type": "error", "error": {"type": "invalid_request_error", "message": ...}}`. Refused +with 400: `thinking` `{"type": "disabled"}` (the template cannot turn reasoning off), `{"type": "enabled"}` with +`budget_tokens` (a thinking budget, which no served graph enforces), `tool_choice` `any` and `tool` (guided decoding), +image and document blocks, `redacted_thinking`, server tools, `mcp_servers`, `container` and `context_management`. + ## Tokenize and detokenize `/tokenize` and `/detokenize` (also under `/v1`) answer as SGLang and vLLM do. Neither runs the model, and the router From 2260b1e685f3f5530bd7d6ff09591878f0ad8813 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 042/135] docs: the Responses and Messages modules in the repository map --- docs/ARCHITECTURE.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 9b2f063..97a327f 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -19,6 +19,7 @@ | `trnrt/serve/api.py` | The OpenAI-compatible server for one engine: request validation, chat template, streaming. | | `trnrt/serve/router.py` | The session-affine, prefix-aware router in front of the engines, with failover. | | `trnrt/serve/parsers.py` | GLM reasoning and tool-call parsing, streamed and whole. | +| `trnrt/serve/responses.py`, `anthropic.py` | The OpenAI Responses API and the Anthropic Messages API, translated into the chat path. | | `trnrt/engine/engine.py` | `LLMEngine`: model build, the step loop, start and close. | | `trnrt/engine/scheduler.py` | Continuous batching, chunked prefill and prefix-cache restores. | | `trnrt/engine/model_runner.py` | Pads each step into a compiled bucket and runs its graph. | From 0a33489fddb412e4b070cf212478ac8c3dc65c93 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:00 +0530 Subject: [PATCH 043/135] changelog: sampling extras, /v1/responses and /v1/messages --- CHANGELOG.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0748e1b..a15725b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,23 @@ and versions follow [Semantic Versioning](https://semver.org/). Rebuild the comp - `/tokenize` and `/detokenize` (also under `/v1`), as SGLang and vLLM serve them. A chat request is counted exactly as `/v1/chat/completions` would send it to the engine. The router sends them to any live engine without touching its routing state ([docs/API.md](docs/API.md)). +- `n` choices on `/v1/chat/completions` and `/v1/completions`, sharing the prompt through the prefix cache; a seed gives + choice i `seed + i`, as in vLLM ([docs/API.md](docs/API.md#sampling-extras)). +- SGLang's `stop_regex`, and keeping the stop text with `include_stop_str_in_output` (vLLM) or `no_stop_trim` (SGLang). +- `skip_special_tokens: false`, and SGLang's `matched_stop` on every finished choice. +- The OpenAI Responses API, `POST /v1/responses`, stateless: reasoning, message and function-call items, and the + Responses stream events ([docs/API.md](docs/API.md#responses-api)). +- The Anthropic Messages API, `POST /v1/messages` and `/v1/messages/count_tokens`: thinking, text and tool-use blocks, + Anthropic's stream events and error shape ([docs/API.md](docs/API.md#anthropic-messages-api)). +- The router routes `/v1/responses` and `/v1/messages` by session and prefix as it routes chat, and sends + `/v1/messages/count_tokens` to any live engine. + +### Changed + +- The stop token that ends a request (EOS or a `stop_token_ids` entry) is no longer part of its text when the tokenizer + does not mark it special, as in vLLM and SGLang. +- `min_new_tokens` (SGLang's name for `min_tokens`) and `spaces_between_special_tokens: false` are refused with 400; + they were ignored. ## [0.1.0] - 2026-11-08 From b7492051e3b6f47b4d6d7dfe2040468e7979d8e7 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:09:28 +0530 Subject: [PATCH 044/135] api: name the scheduler paths that share an n-choice prompt --- trnrt/serve/api.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/trnrt/serve/api.py b/trnrt/serve/api.py index c62b724..72193cb 100644 --- a/trnrt/serve/api.py +++ b/trnrt/serve/api.py @@ -935,10 +935,11 @@ def __init__(self, detail: str, status: int): class Run: """The engine requests of one HTTP request: `n` choices of one prompt (OpenAI's `n`), each its own engine request, - sharing the prompt. The prefix cache makes the copies cheap: the scheduler sees the requests queued together, makes - their shared prompt pages a checkpoint target of the first one it admits, and holds the others until they can resume - from it (engine/scheduler.py: _junctions_ahead, _prefix_in_progress), so the prompt is computed once and each other - choice computes only its last partial page. Their tokens arrive tagged with the choice's index on one queue. + sharing the prompt. The prefix cache makes the copies cheap: their shared prompt pages become a checkpoint target of + the first one the scheduler admits (_junctions_ahead when they are queued together; _mark_running_junctions when one + arrives while the first prefills), and the others wait until they can resume from it (_prefix_in_progress, + engine/scheduler.py), so the prompt is computed once and each other choice computes only its last partial page. + Their tokens arrive tagged with the choice's index on one queue. Each choice has its own detokenizer (stop strings, stop_regex, skip_special_tokens) and, for a chat request, its own GLM parser. A choice the host stops (a stop_regex match) is aborted in the engine. Whatever ends the run early (an From 0733cf3d1ec985dd2cc83954c952c02124cd2896 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:35 +0530 Subject: [PATCH 045/135] stats: request rates, latency percentiles and a time series per engine --- trnrt/serve/stats.py | 363 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 363 insertions(+) create mode 100644 trnrt/serve/stats.py diff --git a/trnrt/serve/stats.py b/trnrt/serve/stats.py new file mode 100644 index 0000000..67d0ad7 --- /dev/null +++ b/trnrt/serve/stats.py @@ -0,0 +1,363 @@ +"""Live serving statistics of one engine: GET /v1/stats, the page GET /dashboard draws (serve/dashboard.py), and the +records the production log keeps (serve/prodlog.py). Host only: nothing here touches the engine or its graphs. + +Cost on the engine thread, O(1) each, under an uncontended lock: +- per finished request, record(): one dict built from the Request, appended to a bounded ring (RING); +- per loop iteration (a step, or an idle wait), tick(): the tokens the step handed out and the running and waiting + counts, added into the current BUCKET_S bucket. +Everything else (windows, percentiles, the time series) is computed by snapshot() off the event loop, from a copy taken +under the lock, and cached for SNAPSHOT_TTL_S so that several viewers share one computation. + +Times are wall-clock seconds (time.time()): the engine's monotonic Request times plus one offset taken at start, so +records reloaded from the production log of an earlier process share the timeline. All are measured inside the server, +from when the engine received the request (after the chat template and tokenization) to its last token, without the +client's network round trip. + +Per request: TTFT = first token - arrival; end-to-end = done - arrival; decode tok/s (per stream) = (output - 1) / +(done - first token), and the mean inter-token latency its inverse (requests with at least 2 output tokens). Rates per +window count requests by when they finished: RPM = requests / minutes; TPM input, cached, output, and total (input + +output). The aggregate output tok/s counts tokens as the steps hand them out (BUCKET_S buckets), so it stays current +while long requests are still running. +""" + +from __future__ import annotations + +import collections +import math +import threading +import time +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + +RING = 10_000 # finished requests kept for the windows and percentiles (~2.5 MB of dicts) +BUCKET_S = 10 +SERIES_BUCKETS = 180 # the time series: 30 min of 10-s points +WINDOWS = (("1m", 60), ("5m", 300), ("15m", 900)) +RECENT = 50 # the newest finished requests listed +SNAPSHOT_TTL_S = 1.0 +OK_FINISH = ("stop", "length", "tool_calls") +RECENT_KEYS = ( + "rid", + "arrival", + "first", + "done", + "prompt", + "cached", + "output", + "finish", + "stream", + "ttft", + "e2e", + "decode_tps", + "itl", + "running_at_arrival", + "waiting_at_arrival", + "spec_proposed", + "spec_accepted", +) + + +def pct(xs: list[float], q: float) -> float | None: + """The q-quantile (0 .. 1) of xs, interpolating linearly between order statistics (numpy's default).""" + if not xs: + return None + s = sorted(xs) + pos = q * (len(s) - 1) + lo = math.floor(pos) + hi = min(lo + 1, len(s) - 1) + return s[lo] + (s[hi] - s[lo]) * (pos - lo) + + +def summary(xs: list[float], low: bool = False) -> dict: + """n, mean, p50, p90 and p99 (and p10 first with low=True: a rate's slowest streams).""" + out: dict = {"n": len(xs), "mean": (sum(xs) / len(xs)) if xs else None} + for q in ((0.1,) if low else ()) + (0.5, 0.9, 0.99): + out[f"p{round(q * 100)}"] = pct(xs, q) + return out + + +def derive(rec: dict) -> dict: + """A record's derived times (seconds) and rate, in place: ttft, e2e, decode_tps, itl (None when undefined).""" + first, done, arr, n = rec.get("first"), rec.get("done"), rec.get("arrival"), rec.get("output", 0) + rec["ttft"] = (first - arr) if first is not None and arr is not None else None + rec["e2e"] = (done - arr) if done is not None and arr is not None else None + span = (done - first) if first is not None and done is not None else None + rec["decode_tps"] = (n - 1) / span if span and span > 0 and n > 1 else None + rec["itl"] = span / (n - 1) if span is not None and n > 1 else None + return rec + + +def _kind(finish: str) -> str: + return "ok" if finish in OK_FINISH else "aborted" if finish == "abort" else "failed" + + +class Stats: + def __init__( + self, ring: int = RING, clock: Callable[[], float] = time.time, mono: Callable[[], float] = time.monotonic + ): + self.lock = threading.Lock() + self.clock = clock + self.offset = clock() - mono() # wall = monotonic + offset + self.started = clock() + self.records: collections.deque = collections.deque(maxlen=ring) + self.counts = {"ok": 0, "aborted": 0, "failed": 0} + self.failed_kinds: dict[str, int] = {} + self.tokens = {"prompt": 0, "cached": 0, "output": 0} + # Earlier processes' records of today (the production log's reload): totals and the first record's time. + self.day: dict | None = None + # BUCKET_S buckets in a ring of SERIES_BUCKETS slots: [bucket id, ticks, running sum, running max, waiting + # sum, waiting max, tokens handed out]. + self.buckets: list[list | None] = [None] * SERIES_BUCKETS + self.cur: list | None = None + self.past: dict[int, dict] = {} # bucket id -> a gauges record reloaded from today's log (before this start) + self.on_bucket: Callable[[dict], object] | None = None # (the closed bucket's gauges record): the production log + self.kv: Callable[[], tuple[int, int]] | None = None # () -> (pages used, pages total), for the gauges records + self._snap: tuple[float, dict] | None = None + self._snap_lock = threading.Lock() # snapshot() runs in worker threads: one computes, the others reuse it + self.handed = 0 # output tokens handed out since this process started + + # -- the engine thread -------------------------------------------------------------------------------------------- + + def record(self, req, stream: bool, at_arrival: tuple[int, int] | None, finish: str | None = None) -> dict: + """A finished request's record: appended to the ring, counted, and returned for the production log.""" + off = self.offset + rec = { + "rid": req.rid, + "arrival": req.arrival_time + off, + "first": None if req.first_token_time is None else req.first_token_time + off, + "done": (req.finish_time if req.finish_time is not None else time.monotonic()) + off, + "prompt": req.num_prompt, + "cached": max(req.num_cached_tokens, 0), + "output": len(req.output_ids), + "finish": finish or req.finish_reason or "unknown", + "stream": bool(stream), + "running_at_arrival": None if at_arrival is None else at_arrival[0], + "waiting_at_arrival": None if at_arrival is None else at_arrival[1], + } + if req.spec_proposed: + rec["spec_proposed"], rec["spec_accepted"] = req.spec_proposed, req.spec_accepted + derive(rec) + kind = _kind(rec["finish"]) + with self.lock: + self.records.append(rec) + self.counts[kind] += 1 + if kind == "failed": + self.failed_kinds[rec["finish"]] = self.failed_kinds.get(rec["finish"], 0) + 1 + for k in self.tokens: + self.tokens[k] += rec[k] + return rec + + def failed(self, kind: str) -> None: + """A request that never ran (a full queue, a refusal) or that an engine failure ended.""" + with self.lock: + self.counts["failed"] += 1 + self.failed_kinds[kind] = self.failed_kinds.get(kind, 0) + 1 + + def tick(self, now_mono: float, tokens: int, running: int, waiting: int) -> None: + """One engine-loop iteration: the tokens its step handed out (0 when idle) and the scheduler's counts.""" + b = int((now_mono + self.offset) // BUCKET_S) + closed = None + with self.lock: + cur = self.cur + if cur is None or cur[0] != b: + closed = cur + cur = self.cur = [b, 0, 0, 0, 0, 0, 0] + self.buckets[b % SERIES_BUCKETS] = cur + cur[1] += 1 + cur[2] += running + cur[3] = max(cur[3], running) + cur[4] += waiting + cur[5] = max(cur[5], waiting) + cur[6] += tokens + self.handed += tokens + if closed is not None and self.on_bucket is not None: + self.on_bucket(self.gauge_record(closed)) + + def gauge_record(self, b: list) -> dict: + """A closed bucket as the production log's gauges record (requests counted by their finish time in it).""" + lo, hi = b[0] * BUCKET_S, (b[0] + 1) * BUCKET_S + done = [] + with self.lock: + for r in reversed(self.records): # newest first: stop at the first one older than the bucket + if r["done"] < lo: + break + if r["done"] < hi and r["finish"] in OK_FINISH: + done.append(r) + rec = { + "t": lo, + "running": b[2] / b[1], + "running_max": b[3], + "waiting": b[4] / b[1], + "waiting_max": b[5], + "out_tok_s": b[6] / BUCKET_S, + "rpm": len(done) * 60 / BUCKET_S, + "input_tpm": sum(r["prompt"] for r in done) * 60 / BUCKET_S, + "output_tpm": sum(r["output"] for r in done) * 60 / BUCKET_S, + } + if self.kv is not None: + used, total = self.kv() + rec["kv_used_pct"] = 100.0 * used / total if total else None + return rec + + # -- the production log's reload, at start ------------------------------------------------------------------------ + + def merge_day(self, records: list[dict], gauges: list[dict], day_start: float) -> None: + """Today's records and gauges from earlier processes: the "today" window, the ring's older end, the series.""" + for r in records: + derive(r) + n_ok = sum(1 for r in records if r.get("finish") in OK_FINISH) + n_abort = sum(1 for r in records if r.get("finish") == "abort") + with self.lock: + maxlen = self.records.maxlen or 0 + older = [r for r in records if r["done"] < self.started][-maxlen:] + room = maxlen - len(self.records) + if room > 0: + self.records.extendleft(reversed(older[-room:])) + self.day = { + "start": day_start, + "first": min((r["arrival"] for r in records if r.get("arrival") is not None), default=None), + "ok": n_ok, + "aborted": n_abort, + "failed": len(records) - n_ok - n_abort, + **{k: sum(r.get(k, 0) for r in records) for k in ("prompt", "cached", "output")}, + } + for g in gauges: + self.past[int(g["t"] // BUCKET_S)] = g + self._snap = None + + # -- the HTTP side ------------------------------------------------------------------------------------------------ + + def snapshot(self, gauges: dict | None = None, now: float | None = None, ttl: float = SNAPSHOT_TTL_S) -> dict: + """GET /v1/stats: the windows, the time series and the recent requests, cached for `ttl` seconds, with `gauges` + (the server's live state, never cached).""" + now = self.clock() if now is None else now + with self._snap_lock: + snap = self._snap + if not (ttl and snap is not None and 0 <= now - snap[0] < ttl): + snap = self._snap = (now, self._compute(now)) + return dict(snap[1], gauges=gauges or {}) + + def _compute(self, now: float) -> dict: + with self.lock: + recs = list(self.records) + buckets = [list(b) for b in self.buckets if b is not None] + counts, failed_kinds, tokens = dict(self.counts), dict(self.failed_kinds), dict(self.tokens) + handed = self.handed + day, past = (dict(self.day) if self.day else None), dict(self.past) + out: dict = { + "now": now, + "started": self.started, + "uptime_s": now - self.started, + "totals": { + "requests": counts, + "failed_kinds": failed_kinds, + "tokens": tokens, + "processed": tokens["prompt"] + tokens["output"], + }, + } + by_bucket = {b[0]: b for b in buckets} + windows = {} + for name, w in WINDOWS: + lo = now - w + span = now - max(lo, self.started) + sel = [r for r in recs if r["done"] >= max(lo, self.started)] + b_lo = math.ceil(lo / BUCKET_S) + toks = sum(b[6] for k, b in by_bucket.items() if k >= b_lo) + tspan = now - max(b_lo * BUCKET_S, self.started) + windows[name] = self._window(sel, span, toks / tspan if tspan > 0 else None) + span = now - self.started + start_sel = [r for r in recs if r["done"] >= self.started] + windows["start"] = self._window( + start_sel, span, handed / span if span > 0 else None, totals=(counts["ok"], tokens) + ) + if day is not None and now - day["start"] < 86400: # a process that ran past midnight (UTC) drops "today" + first = day["first"] if day["first"] is not None else self.started + dspan = now - first + dtok = {k: day[k] + tokens[k] for k in ("prompt", "cached", "output")} + dsel = [r for r in recs if r["done"] >= day["start"]] + # handed out today: earlier processes' gauges (out_tok_s x BUCKET_S each) plus this process's + dhanded = handed + BUCKET_S * sum(g.get("out_tok_s") or 0 for g in past.values() if g["t"] >= day["start"]) + windows["today"] = self._window( + dsel, dspan, dhanded / dspan if dspan > 0 else None, totals=(day["ok"] + counts["ok"], dtok) + ) + windows["today"]["first"] = first + out["windows"] = windows + out["series"] = self._series(recs, by_bucket, past, now) + out["recent"] = [{k: r.get(k) for k in RECENT_KEYS} for r in recs[-RECENT:][::-1]] + return out + + @staticmethod + def _window(sel: list[dict], span: float, out_tok_s: float | None, totals=None) -> dict: + """One window: rates over `span` seconds, and the percentiles of the requests that finished in it.""" + minutes = span / 60 if span > 0 else None + ok = [r for r in sel if r["finish"] in OK_FINISH] + if totals is None: + n_ok = len(ok) + prompt, cached, output = (sum(r[k] for r in sel) for k in ("prompt", "cached", "output")) + else: + n_ok, t = totals + prompt, cached, output = t["prompt"], t["cached"], t["output"] + + def per_min(x: float) -> float | None: + return x / minutes if minutes else None + + itl = [r["itl"] for r in ok if r["itl"] is not None] + return { + "span_s": span, + "n": len(sel), + "ok": n_ok, + "rpm": per_min(n_ok), + "tpm": { + "input": per_min(prompt), + "cached": per_min(cached), + "output": per_min(output), + "total": per_min(prompt + output), + }, + "cache_hit_pct": 100.0 * cached / prompt if prompt else None, + "out_tok_s": out_tok_s, + "ttft": summary([r["ttft"] for r in ok if r["ttft"] is not None]), + "decode_tps": summary([r["decode_tps"] for r in ok if r["decode_tps"] is not None], low=True), + "e2e": summary([r["e2e"] for r in ok if r["e2e"] is not None]), + "itl_mean": sum(itl) / len(itl) if itl else None, + } + + @staticmethod + def _series(recs: list[dict], by_bucket: dict, past: dict, now: float) -> list[dict]: + """SERIES_BUCKETS points ending at the current one: live buckets, else today's reloaded gauges, else gaps.""" + last = int(now // BUCKET_S) + first = last - SERIES_BUCKETS + 1 + done: dict[int, list[dict]] = collections.defaultdict(list) + for r in recs: + k = int(r["done"] // BUCKET_S) + if k >= first and r["finish"] in OK_FINISH: + done[k].append(r) + gauge_keys = ("running", "running_max", "waiting", "waiting_max", "out_tok_s") + out = [] + for k in range(first, last + 1): + t0 = k * BUCKET_S + b, g, rs = by_bucket.get(k), past.get(k), done.get(k, []) + pt: dict = {"t": t0} + if b is not None: + width = min(BUCKET_S, now - t0) if k == last else BUCKET_S + pt.update( + running=b[2] / b[1], + running_max=b[3], + waiting=b[4] / b[1], + waiting_max=b[5], + out_tok_s=b[6] / width if width > 0 else None, + ) + else: + pt.update({x: g.get(x) if g is not None else None for x in gauge_keys}) + if b is not None or rs: + pt["rpm"] = len(rs) * 60 / BUCKET_S + else: + pt["rpm"] = g.get("rpm") if g is not None else None + ttft = [r["ttft"] for r in rs if r["ttft"] is not None] + dec = [r["decode_tps"] for r in rs if r["decode_tps"] is not None] + pt.update( + ttft_p50=pct(ttft, 0.5), ttft_p90=pct(ttft, 0.9), decode_p50=pct(dec, 0.5), decode_p10=pct(dec, 0.1) + ) + out.append(pt) + return out From b369cefa93170dd664ea182dfcdb4e7a475c0786 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 046/135] prodlog: an opt-in on-disk log of request records and gauges --- trnrt/serve/prodlog.py | 292 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 292 insertions(+) create mode 100644 trnrt/serve/prodlog.py diff --git a/trnrt/serve/prodlog.py b/trnrt/serve/prodlog.py new file mode 100644 index 0000000..c5e0c8b --- /dev/null +++ b/trnrt/serve/prodlog.py @@ -0,0 +1,292 @@ +"""The production log: each finished request's record and the engine's gauges, on disk, when TRNRT_PROD_DIR names a +directory (host only; unset, the default, nothing is written). Each engine writes under a directory of its own: + + ///requests.jsonl one line per finished request (serve/stats.py's record) + ///gauges.jsonl one line per 10 s: running, waiting, KV use, output + tok/s, RPM, TPM + +A record holds counts and times only (token counts, timings, the finish reason, the request id), never a prompt or an +output. Each file rotates at the UTC day and at ROTATE_BYTES to ..jsonl, and every rotated file (and any left +uncompressed by an earlier process on an earlier day) is gzipped in the background. At start, today's records and +gauges are read back into the live statistics, so GET /v1/stats's "today" window spans restarts. + +The serving threads only enqueue: put() is a put_nowait on a bounded queue (QUEUE), and a full queue drops the record +and counts it. One writer thread formats and writes, so neither the engine's step loop nor the event loop waits on the +disk. A write error (a full disk) is counted, logged at most once a minute, and the record dropped: serving goes on. +""" + +from __future__ import annotations + +import calendar +import contextlib +import gzip +import json +import logging +import os +import queue +import shutil +import threading +import time +from typing import IO, TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + + from .stats import Stats + +logger = logging.getLogger(__name__) + +ENV = "TRNRT_PROD_DIR" +ROTATE_BYTES = 1 << 30 +QUEUE = 8192 +NAMES = {"requests": "requests.jsonl", "gauges": "gauges.jsonl"} +GAUGES_RELOAD_S = 1800 # the series covers 30 minutes: older gauges are not read back + + +def utc_day(t: float) -> str: + return time.strftime("%Y-%m-%d", time.gmtime(t)) + + +def read_jsonl(path: str) -> list[dict]: + """A .jsonl or .jsonl.gz file's records; a torn line (a write cut short) is skipped.""" + out = [] + opener = gzip.open if path.endswith(".gz") else open + try: + with opener(path, "rt", encoding="utf-8", errors="replace") as f: + for line in f: + try: + rec = json.loads(line) + except ValueError: + continue + if isinstance(rec, dict): + out.append(rec) + except OSError: + pass + return out + + +def day_files(day_dir: str, kind: str) -> list[str]: + """A day directory's files of one kind, oldest first: the rotated segments ..jsonl[.gz], then the + current one.""" + stem = NAMES[kind][: -len(".jsonl")] + try: + names = os.listdir(day_dir) + except OSError: + return [] + segs: list[tuple[float, str]] = [] + for n in names: + base = n.removesuffix(".gz") + if base == NAMES[kind]: + segs.append((float("inf"), n)) + elif base.startswith(stem + ".") and base.endswith(".jsonl"): + mid = base[len(stem) + 1 : -len(".jsonl")] + if mid.isdigit(): + segs.append((int(mid), n)) + return [os.path.join(day_dir, n) for _, n in sorted(segs)] + + +def gzip_file(path: str) -> None: + try: + with open(path, "rb") as f, gzip.open(path + ".gz", "wb") as g: + shutil.copyfileobj(f, g, 1 << 20) + os.remove(path) + except OSError as e: + logger.warning("production log: could not gzip %s: %s", path, e) + + +class ProdLog: + def __init__( + self, + root: str, + stats: Stats | None = None, + rotate_bytes: int = ROTATE_BYTES, + queue_size: int = QUEUE, + clock: Callable[[], float] = time.time, + ): + self.root = os.path.expanduser(root) + self.rotate_bytes = rotate_bytes + self.clock = clock + self.q: queue.Queue = queue.Queue(queue_size) + self.dropped = dict.fromkeys(NAMES, 0) + self.written = dict.fromkeys(NAMES, 0) + self.write_errors = 0 + self.last_error: str | None = None + self._last_log = float("-inf") + self._open: dict[str, tuple[str, str, IO[bytes], int]] = {} # kind -> (day, path, file, size) + self._stop = object() + self._gzips: list[threading.Thread] = [] + self.thread = threading.Thread(target=self._run, args=(stats,), name="trnrt-prodlog", daemon=True) + self.thread.start() + + @classmethod + def from_env(cls, name: str, stats: Stats | None = None) -> ProdLog | None: + """The production log of the engine `name` under TRNRT_PROD_DIR, or None when it is unset.""" + root = os.environ.get(ENV) + if not root: + return None + return cls(os.path.join(root, name), stats=stats) + + # -- any thread --------------------------------------------------------------------------------------------------- + + def put(self, kind: str, rec: dict) -> bool: + """Enqueue one record; never blocks. False when it was dropped (the queue is full).""" + try: + self.q.put_nowait((kind, rec)) + except queue.Full: + self.dropped[kind] += 1 + return False + return True + + def close(self, timeout: float = 10.0) -> None: + """Write what is queued and close the files (at exit).""" + with contextlib.suppress(queue.Full): + self.q.put((self._stop, None), timeout=timeout) + self.thread.join(timeout) + for t in self._gzips: + t.join(timeout) + + def status(self) -> dict: + """GET /v1/stats's view of it: where it writes, today's files, and what was written, dropped and failed.""" + day = utc_day(self.clock()) + d = os.path.join(self.root, day) + files = {} + try: + for n in sorted(os.listdir(d)): + files[n] = os.path.getsize(os.path.join(d, n)) + except OSError: + pass + return { + "on": True, + "dir": self.root, + "day": day, + "files": files, + "written": dict(self.written), + "dropped": dict(self.dropped), + "write_errors": self.write_errors, + "last_error": self.last_error, + "queued": self.q.qsize(), + } + + # -- the writer thread -------------------------------------------------------------------------------------------- + + def _run(self, stats: Stats | None) -> None: + try: + os.makedirs(self.root, exist_ok=True) + except OSError as e: + self._error(f"cannot create {self.root}: {e}") + threading.Thread(target=self._gzip_old_days, name="trnrt-prodlog-gzip", daemon=True).start() + if stats is not None: + try: + self.load_today(stats) + except Exception: # a bad file must not stop the recording + logger.exception("production log: reading today's records back failed") + while True: + batch = [self.q.get()] + while len(batch) < 512: + try: + batch.append(self.q.get_nowait()) + except queue.Empty: + break + stop = False + touched = set() + for kind, rec in batch: + if kind is self._stop: + stop = True + continue + try: + line = (json.dumps(rec, ensure_ascii=False, separators=(",", ":")) + "\n").encode("utf-8") + if self._roll_due(kind, len(line)): + self._roll(kind) + self._write(kind, line) + touched.add(kind) + self.written[kind] += 1 + except OSError as e: + self.dropped[kind] += 1 + self._error(f"writing {kind}: {e}") + except Exception as e: # noqa: BLE001 one bad record must not end the recording + self.dropped[kind] += 1 + self._error(f"formatting {kind}: {e!r}") + for kind in touched: + try: + self._open[kind][2].flush() + except (OSError, KeyError) as e: + self._error(f"flushing {kind}: {e}") + if stop: + for kind in list(self._open): + self._close(kind, gz=False) + return + + def _error(self, msg: str) -> None: + self.write_errors += 1 + self.last_error = msg + now = time.monotonic() + if now - self._last_log >= 60: + self._last_log = now + logger.warning("production log: %s (records dropped; serving continues)", msg) + + def _roll_due(self, kind: str, n: int) -> bool: + cur = self._open.get(kind) + return cur is not None and (cur[0] != utc_day(self.clock()) or cur[3] + n > self.rotate_bytes) + + def _roll(self, kind: str) -> None: + """Close the kind's file (a new UTC day, or the size limit) and gzip it.""" + same_day = self._open[kind][0] == utc_day(self.clock()) + self._close(kind, gz=True, rotate=same_day) + + def _write(self, kind: str, data: bytes) -> None: + cur = self._open.get(kind) + if cur is None: + day = utc_day(self.clock()) + d = os.path.join(self.root, day) + os.makedirs(d, exist_ok=True) + path = os.path.join(d, NAMES[kind]) + f = open(path, "ab") # noqa: SIM115 held open across writes, closed by _close + cur = (day, path, f, f.tell()) + cur[2].write(data) + self._open[kind] = (cur[0], cur[1], cur[2], cur[3] + len(data)) + + def _close(self, kind: str, gz: bool, rotate: bool = False) -> None: + _, path, f, _ = self._open.pop(kind) + try: + f.close() + except OSError as e: + self._error(f"closing {path}: {e}") + if not gz: + return + if rotate: # the same day, past the size limit: the current file becomes the next numbered segment + stem = path[: -len(".jsonl")] + n = 1 + while os.path.exists(f"{stem}.{n}.jsonl") or os.path.exists(f"{stem}.{n}.jsonl.gz"): + n += 1 + dest = f"{stem}.{n}.jsonl" + os.rename(path, dest) + path = dest + t = threading.Thread(target=gzip_file, args=(path,), name="trnrt-prodlog-gzip", daemon=True) + t.start() + self._gzips.append(t) + + def _gzip_old_days(self) -> None: + today = utc_day(self.clock()) + try: + days = [d for d in os.listdir(self.root) if d != today and os.path.isdir(os.path.join(self.root, d))] + except OSError: + return + for d in days: + for n in os.listdir(os.path.join(self.root, d)): + if n.endswith(".jsonl"): + gzip_file(os.path.join(self.root, d, n)) + + def load_today(self, stats: Stats) -> None: + """Today's records and gauges (earlier processes') into the live statistics: the "today" window and the + series.""" + now = self.clock() + day = utc_day(now) + d = os.path.join(self.root, day) + recs = [r for p in day_files(d, "requests") for r in read_jsonl(p) if isinstance(r.get("done"), (int, float))] + gauges = [ + g + for p in day_files(d, "gauges") + for g in read_jsonl(p) + if isinstance(g.get("t"), (int, float)) and g["t"] >= now - GAUGES_RELOAD_S + ] + stats.merge_day(recs, gauges, calendar.timegm(time.strptime(day, "%Y-%m-%d"))) # from UTC midnight From 1423c2662e89711fb453a8aec052f403f7881780 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 047/135] prompt_cache: a chat turn's prompt ids from the previous turn's, exact or off --- trnrt/serve/prompt_cache.py | 283 ++++++++++++++++++++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100644 trnrt/serve/prompt_cache.py diff --git a/trnrt/serve/prompt_cache.py b/trnrt/serve/prompt_cache.py new file mode 100644 index 0000000..f748102 --- /dev/null +++ b/trnrt/serve/prompt_cache.py @@ -0,0 +1,283 @@ +"""The prompt token cache (TRNRT_PROMPT_CACHE_MB, host only, off by default): a chat turn's prompt ids built from an +earlier turn's ids and an encode of the new text's tail, instead of an encode of the whole rendered prompt. + +Why: an agent's next turn is its previous prompt plus a few messages, and the chat endpoint renders and encodes the +whole conversation on every turn. Rendering is the cheap half: on one host CPU core of a trn2.48xlarge, GLM-5.3-Flash's +chat template rendered a ~100K-token conversation in about 5 ms, and its tokenizer then took 195-239 ms to encode +97K-114K-token prompts. An encode holds the GIL, so a wider prompt pool would not run two at once. This keeps each key's +last (text, ids) and encodes only what follows the text both turns share. + +Why a cut is exact. A Hugging Face fast tokenizer splits its added tokens out of the raw text first, then normalizes, +pre-tokenizes and encodes each piece between them on its own. So when no added token strips its neighbours' whitespace +or is matched after normalization (lstrip, rstrip and normalized all false, as for every added token of GLM-5.3-Flash), +the text before an added token encodes the same whatever follows it: + + ids(new) = ids(old)[:i] + encode(new[c:]) + +for an added token T at character c and token index i of the old encoding, provided its match is the same in the new +text. Matching is leftmost-longest, so it is when every match that could start at or overlap c lies wholly inside the +text both share (L characters): c + need(T) <= L, need(T) the longest added token that starts with T's text; and, when +some added token could overlap T from the left (a proper suffix of it a prefix of T, or T inside it), c + M <= L for M +the longest added token. + +The key (the request's session_id, else a hash of the text's first KEY_CHARS characters: the system prompt and tools) +only finds up to BUCKET candidates, of which the one sharing the longest prefix is used; the texts are compared every +time, so a stale entry or a key several chats share costs a miss, never wrong ids. Guards, any of which turns the cache +off for good (logged, and the served path's ids served): +- at construction, a self-check of the cut against the full encode on strings built from the tokenizer's own added + tokens, and a check that the served path (apply_chat_template(tokenize=True)) equals the rendered text's encode on a + sample chat; and off outright for a tokenizer that truncates or pads, or whose added tokens strip or normalize; +- an audit of the first AUDIT_FIRST hits, and of every TRNRT_PROMPT_CACHE_AUDIT-th hit after them (default 64; 0: + none), against the served path's ids for that request. +Used by one thread (the API server's prompt pool), so it takes no lock. +""" + +from __future__ import annotations + +import bisect +import hashlib +import logging +import os +import sys +from array import array +from collections import OrderedDict +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + +logger = logging.getLogger(__name__) + +ENV_MB = "TRNRT_PROMPT_CACHE_MB" +ENV_AUDIT = "TRNRT_PROMPT_CACHE_AUDIT" +KEY_CHARS = 256 # an agent's system prompt and tools: the same on every turn, and on every session of that agent +BUCKET = {"s": 8, "h": 32} # candidates per key: a session id's, and all the sessions behind one content key +AUDIT_FIRST = 8 # every one of a process's first hits is audited: a wrong cut would otherwise serve 63 more + + +def _lcp(a: str, b: str) -> int: + """The length of a's and b's common prefix: one compare when b extends a (the usual case), else a gallop and then a + binary search, comparing chunks in place (b.startswith(a[lo:hi], lo)).""" + if b.startswith(a): + return len(a) + n, lo, step = min(len(a), len(b)), 0, 1024 + while True: # a[:lo] == b[:lo] + hi = min(n, lo + step) + if hi == lo or not b.startswith(a[lo:hi], lo): + break + lo, step = hi, 2 * step + hi -= 1 # a[:hi + 1] != b[:hi + 1]: the prefix length is in [lo, hi] + while lo < hi: + mid = (lo + hi + 1) // 2 + if b.startswith(a[lo:mid], lo): + lo = mid + else: + hi = mid - 1 + return lo + + +class _Entry: + __slots__ = ("bstart", "btok", "ids", "key", "nbytes", "text") + + def __init__(self, text: str, ids: array, bstart: array, btok: array): + self.text, self.ids, self.bstart, self.btok = text, ids, bstart, btok + self.key = "" + self.nbytes = sys.getsizeof(text) + sum(x.itemsize * len(x) for x in (ids, bstart, btok)) + 200 + + +class PromptCache: + def __init__(self, tok, max_bytes: int, audit_every: int = 64): + self.fast = getattr(tok, "_tokenizer", None) # a PreTrainedTokenizerFast's tokenizers.Tokenizer + self.max_bytes, self.audit_every = int(max_bytes), int(audit_every) + self.entries: OrderedDict[int, _Entry] = OrderedDict() # least recently used first + self.buckets: dict[str, list[int]] = {} + self.bytes, self._next = 0, 0 + self.counts = dict.fromkeys(("hits", "misses", "reused_tokens", "tail_tokens", "audits", "evictions"), 0) + self.off: str | None = None # why the cache is off, once it is + self.margin = 0 + self.need: dict[int, int] = {} # token id -> characters past its start both texts must share before a cut there + self.added_ids: frozenset[int] = frozenset() + if self.fast is None: + self.off = "not a fast tokenizer" + logger.warning("prompt cache off: %s", self.off) + return + added = list(self.fast.get_added_tokens_decoder().values()) + texts = [t.content for t in added] + self.margin = max((len(x) for x in texts), default=0) + for x in texts: + over = any(y != x and (x in y or any(x.startswith(y[k:]) for k in range(1, len(y)))) for y in texts) + self.need[self.fast.token_to_id(x)] = self.margin if over else max(len(y) for y in texts if y.startswith(x)) + self.added_ids = frozenset(self.need) + bad = [t.content for t in added if t.lstrip or t.rstrip or t.normalized] + if self.fast.truncation is not None or self.fast.padding is not None: + # transformers sets these per call and a raw encode keeps them: a truncating tokenizer.json would cut the + # cached ids and never the served ones + self.off = f"the tokenizer truncates or pads ({self.fast.truncation}, {self.fast.padding})" + elif bad: + self.off = f"added tokens that strip or normalize: {bad[:4]}" + elif not added: + self.off = "no added tokens to cut at" + else: + self.off = self._self_check(texts) or self._template_check(tok) + if self.off is not None: + logger.warning("prompt cache off: %s", self.off) + + @classmethod + def from_env(cls, tok) -> PromptCache | None: + """The cache TRNRT_PROMPT_CACHE_MB sizes (MiB of host memory), or None when it is unset or 0.""" + mb = float(os.environ.get(ENV_MB, "0") or 0) + if mb <= 0 or tok is None: + return None + return cls(tok, int(mb * 2**20), int(os.environ.get(ENV_AUDIT, "64") or 64)) + + def _full(self, text: str): + return self.fast.encode(text, add_special_tokens=False) # type: ignore[union-attr] # off when it is None + + def _bounds(self, enc, c0: int = 0, t0: int = 0) -> tuple[list[int], list[int]]: + """(character start, token index) of every added token in an encoding of text[c0:], shifted to the whole + text.""" + s, t = [], [] + for j, i in enumerate(enc.ids): + if i in self.added_ids: + s.append(enc.offsets[j][0] + c0) + t.append(j + t0) + return s, t + + def _cut(self, e: _Entry, text: str, shared: int | None = None) -> tuple[int, int] | None: + """(character, token index) to cut at: the last added token of e's text at c with c + need <= the common prefix; + None if there is none.""" + n = _lcp(e.text, text) if shared is None else shared + j = bisect.bisect_left(e.bstart, n) - 1 + while j >= 0: + c, i = e.bstart[j], e.btok[j] + if c + self.need[e.ids[i]] <= n: + return c, i + j -= 1 + return None + + def encode(self, text: str, key: str | None = None, served: Callable[[], list[int]] | None = None) -> list[int]: + """text's ids, as the tokenizer's encode(text, add_special_tokens=False) gives them. served: () -> the ids the + uncached path serves for the same request (api.chat_prompt passes apply_chat_template(tokenize=True)), what an + audit compares with; without it, the text's full encode.""" + if self.off is not None: + return list(served()) if served is not None else self._full(text).ids + if key is not None: + k = f"s:{key}" + else: + k = ( + "h:" + + hashlib.sha1(text[:KEY_CHARS].encode("utf-8", "surrogatepass"), usedforsecurity=False).hexdigest() + ) + best: tuple[int, int] | None = None + for eid in self.buckets.get(k, ()): # the candidate sharing the longest prefix + n = _lcp(self.entries[eid].text, text) + if best is None or n > best[1]: + best = (eid, n) + cut = self._cut(self.entries[best[0]], text, best[1]) if best is not None else None + if best is not None and cut is not None: + e = self.entries[best[0]] + c, i = cut + tail = self._full(text[c:]) + ids = e.ids[:i].tolist() + tail.ids + ts, tt = self._bounds(tail, c, i) + j = bisect.bisect_left(e.btok, i) + bstart, btok = e.bstart[:j].tolist() + ts, e.btok[:j].tolist() + tt + self.counts["hits"] += 1 + self.counts["reused_tokens"] += i + self.counts["tail_tokens"] += len(tail.ids) + h = self.counts["hits"] + if self.audit_every > 0 and (h <= AUDIT_FIRST or h % self.audit_every == 0): + self.counts["audits"] += 1 + want = list(served()) if served is not None else self._full(text).ids + if want != ids: + self.off = f"audit mismatch at hit {h} ({len(ids)} ids cached, {len(want)} served)" + logger.error("prompt cache off: %s", self.off) + self.entries.clear() + self.buckets.clear() + self.bytes = 0 + return want + self._drop(best[0]) # the new text continues that one + else: + enc = self._full(text) + ids = enc.ids + bstart, btok = self._bounds(enc) + self.counts["misses"] += 1 + self._put(k, _Entry(text, array("I", ids), array("q", bstart), array("q", btok))) + return ids + + def _drop(self, eid: int) -> None: + e = self.entries.pop(eid) + self.bytes -= e.nbytes + b = self.buckets[e.key] + b.remove(eid) + if not b: + del self.buckets[e.key] + + def _put(self, k: str, e: _Entry) -> None: + if e.nbytes > self.max_bytes: + return + e.key = k + eid, self._next = self._next, self._next + 1 + self.entries[eid] = e + self.bytes += e.nbytes + b = self.buckets.setdefault(k, []) + b.append(eid) + if len(b) > BUCKET[k[0]]: + self._drop(b[0]) + while self.bytes > self.max_bytes: + self._drop(next(iter(self.entries))) + self.counts["evictions"] += 1 + + def _self_check(self, added: list[str]) -> str | None: + """The cut against the full encode on strings built from the tokenizer's own added tokens, at every cut point + of them: None, or what differed.""" + a, b = added[0], added[min(1, len(added) - 1)] + base = f"System line, café 中文 🙂\n{a}\nuser text with spaces\t{b} more words{a}\n\n tail{b}" + texts = [base + extra for extra in ("", " and then some", "\n" + b + "x", "y" + a + " z", a, a + b)] + pairs = [(old, new) for old in [base[:n] for n in range(1, len(base) + 1, 3)] + texts for new in texts] + # the new text parting inside an added token that punctuation runs into: a cut there would split "!<|" in two + pairs += [(f"x{p}{t}hello", f"x{p}{t[:k]}zz more") for t in added[:4] for p in "!}>" for k in (1, 2)] + for old, new in pairs: + enc = self._full(old) + bs, bt = self._bounds(enc) + cut = self._cut(_Entry(old, array("I", enc.ids), array("q", bs), array("q", bt)), new) + if cut is None: + continue + c, i = cut + if enc.ids[:i] + self._full(new[c:]).ids != self._full(new).ids: + return f"self-check: a cut at character {c} of {old[:40]!r} -> {new[:40]!r} differs" + return None + + def _template_check(self, tok) -> str | None: + """The chat endpoint serves apply_chat_template(tokenize=True); the cache encodes the text it renders. The two + must agree on a sample chat, or a transformers version adding special tokens on one path only would change the + served ids (the audits catch it on real chats too): None, or what differed.""" + if not getattr(tok, "chat_template", None): + return None # the chat endpoint cannot render without one either + m = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Hello, 世界! Ready?"}, + {"role": "assistant", "content": "Yes."}, + {"role": "user", "content": "Then go on."}, + ] + try: + want = tok.apply_chat_template(m, add_generation_prompt=True, tokenize=True) + want = list(want["input_ids"] if hasattr(want, "keys") else want) + text = tok.apply_chat_template(m, add_generation_prompt=True, tokenize=False) + except Exception as e: # noqa: BLE001 any failure of the sample turns the cache off, never the server + return f"template check failed: {type(e).__name__}: {e}"[:200] + if want != self._full(text).ids: + return "template check: apply_chat_template(tokenize=True) differs from the rendered text's encode" + return None + + def stats(self) -> dict: + return dict( + self.counts, + on=self.off is None, + off=self.off, + entries=len(self.entries), + buckets=len(self.buckets), + bytes=self.bytes, + max_bytes=self.max_bytes, + audit_every=self.audit_every, + ) From 6a175a9e416696c95f44886f10b8c55e61e8bf6d Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 048/135] dashboard: one self-contained page that draws /v1/stats --- trnrt/serve/dashboard.py | 350 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 350 insertions(+) create mode 100644 trnrt/serve/dashboard.py diff --git a/trnrt/serve/dashboard.py b/trnrt/serve/dashboard.py new file mode 100644 index 0000000..bc56621 --- /dev/null +++ b/trnrt/serve/dashboard.py @@ -0,0 +1,350 @@ +"""GET /dashboard: one self-contained page (inline CSS and JavaScript, no CDN, so it works where the host has no +internet access) that polls GET /v1/stats every 2 s and draws it. serve/stats.py defines every number on it. An engine +serves it for itself; the router serves it too, with a link to each engine (`?engine=i`), since the engines listen on +127.0.0.1 only.""" + +PAGE = r""" + + + + +TrnRT serving + + + +
+
+

TrnRT serving

+ + + + connecting + +
+ +
+ +

Rates and percentiles by window

+
+ +

Last 30 minutes (10-s points)

+
+ +

Recent requests

+
+ +

Data collection

+
+ +

+All times are measured inside the server, from when the engine receives a request (after the chat template and +tokenization) to its last token. They exclude the client's network round trip, so a client sees a longer TTFT and +end-to-end time. TTFT = first token - arrival. Decode tok/s is per stream: (output tokens - 1) ÷ (last token - first +token); p10 is the slowest tenth of streams. RPM, TPM and the percentiles count requests by when they finished. Output +tok/s (aggregate) counts tokens as the engine hands them out. TPM total = input + output; cached is the part of the input +served from the prefix cache. "Since start" is this engine process; "today" (UTC) includes earlier processes' records +when the production log is on. Each engine has its own numbers: behind the router, pick the engine above. +

+
+ + + + +""" From 0813a998dda95b3fc8ef097fc3bf5177941d0986 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 049/135] api: richer /v1/stats, /dashboard, the production log and the prompt token cache --- trnrt/serve/api.py | 157 ++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 141 insertions(+), 16 deletions(-) diff --git a/trnrt/serve/api.py b/trnrt/serve/api.py index fdccb0b..3c5b5e7 100644 --- a/trnrt/serve/api.py +++ b/trnrt/serve/api.py @@ -6,9 +6,15 @@ - POST /v1/completions: a prompt as text or token ids, through the same generation path; - GET /v1/models; - GET /health: 200 while the engine can serve, 503 with the reason once it cannot; -- GET /v1/stats: the KV-cache gauges the router (serve/router.py) uses to break ties between engines; +- GET /v1/stats: the KV-cache gauges the router (serve/router.py) uses to break ties between engines, the engine's live + state, and request rates and latency percentiles over recent windows (serve/stats.py); +- GET /dashboard: a page drawing /v1/stats (serve/dashboard.py); - GET /metrics: Prometheus counters and histograms (trnrt/metrics.py). +TRNRT_PROD_DIR keeps each finished request's record and the gauges on disk (serve/prodlog.py), and +TRNRT_PROMPT_CACHE_MB builds a chat turn's prompt ids from the previous turn's (serve/prompt_cache.py); both are off +unless set. + A request may carry a `session_id` (a string): the engine keeps that session's prefix referenced in its prefix cache, and the router pins the session to one engine. @@ -23,6 +29,7 @@ import atexit import contextlib import copy +import dataclasses import functools import json import logging @@ -43,12 +50,16 @@ from fastapi import FastAPI, HTTPException from fastapi import Request as HTTPRequest -from fastapi.responses import JSONResponse, PlainTextResponse, Response, StreamingResponse +from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse, Response, StreamingResponse from ..engine.request import Request, SamplingParams, Status from ..engine.scheduler import QueueFull from ..metrics import Metrics, kv_pages, spec_stats +from .dashboard import PAGE as DASHBOARD_PAGE from .parsers import THINK_OPEN, StreamParser, chat_finish_reason, parse_full, template_messages, tool_choice_error +from .prodlog import ProdLog +from .prompt_cache import PromptCache +from .stats import Stats if TYPE_CHECKING: # only annotations: importing this module does not load the engine and its model code from ..engine.engine import LLMEngine @@ -64,6 +75,8 @@ class _Stream: loop: asyncio.AbstractEventLoop out: asyncio.Queue sent: int = 0 + stream: bool = False # the HTTP response streams (GET /v1/stats's records) + at_arrival: tuple[int, int] | None = None # (running, waiting) when the engine took it @dataclass @@ -82,6 +95,7 @@ class _Submit: session_id: str | None = None priority: int | None = None out: asyncio.Queue = field(default_factory=asyncio.Queue) + stream: bool = False # Set on a 503 that answers a full waiting queue (scheduler.QueueFull) or a request displaced from it: the engine is @@ -123,8 +137,18 @@ class EngineLoop: """Owns the engine: one thread runs its step loop and is the only one that touches it. Handlers hand it requests and aborts through `inbox`; it hands each request's tokens to that request's asyncio queue after every step.""" - def __init__(self, engine: LLMEngine): + def __init__(self, engine: LLMEngine, name: str = "engine"): self.engine = engine + # GET /v1/stats and /dashboard (serve/stats.py), and with TRNRT_PROD_DIR the production log (serve/prodlog.py), + # under `name`, this engine's own directory. Fed on this loop's thread through _hook, which never raises. + self.stats = Stats() + self.stats.kv = lambda: kv_pages(engine) + self.stats_errors = 0 + self._stats_err_t = float("-inf") + self.prodlog = ProdLog.from_env(name, self.stats) + if self.prodlog is not None: + prodlog = self.prodlog + self.stats.on_bucket = lambda g: prodlog.put("gauges", g) self._stopping = threading.Event() self.metrics = Metrics() self.vocab = engine.mcfg.vocab_size # the embedding's rows: a prompt id outside them fails the step @@ -153,6 +177,8 @@ def stop(self, wait: float = 120.0) -> None: self.thread.join(wait) if self.thread.is_alive(): logger.warning("trnrt: the engine thread did not stop within %.0f s; closing the engine anyway", wait) + if self.prodlog is not None: + self.prodlog.close() def health_error(self) -> str | None: """Why the engine cannot serve, or None. A worker that died while this loop waits in a step (a device collective @@ -174,7 +200,12 @@ def _put(self, kind: str, payload) -> None: self.inbox.put((kind, payload)) def submit( - self, prompt_ids: list[int], params: SamplingParams, session_id: str | None = None, priority: int | None = None + self, + prompt_ids: list[int], + params: SamplingParams, + session_id: str | None = None, + priority: int | None = None, + stream: bool = False, ) -> _Submit: # Checked here, before the engine sees the request: an id outside the vocabulary fails the step, and at tp > 1 # one failed step fails the engine for every request (_fail). @@ -185,12 +216,38 @@ def submit( if session_id is not None and not isinstance(session_id, str): raise HTTPException(status_code=400, detail="session_id must be a string") sub = _Submit(uuid.uuid4().hex, prompt_ids, params, asyncio.get_running_loop(), session_id, priority) + sub.stream = stream self._put("add", sub) return sub def abort(self, rid: str) -> None: self.inbox.put(("abort", rid)) + def _hook(self, fn, *args) -> None: + """Run a statistics or recording hook on the engine thread: an exception is counted and logged (at most once a + minute), never raised, so recording cannot stop serving (GET /v1/stats shows the count as stats_errors).""" + try: + fn(*args) + except Exception: # the statistics must never fail a step or a request + self.stats_errors += 1 + now = time.monotonic() + if now - self._stats_err_t >= 60: + self._stats_err_t = now + logger.exception("trnrt: a statistics hook failed (%d so far)", self.stats_errors) + + def _tick(self, tokens: int) -> None: + s = self.engine.scheduler + self._hook(self.stats.tick, time.monotonic(), tokens, len(s.running), len(s.waiting)) + + def _finished(self, st: _Stream, finish: str | None = None) -> None: + """A request's end: its record into the live statistics and the production log (O(1), never blocks).""" + self._hook(self._finished_hook, st, finish) + + def _finished_hook(self, st: _Stream, finish: str | None) -> None: + rec = self.stats.record(st.req, st.stream, st.at_arrival, finish) + if self.prodlog is not None: + self.prodlog.put("requests", rec) + def _drain(self, block: bool) -> None: try: item = self.inbox.get(timeout=0.05) if block else self.inbox.get_nowait() @@ -202,24 +259,30 @@ def _drain(self, block: bool) -> None: post = payload.loop.call_soon_threadsafe # priority only when sent: an engine without priority scheduling refuses any value (LLMEngine.add_request) extra = {} if payload.priority is None else {"priority": payload.priority} + sch = self.engine.scheduler + at_arrival = (len(sch.running), len(sch.waiting)) try: req = self.engine.add_request( payload.prompt_ids, payload.params, rid=payload.rid, session_id=payload.session_id, **extra ) except QueueFull as e: post(payload.out.put_nowait, ("error", str(e), 503, True)) + self._hook(self.stats.failed, "queue_full") except ValueError as e: post(payload.out.put_nowait, ("error", str(e), 400)) + self._hook(self.stats.failed, "refused") except Exception as e: # noqa: BLE001 add_request raised before the request reached the scheduler traceback.print_exc() # (its last step): that request fails, the engine keeps serving post(payload.out.put_nowait, ("error", f"request refused: {e!r}", 500)) + self._hook(self.stats.failed, "error") else: - self.streams[payload.rid] = _Stream(req, payload.loop, payload.out) + self.streams[payload.rid] = _Stream(req, payload.loop, payload.out, 0, payload.stream, at_arrival) post(payload.out.put_nowait, ("accepted", req)) elif kind == "abort": st = self.streams.pop(payload, None) if st is not None and st.req.status is not Status.FINISHED: self.engine.abort(st.req) + self._finished(st, "abort") # "wake": nothing to do, the loop re-checks whether it is stopping try: item = self.inbox.get_nowait() @@ -275,6 +338,7 @@ def _fail(self, reason: str) -> None: streams, self.streams = self.streams, {} for st in streams.values(): st.loop.call_soon_threadsafe(st.out.put_nowait, ("error", f"engine failed: {reason}", 503)) + self._hook(self.stats.failed, "engine") while True: try: kind, payload = self.inbox.get_nowait() @@ -292,6 +356,7 @@ def _loop(self) -> None: self._drain(block=not self.engine.has_work()) if not self.engine.has_work(): self._read_local_k() + self._tick(0) continue self._local_k_stale = True try: @@ -307,8 +372,10 @@ def _loop(self) -> None: st.loop.call_soon_threadsafe(st.out.put_nowait, ("error", f"engine step failed: {e}", 500)) if st.req.status is not Status.FINISHED: self.engine.abort(st.req) + self._hook(self.stats.failed, "engine") self.streams.clear() continue + handed = 0 # output tokens handed out by this step (GET /v1/stats: the aggregate output tok/s) batch: dict = {} # one cross-thread call per event loop and step, not one per stream for rid, st in list(self.streams.items()): req = st.req @@ -320,9 +387,11 @@ def _loop(self) -> None: if finished and req.error is not None: # ended unserved: displaced from a full queue batch.setdefault(st.loop, []).append((st.out, ("error", req.error[1], req.error[0], True))) del self.streams[rid] + self._hook(self.stats.failed, "displaced") continue if ready > st.sent or finished: new = out[st.sent : ready] + handed += len(new) lps = req.logprobs[st.sent : ready] if with_lp else None st.sent = ready batch.setdefault(st.loop, []).append( @@ -331,8 +400,10 @@ def _loop(self) -> None: if finished: del self.streams[rid] self.metrics.on_finish(req) + self._finished(st) for loop, items in batch.items(): loop.call_soon_threadsafe(_deliver, items) + self._tick(handed) def _deliver(items) -> None: @@ -719,10 +790,25 @@ def template_kwargs(tok, body: dict) -> tuple[dict, bool | None]: return kw, thinking -def chat_prompt(tok, body: dict) -> tuple[list[int], dict]: +def _cached_ids(tok, messages: list, kwargs: dict, cache: PromptCache, session) -> list[int]: + """_template_ids through the prompt token cache: the template renders the text, and the cache encodes it from the + session's (or the same system prompt's) previous text, giving the ids _template_ids would (or, on an audit that + finds otherwise, those ids, and the cache turns itself off).""" + import jinja2 + + try: + text = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False, **kwargs) + except jinja2.TemplateError as e: + raise _bad(f"the chat template refused the messages: {e}") from None + key = session if isinstance(session, str) else None + return cache.encode(text, key, served=lambda: _template_ids(tok, messages, kwargs)) + + +def chat_prompt(tok, body: dict, cache: PromptCache | None = None) -> tuple[list[int], dict]: """A /v1/chat/completions request's prompt ids, rendered by the model's chat template from its messages (assistant tool_calls arguments as objects: parsers.template_messages), and the StreamParser / parse_full arguments its output - is parsed with. HTTPException 400 on a malformed request.""" + is parsed with. HTTPException 400 on a malformed request. With `cache` (TRNRT_PROMPT_CACHE_MB) the ids come through + the prompt token cache, the same ids.""" messages = body.get("messages") if not messages: raise HTTPException(status_code=400, detail="messages is required") @@ -738,7 +824,10 @@ def chat_prompt(tok, body: dict) -> tuple[list[int], dict]: tools = body.get("tools") if tool_choice != "none" else None if tools: kwargs["tools"] = tools - ids = _template_ids(tok, messages, kwargs) + if cache is not None and cache.off is None: + ids = _cached_ids(tok, messages, kwargs, cache, body.get("session_id")) + else: + ids = _template_ids(tok, messages, kwargs) # With thinking on, the template opens the reasoning block in the prompt and the output starts inside it; with # thinking off it closes the block there, and the output starts with the answer. tail = tok.decode(ids[-4:]).rstrip() @@ -795,15 +884,20 @@ async def _unless_disconnected(http: HTTPRequest, aw): return task.result() -def build_app(engine: LLMEngine, model_name: str) -> FastAPI: +def build_app(engine: LLMEngine, model_name: str, name: str = "engine") -> FastAPI: + """The engine's API. `name` names its directory in the production log (TRNRT_PROD_DIR): `serve` gives each engine + its own, by port.""" app = FastAPI(title="TrnRT") - loop_ = EngineLoop(engine) + loop_ = EngineLoop(engine, name) tok = engine.tokenizer # A chat request's prompt is rendered and tokenized on one worker thread with its own tokenizer, not on the event # loop: GLM-5.3-Flash's template takes tens of milliseconds at 100K tokens, which would stall every other request's # stream, and two threads in one Hugging Face fast tokenizer can raise "Already borrowed". prompt_pool = ThreadPoolExecutor(1, thread_name_prefix="trnrt-prompt") prompt_tok = copy.deepcopy(tok) + # TRNRT_PROMPT_CACHE_MB (serve/prompt_cache.py): a chat prompt's ids from the previous turn's; used on prompt_pool's + # thread only, by the chat endpoint (/tokenize renders without it) + prompt_cache = PromptCache.from_env(prompt_tok) gen_defaults = generation_defaults(getattr(engine, "model_path", None)) if gen_defaults: logger.info("trnrt: sampling defaults from the checkpoint's generation_config.json: %s", gen_defaults) @@ -815,15 +909,44 @@ async def health(): return JSONResponse(status_code=503, content={"status": "failed", "error": err}) return {"status": "ok"} - @app.get("/v1/stats") - async def stats(): - """The router's tie-break between engines: KV pages in use and usable (metrics.kv_pages).""" + def live_gauges() -> dict: + """GET /v1/stats's live state. kv_pages_used and kv_pages_total are the router's tie-break between engines + (metrics.kv_pages); the rest is the scheduler, the speculation, the scheduling settings, the prompt token cache + and the production log. Read without the engine thread: counts and settings only, nothing changed.""" used, total = kv_pages(engine) gauges: dict = {"kv_pages_used": used, "kv_pages_total": total} spec = spec_stats(engine) if spec is not None: gauges["spec"] = spec - return {"gauges": gauges} + s, cfg = engine.scheduler, engine.cfg + sch = getattr(engine, "scheduling", None) + gauges.update( + running=len(s.running), + waiting=len(s.waiting), + seats=getattr(cfg, "max_num_seqs", None), + cap=getattr(cfg, "max_running", None), + kv_used_pct=100.0 * used / total if total else None, + preemptions=getattr(s, "num_preemptions", 0), + model=model_name, + max_model_len=cfg.max_model_len, + tp=cfg.tp, + scheduling=dataclasses.asdict(sch) if sch is not None else None, + prompt_cache=prompt_cache.stats() if prompt_cache is not None else {"on": False}, + prod=loop_.prodlog.status() if loop_.prodlog is not None else {"on": False}, + stats_errors=loop_.stats_errors, + ) + return gauges + + @app.get("/v1/stats") + async def stats(): + """The live gauges, and request rates and percentiles over recent windows, a 30-minute series and the recent + requests (serve/stats.py), computed in a worker thread: sorting up to stats.RING records must not hold the + event loop that delivers every stream's tokens.""" + return JSONResponse(await asyncio.to_thread(loop_.stats.snapshot, live_gauges())) + + @app.get("/dashboard") + async def dashboard(): + return HTMLResponse(DASHBOARD_PAGE) @app.get("/metrics") async def metrics(): @@ -847,7 +970,7 @@ async def _generate( include_usage = isinstance(opts, dict) and _flag(opts, "include_usage") # SGLang's per-request priority (engine/scheduler.py); an engine without priority scheduling refuses it (400) priority = _int(body, "priority", None, -_INT64_MAX - 1) - sub = loop_.submit(prompt_ids, params, body.get("session_id"), priority) + sub = loop_.submit(prompt_ids, params, body.get("session_id"), priority, stream=stream) # Admission errors (a prompt that does not fit, an unsupported parameter, a full queue) arrive before any token: # surface them as HTTP status codes rather than inside a 200 stream. first = await sub.out.get() @@ -1033,7 +1156,9 @@ async def completions(http: HTTPRequest): @app.post("/v1/chat/completions") async def chat(http: HTTPRequest): body = await _body(http) - ids, parse = await asyncio.get_running_loop().run_in_executor(prompt_pool, chat_prompt, prompt_tok, body) + ids, parse = await asyncio.get_running_loop().run_in_executor( + prompt_pool, chat_prompt, prompt_tok, body, prompt_cache + ) # max_tokens defaults to the rest of the context (vLLM's default): reasoning at the template's default effort # (Max) can run long, and a smaller hidden default would cut the answer off room = max(engine.cfg.max_model_len - len(ids), 1) From 0216b16a7c0e245bf2fa958eca69cb90b2843aeb Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 050/135] cli: name each engine's production log directory by its port --- trnrt/__main__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/trnrt/__main__.py b/trnrt/__main__.py index 87d8354..6bb1f3c 100644 --- a/trnrt/__main__.py +++ b/trnrt/__main__.py @@ -636,7 +636,9 @@ def cmd_engine(argv: list[str]) -> int: print(f"bucket warmup {w['seconds']:.1f} s over {w['graphs']} graphs", flush=True) # The router keeps its connections to the engines alive and drops idle ones after 30 s: uvicorn closing them first # (its default keep-alive is 5 s) raced a request onto a closing connection. - serve(build_app(eng, a.served_model_name), host=a.host, port=a.port, log_level="warning", timeout_keep_alive=600) + # name: the engine's own directory in the production log (TRNRT_PROD_DIR), by its port, stable across restarts + app = build_app(eng, a.served_model_name, name=f"engine-{a.port}") + serve(app, host=a.host, port=a.port, log_level="warning", timeout_keep_alive=600) return 0 From 3ad110f6b04f925cc88ebd513de1a8eed1440dcf Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 051/135] router: each engine's /v1/stats and its dashboard --- trnrt/serve/router.py | 55 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 53 insertions(+), 2 deletions(-) diff --git a/trnrt/serve/router.py b/trnrt/serve/router.py index e1741a0..a308398 100644 --- a/trnrt/serve/router.py +++ b/trnrt/serve/router.py @@ -55,7 +55,11 @@ sessions go to the live ones, so a load balancer in front keeps sending); 503 {"status": "down"} once none is. Each body names the workers down ("down": their URLs) and carries the routing counters. Up and down are the router's view (above): the /health polls and the requests' refused connections and 503s; -- GET /v1/models: the answer of the first live worker that gives one (least loaded first), 503 if none does. +- GET /v1/models: the answer of the first live worker that gives one (least loaded first), 503 if none does; +- GET /v1/stats: every worker's /v1/stats ({"engines": [{"url", "up", "stats"}], "router": the /health counters}), or + with ?engine=i worker i's alone, with a "router" block naming the engines; +- GET /dashboard: the engines' page (serve/dashboard.py), one engine at a time (?engine=i): the engines listen on + 127.0.0.1 only, so the router is where a browser can reach it. python -m trnrt.serve.router --worker-urls http://127.0.0.1:8111,http://127.0.0.1:8112 --port 8000 """ @@ -73,9 +77,10 @@ import httpx from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse, Response, StreamingResponse +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response, StreamingResponse from .. import log +from .dashboard import PAGE as DASHBOARD_PAGE logger = log.get("router") @@ -620,6 +625,52 @@ async def models(): return JSONResponse(r.json()) raise HTTPException(status_code=503, detail="no engine answered /v1/models") + def engines() -> list[dict]: + return [{"url": u, "up": ok} for u, ok in zip(router.workers, router.up)] + + async def worker_stats(w: int) -> dict: + """Worker w's GET /v1/stats; HTTPException 503 when it does not answer with one.""" + try: + r = await client.get(router.workers[w] + "/v1/stats", timeout=health_timeout_s) + if r.status_code == 200: + body = r.json() + if isinstance(body, dict): + return body + except (httpx.HTTPError, ValueError): + pass + raise HTTPException(status_code=503, detail=f"engine {w} did not answer /v1/stats") + + def engine_index(raw: str | None) -> int | None: + if raw is None: + return None + if not raw.isdigit() or int(raw) >= len(router.workers): + raise HTTPException(status_code=400, detail=f"engine must be an integer in [0, {len(router.workers)})") + return int(raw) + + @app.get("/v1/stats") + async def stats(engine: str | None = None): + """Every engine's statistics (serve/stats.py), or engine i's with ?engine=i. Read only: no routing state.""" + w = engine_index(engine) + if w is not None: + body = await worker_stats(w) + body["router"] = {"engine": w, "engines": engines()} + return JSONResponse(body) + out = [] + for i, e in enumerate(engines()): + try: + out.append({**e, "stats": await worker_stats(i)}) + except HTTPException as x: + out.append({**e, "stats": None, "error": x.detail}) + return JSONResponse({"engines": out, "router": router.health()}) + + @app.get("/dashboard") + async def dashboard(engine: str | None = None): + """The engines' page, for one engine: without ?engine, the first live one's.""" + if engine_index(engine) is None: + live = router._live() + return RedirectResponse(f"dashboard?engine={live[0] if live else 0}", status_code=307) + return HTMLResponse(DASHBOARD_PAGE) + async def proxy(path: str, http: Request): body = await _json_object(http) sid = router.session_of(body) From f79e05f5d7db09b931722e3c835353e69467e379 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 052/135] tests: /v1/stats keeps the router's gauges among the new fields --- tests/test_serve_api.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_serve_api.py b/tests/test_serve_api.py index 66abcd4..45e7851 100644 --- a/tests/test_serve_api.py +++ b/tests/test_serve_api.py @@ -364,7 +364,8 @@ def test_an_sdk_request_with_every_neutral_field_is_served(serve): def test_health_stats_models_and_metrics(serve): client, _ = serve(PLAIN) assert client.get("/health").json() == {"status": "ok"} - assert client.get("/v1/stats").json() == {"gauges": {"kv_pages_used": 6, "kv_pages_total": 100}} + gauges = client.get("/v1/stats").json()["gauges"] + assert (gauges["kv_pages_used"], gauges["kv_pages_total"]) == (6, 100) and "spec" not in gauges # the router's assert client.get("/v1/models").json()["data"][0]["id"] == "glm" client.post("/v1/completions", json={"prompt": [1, 2]}) m = client.get("/metrics").text From 08066bb28e30efca3b2374305421a0bbcbf92dd9 Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 053/135] tests: the prompt token cache gives the served ids turn by turn, or turns off --- tests/test_prompt_cache.py | 487 +++++++++++++++++++++++++++++++++++++ 1 file changed, 487 insertions(+) create mode 100644 tests/test_prompt_cache.py diff --git a/tests/test_prompt_cache.py b/tests/test_prompt_cache.py new file mode 100644 index 0000000..905982a --- /dev/null +++ b/tests/test_prompt_cache.py @@ -0,0 +1,487 @@ +"""The prompt token cache (serve/prompt_cache.py): every id list it returns must equal a full encode of the same text, +and the served path's (apply_chat_template(tokenize=True)) ids, in every case it serves, or it must turn itself off and +serve those. + +The tokenizer is a byte-level BPE trained here on chat-like text (merges across spaces, punctuation, digits and +multi-byte characters, as GLM-5.3-Flash's does), with GLM-5.3-Flash's tags as added tokens and its own chat template +(tests/fixtures). Conversations grow turn by turn as an agent's do (tool calls and their results, reasoning, sessions +that branch, contents that contain the tags' text or a part of it), and every turn's ids are compared. A second +tokenizer adds tokens that overlap one another, the case the cut's margin exists for. The real GLM-5.3-Flash tokenizer +runs too when it is in the local Hugging Face cache (tests/test_glm53_chat.py's fixture).""" + +from __future__ import annotations + +import json +import os +import random + +import pytest + +pytest.importorskip("tokenizers") +pytest.importorskip("transformers") + +from tests import test_glm53_chat as glm +from trnrt.serve import prompt_cache as PC + +TEMPLATE = os.path.join(os.path.dirname(__file__), "fixtures", "GLM-5.3-Flash-chat_template.jinja") +TAGS = [ + "[gMASK]", + "", + "<|system|>", + "<|user|>", + "<|assistant|>", + "<|observation|>", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", +] +WORDS = [ + "the", + "a", + "of", + "to", + "and", + "in", + "is", + "it", + "you", + "that", + "was", + "for", + "on", + "are", + "with", + "as", + "I", + "his", + "they", + "be", + "at", + "one", + "have", + "this", + "from", + "or", + "had", + "by", + "hot", + "word", + "but", + "what", + "some", + "we", + "can", + "out", + "other", + "were", + "all", + "there", + "when", + "up", + "use", + "your", + "how", + "said", + "an", + "each", + "she", + "which", + "do", + "their", + "time", + "if", + "will", + "way", + "about", + "many", + "then", + "them", + "write", + "would", + "like", + "so", + "these", + "her", + "long", + "make", + "thing", + "see", + "him", + "two", + "has", + "look", + "more", + "day", + "could", + "go", + "come", + "did", + "number", + "sound", + "no", + "most", + "people", + "my", + "over", + "know", + "water", + "than", + "call", + "first", + "who", + "may", + "down", + "side", + "been", + "now", + "find", + "def", + "class", + "return", + "import", + "self", + "None", + "True", + "False", + "print", + "json", + "path", + "file", + "error", + "test", + "pytest", + "assert", + "value", + "key", +] +EXTRA = [ + "中文", + "café", + "naïve", + "🙂", + "Ω", + "\t", + "\n", + " ", + "\n\n", + "!!", + "...", + "{}", + '{"a": 1}', + "x=1;", + "€100", + "<", + ">", + "|", + "<|", + "|>", + "", +] + + +def _corpus(rng, n=1500): + return [" ".join(rng.choice(WORDS + EXTRA) for _ in range(rng.randrange(3, 30))) for _ in range(n)] + + +def _tokenizer(added=TAGS, template=True): + from tokenizers import AddedToken, Tokenizer, decoders, models, pre_tokenizers, trainers + from transformers import PreTrainedTokenizerFast + + tk = Tokenizer(models.BPE()) + tk.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + tk.decoder = decoders.ByteLevel() + trainer = trainers.BpeTrainer( + vocab_size=700, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False + ) + tk.train_from_iterator(_corpus(random.Random(0)), trainer) + tok = PreTrainedTokenizerFast(tokenizer_object=tk, clean_up_tokenization_spaces=False) + tok.add_tokens([AddedToken(t, normalized=False) for t in added]) + if template: + with open(TEMPLATE) as f: + tok.chat_template = f.read() + return tok + + +@pytest.fixture(scope="module") +def tok(): + return _tokenizer() + + +def full(tok, text): + return tok._tokenizer.encode(text, add_special_tokens=False).ids + + +def served(tok, messages, **kw): + ids = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, **kw) + return list(ids["input_ids"] if hasattr(ids, "keys") else ids) + + +def text_of(tok, messages, **kw): + return tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False, **kw) + + +def _content(rng): + """A message's text: words, multi-byte characters, whitespace runs, and now and then a tag's text or a part of one + (a user pasting a transcript).""" + parts = [] + for _ in range(rng.randrange(1, 25)): + r = rng.random() + if r < 0.06: + parts.append(rng.choice(TAGS)) + elif r < 0.12: + t = rng.choice(TAGS) + parts.append(t[: rng.randrange(1, len(t))]) + elif r < 0.25: + parts.append(rng.choice(EXTRA)) + else: + parts.append(rng.choice(WORDS)) + return rng.choice(["", " ", "\n"]).join(parts) + + +TOOLS = [ + { + "type": "function", + "function": { + "name": "bash", + "description": "Run a command", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}, + }, + } +] + + +def _turn(rng, messages, k): + """The agent's next turn: its reply (text, reasoning, tool calls), the tool results, and the next user message.""" + msgs = list(messages) + reply = {"role": "assistant", "content": _content(rng)} + if rng.random() < 0.5: + reply["reasoning_content"] = _content(rng) + if rng.random() < 0.5: + calls = [ + { + "id": f"call_{k}_{j}", + "type": "function", + "function": {"name": "bash", "arguments": json.dumps({"command": _content(rng)})}, + } + for j in range(rng.randrange(1, 3)) + ] + reply["tool_calls"] = calls + msgs.append(reply) + msgs.extend({"role": "tool", "tool_call_id": c["id"], "content": _content(rng)} for c in calls) + else: + msgs.append(reply) + if rng.random() < 0.8 or msgs[-1]["role"] == "assistant": + msgs.append({"role": "user", "content": _content(rng)}) + return msgs + + +def _template_messages(messages): + from trnrt.serve.parsers import template_messages + + return template_messages(messages) + + +def test_the_cache_is_on_for_a_byte_level_bpe_with_plain_added_tokens(tok): + cache = PC.PromptCache(tok, 1 << 26) + assert cache.off is None, cache.off + + +@pytest.mark.parametrize("seed", range(6)) +def test_agent_sessions_turn_by_turn_give_the_served_ids(tok, seed): + """Several sessions, interleaved, each growing turn by turn; some requests carry their session id and some do not + (the content key), and now and then a session branches from an earlier turn (an agent retrying). Every turn's ids + are the served path's, and the cache hits on most turns (it is not vacuously falling back).""" + rng = random.Random(seed) + cache = PC.PromptCache(tok, 1 << 26, audit_every=0) # no audits: the comparison here is the check + system = {"role": "system", "content": "You are a coding agent. " + _content(rng)} + sessions = {f"s{i}": [system, {"role": "user", "content": _content(rng)}] for i in range(4)} + history: dict[str, list] = {s: [] for s in sessions} + for k in range(60): + sid = rng.choice(list(sessions)) + if history[sid] and rng.random() < 0.1: # a branch: retry from an earlier turn + sessions[sid] = rng.choice(history[sid]) + msgs = sessions[sid] = _turn(rng, sessions[sid], k) + history[sid].append(msgs) + tm = _template_messages(msgs) + kw = {"tools": TOOLS} if seed % 2 else {} + text = text_of(tok, tm, **kw) + key = sid if rng.random() < 0.7 else None + got = cache.encode(text, key) + assert got == served(tok, tm, **kw) == full(tok, text), (seed, k) + assert cache.off is None + assert cache.counts["hits"] > cache.counts["misses"], cache.counts + assert cache.counts["reused_tokens"] > cache.counts["tail_tokens"], cache.counts + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_text_pairs_around_the_added_tokens(tok, seed): + """Texts built of words, tags and tag fragments, each next one sharing a random prefix of the last (cut anywhere: + inside a tag, right after one, inside a multi-byte run); every encode equals the full encode.""" + rng = random.Random(100 + seed) + cache = PC.PromptCache(tok, 1 << 24, audit_every=0) + text = _content(rng) + for k in range(400): + cut = rng.randrange(0, len(text) + 1) + text = text[:cut] + _content(rng) if rng.random() < 0.8 else _content(rng) + text + assert cache.encode(text, "one") == full(tok, text), (seed, k, text[-80:]) + assert cache.counts["hits"] > 100 + + +def test_overlapping_added_tokens_keep_the_margin(): + """Added tokens that overlap one another (one a prefix, a suffix or a part of another): the cut needs the longer + shared margin. Random texts over their fragments, every encode the full one. With the margin reduced to the token's + own length (the mutant need = len), the same texts produce a wrong cut.""" + added = ["", "b", "b", "", "xxyy", "<|z|>", "qzz", "w<|z|>w"] + t = _tokenizer(added, template=False) + frags = [*added, "<", ">", "a", "b", "c", "x", "y", " ", "word", "", "b<", "<|", "|>", "z", "q", "w", "zz"] + + def run(cache, rng): + text = "".join(rng.choice(frags) for _ in range(30)) + for _ in range(1500): + text = text[: rng.randrange(0, len(text) + 1)] + "".join( + rng.choice(frags) for _ in range(rng.randrange(1, 12)) + ) + if cache.encode(text, "k") != full(t, text): + return False + return True + + for seed in range(3): + cache = PC.PromptCache(t, 1 << 24, audit_every=0) + assert cache.off is None, cache.off + assert run(cache, random.Random(seed)), seed + + def naive(): + c = PC.PromptCache(t, 1 << 24, audit_every=0) + c.need = {i: len(t.convert_ids_to_tokens(i)) for i in c.need} # no margin past the token itself + return c + + assert not all(run(naive(), random.Random(seed)) for seed in range(3)), "the control: some cut must be wrong" + + +def test_sessions_and_content_keys_find_their_own_entries(tok): + cache = PC.PromptCache(tok, 1 << 24, audit_every=0) + long = " ".join(WORDS) # a system prompt longer than KEY_CHARS: the content key is the same on the next turn + a = text_of(tok, [{"role": "system", "content": "Alpha. " + long}, {"role": "user", "content": "first task here"}]) + b = text_of(tok, [{"role": "system", "content": "Beta. " + long}, {"role": "user", "content": "other task"}]) + cache.encode(a, "s1") + cache.encode(b, None) + a2 = a + "more words of the first session <|user|> again" + b2 = b + " and the next turn" + assert cache.encode(a2, "s1") == full(tok, a2) and cache.counts["hits"] == 1 + assert ( + cache.encode(b2, None) == full(tok, b2) and cache.counts["hits"] == 2 + ) # same first KEY_CHARS: the content key + assert ( + cache.encode(a2 + " x", "s2") == full(tok, a2 + " x") and cache.counts["hits"] == 2 + ) # another session: a miss + + +def test_memory_bound_and_buckets(tok): + cache = PC.PromptCache(tok, 12_000, audit_every=0) + rng = random.Random(3) + for i in range(200): + t = text_of(tok, [{"role": "user", "content": f"{i} " + _content(rng)}]) + assert cache.encode(t, f"s{i % 50}") == full(tok, t) + assert cache.bytes <= cache.max_bytes + assert cache.counts["evictions"] > 0 and all(len(b) <= PC.BUCKET["s"] for b in cache.buckets.values()) + + +def test_an_audit_that_disagrees_turns_the_cache_off_and_serves_the_served_ids(tok): + """The served path is the reference: if it ever differs from the cut (a tokenizer the checks could not see), the + cache serves the served ids, turns off for good, and passes every later request through.""" + cache = PC.PromptCache(tok, 1 << 24, audit_every=1) + t1 = text_of(tok, [{"role": "user", "content": "hello there"}]) + cache.encode(t1, "s") + t2 = t1 + "general kenobi <|user|> next" + wrong = [*full(tok, t2), 5] + assert cache.encode(t2, "s", served=lambda: wrong) == wrong + assert cache.off is not None and "audit mismatch" in cache.off and not cache.entries + t3 = t2 + " more" + assert cache.encode(t3, "s", served=lambda: full(tok, t3)) == full(tok, t3) + assert cache.counts["hits"] == 1 + + +def test_audits_run_on_the_first_hits_and_then_every_nth(tok): + calls = [] + cache = PC.PromptCache(tok, 1 << 24, audit_every=4) + t = text_of(tok, [{"role": "user", "content": "start"}]) + cache.encode(t, "s") + for i in range(20): + t = t + f" step {i} <|observation|> ok" + cache.encode(t, "s", served=lambda t=t: calls.append(1) or full(tok, t)) + assert cache.counts["hits"] == 20 and cache.counts["audits"] == len(calls) + assert len(calls) == sum(1 for h in range(1, 21) if h <= PC.AUDIT_FIRST or h % 4 == 0) + + +def test_the_cache_stays_off_for_tokenizers_it_cannot_cut_exactly(): + from tokenizers import AddedToken + + strip = _tokenizer(TAGS[:4]) + strip.add_tokens([AddedToken("<|strip|>", lstrip=True, normalized=False)]) + assert "strip or normalize" in PC.PromptCache(strip, 1 << 20).off + trunc = _tokenizer(TAGS[:4]) + trunc._tokenizer.enable_truncation(64) + assert "truncates or pads" in PC.PromptCache(trunc, 1 << 20).off + none = _tokenizer([]) + assert PC.PromptCache(none, 1 << 20).off == "no added tokens to cut at" + assert PC.PromptCache(object(), 1 << 20).off == "not a fast tokenizer" + off = PC.PromptCache(strip, 1 << 20) + text = "plain <|user|> text" + assert off.encode(text, "s", served=lambda: [1, 2, 3]) == [1, 2, 3] and not off.entries # off: the served ids + + +def test_the_template_check_compares_the_served_path(tok, monkeypatch): + """At construction the served path (apply_chat_template(tokenize=True)) must equal the rendered text's encode on a + sample chat; a transformers that added a special token on one path only turns the cache off.""" + t = _tokenizer() + real = t.apply_chat_template + + def plus_bos(*a, **kw): + out = real(*a, **kw) + if kw.get("tokenize"): + ids = list(out["input_ids"] if hasattr(out, "keys") else out) + return [0, *ids] + return out + + monkeypatch.setattr(t, "apply_chat_template", plus_bos) + assert "template check" in PC.PromptCache(t, 1 << 20).off + + +def test_from_env(tok, monkeypatch): + monkeypatch.delenv(PC.ENV_MB, raising=False) + assert PC.PromptCache.from_env(tok) is None + monkeypatch.setenv(PC.ENV_MB, "0") + assert PC.PromptCache.from_env(tok) is None + monkeypatch.setenv(PC.ENV_MB, "1.5") + monkeypatch.setenv(PC.ENV_AUDIT, "0") + c = PC.PromptCache.from_env(tok) + assert c is not None and c.max_bytes == int(1.5 * 2**20) and c.audit_every == 0 + assert PC.PromptCache.from_env(None) is None + + +# The real GLM-5.3-Flash tokenizer and template, when cached locally (tests/test_glm53_chat.py's fixture). +glm_tok = glm.tok + + +def test_the_real_tokenizer_turn_by_turn(glm_tok): + cache = PC.PromptCache(glm_tok, 1 << 26, audit_every=0) + assert cache.off is None, cache.off + rng = random.Random(1) + msgs = [{"role": "system", "content": "You are a coding agent."}, {"role": "user", "content": "Fix the test."}] + for k in range(25): + msgs = _turn(rng, msgs, k) + tm = _template_messages(msgs) + text = text_of(glm_tok, tm, tools=glm.TOOLS) + assert cache.encode(text, "s") == served(glm_tok, tm, tools=glm.TOOLS) == full(glm_tok, text), k + assert cache.counts["hits"] >= 20 From 384dc7bbf1fdde7411bb3a8fb8ba7e72a2f23eba Mon Sep 17 00:00:00 2001 From: AmanSwar Date: Mon, 9 Nov 2026 17:14:36 +0530 Subject: [PATCH 054/135] tests: the statistics, the production log, the dashboard and their router views --- tests/test_serve_stats.py | 339 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 339 insertions(+) create mode 100644 tests/test_serve_stats.py diff --git a/tests/test_serve_stats.py b/tests/test_serve_stats.py new file mode 100644 index 0000000..b144fe5 --- /dev/null +++ b/tests/test_serve_stats.py @@ -0,0 +1,339 @@ +"""GET /v1/stats's statistics (serve/stats.py), the production log (serve/prodlog.py), GET /dashboard +(serve/dashboard.py), the prompt token cache behind the chat endpoint (serve/prompt_cache.py), and the router's views +of all of them (serve/router.py). + +The statistics run on fake clocks, so every rate and percentile is checked against a hand count. The HTTP tests use +tests/test_serve_api.py's scripted engine.""" + +from __future__ import annotations + +import gzip +import json +import os +import time +import types + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from tests.test_serve_api import PLAIN, ScriptedEngine +from trnrt.serve import prodlog as PL +from trnrt.serve import stats as S + +T0 = 1_800_000_000.0 # a wall-clock second on a bucket boundary + + +class Clock: + def __init__(self): + self.mono = 1000.0 + + def wall(self) -> float: + return T0 + self.mono - 1000.0 + + def mono_(self) -> float: + return self.mono + + +def req(rid, arrival, first, done, prompt=100, cached=40, output=11, finish="stop"): + return types.SimpleNamespace( + rid=rid, + arrival_time=arrival, + first_token_time=first, + finish_time=done, + num_prompt=prompt, + num_cached_tokens=cached, + output_ids=[1] * output, + finish_reason=finish, + spec_proposed=0, + spec_accepted=0, + ) + + +def stats(): + c = Clock() + return S.Stats(clock=c.wall, mono=c.mono_), c + + +def test_percentiles_and_derived_times(): + assert S.pct([], 0.5) is None and S.pct([3.0], 0.9) == 3.0 + assert S.pct([1.0, 2.0, 3.0, 4.0], 0.5) == 2.5 and S.pct([0.0, 10.0], 0.9) == 9.0 # numpy's linear + rec = S.derive({"arrival": 10.0, "first": 10.5, "done": 12.5, "output": 11}) + assert (rec["ttft"], rec["e2e"], rec["decode_tps"], rec["itl"]) == (0.5, 2.5, 5.0, 0.2) + one = S.derive({"arrival": 1.0, "first": 2.0, "done": 2.0, "output": 1}) + assert one["decode_tps"] is None and one["itl"] is None + + +def test_windows_count_requests_by_their_finish_and_tokens_as_handed_out(): + st, c = stats() + # three finished requests: two ok in the last minute, one ok 4 minutes ago; one aborted, one failed + c.mono = 1000.0 + 600 # 10 minutes after start + for rid, arr, first, done, out in (("a", 1590, 1590.5, 1592.5, 11), ("b", 1580, 1581, 1585, 21)): + st.record(req(rid, arr, first, done, output=out), stream=True, at_arrival=(1, 0)) + st.record(req("old", 1350, 1351, 1360, output=5), stream=False, at_arrival=None) + st.record(req("gone", 1595, None, 1596, output=0, finish="abort"), stream=True, at_arrival=(2, 1)) + st.failed("queue_full") + for k in range(6): # 6 ticks of 10 tokens in the last minute + st.tick(1550.0 + 10 * k, 10, 2, 1) + snap = st.snapshot({"x": 1}, now=c.wall(), ttl=0) + w1, w5 = snap["windows"]["1m"], snap["windows"]["5m"] + assert (w1["n"], w1["ok"], w1["rpm"]) == (3, 2, 2.0) + assert (w5["ok"], w5["rpm"]) == (3, 3 / 5) + assert w1["tpm"]["input"] == 300 and w1["tpm"]["output"] == 32 and w1["cache_hit_pct"] == 40.0 + assert w1["ttft"]["n"] == 2 and w1["ttft"]["p50"] == pytest.approx(0.75) # 0.5 and 1.0 + assert w1["decode_tps"]["mean"] == pytest.approx((5.0 + 5.0) / 2) # 10 / 2 s and 20 / 4 s + assert w1["out_tok_s"] == pytest.approx(60 / 60) + assert snap["totals"]["requests"] == {"ok": 3, "aborted": 1, "failed": 1} + assert snap["totals"]["failed_kinds"] == {"queue_full": 1} and snap["gauges"] == {"x": 1} + assert [r["rid"] for r in snap["recent"]] == ["gone", "old", "b", "a"] # newest finished first + assert snap["recent"][0]["running_at_arrival"] == 2 and snap["recent"][0]["stream"] is True + series = snap["series"] + assert len(series) == S.SERIES_BUCKETS and series[-1]["t"] == (c.wall() // S.BUCKET_S) * S.BUCKET_S + assert sum(p["rpm"] or 0 for p in series) * S.BUCKET_S / 60 == 3 # each ok request in its finish bucket + + +def test_the_snapshot_is_shared_for_its_ttl_but_the_gauges_are_live(): + st, _c = stats() + a = st.snapshot({"g": 1}, now=T0 + 5, ttl=1.0) + st.record(req("a", 1001, 1002, 1003), stream=False, at_arrival=None) + b = st.snapshot({"g": 2}, now=T0 + 5.5, ttl=1.0) + assert b["recent"] == a["recent"] == [] and b["gauges"] == {"g": 2} + assert len(st.snapshot({}, now=T0 + 6.5, ttl=1.0)["recent"]) == 1 + + +def test_a_closed_bucket_becomes_a_gauges_record(): + st, _ = stats() + out = [] + st.on_bucket = out.append + st.kv = lambda: (25, 100) + st.record(req("a", 1001, 1002, 1005, prompt=60, output=6), stream=False, at_arrival=None) + st.tick(1001.0, 30, 3, 1) + st.tick(1004.0, 10, 1, 3) + st.tick(1012.0, 0, 0, 0) # the next bucket: the first one closes + (g,) = out + assert g["t"] == T0 and (g["running"], g["running_max"], g["waiting"], g["waiting_max"]) == (2, 3, 2, 3) + assert (g["out_tok_s"], g["rpm"], g["input_tpm"], g["output_tpm"], g["kv_used_pct"]) == ( + 4.0, + 6.0, + 360.0, + 36.0, + 25.0, + ) + + +def test_the_production_log_writes_reloads_and_rotates(tmp_path): + st, c = stats() + log = PL.ProdLog(str(tmp_path / "engine-8001"), stats=st, clock=c.wall, rotate_bytes=400) + for i in range(6): + log.put("requests", st.record(req(f"r{i}", 1001, 1002, 1003 + i), stream=False, at_arrival=(0, 0))) + log.put("gauges", {"t": T0, "running": 1.0, "out_tok_s": 3.0}) + log.close() + day = tmp_path / "engine-8001" / PL.utc_day(T0) + reqs = [r for p in PL.day_files(str(day), "requests") for r in PL.read_jsonl(p)] + assert [r["rid"] for r in reqs] == [f"r{i}" for i in range(6)] + assert any(n.endswith(".gz") for n in os.listdir(day)), "a rotated segment is gzipped" + assert set(reqs[0]) >= {"prompt", "cached", "output", "ttft", "finish"} and "prompt_ids" not in reqs[0] + assert log.written == {"requests": 6, "gauges": 1} and log.write_errors == 0 + # a new process: today's records come back into the "today" window + st2, c2 = stats() + c2.mono = 1100.0 + log2 = PL.ProdLog(str(tmp_path / "engine-8001"), stats=st2, clock=c2.wall) + log2.close() + today = st2.snapshot({}, now=c2.wall(), ttl=0)["windows"]["today"] + assert today["ok"] == 6 and today["tpm"]["input"] is not None + assert log2.status()["on"] and "requests.jsonl" in "".join(log2.status()["files"]) + + +def test_a_torn_line_and_a_full_queue_lose_only_themselves(tmp_path): + p = tmp_path / "x.jsonl.gz" + with gzip.open(p, "wt") as f: + f.write('{"a": 1}\n{"b": 2') + assert PL.read_jsonl(str(p)) == [{"a": 1}] + log = PL.ProdLog.__new__(PL.ProdLog) + import queue + + log.q, log.dropped = queue.Queue(1), dict.fromkeys(PL.NAMES, 0) + assert log.put("requests", {}) and not log.put("requests", {}) and log.dropped["requests"] == 1 + + +def test_the_production_log_is_off_unless_set(monkeypatch, tmp_path): + monkeypatch.delenv(PL.ENV, raising=False) + assert PL.ProdLog.from_env("engine-8001") is None + monkeypatch.setenv(PL.ENV, str(tmp_path)) + log = PL.ProdLog.from_env("engine-8001") + assert log is not None and log.root == str(tmp_path / "engine-8001") + log.close() + + +# -- over HTTP --------------------------------------------------------------------------------------------------------- + + +@pytest.fixture +def app(): + from fastapi.testclient import TestClient + + from trnrt.serve import api + + started = [] + + def start(eng, **kw): + client = TestClient(api.build_app(eng, "glm", **kw)) + client.__enter__() + started.append((client, eng)) + return client + + yield start + for client, eng in started: + client.__exit__(None, None, None) + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def _wait(cond, t=10.0): + end = time.monotonic() + t + while not cond(): + assert time.monotonic() < end + time.sleep(0.02) + + +def test_stats_over_http_keep_the_routers_gauges_and_count_what_ran(app): + eng = ScriptedEngine(PLAIN) + client = app(eng) + before = client.get("/v1/stats").json() + g = before["gauges"] + assert (g["kv_pages_used"], g["kv_pages_total"]) == (6, 100) # serve/router.py reads these two + assert g["running"] == 0 and g["waiting"] == 0 and g["model"] == "glm" and g["prompt_cache"] == {"on": False} + assert g["prod"] == {"on": False} and g["stats_errors"] == 0 + assert client.post("/v1/completions", json={"prompt": [1, 2], "max_tokens": 3}).status_code == 200 + events = client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "x"}], "stream": True}) + assert events.status_code == 200 + assert client.post("/v1/completions", json={"prompt": [1] * 70}).status_code == 400 # the engine refuses it + _wait(lambda: client.get("/v1/stats").json()["totals"]["requests"]["ok"] == 2) + snap = client.get("/v1/stats").json() + assert snap["totals"]["failed_kinds"] == {"refused": 1} + recent = snap["recent"] + assert len(recent) == 2 and {r["stream"] for r in recent} == {True, False} + plain = next(r for r in recent if not r["stream"]) + assert (plain["prompt"], plain["output"], plain["cached"], plain["finish"]) == (2, 3, 2, "length") + assert snap["windows"]["start"]["ok"] == 2 and snap["windows"]["1m"]["tpm"]["output"] is not None + + +def test_the_dashboard_is_one_self_contained_page(app): + client = app(ScriptedEngine(PLAIN)) + r = client.get("/dashboard") + assert r.status_code == 200 and r.headers["content-type"].startswith("text/html") + page = r.text + assert "TrnRT serving" in page and "v1/stats" in page + assert "http://" not in page and "https://" not in page and "