diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c1aaa3..3d51364 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,58 @@ 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)). +- 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)). +- `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. +- `stop_regex` refuses a pattern whose search can take exponential or polynomial time, and matches within the last + 2048 characters of output. +- SGLang's `no_logs` on `/generate`: the request's record stays out of the production log. +- 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. +- Scheduling flags with SGLang's names, host only and per engine: `--schedule-policy` (`fcfs`, `lpm`, `lof`, `spf`), + priority scheduling with the request field `priority` and priority preemption, `--max-queued-requests` and + `--max-queued-tokens` (a full queue answers 503), and `--radix-eviction-policy` (`lru`, `lfu`, `slru`, `priority`, + `tlru`). The defaults schedule as before ([docs/CONFIGURATION.md](docs/CONFIGURATION.md#scheduling)). +- `GET /v1/stats` adds the engine's live state and request rates, TTFT, decode-speed and end-to-end percentiles over + recent windows, and `GET /dashboard` draws them; the router serves both for each engine + ([docs/API.md](docs/API.md#statistics-and-the-dashboard)). The fields the router reads are unchanged. +- `TRNRT_PROD_DIR`: an opt-in on-disk log of request records (counts and times only) and gauges. +- `TRNRT_PROMPT_CACHE_MB`: an opt-in prompt token cache that builds a chat turn's prompt ids from the previous turn's, + with the same ids as a full encode or none. +- Structured outputs through xgrammar, an optional dependency: `response_format` `json_object` and `json_schema`, vLLM's + `guided_*` and `structured_outputs`, and SGLang's `json_schema`, `regex` and `ebnf`. Also `min_tokens` and + `thinking_token_budget`. They need `--structured-outputs`, which adds graphs that take a token mask + ([docs/API.md](docs/API.md#structured-outputs-penalties-and-thinking-budgets)). Tested on CPU; not yet run on + Trainium2, and the check that the default graphs are unchanged is pending. +- Repetition, frequency and presence penalties, with `--penalties` and `--penalty-buckets`, which add graphs that take + penalty inputs ([docs/CONFIGURATION.md](docs/CONFIGURATION.md#structured-outputs-penalties-and-thinking-budgets)). + Tested on CPU; not yet run on Trainium2. + +### 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. +- `service_tier` other than `"auto"` or `"default"` is refused with 400 on `/v1/chat/completions`, `/v1/completions` + and `/v1/responses`. + ## [0.1.0] - 2026-11-08 The first public release. diff --git a/docs/API.md b/docs/API.md index 3b401a0..24292b7 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)), and serves each engine's statistics and dashboard ([below](#statistics-and-the-dashboard)). 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). It also serves the OpenAI Responses API ([below](#responses-api)) and the Anthropic Messages API ([below](#anthropic-messages-api)). ## Example @@ -30,7 +30,10 @@ 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/responses`, `/v1/messages`, + `/v1/messages/count_tokens`, `/v1/models`, `/health`, `/tokenize`, `/detokenize`, `/v1/stats`, `/dashboard`, 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)). @@ -39,9 +42,224 @@ curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' - - The router keeps an idle client connection open for 75 s (`--timeout-keep-alive`). A proxy or client connection pool 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. +- Structured outputs, penalties, `min_tokens` and thinking budgets need an engine started with `--structured-outputs` or + `--penalties` ([below](#structured-outputs-penalties-and-thinking-budgets)); otherwise they are answered with 400, + naming the flag. +- Not supported, answered with 400: `tool_choice` `"required"` or a named tool, prompt logprobs (`prompt_logprobs`, + `echo` with `logprobs`), and SGLang's `min_new_tokens` on the OpenAI endpoints (use `min_tokens`). - 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`. + 503 once the engine has failed, with the reason on `/health`, and 503 for a full queue + ([Scheduling](#scheduling-and-priority)). + +## Scheduling and priority + +Each engine schedules its own requests; the flags are in [CONFIGURATION.md](CONFIGURATION.md#scheduling). The default +is first come, first served, with no queue limit. + +| Field or answer | Behaviour | +|---|---| +| `priority` (integer, on `/v1/chat/completions` and `/v1/completions`) | SGLang's request field. With `--enable-priority-scheduling`, a waiting request with a higher value starts first (`--schedule-low-priority-values-first`: a lower one), and a request without one starts last. With `--radix-eviction-policy priority`, the prefix cache keeps the prefixes of high-priority requests longer. Without either flag the field is refused with 400, since nothing would use it | +| A full queue | Past `--max-queued-requests` or `--max-queued-tokens` waiting on an engine, a request is answered 503 `{"detail": "The request queue is full."}`, as SGLang answers it, with the header `x-trnrt-queue-full: 1`. The router passes it on and keeps routing to that engine; OpenAI clients retry a 503 | +| Displaced from a full queue | With priority scheduling, a request better than the worst one waiting takes its place. The displaced request gets 503 `The request is aborted by a higher priority request.` (as an error event if it streams). A request that has already generated tokens is never displaced | +| Preempted for priority | With priority scheduling, when every seat is taken, a waiting request whose priority beats a running one's by more than `--priority-scheduling-preemption-threshold` (10) takes that one's seat. The preempted request waits and later resumes from its cached prefix; its client sees only a pause | + +## Statistics and the dashboard + +| Endpoint | Where | Answer | +|---|---|---| +| `GET /v1/stats` | each engine | `gauges` (live: `kv_pages_used`, `kv_pages_total`, `running`, `waiting`, `seats`, `cap`, `kv_used_pct`, `preemptions`, `scheduling`, `prompt_cache`, `prod`, and `spec` under speculation); `windows` (`1m`, `5m`, `15m`, `start`, and `today` with the production log on: requests, RPM, TPM, prefix-cache hit %, output tok/s, TTFT, per-stream decode tok/s and end-to-end percentiles); `totals`; `series` (30 minutes of 10-s points); `recent` (the last 50 finished requests) | +| `GET /v1/stats` | the router | every engine's answer, as `{"engines": [{"url", "up", "stats"}], "router": ...}`; with `?engine=i`, engine i's alone, with a `router` block naming the engines | +| `GET /dashboard` | each engine, and the router | one self-contained page (no CDN) that draws `/v1/stats` every 2 s. The router's shows one engine at a time, with a link to each (`?engine=i`) | + +The times are measured inside the engine, from when it receives a request (after the chat template and tokenization) +to its last token, so they exclude the client's network round trip. A record holds counts and times, never a prompt or +an output. Like the rest of the router, these endpoints have no authentication. + +## 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"`. A match must start within 2048 characters of the newest text. A pattern whose search can take exponential or polynomial time (nested or overlapping repeats, such as `(a+)+` or `a*a*`) is refused with 400. | +| `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, `reasoning` fields other than +`effort` and `summary`, `service_tier` other than `"auto"` or `"default"`, item references, non-text input parts and +tools other than functions. `metadata`, `user`, `safety_identifier` and `prompt_cache_key` are accepted and change +nothing: the prefix cache matches every prompt by its tokens, with no key. `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` on an engine started without `--structured-outputs` (with it, the budget is enforced as +`thinking_token_budget`), `tool_choice` `any` and `tool` (a tool-call grammar for GLM's format is not built), +image and document blocks, `redacted_thinking`, a `tool_result` with `is_error` true (GLM's chat template cannot mark a +failed call; put the error in the result's content), 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 +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?"}]}' +``` + +## 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`; with `--structured-outputs`, `json_schema`, `regex`, `ebnf` and `min_new_tokens`; with `--penalties`, `frequency_penalty`, `presence_penalty` and `repetition_penalty`. 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`. | +| `no_logs` | `true` keeps the request's record out of the production log (`TRNRT_PROD_DIR`); it is still counted in `/v1/stats`. | + +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. | + +Guided decoding (`json_schema`, `regex`, `ebnf`), `min_new_tokens` and the penalties are served as on the OpenAI +routes, by an engine started with `--structured-outputs` or `--penalties` +([below](#structured-outputs-penalties-and-thinking-budgets)); otherwise they are answered with 400, naming the flag. + +Answered with 400, naming the field: prompt logprobs (`logprob_start_len` >= 0 with `return_logprob`); in +`sampling_params`, `n` > 1, `structural_tag`, `stop_regex`, `no_stop_trim`, `skip_special_tokens` false, +`spaces_between_special_tokens` false, `logit_bias`, `stream_interval` or `beam_width` other than 1, and +`custom_params`; and 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. An engine's 503 counts toward the router's `down_after`, as a stalled `/health` does, rather than marking it down at once | +| `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. An engine with no answer after `timeout` plus the health timeout is named "no answer" in the 400 | + +None of these returns a secret or a host path the operator did not pass. + +## Structured outputs, penalties and thinking budgets + +Each needs graphs that take inputs from the host. An engine serves them only when started with the flag +([CONFIGURATION.md](CONFIGURATION.md#structured-outputs-penalties-and-thinking-budgets)). + +| Field | Flag | Behaviour | +|---|---|---| +| `response_format` `{"type": "json_object"}` | `--structured-outputs` | a JSON object (the schema `{"type": "object"}`, as SGLang) | +| `response_format` `{"type": "json_schema", "json_schema": {"schema": ...}}` | `--structured-outputs` | JSON matching the schema | +| vLLM `guided_json`, `guided_regex`, `guided_choice`, `guided_grammar`; `structured_outputs` `{json, regex, choice, grammar, json_object}` | `--structured-outputs` | a schema, a regular expression, one of the strings, an EBNF grammar | +| SGLang `json_schema` (a string), `regex`, `ebnf` | `--structured-outputs` | the same | +| `min_tokens` | `--structured-outputs` | no end-of-sequence or `stop_token_ids` id among the first `min_tokens` tokens; a stop string that ends among them does not end the request (vLLM) | +| `thinking_token_budget` | `--structured-outputs` | at most this many reasoning tokens after ``; then `` is forced | +| `repetition_penalty`, `frequency_penalty`, `presence_penalty` | `--penalties` | vLLM's formulas: repetition over prompt and output tokens, frequency and presence over the output | + +- One constraint per request. Two (say `response_format` and `regex`) are refused with 400, as SGLang refuses them. +- Guided decoding uses [xgrammar](https://github.com/mlc-ai/xgrammar), an optional dependency: + `pip install 'trnrt[grammar]'`. Without it, guided requests are answered with 400. +- GLM-5.3-Flash reasons first. As in vLLM, the constraint applies to the answer after ``, not to the reasoning. +- Logprobs stay the raw model's, before the mask and the penalties, as vLLM reports them by default. +- A request with penalties, a grammar, or a thinking budget near its end waits for its newest token before each step, + so it gives up the overlap of host and device work; the requests beside it keep it. Under the dynamic `mtp-*` presets + the engine stops speculating while a request that uses any of these runs. +- If a grammar can continue only with text no token spells, the request ends there with `finish_reason` `"stop"`, and a + warning is logged. diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 9b2f063..957559d 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -17,8 +17,13 @@ | `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/serve/responses.py`, `anthropic.py` | The OpenAI Responses API and the Anthropic Messages API, translated into the chat path. | +| `trnrt/serve/stats.py`, `trnrt/serve/dashboard.py` | `GET /v1/stats`'s request rates, latency percentiles and time series per engine, and the page `GET /dashboard` draws from them. | +| `trnrt/serve/prodlog.py` | The opt-in production log (`TRNRT_PROD_DIR`): request records and gauges on disk. | +| `trnrt/serve/prompt_cache.py` | The opt-in prompt token cache (`TRNRT_PROMPT_CACHE_MB`): a chat turn's prompt ids from the previous turn's. | | `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. | diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 5d13aaa..86cd81b 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -68,6 +68,26 @@ arrival order. It changes no graph, so it runs on the compile cache of the confi docker run ... trnrt serve --model /model --port 8000 --preset batch8 --max-running 4 # the mounts and variables of DEPLOY.md, step 4 ``` +## Scheduling + +Each engine schedules its own requests. These flags are host only: they change no graph, so they run on the compile +cache of the configuration they are added to, and `serve` hands them to every engine. The names are SGLang's, except +`--max-queued-tokens` and the policies SGLang does not have (`spf`, `tlru`). The defaults schedule and evict as an +engine without these flags. + +| Argument | Default | | +|---|---|---| +| `--schedule-policy` | fcfs | the order waiting requests start in: `fcfs` arrival, `lpm` the longest cached prefix first (arrival order above 128 waiting, as SGLang), `lof` the longest `max_tokens` first, `spf` the shortest prompt first; `priority` is `fcfs` with `--enable-priority-scheduling` | +| `--enable-priority-scheduling` | off | each request's `priority` first ([API.md](API.md#scheduling-and-priority)), higher values first; KV pressure preempts the worst priority first. Takes `fcfs`, `lof` or `spf`, not `lpm` (SGLang's rule) | +| `--schedule-low-priority-values-first` | off | lower priority values first | +| `--priority-scheduling-preemption-threshold` | 10 | with every seat taken, a waiting request better than a running one by more than this takes its seat | +| `--disable-priority-preemption` | off | never preempt a running request for a waiting one | +| `--max-queued-requests`, `--max-queued-tokens` | no limit | requests, and prompt tokens, waiting per engine at most; past either a request is answered 503 | +| `--radix-eviction-policy` | lru | which cached prefix the prefix cache drops first: `lru` the least recently used; `lfu` the fewest hits; `slru` those with fewer than `protected_threshold` hits (2) before the rest; `priority` the lowest request priority; `tlru` (Tail-Optimized LRU, [arXiv 2510.15152](https://arxiv.org/abs/2510.15152)) first the tail a session's next turn could recompute within `threshold - next_prompt_estimate` tokens. Prefixes a session still references go last under every policy | +| `--radix-eviction-policy-config` | none | the policy's settings as JSON: `{"protected_threshold": N}` for `slru`; `{"threshold": N, "next_prompt_estimate": M}` for `tlru`, which needs both | + +Behind the router, a session stays on one engine, so its requests are ordered by that engine's policy. + ## Speculative decoding: the `mtp-*` presets The model's own multi-token-prediction (MTP) layer drafts 2 tokens a step and the target model verifies them in the next @@ -114,6 +134,33 @@ other than a preset's changes the graphs and needs its own capture. The `mtp-*` which changes the kernels' DMA form. A speculative engine serves temperature and top-p sampling as the plain engine does (verification is the rejection sampler's, so the output distribution is unchanged), and refuses the requests the plain engine refuses. +## Structured outputs, penalties and thinking budgets + +These need graphs that take inputs from the host, so an engine serves them only when started with the flag. Each flag adds +graphs to warmup and to the capture: capture with the flags you serve with. With both off, the graphs are meant to be +the ones an engine without the flags compiles; the graph-equivalence check of that, and a Trainium2 run of the variants, +are pending (they are tested on CPU). They are refused with speculation at every batch size (`mtp-always`): its verify graph +takes no mask or penalties. Under the other `mtp-*` presets the engine steps every request plain while one that uses +them runs. + +| Argument | Default | | +|---|---|---| +| `--structured-outputs` | off | guided decoding (needs the `xgrammar` package: `pip install 'trnrt[grammar]'`), `thinking_token_budget` and `min_tokens`: decode and prefill graphs that take a token mask | +| `--penalties` | off | repetition, frequency and presence penalties: decode and prefill graphs that take penalty inputs | +| `--penalty-buckets` | 1024,16384 | with `--penalties`: the distinct-token counts per request the penalty graphs are compiled at; the vocabulary size is always added as the last | + +Graphs each flag adds (decode graphs at the plain decode buckets times the page buckets; prefill graphs are the sampling +step only, one per prefill bucket): + +| Configuration | `--structured-outputs` | `--penalties` (3 rungs) | both | +|---|---|---|---| +| default (batch 12) | 18 decode, 4 prefill | 54 decode, 12 prefill | 72 decode, 16 prefill | +| `--preset batch8` | 15 decode, 4 prefill | 45 decode, 12 prefill | 60 decode, 16 prefill | +| `mtp-batch8`, `mtp-batch12` | 9 decode, 4 prefill | 27 decode, 12 prefill | 36 decode, 16 prefill | +| `mtp-batch24` | 6 decode, 4 prefill | 18 decode, 12 prefill | 24 decode, 16 prefill | + +These counts follow from the bucket ladders; the variants' device cost (compile time, HBM, step time) is not measured yet. + ## Supported scope The supported configurations are the default (batch 12), `--preset batch8`, the four `mtp-*` presets, and any of them @@ -135,5 +182,8 @@ Set them on the server's container (`docker run -e NAME=value ...`); its engines | `TRNRT_SKIP_DEVICE_CHECK` | unset | `1`: start without checking that the engine's Neuron devices are free | | `TRNRT_SKIP_HOST_MEM_CHECK` | unset | `1`: start without checking the host's free contiguous memory | | `TRNRT_LOG_LEVEL` | `INFO` | the level of the server's, engines' and router's log messages, on stderr (`DEBUG`, `INFO`, `WARNING`, `ERROR`) | +| `TRNRT_PROD_DIR` | unset | a directory: each engine writes every finished request's record (token counts, times, the finish reason; never a prompt or an output) and a 10-s gauges record there, under `engine-//`, rotated daily and at 1 GiB and gzipped; `/v1/stats`'s `today` window reads them back after a restart. Unset, nothing is written | +| `TRNRT_PROMPT_CACHE_MB` | unset | MiB of host memory per engine for the prompt token cache: a chat turn's prompt ids built from the session's previous turn and an encode of the new tail, instead of an encode of the whole prompt (195-239 ms for 97K-114K-token prompts, GLM-5.3-Flash's tokenizer on one host CPU core of a trn2.48xlarge). It returns exactly the ids of a full encode, or turns itself off: it checks itself at start, and audits its first 8 hits and every 64th after against a full render and encode | +| `TRNRT_PROMPT_CACHE_AUDIT` | 64 | with the prompt token cache, audit every this many hits after the first 8 (0: only the checks at start) | | `TRNRT_ROUTER_LOG_SESSIONS` | unset | `1`, for debugging: the router's log names session ids: a line per new pin and per side request, and the id in its failover and return-home lines; unset, it logs no session id | | `NEURON_LIBTORCH_CACHE_ROOT` | `~/.cache/neuron_libtorch` | the compile cache's root ([DEPLOY.md](DEPLOY.md), step 4) | diff --git a/pyproject.toml b/pyproject.toml index 6821207..d4697e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,8 @@ dependencies = [ [project.optional-dependencies] test = ["pytest>=8", "openai>=1.40", "ml_dtypes"] +# Guided decoding (--structured-outputs): the grammar engine. Optional: without it guided requests get a 400. +grammar = ["xgrammar>=0.1.19"] lint = ["ruff==0.17.3", "pre-commit>=4", "mypy==2.4.1"] dev = ["trnrt[test,lint]"] 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"} diff --git a/tests/test_engine_priority.py b/tests/test_engine_priority.py new file mode 100644 index 0000000..a318b00 --- /dev/null +++ b/tests/test_engine_priority.py @@ -0,0 +1,119 @@ +"""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 drain(eng, limit=8000): + """Step until the engine has no work, failing (not hanging) if it never drains.""" + steps = 0 + while eng.has_work(): + eng.step() + steps += 1 + assert steps < limit, "the engine never drained its work" + + +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() + drain(eng) + 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" + drain(eng) + 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 + drain(eng) + check_done(eng, [r]) + + +class SpyAdmit: + """The engine's host sampling, recording each request it set up (a grammar matcher, in the real one), and refusing + the ones `refuse` names as admit refuses a spec xgrammar cannot compile.""" + + def __init__(self, refuse=()): + self.admitted: list[str] = [] + self.refuse = set(refuse) + + def admit(self, req): + if req.rid in self.refuse: + raise ValueError("invalid regex for guided decoding") + self.admitted.append(req.rid) + + def advance(self, seqs, finish_now): + return [] + + +def test_a_request_the_queue_refuses_is_refused_before_admit_builds_its_matcher(): + eng = with_scheduling(make_engine(seats=2), priority_scheduling=True, max_queued_requests=2) + eng.host_sampling = spy = SpyAdmit(refuse={"bad"}) + 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() + low = [eng.add_request(p, sp, priority=0) for p in prompts(2, seed=6, shared=False)] + before = list(spy.admitted) + with pytest.raises(QueueFull): + eng.add_request([10, 11, 12], sp, rid="full", priority=0) + assert spy.admitted == before, "admit ran for a request the full queue refused" + with pytest.raises(ValueError, match="invalid regex"): + eng.add_request([10, 11, 12], sp, rid="bad", priority=9) # would displace, but admit refuses it first + assert all(r.status is Status.WAITING for r in low), "a refused request displaced a waiting one" + high = eng.add_request([10, 11, 12], sp, rid="high", priority=9) + assert spy.admitted == [*before, "high"] and low[1].error == (503, DISPLACED) + drain(eng) + check_done(eng, [*running, low[0], high]) diff --git a/tests/test_grammar_engine.py b/tests/test_grammar_engine.py new file mode 100644 index 0000000..3c77464 --- /dev/null +++ b/tests/test_grammar_engine.py @@ -0,0 +1,174 @@ +"""Guided decoding through xgrammar on the real engine (engine/grammar.py, engine/guided.py, engine/host_inputs.py): a +tiny random GLM-5.3-Flash (fp32, CPU) whose word-level tokenizer holds the pieces of a little JSON, so a grammar can +finish in a few tokens. The engine accepts every sampled token into the request's matcher and raises if one is refused, +so every test here also checks that no forbidden token was ever sampled. + +- a choice, a JSON schema and a regex come out exactly as constrained, greedy and sampled, ending on EOS ("stop"); +- a free JSON object is always a valid prefix of one, and parses when the grammar finished; +- under a reasoning block the prompt opened, the grammar constrains only after the reasoning's end (here forced by a + thinking budget); +- the requests decoding beside a guided one are unchanged; +- under dynamic speculation a guided request turns the engine to plain decode steps while it runs, and greedy output + of the others still equals plain greedy. +""" + +from __future__ import annotations + +import json +import re + +import pytest + +pytest.importorskip("transformers.models.glm5_next.modeling_glm5_next") +pytest.importorskip("xgrammar") + +from tests import test_glm5_next as T +from tests.test_host_sampling_engine import THINK, save_tokenizer +from trnrt.engine.grammar import GrammarSpec +from trnrt.engine.request import SamplingParams + +PIECES = ["{", "}", '"', ":", ",", "a", "b", "x", "yes", "no", "y", "[", "]", "0", "1", "true", *THINK] +WORDS = {w: 300 + i for i, w in enumerate(PIECES)} +BASE = {"max_num_seqs": 4, "max_model_len": 256, "num_pages": 256, "structured_outputs": True} +SCHEMA = {"type": "object", "properties": {"a": {"enum": ["x", "b"]}}, "required": ["a"], "additionalProperties": False} + + +def prompts(seed, lengths): + """Random prompts drawn below the word ids (300 on), so none holds a reasoning tag: a prompt that opened a reasoning + block would hold its grammar back until a close tag (guided.in_reasoning).""" + return T.prompts(seed, lengths, vocab=min(WORDS.values())) + + +@pytest.fixture(scope="module") +def model(tmp_path_factory): + path = str(tmp_path_factory.mktemp("glm5_next_grammar")) + T.build(path, seed=5) + save_tokenizer(path, 384, WORDS) + return path + + +def run(path, prompts, params, **kw): + eng = T.engine(path, **{**BASE, **kw}) + try: + reqs = eng.generate(prompts, params) + return eng, reqs + finally: + eng.close() + + +def text(eng, r) -> str: + return eng.tokenizer.decode(r.output_ids, skip_special_tokens=True) + + +@pytest.mark.parametrize("temperature", [0.0, 1.0]) +def test_a_choice_a_schema_and_a_regex_come_out_as_constrained(model, temperature): + ps = prompts(31, (9, 14, 7, 11)) + specs = [ + GrammarSpec.choice(["no", "a,b"]), + GrammarSpec("json_schema", json.dumps(SCHEMA)), + GrammarSpec("regex", "(ab|ba){1,2}x"), + GrammarSpec.choice(["true", "[1]"]), + ] + params = [SamplingParams(max_new_tokens=24, temperature=temperature, seed=3, grammar=g) for g in specs] + eng, got = run(model, ps, params) + outs = [text(eng, r) for r in got] + assert outs[0] in ("no", "a,b") and outs[3] in ("true", "[1]"), outs + assert json.loads(outs[1]) in ({"a": "x"}, {"a": "b"}), outs + assert re.fullmatch("(ab|ba){1,2}x", outs[2]), outs + for r in got: + assert r.finish_reason == "stop" and r.output_ids[-1] == 1 and r.matcher.is_terminated(), r.output_ids + + +def test_a_free_json_object_is_always_a_valid_prefix(model): + ps = prompts(32, (8, 12, 10)) + params = [ + SamplingParams( + max_new_tokens=30, temperature=1.0, seed=s, grammar=GrammarSpec("json_schema", '{"type": "object"}') + ) + for s in range(3) + ] + eng, got = run(model, ps, params) + for r in got: # the engine raises on a token its matcher refuses; a finished grammar parses + t = text(eng, r) + assert t.startswith("{"), t + if r.finish_reason == "stop": + assert isinstance(json.loads(t), dict), t + + +def test_under_reasoning_the_grammar_waits_for_the_reasoning_end(model): + open_, close = WORDS[THINK[0]], WORDS[THINK[1]] + ps = [[*p, open_] for p in prompts(33, (9, 12))] + g = GrammarSpec.choice(["no", "a,b"]) + params = [SamplingParams(max_new_tokens=24, temperature=0.0, thinking_token_budget=b, grammar=g) for b in (0, 4)] + eng, got = run(model, ps, params) + for r, b in zip(got, (0, 4)): + out = r.output_ids + assert out[b] == close and r.think_count == b, out + assert eng.tokenizer.decode(out[b + 1 :], skip_special_tokens=True) in ("no", "a,b"), out + assert r.finish_reason == "stop" + + +def test_the_requests_beside_a_guided_one_are_unchanged(model): + ps = prompts(34, (9, 14, 7)) + plain = SamplingParams(max_new_tokens=14, temperature=0.0, ignore_eos=True) + guided = SamplingParams(max_new_tokens=14, temperature=0.0, grammar=GrammarSpec("json_schema", json.dumps(SCHEMA))) + _, alone = run(model, ps, plain, structured_outputs=False) + eng, got = run(model, ps, [plain, guided, plain]) + assert [got[0].output_ids, got[2].output_ids] == [alone[0].output_ids, alone[2].output_ids] + assert json.loads(text(eng, got[1])) in ({"a": "x"}, {"a": "b"}) + + +@pytest.mark.parametrize("ignore_eos", [False, True]) +def test_a_grammar_the_vocabulary_cannot_finish_ends_the_request_not_the_engine(model, ignore_eos): + """xgrammar allows a token whose bytes start a match without knowing whether the vocabulary can finish it: "y" for + the regex "yz" (no token spells "z"). The request then ends on EOS ("stop"), also under ignore_eos (the commit does + not stop on that EOS: the dead end does), and the request beside it is unaffected.""" + ps = prompts(36, (9, 11)) + plain = SamplingParams(max_new_tokens=10, temperature=0.0, ignore_eos=True) + grammar = GrammarSpec("regex", "yz") + stuck = SamplingParams(max_new_tokens=10, temperature=0.0, ignore_eos=ignore_eos, grammar=grammar) + _, alone = run(model, ps, plain, structured_outputs=False) + eng, got = run(model, ps, [plain, stuck]) + assert got[0].output_ids == alone[0].output_ids + assert text(eng, got[1]) == "y" and got[1].finish_reason == "stop" and got[1].grammar_dead + assert got[1].output_ids == [WORDS["y"], 1], got[1].output_ids # one EOS, not one a step to the length + + +def test_a_malformed_grammar_is_refused_at_admission(model): + eng = T.engine(model, **BASE) + try: + with pytest.raises(ValueError, match="invalid regex"): + eng.add_request([5, 6], SamplingParams(grammar=GrammarSpec("regex", "(unclosed"))) + assert not eng.has_work() + finally: + eng.close() + + +@pytest.fixture(scope="module") +def mtp_model(tmp_path_factory): + from tests.test_spec_model import add_mtp + + path = str(tmp_path_factory.mktemp("glm5_next_grammar_mtp")) + T.build(path, seed=2, index_topk=16) + add_mtp(path) + save_tokenizer(path, 384, WORDS) + return path + + +def test_under_dynamic_speculation_a_guided_request_never_speculates(mtp_model, monkeypatch): + """Two requests, within spec_max_batch: the engine would speculate both, but runs plain while the guided one runs + (a verify takes no mask), then speculates the other alone; its greedy output equals plain greedy.""" + from trnrt.models import decoder + + monkeypatch.setattr(decoder, "MTP_CHUNK_Q", decoder.MTP_CHUNK_Q) # the engine writes its routes (test_spec_model) + monkeypatch.setattr(decoder, "DYN_MTP_KV", decoder.DYN_MTP_KV) + ps = prompts(35, (9, 23)) + plain = SamplingParams(max_new_tokens=30, temperature=0.0, ignore_eos=True) + guided = SamplingParams(max_new_tokens=20, temperature=0.0, grammar=GrammarSpec("json_schema", json.dumps(SCHEMA))) + kw = {"max_prefill_tokens": 16, "max_model_len": 512, "num_pages": 512} + _, want = run(mtp_model, ps[:1], plain, structured_outputs=False, **kw) + spec = {"spec_method": "mtp", "spec_k": 2, "spec_max_batch": 2, "dyn_mtp_kv": True} + eng, got = run(mtp_model, ps, [plain, guided], **kw, **spec) + assert got[0].output_ids == want[0].output_ids + assert json.loads(text(eng, got[1])) in ({"a": "x"}, {"a": "b"}) + assert eng.dyn_plain_steps > 0 and eng.dyn_spec_steps > 0, "plain while the guided request ran, then speculative" diff --git a/tests/test_guided_host.py b/tests/test_guided_host.py new file mode 100644 index 0000000..70b29f1 --- /dev/null +++ b/tests/test_guided_host.py @@ -0,0 +1,209 @@ +"""The host side of guided decoding and thinking budgets (engine/guided.py HostSampling) and the stop-string check under +min_tokens (engine/engine.py), driven token by token with a stand-in tokenizer and matcher: no model, no xgrammar. + +- a grammar's dead end ends its request as "stop", with or without ignore_eos; +- a thinking budget reached on a partial close tag forces only the rest of the tag; +- a stop string that ended among the first min_tokens tokens never ends the request later; +- the compiled-grammar cache keeps the most recently used grammars, CACHE_SIZE at most (a stand-in xgrammar). +""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest + +from trnrt.engine import grammar +from trnrt.engine.engine import LLMEngine +from trnrt.engine.grammar import GrammarBackend, GrammarSpec +from trnrt.engine.guided import HostSampling +from trnrt.engine.request import Request, SamplingParams, Status + +EOS = 1 +# "" is two ids and "" three, so a close tag can be cut short at the budget. +TAGS = {"": [10, 11], "": [12, 13, 14]} + + +def tokenizer(text, add_special_tokens=False): + return {"input_ids": TAGS[text]} + + +class DeadMatcher: + """A grammar at a dead end: it refuses every token and never terminates (host_inputs._mask_row then allows EOS).""" + + def __init__(self): + self.calls = 0 + + def is_terminated(self): + return False + + def accept_token(self, t): + self.calls += 1 + return False + + +def running(prompt, params, sampling: HostSampling | None = None) -> Request: + r = Request("r", prompt, params) + if sampling is not None: + sampling.admit(r) + r.status = Status.RUNNING + return r + + +def emit(r: Request, *ids: int) -> list: + r.token_ids.extend(ids) + return [SimpleNamespace(req=r, sample=True, draft=False)] + + +@pytest.mark.parametrize("ignore_eos", [False, True]) +def test_a_grammar_dead_end_ends_the_request_as_stop(ignore_eos): + hs = HostSampling(tokenizer, 64, (EOS,), runner=None) + r = running([5, 6], SamplingParams(ignore_eos=ignore_eos)) + r.matcher, r.grammar_dead = DeadMatcher(), True + finished = [] + + def finish_now(req, why): + finished.append((req, why)) + req.status = Status.FINISHED + + done = hs.advance(emit(r, EOS), finish_now) + assert done == [r] and finished == [(r, "stop")] + assert r.grammar_seen == 1 and r.matcher.calls == 1 + + +def test_a_dead_end_the_commit_already_stopped_is_not_finished_twice(): + hs = HostSampling(tokenizer, 64, (EOS,), runner=None) + r = running([5, 6], SamplingParams()) + r.matcher, r.grammar_dead = DeadMatcher(), True + r.status = Status.FINISHED # the commit stopped on the EOS (no ignore_eos) + assert hs.advance(emit(r, EOS), lambda req, why: pytest.fail("finished twice")) == [] + + +def test_a_token_its_grammar_refuses_without_a_dead_end_still_fails_loudly(): + hs = HostSampling(tokenizer, 64, (EOS,), runner=None) + r = running([5, 6], SamplingParams()) + r.matcher = DeadMatcher() # refuses, but no dead end was reached: the mask was not applied + with pytest.raises(RuntimeError, match="rejected by its own grammar"): + hs.advance(emit(r, EOS), lambda req, why: None) + + +def think(budget: int, *outs: list[int]) -> Request: + """A request whose prompt opened a reasoning block, fed `outs` one commit at a time.""" + hs = HostSampling(tokenizer, 64, (EOS,), runner=None) + r = running([5, *TAGS[""]], SamplingParams(thinking_token_budget=budget), hs) + assert r.think_in + for ids in outs: + hs.advance(emit(r, *ids), lambda req, why: None) + return r + + +def test_a_budget_reached_on_a_partial_close_tag_forces_only_the_rest_of_it(): + r = think(3, [20], [21], [12]) # the third reasoning token is the close tag's first id + assert r.think_force == [13, 14] and not r.think_done + r = think(3, [20], [21], [12], [13], [14]) # the forced rest, one token a step: the tag is whole once + assert r.think_done and not r.think_force and r.output_ids == [20, 21, 12, 13, 14] + + +def test_a_budget_reached_elsewhere_forces_the_whole_close_tag(): + r = think(2, [20], [21]) + assert r.think_force == [12, 13, 14] + r = think(3, [20], [12, 21]) # a close tag's first id that the next token showed was no tag + assert r.think_force == [12, 13, 14] + r = think(6, [20], [21], [22], [12, 13, 14]) # the model closes the block itself, under the budget + assert r.think_done and not r.think_force + + +class _Tok: + """Detokenizes each id to its own bracketed text, so a stop string names exactly one token.""" + + @staticmethod + def decode(ids): + return "".join(f"<{i}>" for i in ids) + + +class StopEngine: + """The engine's stop-string check and speculative trim, over _Tok, recording what it finishes.""" + + _check_stop_strings = LLMEngine._check_stop_strings + _stop_in = LLMEngine._stop_in + _trim_past_stop = LLMEngine._trim_past_stop + + def __init__(self): + self.tokenizer = _Tok() + self.finished: list[str] = [] + self.scheduler = SimpleNamespace(finish_now=lambda req, why: self.finished.append(why)) + + +def stop_check(r: Request, draft=False) -> list[Request]: + eng = StopEngine() + done = eng._check_stop_strings([SimpleNamespace(req=r, sample=True, draft=draft)]) + assert len(done) == len(eng.finished) and all(w == "stop" for w in eng.finished) + return done + + +def test_a_stop_string_among_the_first_min_tokens_never_ends_the_request(): + """vLLM checks stop strings from min_tokens on, over the newly decoded text only: a stop string that ended among the + protected tokens stays in the output and is not matched later.""" + r = running([5, 6], SamplingParams(stop=("<7>",), min_tokens=3)) + for t in (7, 8, 9, 20, 21): # the stop token is the first output token, protected + emit(r, t) + assert stop_check(r) == [], r.output_ids + emit(r, 7) # past min_tokens: it ends the request + assert stop_check(r) == [r] + + +def test_a_stop_string_completed_by_the_token_past_min_tokens_ends_the_request(): + r = running([5, 6], SamplingParams(stop=("<8><9>",), min_tokens=2)) + emit(r, 30, 8) + assert stop_check(r) == [] + emit(r, 9) # the stop string straddles the boundary and ends past it + assert stop_check(r) == [r] + + +def test_a_speculative_step_past_min_tokens_trims_to_its_own_stop_string_not_a_protected_one(): + r = running([5, 6], SamplingParams(stop=("<7>",), min_tokens=3)) + emit(r, 7, 8, 9) + assert stop_check(r) == [] + emit(r, 20, 7, 21) # one speculative step of three tokens: its stop token is the middle one + assert stop_check(r, draft=True) == [r] + assert r.output_ids == [7, 8, 9, 20, 7], "trimmed past the protected stop token" + + +class FakeXgr: + """The parts of xgrammar GrammarBackend calls, counting compilations.""" + + def __init__(self): + self.compiled: list[str] = [] + self.compiler_kw: dict = {} + outer = self + + class TokenizerInfo: + @staticmethod + def from_huggingface(tokenizer, vocab_size, stop_token_ids): + return "info" + + class GrammarCompiler: + def __init__(self, info, **kw): + outer.compiler_kw = kw + + def compile_regex(self, value): + outer.compiled.append(value) + return ("compiled", value) + + self.TokenizerInfo, self.GrammarCompiler = TokenizerInfo, GrammarCompiler + + +def test_the_compiled_grammar_cache_is_a_bounded_lru(monkeypatch): + xgr = FakeXgr() + monkeypatch.setitem(sys.modules, "xgrammar", xgr) + monkeypatch.setattr(grammar, "CACHE_SIZE", 3) + be = GrammarBackend(object(), 64, (EOS,)) + assert xgr.compiler_kw == {"cache_enabled": False}, "xgrammar's own unbounded cache stays off" + spec = [GrammarSpec("regex", f"r{i}") for i in range(5)] + for i in (0, 1, 2, 0, 3): # r0 used again before r3 arrives: r1 is the least recent, and goes + be.compile(spec[i]) + assert xgr.compiled == ["r0", "r1", "r2", "r3"] and len(be._cache) == 3 + be.compile(spec[0]) + be.compile(spec[1]) # recompiled: it was dropped + assert xgr.compiled == ["r0", "r1", "r2", "r3", "r1"] and len(be._cache) == 3 diff --git a/tests/test_host_sampling_engine.py b/tests/test_host_sampling_engine.py new file mode 100644 index 0000000..84e6b2e --- /dev/null +++ b/tests/test_host_sampling_engine.py @@ -0,0 +1,208 @@ +"""Host-input sampling on the real engine: penalties, min_tokens and thinking budgets on a tiny random GLM-5.3-Flash +(fp32, CPU), with the engine's own scheduler, overlap and graphs (engine/host_inputs.py, engine/guided.py). Guided +decoding through xgrammar is tests/test_grammar_engine.py. + +- with the flags off the engine refuses what needs them, and with them on but no request using them every call is the + graph (and every token the one) the engine without them gives; +- a penalized request follows vLLM's formula, step by step, against its own raw logprobs, while the requests beside it + decode exactly as alone; +- min_tokens keeps the end-of-sequence and stop ids out of the first min_tokens tokens; +- a thinking budget forces the reasoning's end at the budget. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("transformers.models.glm5_next.modeling_glm5_next") + +from tests import test_glm5_next as T +from trnrt.engine.request import SamplingParams + +THINK = ("", "") + + +def save_tokenizer(path: str, vocab: int, words: dict[str, int], eos: int = 1) -> None: + """A word-level tokenizer over the model's whole vocabulary: id i is the word t, except the ids `words` names, + decoded by joining the words with no space (tests/test_e2e.py's).""" + from tokenizers import Tokenizer, decoders, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast + + table = {f"t{i}": i for i in range(vocab)} + for w, i in words.items(): + table.pop(f"t{i}") + table[w] = i + tk = Tokenizer(models.WordLevel(table, unk_token="t0")) + tk.pre_tokenizer = pre_tokenizers.WhitespaceSplit() + tk.decoder = decoders.Fuse() + t = PreTrainedTokenizerFast(tokenizer_object=tk, unk_token="t0", eos_token=f"t{eos}") + t.save_pretrained(path) + + +@pytest.fixture(scope="module") +def model(tmp_path_factory): + """The tiny random GLM, with a tokenizer that has GLM's reasoning tags at ids 300 and 301.""" + path = str(tmp_path_factory.mktemp("glm5_next_host")) + T.build(path, seed=4) + save_tokenizer(path, 384, {THINK[0]: 300, THINK[1]: 301}) + return path + + +BASE = {"max_num_seqs": 4, "max_model_len": 256, "num_pages": 256} +FLAGS = {"structured_outputs": True, "penalties": True, "penalty_buckets": (16, 64)} + + +def run(path, prompts, params, **kw): + eng = T.engine(path, **{**BASE, **kw}) + try: + reqs = eng.generate(prompts, params) + return reqs, set(eng.runner.calls) + finally: + eng.close() + + +def test_the_flags_off_engine_refuses_what_needs_them(model): + eng = T.engine(model, **BASE) + try: + for p, flag in ( + (SamplingParams(presence_penalty=0.3), "--penalties"), + (SamplingParams(min_tokens=2), "--structured-outputs"), + (SamplingParams(thinking_token_budget=4), "--structured-outputs"), + ): + with pytest.raises(ValueError, match=flag): + eng.add_request([5, 6, 7], p) + assert not eng.has_work() + finally: + eng.close() + + +def test_with_the_flags_on_and_no_request_using_them_every_call_is_todays(model): + """The same tokens, greedy and sampled, and the same call keys: no variant was launched.""" + ps = T.prompts(21, (9, 17, 12)) + params = [ + SamplingParams(max_new_tokens=12, temperature=0.0, ignore_eos=True), + SamplingParams(max_new_tokens=12, temperature=0.9, top_p=0.9, seed=7, ignore_eos=True), + SamplingParams(max_new_tokens=12, temperature=0.0, ignore_eos=True, logprobs=3), + ] + want, keys_off = run(model, ps, params) + got, keys_on = run(model, ps, params, **FLAGS) + assert [r.output_ids for r in got] == [r.output_ids for r in want] + assert [r.logprobs for r in got] == [r.logprobs for r in want] + assert keys_on == keys_off and all(len(k) == 3 for k in keys_on if k[0] in ("decode", "prefill")), sorted(keys_on) + + +def test_a_penalized_request_follows_the_formula_beside_unchanged_neighbours(model): + """Greedy with frequency and presence penalties: each token is the argmax of its raw logprobs (logprobs=20, before + penalties) minus vLLM's penalties over the output so far, and the plain requests decoding beside it (every other + step it waits for its token) produce what they produce alone.""" + ps = T.prompts(22, (11, 8, 15)) + plain = SamplingParams(max_new_tokens=16, temperature=0.0, ignore_eos=True) + pen = SamplingParams( + max_new_tokens=16, temperature=0.0, ignore_eos=True, logprobs=20, frequency_penalty=1.5, presence_penalty=1.5 + ) + alone, _ = run(model, ps, [plain, plain, plain]) + got, keys = run(model, ps, [plain, pen, plain], **FLAGS) + assert [got[0].output_ids, got[2].output_ids] == [alone[0].output_ids, alone[2].output_ids] + out = got[1].output_ids + assert out != alone[1].output_ids, "the penalties changed nothing: the case does not test them" + for i, (tok, (_, ids, lps)) in enumerate(zip(out, got[1].logprobs)): + seen = out[:i] + score = {t: lp - 1.5 * seen.count(t) - 1.5 * (t in seen) for t, lp in zip(ids, lps)} + assert tok in score, f"step {i}: the penalized argmax fell outside the raw top 20" + assert tok == max(score, key=score.__getitem__), (i, tok) + tags = {k[3:] for k in keys if len(k) > 3} + assert ("grammar", "pen16") in tags, sorted(keys) + + +def test_a_repetition_penalty_counts_the_prompt(model): + """repetition_penalty divides positive and multiplies negative logits of every token in the prompt or the output: at + a large value the greedy output never takes such a token again (the unpenalized output does: the control).""" + (p,) = T.prompts(23, (24,)) + free, _ = run(model, [p], SamplingParams(max_new_tokens=24, temperature=0.0, ignore_eos=True)) + out = free[0].output_ids + assert any(t in p or t in out[:i] for i, t in enumerate(out)), "control: the free output repeats" + got, _ = run( + model, + [p], + SamplingParams(max_new_tokens=24, temperature=0.0, ignore_eos=True, repetition_penalty=60.0), + penalties=True, + ) + out = got[0].output_ids + assert len(out) == 24 and not [t for i, t in enumerate(out) if t in p or t in out[:i]] + + +def test_min_tokens_keeps_eos_and_stop_ids_out_of_the_first_tokens(model): + ps = T.prompts(24, (10, 13)) + free, _ = run(model, ps, SamplingParams(max_new_tokens=30, temperature=0.0, ignore_eos=True)) + # The stop ids: the free output's first tokens, so a request that may stop on them would stop at once. + stops = tuple({free[0].output_ids[0], free[1].output_ids[0], free[0].output_ids[3]}) + eos = 1 + params = SamplingParams(max_new_tokens=30, temperature=0.0, min_tokens=8, stop_token_ids=stops) + got, keys = run(model, ps, params, structured_outputs=True) + for r in got: + assert not set(r.output_ids[:8]) & {*stops, eos}, r.output_ids + assert len(r.output_ids) >= 8 + first = SamplingParams(max_new_tokens=30, temperature=0.0, stop_token_ids=stops) + ctl, _ = run(model, ps, first, structured_outputs=True) + assert len(ctl[0].output_ids) == 1, "control: without min_tokens the first stop id ends it at once" + assert any(k[3:] == ("grammar",) for k in keys) + + +def test_a_thinking_budget_forces_the_reasoning_end_at_the_budget(model): + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(model) + close = tok.convert_tokens_to_ids(THINK[1]) + ps = [[*p, 300] for p in T.prompts(25, (9, 14, 6))] # each prompt opens the reasoning block + free, _ = run(model, ps, SamplingParams(max_new_tokens=20, temperature=0.0, ignore_eos=True)) + assert all(close not in r.output_ids[:12] for r in free), "control: the free model does not close it itself" + budgets = [0, 3, 7] + params = [ + SamplingParams(max_new_tokens=20, temperature=0.0, ignore_eos=True, thinking_token_budget=b) for b in budgets + ] + plain = SamplingParams(max_new_tokens=20, temperature=0.0, ignore_eos=True) + got, _ = run(model, [*ps, ps[0]], [*params, plain], structured_outputs=True) + for r, b, f in zip(got, budgets, free): + assert r.output_ids[:b] == f.output_ids[:b], "the reasoning before the budget is the free model's" + assert r.output_ids[b] == close and r.think_done and r.think_count == b, (b, r.output_ids) + assert got[3].output_ids == free[0].output_ids, "the request beside them is unchanged" + + +def test_a_static_speculative_engine_refuses_the_flags(): + from trnrt.config import EngineConfig + + for flag in ("structured_outputs", "penalties"): + with pytest.raises(ValueError, match="spec_max_batch"): + EngineConfig(model_path="m", spec_method="mtp", **{flag: True}) + EngineConfig(model_path="m", spec_method="mtp", spec_max_batch=2, **{flag: True}) # dynamic: served + + +def test_a_host_bound_row_never_launches_with_its_newest_token_in_flight(model): + """The ordering rule (engine/scheduler.py): a decode of a penalized request carries its newest token from the host, + never the board's placeholder (its penalties would miss that token), while the plain requests beside it keep the + overlap: some of their decodes take the token from the board.""" + from trnrt.engine.scheduler import PENDING + + eng = T.engine(model, **BASE, **FLAGS) + seen = {"bound": 0, "overlap": 0} + orig = eng.runner.decode + + def spy(seqs): + for s in seqs: + inflight = s.req.token_ids[s.start] == PENDING + if s.req.params.has_penalties: + assert not inflight, f"{s.req.rid} launched with its newest token in flight" + seen["bound"] += 1 + else: + seen["overlap"] += inflight + return orig(seqs) + + eng.runner.decode = spy + try: + ps = T.prompts(26, (9, 12, 7)) + plain = SamplingParams(max_new_tokens=12, temperature=0.0, ignore_eos=True) + pen = SamplingParams(max_new_tokens=12, temperature=0.0, ignore_eos=True, presence_penalty=0.5) + eng.generate(ps, [plain, pen, plain]) + finally: + eng.close() + assert seen["bound"] >= 11 and seen["overlap"] > 0, seen diff --git a/tests/test_host_sampling_units.py b/tests/test_host_sampling_units.py new file mode 100644 index 0000000..e71de6d --- /dev/null +++ b/tests/test_host_sampling_units.py @@ -0,0 +1,269 @@ +"""The pieces of host-input sampling that need no model: the packed token mask and the sampler's masked, penalized draw +(engine/sampler.py shaped_logits), the per-call inputs and variant tags (engine/host_inputs.py), a request's host-bound +rules (engine/request.py), and the API's constraint fields (serve/structured.py).""" + +from __future__ import annotations + +import json +import types + +import numpy as np +import pytest +import torch + +from trnrt.engine import grammar, host_inputs +from trnrt.engine.grammar import GrammarSpec +from trnrt.engine.request import Request, SamplingParams, unsupported +from trnrt.engine.sampler import NUM_TOP_LOGPROBS, sample, shaped_logits, unpack_bitmask +from trnrt.serve.structured import OBJECT, grammar_spec + +V = 300 +N = NUM_TOP_LOGPROBS + + +def _draw(logits, temperature, seed, bitmask=None, penalties=None, k=64): + B = logits.shape[0] + g = torch.Generator().manual_seed(seed) + return sample( + logits, + torch.full((B,), float(temperature)), + torch.full((B,), 0.95), + torch.zeros(B, dtype=torch.int64), + torch.zeros(B), + torch.rand(B, k, generator=g), + bitmask, + penalties, + ) + + +def test_the_packed_mask_round_trips_through_the_graphs_unpacking(): + m = grammar.full_mask(3, V) + assert m.shape == (3, grammar.words(V)) and unpack_bitmask(torch.from_numpy(m), V).all() + grammar.allow_only(m[0], 37) + grammar.forbid(m[1], [0, 31, 32, 299]) + want = np.ones((3, V), bool) + want[0] = False + want[0, 37] = True + want[1, [0, 31, 32, 299]] = False + assert np.array_equal(unpack_bitmask(torch.from_numpy(m), V).numpy(), want) + assert grammar.allowed(m[0], 37) and not grammar.allowed(m[0], 36) and not grammar.allowed(m[1], 31) + + +@pytest.mark.parametrize("temperature", [0.0, 0.7, 1.5]) +def test_a_masked_draw_never_takes_a_forbidden_token_and_reports_raw_logprobs(temperature): + g = torch.Generator().manual_seed(1) + logits = torch.randn(6, V, generator=g) * 3 + mask = grammar.full_mask(6, V) + rng = np.random.default_rng(0) + allowed = [] + for i in range(6): + keep = rng.choice(V, size=[1, 3, 40, 200, 299, 5][i], replace=False) + mask[i] = 0 + for t in keep: + mask[i].view(np.uint32)[t // 32] |= np.uint32(1 << (t % 32)) + allowed.append(set(keep.tolist())) + raw = _draw(logits, temperature, 0) + for seed in range(20): + out = _draw(logits, temperature, seed, torch.from_numpy(mask)) + tok = out[:, 0].long().tolist() + assert all(t in allowed[i] for i, t in enumerate(tok)), (seed, tok) + lse = torch.logsumexp(logits, -1) + assert torch.allclose(out[:, 1], logits[torch.arange(6), out[:, 0].long()] - lse, atol=1e-5) + assert torch.equal(out[:, 2:], raw[:, 2:]), "top logprobs are the raw (unmasked) ones" + if temperature == 0.0: # greedy: the best allowed token + best = [max(allowed[i], key=lambda t, i=i: float(logits[i, t])) for i in range(6)] + assert _draw(logits, 0.0, 0, torch.from_numpy(mask))[:, 0].long().tolist() == best + + +def _reference_penalized(logits, prompt, out, rep, freq, pres): + """vLLM's penalties, written out per token: repetition over prompt and output tokens, then frequency and presence + over the output alone.""" + x = logits.clone() + counts = {t: out.count(t) for t in set(out)} + for t in set(prompt) | set(out) if rep != 1.0 else set(out): + if rep != 1.0: + x[t] = x[t] / rep if x[t] > 0 else x[t] * rep + c = counts.get(t, 0) + x[t] -= freq * c + pres * (c > 0) + return x + + +@pytest.mark.parametrize(("rep", "freq", "pres"), [(2.0, 0.0, 0.0), (1.0, 1.5, 0.0), (1.0, 0.0, 3.0), (1.5, -0.4, 2.0)]) +def test_penalties_shape_the_draw_as_the_reference_formula_says(rep, freq, pres): + g = torch.Generator().manual_seed(2) + logits = torch.randn(4, V, generator=g) * 2 + rng = np.random.default_rng(3) + reqs = [] + for i in range(4): + prompt = rng.integers(0, V, 30 + i).tolist() + top = np.argsort(-logits[i].numpy()) + out = [*top[:3].tolist() * 2, *rng.choice(top[3:40], 4).tolist()] # the best three twice, then a few more + r = Request( + f"r{i}", + prompt, + SamplingParams(repetition_penalty=rep, frequency_penalty=freq, presence_penalty=pres), + ) + r.token_ids += out + reqs.append((r, prompt, out)) + m = 256 + pen = tuple(torch.from_numpy(a) for a in host_inputs.penalty_inputs([r for r, _, _ in reqs], m)) + got = _draw(logits, 0.0, 0, penalties=pen)[:, 0].long().tolist() + want = [int(_reference_penalized(logits[i], p, o, rep, freq, pres).argmax()) for i, (_, p, o) in enumerate(reqs)] + assert got == want, "MISMATCH" + unpen = logits.argmax(-1).tolist() + assert all(g != u for g, u in zip(got, unpen)), "NOCHANGE: the penalties left a row's best token: no test of them" + + +def test_a_penalty_never_brings_a_forbidden_token_back(): + logits = torch.zeros(1, V) + logits[0, 5] = 1.0 + mask = grammar.full_mask(1, V) + grammar.forbid(mask[0], [7]) + r = Request("r", [1, 2], SamplingParams(frequency_penalty=-100.0)) + r.token_ids += [7, 7] # a negative frequency penalty lifts token 7 by 200 + pen = tuple(torch.from_numpy(a) for a in host_inputs.penalty_inputs([r], 8)) + shaped = shaped_logits(logits, None, pen) + assert int(shaped.argmax()) == 7, "control: the penalty alone picks 7" + assert int(_draw(logits, 0.0, 0, torch.from_numpy(mask), pen)[0, 0]) == 5 + + +def _ecfg(structured=False, penalties=False, buckets=(16, 64)): + return types.SimpleNamespace(structured_outputs=structured, penalties=penalties, penalty_buckets=buckets) + + +def test_variant_tags_follow_the_flags(): + assert host_inputs.variant_tags(_ecfg(), 100) == [] + assert host_inputs.variant_tags(_ecfg(True), 100) == [("grammar",)] + assert host_inputs.variant_tags(_ecfg(False, True), 100) == [("pen16",), ("pen64",), ("pen100",)] + assert host_inputs.variant_tags(_ecfg(True, True, (16, 500)), 100) == [ + ("grammar",), + ("grammar", "pen16"), + ("grammar", "pen100"), + ] + + +def test_a_call_takes_a_variant_only_when_a_row_needs_it(): + plain = Request("a", [1, 2, 3], SamplingParams()) + pen = Request("b", [1, 2, 3], SamplingParams(presence_penalty=0.5)) + mt = Request("c", [4, 5], SamplingParams(min_tokens=2, stop_token_ids=(9,))) + eos = (1,) + assert host_inputs.extras(_ecfg(True, True), 100, eos, None, [plain, None]) == ([], ()) + extra, tag = host_inputs.extras(_ecfg(True, False), 100, eos, None, [plain, mt, None]) + assert tag == ("grammar",) and extra[1] is None + rows = unpack_bitmask(torch.from_numpy(extra[0]), 100) + assert rows[0].all() and rows[2].all() and not rows[1, 1] and not rows[1, 9] and int(rows[1].sum()) == 98 + mt.token_ids += [7, 8] # min_tokens reached: no mask any more + assert host_inputs.extras(_ecfg(True, False), 100, eos, None, [plain, mt]) == ([], ()) + extra, tag = host_inputs.extras(_ecfg(False, True), 100, eos, None, [plain, pen]) + assert tag == ("pen16",) and extra[0] is None + extra, tag = host_inputs.extras(_ecfg(True, True), 100, eos, None, [plain, pen]) + assert tag == ("grammar", "pen16") and unpack_bitmask(torch.from_numpy(extra[0]), 100).all() + many = Request("d", list(range(40)), SamplingParams(repetition_penalty=1.2)) + assert host_inputs.extras(_ecfg(False, True), 100, eos, None, [many])[1] == ("pen64",) + many.token_ids += list(range(40, 90)) + assert host_inputs.extras(_ecfg(False, True), 100, eos, None, [many])[1] == ("pen100",) + + +def test_min_tokens_keeps_eos_under_ignore_eos_and_counts_tokens_in_flight(): + r = Request("r", [4, 5], SamplingParams(min_tokens=3, ignore_eos=True, stop_token_ids=(9,))) + extra, _ = host_inputs.extras(_ecfg(True), 100, (1,), None, [r]) + row = unpack_bitmask(torch.from_numpy(extra[0]), 100)[0] + assert row[1] and not row[9] + r.token_ids += [6, -1] # one known, one in flight: the next token is output index 2, still below 3 + assert host_inputs.extras(_ecfg(True), 100, (1,), None, [r])[1] == ("grammar",) + r.token_ids.append(-1) + assert host_inputs.extras(_ecfg(True), 100, (1,), None, [r]) == ([], ()) + + +def test_a_forced_token_is_the_only_one_allowed(): + r = Request("r", [4, 5], SamplingParams(thinking_token_budget=2, min_tokens=5)) + r.think_force = [42, 43] + extra, tag = host_inputs.extras(_ecfg(True), 100, (42,), None, [r]) + row = unpack_bitmask(torch.from_numpy(extra[0]), 100)[0] + assert tag == ("grammar",) and row.nonzero().flatten().tolist() == [42] + + +def test_host_bound_rows_and_rows_that_never_speculate(): + assert not Request("a", [1], SamplingParams()).host_bound + assert Request("a", [1], SamplingParams(frequency_penalty=0.1)).host_bound + mt = Request("a", [1], SamplingParams(min_tokens=2)) + assert not mt.host_bound and mt.needs_host_inputs, "min_tokens depends on the count alone, not the tokens" + b = Request("a", [1], SamplingParams(thinking_token_budget=10)) + b.think_in = True + assert not b.host_bound and b.needs_host_inputs + b.think_count = 8 + assert not b.host_bound + b.think_count = 9 + assert b.host_bound, "one reasoning token left: the next one may be the last before the force" + b.think_done = True + assert not b.host_bound and not b.needs_host_inputs + m = types.SimpleNamespace(done=False, is_terminated=lambda: m.done) + g = Request("a", [1], SamplingParams(grammar=GrammarSpec("regex", "a"))) + g.matcher = m + assert g.host_bound and g.needs_host_inputs + m.done = True + assert not g.host_bound and not g.needs_host_inputs + + +def test_the_engine_refuses_what_it_was_not_started_with(): + p = SamplingParams(grammar=GrammarSpec("regex", "a")) + assert "--structured-outputs" in unsupported(p) + assert unsupported(p, structured=True) is None + for q in (SamplingParams(min_tokens=1), SamplingParams(thinking_token_budget=0)): + assert "--structured-outputs" in unsupported(q) and unsupported(q, structured=True) is None + pen = SamplingParams(repetition_penalty=1.1) + assert "--penalties" in unsupported(pen, structured=True) and unsupported(pen, penalties=True) is None + assert unsupported(SamplingParams(repetition_penalty=1.0, frequency_penalty=0.0, presence_penalty=0.0)) is None + + +@pytest.mark.parametrize( + ("body", "want"), + [ + ({}, None), + ({"response_format": {"type": "text"}}, None), + ({"response_format": {"type": "json_object"}}, GrammarSpec("json_schema", OBJECT)), + ( + {"response_format": {"type": "json_schema", "json_schema": {"name": "x", "schema": {"type": "integer"}}}}, + GrammarSpec("json_schema", '{"type": "integer"}'), + ), + ( + {"response_format": {"type": "json_schema", "schema": {"type": "null"}}}, + GrammarSpec("json_schema", '{"type": "null"}'), + ), + ({"guided_json": {"type": "string"}}, GrammarSpec("json_schema", '{"type": "string"}')), + ({"guided_json": '{"type": "string"}'}, GrammarSpec("json_schema", '{"type": "string"}')), + ({"guided_regex": "[a-z]+"}, GrammarSpec("regex", "[a-z]+")), + ({"guided_choice": ["a.b", "c"]}, GrammarSpec("regex", r"(a\.b|c)")), + ({"guided_grammar": 'root ::= "x"'}, GrammarSpec("ebnf", 'root ::= "x"')), + ({"structured_outputs": {"regex": "x+"}}, GrammarSpec("regex", "x+")), + ({"structured_outputs": {"json_object": True}}, GrammarSpec("json_schema", OBJECT)), + ({"json_schema": '{"type": "boolean"}'}, GrammarSpec("json_schema", '{"type": "boolean"}')), + ({"regex": "y"}, GrammarSpec("regex", "y")), + ({"ebnf": 'root ::= "z"'}, GrammarSpec("ebnf", 'root ::= "z"')), + ], +) +def test_every_constraint_field_maps_to_one_grammar(body, want): + got = grammar_spec(body) + assert got == want + if got is not None and got.kind == "json_schema": + json.loads(got.value) + + +@pytest.mark.parametrize( + "body", + [ + {"response_format": {"type": "json_object"}, "regex": "x"}, + {"guided_json": {}, "guided_regex": "x"}, + {"response_format": {"type": "structural_tag"}}, + {"response_format": {"type": "json_schema"}}, + {"response_format": "json"}, + {"guided_choice": []}, + {"guided_choice": ["a", 1]}, + {"guided_regex": ""}, + {"json_schema": "{not json"}, + {"structured_outputs": {"structural_tag": "x"}}, + ], +) +def test_a_malformed_or_doubled_constraint_is_refused(body): + with pytest.raises(ValueError): # noqa: PT011 the message varies with the field; the API answers any with 400 + grammar_spec(body) diff --git a/tests/test_host_sampling_warmup.py b/tests/test_host_sampling_warmup.py new file mode 100644 index 0000000..0683381 --- /dev/null +++ b/tests/test_host_sampling_warmup.py @@ -0,0 +1,85 @@ +"""The graphs --structured-outputs and --penalties add to warmup (and so to a capture), per shipped preset, on +tests/test_spec_runner_stub.py's recording stub model: the variants exactly at the plain decode family's buckets and page +rungs and at every prefill bucket and rung, none beside a verify, and nothing at all without the flags.""" + +from __future__ import annotations + +import types + +import pytest +import torch + +from tests import test_spec_runner_stub as S +from trnrt import __main__ as cli +from trnrt.config import prefill_pages +from trnrt.engine import model_runner as mr + +VOCAB = 154_880 # GLM-5.3-Flash's +PRESETS = ["batch12", "batch8", "mtp-batch8", "mtp-batch12", "mtp-batch24"] + + +def build(preset, monkeypatch, *extra): + ecfg = cli.engine_config(["--model", "m", "--preset", preset, "--device", "cpu", "--tp", "1", *extra]) + log: list = [] + monkeypatch.setattr(mr, "_piecewise", lambda *a, **k: S._recorders(log)) + mcfg = types.SimpleNamespace(vocab_size=VOCAB, eos_token_ids=(1,)) + return mr.ModelRunner(S.StubModel(), mcfg, ecfg, 8, torch.device("cpu")), ecfg + + +def tagged(runner): + """The warmed keys that carry a variant tag.""" + return {k for k in runner.compile_seconds if any(isinstance(x, str) for x in k[1:])} + + +def expected(ecfg, tags): + dec = ecfg.resolved_decode_batch_buckets() + plain = [b for b in dec if not ecfg.plain_bucket_min or b >= ecfg.plain_bucket_min] + out = {("decode", B, P, *t) for B in plain for P in ecfg.resolved_page_buckets() for t in tags} + for C in ecfg.resolved_prefill_token_buckets(): + for P in prefill_pages(ecfg.resolved_prefill_page_buckets(), ecfg.prefill_page_min_rows, C): + out |= {("prefill", C, P, *t) for t in tags} + return out + + +@pytest.mark.parametrize("preset", PRESETS) +def test_without_the_flags_warmup_adds_nothing(preset, monkeypatch): + runner, _ = build(preset, monkeypatch) + runner.warmup() + assert tagged(runner) == set() + + +@pytest.mark.parametrize("preset", PRESETS) +@pytest.mark.parametrize( + ("flags", "tags"), + [ + (["--structured-outputs"], [("grammar",)]), + (["--penalties"], [("pen1024",), ("pen16384",), (f"pen{VOCAB}",)]), + ( + ["--structured-outputs", "--penalties"], + [("grammar",), ("grammar", "pen1024"), ("grammar", "pen16384"), ("grammar", f"pen{VOCAB}")], + ), + ], +) +def test_the_flags_warm_exactly_their_variants(preset, flags, tags, monkeypatch): + runner, ecfg = build(preset, monkeypatch, *flags) + runner.warmup() + got, want = tagged(runner), expected(ecfg, tags) + assert got == want, (sorted(got - want)[:4], sorted(want - got)[:4]) + assert not [k for k in got if k[0] == "verify"] + + +def test_batch12_counts(monkeypatch): + """Literal counts for the served preset with --structured-outputs: decode buckets 1 2 4 6 8 12 x pages 512 1024 4096, + and prefill 512 1024 (the 4096-page rung only, below 2048 rows) and 2048 4096 (both rungs).""" + runner, _ = build("batch12", monkeypatch, "--structured-outputs") + runner.warmup() + keys = tagged(runner) + assert len([k for k in keys if k[0] == "decode"]) == 18 + assert sorted(k[1:3] for k in keys if k[0] == "prefill") == [ + (512, 4096), + (1024, 4096), + (2048, 1024), + (2048, 4096), + (4096, 1024), + (4096, 4096), + ] diff --git a/tests/test_prompt_cache.py b/tests/test_prompt_cache.py new file mode 100644 index 0000000..25d44df --- /dev/null +++ b/tests/test_prompt_cache.py @@ -0,0 +1,493 @@ +"""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, normalize or match single words" in PC.PromptCache(strip, 1 << 20).off + # a whole-word token matches by what follows it: cached "x foo bar", a cut at "foo" would keep "x " apart from + # it, while in "x foobar" no token matches and " f" merges. The self-check does not catch it; the guard must. + word = _tokenizer(TAGS[:4]) + word.add_tokens([AddedToken("foo", single_word=True, normalized=False)]) + assert "match single words: ['foo']" in PC.PromptCache(word, 1 << 20).off + assert PC.PromptCache(word, 1 << 20).encode("x foobar", "s") == full(word, "x foobar") + 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 diff --git a/tests/test_radix_eviction.py b/tests/test_radix_eviction.py new file mode 100644 index 0000000..e376210 --- /dev/null +++ b/tests/test_radix_eviction.py @@ -0,0 +1,168 @@ +"""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_and_a_sibling_does_not(): + 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 + # Another branch under it: the evicted sequence's next turn must still recompute its 4 tokens. + rc.insert([3] * 4 + [6] * 4, [parent.pages[0], *pool.alloc(1)]) + assert parent.tail_evicted == 4 + rc.insert([3] * 4 + [4] * 4, [parent.pages[0], *pool.alloc(1)]) # the evicted sequence itself, back + assert parent.tail_evicted == 0 + + +def test_a_sibling_insert_keeps_the_shared_prefix_unsafe(): + """Budget 8 tokens. A 4-token shared prefix whose 8-token branch was evicted must not look safe (4 + 8 > 8) once a + short sibling branch is inserted and evicted under it: it outlives a newer safe prefix.""" + pool = PagePool(64) + rc = RadixCache(pool, PS, "tlru", {"threshold": 10, "next_prompt_estimate": 2}) + shared = [3] * 4 + rc.insert(shared, pool.alloc(1)) + page = rc.match_prefix(shared).pages[0] + rc.insert(shared + [4] * 8, [page, *pool.alloc(2)]) + rc.evict(1) # the 8-token branch + rc.insert(shared + [6] * 4, [page, *pool.alloc(1)]) + rc.insert([2] * 8, pool.alloc(2)) # newer, 8 tokens: safe + rc.evict(1) # the sibling (4 tokens below an evicted 0: safe, the oldest); a match would refresh the prefix + assert rc.total_pages() == 3 + rc.evict(1) # the shared prefix (4 + 8 evicted, unsafe) outlives the newer safe one + assert rc.match_prefix(shared).num_pages == 1 and rc.match_prefix([2] * 8).num_pages == 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) diff --git a/tests/test_router.py b/tests/test_router.py index 1e55155..a8700c8 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,9 @@ 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", "/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 @@ -277,7 +287,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) @@ -493,3 +539,137 @@ 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_a_session_key_reads_only_the_start_of_a_long_prompt(monkeypatch): + """A session request's key hashes the first SESSION_KEY_CHARS characters of its prompt's JSON text, on the event + loop: written from the prompt's start only, never the whole prompt (a 200K-token input_ids list), with the same key + as the whole prompt's text cut there.""" + import hashlib + import json + + from trnrt.serve import router as R + + K = R.SESSION_KEY_CHARS + h = lambda b: hashlib.sha1(b.encode("utf-8", "replace"), usedforsecurity=False).hexdigest() # noqa: E731 + ids = list(range(200_000)) + prompts = [ + ("/generate", {"input_ids": ids}, ids), + ("/generate", {"input_ids": [ids, ids]}, [ids, ids]), + ("/generate", {"input_ids": [[7] * 3, ids]}, [[7] * 3, ids]), # a short first prompt: the next one is read too + ("/generate", {"text": ["x" * 100_000, "y"]}, ["x" * 100_000, "y"]), + ("/v1/completions", {"prompt": ["\u00e9\U0001f642" * 50_000]}, ["\u00e9\U0001f642" * 50_000]), + ] + want = [(h(""), h(json.dumps(p)[:K])) for _, _, p in prompts] + content = [{"type": "text", "text": "z" * 100_000}, {"type": "text", "text": "w"}] + chat = {"messages": [{"role": "system", "content": "s"}, {"role": "user", "content": content}]} + want_chat = h("user\x00" + json.dumps(content, ensure_ascii=False)[:K]) + sizes: list[int] = [] + real = json.dumps + monkeypatch.setattr(R.json, "dumps", lambda *a, **kw: sizes.append(len(out := real(*a, **kw))) or out) + assert [R._session_key(path, body) for path, body, _ in prompts] == want + assert R._session_key("/v1/chat/completions", chat)[1] == want_chat + assert max(sizes) < 16 * K # a cut prompt's text (up to 12 characters an item): never the 1.3 MB of the list + + +@pytest.mark.parametrize("raw", ["\u00b2", "\u0663", "1" * 5000, "2", "-1", "x", ""]) +def test_an_engine_index_that_is_not_one_is_a_400(raw): + """?engine= takes an ASCII integer below the engine count: a superscript or Arabic-Indic digit (str.isdigit's, not + int's), or more digits than int reads, is a 400 like any other bad index, never a 500.""" + from fastapi.testclient import TestClient + + from trnrt.serve.router import build_router_app + + with TestClient(build_router_app(Router(["http://127.0.0.1:9", "http://127.0.0.1:9"]), poll_s=0)) as c: + for path in ("/v1/stats", "/dashboard"): + r = c.get(path, params={"engine": raw}, follow_redirects=False) + assert r.status_code == 400 and "engine must be an integer in [0, 2)" in r.text, (path, r.text) + + +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()) diff --git a/tests/test_router_native.py b/tests/test_router_native.py new file mode 100644 index 0000000..b1d23aa --- /dev/null +++ b/tests/test_router_native.py @@ -0,0 +1,254 @@ +"""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, flush_s: float = 0.0): + """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)) + await asyncio.sleep(flush_s) + hold = getattr(app.state, "hold", None) # an asyncio.Event a test sets: a flush stuck until then + if hold is not None: + await hold.wait() + 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(): + """A worker's 503 (no step of its engine ended in time: maybe a long prefill on an engine still serving) names it + down in the answer, and counts toward down_after as a stalled /health does: it is marked down only after that many + in a row, so one slow probe moves no session pinned there.""" + + 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, True] and router.fails == [0, 1] + for _ in range(router.down_after - 1): + await c.get("/health_generate") + assert 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) + + +def test_flush_cache_bounds_its_wait_by_the_timeout_it_passes_on(): + """An engine answers /flush_cache once it flushed or its `timeout` passed, so one that has not answered timeout + + health_timeout_s seconds after it was asked is stuck: the router answers 400 naming it ("no answer") instead of + waiting on it for ever, and still waits out a timeout an engine may use.""" + import time + + import httpx + + async def main() -> None: + ports = [_free_port() for _ in range(4)] + workers = [_native_worker(0, flush_s=0.6), _native_worker(1)] + workers[1].state.hold = asyncio.Event() + router = Router([f"http://127.0.0.1:{p}" for p in ports[:2]]) + apps = [*zip(workers, ports), (build_router_app(router, poll_s=0, health_timeout_s=0.3), ports[2])] + servers, tasks = await _start(apps) + try: + async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{ports[2]}", timeout=60) as c: + t0 = time.monotonic() + r = await c.post("/flush_cache?timeout=0.5") + took = time.monotonic() - t0 + assert r.status_code == 400 and r.text.endswith("engine 1: no answer\n"), r.text + assert "engine 0" not in r.text # it took 0.6 s, within its 0.5 s timeout + 0.3 s + assert 0.75 <= took < 10, took + assert router.up == [True, True] # a stuck flush is not a failed engine + finally: + workers[1].state.hold.set() + await _stop(servers, tasks) + + asyncio.run(main()) 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()] diff --git a/tests/test_scheduler_policy.py b/tests/test_scheduler_policy.py new file mode 100644 index 0000000..432ee7d --- /dev/null +++ b/tests/test_scheduler_policy.py @@ -0,0 +1,387 @@ +"""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, LPM_MAX_QUEUE, 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) + with pytest.raises(ValueError, match="priority_preemption_threshold must be at least 0"): + SchedulerConfig(1, 1, 1, 2, **Scheduling(priority_preemption_threshold=-1).scheduler_fields()) + assert SchedulerConfig(1, 1, 1, 2, priority_preemption_threshold=0).priority_preemption_threshold == 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_lpm_past_its_queue_limit_returns_to_arrival_order(): + """Above LPM_MAX_QUEUE waiting requests lpm stops matching prefixes and orders by arrival (SGLang's fcfs then); a + queue an earlier round sorted by cache hits does not stay in that order.""" + sched = policy(make(), policy="lpm") + reqs = [req(f"r{i}", 4) for i in range(LPM_MAX_QUEUE + 1)] + for i, r in enumerate(reqs): + r.arrival_time = float(i) + shuffled = reqs[:] + random.Random(0).shuffle(shuffled) + sched.waiting = deque(shuffled) # as a cache-hit order an earlier round left + sched._order_waiting() + assert list(sched.waiting) == reqs + + +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 + + +BIG = 2**60 # the API takes signed 64-bit priorities; floats step by 256 here, so a float key would tie these + + +def test_priorities_past_2_53_order_admission_exactly(): + hi = policy(make(max_seqs=1, budget=64), priority_scheduling=True) + assert admitted(hi, [req("a", 4, priority=BIG), req("b", 4, priority=BIG + 1)]) == ["b", "a"] + lo = policy(make(max_seqs=1, budget=64), priority_scheduling=True, low_priority_first=True) + assert admitted(lo, [req("a", 4, priority=-BIG), req("b", 4, priority=-BIG - 1)]) == ["b", "a"] + sched = policy(make(), low_priority_first=True) + assert sched._importance(req("a", 4, priority=-BIG - 1)) > sched._importance(req("b", 4, priority=-BIG)) + + +@pytest.mark.parametrize(("prio", "preempts"), [(BIG + 11, True), (BIG + 10, False)]) +def test_the_preemption_threshold_holds_exactly_past_2_53(prio, preempts): + sched, low = _seats_full() + for r in low: + r.priority = BIG + 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) + + +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 diff --git a/tests/test_serve_anthropic.py b/tests/test_serve_anthropic.py new file mode 100644 index 0000000..35652c6 --- /dev/null +++ b/tests/test_serve_anthropic.py @@ -0,0 +1,312 @@ +"""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, + "messages": [ + {"role": "user", "content": "x"}, + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": None}, {"type": "text", "text": "a"}], + }, + {"role": "user", "content": "y"}, + ], + }, + { + **ASK, + "messages": [ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}]}, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "boom", "is_error": True}], + }, + ], + }, + {**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 eng.added == [] and eng.scheduler.running == [] # a refusal never reaches the engine + + +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"}], + "is_error": False, + }, + {"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 + ] + + +def test_a_whitespace_only_answer_makes_no_text_block(client): + """An answer of whitespace alone (thinking off, so parse_full keeps it) makes no text block, streamed or not.""" + + def newline(eng) -> None: + eng.tokenizer.add_tokens(["\n"]) + eng.script = [eng.tokenizer.convert_tokens_to_ids("\n")] + + c, _ = client(PLAIN, newline) + body = {**ASK, "thinking": {"type": "disabled"}} + assert c.post("/v1/messages", json=body).json()["content"] == [] + assert _blocks(_sse(c, body)) == [] diff --git a/tests/test_serve_api.py b/tests/test_serve_api.py index ebce5ad..f1904ab 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) @@ -349,7 +353,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), @@ -364,7 +367,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 @@ -420,8 +424,9 @@ def test_request_fields_are_checked_before_the_engine(serve): ("prompt", [-1]), ("prompt", [len(WORDS)]), ("prompt", [1, True]), - ("min_tokens", 4), - ("include_stop_str_in_output", True), + ("min_tokens", -1), + ("min_tokens", 100), # above max_tokens (16, the completions default) + ("include_stop_str_in_output", "yes"), ("thinking_token_budget", "lots"), ("temperature", 10**400), ] @@ -628,3 +633,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 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) diff --git a/tests/test_serve_native.py b/tests/test_serve_native.py new file mode 100644 index 0000000..f4212bd --- /dev/null +++ b/tests/test_serve_native.py @@ -0,0 +1,468 @@ +"""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"), + # served through the engine's guided decoding, which an engine without --structured-outputs refuses + ({"input_ids": [1], "sampling_params": {"json_schema": "{}"}}, "--structured-outputs"), + ({"input_ids": [1], "sampling_params": {"no_stop_trim": True}}, "not supported: no_stop_trim"), + ({"input_ids": [1], "sampling_params": {"skip_special_tokens": False}}, "not supported: skip_special_tokens"), + ({"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 + + from trnrt.serve import api + + 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")) + (lp,) = [lp for lp in api._LOOPS if lp.engine is eng] + _wait_for(lambda: lp._flush is not None) # the engine thread holds the first flush: the next one is refused + 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. /v1/stats counts neither probe.""" + import httpx + + from trnrt.serve import api, 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]) + # neither probe is a request in /v1/stats or the production log, the aborted one included: the engine thread + # handled the abort before this later request, so once its record is in, an abort record would be too + eng.step_delay = 0.0 + assert c.post("/v1/completions", json={"prompt": [1]}).status_code == 200 + (lp,) = [lp for lp in api._LOOPS if lp.engine is eng] + _wait_for(lambda: lp.stats.counts["ok"] == 3) + assert lp.stats.counts == {"ok": 3, "aborted": 0, "failed": 0} + probes = {eng.added[0].rid, eng.added[3].rid} + assert len(lp.stats.records) == 3 and not probes & {r["rid"] for r in lp.stats.records} + + +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"] + + +@pytest.mark.parametrize("stream", [False, True]) +def test_a_full_queue_503_from_generate_keeps_the_queue_full_header(stream): + """A request the full queue refuses, and one a better request displaces from it after admission: both 503s carry + QUEUE_FULL_HEADER, as on the OpenAI routes, so the router does not take the engine for failed.""" + from fastapi.testclient import TestClient + + from tests.test_serve_sched import QueueEngine + from trnrt.engine.scheduler import DISPLACED, Scheduling + from trnrt.serve import api + + eng = QueueEngine(PLAIN, Scheduling(priority_scheduling=True), limit=0) + try: + with TestClient(api.build_app(eng, "glm")) as c: + r = c.post("/generate", json={"input_ids": [1], "stream": stream}) + assert r.status_code == 503 and r.json()["detail"] == "The request queue is full.", r.text + assert r.headers[api.QUEUE_FULL_HEADER] == "1" + if stream: + return + eng.limit = 1 + waiting = _background(lambda: c.post("/generate", json={"input_ids": [1]})) + _wait_for(lambda: eng.parked) + better = c.post("/v1/completions", json={"prompt": [1], "priority": 9}) + assert better.status_code == 200, better.text + r = waiting() + assert r.status_code == 503 and r.json()["detail"] == DISPLACED and r.headers[api.QUEUE_FULL_HEADER] == "1" + finally: + for lp in [lp for lp in api._LOOPS if lp.engine is eng]: + lp.stop() + + +def test_no_logs_keeps_a_request_out_of_the_production_log(serve, monkeypatch, tmp_path): + """SGLang's no_logs: true writes nothing of the request to disk (TRNRT_PROD_DIR); /v1/stats still counts it, as + /metrics does (log_metrics is the field for those).""" + from trnrt.serve import api + from trnrt.serve import prodlog as PL + + monkeypatch.setenv(PL.ENV, str(tmp_path)) + client, eng = serve(PLAIN) + for rid, no_logs in (("quiet", True), ("logged", False)): + r = client.post("/generate", json={"input_ids": [1], "rid": rid, "no_logs": no_logs}) + assert r.status_code == 200, r.text + (lp,) = [lp for lp in api._LOOPS if lp.engine is eng] + _wait_for(lambda: lp.stats.counts["ok"] == 2 and lp.prodlog.written["requests"] == 1) + lp.stop() + (path,) = tmp_path.rglob("requests.jsonl") + assert [rec["rid"] for rec in PL.read_jsonl(str(path))] == ["logged"] + assert client.post("/generate", json={"input_ids": [1], "no_logs": "yes"}).status_code == 400 + + +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) diff --git a/tests/test_serve_responses.py b/tests/test_serve_responses.py new file mode 100644 index 0000000..11148f6 --- /dev/null +++ b/tests/test_serve_responses.py @@ -0,0 +1,309 @@ +"""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", + ] + # each item's deltas in their place: a run of one delta kind counts once (the script sends one or more) + mid = kinds[2:-1] + assert [k for i, k in enumerate(mid) if i == 0 or not (k.endswith(".delta") and k == mid[i - 1])] == expected + 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", "why"), + [ + ({"store": True}, "store is not supported"), + ({"previous_response_id": "resp_1"}, "previous_response_id is not supported"), + ({"background": True}, "background is not supported"), + ({"truncation": "auto"}, "truncation is not supported"), + ({"include": ["file_search_call.results"]}, "include supports only"), + ({"include": ""}, "include supports only"), # a falsy non-list is malformed, not an empty list + ({"include": False}, "include supports only"), + ({"include": {}}, "include supports only"), + ({"text": {"format": {"type": "json_schema", "name": "x", "schema": {}}}}, "text.format"), + ({"tools": [{"type": "web_search"}]}, "tool type 'web_search' is not supported"), + ({"tool_choice": "required"}, "tool_choice"), + ({"tool_choice": {"type": "function", "name": "get_weather"}}, "tool_choice"), + ({"top_logprobs": 2}, "top_logprobs is not supported"), + ({"max_tool_calls": 1}, "max_tool_calls is not supported"), + ({"n": 2}, "n is not a /v1/responses field"), + ({"max_tokens": 5}, "max_tokens is not a /v1/responses field"), + ({"messages": []}, "messages is not a /v1/responses field"), + ({"reasoning": {"effort": "medium"}}, "reasoning_effort"), + ({"reasoning": {"summary": "long"}}, "reasoning.summary must be"), + ({"reasoning": {"summary": "auto", "budget": 5}}, "reasoning.budget is not supported"), + ({"service_tier": "flex"}, "service_tier is not supported"), + ({"service_tier": "priority"}, "service_tier is not supported"), + ({"input": [{"role": "user", "content": [{"type": "input_image", "image_url": "x"}]}]}, "'input_image'"), + ({"input": [{"type": "item_reference", "id": "x"}]}, "item_reference is not supported"), + ({"input": []}, "input must be"), + ({"input": 3}, "input must be"), + ({"stream": "yes"}, "stream must be true or false"), + ], + ids=repr, +) +def test_what_cannot_be_served_is_refused(client, extra, why): + """Each refusal is a 400 whose detail says why (AGENTS.md), and nothing reaches the engine.""" + c, eng = client(PLAIN) + r = c.post("/v1/responses", json={"input": "hi", "tools": R_TOOLS, **extra}) + assert r.status_code == 400, r.text + assert why in r.json()["detail"], 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", + "safety_identifier": "user-hash-1", + "reasoning": {"summary": None, "generate_summary": None}, + } + 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 fastapi import HTTPException + + 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(HTTPException, match="not supported") as e: # a part without a type is not text + input_messages(body) + assert e.value.status_code == 400 + 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) + + +def test_a_whitespace_only_answer_makes_no_message_item(client): + """An answer of whitespace alone (thinking off, so parse_full keeps it) makes no message item, streamed or not.""" + + def newline(eng) -> None: + eng.tokenizer.add_tokens(["\n"]) + eng.script = [eng.tokenizer.convert_tokens_to_ids("\n")] + + c, _ = client(PLAIN, newline) + body = {"input": "hi", "chat_template_kwargs": {"enable_thinking": False}} + whole = c.post("/v1/responses", json=body).json()["output"] + assert whole == [] + assert _sse(c, body)[-1]["response"]["output"] == [] diff --git a/tests/test_serve_sampling.py b/tests/test_serve_sampling.py new file mode 100644 index 0000000..ba72406 --- /dev/null +++ b/tests/test_serve_sampling.py @@ -0,0 +1,356 @@ +"""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"}, + {"include_stop_str_in_output": True, "no_stop_trim": "yes"}, # checked even when the other keeps the stop + {"skip_special_tokens": "no"}, + {"service_tier": "flex"}, + {"service_tier": "priority"}, + ], + 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, + "service_tier": "auto", + } + 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. The engine takes 10 ms a + token, so the abort reaches it before it ends the request on its own.""" + client, eng = app(PLAIN * 10, setup=lambda e: setattr(e, "step_delay", 0.01)) + 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 back over the text (STOP_REGEX_SPAN characters before the new + 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( + "pattern", + [ + r"^(a+)+b$", # exponential: a repeat over a part that matches more than one way + r"(a|a)*b", + r"(a{1,3}){1,10}b", + r"(?:\w+\s?)+$", + r"a*a*b", # a power of the length: two repeats in a row that take the same characters + r"\d+\s*\d+", + r".*x.*y", + r"(?:)+", + ], +) +def test_stop_regex_refuses_patterns_that_backtrack(app, pattern): + """CPython's re holds the GIL through a search and cannot be interrupted, so a pattern whose search can take + exponential time, or a power of the text's length, is refused with the reason before it runs on the event loop.""" + client, eng = app(PLAIN) + r = client.post("/v1/completions", json={"prompt": [1], "stop_regex": ["again", pattern]}) + assert r.status_code == 400 and r.json()["detail"].startswith(f"stop_regex {pattern!r} is refused: it "), r.text + assert eng.added == [] + + +def test_stop_regex_admits_ordinary_patterns(): + from trnrt.serve.api import _backtracking + + ordinary = [ + r"\n\n+", + r"", + r"Answer:\s*\d+", + r"\d+\.\d+", + r"\$\d+(?:\.\d\d)?", + r"\bEND\b", + r"(?:Human|User):", + r"(?:foo|bar)+", + r"[.!?]\s*$", + r"\w+@\w+\.\w+", + r"(?s).*", + r"(?i)stop here", + r"[a-f0-9]{40}", + r"(\w+)\s\1", + r"re\s?ag", + r"!.*again", + ] + assert {p: _backtracking(p) for p in ordinary} == dict.fromkeys(ordinary) + + +def test_stop_regex_search_reaches_back_at_most_the_span(monkeypatch): + """The search for an unbounded pattern starts STOP_REGEX_SPAN (+ 8) characters before the text already searched, + which bounds its cost: a match starting further back is not found (the limitation, pinned), a later one is.""" + import re + + from trnrt.serve import api + + class Chars: # token i is the character chr(i) + def decode(self, ids, skip_special_tokens=True): + return "".join(map(chr, ids)) + + def run(span): + monkeypatch.setattr(api, "STOP_REGEX_SPAN", span) + d = api._Detok(Chars(), regexes=((re.compile(r"<.*>"), None),)) + text = "<" + "abcdefghijklmnopqrstuvwxyz" + "> then and more" + ids = [ord(x) for x in text] + out = "".join(d.push(ids[i : i + 2], i + 2 >= len(ids)) for i in range(0, len(ids), 2)) + return out, d.matched + + assert run(2048) == ("", r"<.*>") + assert run(4) == (" then ", r"<.*>") + + +@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!" 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()) diff --git a/tests/test_serve_stats.py b/tests/test_serve_stats.py new file mode 100644 index 0000000..9430fdd --- /dev/null +++ b/tests/test_serve_stats.py @@ -0,0 +1,428 @@ +"""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_start_and_today_windows_count_past_the_ring(): + """ "start" and "today" count every request, as their ok and token totals do, not only the ring's newest.""" + c = Clock() + st = S.Stats(ring=3, clock=c.wall, mono=c.mono_) + earlier = [ + S.derive( + { + "rid": f"e{i}", + "arrival": T0 - 50, + "first": None, + "done": T0 - 40, + "prompt": 1, + "cached": 0, + "output": 0, + "finish": "abort", + } + ) + for i in range(4) + ] + st.merge_day(earlier, [], T0 - 3600) + c.mono = 1060.0 + for i in range(5): + st.record(req(f"r{i}", 1001, 1002, 1003 + i), stream=False, at_arrival=None) + w = st.snapshot({}, now=c.wall(), ttl=0)["windows"] + assert len(st.records) == 3 and w["1m"]["n"] == 3 # a window inside the ring counts what it holds + assert (w["start"]["n"], w["start"]["ok"]) == (5, 5) + assert (w["today"]["n"], w["today"]["ok"]) == (9, 5) + + +def test_the_current_series_point_rates_over_its_elapsed_seconds(): + """The newest point is a bucket still filling: its RPM, like its output tok/s, divides by the seconds it has run.""" + st, _ = stats() + st.tick(1000.5, 20, 1, 0) + st.record(req("a", 1000.2, 1000.4, 1001.0), stream=False, at_arrival=None) + last = st.snapshot({}, now=T0 + 2.5, ttl=0)["series"][-1] + assert last["t"] == T0 and last["out_tok_s"] == pytest.approx(20 / 2.5) + assert last["rpm"] == pytest.approx(60 / 2.5) + + +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_today_counts_the_output_of_earlier_processes_all_day(tmp_path): + """After a restart, "today"'s output tok/s counts every gauge of the day, not only the 30 minutes the series reads + back.""" + day = tmp_path / "engine-8001" / PL.utc_day(T0) + day.mkdir(parents=True) + rec = {"rid": "a", "arrival": T0 - 3700, "first": T0 - 3699, "done": T0 - 3690, "prompt": 5, "cached": 0} + (day / "requests.jsonl").write_text(json.dumps({**rec, "output": 3, "finish": "stop"}) + "\n") + gauges = [{"t": T0 - 3600, "out_tok_s": 5.0}, {"t": T0 - 60, "out_tok_s": 2.0}] # an hour ago, and a minute ago + (day / "gauges.jsonl").write_text("".join(json.dumps(g) + "\n" for g in gauges)) + st, c = stats() + c.mono = 1100.0 + PL.ProdLog(str(tmp_path / "engine-8001"), stats=st, clock=c.wall).close() + today = st.snapshot({}, now=c.wall(), ttl=0)["windows"]["today"] + assert today["first"] == T0 - 3700 + assert today["out_tok_s"] == pytest.approx((50 + 20) / (T0 + 100 - (T0 - 3700))) + assert set(st.past) == {int((T0 - 60) // S.BUCKET_S)} # the series still reads back 30 minutes only + + +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}] + cut = tmp_path / "cut.jsonl.gz" # a gzip stream cut short, or a disk that lost its tail + with gzip.open(cut, "wt") as f: + f.writelines(json.dumps({"i": i, "pad": os.urandom(64).hex()}) + "\n" for i in range(200)) + data = cut.read_bytes() + cut.write_bytes(data[: len(data) // 2]) + got = PL.read_jsonl(str(cut)) + assert 0 < len(got) < 200 and [r["i"] for r in got] == list(range(len(got))) + 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_a_segment_is_read_once_while_gzip_runs_or_after_it_was_cut_short(tmp_path, monkeypatch): + """gzip_file publishes the .gz only once complete; a restart between its rename and its remove leaves both copies, + and day_files reads one.""" + seg = tmp_path / "requests.1.jsonl" + seg.write_text('{"rid": "a"}\n{"rid": "b"}\n') + + def cut_short(f, g, n): + g.write(f.read(5)) + assert not os.path.exists(str(seg) + ".gz"), "a partial .gz is never visible under its name" + raise OSError("killed") + + monkeypatch.setattr(PL.shutil, "copyfileobj", cut_short) + PL.gzip_file(str(seg)) + assert sorted(os.listdir(tmp_path)) == ["requests.1.jsonl"] # the source whole, no partial copy left + monkeypatch.undo() + with gzip.open(str(seg) + ".gz", "wb") as g: # the rename done, the remove not + g.write(seg.read_bytes()) + (tmp_path / "requests.jsonl").write_text('{"rid": "c"}\n') + files = PL.day_files(str(tmp_path), "requests") + assert files == [str(seg), str(tmp_path / "requests.jsonl")] + assert [r["rid"] for p in files for r in PL.read_jsonl(p)] == ["a", "b", "c"] + PL.gzip_file(str(seg)) + assert sorted(os.listdir(tmp_path)) == ["requests.1.jsonl.gz", "requests.jsonl"] + + +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 " + + +""" diff --git a/trnrt/serve/native.py b/trnrt/serve/native.py new file mode 100644 index 0000000..f1e7633 --- /dev/null +++ b/trnrt/serve/native.py @@ -0,0 +1,621 @@ +"""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). Guided decoding, min_new_tokens and penalties go to the engine, which serves them +when started with --structured-outputs or --penalties and refuses them otherwise. What the served graphs cannot compute +(prompt logprobs, 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, cast + +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, _error, _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", + "logit_bias": "logit_bias", + "min_new_tokens": "min_tokens", + "json_schema": "guided_json", + "regex": "guided_regex", + "ebnf": "guided_grammar", +} +# SGLang sampling parameters with no counterpart here, or not served on /generate (n choices, kept stop text and special +# tokens in the text are the OpenAI endpoints' Run path), accepted only at these values. +_SAMPLING_NEUTRAL: dict[str, tuple[object, ...]] = { + "n": (None, 1), + "no_stop_trim": (None, False), + "skip_special_tokens": (None, True), + "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 cast("asyncio.Queue", self.sub.out).get() + if kind == "error": + raise _error(*rest) # a displaced request's 503 keeps QUEUE_FULL_HEADER + 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 cast("asyncio.Queue", 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, logs = _flag(body, "log_metrics", True), not _flag(body, "no_logs") + 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, logs=logs) + 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. + # /generate's requests own their queues (api._Submit.out is a _Tagged view only for `n` > 1 choices) + firsts = [await cast("asyncio.Queue", 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 _error(*bad[1:]) # a full queue's 503 keeps QUEUE_FULL_HEADER: the router keeps the engine up + 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 cast("asyncio.Queue", 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() diff --git a/trnrt/serve/prodlog.py b/trnrt/serve/prodlog.py new file mode 100644 index 0000000..484c868 --- /dev/null +++ b/trnrt/serve/prodlog.py @@ -0,0 +1,305 @@ +"""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 +import zlib +from typing import IO, TYPE_CHECKING + +from .stats import BUCKET_S + +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, and a truncated or corrupt .gz + (EOFError, zlib.error) keeps the records read before it.""" + 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, EOFError, zlib.error): + 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 = set(os.listdir(day_dir)) + except OSError: + return [] + segs: list[tuple[float, str]] = [] + for n in names: + base = n.removesuffix(".gz") + if n != base and base in names: # a restart between gzip_file's rename and its remove: read one copy + continue + 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: + """path as path.gz, which appears only once complete (written as a .tmp, then renamed), then path removed. A + process killed meanwhile leaves the source whole, and day_files reads it alone.""" + tmp = path + ".gz.tmp" + try: + with open(path, "rb") as f, gzip.open(tmp, "wb") as g: + shutil.copyfileobj(f, g, 1 << 20) + os.replace(tmp, path + ".gz") + os.remove(path) + except OSError as e: + logger.warning("production log: could not gzip %s: %s", path, e) + with contextlib.suppress(OSError): + os.remove(tmp) + + +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) + start = calendar.timegm(time.strptime(day, "%Y-%m-%d")) # UTC midnight + 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))] + # the "today" output tok/s counts every gauge of the day (a bucket closed just after midnight is yesterday's); + # the series reads back only its own 30 minutes + out_tok_s = (g.get("out_tok_s") for g in gauges if g["t"] >= start) + handed = BUCKET_S * sum(x for x in out_tok_s if isinstance(x, (int, float))) + recent = [g for g in gauges if g["t"] >= now - GAUGES_RELOAD_S] + stats.merge_day(recs, recent, start, handed=handed) diff --git a/trnrt/serve/prompt_cache.py b/trnrt/serve/prompt_cache.py new file mode 100644 index 0000000..b49804a --- /dev/null +++ b/trnrt/serve/prompt_cache.py @@ -0,0 +1,286 @@ +"""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, +is matched after normalization or matches only as a whole word (lstrip, rstrip, normalized and single_word 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, normalize or + match single words (their match depends on the characters around them, which the self-check samples only); +- 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) + # each makes a token's match depend on the characters around it. single_word: cached "x foo bar", a cut at + # "foo" keeps "x " apart from it, while in "x foobar" no token matches and " f" merges; the self-check passes + bad = [t.content for t in added if t.lstrip or t.rstrip or t.normalized or t.single_word] + 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, normalize or match single words: {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, + ) diff --git a/trnrt/serve/responses.py b/trnrt/serve/responses.py new file mode 100644 index 0000000..6dccb07 --- /dev/null +++ b/trnrt/serve/responses.py @@ -0,0 +1,552 @@ +"""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. Four are dropped because they +# change nothing a client gets back: `metadata` and `user` (echoed), `safety_identifier` (names the end user for OpenAI's +# abuse monitoring, as `user` does) and `prompt_cache_key` (a hint for OpenAI's prompt-cache routing; the prefix cache +# here matches every prompt by its tokens, and the router routes by session_id or prefix). `service_tier` is served +# only as "auto" or "default" (to_chat). +_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") + # flex, priority and scale are processing classes with their own latency, which one engine's queue does not have + if body.get("service_tier") not in (None, "auto", "default"): + raise _bad('service_tier is not supported; send "auto" or "default", or leave it out') + include = body.get("include") + if include is None: + include = [] + 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 {} + extra = sorted(k for k in reasoning if k not in ("effort", "summary", "generate_summary")) + if extra: + raise _bad(f"reasoning.{extra[0]} is not supported: reasoning takes effort and summary") + 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. A + whitespace-only answer makes no message item, as the stream opens none for it (_Stream).""" + out = [] + if msg.get("reasoning_content"): + out.append(_reasoning_item(_id("rs_"), msg["reasoning_content"], info)) + if msg.get("content") and msg["content"].strip(): + 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 e.http() 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)" + ) diff --git a/trnrt/serve/router.py b/trnrt/serve/router.py index 3085e1b..e6d9c67 100644 --- a/trnrt/serve/router.py +++ b/trnrt/serve/router.py @@ -50,12 +50,31 @@ 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 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 | +- 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 """ @@ -67,15 +86,24 @@ import contextlib import hashlib import json +import math import os import threading import time import httpx from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse, Response, StreamingResponse +from fastapi.responses import ( + HTMLResponse, + JSONResponse, + PlainTextResponse, + RedirectResponse, + Response, + StreamingResponse, +) from .. import log +from .dashboard import PAGE as DASHBOARD_PAGE logger = log.get("router") @@ -86,7 +114,16 @@ # 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 (_prompt), routed by its session or its prefix. +ROUTES = ("/v1/completions", "/v1/chat/completions", "/generate", "/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", "/v1/messages/count_tokens") +# 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: @@ -344,6 +381,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]) @@ -460,16 +502,76 @@ 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]" +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 _json_head(x, k: int): + """x cut to its first k items and string characters, counted in document order at every level: its JSON text + starts with the same k characters as x's, since each one kept writes at least one character before the first that + differs. A key reads that much of a prompt of any size (a 200K-token input_ids list, a batch of them) without + writing all of it on the event loop.""" + left = k + + def cut(v): + nonlocal left + if isinstance(v, str): + n = max(0, min(len(v), left)) + left -= n + return v[:n] + if isinstance(v, list): + out = [] + for item in v: + if left <= 0: + break + left -= 1 + out.append(cut(item)) + return out + if isinstance(v, dict): + d = {} + for key, item in v.items(): + if left <= 0: + break + left -= 1 + d[key] = cut(item) + return d + return v + + return cut(x) + + 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 @@ -477,21 +579,25 @@ 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"): - 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 = _conversation(path, body) + if msgs is None: + p = _prompt(path, body) + p = p if isinstance(p, str) else json.dumps(_json_head(p, SESSION_KEY_CHARS)) + return h(""), h(p[:SESSION_KEY_CHARS]) + 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") - c = c if isinstance(c, str) else "" if c is None else json.dumps(c, ensure_ascii=False) + if not isinstance(c, str): + c = "" if c is None else json.dumps(_json_head(c, SESSION_KEY_CHARS), ensure_ascii=False) return h(json.dumps(msgs[:k], ensure_ascii=False)), h(str(first.get("role", "")) + "\x00" + c[:SESSION_KEY_CHARS]) 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", "") + msgs = _conversation(path, body) + if msgs is not None: + return json.dumps(msgs, ensure_ascii=False) + p = _prompt(path, body) return p if isinstance(p, str) else json.dumps(p) @@ -596,13 +702,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 @@ -611,7 +716,172 @@ 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, hard_503: bool = True, **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), or with hard_503 False counts toward + down_after as a stalled /health does (Router.note_failure).""" + + 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=hard_503) + 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"). An + engine answers 503 when no step of its ended within HEALTH_GENERATE_TIMEOUT_S, which a long prefill can do to + an engine still serving: that 503 counts toward down_after, as a stalled /health does, and does not move the + sessions pinned there at once. A failed engine's /health answers 503 too, which the poll takes as down.""" + from .native import HEALTH_GENERATE_TIMEOUT_S + + rs = await each("GET", "/health_generate", HEALTH_GENERATE_TIMEOUT_S + 10, hard_503=False) + 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. An engine that has not answered `timeout` + health_timeout_s seconds + after it was asked is one that did not flush ("no answer"): an engine answers once it flushed or its timeout + passed, so only one that is stuck takes longer.""" + query = http.url.query + try: + wait = float(http.query_params.get("timeout", "0")) + except ValueError: + wait = 0.0 # each engine answers a malformed timeout with 400 at once + wait = wait if math.isfinite(wait) and wait >= 0 else 0.0 + rs = await each(http.method, "/flush_cache" + (f"?{query}" if query else ""), wait + health_timeout_s) + 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 "") + + 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 + # ASCII digits only ("²" is a digit to isdigit and not to int), and few: int refuses above 4300 digits + if not (raw.isascii() and raw.isdigit()) or len(raw) > 9 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) @@ -654,7 +924,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): @@ -691,8 +970,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"]) + 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 diff --git a/trnrt/serve/stats.py b/trnrt/serve/stats.py new file mode 100644 index 0000000..4abf917 --- /dev/null +++ b/trnrt/serve/stats.py @@ -0,0 +1,373 @@ +"""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.recorded = 0 # finished requests recorded since start (the ring keeps only the newest RING) + 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.recorded += 1 + 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, handed: float = 0.0) -> None: + """Today's records and gauges from earlier processes: the "today" window, the ring's older end, the series. + `handed`: the output tokens their gauges counted all day (gauges holds only the series' last 30 minutes).""" + 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, + "n": len(records), + "handed": handed, + **{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, recorded = self.handed, self.recorded + 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=(recorded, 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"]] + dhanded = handed + day["handed"] # handed out today: this process's, plus earlier processes' all day + windows["today"] = self._window( + dsel, + dspan, + dhanded / dspan if dspan > 0 else None, + totals=(day["n"] + recorded, 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. `totals` + (requests, ok requests, tokens) counts a window longer than the ring holds ("start", "today").""" + 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, n_ok = len(sel), len(ok) + prompt, cached, output = (sum(r[k] for r in sel) for k in ("prompt", "cached", "output")) + else: + n, 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": n, + "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} + width = min(BUCKET_S, now - t0) if k == last else BUCKET_S # the current bucket has run `width` s so far + if b is not None: + 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 / width if width > 0 else None + 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 diff --git a/trnrt/serve/structured.py b/trnrt/serve/structured.py new file mode 100644 index 0000000..2744db0 --- /dev/null +++ b/trnrt/serve/structured.py @@ -0,0 +1,126 @@ +"""A request's output constraint (engine/grammar.py GrammarSpec) from the fields the OpenAI API, vLLM and SGLang use: + +| Field | From | Constraint | +|---|---|---| +| `response_format` `{"type": "json_object"}` | OpenAI | a JSON object (SGLang's schema `{"type": "object"}`) | +| `response_format` `{"type": "json_schema", "json_schema": {"schema": ...}}` | OpenAI | the schema | +| `guided_json`, `guided_regex`, `guided_choice`, `guided_grammar` | vLLM (older) | schema, regex, one of the strings, EBNF | +| `structured_outputs` `{json, regex, choice, grammar, json_object}` | vLLM | the same | +| `json_schema` (a string), `regex`, `ebnf` | SGLang | schema, regex, EBNF | + +At most one constraint per request, as SGLang requires ("Only one of json_schema, regex, ebnf, or structural_tag can be +set"). A malformed field is a ValueError, which the server answers with 400; a constraint xgrammar cannot compile is +refused by the engine when the request is added (also a 400). +""" + +from __future__ import annotations + +import json + +from ..engine.grammar import GrammarSpec + +OBJECT = '{"type": "object"}' +_VLLM_KEYS = ("json", "regex", "choice", "grammar", "json_object") +# vLLM's other StructuredOutputsParams fields (vllm/sampling_params.py), accepted only at their defaults: a client that +# sends the whole dataclass sends them. Any other key is refused whatever its value. +_VLLM_NEUTRAL: dict[str, tuple[object, ...]] = { + "disable_any_whitespace": (None, False), + "disable_additional_properties": (None, False), + "whitespace_pattern": (None,), + "structural_tag": (None,), +} + + +def _schema_text(v, where: str) -> str: + if isinstance(v, str): + try: + json.loads(v) + except ValueError: + raise ValueError(f"{where} must be a JSON schema (an object, or its JSON text)") from None + return v + if isinstance(v, (dict, bool)): + return json.dumps(v) + raise ValueError(f"{where} must be a JSON schema (an object, or its JSON text)") + + +def _text(v, where: str) -> str: + if not isinstance(v, str) or not v: + raise ValueError(f"{where} must be a non-empty string") + return v + + +def _choice(v, where: str) -> GrammarSpec: + if not isinstance(v, list) or not v or not all(isinstance(x, str) and x for x in v): + raise ValueError(f"{where} must be a non-empty list of non-empty strings") + return GrammarSpec.choice(v) + + +def _response_format(rf) -> GrammarSpec | None: + if rf is None: + return None + if not isinstance(rf, dict): + raise ValueError("response_format must be an object") + kind = rf.get("type", "text") + if kind == "text": + return None + if kind == "json_object": + return GrammarSpec("json_schema", OBJECT) + if kind == "json_schema": + js = rf.get("json_schema") + if js is None and "schema" in rf: # SGLang also takes the schema at the top of response_format + js = {"schema": rf["schema"]} + if not isinstance(js, dict) or js.get("schema") is None: + raise ValueError('response_format json_schema needs "json_schema": {"name": ..., "schema": {...}}') + return GrammarSpec("json_schema", _schema_text(js["schema"], "response_format.json_schema.schema")) + raise ValueError(f"response_format type {kind!r} is not supported: text, json_object or json_schema") + + +def _vllm(so) -> list[GrammarSpec]: + if so is None: + return [] + if not isinstance(so, dict): + raise ValueError("structured_outputs must be an object") + unknown = sorted(k for k, v in so.items() if k not in _VLLM_KEYS and v not in _VLLM_NEUTRAL.get(k, ())) + if unknown: + raise ValueError(f"structured_outputs.{unknown[0]} is not supported: {', '.join(_VLLM_KEYS)}") + out = [] + if so.get("json") is not None: + out.append(GrammarSpec("json_schema", _schema_text(so["json"], "structured_outputs.json"))) + if so.get("regex") is not None: + out.append(GrammarSpec("regex", _text(so["regex"], "structured_outputs.regex"))) + if so.get("choice") is not None: + out.append(_choice(so["choice"], "structured_outputs.choice")) + if so.get("grammar") is not None: + out.append(GrammarSpec("ebnf", _text(so["grammar"], "structured_outputs.grammar"))) + json_object = so.get("json_object") + if json_object is not None and not isinstance(json_object, bool): # the string "false" must not turn it on + raise ValueError("structured_outputs.json_object must be true or false") + if json_object: + out.append(GrammarSpec("json_schema", OBJECT)) + return out + + +def grammar_spec(body: dict) -> GrammarSpec | None: + """The request's constraint, or None. ValueError for a malformed field or more than one constraint.""" + specs: list[tuple[str, GrammarSpec]] = [] + rf = _response_format(body.get("response_format")) + if rf is not None: + specs.append(("response_format", rf)) + specs += [("structured_outputs", s) for s in _vllm(body.get("structured_outputs"))] + fields = ( + ("guided_json", lambda v, k: GrammarSpec("json_schema", _schema_text(v, k))), + ("guided_regex", lambda v, k: GrammarSpec("regex", _text(v, k))), + ("guided_choice", _choice), + ("guided_grammar", lambda v, k: GrammarSpec("ebnf", _text(v, k))), + ("json_schema", lambda v, k: GrammarSpec("json_schema", _schema_text(v, k))), + ("regex", lambda v, k: GrammarSpec("regex", _text(v, k))), + ("ebnf", lambda v, k: GrammarSpec("ebnf", _text(v, k))), + ) + for key, make in fields: + v = body.get(key) + if v is not None: + specs.append((key, make(v, key))) + if len(specs) > 1: + names = ", ".join(dict.fromkeys(k for k, _ in specs)) + raise ValueError(f"at most one output constraint per request, got {names}") + return specs[0][1] if specs else None