diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index c4cb84d6bbf8..63b2dc9644c2 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -84,6 +84,24 @@ def _init_hf_modules(): _init_hf_modules() +def _request_stop_words(request: GenerationRequest) -> List[List[int]]: + """Runtime stop words for a request. + + Combines the stop words derived from ``sampling_params`` (stop token ids + plus tokenizer-processed stop strings) with the request's pre-tokenized + ``stop_token_sequences``, which callers set when they tokenize stop + strings themselves. ``ignore_eos`` suppresses all of them, matching the + legacy behavior for sampling-params-derived stops. + """ + if request.sampling_params.ignore_eos: + return [] + stop_words = request.sampling_params._get_stop_words() + stop_token_sequences = getattr(request, "stop_token_sequences", None) + if stop_token_sequences: + stop_words = stop_words + [list(seq) for seq in stop_token_sequences] + return stop_words + + class BaseWorker(GenerationExecutor): class WorkerExit(GeneratorExit): @@ -546,8 +564,7 @@ def _deduce_max_tokens(request: GenerationRequest, guided_decoding_params=request.sampling_params. _get_guided_decoding_params(), bad_words=request.sampling_params._get_bad_words(), - stop_words=[] if request.sampling_params.ignore_eos else - request.sampling_params._get_stop_words(), + stop_words=_request_stop_words(request), embedding_bias=request.sampling_params.embedding_bias, lora_config=lora_config, prompt_tuning_config=prompt_tuning_config, @@ -1337,12 +1354,20 @@ def _compute_pytorch_prompt_logprobs( prompt=cached, generation=None ) # generation logprobs, if requested, is provided directly in response.result.log_probs from the sampler. context_logits = response.result.context_logits - assert context_logits is not None, "context_logits must not be None when prompt_logprobs is requested." result = response.result.get_result() - assert result is not None, "result must not be None when prompt_logprobs is requested." - # Single element list - first_generation_token = result.output_token_ids[0][:1] - assert first_generation_token, "first generation token must not be empty when prompt_logprobs is requested." + first_generation_token = (result.output_token_ids[0][:1] if + (result is not None and result.output_token_ids + and result.output_token_ids[0]) else []) + if context_logits is None or not first_generation_token: + # A final response with no generated token (e.g. a request cancelled + # before its first token) legitimately has nothing to compute prompt + # logprobs from. Skip instead of asserting so the terminal response + # still reaches the consumer. + logger.debug( + f"Skipping prompt-logprob computation for client_id=" + f"{response.client_id}: response carries no generated token or " + f"context logits.") + return None # Pass prompt_token_ids with an offset of 1 for correct mapping to the context logits prompt_token_ids = generation_result._generation_request.prompt_token_ids[ 1:] + first_generation_token diff --git a/tensorrt_llm/executor/engine_client/E5B_CUTOVER_DESIGN.md b/tensorrt_llm/executor/engine_client/E5B_CUTOVER_DESIGN.md new file mode 100644 index 000000000000..132f291347ca --- /dev/null +++ b/tensorrt_llm/executor/engine_client/E5B_CUTOVER_DESIGN.md @@ -0,0 +1,575 @@ +# E5b cut-over: contract-native submission and rank-0 result removal + +Status: design for gate ①c on branch `prototype_1`. + +This change cuts eligible engine-client requests over from the proxy's legacy +`GenerationResult` path to a contract-native path. After the cut-over, a contract request has +no rank-0 `GenerationResult` and no entry in `GenerationExecutorProxy._results`. Legacy requests +on the same proxy continue to use `GenerationExecutorProxy.submit`, rank-0 `GenerationResult`, +and the existing queue delivery behavior without semantic changes. + +The current path is: + +```text +LocalProcessEngineClient.submit + -> engine_request_to_generation_request + -> EngineFrameRouter.register_pending + -> GenerationExecutorProxy.submit + -> allocate client id + -> EngineFrameRouter.observe_submit + -> construct GenerationResult + -> proxy._results[client_id] = result + -> enqueue GenerationRequest + -> dispatch_result_task + -> EngineFrameRouter.on_response (observe) + -> proxy._results delivery (legacy) +``` + +The E5b path is: + +```text +LocalProcessEngineClient.submit + -> GenerationExecutorProxy.submit_contract + -> EngineService.submit_contract + -> engine_request_to_generation_request + -> allocate client id from the proxy's shared allocator + -> bind ContractRequestState + -> enqueue GenerationRequest + -> dispatch_result_task + -> EngineService.route_response (claim and consume contract ids) + -> legacy proxy._results delivery only for non-contract ids +``` + +`GenerationExecutorProxy.submit` remains the legacy entry point. Contract traffic must not call it +after E5b. + +## 1. Contract request state + +### Decision + +`EngineFrameRouter.RequestBinding` already is the rank-0 contract request state. E5b should not add +a second `ContractRequestState` registry beside it. Either rename `RequestBinding` to +`ContractRequestState` and keep `RequestBinding` as a private compatibility alias, or retain the +name and document that it is the canonical state. The important constraint is one state object and +one router-owned registry per contract request. + +The existing binding already owns request/delivery lifecycle state under +`EngineFrameRouter._lock`: frontend request id, proxy client id, prompt token ids, stop-reason +associations, delivery, runtime-started/terminal/ended flags, abort state, event sequence, +completion and cached-token accounting, prompt-logprob delivery state, recent tokens, final +metrics/status, and stream-open state. + +The canonical field model for E5b is: + +| Field | Required representation and purpose | Current `RequestBinding` status | +|---|---|---| +| Frontend request id | `request_id: str`; the contract-visible identity and stream lookup key | Present | +| Proxy client id | `client_id: int`; allocated before binding and used for worker responses and abort | Present but initially optional because of the pending/observe hook; E5b binds it at construction | +| Prompt | `prompt_token_ids: tuple[int, ...]` plus a derived `prompt_length`; the ids are required by `normalize_response` for map-shaped prompt logprobs, so storing only the length is insufficient | Token ids present; length is currently repeatedly derived with `len(...)` | +| Per-sequence state | `sequence_states: dict[int, SequenceState]`, where each state owns `runtime_started`, `terminal_emitted`, `completion_tokens`, `recent_tokens`, prompt-logprob sent/held state, cached tokens, pending final metrics, and final status | These are currently scalar fields on the binding. That is correct only because V0 rejects `n > 1`; E5b may initialize sequence `0` eagerly or lazily, but must not imply that the current scalars support multiple sequences | +| Stop reason state | Ordered `(stop_token_sequence, visible_reason)` pairs used by `_resolve_stop_reason` | Present as `stop_reasons` | +| Delivery state | Bounded `_Delivery`, stream-open flag, request-ended flag, and request-level ending state | Present | +| Abort state | `abort_requested: bool` | Present | +| Event ordering | Request-wide `event_seq: int`, incremented only through `next_seq()` while holding the router lock | Present | +| Routing classification | `_by_client` for active or delivery-retained bindings, plus a bounded ordered recently-retired proxy-client-id structure used only for late/duplicate accounting | `_by_client` is present, but the current service-lifetime `_owned_client_ids` set must be replaced by the bounded structure | +| Request-id lifecycle | The frontend id remains reserved while its binding is active or its delivery is retained; it becomes reusable after the request has ended and that delivery is retired | The active and delivery indexes exist, but the current service-lifetime `_seen_request_ids` set must be removed and stream consumption/close must retire the delivery | + +Delivery retirement is an explicit lifecycle transition under `EngineFrameRouter._lock`. It occurs +when the consumer reads through the request-ending frame, explicitly closes the stream, or the +bounded tombstone policy evicts the retained delivery. Duplicate frontend ids are rejected only +while the prior binding is active or its delivery remains in `_delivery_index`; after retirement, +reusing that frontend id is a required positive replay case. The recently-retired proxy-client-id +structure is independent of frontend-id reuse and is capped by `tombstone_limit` +(`DEFAULT_TOMBSTONE_LIMIT`, currently 4096). No per-request service-lifetime ownership set is part +of the canonical state. + +The per-sequence refactor is structurally preferable because the wire contract says `Terminal` is +exactly once per started sequence. It is not a scope expansion: V0 still admits exactly one +sequence. If E5b keeps the scalar representation to minimize the gate, the state must be named and +documented explicitly as sequence-0 state, and any envelope with `sequence_index != 0` must fail the +request rather than mix token/terminal state across sequences. + +### Sampling and logprob state + +Do **not** put a second `SamplingParams` or `LogprobParams` copy in the rank-0 contract state. +`engine_request_to_generation_request` constructs a `GenerationRequest` whose synthetic +`SamplingParams` is what the worker receives. For `_get_logprob_params`, the relevant values are +`logprobs`, `prompt_logprobs`, `logprobs_simple_format`, +`prompt_logprobs_simple_format`, `_need_return_context_logits`, and +`_need_return_generation_logits` (plus the worker's postprocess-worker count). The contract +translation explicitly restores the two requested logprob counts. The other four values use their +synthetic `SamplingParams` defaults: both simple-format flags are false and both explicit +return-logit flags are false. That is correct for the current eligible surface: raw logits are +rejected and `num_postprocess_workers` must be zero, while `normalize_response` accepts the +map-shaped logprobs produced when simple format is false. If the contract later promises a specific +worker logprob representation, those booleans must be added to `EngineSamplingConfig`; they still +must not be stored in a rank-0 `GenerationResult`. + +The rank-0 router needs only the prompt token ids to normalize map-shaped prompt logprobs and the +contract bookkeeping listed above. It must not call `_get_logprob_params`; doing so would preserve +rank-0-only state that E5b is intended to remove. + +## 2. `submit_contract` on the proxy + +### Interface and ownership + +Add a distinct internal entry point: + +```python +GenerationExecutorProxy.submit_contract( + engine_request: EngineRequest, + *, + stop_reasons: tuple = (), +) -> int +``` + +It returns the allocated proxy `client_id`, not a `GenerationResult`. +`LocalProcessEngineClient.submit` continues to return the frontend `request_id`; it uses the integer +only as an internal submission result. The proxy method is a narrow facade over +`EngineService.submit_contract`, described in section 5. + +Move the call to `engine_request_to_generation_request` behind this contract-native entry point. +This makes the service boundary accept the contract request rather than an in-process +`GenerationRequest`, and keeps later transport detachment from changing the service internals. + +### Required order + +Submission and shutdown need a proxy-level submission/lifecycle lock. Both legacy `submit` and +contract `submit_contract` use it. This lock serializes allocation plus enqueue with +`pre_shutdown`; it does not own router state. The lock order is always proxy lifecycle lock, then +`EngineFrameRouter._lock`. The dispatch thread never takes the proxy lifecycle lock. + +The contract flow is: + +1. Perform type/capability checks and `engine_request_to_generation_request` before allocating an + id or mutating a registry. A conversion rejection therefore has no engine-side state. +2. Start the result dispatch thread as `GenerationExecutorProxy.submit` does today. +3. Acquire the proxy submission/lifecycle lock. Reject if `doing_shutdown`, `_fatal_error`, or the + service fatal latch is set. +4. Allocate the proxy client id with the same `_get_next_client_id()` used by legacy submission, + under this lock, and call `generation_request.set_id(client_id)`. +5. Under the router's single `_lock`, reject a duplicate frontend request id and install the fully + bound state in the request-id, client-id, delivery, and routing-ownership registries. There is no + `_pending[id(generation_request)]` interval and no `observe_submit` matching by Python object + identity. +6. Still under the proxy lifecycle lock, enqueue the `GenerationRequest`. Binding is visible before + the worker can return a response, while shutdown cannot put its sentinel between binding and + enqueue. +7. If enqueue raises, call `EngineService.on_submit_enqueue_failed(client_id)` before releasing + the lifecycle lock. It transitions the binding to a standalone + `ErrorFrame(error_code="enqueue_failed")`. The request was never sent, so no runtime abort is + needed. Re-raise the enqueue exception. +8. Release the lifecycle lock, run the existing background-error check, and return `client_id`. + If that check discovers a fatal error after enqueue, the normal `pre_shutdown` path poisons the + binding. + +The binding replaces `proxy._results[client_id] = result`. In particular, the contract path must +not construct `GenerationResult`, call the proxy's `_get_logprob_params`, allocate a result queue, +or insert anything into `GenerationExecutorProxy._results`. + +The current `register_pending` -> `observe_submit` protocol can then be removed from the contract +path. Keeping it available only for transitional tests is acceptable temporarily, but production +contract submission must have one explicit method so it cannot accidentally fall back through +legacy `submit`. + +### Why worker-side response wrapping still works + +The rank-0 and worker registries are independent. In `worker_main`, rank 0 dequeues the +`GenerationRequest` and calls `BaseWorker.submit`. That method: + +1. preserves the proxy-assigned `request.id` as `client_id`; +2. calls `BaseWorker._get_logprob_params(request)`; +3. constructs a worker-local `GenerationResult` with those parameters; +4. inserts it into `BaseWorker._results[client_id]`; and +5. only then calls `_enqueue_request`. + +The await-response path consults this worker-local registry in `_get_logprobs`, +`_compute_pytorch_prompt_logprobs`, and `_get_params_for_first_rsp`; `_maybe_wrap_response` can +therefore continue to build `ResponseWrapper` with prompt logprobs and request metrics. `_send_rsp` +also continues to pop the worker-local result on final/error. This is V0 behavior and must remain. + +Tests for E5b must spy on the two processes/objects separately: a contract request is absent from +`GenerationExecutorProxy._results` while present in `BaseWorker._results` during execution, and +prompt-logprob `ResponseWrapper` output remains identical to the pre-cut-over path. + +## 3. Unified response routing in `dispatch_result_task` + +The router must change from an observer to an exclusive route claimant. Add an operation such as +`EngineService.route_response(raw) -> bool`, where `True` means the client id is owned by the +contract population and the response has been consumed or deliberately absorbed. + +For every unbatched response item, `GenerationExecutorProxy.dispatch_result_task` must use this +exact order: + +```python +client_id = raw.client_id + +if engine_service is not None and engine_service.route_response(raw): + continue + +result = self._results.get(client_id) +if result is None: + continue + +# Existing legacy queue delivery, async notification, and final/error pop. +deliver_to_generation_result(result, raw) +if is_final_or_error(raw): + self._results.pop(client_id, None) +``` + +Contract routing is checked before `proxy._results.get`. A claimed contract response is never +delivered to a legacy queue and never enters the legacy final-pop block. Legacy ids take the exact +existing path. An id that belongs to neither population is dropped as today. + +`EngineFrameRouter.on_response` currently returns no ownership result and the proxy always calls +the legacy `process_res` afterward. E5b must replace that tap behavior; merely relying on the +absence of a contract id from `_results` is not a sufficient registry contract. + +### Atomic classification and concurrency + +`route_response` uses the router's existing single `_lock` for the initial client-id lookup and +claim. Contract-owned classification consists of a binding in `_by_client` plus the bounded +recently-retired proxy-client-id structure. If it finds active or delivery-retained state, it may +release the lock while `normalize_response` snapshots the raw response, but it has already claimed +the population. On reacquisition it either advances that state or absorbs the response if another +thread ended it. If it finds a recently retired id, it counts and absorbs the late response and +returns `True`. It returns `False` when the id is in neither structure. + +This preserves the concurrency specification in `engine_client/router.py`: + +- The submit thread binds under the router lock before enqueue. +- The dispatch thread claims and transitions under the same lock. +- Consumer threads mark abort/close state under the same lock, and issue runtime abort outside it. +- The shutdown thread latches failure and ends states under the same lock. +- `_Delivery` keeps its inner lock and never calls back into the router. Router state is never + accessed while holding a delivery lock. + +The relevant races all have one outcome: + +- If dispatch wins before shutdown, it emits the response transition; shutdown sees the updated or + ended state. +- If shutdown wins, `fail_all` emits the typed ending; dispatch recognizes the retained contract + ownership and absorbs the late frame. +- If abort and final race, `abort_requested` and the final transition still meet at the existing + idempotent ending functions. +- A worker response cannot beat registration because the request is not enqueued until after the + client-id binding is visible. + +The shared allocator, rather than unbounded router ownership state, makes bounded classification +safe. Proxy client ids increase monotonically across both legacy and contract submission and are +never reused within a proxy lifetime; wrap or any collision is a submission error. Therefore, once +an id ages out of the recently-retired structure, a very late frame may return `False`, fall through +to the legacy `_results.get`, find no result, and be dropped exactly as an unknown legacy late frame +is dropped today. It cannot be delivered to a newer legacy request because no newer request can +receive that id. + +## 4. Abort, shutdown, and crash behavior across both populations + +### Per-request abort + +Legacy `GenerationResult.abort()` remains unchanged. Contract abort continues to look up the +frontend request id, set `abort_requested` under the router lock, and call +`proxy.abort_request(client_id)` outside that lock. A final response, a cancellation response, or a +later service failure resolves through the same exactly-once state machine. + +### Abort-all and typed poisoning + +`GenerationExecutorProxy._abort_all_requests` currently enumerates only `proxy._results`, so it +sees only legacy requests after E5b. Change the shutdown operation to cover both populations: + +1. Snapshot legacy `GenerationResult` objects as today. +2. Through `EngineService.fail_all(reason)`, atomically latch the service failure, mark every active + contract binding aborted, emit its appropriate typed ending, and return the active contract + client ids that need runtime cancellation. +3. Outside the router lock, send `CancellingRequest` for those contract client ids on a best-effort + basis. +4. Call `result.abort()` for the legacy snapshot. +5. Put the worker shutdown sentinel only after the submission/lifecycle lock has excluded new + enqueues. + +Poisoning is guaranteed; runtime cancellation is best effort because the worker may already have +crashed. `fail_all` must be idempotent, and the list of ids to cancel must be captured before the +bindings are pruned. Calling the current `router.fail_all` and then trying to enumerate +`_by_request` is too late. + +The current `EngineFrameRouter.fail_all` only emits endings; it does not call the abort function. +That must change at the `EngineService` level. It also fixes the current mismatch in +`LocalProcessEngineClient.close_client`: its docstring promises to abort this client's in-flight +requests, but the implementation only calls `router.fail_all("client closed")`, allowing engine +work to continue. + +### Every shutdown/crash entry + +`GenerationExecutorProxy.pre_shutdown` is the common non-blocking transition and already calls +`router.fail_all("executor shutdown")`. The implementation currently reaches `pre_shutdown` from: + +- explicit `GenerationExecutorProxy.shutdown`; +- `check_health` through `_drain_error_queue` or `_check_mpi_futures`; +- the background error monitor through those same helpers; +- the MPI-future done path after the monitor observes its queued/future error; +- `_handle_background_error`, which calls `shutdown` for a system error; and +- the registered interpreter/threading atexit callback. + +E5b replaces the direct router call with the combined EngineService fail-and-abort operation above. +The first `pre_shutdown` invocation owns the transition under the lifecycle lock; later invocations +are no-ops. This gives check-health-driven, monitor-driven, explicit, and exit-driven shutdown the +same behavior for both request populations. + +Client detach is separate from engine shutdown. `LocalProcessEngineClient.close_client` must call +an EngineService operation that poisons and aborts this client's active bindings without shutting +down a shared engine. V0 attaches only one service/client to a proxy; attachment of a second service +must be rejected rather than silently overwriting `_engine_frame_router` as +`attach_engine_frame_router` does today. + +### Enqueue and submission failures + +- A failure before binding (validation, conversion, dispatch-thread startup, shutdown check, or + duplicate request id) creates no state and no stream. +- A failure after binding but before a successful queue put calls + `on_submit_enqueue_failed(client_id)`, producing exactly one standalone `ErrorFrame`; no + `GenerationResult` exists to leak and no abort is sent for a request that was not enqueued. +- A system error discovered after queue put flows through `pre_shutdown` and `fail_all`. +- A per-request worker admission error returns as `ErrorResponse` and is exclusively routed by the + contract state machine. + +The current generic `GenerationExecutorProxy.submit` creates and registers its rank-0 result before +`request_queue.put` and does not remove it if put raises. E5b must not reproduce that stale-entry +behavior in `submit_contract`; changing the legacy behavior is a separate cleanup. + +## 5. `EngineService` factoring and remote-detach seam + +Introduce one server-side `EngineService` component in the engine-client package. It owns the +contract execution endpoint, not frontend rendering and not the legacy population. + +Its cohesive responsibilities are: + +- validate/translate `EngineRequest` with `engine_request_to_generation_request`; +- submit through the proxy's shared allocator and request queue without a rank-0 + `GenerationResult`; +- own `ContractRequestState`/`RequestBinding`, delivery, tombstones, durable id ownership, and the + router counters; +- exclusively route contract responses; +- implement request abort, client detach, and `fail_all`; +- expose capabilities/health and typed `IterationStatsBatch` / `KvCacheEventsBatch` query results; + and +- hold no tokenizer, HTTP request, formatter callback, event loop, or other frontend object. + +The current pieces map as follows: + +| Existing piece | E5b mapping | +|---|---| +| `EngineFrameRouter` | Becomes the state-machine core owned by `EngineService`; its lock remains the sole contract-state lock | +| `LocalProcessEngineClient` | Becomes a thin in-process client/transport adapter: pre-submit contract checks, calls into the service-facing proxy methods, and exposes `FrameStream`; direct executor stats/KV calls move behind the service | +| `engine_request_to_generation_request` call in `LocalProcessEngineClient.submit` | Moves into `EngineService.submit_contract` | +| `attach_engine_frame_router` | Replaced by a one-time `attach_engine_service` (or equivalent constructor injection); duplicate attachment is an error | +| `observe_submit` and pending object-identity matching | Removed from production submission; explicit `submit_contract` binds the assigned id directly | +| `on_submit_enqueue_failed` | Retained as an internal EngineService/router transition | +| `dispatch_result_task` response tap | Replaced by exclusive `EngineService.route_response` before legacy lookup | +| `pre_shutdown` direct `router.fail_all` | Replaced by `EngineService.fail_all` plus best-effort runtime cancellation | +| `LocalProcessEngineClient.get_stats/get_kv_events/health` direct proxy calls | Delegate to EngineService query methods that return the existing typed contract results | + +Removing the pending/object-identity protocol also fixes the service's active-count definition. +Today `EngineFrameRouter.active_request_count` adds `len(_by_request)` and `len(_pending)`, even +though `register_pending` inserts the same binding into both, so a pending submit is counted twice. +EngineService should count each non-ended binding once. + +The service should depend on a narrow runtime adapter supplied by the proxy: start dispatch, +allocate id, enqueue, abort by client id, fetch stats/KV events, and inspect health. It must not know +about ZMQ socket construction or MPI process management. + +This produces the later detach seam: + +```text +frontend EngineClient + | + | local calls now; encoded transport later + v +transport terminator + | + v +EngineService -> proxy runtime adapter -> worker +``` + +For V0, `LocalProcessEngineClient` is effectively the pass-through terminator. A later remote +transport terminator validates/decodes the same `EngineRequest`, submit metadata, abort, stats, KV, +and health operations and invokes the same EngineService methods. Nothing inside EngineService or +the worker-facing path changes. The ordered `stop_reasons` submit metadata must remain +primitive/callable-free and explicitly encoded by that terminator; `FrontendOutputConfig` itself +stays frontend-owned and is not passed wholesale into the engine. + +## 6. `LlmArgs` option and environment precedence + +Add a user-facing PyTorch-only field to `TorchLlmArgs`: + +```python +experimental_engine_client: bool = Field( + default=False, + description="Enable the experimental engine-client path for eligible serving requests.", + status="prototype", +) +``` + +The environment variable remains `TLLM_EXPERIMENTAL_ENGINE_CLIENT`. Resolve it once with this +precedence: + +| Environment | `TorchLlmArgs.experimental_engine_client` | Effective value | +|---|---:|---:| +| unset | `False` | `False` | +| unset | `True` | `True` | +| `0` | either | `False` | +| `1` | either | `True` | + +Thus the environment wins in both directions. Presence, not truthiness, decides whether it is an +override. Values other than the documented `0` and `1` should fail closed with a clear +configuration error rather than silently acting as false. + +`EngineClientServing.create_if_enabled` and `EngineClientConfig` must consume the same resolved +value. The current split behavior is unsafe for this requirement: +`engine_client_flag_enabled()` reads only the environment, while +`EngineClientConfig._validate_config` ignores the environment whenever `flag_enabled` is not +`None`. In particular, an explicit `flag_enabled=True` currently defeats an environment value of +`0`. Remove that ambiguity; `EngineClientConfig.flag_enabled` should be the already-resolved value +or should use the shared resolver unconditionally. + +Adding the field changes the generated `LLM.__init__` surface. The same change must update +`tests/unittest/api_stability/references/llm.yaml` with annotation `bool`, default `False`, and +status `prototype`. It must not be added to the committed/stable reference because it is a new +prototype option. API-stability and env-precedence tests must cover all four rows above, including +explicit-true/env-zero and explicit-false/env-one. + +## 7. Risks, required safeguards, and measurement plan + +### Client-id collision between contract and legacy traffic + +Both current populations ultimately use `GenerationExecutor._get_next_client_id` on the same +`GenerationExecutorProxy`; `BaseWorker.submit` preserves the id that the proxy assigned. E5b must +keep this single allocator. It must not add a router-local counter. + +The current allocator mutation is not explicitly synchronized and wraps at unsigned 64-bit. +Because legacy and contract submit threads can run concurrently, the proxy lifecycle/submission +lock must cover allocation for both paths. Make the allocator monotonic for the proxy lifetime and +reject unsigned-64 wrap rather than returning the wrapped value. Before binding, also reject any +collision with active legacy results, `_by_client`, or the bounded recently-retired id structure. +Proxy client ids are never reused in one proxy lifetime. The practical wrap interval is enormous, +but correctness must not depend on probability. + +### Late frames after cut-over + +Final/error removes detailed execution state. While the binding or delivery is retained, +`_by_client` still classifies late contract frames; after delivery retirement, move the client id +to the bounded recently-retired structure and use it only to count and absorb late/duplicate +frames. That structure is capped by `tombstone_limit` (`DEFAULT_TOMBSTONE_LIMIT` is 4096), evicting +oldest entries. Once an id is evicted, a very late frame may fall through to legacy lookup and be +dropped because the shared allocator never reuses the id during the proxy lifetime. Record claimed +late/duplicate frames and unknown-id legacy drops separately enough to diagnose runtime behavior. + +The submission cut-over itself must occur before any contract request is admitted. Do not switch a +request that already has a rank-0 `GenerationResult` to contract ownership mid-flight. Flag changes +apply when the LLM/proxy/service is constructed, not dynamically while it is serving. + +### Frontend request-id reuse + +Reject a frontend request id only while its prior binding is active or its delivery is retained. +Once that request has ended and the delivery is retired by consuming through the ending frame, +explicitly closing the stream, or bounded tombstone eviction, remove its request-id indexes and +allow the id to be submitted again. AC-4 requires this reuse-after-retirement path as a positive +replay case, not merely as an eviction side effect. A stale stream cannot overlap the reused id +because consumption or close retired it; a late runtime frame carries the old, never-reused proxy +client id and therefore cannot bind to the replayed request. Remove `_seen_request_ids`; retaining +frontend-id ownership for the service lifetime violates the bounded-state requirement. + +### Bounded retirement and memory baseline + +Router classification and delivery retention must remain bounded in request count. `_by_request` +contains only active bindings; `_by_client` contains active or delivery-retained bindings; +`_delivery_index` and any request-id tombstones contain only retained deliveries; the ordered +recently-retired proxy-client-id structure contains at most +`tombstone_limit` entries (4096 by default) and carries no delivery or request state. Neither +`_owned_client_ids` nor `_seen_request_ids`, nor an equivalent unbounded per-request history, is +permitted. + +Add a stress assertion with `N` substantially greater than `tombstone_limit`. After all `N` +requests complete and their deliveries retire, `_pending`, `_by_request`, `_by_client`, +`_delivery_index`, and request-id tombstones must return to their empty baseline; buffered delivery +memory must be released; and every recently-retired/accounting structure must be at or below its +configured bound. Repeating the stress cycle must not increase any per-request structure beyond +those baselines/bounds. This is the AC-4 memory-baseline assertion. + +### Rank-0 CPU and streamed-token measurement + +The performance claim for E5b is specifically rank-0 overhead removal, not GPU throughput. Measure +legacy and contract-native modes on the same commit, model, GPU, tokenizer, sampling parameters, +prompt/output lengths, `stream_interval=1`, and server configuration. Disable unrelated telemetry +and pin the rank-0 proxy process to the same CPU set. Alternate the mode order by repetition +(`A/B`, then `B/A`) so thermal/load drift is paired rather than confounded. + +Use at least these streaming concurrency points: 1, 8, 32, and 128, capped only if a validated +configuration has a lower maximum. At each point: + +- run at least 5 warm-up repetitions that are discarded; +- run at least 30 paired measured repetitions per mode; +- make each measured repetition contain at least 10,000 emitted token ids and last at least 30 + seconds, increasing request count if necessary; and +- reject/repeat a pair if request counts, emitted tokens, output-length distribution, or errors + differ between modes. + +Collect: + +- rank-0 process user plus system CPU seconds from process-scoped counters; +- rank-0 CPU seconds per completed request; +- rank-0 CPU microseconds per emitted token id (CPU delta divided by the sum of + `TokenDelta.new_token_ids`, not by frame count); +- proxy dispatch/service routing CPU time per response and per emitted token, with p50/p95/p99; +- peak and steady-state rank-0 RSS at each concurrency; and +- proof counters: zero rank-0 `GenerationResult` constructions/`proxy._results` entries for + contract requests, with legacy counts unchanged. + +Report the paired median difference and a two-sided 95% bootstrap confidence interval for every +concurrency point and primary CPU metric. Pre-register the acceptance threshold before running the +experiment; for gate ①c, use an upper 95% confidence bound of no more than a 5% regression in +rank-0 CPU seconds per emitted token at any concurrency, while expecting an improvement in CPU +and/or RSS from removing rank-0 `GenerationResult`. Do not claim an improvement when the 95% +interval crosses zero. Publish raw per-repetition data, environment/configuration, commit SHA, and +the analysis script so the result is repeatable. + +### Pre-registered plan-DEC-3 gate protocol for this loop + +For the plan-DEC-3 gate in this implementation loop, this subsection is authoritative. The fuller +protocol above remains the standard for publishable performance claims. + +- Measure streaming concurrencies `{1, 8, 32}`. +- At each concurrency, discard two warm-up pairs, then record at least 10 paired repetitions. Each + pair runs both legacy and contract-native modes, alternating order across pairs (`A/B`, then + `B/A`). +- The metric is rank-0 CPU seconds per emitted token id: rank-0 process user-plus-system CPU seconds + divided by the number of emitted token ids, never by response or frame count. +- For each concurrency, compute the per-pair contract-minus-legacy delta and use its median as the + decision statistic. Compute a percentile-bootstrap two-sided 95% confidence interval for that + paired median with at least 2000 resamples of the paired observations. +- The gate passes only when the upper confidence bound does not show a regression (that is, it is + less than or equal to zero for a contract-minus-legacy delta). Otherwise, the result requires an + explicit written acceptance. +- Generate the exit-report table directly from the recorded artifact values; never hand-copy + measurement values into the table. + +### Cut-over safety conditions + +E5b is not safe to enable until all of the following current behaviors are changed and tested: + +- contract submission no longer calls `GenerationExecutorProxy.submit` and never constructs or + registers a rank-0 `GenerationResult`; +- contract ownership is checked before the legacy `_results.get` path and claimed responses never + enter legacy delivery/pop logic; +- submission/id allocation is serialized with shutdown for both populations; +- contract classification uses `_by_client` plus only the bounded recently-retired id structure, + and very late untracked ids fall through safely because proxy client ids are never reused; +- frontend request-id reuse after request ending and delivery retirement is a tested positive replay + case, while reuse during active or retained delivery is rejected; +- after stress beyond the configured retirement bound, all per-request structures satisfy the AC-4 + empty-baseline/configured-bound assertions; +- `fail_all` both poisons contract streams and snapshots ids for best-effort runtime abort; +- `close_client` actually aborts its in-flight engine work; +- every `pre_shutdown` entry treats legacy and contract populations symmetrically; +- enqueue failure produces one typed ending without leaving rank-0 state; and +- worker-local `_results` registration and logprob/first-response behavior are retained and covered + by an IPC worker-path test. diff --git a/tensorrt_llm/executor/engine_client/E5B_EXIT_REPORT.md b/tensorrt_llm/executor/engine_client/E5B_EXIT_REPORT.md new file mode 100644 index 000000000000..0c98cbb61e70 --- /dev/null +++ b/tensorrt_llm/executor/engine_client/E5B_EXIT_REPORT.md @@ -0,0 +1,202 @@ +# E5b exit report (gate ①c) + +Status: cut-over delivered on branch `prototype_1`; V0-scope eligible OpenAI +streaming traffic runs contract-native. This report records the exit +evidence and enumerates the remaining Step-② (remote detach) work honestly — +the "transport-only" posture below is claimed for the **request path only**. + +## Cut-over evidence + +- **Zero rank-0 `GenerationResult` for contract traffic**: asserted + programmatically (`test_e5b_cutover.py::TestContractNativeSubmit:: + test_no_rank0_generation_result`); the router binding is the rank-0 + request state. Legacy traffic on the same proxy is untouched + (`TestExclusiveRouting`, `TestConcurrentMixedTraffic`). +- **Exclusive routing**: contract-owned client ids are claimed before the + legacy `_results` lookup and never enter legacy delivery or the + pop-on-final logic; ownership survives tombstone eviction, so late frames + are absorbed, never legacy-delivered. +- **Shutdown symmetry**: every `pre_shutdown` entry poisons contract streams + typed AND issues best-effort runtime cancellation alongside the legacy + abort-all (`TestShutdownSymmetry`). +- **Flag**: `TorchLlmArgs.experimental_engine_client` (prototype status, + `api_stability` suite green with the reference update); the env var + overrides in both directions, malformed values fail closed + (`TestFlagResolver`, four-row precedence matrix). + +## Real-GPU validation (the approved reproducible job) + +Artifact: **one merged** `gpu-e2e-report.json` across both pytest +invocations (main suite + destructive worker-crash test), containing EVERY +case result (plain streaming, stop-token, stop-string boundary, logprobs, +prompt logprobs, abort, stream close, usage, shadow overhead, measurement, +worker crash) AND the per-test pytest verdicts (recorded by the suite's +`conftest.py` hook); the loop's copy sits in the run directory. The +artifact's `commit` field records the **committed** code the suite ran on; +any later commits touch only documentation regenerated from this artifact. +Pinned configuration recorded in the artifact: commit SHA, model +`Qwen2.5-0.5B-Instruct` with a content-pinned revision +(`model_revision: sha256:0d1a2291cb93371d` over config + tokenizer +manifest), NVIDIA H100 (SM90 build, `cuda_architectures=90-real`), +execution backend `pytorch` and attention backend TRTLLM (recorded as +separate fields), greedy sampling (`top_k=1`), TP1 over the IPC proxy, KV +block reuse disabled (with partial block reuse, a repeated prompt gets +truncated context logits and therefore truncated prompt logprobs — a known +runtime limitation that the legacy logprobs harness also works around by +disabling reuse; the contract path surfaces that truncation as a typed +`logprob_mismatch` failure instead of delivering silently short values). +Reproducible commands are in the `test_gpu_e2e.py` module docstring. + +Results (11/11 green across the two invocations, per-test verdicts in the +artifact): + +- plain-streaming parity, stop-token, stop-string crossing a token + boundary, logprobs, **prompt logprobs** (position-by-position parity with + the legacy map-shaped output), abort mid-stream, usage — all parity + cases green; +- **stream-close cancellation is observed directly**: a spy on the + router's abort hook records the runtime cancellation issued by + `aclose()`, and the engine stays healthy and serviceable afterwards; +- **worker crash**: the test first gives the executor's own detector a + 120 s window after SIGKILL with **no** test interference. Recorded + evidence (artifact `worker_crash` case): `mpi_futures` never complete + and `_error_queue` stays empty (`native_signal_observed: false`). A + dedicated 300 s probe (`crash-native-detection-probe.{py,log}` in the + run directory) confirms this definitively: the worker process is + observed dead at t=1.0 s, the mpi4py `_manager_spawn` pool thread stays + alive, and for the full 300 s every future reports `done=False` with an + empty error queue while the proxy's error monitor polls with nothing to + observe — the propagation gap is in the mpi4py-futures layer beneath + the executor, and a companion probe + (`crash-legacy-gap-probe.{py,log}`) shows the LEGACY path has the + identical gap (a post-kill legacy request receives no result or typed + failure either), i.e. this is pre-existing runtime behavior, not a + contract-path deficiency, and its remediation belongs to the Step-② + process-lifecycle work already inventoried below. The test then + exercises everything the executor CAN own — the failure channel + (`_error_queue` → monitor → `pre_shutdown` → typed poisoning) — via a + **recorded, justified** injection: typed ending in **4.85 s**, inside + the frozen DEC-4 bound measured from the point the failure channel + observes the crash. The mode, window, and justification are stored in + the artifact so this fallback is auditable. + (Harness note: after the test passes and the artifact is written, the + pytest interpreter can hang at exit on the SIGKILLed MPI runtime's + teardown; the documented outer `timeout` reaps it and its 124/143 exit + code is expected — the verdict is the pytest summary and the artifact.) + +## Shadow-overhead report (draft DEC-5, report-only) + +Post-cutover, routing is exclusive — a dual-consumption shadow tap no +longer exists — so the shadow number that still exists is the cost the +attached contract machinery imposes on **legacy** traffic (the per-response +contract-population check in the dispatch loop). Measured at concurrency 8 +(2 warm-ups, 10 reps per arm, arms sequential because attach is +irreversible): **+2.79 %** legacy rank-0 CPU per streamed token in the +final run (+0.02 % in the earlier run — the sequential-arm design makes +this metric drift-sensitive, which is why it is report-only per DEC-5 and +carries no claim; artifact `shadow_overhead` case has the raw per-rep +values for both arms). + +## Rank-0 measurement (contract DEC-5 / plan DEC-3 rule) + +The pre-registered loop protocol of `E5B_CUTOVER_DESIGN.md` §7 was +executed: streaming concurrencies {1, 8, 32}, two discarded warm-up pairs, +alternating legacy/contract pairs per concurrency (protocol minimum 10; the +final run uses 30), metric = rank-0 process CPU seconds per emitted token +id, decision statistic = median paired delta with a percentile-bootstrap +two-sided 95 % CI (2000 resamples, fixed seed). Full disclosure: the +protocol ran twice — a first 10-pair run (c=1 −19.3 % [−50.5, −27.5] µs; +c=8 −1.25 % [−2.54, +0.19] µs; c=32 −5.5 % [−4.53, −2.58] µs) and this +final 30-pair run on the committed final code; the c=8 upper bound crossed +zero in BOTH runs, so more repetitions did not (and were not used to) +manufacture significance. Raw per-repetition values live in the artifact +(`rank0_measurement.rows`); the table below is **generated from the +artifact** by the test itself (`rank0_measurement.md`) and inserted here +mechanically, never hand-copied: + + +| concurrency | pairs | legacy median | contract median | Δ median | Δ% | 95% CI of Δ (bootstrap) | +|---|---|---|---|---|---|---| +| 1 | 30 | 189.18 | 157.45 | -33.81 | -17.87% | [-36.82, -28.80] | +| 8 | 30 | 60.80 | 60.01 | -1.16 | -1.91% | [-2.31, +0.30] | +| 32 | 30 | 63.56 | 60.79 | -3.06 | -4.82% | [-3.60, -2.28] | + + +Gate rule application ("upper 95 % CI bound must not show a regression, +otherwise an explicit written acceptance"): + +- concurrency 1: upper bound −28.80 µs/token → **pass** (clear + improvement, CI does not cross zero); +- concurrency 32: upper bound −2.28 µs/token → **pass**; +- concurrency 8: median delta −1.16 µs/token (−1.91 %) but the upper bound + is **+0.30 µs/token (+0.49 %)** — the CI marginally crosses zero, so the + automatic gate does not pass at this point and, per the pre-registered + rule, this concurrency **requires an explicit written acceptance** + (PENDING — to be granted or refused by the plan owner; not + self-accepted here). The point estimate is an improvement in both runs + and no regression is detectable at any concurrency. + +No performance *improvement claim* is made where the CI crosses zero +(concurrency 8). The structural change this loop stands on is the removed +rank-0 `GenerationResult` construction/registration per contract request. + +## Flag-off sweep (AC-8 final evidence, Round 1) + +Flag-off (`TLLM_EXPERIMENTAL_ENGINE_CLIENT` unset), `LLM_MODELS_ROOT` +populated (Qwen2.5-0.5B-Instruct + TinyLlama-1.1B-Chat-v1.0): + +- `tests/unittest/executor/` (excluding `engine_client/`, **including** + `test_rpc.py`, `test_rpc_proxy.py`, `test_rpc_worker.py`): **206 passed / + 1 failed / 8 skipped**. The single failure + (`test_rpc.py::TestRpcCorrectness::test_incremental_task_async`) fails + identically at the base commit with base `tensorrt_llm/` sources — + pre-existing, base-equivalence recorded. +- `tests/unittest/executor/engine_client/` (CPU): **370 passed** clean. +- `tests/unittest/llmapi/test_llm_args.py`: **190 passed / 3 failed** when + run concurrently with the GPU-holding executor sweep; the 3 failures + (`test_build_config_from_engine`, `test_runtime_sizes`, + `test_model_kwargs_with_num_hidden_layers`) all **pass in an isolated + rerun at head** with the GPU idle (and also pass at the base commit) — + GPU-contention flakes of the concurrent run, not code. +- `tests/unittest/api_stability/`: **64 passed**. + +## Remaining Step-② work (the honest inventory) + +Request path — enumerable as transport only: + +- Socket + codec-on-wire delivery (`encode`/`decode` are frozen and + round-trip-tested; nothing is "in-process only"). +- ROUTER/DEALER ingress, `(client_identity, request_id)` identity and the + three-ID mapping; `event_seq` end-to-end delivery semantics (stamped from + V0; e2e duplicate/gap detection is ② scope per contract draft DEC-3). +- Ready-handshake delivery of `FrontendModelContext` (built locally today; + consumers already read only through it on the contract path). +- Heartbeat/liveness, bounded queues + credit backpressure, and + disconnect/crash semantics symmetric across frontend / engine server / + worker. +- Process lifecycle: engine process owns the MPI session; shutdown + ordering; the serve entrypoint. + +NOT transport-only — server-wide surfaces still coupled to the in-process +`LLM` (unchanged for legacy and fallback traffic): + +- Engine-side sampling normalization (V0 deliberately runs it + frontend-side; the remote target moves it after receipt). +- The serve layer's remaining reaches: ~29 `generator.args.*` reads, the + `generator._executor.resource_governor_queue` reach, `_hf_model_dir`, + and the `_torch.pyexecutor.config_utils` import — all on the legacy/ + control paths, not the contract request path. +- Control endpoints and model lifecycle: health/health_generate, metrics, + KV-cache events surface, weight-update/memory operations, cluster + registration, disagg role endpoints. +- Non-streaming and the public `LLM` API (stay legacy in V0 by decision), + ineligible-request fallback (requires the in-process legacy path by + design), chat templating and input preprocessing (server-side). + +## ②-readiness guardrails (held throughout V0) + +- Every wire type round-trips the strict codec (golden fixtures checked in). +- The `EngineClient`/`EngineService` surfaces expose no proxy or executor + internals; frame payloads and binding state are primitives-only. +- `EngineService` is factored so a remote transport terminator fronts it + without internal changes (see `E5B_CUTOVER_DESIGN.md` §5). diff --git a/tensorrt_llm/executor/engine_client/ENGINE_CONTRACT.md b/tensorrt_llm/executor/engine_client/ENGINE_CONTRACT.md new file mode 100644 index 000000000000..73c0c6bc32f0 --- /dev/null +++ b/tensorrt_llm/executor/engine_client/ENGINE_CONTRACT.md @@ -0,0 +1,263 @@ +# Engine Client Contract (V0, experimental) + +Typed, versioned, iteration-level contract between a serving frontend and the +generation engine. Gated by `TLLM_EXPERIMENTAL_ENGINE_CLIENT=1`; with the flag +off, legacy behavior is byte-identical (see Divergence notes for the single +approved exception). Protocol version: `ENGINE_CONTRACT_VERSION = 1`. + +This document is the living specification companion: it tracks the scope +matrix, the rejection→test table, and every divergence from the design draft +(Part V of `serve → runtime layering`, frozen decisions draft DEC-1..9). It is +updated by every change to this package. + +## Package layout + +| Module | Owns | Status | +|---|---|---| +| `contract.py` | wire types + construction-time validation + `validate_no_callables` | delivered | +| `codec.py` | strict msgpack encode/decode for every wire type | delivered | +| `conversion.py` | legacy inputs → `EngineRequest` + `FrontendOutputConfig`; first capability gate; eligibility matrix; compositional runtime translation | delivered | +| `envelope.py` | raw worker response → normalized envelope; typed rejection of out-of-scope shapes | delivered | +| `router.py` | proxy-attached fork target; per-request exactly-once terminal state machine; concurrency spec | delivered | +| `local_client.py` | in-process `EngineClient`; setup gates; pre-submit checks; async `FrameStream` | delivered | +| `assembler.py` | frontend response assembly: accumulation, detok, stop handling, presentation | delivered | +| `invariants.py` | pass-through stream-shape checker (test utility) | delivered | +| `../../serve/engine_client_serving.py` | flag-on serving path: eligibility fallback + counters, model context, SSE via the real formatters | delivered | +| `service.py` | contract execution endpoint: contract-native submit (no rank-0 `GenerationResult`), exclusive routing, abort/detach/fail_all, typed control plane | delivered | +| `E5B_CUTOVER_DESIGN.md` / `E5B_EXIT_REPORT.md` | cut-over design (Codex-reviewed) and the gate-①c exit evidence + Step-② remaining-work inventory | delivered | + +## Scope matrix (V0) + +**Supported:** PyTorch backend · IPC-proxy transport (`GenerationExecutorProxy`) +· `num_postprocess_workers=0` · streaming · n=1 · pre-tokenized input · plain +text · chosen-token logprobs (`logprobs <= 1`, `prompt_logprobs <= 1`) · +stop-token / stop-token-sequence / stop-string (incl. stop-string ⇒ +engine-abort control edge) · abort · adapter-observed usage (incl. cached-token +accounting) · opt-in flag with legacy byte-identical when off. + +**Config-gated at client construction (typed config error):** non-PyTorch +backends (`_autodeploy` explicitly rejected — do not gate on +`_is_pytorch_backend`) · RPC / Ray / single-process transports · +`num_postprocess_workers > 0` · global `post_processor_hook` · +`speculative_config` · early-first-token response mode · topology beyond the +validated set (TP1 over the IPC proxy) · `trust_remote_code` tokenizers · flag +off. + +**Rejected pre-submit (typed request rejection):** multimodal · n>1 / best_of · +beam search · LoRA · disaggregated serving · scheduling params · per-request +logits-processor callables · response-formatter callables (`PostprocParams`) · +`logprobs > 1` · `prompt_logprobs > 1` · non-streaming mode · guided decoding +(schema field exists, capability-gated until the engine-path validation) · +`ignore_eos` · bad words · `min_p` · prompt-ignore-length · embedding/logit +bias · output flags beyond logprobs · trace-header-bearing requests (avoids +silent telemetry loss) · `cache_salt` · thinking-token-budget · BART-family +models (normalization would create a callable) · `truncate_prompt_tokens` · +`echo` · `top_logprobs`. + +The machine-readable eligibility matrix (delivered with `conversion.py`) is +the authoritative per-field classification; this section is its prose summary. + +## Wire types + +Defined in `contract.py`; every type is a frozen, primitives-only, +callable-free dataclass carrying `protocol_version`, and every type +round-trips `codec.encode`/`codec.decode` (nothing is "in-process only"). +`OutputFrame = TokenDelta | Terminal | RequestComplete | ErrorFrame`; every +frame carries per-request monotonic `event_seq` starting at 0 (draft DEC-3). +`Terminal` is the sole carrier of finish/stop reasons; `TokenDelta` is +non-empty by construction (a no-token response produces no frame, draft +DEC-6); `RequestComplete` usage is adapter-observed (draft DEC-2) with +`cached_tokens` as the final cached-token accounting (divergence note 5). +`FrontendOutputConfig` never crosses the contract; its +`stop_sequence_reasons` is an ordered association list (divergence note 11). +`FrontendModelContext`/`TokenizerSpec` are frozen from V0 (divergence note +10) even though their delivery is local until the remote detach. + +## Codec rules + +Strict msgpack (§2.4 of the draft): any value outside +`None | bool | int | float | str | list | dict` is a typed encode error — no +pickle fallback under any flag. Canonical encoding: dataclass fields in +declaration order, mapping keys sorted — equal values produce identical bytes +(golden fixtures in `tests/unittest/executor/engine_client/test_codec.py`). +Decode rejects with typed errors (`DecodeError.reason`): `unknown_kind`, +`unknown_field`, `missing_field`, `invalid_content`, `version_unsupported`, +`limit_exceeded`, `malformed`. Resource limits: 16 MiB message, depth 16, +2 Mi array items, 4 Ki map entries, 4 Mi-char strings, signed 64-bit ints, +finite floats, msgpack ext/bin rejected. Limits live in `contract.py` / +`codec.py` as constants. + +## Draft frozen decisions (DEC-1..9) — implementation status + +| Decision | Summary | Status | +|---|---|---| +| draft DEC-1 | chosen-token logprobs only; `logprobs > 1` rejected pre-submit; both runtime shapes normalized by exact token-id lookup | implemented (`conversion`, `envelope`; typed `logprob_mismatch` errors) | +| draft DEC-2 | usage adapter-observed; runtime-native usage is a later additive field | implemented (router-observed counts; user-visible usage counts assembled tokens — divergence note 14) | +| draft DEC-3 | exactly-once = server-side emission; `event_seq` stamped from V0 | implemented (router stamps per-request monotonic `event_seq` from 0) | +| draft DEC-4 | failure ownership matrix (worker crash ≤ ~5 s; per-stream bounds; early close ⇒ abort) | implemented (fail_all latch on `pre_shutdown`; `FrameStream.aclose()` ⇒ abort-if-incomplete); mechanism amended — see divergence note 7 | +| draft DEC-5 | cut-over at E5b; shadow-mode numbers are not performance claims | implemented (contract-native submit via `EngineService`; measurement + verdict in `E5B_EXIT_REPORT.md`) | +| draft DEC-6 | runtime-started predicate; `ErrorFrame` iff nothing runtime-started | implemented (router state machine; replay-matrix tested) | +| draft DEC-7 | native stop-token-sequence carrier on the worker-facing request path | implemented (`GenerationRequest.stop_token_sequences` merged by the worker's stop-word derivation) | +| draft DEC-8 | worker metrics enum keys convert via `.value` to `Mapping[str, float]` | implemented (`envelope`, typed `metrics_mismatch` errors) | +| draft DEC-9 | `TIMED_OUT` crosses as `Terminal(error, stop_reason="timeout")`; assembler renders legacy-visible reasons | implemented (router table + assembler presentation map; wire markers not presented as user stop_reason) | + +## Divergence notes (amendments to the draft, all converged at plan time) + +1. ①c cut-over scope = eligible OpenAI streaming chat/completions only; the + public LLM API reroute is future work. +2. Ineligible requests under the flag fall back to legacy transparently, with + per-axis observability counters; conversion still rejects typed internally. +3. V0 sampling normalization runs frontend-side exactly once (equivalent of + `LLM._prepare_sampling_params` minus callable-producing steps); + engine-side normalization is the remote-detach target. +4. `pad_id` added to `EngineSamplingConfig` (additive to draft §2.2) so the + runtime request is reproducible from the encoded `EngineRequest` alone. +5. Cached-token accounting: reserved `TokenDelta.metrics` key + `cached_tokens` pre-completion + typed `RequestComplete.cached_tokens` + final field. Absent = never reported; 0 = reported zero; final must equal + the last per-delta value when any delta carried the key; a zero-token + request may report the value only on `RequestComplete`. +6. `FrameStream` is async-first with explicit `aclose()`. +7. Fatal-error propagation is a non-consuming sticky latch triggering + `router.fail_all()`; the router never drains the shared `_error_queue` + (draft DEC-4 semantics and bounds unchanged). +8. The worker's `_compute_pytorch_prompt_logprobs` gains a guard so a + no-token `CANCELLED` final degrades gracefully instead of asserting. + **This is the sole approved flag-off behavior divergence** (the affected + path crashes today); it has its own regression test. +9. Config gates (client construction) are distinct from per-request + rejections; see the scope matrix. +10. `FrontendModelContext` + `TokenizerSpec` schemas frozen in V0 with full + provenance; `trust_remote_code` tokenizers are config-gated out. +11. `FrontendOutputConfig.stop_sequence_reasons` is an ordered association + list: configuration order, first match wins, collisions resolved by + order. +12. "Transport-only remote detach" is claimed for the request path only; the + exit report enumerates all remaining reaches and control surfaces. +13. The post-verification ZMQ transport spike is out of scope. +14. User-visible usage counts the assembled (post-trim) tokens for legacy + parity; `RequestComplete.completion_tokens` on the wire stays the raw + adapter-observed engine count (the draft DEC-2 posture). +15. The legacy cross-chunk stop-string prefix leak is preserved (parity by + oracle); a deliberate semantic fix is future work. +16. Prompt logprobs under KV block reuse: with (partial) block reuse + enabled, the runtime returns context logits — and therefore prompt + logprobs — only for the non-reused prompt suffix (a known runtime + limitation; `llm_return_logprobs_test_harness` disables reuse for the + same reason). Legacy delivers the truncated list silently; the contract + path ends the request with a typed `logprob_mismatch` failure instead + of delivering silently short values (the truncation check is + unit-tested in `test_envelope.py`). The GPU suite therefore pins + `enable_block_reuse=False` as the supported configuration for + prompt-logprob traffic (`test_gpu_e2e.py` fixture note). + +## Rejection → test table + +One row per scope-matrix rejection axis; every row must have a passing test +and every rejection test must have a row (bidirectional check at audit time). + +| Axis | Rejection point | Typed error | Test | +|---|---|---|---| +| unknown wire kind | codec decode | `DecodeError(unknown_kind)` | `test_codec.py::TestDecodeRejects::test_unknown_kind` | +| unknown field | codec decode | `DecodeError(unknown_field)` | `test_codec.py::TestDecodeRejects::test_unknown_field` | +| missing required field | codec decode | `DecodeError(missing_field)` | `test_codec.py::TestDecodeRejects::test_missing_required_field` | +| wrong item type / bool-as-int | codec decode | `DecodeError(invalid_content)` | `test_codec.py::TestDecodeRejects::test_wrong_item_type`, `test_bool_is_not_int` | +| newer protocol version (top-level + nested) | codec decode | `DecodeError(version_unsupported)` | `test_codec.py::TestDecodeRejects::test_newer_protocol_version`, `test_nested_version_check` | +| NaN/Inf | construction + codec | `ContractConstructionError` / `EncodeError(not_finite)` / `DecodeError(invalid_content)` | `test_contract.py::TestConstructionRejects`, `test_codec.py::TestEncodeRejects::test_smuggled_nan_rejected`, `TestDecodeRejects::test_nan_in_payload` | +| msgpack ext / bin / malformed | codec decode | `DecodeError` | `test_codec.py::TestDecodeRejects::test_extension_type_rejected`, `test_bytes_payload_rejected`, `test_malformed_bytes` | +| oversized message / nesting depth | codec | `EncodeError/DecodeError(limit_exceeded)` | `test_codec.py::TestDecodeRejects::test_oversized_message`, `test_nesting_depth_limit` | +| callable / object / bytes on the wire | construction + codec encode | `ContractConstructionError` / `EncodeError(non_primitive)` | `test_contract.py::TestNoCallables`, `test_codec.py::TestEncodeRejects` | +| non-registered type encode | codec encode | `EncodeError(unknown_type)` | `test_codec.py::TestEncodeRejects::test_unregistered_type` | +| `n_gt_1` (`n`, `best_of`) | conversion pre-submit | `RequestIneligibleError("n_gt_1")` | `test_conversion.py::TestRejections` case `n_gt_1`; facade fallback + counter via `test_serving_sse_full.py::TestEveryFallbackAxisThroughFacade` | +| `beam_search` (`use_beam_search`, `beam_width_array`, diversity/length/early-stopping) | conversion pre-submit | `RequestIneligibleError("beam_search")` | `TestRejections` case `beam_search`; facade recipe | +| `top_logprobs` (`logprobs > 1`) | conversion pre-submit | `RequestIneligibleError("top_logprobs")` | `TestRejections` case `top_logprobs`; facade recipe | +| `prompt_top_logprobs` (`prompt_logprobs > 1`) | conversion pre-submit | `RequestIneligibleError("prompt_top_logprobs")` | `TestRejections` case `prompt_top_logprobs`; facade recipe | +| `logprobs_mode` (non-RAW) | conversion pre-submit | `RequestIneligibleError("logprobs_mode")` | `test_conversion.py::TestRejections::test_logprobs_mode_rejected`; facade recipe | +| `logits_processor` (incl. batched) | conversion pre-submit | `RequestIneligibleError("logits_processor")` | `TestRejections` case `logits_processor`; facade recipe | +| `embedding_bias` | conversion pre-submit | `RequestIneligibleError("embedding_bias")` | `TestRejections` case `embedding_bias`; facade recipe | +| `bad_words` (`bad`, `bad_token_ids`) | conversion pre-submit | `RequestIneligibleError("bad_words")` | `TestRejections` case `bad_words`; facade recipe | +| `ignore_eos` | conversion pre-submit | `RequestIneligibleError("ignore_eos")` | `TestRejections` case `ignore_eos`; facade recipe | +| `min_p` | conversion pre-submit | `RequestIneligibleError("min_p")` | `TestRejections` case `min_p`; facade recipe | +| `top_p_extras` (`top_p_min`/`top_p_reset_ids`/`top_p_decay`) | conversion pre-submit | `RequestIneligibleError("top_p_extras")` | `TestRejections` case `top_p_extras`; facade recipe | +| `no_repeat_ngram` | conversion pre-submit | `RequestIneligibleError("no_repeat_ngram")` | `TestRejections` case `no_repeat_ngram`; facade recipe | +| `prompt_ignore_length` | conversion pre-submit | `RequestIneligibleError("prompt_ignore_length")` | `TestRejections` case `prompt_ignore_length`; facade recipe | +| `return_logits` (context/generation/encoder/additional outputs) | conversion pre-submit | `RequestIneligibleError("return_logits")` | `TestRejections` case `return_logits`; facade recipe | +| `exclude_input_from_output` (non-default) | conversion pre-submit | `RequestIneligibleError("exclude_input_from_output")` | `TestRejections` case `exclude_input_from_output`; facade recipe | +| `truncate_prompt_tokens` | conversion pre-submit | `RequestIneligibleError("truncate_prompt_tokens")` | `TestRejections` case `truncate_prompt_tokens`; facade recipe | +| `lookahead_config` | conversion pre-submit | `RequestIneligibleError("lookahead_config")` | `TestRejections` (direct conversion test); facade recipe | +| `thinking_token_budget` | context-only normalization | `RequestIneligibleError("thinking_token_budget")` | `TestRejections` case; facade recipe | +| `bart_forced_tokens` (`model_type == "bart"`) | context-only normalization | `RequestIneligibleError("bart_forced_tokens")` | `test_serving_e5a.py::TestContextOnlyBoundary` (facade, counter) | +| `echo` | conversion pre-submit | `RequestIneligibleError("echo")` | `TestRejections` case `echo`; facade recipe | +| `non_streaming` | conversion pre-submit | `RequestIneligibleError("non_streaming")` | `TestRejections` case `non_streaming`; facade recipe (`streaming=False`) | +| `multimodal` | conversion pre-submit | `RequestIneligibleError("multimodal")` | `TestRejections` case `multimodal`; facade recipe | +| `lora` | conversion pre-submit | `RequestIneligibleError("lora")` | `TestRejections` case `lora`; facade recipe | +| `prompt_adapter` | conversion pre-submit | `RequestIneligibleError("prompt_adapter")` | `TestRejections` case `prompt_adapter`; facade recipe | +| `disaggregated` | conversion pre-submit | `RequestIneligibleError("disaggregated")` | `TestRejections` case `disaggregated`; facade recipe | +| `scheduling_params` | conversion pre-submit | `RequestIneligibleError("scheduling_params")` | `TestRejections` case `scheduling_params`; facade recipe | +| `conversation_params` | conversion pre-submit | `RequestIneligibleError("conversation_params")` | `TestRejections` case `conversation_params`; facade recipe | +| `postproc_params` | conversion pre-submit (defense in depth — unreachable from the facade: both endpoints construct `PostprocParams` strictly after the contract branch) | `RequestIneligibleError("postproc_params")` | `TestRejections` case `postproc_params` | +| `trace_headers` (requests actually carrying one) | conversion pre-submit | `RequestIneligibleError("trace_headers")` | `TestRejections` case `trace_headers`; facade recipe; empty-mapping normalization tested in `test_serving_e5a.py` | +| `cache_salt` | conversion pre-submit | `RequestIneligibleError("cache_salt")` | `TestRejections` case `cache_salt`; facade recipe | +| `query_token_ids` | conversion pre-submit | `RequestIneligibleError("query_token_ids")` | `TestRejections` case `query_token_ids`; facade recipe | +| `encoder_input` | conversion pre-submit | `RequestIneligibleError("encoder_input")` | `TestRejections` case `encoder_input`; facade recipe | +| `priority` | conversion pre-submit | `RequestIneligibleError("priority")` | `TestRejections` case `priority`; facade recipe | +| `kv_cache_retention` | conversion pre-submit | `RequestIneligibleError("kv_cache_retention")` | `TestRejections` case `kv_cache_retention`; facade recipe | +| config gate: non-PyTorch backend (`_autodeploy` explicit) | client construction | `EngineClientConfigError` | `test_router_client.py::TestSetupGates` case `backend` | +| config gate: non-IPC transport (RPC/Ray/single-process) | client construction | `EngineClientConfigError` | `TestSetupGates` case `transport` | +| config gate: `num_postprocess_workers > 0` | client construction | `EngineClientConfigError` | `TestSetupGates` case `postproc` | +| config gate: global `post_processor_hook` | client construction | `EngineClientConfigError` | `TestSetupGates` case `post_processor_hook` | +| config gate: `speculative_config` | client construction | `EngineClientConfigError` | `TestSetupGates` case `speculative` | +| config gate: early-first-token response mode | client construction | `EngineClientConfigError` | `TestSetupGates` case `early_first_token` | +| config gate: topology beyond TP1-over-IPC | client construction | `EngineClientConfigError` | `TestSetupGates` case `topology` | +| config gate: `trust_remote_code` tokenizer | client construction | `EngineClientConfigError` | `TestSetupGates` case `trust_remote_code`; reload refusal in `load_tokenizer_from_spec` | +| config gate: flag off | client construction | `EngineClientConfigError` | `TestSetupGates` case `flag`; precedence matrix in `test_e5b_cutover.py::TestFlagResolver` | + +The machine-readable sources are `conversion.ELIGIBILITY_MATRIX` and the +setup-gate list in `local_client._validate_config`; bidirectional +completeness between the matrix and the parametrized rejection cases is +enforced by `test_every_ineligible_axis_has_a_rejection_case`, +`TestEligibilityMatrixCompleteness`, and (facade side) +`test_every_request_level_axis_is_driven_through_the_facade`. +| unsupported response shapes (C++-engine, postproc-parallel, unknown) | envelope | `EnvelopeError(reason)` | `test_envelope.py::TestRejectedShapes` | +| pre-submit request checks (newer protocol, forged required_features, capability subset, duplicate/reused id) | client submit | `RequestRejectedError` | `test_router_client.py::TestPreSubmitChecks`, `TestStreamLifecycle` | +| serving fallback (ineligible under the flag ⇒ legacy + per-axis counter) | serving facade | counted fallback; the request is served by the unchanged legacy path (user-visible equivalence is inherited from legacy, exercised end-to-end by the GPU suite) | `test_serving_e5a.py::TestFallbackCounters` | +| submission during shutdown / after service close | service submit | `EngineServiceError` | `test_e5b_cutover.py::TestContractNativeSubmit::test_submission_rejected_during_shutdown` | +| duplicate engine-service attachment | proxy attach | `RuntimeError` | `test_e5b_cutover.py::TestContractNativeSubmit::test_second_service_attachment_rejected` | +| abort unknown id / double stream open / unknown stream | client ops | `UnknownRequestError` / `RequestRejectedError` | `test_router_client.py::TestAbort`, `TestStreamLifecycle` | +| malformed stop-reason carrier (non-primitive) | router bind | `RouterError` | `test_audit_remediations.py::TestRouterRemediations::test_malformed_stop_reasons_rejected_at_bind` | +| unprepared / untokenized conversion inputs | conversion | `ConversionError` | `test_conversion.py::TestConvertRequest::test_unprepared_params_rejected`, `test_untokenized_stop_strings_rejected` | +| multi-sequence response (`sequence_index != 0`) | router | request fails typed (`router_error`) | `test_audit_remediations.py::TestRouterRemediations::test_nonzero_sequence_index_fails_request` | +| mixed/non-string mapping keys on encode; non-finite metrics at the envelope | codec / envelope | `EncodeError(non_primitive)` / `EnvelopeError(metrics_mismatch)` | `test_audit_remediations.py::TestCodecRemediation`, `TestEnvelopeRemediation` | + +## Normalization posture (V0) + +Sampling normalization for the contract path runs frontend-side, exactly +once, at the CONTEXT-ONLY boundary — never via live model-config reaches: + +- `conversion.prepare_sampling_params_from_context` is the sole normalizer + on the serving contract path: EOS/PAD defaulting, generation-config stop + ids, and model-type gates come from the frozen `FrontendModelContext`; + stop strings are tokenized exactly once with the spec-reloaded + tokenizer. Callable-producing legacy steps (BART forced tokens, thinking + budget) raise `RequestIneligibleError`. +- Both endpoints capture a PRISTINE `SamplingParams` copy before any live + `LLM.preprocess` / `input_processor` call and hand that copy to + `try_stream`, which normalizes it unconditionally through the + context-only boundary. The live-prepared original continues to serve the + legacy/fallback path unchanged. A negative test poisons + `prepare_sampling_params` and `LLM._prepare_sampling_params` and runs a + full contract stream + (`test_serving_sse_full.py::TestNoLiveNormalizationReaches`). +- `conversion.prepare_sampling_params` (live-input variant) remains only + for direct library callers outside serving (e.g. the GPU suite driving + the client directly); the serving glue never calls it. +- Downstream of conversion, the compositional translation + (`engine_request_to_generation_request`) rebuilds the worker-facing + request purely from the encoded `EngineRequest`; the synthetic params + carry no stop strings, so nothing can re-run `_setup` (guard-tested). +- Tokenizer provenance is enforced at reload: `load_tokenizer_from_spec` + verifies every manifest content hash and the recorded added/special-token + configuration, failing typed + (`test_serving_sse_full.py::TestTokenizerSpecProvenance`). + +Engine-side normalization remains the remote-detach target (future work). diff --git a/tensorrt_llm/executor/engine_client/__init__.py b/tensorrt_llm/executor/engine_client/__init__.py new file mode 100644 index 000000000000..6ee00773e430 --- /dev/null +++ b/tensorrt_llm/executor/engine_client/__init__.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Experimental typed engine-client contract (gated by TLLM_EXPERIMENTAL_ENGINE_CLIENT). + +Internal package: not part of the public LLM API surface. See +``ENGINE_CONTRACT.md`` for the contract specification, scope matrix, and +divergence notes. +""" + +from .codec import (CodecError, DecodeError, EncodeError, decode, # noqa: F401 + encode) +from .contract import (ENGINE_CONTRACT_VERSION, ContractConstructionError, # noqa: F401 + ContractError, EngineCapabilities, EngineHealth, + EngineRequest, EngineSamplingConfig, ErrorFrame, + FrontendModelContext, FrontendOutputConfig, + GuidedDecodingSpec, IterationStatsBatch, + KvCacheEventsBatch, OutputFrame, RequestComplete, + Terminal, TokenDelta, TokenizerSpec, + validate_no_callables) + +__all__ = [ + "ENGINE_CONTRACT_VERSION", + "ContractError", + "ContractConstructionError", + "CodecError", + "EncodeError", + "DecodeError", + "encode", + "decode", + "EngineCapabilities", + "EngineRequest", + "EngineSamplingConfig", + "GuidedDecodingSpec", + "TokenDelta", + "Terminal", + "RequestComplete", + "ErrorFrame", + "OutputFrame", + "FrontendOutputConfig", + "TokenizerSpec", + "FrontendModelContext", + "EngineHealth", + "IterationStatsBatch", + "KvCacheEventsBatch", + "validate_no_callables", +] diff --git a/tensorrt_llm/executor/engine_client/assembler.py b/tensorrt_llm/executor/engine_client/assembler.py new file mode 100644 index 000000000000..0f395f18fdc2 --- /dev/null +++ b/tensorrt_llm/executor/engine_client/assembler.py @@ -0,0 +1,368 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Frontend response assembly: detokenization, stop handling, presentation. + +``FrontendResponseAssembler`` consumes the typed frame stream for one +request and produces presentation-ready updates. All formatting lives here, +on the frontend — no callable ever crosses to the engine. + +Semantics (legacy-parity by design; the replay oracle is the gate): + +- **Incremental detokenization** uses the tokenizer's incremental decoder + and is O(total tokens) unconditionally: tokenizers without incremental + support are rejected at setup (never a per-delta full re-decode). The + single full re-decode after an engine stop-word trim runs at most once + per request. +- **Stop-string detection** mirrors the legacy algorithm: the ordered stop + reasons are checked in configuration order and the first match wins; the + cumulative text is trimmed at the match (honoring + ``include_stop_str_in_output``), the user-visible finish reason becomes + ``"stop"`` with the string as ``stop_reason``, and the engine abort is + issued as an **internal control action** — the user-visible finish reason + never becomes ``"cancelled"`` for a stop hit, and post-stop frames are + absorbed. The scan uses a bounded tail window, which yields the same + result as the legacy full-text scan because every prior window was + already checked on earlier deltas. +- **Engine stop-word trims**: a ``Terminal(stop)`` whose ``stop_reason`` + matches a configured stop sequence trims the sequence's tokens (and + re-derives the text) when ``include_stop_str_in_output`` is false — + mirroring the legacy token-level trim. This requires the final + ``TokenDelta`` and its ``Terminal`` to be processed in one + ``process_frames`` batch (the router enqueues them together). +- **Presentation map** (contract draft DEC-9): wire ``Terminal(error, + stop_reason="timeout")`` renders as user-visible ``finish_reason= + "timeout"``; wire ``abort`` renders as ``"cancelled"``. +- **Usage**: the wire-level ``RequestComplete.completion_tokens`` stays the + raw adapter-observed engine count (contract draft DEC-2); the + **user-visible** ``completion_tokens`` reports the assembled (post-trim) + token count, matching the legacy accounting that derives usage from the + trimmed output. This trim-vs-raw split is the documented decision: raw on + the wire, assembled at presentation. +""" + +import dataclasses +from typing import List, Optional, Sequence, Union + +from tensorrt_llm.logger import logger + +from .contract import (CACHED_TOKENS_METRIC_KEY, ContractError, ErrorFrame, + FrontendOutputConfig, RequestComplete, Terminal, + TokenDelta) + +__all__ = ["AssemblyUpdate", "FrontendResponseAssembler", "AssemblerError"] + + +class AssemblerError(ContractError): + """The assembler received frames it cannot present.""" + + +@dataclasses.dataclass +class AssemblyUpdate: + """One presentation-ready update produced from a frame batch. + + ``kind`` is one of ``"delta"`` (new text/tokens), ``"finish"`` + (user-visible finish reason resolved), ``"complete"`` (usage available), + ``"error"`` (typed failure). + """ + + kind: str + text_diff: str = "" + new_token_ids: tuple = () + logprobs: Optional[tuple] = None + prompt_logprobs: Optional[tuple] = None + finish_reason: Optional[str] = None + stop_reason: Union[int, str, None] = None + usage: Optional[dict] = None + error_code: Optional[str] = None + error_message: Optional[str] = None + metrics_cached_tokens: Optional[int] = None + + +class FrontendResponseAssembler: + """Assembles one request's frame stream into presentation updates.""" + + def __init__(self, + request_id: str, + output_config: FrontendOutputConfig, + tokenizer=None, + abort_callback=None, + stream_interval: int = 1): + self.request_id = request_id + self.config = output_config + if (output_config.detokenize and tokenizer is not None + and not hasattr(tokenizer, "decode_incrementally")): + # O(total tokens) is a hard requirement: a tokenizer without + # incremental decoding would force per-delta full re-decodes + # (O(n^2)), so it is rejected at setup rather than degraded. + raise AssemblerError( + f"tokenizer {type(tokenizer).__name__} does not support " + "incremental decoding (decode_incrementally); unsupported for " + "streaming assembly") + self.tokenizer = tokenizer + self._abort_callback = abort_callback + self._stream_interval = stream_interval + + self.token_ids: List[int] = [] + self.text = "" + self.logprobs: List[float] = [] + self.prompt_logprobs: Optional[tuple] = None + self._incremental_states = None + self._flushed_text_len = 0 + + self.finish_reason: Optional[str] = None + self.stop_reason: Union[int, str, None] = None + self.usage: Optional[dict] = None + self._stopped_by_string = False + self._finish_presented = False + self._completed = False + max_stop = max((len(s) for s in output_config.stop_strings), default=0) + self._stop_window = max_stop + + # ------------------------------------------------------------------ # + + def process_frames(self, frames: Sequence) -> List[AssemblyUpdate]: + """Process one batch of frames atomically. + + Trims implied by a ``Terminal`` in the batch apply before any text + from the same batch is surfaced, mirroring the legacy behavior where + final tokens and finish state arrive as one response. + """ + updates: List[AssemblyUpdate] = [] + for position, frame in enumerate(frames): + if isinstance(frame, TokenDelta): + # An engine stop-word finish takes precedence over frontend + # stop-string scanning for ITS OWN final delta — mirroring + # legacy ordering, where the final response trims stop tokens + # before the text scan ever sees them. Earlier deltas in the + # batch still scan normally. + next_frame = frames[position + 1] if position + 1 < len(frames) else None + engine_stops_this_delta = (isinstance(next_frame, Terminal) + and next_frame.finish_reason == "stop") + update = self._process_delta(frame, + scan_stops=not engine_stops_this_delta) + if update is not None: + updates.append(update) + elif isinstance(frame, Terminal): + update = self._process_terminal(frame, updates) + if update is not None: + updates.append(update) + elif isinstance(frame, RequestComplete): + updates.append(self._process_complete(frame)) + elif isinstance(frame, ErrorFrame): + updates.append( + AssemblyUpdate(kind="error", error_code=frame.error_code, + error_message=frame.message)) + self._completed = True + else: + raise AssemblerError(f"unknown frame type {type(frame).__name__}") + return updates + + # ------------------------------------------------------------------ # + + def _decode_new_tokens(self, new_token_ids, flush: bool) -> str: + if not self.config.detokenize or self.tokenizer is None: + return "" + kwargs = dict( + skip_special_tokens=self.config.skip_special_tokens, + spaces_between_special_tokens=self.config.spaces_between_special_tokens) + previous_text = self.text + self.text, self._incremental_states = self.tokenizer.decode_incrementally( + list(new_token_ids), + prev_text=previous_text, + states=self._incremental_states, + flush=flush, + stream_interval=self._stream_interval, + **kwargs) + return self.text[len(previous_text):] + + def _scan_stop_strings(self) -> Optional[str]: + if not self.config.stop_strings: + return None + search_start = max(0, self._flushed_text_len - self._stop_window + 1) + window = self.text[search_start:] + for _, reason in self.config.stop_sequence_reasons: + if isinstance(reason, str) and reason in window: + return reason + for reason in self.config.stop_strings: + if reason in window: + return reason + return None + + def _process_delta(self, frame: TokenDelta, + scan_stops: bool = True) -> Optional[AssemblyUpdate]: + if self._stopped_by_string or self._finish_presented: + return None # post-stop frames are absorbed + self.token_ids.extend(frame.new_token_ids) + if frame.logprobs is not None: + self.logprobs.extend(frame.logprobs) + if frame.prompt_logprobs is not None: + self.prompt_logprobs = frame.prompt_logprobs + text_diff = self._decode_new_tokens(frame.new_token_ids, flush=False) + + stop_hit = self._scan_stop_strings() if scan_stops else None + finish_reason = None + stop_reason = None + if stop_hit is not None: + stop_position = self.text.find( + stop_hit, max(0, self._flushed_text_len - self._stop_window + 1)) + if stop_position < 0: + stop_position = self.text.find(stop_hit) + if self.config.include_stop_str_in_output: + self.text = self.text[:stop_position + len(stop_hit)] + else: + self.text = self.text[:stop_position] + text_diff = self.text[self._flushed_text_len:] + self._stopped_by_string = True + self._finish_presented = True + self.finish_reason = finish_reason = "stop" + self.stop_reason = stop_reason = stop_hit + if self._abort_callback is not None: + try: + # Internal control action: the user-visible finish reason + # stays "stop"; the abort only stops background generation. + self._abort_callback(self.request_id) + except Exception as e: + logger.warning(f"assembler: engine abort after stop-string hit " + f"failed for {self.request_id!r}: {e!r}") + cached = None + if frame.metrics is not None: + raw_cached = frame.metrics.get(CACHED_TOKENS_METRIC_KEY) + if raw_cached is not None: + cached = int(raw_cached) + self._flushed_text_len += len(text_diff) + return AssemblyUpdate(kind="delta", text_diff=text_diff, + new_token_ids=frame.new_token_ids, + logprobs=frame.logprobs, + prompt_logprobs=frame.prompt_logprobs, + finish_reason=finish_reason, stop_reason=stop_reason, + metrics_cached_tokens=cached) + + def _trim_engine_stop_tokens(self, stop_reason, + updates: List[AssemblyUpdate]) -> None: + """Token-level trim for an engine stop-word finish (legacy parity).""" + if self.config.include_stop_str_in_output: + return + matched_sequence = None + for sequence, reason in self.config.stop_sequence_reasons: + if reason == stop_reason: + matched_sequence = sequence + break + if matched_sequence is None: + return + n = len(matched_sequence) + if n == 0 or tuple(self.token_ids[-n:]) != tuple(matched_sequence): + return + del self.token_ids[-n:] + if self.logprobs: + del self.logprobs[-n:] + if self.config.detokenize and self.tokenizer is not None: + # One full re-decode at end-of-request: O(total tokens), once. + kwargs = dict( + skip_special_tokens=self.config.skip_special_tokens, + spaces_between_special_tokens=self.config. + spaces_between_special_tokens) + self.text = self.tokenizer.decode(self.token_ids, **kwargs) + # Adjust the batch's delta updates: the stop sequence may span + # several buffered deltas, so trim tokens walking backwards, and + # rebase the batch's text so nothing past the trim is presented. + delta_updates = [u for u in updates if u.kind == "delta"] + remaining = n + for update in reversed(delta_updates): + if remaining == 0: + break + take = min(remaining, len(update.new_token_ids)) + if take: + keep = len(update.new_token_ids) - take + update.new_token_ids = tuple(update.new_token_ids[:keep]) + if update.logprobs: + update.logprobs = tuple(update.logprobs[:keep]) + remaining -= take + if delta_updates: + batch_start = self._flushed_text_len - sum( + len(u.text_diff) for u in delta_updates) + # The batch renders together; concentrate the trimmed tail on the + # first delta and blank the rest so concatenation is exact. + delta_updates[0].text_diff = self.text[batch_start:] + for update in delta_updates[1:]: + update.text_diff = "" + self._flushed_text_len = batch_start + len(delta_updates[0].text_diff) + + def _process_terminal(self, frame: Terminal, + updates: List[AssemblyUpdate]) -> Optional[AssemblyUpdate]: + if self._finish_presented: + return None # stop-string hit already presented the finish + # Flush any pending detokenization state. + if self.config.detokenize and self.tokenizer is not None and hasattr( + self.tokenizer, "decode_incrementally"): + try: + text_diff = self._decode_new_tokens((), flush=True) + except Exception: + text_diff = "" + if text_diff: + self._flushed_text_len += len(text_diff) + for update in reversed(updates): + if update.kind == "delta": + update.text_diff += text_diff + break + else: + updates.append(AssemblyUpdate(kind="delta", text_diff=text_diff)) + + if frame.finish_reason == "stop": + self._trim_engine_stop_tokens(frame.stop_reason, updates) + finish_reason = "stop" + elif frame.finish_reason == "length": + finish_reason = "length" + elif frame.finish_reason == "abort": + finish_reason = "cancelled" + elif frame.finish_reason == "error" and frame.stop_reason == "timeout": + finish_reason = "timeout" + else: + self._finish_presented = True + self.finish_reason = "error" + return AssemblyUpdate(kind="error", error_code="engine_error", + error_message=str(frame.stop_reason or "error")) + self._finish_presented = True + self.finish_reason = finish_reason + # Only a real stop carries a user-visible stop_reason; the wire-level + # "timeout" marker is a rendering input, not presented (legacy parity). + presented_stop_reason = frame.stop_reason if finish_reason == "stop" else None + self.stop_reason = presented_stop_reason + return AssemblyUpdate(kind="finish", finish_reason=finish_reason, + stop_reason=presented_stop_reason) + + def _process_complete(self, frame: RequestComplete) -> AssemblyUpdate: + self._completed = True + # User-visible usage counts the assembled (post-trim) tokens for + # legacy parity; the wire frame keeps the raw engine count. See the + # module docstring for the documented trim-vs-raw decision. + completion_tokens = len(self.token_ids) + self.usage = { + "prompt_tokens": frame.prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": frame.prompt_tokens + completion_tokens, + } + if frame.cached_tokens is not None: + self.usage["cached_tokens"] = frame.cached_tokens + if frame.status == "failed" and not self._finish_presented: + self._finish_presented = True + return AssemblyUpdate(kind="error", error_code="request_failed", + error_message="request failed", usage=self.usage) + return AssemblyUpdate(kind="complete", usage=self.usage, + finish_reason=self.finish_reason, + stop_reason=self.stop_reason) + + @property + def completed(self) -> bool: + return self._completed diff --git a/tensorrt_llm/executor/engine_client/codec.py b/tensorrt_llm/executor/engine_client/codec.py new file mode 100644 index 000000000000..779d4110732f --- /dev/null +++ b/tensorrt_llm/executor/engine_client/codec.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Strict msgpack codec for the engine contract wire types. + +Every wire type in ``contract`` encodes to a msgpack envelope +``{"k": , "d": {: , ...}}`` and decodes back through +the contract constructors, so construction-time validation applies to +every decoded payload. The codec is strictly primitive: any value outside +``None | bool | int | float | str | list | dict`` is a typed encode error, +and there is **no pickle fallback under any flag**. Decode rejects, with +typed errors, unknown kinds, unknown fields, missing required fields, +wrong item types, newer protocol versions, msgpack extension/bin types, +non-finite floats, and payloads exceeding the resource limits below. + +Encoding is canonical: dataclass fields are written in declaration order +and mapping-valued fields are written in sorted key order, so equal values +produce identical bytes (the basis of the checked-in golden fixtures). +""" + +import dataclasses +import math +from collections.abc import Mapping, Sequence + +import msgpack + +from .contract import (ENGINE_CONTRACT_VERSION, INT64_MAX, INT64_MIN, + MAX_STRING_CHARS, ContractConstructionError, + ContractError, EngineCapabilities, EngineHealth, + EngineRequest, EngineSamplingConfig, ErrorFrame, + FrontendModelContext, FrontendOutputConfig, + GuidedDecodingSpec, IterationStatsBatch, + KvCacheEventsBatch, RequestComplete, Terminal, + TokenDelta, TokenizerSpec) + +__all__ = [ + "MAX_MESSAGE_BYTES", + "MAX_NESTING_DEPTH", + "MAX_ARRAY_ITEMS", + "MAX_MAP_ITEMS", + "CodecError", + "EncodeError", + "DecodeError", + "KIND_BY_TYPE", + "TYPE_BY_KIND", + "encode", + "decode", +] + +MAX_MESSAGE_BYTES = 16 * 1024 * 1024 +MAX_NESTING_DEPTH = 16 +MAX_ARRAY_ITEMS = 2 * 1024 * 1024 +MAX_MAP_ITEMS = 4096 + + +class CodecError(ContractError): + """Base class for codec errors. Carries a stable ``reason`` code.""" + + def __init__(self, reason: str, message: str): + super().__init__(f"{reason}: {message}") + self.reason = reason + + +class EncodeError(CodecError): + """A value cannot be represented on the strict primitive wire.""" + + +class DecodeError(CodecError): + """A payload violates the wire schema, limits, or version rules.""" + + +KIND_BY_TYPE = { + EngineCapabilities: "engine_capabilities", + EngineSamplingConfig: "engine_sampling_config", + GuidedDecodingSpec: "guided_decoding_spec", + EngineRequest: "engine_request", + TokenDelta: "token_delta", + Terminal: "terminal", + RequestComplete: "request_complete", + ErrorFrame: "error_frame", + FrontendOutputConfig: "frontend_output_config", + TokenizerSpec: "tokenizer_spec", + FrontendModelContext: "frontend_model_context", + EngineHealth: "engine_health", + IterationStatsBatch: "iteration_stats_batch", + KvCacheEventsBatch: "kv_cache_events_batch", +} +TYPE_BY_KIND = {kind: cls for cls, kind in KIND_BY_TYPE.items()} + +# Fields whose values are nested wire types (decoded recursively through +# their own constructors). ``None`` stays ``None`` for optional fields. +_NESTED_FIELDS = { + EngineRequest: { + "sampling": EngineSamplingConfig, + "guided_decoding": GuidedDecodingSpec, + }, + FrontendModelContext: { + "tokenizer": TokenizerSpec, + "capabilities": EngineCapabilities, + }, +} + + +def _value_to_wire(value, depth: int): + if depth > MAX_NESTING_DEPTH: + raise EncodeError("limit_exceeded", f"nesting depth exceeds {MAX_NESTING_DEPTH}") + if value is None or isinstance(value, bool): + return value + if isinstance(value, int): + if not INT64_MIN <= value <= INT64_MAX: + raise EncodeError("limit_exceeded", "integer outside signed 64-bit range") + return value + if isinstance(value, float): + if not math.isfinite(value): + raise EncodeError("not_finite", "NaN/Inf floats are not encodable") + return value + if isinstance(value, str): + if len(value) > MAX_STRING_CHARS: + raise EncodeError("limit_exceeded", f"string exceeds {MAX_STRING_CHARS} characters") + return value + if dataclasses.is_dataclass(value) and not isinstance(value, type): + if type(value) not in KIND_BY_TYPE: + raise EncodeError("non_primitive", f"{type(value).__name__} is not a registered wire type") + return { + field.name: _value_to_wire(getattr(value, field.name), depth + 1) + for field in dataclasses.fields(value) + } + if isinstance(value, Mapping): + if len(value) > MAX_MAP_ITEMS: + raise EncodeError("limit_exceeded", f"mapping exceeds {MAX_MAP_ITEMS} entries") + keys = list(value.keys()) + for key in keys: + if not isinstance(key, str): + raise EncodeError("non_primitive", "mapping keys must be strings") + wire = {} + for key in sorted(keys): + wire[key] = _value_to_wire(value[key], depth + 1) + return wire + if isinstance(value, (bytes, bytearray)): + raise EncodeError("non_primitive", "raw bytes are not part of the contract") + if isinstance(value, Sequence): + if len(value) > MAX_ARRAY_ITEMS: + raise EncodeError("limit_exceeded", f"sequence exceeds {MAX_ARRAY_ITEMS} items") + return [_value_to_wire(item, depth + 1) for item in value] + raise EncodeError("non_primitive", f"{type(value).__name__} is not encodable on the wire") + + +def encode(obj) -> bytes: + """Encode a registered wire type to canonical msgpack bytes. + + Raises: + EncodeError: for unregistered types, non-primitive values, + non-finite floats, or payloads exceeding the resource limits. + """ + kind = KIND_BY_TYPE.get(type(obj)) + if kind is None: + raise EncodeError("unknown_type", f"{type(obj).__name__} is not a registered wire type") + payload = {"k": kind, "d": _value_to_wire(obj, 1)} + data = msgpack.packb(payload, use_bin_type=True) + if len(data) > MAX_MESSAGE_BYTES: + raise EncodeError("limit_exceeded", f"message exceeds {MAX_MESSAGE_BYTES} bytes") + return data + + +def _check_raw_depth(value, depth: int) -> None: + if depth > MAX_NESTING_DEPTH: + raise DecodeError("limit_exceeded", f"nesting depth exceeds {MAX_NESTING_DEPTH}") + if isinstance(value, dict): + for key, item in value.items(): + _check_raw_depth(item, depth + 1) + elif isinstance(value, list): + for item in value: + _check_raw_depth(item, depth + 1) + + +def _decode_dataclass(cls, raw, path: str): + if raw is None and cls in (GuidedDecodingSpec,): + return None + if not isinstance(raw, dict): + raise DecodeError("invalid_content", f"{path}: expected a field map") + field_map = {field.name: field for field in dataclasses.fields(cls)} + unknown = set(raw.keys()) - set(field_map.keys()) + if unknown: + raise DecodeError("unknown_field", f"{path}: unknown field(s) {sorted(unknown)!r}") + version = raw.get("protocol_version", ENGINE_CONTRACT_VERSION) + if isinstance(version, bool) or not isinstance(version, int): + raise DecodeError("invalid_content", f"{path}.protocol_version: expected int") + if version > ENGINE_CONTRACT_VERSION: + raise DecodeError( + "version_unsupported", + f"{path}.protocol_version {version} > supported {ENGINE_CONTRACT_VERSION}") + nested = _NESTED_FIELDS.get(cls, {}) + kwargs = {} + for name, field in field_map.items(): + if name not in raw: + has_default = (field.default is not dataclasses.MISSING + or field.default_factory is not dataclasses.MISSING) + if not has_default: + raise DecodeError("missing_field", f"{path}: missing required field {name!r}") + continue + value = raw[name] + if name in nested: + value = _decode_dataclass(nested[name], value, f"{path}.{name}") + kwargs[name] = value + try: + return cls(**kwargs) + except ContractConstructionError as e: + raise DecodeError("invalid_content", f"{path}: {e}") from e + + +def decode(data: bytes): + """Decode msgpack bytes back into a registered wire type. + + Raises: + DecodeError: for malformed msgpack, extension/bin types, unknown + kinds/fields, missing required fields, wrong types, newer + protocol versions, or payloads exceeding the resource limits. + """ + if not isinstance(data, (bytes, bytearray)): + raise DecodeError("invalid_content", f"expected bytes, got {type(data).__name__}") + if len(data) > MAX_MESSAGE_BYTES: + raise DecodeError("limit_exceeded", f"message exceeds {MAX_MESSAGE_BYTES} bytes") + try: + payload = msgpack.unpackb( + data, + raw=False, + strict_map_key=True, + use_list=True, + max_str_len=4 * MAX_STRING_CHARS, + max_bin_len=0, + max_array_len=MAX_ARRAY_ITEMS, + max_map_len=MAX_MAP_ITEMS, + max_ext_len=0, + ) + except Exception as e: # msgpack raises several unpack error types. + raise DecodeError("malformed", f"msgpack unpack failed: {e}") from e + if not isinstance(payload, dict) or set(payload.keys()) != {"k", "d"}: + raise DecodeError("malformed", "expected an envelope with exactly the keys 'k' and 'd'") + kind = payload["k"] + if not isinstance(kind, str) or kind not in TYPE_BY_KIND: + raise DecodeError("unknown_kind", f"unknown wire kind {kind!r}") + _check_raw_depth(payload["d"], 1) + return _decode_dataclass(TYPE_BY_KIND[kind], payload["d"], kind) diff --git a/tensorrt_llm/executor/engine_client/contract.py b/tensorrt_llm/executor/engine_client/contract.py new file mode 100644 index 000000000000..5d8a307567fc --- /dev/null +++ b/tensorrt_llm/executor/engine_client/contract.py @@ -0,0 +1,705 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Typed wire contract between a serving frontend and the generation engine. + +Every type in this module is a frozen, primitives-only, callable-free +dataclass carrying ``protocol_version``. Construction validates types +strictly: ``bool`` is never accepted where ``int`` is required, identifiers +must be non-empty strings, token ids are 64-bit int tuples, and +mapping-valued fields are converted to immutable views. Nothing here may +hold a live Python object graph — these types define the future +inter-process wire and must survive ``codec`` round-trips unchanged. + +See ``ENGINE_CONTRACT.md`` in this package for the scope matrix, the +rejection-to-test table, and divergence notes against the design draft. +""" + +import dataclasses +import math +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Optional, Union + +__all__ = [ + "ENGINE_CONTRACT_VERSION", + "FINISH_REASONS", + "REQUEST_STATUSES", + "GUIDED_DECODING_MODES", + "MAX_PROMPT_TOKENS", + "MAX_STOP_SEQUENCES", + "MAX_STOP_SEQUENCE_TOKENS", + "MAX_STRING_CHARS", + "MAX_METRICS_ENTRIES", + "INT64_MIN", + "INT64_MAX", + "ContractError", + "ContractConstructionError", + "EngineCapabilities", + "EngineSamplingConfig", + "GuidedDecodingSpec", + "EngineRequest", + "TokenDelta", + "Terminal", + "RequestComplete", + "ErrorFrame", + "OutputFrame", + "OUTPUT_FRAME_TYPES", + "FrontendOutputConfig", + "TokenizerSpec", + "FrontendModelContext", + "EngineHealth", + "IterationStatsBatch", + "KvCacheEventsBatch", + "CACHED_TOKENS_METRIC_KEY", + "validate_no_callables", +] + +ENGINE_CONTRACT_VERSION = 1 +"""int: Current engine contract protocol version.""" + +FINISH_REASONS = ("stop", "length", "abort", "error") +"""Valid ``Terminal.finish_reason`` values (sole carrier of finish state).""" + +REQUEST_STATUSES = ("ok", "aborted", "failed") +"""Valid ``RequestComplete.status`` values.""" + +GUIDED_DECODING_MODES = ("json_schema", "json_object", "regex", "grammar", "structural_tag") +"""Valid ``GuidedDecodingSpec.mode`` values (schema-as-data; execution is engine-side).""" + +CACHED_TOKENS_METRIC_KEY = "cached_tokens" +"""Reserved ``TokenDelta.metrics`` key carrying pre-completion cached-token accounting.""" + +# Resource limits enforced at construction time (the codec enforces its own +# byte-level limits on top of these; see codec.py). +MAX_PROMPT_TOKENS = 1_048_576 +MAX_STOP_SEQUENCES = 1_024 +MAX_STOP_SEQUENCE_TOKENS = 256 +MAX_STRING_CHARS = 4 * 1024 * 1024 +MAX_METRICS_ENTRIES = 256 +INT64_MIN = -(2**63) +INT64_MAX = 2**63 - 1 + + +class ContractError(Exception): + """Base class for all engine-contract errors.""" + + +class ContractConstructionError(ContractError): + """A wire type was constructed with invalid content.""" + + +def _fail(type_name: str, field: str, message: str) -> "ContractConstructionError": + return ContractConstructionError(f"{type_name}.{field}: {message}") + + +def _check_int(type_name: str, field: str, value, *, minimum: Optional[int] = None) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise _fail(type_name, field, f"expected int, got {type(value).__name__}") + if not INT64_MIN <= value <= INT64_MAX: + raise _fail(type_name, field, "outside signed 64-bit range") + if minimum is not None and value < minimum: + raise _fail(type_name, field, f"must be >= {minimum}, got {value}") + return value + + +def _check_opt_int(type_name: str, field: str, value, *, minimum: Optional[int] = None) -> Optional[int]: + if value is None: + return None + return _check_int(type_name, field, value, minimum=minimum) + + +def _check_float(type_name: str, field: str, value) -> float: + if isinstance(value, bool): + raise _fail(type_name, field, "expected float, got bool") + if isinstance(value, int): + value = float(value) + if not isinstance(value, float): + raise _fail(type_name, field, f"expected float, got {type(value).__name__}") + if not math.isfinite(value): + raise _fail(type_name, field, "must be finite (NaN/Inf rejected)") + return value + + +def _check_opt_float(type_name: str, field: str, value) -> Optional[float]: + if value is None: + return None + return _check_float(type_name, field, value) + + +def _check_str(type_name: str, field: str, value, *, allow_empty: bool = False) -> str: + if not isinstance(value, str): + raise _fail(type_name, field, f"expected str, got {type(value).__name__}") + if not allow_empty and not value: + raise _fail(type_name, field, "must be a non-empty string") + if len(value) > MAX_STRING_CHARS: + raise _fail(type_name, field, f"exceeds {MAX_STRING_CHARS} characters") + return value + + +def _check_opt_str(type_name: str, field: str, value, *, allow_empty: bool = True) -> Optional[str]: + if value is None: + return None + return _check_str(type_name, field, value, allow_empty=allow_empty) + + +def _check_bool(type_name: str, field: str, value) -> bool: + if not isinstance(value, bool): + raise _fail(type_name, field, f"expected bool, got {type(value).__name__}") + return value + + +def _int_tuple(type_name: str, field: str, value, *, allow_empty: bool = True, + max_items: Optional[int] = None) -> tuple: + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise _fail(type_name, field, f"expected a sequence of ints, got {type(value).__name__}") + items = tuple(value) + if not allow_empty and not items: + raise _fail(type_name, field, "must not be empty") + if max_items is not None and len(items) > max_items: + raise _fail(type_name, field, f"exceeds {max_items} items") + for i, item in enumerate(items): + if isinstance(item, bool) or not isinstance(item, int): + raise _fail(type_name, field, f"item {i}: expected int, got {type(item).__name__}") + if not INT64_MIN <= item <= INT64_MAX: + raise _fail(type_name, field, f"item {i}: outside signed 64-bit range") + return items + + +def _float_tuple(type_name: str, field: str, value) -> Optional[tuple]: + if value is None: + return None + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise _fail(type_name, field, f"expected a sequence of floats, got {type(value).__name__}") + return tuple(_check_float(type_name, f"{field}[{i}]", item) for i, item in enumerate(value)) + + +def _str_tuple(type_name: str, field: str, value, *, allow_empty_items: bool = False) -> tuple: + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise _fail(type_name, field, f"expected a sequence of strings, got {type(value).__name__}") + return tuple( + _check_str(type_name, f"{field}[{i}]", item, allow_empty=allow_empty_items) + for i, item in enumerate(value)) + + +def _freeze_metrics(type_name: str, field: str, value) -> Optional[Mapping]: + if value is None: + return None + if not isinstance(value, Mapping): + raise _fail(type_name, field, f"expected a mapping, got {type(value).__name__}") + if len(value) > MAX_METRICS_ENTRIES: + raise _fail(type_name, field, f"exceeds {MAX_METRICS_ENTRIES} entries") + frozen = {} + for key in value: + if not isinstance(key, str) or not key: + raise _fail(type_name, field, "metric keys must be non-empty strings") + frozen[key] = _check_float(type_name, f"{field}[{key!r}]", value[key]) + return MappingProxyType(frozen) + + +def _check_protocol_version(type_name: str, value) -> int: + version = _check_int(type_name, "protocol_version", value, minimum=1) + return version + + +def validate_no_callables(obj, _path: str = "value") -> None: + """Reject any callable anywhere in a wire value graph. + + Walks dataclasses, mappings, sequences, and primitives. Raises + ``ContractConstructionError`` when a callable (or an unrecognized + object that could smuggle behavior) is found. + """ + if obj is None or isinstance(obj, (bool, int, float, str)): + return + if callable(obj): + raise ContractConstructionError(f"{_path}: callables are forbidden on the wire") + if dataclasses.is_dataclass(obj) and not isinstance(obj, type): + for field in dataclasses.fields(obj): + validate_no_callables(getattr(obj, field.name), f"{_path}.{field.name}") + return + if isinstance(obj, Mapping): + for key, item in obj.items(): + validate_no_callables(key, f"{_path} key") + validate_no_callables(item, f"{_path}[{key!r}]") + return + if isinstance(obj, (bytes, bytearray)): + raise ContractConstructionError(f"{_path}: raw bytes are not part of the contract") + if isinstance(obj, Sequence): + for i, item in enumerate(obj): + validate_no_callables(item, f"{_path}[{i}]") + return + raise ContractConstructionError( + f"{_path}: {type(obj).__name__} is not a primitive contract value") + + +@dataclasses.dataclass(frozen=True) +class EngineCapabilities: + """Handshake payload naming what the engine supports. + + A request whose ``required_features`` are not a subset of ``features`` + fails before submit. + """ + + features: tuple = () + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "features", _str_tuple(name, "features", self.features)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class EngineSamplingConfig: + """Canonical generation data in wire form. + + Generation semantics live here; result-assembly concerns (stop strings, + formatting, candidate selection) have no fields here — they belong to + ``FrontendOutputConfig``. ``pad_id`` is carried so the runtime request is + reproducible from the encoded form alone (contract divergence note 4). + """ + + max_new_tokens: int + end_id: Optional[int] = None + pad_id: Optional[int] = None + stop_token_ids: tuple = () + stop_token_sequences: tuple = () + min_tokens: Optional[int] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + top_k: Optional[int] = None + seed: Optional[int] = None + repetition_penalty: Optional[float] = None + presence_penalty: Optional[float] = None + frequency_penalty: Optional[float] = None + num_logprobs: Optional[int] = None + num_prompt_logprobs: Optional[int] = None + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "max_new_tokens", + _check_int(name, "max_new_tokens", self.max_new_tokens, minimum=1)) + object.__setattr__(self, "end_id", _check_opt_int(name, "end_id", self.end_id)) + object.__setattr__(self, "pad_id", _check_opt_int(name, "pad_id", self.pad_id)) + object.__setattr__(self, "stop_token_ids", + _int_tuple(name, "stop_token_ids", self.stop_token_ids)) + if isinstance(self.stop_token_sequences, (str, bytes)) or \ + not isinstance(self.stop_token_sequences, Sequence): + raise _fail(name, "stop_token_sequences", "expected a sequence of int sequences") + if len(self.stop_token_sequences) > MAX_STOP_SEQUENCES: + raise _fail(name, "stop_token_sequences", f"exceeds {MAX_STOP_SEQUENCES} sequences") + object.__setattr__( + self, "stop_token_sequences", + tuple( + _int_tuple(name, f"stop_token_sequences[{i}]", seq, allow_empty=False, + max_items=MAX_STOP_SEQUENCE_TOKENS) + for i, seq in enumerate(self.stop_token_sequences))) + object.__setattr__(self, "min_tokens", + _check_opt_int(name, "min_tokens", self.min_tokens, minimum=0)) + object.__setattr__(self, "temperature", _check_opt_float(name, "temperature", self.temperature)) + object.__setattr__(self, "top_p", _check_opt_float(name, "top_p", self.top_p)) + object.__setattr__(self, "top_k", _check_opt_int(name, "top_k", self.top_k)) + object.__setattr__(self, "seed", _check_opt_int(name, "seed", self.seed)) + object.__setattr__(self, "repetition_penalty", + _check_opt_float(name, "repetition_penalty", self.repetition_penalty)) + object.__setattr__(self, "presence_penalty", + _check_opt_float(name, "presence_penalty", self.presence_penalty)) + object.__setattr__(self, "frequency_penalty", + _check_opt_float(name, "frequency_penalty", self.frequency_penalty)) + object.__setattr__(self, "num_logprobs", + _check_opt_int(name, "num_logprobs", self.num_logprobs, minimum=0)) + object.__setattr__(self, "num_prompt_logprobs", + _check_opt_int(name, "num_prompt_logprobs", self.num_prompt_logprobs, minimum=0)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class GuidedDecodingSpec: + """Structured-output request as data; grammar compile + masking are engine-side. + + Present in the schema from V0 but capability-gated (rejected pre-submit) + until the engine-path validation lands. + """ + + mode: str + payload: Optional[str] = None + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + mode = _check_str(name, "mode", self.mode) + if mode not in GUIDED_DECODING_MODES: + raise _fail(name, "mode", f"must be one of {GUIDED_DECODING_MODES}, got {mode!r}") + payload = _check_opt_str(name, "payload", self.payload, allow_empty=False) + if payload is None and mode != "json_object": + raise _fail(name, "payload", f"required for mode {mode!r}") + object.__setattr__(self, "mode", mode) + object.__setattr__(self, "payload", payload) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class EngineRequest: + """Southbound request: token ids in, sampling data, no Python object graph.""" + + request_id: str + prompt_token_ids: tuple + sampling: EngineSamplingConfig + guided_decoding: Optional[GuidedDecodingSpec] = None + required_features: tuple = () + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "request_id", _check_str(name, "request_id", self.request_id)) + object.__setattr__( + self, "prompt_token_ids", + _int_tuple(name, "prompt_token_ids", self.prompt_token_ids, allow_empty=False, + max_items=MAX_PROMPT_TOKENS)) + if not isinstance(self.sampling, EngineSamplingConfig): + raise _fail(name, "sampling", "expected EngineSamplingConfig") + if self.guided_decoding is not None and not isinstance(self.guided_decoding, GuidedDecodingSpec): + raise _fail(name, "guided_decoding", "expected GuidedDecodingSpec or None") + object.__setattr__(self, "required_features", + _str_tuple(name, "required_features", self.required_features)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + validate_no_callables(self, name) + + +def _check_frame_common(name: str, request_id, event_seq) -> tuple: + return (_check_str(name, "request_id", request_id), + _check_int(name, "event_seq", event_seq, minimum=0)) + + +@dataclasses.dataclass(frozen=True) +class TokenDelta: + """Per-iteration token emission. Never carries completion state. + + ``new_token_ids`` is non-empty by construction: a runtime response with + no new tokens produces no ``TokenDelta`` frame. ``logprobs``, when + present, aligns 1:1 with ``new_token_ids``. ``prompt_logprobs`` arrives + at most once per sequence. The reserved metrics key + ``CACHED_TOKENS_METRIC_KEY`` carries pre-completion cached-token + accounting. + """ + + request_id: str + sequence_id: int + new_token_ids: tuple + logprobs: Optional[tuple] = None + prompt_logprobs: Optional[tuple] = None + metrics: Optional[Mapping] = None + event_seq: int = 0 + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + request_id, event_seq = _check_frame_common(name, self.request_id, self.event_seq) + object.__setattr__(self, "request_id", request_id) + object.__setattr__(self, "event_seq", event_seq) + object.__setattr__(self, "sequence_id", _check_int(name, "sequence_id", self.sequence_id, minimum=0)) + object.__setattr__(self, "new_token_ids", + _int_tuple(name, "new_token_ids", self.new_token_ids, allow_empty=False)) + logprobs = _float_tuple(name, "logprobs", self.logprobs) + if logprobs is not None and len(logprobs) != len(self.new_token_ids): + raise _fail(name, "logprobs", + f"length {len(logprobs)} != new_token_ids length {len(self.new_token_ids)}") + object.__setattr__(self, "logprobs", logprobs) + object.__setattr__(self, "prompt_logprobs", + _float_tuple(name, "prompt_logprobs", self.prompt_logprobs)) + object.__setattr__(self, "metrics", _freeze_metrics(name, "metrics", self.metrics)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class Terminal: + """Sole carrier of finish/stop reasons; exactly once per started sequence.""" + + request_id: str + sequence_id: int + finish_reason: str + stop_reason: Union[int, str, None] = None + event_seq: int = 0 + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + request_id, event_seq = _check_frame_common(name, self.request_id, self.event_seq) + object.__setattr__(self, "request_id", request_id) + object.__setattr__(self, "event_seq", event_seq) + object.__setattr__(self, "sequence_id", _check_int(name, "sequence_id", self.sequence_id, minimum=0)) + finish_reason = _check_str(name, "finish_reason", self.finish_reason) + if finish_reason not in FINISH_REASONS: + raise _fail(name, "finish_reason", f"must be one of {FINISH_REASONS}, got {finish_reason!r}") + object.__setattr__(self, "finish_reason", finish_reason) + if self.stop_reason is not None: + if isinstance(self.stop_reason, bool) or not isinstance(self.stop_reason, (int, str)): + raise _fail(name, "stop_reason", "expected int, str, or None") + if isinstance(self.stop_reason, int): + _check_int(name, "stop_reason", self.stop_reason) + else: + _check_str(name, "stop_reason", self.stop_reason, allow_empty=False) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class RequestComplete: + """Exactly-once request ending frame (after every started sequence's Terminal). + + ``prompt_tokens``/``completion_tokens`` are adapter-observed (draft + DEC-2), not engine-native usage. ``cached_tokens`` is the final + cached-token accounting: ``None`` means the source never reported it; + when at least one ``TokenDelta`` carried ``CACHED_TOKENS_METRIC_KEY``, + this field must equal the last per-delta value. + """ + + request_id: str + status: str + prompt_tokens: int + completion_tokens: int + cached_tokens: Optional[int] = None + metrics: Optional[Mapping] = None + event_seq: int = 0 + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + request_id, event_seq = _check_frame_common(name, self.request_id, self.event_seq) + object.__setattr__(self, "request_id", request_id) + object.__setattr__(self, "event_seq", event_seq) + status = _check_str(name, "status", self.status) + if status not in REQUEST_STATUSES: + raise _fail(name, "status", f"must be one of {REQUEST_STATUSES}, got {status!r}") + object.__setattr__(self, "status", status) + object.__setattr__(self, "prompt_tokens", + _check_int(name, "prompt_tokens", self.prompt_tokens, minimum=0)) + object.__setattr__(self, "completion_tokens", + _check_int(name, "completion_tokens", self.completion_tokens, minimum=0)) + object.__setattr__(self, "cached_tokens", + _check_opt_int(name, "cached_tokens", self.cached_tokens, minimum=0)) + object.__setattr__(self, "metrics", _freeze_metrics(name, "metrics", self.metrics)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class ErrorFrame: + """Standalone request ending, only for failures before any sequence started.""" + + request_id: str + error_code: str + message: str = "" + event_seq: int = 0 + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + request_id, event_seq = _check_frame_common(name, self.request_id, self.event_seq) + object.__setattr__(self, "request_id", request_id) + object.__setattr__(self, "event_seq", event_seq) + object.__setattr__(self, "error_code", _check_str(name, "error_code", self.error_code)) + object.__setattr__(self, "message", _check_str(name, "message", self.message, allow_empty=True)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +OUTPUT_FRAME_TYPES = (TokenDelta, Terminal, RequestComplete, ErrorFrame) +OutputFrame = Union[TokenDelta, Terminal, RequestComplete, ErrorFrame] +"""One typed output frame; a stream is a multiplex of these per request.""" + + +@dataclasses.dataclass(frozen=True) +class FrontendOutputConfig: + """Frontend-only result-assembly configuration, keyed by ``request_id``. + + Never crosses the frontend↔engine contract; typed and codec-tested so + frontend state stays serializable. ``stop_sequence_reasons`` is an + ordered association list of ``(stop_token_sequence, user_visible_reason)`` + pairs: configuration order, first match wins, collisions resolved by + order (contract divergence note 11). + """ + + detokenize: bool = True + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True + stop_strings: tuple = () + include_stop_str_in_output: bool = False + stop_sequence_reasons: tuple = () + end_id: Optional[int] = None + num_logprobs: Optional[int] = None + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "detokenize", _check_bool(name, "detokenize", self.detokenize)) + object.__setattr__(self, "skip_special_tokens", + _check_bool(name, "skip_special_tokens", self.skip_special_tokens)) + object.__setattr__( + self, "spaces_between_special_tokens", + _check_bool(name, "spaces_between_special_tokens", self.spaces_between_special_tokens)) + stop_strings = _str_tuple(name, "stop_strings", self.stop_strings, allow_empty_items=False) + object.__setattr__(self, "stop_strings", stop_strings) + object.__setattr__(self, "include_stop_str_in_output", + _check_bool(name, "include_stop_str_in_output", self.include_stop_str_in_output)) + if isinstance(self.stop_sequence_reasons, (str, bytes)) or \ + not isinstance(self.stop_sequence_reasons, Sequence): + raise _fail(name, "stop_sequence_reasons", "expected an ordered sequence of pairs") + pairs = [] + for i, pair in enumerate(self.stop_sequence_reasons): + if isinstance(pair, (str, bytes)) or not isinstance(pair, Sequence) or len(pair) != 2: + raise _fail(name, f"stop_sequence_reasons[{i}]", "expected a (sequence, reason) pair") + sequence = _int_tuple(name, f"stop_sequence_reasons[{i}][0]", pair[0], allow_empty=False, + max_items=MAX_STOP_SEQUENCE_TOKENS) + reason = pair[1] + if isinstance(reason, bool) or not isinstance(reason, (int, str)): + raise _fail(name, f"stop_sequence_reasons[{i}][1]", "expected int or str reason") + pairs.append((sequence, reason)) + object.__setattr__(self, "stop_sequence_reasons", tuple(pairs)) + object.__setattr__(self, "end_id", _check_opt_int(name, "end_id", self.end_id)) + object.__setattr__(self, "num_logprobs", + _check_opt_int(name, "num_logprobs", self.num_logprobs, minimum=0)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class TokenizerSpec: + """Data-only tokenizer provenance sufficient to reload without guessing. + + ``files_manifest`` is an ordered tuple of ``(relative_path, sha256_hex)`` + pairs. JSON-valued configuration (added tokens, special tokens, + normalizer, pre-tokenizer) crosses as source strings, never live + objects: a live tokenizer handle could not survive a process boundary. + """ + + uri: str + files_manifest: tuple = () + revision: Optional[str] = None + fast: bool = True + trust_remote_code: bool = False + added_tokens_json: Optional[str] = None + special_tokens_json: Optional[str] = None + normalizer_json: Optional[str] = None + pre_tokenizer_json: Optional[str] = None + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "uri", _check_str(name, "uri", self.uri)) + if isinstance(self.files_manifest, (str, bytes)) or not isinstance(self.files_manifest, Sequence): + raise _fail(name, "files_manifest", "expected a sequence of (path, sha256) pairs") + manifest = [] + for i, pair in enumerate(self.files_manifest): + if isinstance(pair, (str, bytes)) or not isinstance(pair, Sequence) or len(pair) != 2: + raise _fail(name, f"files_manifest[{i}]", "expected a (path, sha256) pair") + manifest.append((_check_str(name, f"files_manifest[{i}][0]", pair[0]), + _check_str(name, f"files_manifest[{i}][1]", pair[1]))) + object.__setattr__(self, "files_manifest", tuple(manifest)) + object.__setattr__(self, "revision", _check_opt_str(name, "revision", self.revision, + allow_empty=False)) + object.__setattr__(self, "fast", _check_bool(name, "fast", self.fast)) + object.__setattr__(self, "trust_remote_code", + _check_bool(name, "trust_remote_code", self.trust_remote_code)) + for field in ("added_tokens_json", "special_tokens_json", "normalizer_json", + "pre_tokenizer_json"): + object.__setattr__(self, field, _check_opt_str(name, field, getattr(self, field))) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class FrontendModelContext: + """Data-only model context the frontend reads instead of engine internals. + + Built locally from the in-process engine in V0; a later remote detach + only swaps the delivery (ready-handshake) without changing consumers. + """ + + tokenizer: TokenizerSpec + capabilities: EngineCapabilities + chat_template_source: Optional[str] = None + chat_template_version: Optional[str] = None + eos_id: Optional[int] = None + pad_id: Optional[int] = None + max_context_length: Optional[int] = None + # Normalization inputs so the frontend never reaches into live model + # config objects: extra stop ids from the model's generation config and + # the architecture family (used by the callable-producing-normalization + # eligibility gates). + generation_stop_token_ids: tuple = () + model_type: Optional[str] = None + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + if not isinstance(self.tokenizer, TokenizerSpec): + raise _fail(name, "tokenizer", "expected TokenizerSpec") + if not isinstance(self.capabilities, EngineCapabilities): + raise _fail(name, "capabilities", "expected EngineCapabilities") + object.__setattr__(self, "chat_template_source", + _check_opt_str(name, "chat_template_source", self.chat_template_source)) + object.__setattr__(self, "chat_template_version", + _check_opt_str(name, "chat_template_version", self.chat_template_version, + allow_empty=False)) + object.__setattr__(self, "eos_id", _check_opt_int(name, "eos_id", self.eos_id)) + object.__setattr__(self, "pad_id", _check_opt_int(name, "pad_id", self.pad_id)) + object.__setattr__(self, "max_context_length", + _check_opt_int(name, "max_context_length", self.max_context_length, minimum=1)) + object.__setattr__( + self, "generation_stop_token_ids", + _int_tuple(name, "generation_stop_token_ids", + self.generation_stop_token_ids)) + object.__setattr__(self, "model_type", + _check_opt_str(name, "model_type", self.model_type, + allow_empty=False)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + validate_no_callables(self, name) + + +@dataclasses.dataclass(frozen=True) +class EngineHealth: + """Typed health probe result.""" + + healthy: bool + detail: str = "" + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "healthy", _check_bool(name, "healthy", self.healthy)) + object.__setattr__(self, "detail", _check_str(name, "detail", self.detail, allow_empty=True)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class IterationStatsBatch: + """Iteration stats entries; each entry is a self-contained JSON document string.""" + + entries: tuple = () + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "entries", _str_tuple(name, "entries", self.entries)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) + + +@dataclasses.dataclass(frozen=True) +class KvCacheEventsBatch: + """KV-cache event entries; each entry is a self-contained JSON document string.""" + + entries: tuple = () + protocol_version: int = ENGINE_CONTRACT_VERSION + + def __post_init__(self): + name = type(self).__name__ + object.__setattr__(self, "entries", _str_tuple(name, "entries", self.entries)) + object.__setattr__(self, "protocol_version", _check_protocol_version(name, self.protocol_version)) diff --git a/tensorrt_llm/executor/engine_client/conversion.py b/tensorrt_llm/executor/engine_client/conversion.py new file mode 100644 index 000000000000..f24f58ab960e --- /dev/null +++ b/tensorrt_llm/executor/engine_client/conversion.py @@ -0,0 +1,654 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion between legacy request inputs and the engine contract. + +Three responsibilities, in pipeline order: + +1. ``prepare_sampling_params`` — the one-shot normalization boundary. EOS/PAD + defaulting, generation-config stop ids, and stop-string tokenization run + exactly once here (the documented equivalent of + ``LLM._prepare_sampling_params`` minus its callable-producing steps, which + make a request ineligible instead). +2. ``convert_request`` — prepared ``SamplingParams`` → ``EngineRequest`` + + ``FrontendOutputConfig``. The first capability gate: everything the + primitive contract cannot represent is rejected here with a typed error + naming the axis, before the engine could ever see the request. +3. ``engine_request_to_generation_request`` — the compositional runtime + translation. Consumes ONLY the ``EngineRequest`` (no unencoded sidecar), + so the worker-facing request is reproducible from the encoded wire form + alone; pre-tokenized stop sequences ride ``GenerationRequest. + stop_token_sequences`` and ``SamplingParams._setup`` is never re-run. + +``ELIGIBILITY_MATRIX`` is the machine-readable classification of every +``SamplingParams`` field, request-level input, and engine/server config axis; +tests enforce its completeness against the live ``SamplingParams`` surface. +""" + +import dataclasses +import json as json_module +from typing import List, Optional, Tuple + +from pydantic import BaseModel + +from ...sampling_params import GuidedDecodingParams, SamplingParams +from ..request import DEFAULT_REQUEST_PRIORITY, GenerationRequest +from .contract import (ContractError, EngineRequest, EngineSamplingConfig, + FrontendOutputConfig, GuidedDecodingSpec) + +__all__ = [ + "ConversionError", + "RequestIneligibleError", + "EligibilityRule", + "ELIGIBILITY_MATRIX", + "GUIDED_DECODING_FEATURE", + "prepare_sampling_params", + "derive_required_features", + "convert_request", + "engine_request_to_generation_request", +] + +GUIDED_DECODING_FEATURE = "guided_decoding" + +# Frozen V0 disposition of every OpenAI request field (the union of the +# chat and completions request models). Machine-checked for completeness +# against the live protocol models. Values: +# "preprocessing" — consumed by shared frontend preprocessing before +# conversion (templates, tokenization, metadata) +# "supported" — carried on the contract (directly or via +# SamplingParams fields classified in the matrix) +# "frontend" — frontend-only presentation/assembly concern +# "capability_gated" — representable data, rejected pre-submit until the +# capability is validated (guided decoding) +# "axis:" — maps to that ineligible/config-gate matrix axis +OPENAI_REQUEST_FIELD_DISPOSITION = { + "model": "preprocessing", + "messages": "preprocessing", + "prompt": "preprocessing", + "prompt_token_ids": "preprocessing", + "add_generation_prompt": "preprocessing", + "add_special_tokens": "preprocessing", + "chat_template": "preprocessing", + "chat_template_kwargs": "preprocessing", + "documents": "preprocessing", + "media_io_kwargs": "axis:multimodal", + "mm_processor_kwargs": "axis:multimodal", + "user": "preprocessing", + "reasoning_effort": "preprocessing", + "stream": "supported", + "stream_options": "frontend", + "suffix": "frontend", + "echo": "axis:echo", + "max_tokens": "supported", + "max_completion_tokens": "supported", + "temperature": "supported", + "top_p": "supported", + "top_k": "supported", + "seed": "supported", + "min_tokens": "supported", + "presence_penalty": "supported", + "frequency_penalty": "supported", + "repetition_penalty": "supported", + "stop": "supported", + "stop_token_ids": "supported", + "include_stop_str_in_output": "frontend", + "detokenize": "frontend", + "skip_special_tokens": "frontend", + "spaces_between_special_tokens": "frontend", + "logprobs": "supported", + "top_logprobs": "axis:top_logprobs", + "logit_bias": "axis:embedding_bias", + "n": "axis:n_gt_1", + "best_of": "axis:n_gt_1", + "use_beam_search": "axis:beam_search", + "early_stopping": "axis:beam_search", + "length_penalty": "axis:beam_search", + "min_p": "axis:min_p", + "top_p_min": "axis:top_p_extras", + "ignore_eos": "axis:ignore_eos", + "prompt_ignore_length": "axis:prompt_ignore_length", + "return_context_logits": "axis:return_logits", + "truncate_prompt_tokens": "axis:truncate_prompt_tokens", + "thinking_token_budget": "axis:thinking_token_budget", + "response_format": "capability_gated", + # Tools render through the chat template and parse in the frontend + # formatters; strict-mode tools add constrained decoding, which flows + # into guided_decoding and hits the capability gate automatically. + "tools": "frontend", + "tool_choice": "frontend", + "lora_request": "axis:lora", + "disaggregated_params": "axis:disaggregated", + "conversation_params": "axis:conversation_params", + "cache_salt": "axis:cache_salt", + "agent_hierarchy": "axis:scheduling_params", +} + + +class ConversionError(ContractError): + """A legacy input cannot be converted onto the contract.""" + + +class RequestIneligibleError(ConversionError): + """Typed pre-submit rejection naming the eligibility axis. + + The serving facade treats this as "fall back to the legacy path" (and + counts it per axis); it must be raised before any engine enqueue. + """ + + def __init__(self, axis: str, message: str): + super().__init__(f"[{axis}] {message}") + self.axis = axis + + +@dataclasses.dataclass(frozen=True) +class EligibilityRule: + """One row of the machine-readable V0 eligibility matrix. + + ``classification`` is one of: + - ``supported``: representable and carried on the contract. + - ``normalized``: consumed by the normalization boundary, not carried. + - ``ineligible``: request-level typed rejection (axis = ``axis``). + - ``config_gate``: rejected at client construction, not per request. + """ + + axis: str + classification: str + sources: tuple + notes: str = "" + + +ELIGIBILITY_MATRIX: Tuple[EligibilityRule, ...] = ( + # --- supported sampling surface (crosses as EngineSamplingConfig) --- + EligibilityRule("core_sampling", "supported", + ("max_tokens", "end_id", "pad_id", "stop_token_ids", "min_tokens", + "temperature", "top_p", "top_k", "seed", "repetition_penalty", + "presence_penalty", "frequency_penalty"), + "direct EngineSamplingConfig fields"), + EligibilityRule("chosen_token_logprobs", "supported", ("logprobs", "prompt_logprobs"), + "eligible only when <= 1; higher values are the top_logprobs axis"), + EligibilityRule("logprob_wire_shape", "supported", + ("logprobs_mode", "logprobs_simple_format", "prompt_logprobs_simple_format"), + "RAW mode only; both legacy result shapes normalized by the envelope"), + EligibilityRule("stop_strings", "supported", ("stop", "include_stop_str_in_output"), + "strings stay frontend-side; tokenized sequences cross as " + "stop_token_sequences with ordered reasons"), + EligibilityRule("frontend_detok", "supported", + ("detokenize", "skip_special_tokens", "spaces_between_special_tokens"), + "frontend-only assembly configuration (FrontendOutputConfig)"), + EligibilityRule("perf_metrics", "supported", ("return_perf_metrics", ), + "metrics keys convert enum->str at the envelope"), + EligibilityRule("guided_decoding_schema", "supported", ("guided_decoding", ), + "schema-as-data crosses; capability-gated at submit until the " + "engine-path validation, so V0 rejects it pre-submit"), + # --- consumed by normalization, never carried --- + EligibilityRule("normalization_inputs", "normalized", + ("add_special_tokens", "_stream_interval"), + "add_special_tokens is a preprocessing input; stream interval only " + "affects legacy chunking / delta granularity"), + EligibilityRule("derived_candidate_count", "normalized", ("n", "best_of"), + "must equal 1 in V0; the n>1 axis rejects everything else"), + # --- request-level ineligible axes (typed rejection pre-submit) --- + EligibilityRule("n_gt_1", "ineligible", ("n", "best_of"), "n>1 / best_of>1"), + EligibilityRule("beam_search", "ineligible", + ("use_beam_search", "beam_width_array", "beam_search_diversity_rate", + "length_penalty", "early_stopping"), + "beam search and beam-only controls"), + EligibilityRule("top_logprobs", "ineligible", ("logprobs", ), "logprobs > 1"), + EligibilityRule("prompt_top_logprobs", "ineligible", ("prompt_logprobs", ), + "prompt_logprobs > 1"), + EligibilityRule("logprobs_mode", "ineligible", ("logprobs_mode", ), + "non-RAW logprob modes"), + EligibilityRule("logits_processor", "ineligible", + ("logits_processor", "apply_batched_logits_processor"), + "per-request callables cannot cross the wire"), + EligibilityRule("embedding_bias", "ineligible", ("embedding_bias", ), + "tensor payloads need the tensor channel (covers OpenAI logit_bias)"), + EligibilityRule("bad_words", "ineligible", ("bad", "bad_token_ids"), ""), + EligibilityRule("ignore_eos", "ineligible", ("ignore_eos", ), ""), + EligibilityRule("min_p", "ineligible", ("min_p", ), "not in EngineSamplingConfig"), + EligibilityRule("top_p_extras", "ineligible", + ("top_p_min", "top_p_reset_ids", "top_p_decay"), ""), + EligibilityRule("no_repeat_ngram", "ineligible", ("no_repeat_ngram_size", ), ""), + EligibilityRule("prompt_ignore_length", "ineligible", ("prompt_ignore_length", ), ""), + EligibilityRule("return_logits", "ineligible", + ("return_context_logits", "return_generation_logits", + "return_encoder_output", "additional_model_outputs", + "_return_log_probs", "_context_logits_auto_enabled", + "_generation_logits_auto_enabled"), + "raw logits/tensor outputs need the tensor channel"), + EligibilityRule("exclude_input_from_output", "ineligible", ("exclude_input_from_output", ), + "non-default echo-style output"), + EligibilityRule("truncate_prompt_tokens", "ineligible", ("truncate_prompt_tokens", ), ""), + EligibilityRule("lookahead_config", "ineligible", ("lookahead_config", ), ""), + EligibilityRule("thinking_token_budget", "ineligible", ("thinking_token_budget", ), + "normalization would attach a logits-processor callable"), + EligibilityRule("bart_forced_tokens", "ineligible", ("model_type", ), + "BART normalization would attach a logits-processor callable"), + EligibilityRule("echo", "ineligible", ("echo", ), + "prompt echo needs legacy prompt-text handling in the " + "postprocessors"), + EligibilityRule("non_streaming", "ineligible", ("streaming", ), "V0 is streaming-only"), + EligibilityRule("multimodal", "ineligible", ("multimodal_params", ), "request-level input"), + EligibilityRule("lora", "ineligible", ("lora_request", ), "request-level input"), + EligibilityRule("prompt_adapter", "ineligible", ("prompt_adapter_request", ), + "request-level input"), + EligibilityRule("disaggregated", "ineligible", ("disaggregated_params", ), + "request-level input"), + EligibilityRule("scheduling_params", "ineligible", ("scheduling_params", ), + "request-level input"), + EligibilityRule("conversation_params", "ineligible", ("conversation_params", ), + "request-level input"), + EligibilityRule("postproc_params", "ineligible", ("postproc_params", ), + "formatter callables never cross; assembly is frontend-owned"), + EligibilityRule("trace_headers", "ineligible", ("trace_headers", ), + "no telemetry propagation on the contract yet; rejecting avoids " + "silent trace loss"), + EligibilityRule("cache_salt", "ineligible", ("cache_salt", ), "request-level input"), + EligibilityRule("query_token_ids", "ineligible", ("query_token_ids", ), + "star-attention workflow"), + EligibilityRule("encoder_input", "ineligible", ("encoder_input_token_ids", ), + "encoder-decoder models"), + EligibilityRule("priority", "ineligible", ("priority", ), + "non-default request priority is not carried in V0"), + EligibilityRule("kv_cache_retention", "ineligible", ("kv_cache_retention_config", ), + "request-level input"), + # --- config gates (client construction; enforced by local_client) --- + EligibilityRule("backend", "config_gate", ("backend", ), + "pytorch only; the transitional _autodeploy is explicitly rejected"), + EligibilityRule("transport", "config_gate", ("transport", ), "IPC proxy only"), + EligibilityRule("postproc_workers", "config_gate", ("num_postprocess_workers", ), + "must be 0"), + EligibilityRule("post_processor_hook", "config_gate", ("post_processor_hook", ), + "global output hook can rewrite/suppress/terminate output"), + EligibilityRule("speculative_config", "config_gate", ("speculative_config", ), + "unvalidated interaction surface"), + EligibilityRule("early_first_token", "config_gate", + ("enable_early_first_token_response", ), + "changes the per-step response shape; unvalidated"), + EligibilityRule("topology", "config_gate", ("world_size", ), + "the setup gate admits exactly the GPU-validated topology (TP1)"), + EligibilityRule("trust_remote_code_tokenizer", "config_gate", ("trust_remote_code", ), + "tokenizer provenance cannot be pinned"), + EligibilityRule("flag", "config_gate", ("TLLM_EXPERIMENTAL_ENGINE_CLIENT", ), + "client construction requires the experimental flag"), +) + + +def prepare_sampling_params(sampling_params: SamplingParams, + *, + tokenizer, + hf_model_config=None, + generation_config=None, + model_type: Optional[str] = None, + stream_interval: int = 1, + force_return_perf_metrics: bool = False) -> SamplingParams: + """One-shot normalization boundary for the contract path. + + Mirrors ``LLM._prepare_sampling_params`` for the V0-eligible surface: + EOS/PAD defaulting, generation-config stop ids, and stop-string + tokenization run exactly once (via ``SamplingParams._setup`` when + ``end_id`` is unset, identical to the legacy condition). The legacy + callable-producing steps (BART forced-token processors, thinking-budget + processors) are NOT replicated: requests that would need them are + rejected as ineligible instead. After this call the params must never be + passed through ``_setup`` again — downstream translation carries + tokenized stop sequences on the request envelope. + """ + if not isinstance(sampling_params, SamplingParams): + raise ConversionError( + f"expected SamplingParams, got {type(sampling_params).__name__}") + if model_type == "bart" and generation_config is not None and ( + getattr(generation_config, "forced_bos_token_id", None) is not None + or getattr(generation_config, "forced_eos_token_id", None) is not None): + raise RequestIneligibleError( + "bart_forced_tokens", + "BART forced-token normalization would attach a logits-processor callable") + if sampling_params.thinking_token_budget is not None: + raise RequestIneligibleError( + "thinking_token_budget", + "thinking-budget normalization would attach a logits-processor callable") + if sampling_params.end_id is None: + if tokenizer is None: + raise ConversionError( + "tokenizer is required to derive end_id/pad_id when end_id is None") + sampling_params._setup(tokenizer, hf_model_config, generation_config) + if sampling_params._stream_interval is None: + sampling_params._stream_interval = stream_interval + sampling_params.return_perf_metrics = (sampling_params.return_perf_metrics + or force_return_perf_metrics) + return sampling_params + + +def prepare_sampling_params_from_context(sampling_params: SamplingParams, + *, + context, + tokenizer, + stream_interval: int = 1, + force_return_perf_metrics: bool = False + ) -> SamplingParams: + """Context-only normalization boundary: no live model-config reaches. + + Mirrors ``prepare_sampling_params`` but consumes only the data-only + ``FrontendModelContext`` (eos/pad ids, generation-config stop ids, + model type) plus the spec-reloaded tokenizer for stop-string + tokenization. This is the serving glue's boundary; a remote detach + swaps the context's delivery without touching this code. + """ + if not isinstance(sampling_params, SamplingParams): + raise ConversionError( + f"expected SamplingParams, got {type(sampling_params).__name__}") + if context.model_type == "bart": + # Conservative context-only stand-in for the legacy forced-token + # check: the BART family normalization attaches callables. + raise RequestIneligibleError( + "bart_forced_tokens", + "BART-family normalization would attach a logits-processor callable") + if sampling_params.thinking_token_budget is not None: + raise RequestIneligibleError( + "thinking_token_budget", + "thinking-budget normalization would attach a logits-processor callable") + if sampling_params.end_id is None: + if context.eos_id is None: + raise ConversionError( + "the model context carries no eos id; cannot default end_id") + sampling_params.end_id = context.eos_id + sampling_params.pad_id = (context.pad_id + if context.pad_id is not None else context.eos_id) + # Stop strings are tokenized exactly once, with the spec-reloaded + # tokenizer (mirroring SamplingParams._setup's encoding call). + if sampling_params.stop is not None and sampling_params._stop_word_ids is None: + strings = ([sampling_params.stop] + if isinstance(sampling_params.stop, str) else + list(sampling_params.stop)) + sampling_params._stop_word_ids = [ + _encode_stop(tokenizer, text) for text in strings + ] + # Generation-config stop ids merge (mirrors _setup). + if context.generation_stop_token_ids: + stop_ids = list(sampling_params.stop_token_ids or []) + for token_id in context.generation_stop_token_ids: + if token_id != sampling_params.end_id and token_id not in stop_ids: + stop_ids.append(token_id) + sampling_params.stop_token_ids = stop_ids or None + if sampling_params._stream_interval is None: + sampling_params._stream_interval = stream_interval + sampling_params.return_perf_metrics = (sampling_params.return_perf_metrics + or force_return_perf_metrics) + return sampling_params + + +def _encode_stop(tokenizer, text: str): + try: + return tokenizer.encode(text, add_special_tokens=False) + except TypeError: + return tokenizer.encode(text) + + +def _check_sampling_eligibility(sp: SamplingParams) -> None: + if sp.n != 1 or (sp.best_of or sp.n) != 1: + raise RequestIneligibleError("n_gt_1", f"n={sp.n}, best_of={sp.best_of}") + if sp.use_beam_search or sp.beam_width_array is not None or \ + sp.beam_search_diversity_rate is not None or sp.length_penalty is not None or \ + sp.early_stopping is not None: + raise RequestIneligibleError("beam_search", "beam search / beam-only controls set") + if sp.logprobs is not None and sp.logprobs > 1: + raise RequestIneligibleError("top_logprobs", f"logprobs={sp.logprobs}") + if sp.prompt_logprobs is not None and sp.prompt_logprobs > 1: + raise RequestIneligibleError("prompt_top_logprobs", + f"prompt_logprobs={sp.prompt_logprobs}") + if getattr(sp.logprobs_mode, "value", sp.logprobs_mode) not in ("raw", "RAW"): + raise RequestIneligibleError("logprobs_mode", f"logprobs_mode={sp.logprobs_mode}") + if sp.logits_processor is not None or sp.apply_batched_logits_processor: + raise RequestIneligibleError("logits_processor", + "per-request logits processors cannot cross the wire") + if sp.embedding_bias is not None: + raise RequestIneligibleError("embedding_bias", "embedding/logit bias tensor set") + if sp.bad is not None or sp.bad_token_ids: + raise RequestIneligibleError("bad_words", "bad words set") + if sp.ignore_eos: + raise RequestIneligibleError("ignore_eos", "ignore_eos set") + if sp.min_p is not None: + raise RequestIneligibleError("min_p", "min_p set") + if sp.top_p_min is not None or sp.top_p_reset_ids is not None or sp.top_p_decay is not None: + raise RequestIneligibleError("top_p_extras", "top_p decay controls set") + if sp.no_repeat_ngram_size is not None: + raise RequestIneligibleError("no_repeat_ngram", "no_repeat_ngram_size set") + if sp.prompt_ignore_length is not None: + raise RequestIneligibleError("prompt_ignore_length", "prompt_ignore_length set") + if sp.return_context_logits or sp.return_generation_logits or sp.return_encoder_output \ + or sp.additional_model_outputs is not None or sp._return_log_probs: + raise RequestIneligibleError("return_logits", "raw logits/tensor outputs requested") + if not sp.exclude_input_from_output: + raise RequestIneligibleError("exclude_input_from_output", + "echo-style output requested") + if sp.truncate_prompt_tokens is not None: + raise RequestIneligibleError("truncate_prompt_tokens", "prompt truncation requested") + if sp.lookahead_config is not None: + raise RequestIneligibleError("lookahead_config", "lookahead decoding config set") + if sp.thinking_token_budget is not None: + raise RequestIneligibleError("thinking_token_budget", "thinking token budget set") + + +_REQUEST_LEVEL_AXES = ( + ("multimodal_params", "multimodal"), + ("lora_request", "lora"), + ("prompt_adapter_request", "prompt_adapter"), + ("disaggregated_params", "disaggregated"), + ("scheduling_params", "scheduling_params"), + ("conversation_params", "conversation_params"), + ("postproc_params", "postproc_params"), + ("trace_headers", "trace_headers"), + ("cache_salt", "cache_salt"), + ("query_token_ids", "query_token_ids"), + ("encoder_input_token_ids", "encoder_input"), + ("kv_cache_retention_config", "kv_cache_retention"), +) + + +def _guided_decoding_to_spec(params: GuidedDecodingParams) -> GuidedDecodingSpec: + if params.json_object: + return GuidedDecodingSpec(mode="json_object") + if params.json is not None: + schema = params.json + if isinstance(schema, BaseModel): + schema = schema.model_json_schema() + if isinstance(schema, dict): + schema = json_module.dumps(schema) + return GuidedDecodingSpec(mode="json_schema", payload=schema) + if params.regex is not None: + return GuidedDecodingSpec(mode="regex", payload=params.regex) + if params.grammar is not None: + return GuidedDecodingSpec(mode="grammar", payload=params.grammar) + if params.structural_tag is not None: + return GuidedDecodingSpec(mode="structural_tag", payload=params.structural_tag) + raise ConversionError("guided_decoding set but no guide field populated") + + +def derive_required_features(engine_request: EngineRequest) -> tuple: + """Derive the feature set a request actually needs from its own fields. + + The client re-derives and validates this at submit time; caller-supplied + ``required_features`` are checked against it, never trusted. + """ + features = [] + if engine_request.guided_decoding is not None: + features.append(GUIDED_DECODING_FEATURE) + return tuple(features) + + +def _split_stop_handling(sp: SamplingParams): + """Split prepared stop configuration by representation. + + Stop token ids and tokenized stop-string sequences cross to the engine; + the stop strings themselves stay frontend-side, with an ordered + ``(sequence, user_visible_reason)`` association mirroring the legacy + stop-reason resolution (configuration order, first match wins). + """ + stop_token_ids = tuple(sp.stop_token_ids or ()) + reasons = [((token_id, ), token_id) for token_id in stop_token_ids] + stop_strings: List[str] = [] + stop_sequences = [] + if sp.stop is not None: + if sp._stop_word_ids is None: + raise ConversionError( + "stop strings present but not tokenized; prepare_sampling_params " + "must run before conversion") + stop_strings = [sp.stop] if isinstance(sp.stop, str) else list(sp.stop) + if len(sp._stop_word_ids) != len(stop_strings): + raise ConversionError( + f"tokenized stop sequences ({len(sp._stop_word_ids)}) do not match " + f"stop strings ({len(stop_strings)})") + for stop_string, word_ids in zip(stop_strings, sp._stop_word_ids): + sequence = tuple(word_ids) + stop_sequences.append(sequence) + reasons.append((sequence, stop_string)) + return stop_token_ids, tuple(stop_sequences), tuple(stop_strings), tuple(reasons) + + +def convert_request(request_id: str, + prompt_token_ids, + sampling_params: SamplingParams, + *, + streaming: bool = True, + multimodal_params=None, + lora_request=None, + prompt_adapter_request=None, + disaggregated_params=None, + scheduling_params=None, + conversation_params=None, + postproc_params=None, + trace_headers=None, + cache_salt=None, + query_token_ids=None, + encoder_input_token_ids=None, + kv_cache_retention_config=None, + priority=None, + echo=False): + """Convert prepared legacy inputs into ``(EngineRequest, FrontendOutputConfig)``. + + This is the first capability gate: every input the V0 contract cannot + carry raises ``RequestIneligibleError`` (with the axis name) before any + engine enqueue. ``sampling_params`` must already be normalized by + ``prepare_sampling_params``. + """ + if not streaming: + raise RequestIneligibleError("non_streaming", "V0 is streaming-only") + if echo: + raise RequestIneligibleError( + "echo", "prompt echo needs legacy prompt-text handling in the " + "postprocessors") + local_values = dict(multimodal_params=multimodal_params, + lora_request=lora_request, + prompt_adapter_request=prompt_adapter_request, + disaggregated_params=disaggregated_params, + scheduling_params=scheduling_params, + conversation_params=conversation_params, + postproc_params=postproc_params, + trace_headers=trace_headers, + cache_salt=cache_salt, + query_token_ids=query_token_ids, + encoder_input_token_ids=encoder_input_token_ids, + kv_cache_retention_config=kv_cache_retention_config) + for kwarg_name, axis in _REQUEST_LEVEL_AXES: + if local_values[kwarg_name] is not None: + raise RequestIneligibleError(axis, f"{kwarg_name} is set") + if priority is not None and priority != DEFAULT_REQUEST_PRIORITY: + raise RequestIneligibleError("priority", f"non-default priority {priority}") + _check_sampling_eligibility(sampling_params) + if sampling_params.end_id is None: + raise ConversionError( + "end_id is unset; prepare_sampling_params must run before conversion") + + stop_token_ids, stop_sequences, stop_strings, stop_reasons = \ + _split_stop_handling(sampling_params) + + sampling = EngineSamplingConfig( + max_new_tokens=sampling_params.max_tokens, + end_id=sampling_params.end_id, + pad_id=sampling_params.pad_id, + stop_token_ids=stop_token_ids, + stop_token_sequences=stop_sequences, + min_tokens=sampling_params.min_tokens, + temperature=sampling_params.temperature, + top_p=sampling_params.top_p, + top_k=sampling_params.top_k, + seed=sampling_params.seed, + repetition_penalty=sampling_params.repetition_penalty, + presence_penalty=sampling_params.presence_penalty, + frequency_penalty=sampling_params.frequency_penalty, + num_logprobs=sampling_params.logprobs, + num_prompt_logprobs=sampling_params.prompt_logprobs, + ) + guided_decoding = None + if sampling_params.guided_decoding is not None: + guided_decoding = _guided_decoding_to_spec(sampling_params.guided_decoding) + + engine_request = EngineRequest(request_id=request_id, + prompt_token_ids=tuple(prompt_token_ids), + sampling=sampling, + guided_decoding=guided_decoding) + engine_request = dataclasses.replace( + engine_request, required_features=derive_required_features(engine_request)) + + output_config = FrontendOutputConfig( + detokenize=sampling_params.detokenize, + skip_special_tokens=sampling_params.skip_special_tokens, + spaces_between_special_tokens=sampling_params.spaces_between_special_tokens, + stop_strings=stop_strings, + include_stop_str_in_output=sampling_params.include_stop_str_in_output, + stop_sequence_reasons=stop_reasons, + end_id=sampling_params.end_id, + num_logprobs=sampling_params.logprobs, + ) + return engine_request, output_config + + +def engine_request_to_generation_request(engine_request: EngineRequest) -> GenerationRequest: + """Compositional runtime translation: encoded ``EngineRequest`` → ``GenerationRequest``. + + Consumes only the ``EngineRequest``: the synthetic ``SamplingParams`` is + rebuilt purely from ``EngineSamplingConfig`` wire fields (which is why + ``pad_id`` is on the wire), and tokenized stop sequences ride the + request's ``stop_token_sequences`` carrier — ``SamplingParams._setup`` is + never involved, so nothing can re-tokenize the stops. + """ + if not isinstance(engine_request, EngineRequest): + raise ConversionError( + f"expected EngineRequest, got {type(engine_request).__name__}") + if engine_request.guided_decoding is not None: + # V0 capability-gates guided decoding at submit; translation refuses it + # too so a bypassed gate cannot smuggle it into the runtime untested. + raise ConversionError("guided_decoding is capability-gated in V0") + sampling = engine_request.sampling + synthetic = SamplingParams( + max_tokens=sampling.max_new_tokens, + end_id=sampling.end_id, + pad_id=sampling.pad_id, + stop_token_ids=list(sampling.stop_token_ids) or None, + min_tokens=sampling.min_tokens, + temperature=sampling.temperature, + top_p=sampling.top_p, + top_k=sampling.top_k, + seed=sampling.seed, + repetition_penalty=sampling.repetition_penalty, + presence_penalty=sampling.presence_penalty, + frequency_penalty=sampling.frequency_penalty, + logprobs=sampling.num_logprobs, + prompt_logprobs=sampling.num_prompt_logprobs, + ) + stop_token_sequences = [list(seq) for seq in sampling.stop_token_sequences] or None + return GenerationRequest(prompt_token_ids=list(engine_request.prompt_token_ids), + sampling_params=synthetic, + streaming=True, + stop_token_sequences=stop_token_sequences) diff --git a/tensorrt_llm/executor/engine_client/envelope.py b/tensorrt_llm/executor/engine_client/envelope.py new file mode 100644 index 000000000000..93ea50c194eb --- /dev/null +++ b/tensorrt_llm/executor/engine_client/envelope.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Normalization of raw worker responses into a uniform envelope. + +The proxy's result-dispatch loop delivers several shapes: the PyTorch +path's ``LlmResponse``, ``ResponseWrapper`` (post-computed logprobs and +perf metrics attached by the worker), and ``ErrorResponse``. This module +normalizes the PyTorch shapes into ``RuntimeResponseEnvelope`` — immutable +snapshots of exactly the fields the contract needs — and rejects, with +typed errors, the shapes V0 does not support (the C++-engine +``tllm.Response`` and the postprocess-parallel ``PostprocWorker.Output``). + +Normalization rules (contract DEC-1/DEC-8 and divergence note 5): +- Both runtime logprob shapes are accepted: chosen-token float lists pass + through; token-id→``Logprob`` maps resolve by exact token-id lookup, and + any missing key, misalignment, or length mismatch is a typed error — + never a silent drop or pad. +- Perf-metric keys convert enum→string via ``.value``. +- Cached-token accounting is read from the same runtime field the legacy + usage path reads. +- A no-token response produces an envelope with no new tokens; the router + decides what (if any) frame that becomes. +""" + +import dataclasses +import math +from typing import Optional + +from ..._torch.pyexecutor.llm_request import LlmResponse +from ..postproc_worker import PostprocWorker +from ..result import ResponseWrapper +from ..utils import ErrorResponse, is_llm_response +from .contract import ContractError + +__all__ = [ + "EnvelopeError", + "RuntimeResponseEnvelope", + "normalize_response", +] + + +class EnvelopeError(ContractError): + """A worker response cannot be normalized. Carries a stable ``reason``.""" + + def __init__(self, reason: str, message: str): + super().__init__(f"{reason}: {message}") + self.reason = reason + + +@dataclasses.dataclass(frozen=True) +class RuntimeResponseEnvelope: + """Immutable snapshot of one worker response, normalized for the router. + + Not a wire type: this is the rank-0-internal handoff between the proxy + tap and the frame router. ``new_token_ids`` is this response's delta + (possibly empty for a no-token final); ``finish_reason_name`` is the raw + engine reason name (e.g. ``"END_ID"``) or ``None`` when the sequence is + not finished. + """ + + client_id: int + is_final: bool + error_msg: Optional[str] = None + sequence_index: int = 0 + new_token_ids: tuple = () + finish_reason_name: Optional[str] = None + logprobs: Optional[tuple] = None + prompt_logprobs: Optional[tuple] = None + # Map-shaped prompt logprobs that arrived on a no-token response cannot + # be normalized yet (the first generated token supplies the last lookup + # key); they are carried raw for the router to hold and normalize on the + # next token-carrying delta via ``resolve_held_prompt_logprobs``. + raw_prompt_logprob_entries: Optional[tuple] = None + metrics: Optional[dict] = None + cached_tokens: Optional[int] = None + + +def _snapshot_float(value, context: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise EnvelopeError("logprob_mismatch", f"{context}: expected a number, got " + f"{type(value).__name__}") + value = float(value) + if not math.isfinite(value): + raise EnvelopeError("logprob_mismatch", f"{context}: non-finite value") + return value + + +def _logprob_entry(entry, expected_token_id: int, context: str) -> float: + """Normalize one logprob entry: float passes, mapping needs exact lookup.""" + if isinstance(entry, dict): + if expected_token_id not in entry: + raise EnvelopeError( + "logprob_mismatch", + f"{context}: token id {expected_token_id} missing from logprob map " + f"(keys={sorted(entry.keys())[:8]!r})") + item = entry[expected_token_id] + value = getattr(item, "logprob", item) + return _snapshot_float(value, context) + return _snapshot_float(entry, context) + + +def _normalize_generation_logprobs(log_probs, new_token_ids, client_id) -> Optional[tuple]: + if log_probs is None: + return None + entries = list(log_probs) + if len(entries) != len(new_token_ids): + raise EnvelopeError( + "logprob_mismatch", + f"client_id={client_id}: {len(entries)} logprob entries for " + f"{len(new_token_ids)} new tokens") + return tuple( + _logprob_entry(entry, token_id, f"client_id={client_id} logprobs[{i}]") + for i, (entry, token_id) in enumerate(zip(entries, new_token_ids))) + + +def _normalize_prompt_logprobs(prompt_entries, prompt_token_ids, new_token_ids, + client_id) -> Optional[tuple]: + if prompt_entries is None: + return None + entries = list(prompt_entries) + if all(not isinstance(entry, dict) for entry in entries): + return tuple( + _snapshot_float(entry, f"client_id={client_id} prompt_logprobs[{i}]") + for i, entry in enumerate(entries)) + # Map-shaped prompt logprobs need the exact per-position token ids: the + # prompt shifted by one plus the first generated token (mirroring the + # worker-side computation offset). + if prompt_token_ids is None: + raise EnvelopeError( + "logprob_mismatch", + f"client_id={client_id}: map-shaped prompt logprobs need the bound " + "prompt token ids") + if not new_token_ids: + # Legitimately possible: the worker re-attaches cached prompt + # logprobs to later (even terminal-only) responses. The caller holds + # the raw entries or drops them if they were already delivered. + return None + expected = list(prompt_token_ids[1:]) + [new_token_ids[0]] + if len(entries) != len(expected): + raise EnvelopeError( + "logprob_mismatch", + f"client_id={client_id}: {len(entries)} prompt-logprob entries for " + f"{len(expected)} expected positions") + return tuple( + _logprob_entry(entry, token_id, f"client_id={client_id} prompt_logprobs[{i}]") + for i, (entry, token_id) in enumerate(zip(entries, expected))) + + +def _normalize_metrics(raw_metrics, client_id) -> Optional[dict]: + if raw_metrics is None: + return None + if not isinstance(raw_metrics, dict): + raise EnvelopeError( + "metrics_mismatch", + f"client_id={client_id}: expected a metrics dict, got " + f"{type(raw_metrics).__name__}") + metrics = {} + for key, value in raw_metrics.items(): + name = getattr(key, "value", key) + if not isinstance(name, str): + name = str(name) + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise EnvelopeError( + "metrics_mismatch", + f"client_id={client_id}: metric {name!r} has non-numeric value " + f"{type(value).__name__}") + value = float(value) + if not math.isfinite(value): + raise EnvelopeError( + "metrics_mismatch", + f"client_id={client_id}: metric {name!r} is not finite") + metrics[name] = value + return metrics + + +def normalize_response(raw, prompt_token_ids=None) -> RuntimeResponseEnvelope: + """Normalize one raw dispatch-loop item into an envelope. + + Args: + raw: One item as delivered to the proxy result-dispatch loop (after + batched lists are unrolled by the caller). + prompt_token_ids: The bound request's prompt token ids, required only + to resolve map-shaped prompt logprobs. + + Raises: + EnvelopeError: for unsupported shapes (C++-engine ``tllm.Response``, + ``PostprocWorker.Output``, unknown objects) and for logprob or + metric content that cannot be normalized losslessly. + """ + wrapper_metrics = None + logprobs_result = None + if isinstance(raw, ResponseWrapper): + wrapper_metrics = raw.request_perf_metrics + logprobs_result = raw.logprobs + raw = raw._response + + if isinstance(raw, ErrorResponse): + return RuntimeResponseEnvelope(client_id=raw.client_id, is_final=True, + error_msg=str(raw.error_msg)) + if isinstance(raw, PostprocWorker.Output): + raise EnvelopeError( + "postproc_parallel_shape", + "PostprocWorker.Output requires num_postprocess_workers=0 in V0") + if not isinstance(raw, LlmResponse): + if is_llm_response(raw): + raise EnvelopeError( + "cpp_engine_shape", + f"{type(raw).__name__} (C++-engine response) is not supported in V0") + raise EnvelopeError("unknown_shape", + f"{type(raw).__name__} is not a recognized response shape") + + client_id = raw.client_id + if raw.has_error(): + return RuntimeResponseEnvelope(client_id=client_id, is_final=True, + error_msg=str(raw.error_msg)) + + result = raw.result + if hasattr(result, "_result") and isinstance(getattr(result, "_result"), bytes): + result.deserialize() + + sequence_index = getattr(result, "sequence_index", 0) or 0 + output_token_ids = result.output_token_ids + source_tokens = output_token_ids[0] if output_token_ids else [] + new_token_ids = tuple(int(token) for token in source_tokens) + + finish_reasons = result.finish_reasons + finish_reason_name = None + if finish_reasons: + reason = finish_reasons[0] + finish_reason_name = getattr(reason, "name", str(reason)) + + log_probs = result.log_probs + logprobs = _normalize_generation_logprobs( + log_probs[0] if log_probs else None, new_token_ids, client_id) + raw_prompt_entries = (logprobs_result.prompt + if logprobs_result is not None else None) + prompt_logprobs = _normalize_prompt_logprobs(raw_prompt_entries, + prompt_token_ids, + new_token_ids, client_id) + raw_prompt_logprob_entries = None + if raw_prompt_entries is not None and prompt_logprobs is None: + # Map-shaped entries on a no-token response: not normalizable yet. + raw_prompt_logprob_entries = tuple(raw_prompt_entries) + + cached_tokens = getattr(result, "cached_tokens", None) + if cached_tokens is not None and (isinstance(cached_tokens, bool) + or not isinstance(cached_tokens, int)): + cached_tokens = None + + return RuntimeResponseEnvelope( + client_id=client_id, + is_final=bool(result.is_final), + sequence_index=int(sequence_index), + new_token_ids=new_token_ids, + finish_reason_name=finish_reason_name, + logprobs=logprobs, + prompt_logprobs=prompt_logprobs, + raw_prompt_logprob_entries=raw_prompt_logprob_entries, + metrics=_normalize_metrics(wrapper_metrics, client_id), + cached_tokens=cached_tokens, + ) + + +def resolve_held_prompt_logprobs(raw_entries, prompt_token_ids, first_token_id, + client_id) -> tuple: + """Normalize held map-shaped prompt logprobs once the first token is known.""" + return _normalize_prompt_logprobs(list(raw_entries), prompt_token_ids, + (first_token_id, ), client_id) diff --git a/tensorrt_llm/executor/engine_client/invariants.py b/tensorrt_llm/executor/engine_client/invariants.py new file mode 100644 index 000000000000..6dc000485daf --- /dev/null +++ b/tensorrt_llm/executor/engine_client/invariants.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pass-through stream-shape checker (test utility). + +``InvariantCheckingStream`` wraps an async frame iterator and enforces the +contract stream invariants: frames belong to one request; per-sequence +ordering; at most one ``Terminal`` per sequence; no ``TokenDelta`` after +its sequence's ``Terminal``; exactly one request-ending frame +(``RequestComplete`` after all started sequences terminated, or a +standalone ``ErrorFrame`` iff nothing started); nothing after the ending +frame; monotonically increasing ``event_seq``; a stream that ends without +an ending frame is a violation. +""" + +from .contract import (ContractError, ErrorFrame, RequestComplete, Terminal, + TokenDelta) + +__all__ = ["StreamInvariantViolation", "InvariantCheckingStream"] + + +class StreamInvariantViolation(ContractError): + """A frame stream violated the contract's stream invariants.""" + + +class InvariantCheckingStream: + """Async pass-through wrapper enforcing the §2.3 stream invariants.""" + + def __init__(self, stream, request_id=None): + self._stream = stream + self._request_id = request_id + self._terminated_sequences = set() + self._started_sequences = set() + self._ended = False + self._last_event_seq = -1 + self._frames_seen = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + try: + frame = await self._stream.__anext__() + except StopAsyncIteration: + if not self._ended: + raise StreamInvariantViolation( + f"stream for {self._request_id!r} ended without a " + "request-ending frame") + raise + self._check(frame) + return frame + + async def aclose(self): + close = getattr(self._stream, "aclose", None) + if close is not None: + await close() + + def _fail(self, message: str): + raise StreamInvariantViolation(f"request {self._request_id!r}: {message}") + + def _check(self, frame): + self._frames_seen += 1 + if self._request_id is None: + self._request_id = frame.request_id + if frame.request_id != self._request_id: + self._fail(f"frame for foreign request {frame.request_id!r}") + if self._ended: + self._fail(f"{type(frame).__name__} after the request-ending frame") + if frame.event_seq <= self._last_event_seq: + self._fail(f"event_seq {frame.event_seq} not increasing " + f"(last {self._last_event_seq})") + self._last_event_seq = frame.event_seq + + if isinstance(frame, TokenDelta): + if frame.sequence_id in self._terminated_sequences: + self._fail(f"TokenDelta after Terminal for sequence " + f"{frame.sequence_id}") + self._started_sequences.add(frame.sequence_id) + elif isinstance(frame, Terminal): + if frame.sequence_id in self._terminated_sequences: + self._fail(f"duplicate Terminal for sequence {frame.sequence_id}") + self._terminated_sequences.add(frame.sequence_id) + self._started_sequences.add(frame.sequence_id) + elif isinstance(frame, RequestComplete): + unterminated = self._started_sequences - self._terminated_sequences + if unterminated: + self._fail(f"RequestComplete before Terminal for sequences " + f"{sorted(unterminated)}") + self._ended = True + elif isinstance(frame, ErrorFrame): + if self._started_sequences: + self._fail("standalone ErrorFrame on a request with started " + "sequences (must use Terminal + RequestComplete)") + self._ended = True + else: + self._fail(f"unknown frame type {type(frame).__name__}") diff --git a/tensorrt_llm/executor/engine_client/local_client.py b/tensorrt_llm/executor/engine_client/local_client.py new file mode 100644 index 000000000000..cd5c3a07d249 --- /dev/null +++ b/tensorrt_llm/executor/engine_client/local_client.py @@ -0,0 +1,345 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""In-process ``EngineClient``: setup gates, pre-submit checks, async streams. + +``LocalProcessEngineClient`` is the V0 (in-process) implementation of the +engine-client interface: it attaches an ``EngineFrameRouter`` to the live +``GenerationExecutorProxy``, translates ``EngineRequest`` compositionally +into the worker-facing request, and exposes the typed frame stream as an +async-first ``FrameStream``. Nothing on this surface leaks in-process +affordances: liveness and failure signaling are behind the interface, and +every value that crosses it is a wire type or a typed control-plane result. +""" + +import asyncio +import dataclasses +import json +import os +from typing import Optional + +from .contract import (ENGINE_CONTRACT_VERSION, ContractError, + EngineCapabilities, EngineHealth, EngineRequest, + ErrorFrame, FrontendOutputConfig, IterationStatsBatch, + KvCacheEventsBatch, RequestComplete, Terminal) +from .conversion import (ConversionError, RequestIneligibleError, + derive_required_features) +from .router import RequestBinding, RouterError + +__all__ = [ + "ENGINE_CLIENT_FLAG_ENV", + "V0_CAPABILITY_FEATURES", + "EngineClientConfigError", + "RequestRejectedError", + "EngineClientConfig", + "FrameStream", + "LocalProcessEngineClient", + "engine_client_flag_enabled", +] + +ENGINE_CLIENT_FLAG_ENV = "TLLM_EXPERIMENTAL_ENGINE_CLIENT" + +V0_CAPABILITY_FEATURES = ("streaming", "logprobs", "prompt_logprobs", + "stop_token_sequences", "abort", "usage") + + +class EngineClientConfigError(ContractError): + """Unsupported configuration detected at client construction.""" + + +class RequestRejectedError(ContractError): + """Typed pre-submit rejection (capability, duplicate, or malformed).""" + + +def resolve_engine_client_flag(args_value: bool = False) -> bool: + """Resolve the effective flag: the environment wins in both directions. + + Presence, not truthiness, decides the override: unset env defers to the + configured value; ``"1"``/``"0"`` force the path on/off regardless of + configuration; anything else fails closed with a typed config error. + """ + env = os.environ.get(ENGINE_CLIENT_FLAG_ENV) + if env is None: + return bool(args_value) + if env == "1": + return True + if env == "0": + return False + raise EngineClientConfigError( + f"flag: {ENGINE_CLIENT_FLAG_ENV} must be '0' or '1', got {env!r}") + + +def engine_client_flag_enabled() -> bool: + return resolve_engine_client_flag(False) + + +@dataclasses.dataclass(frozen=True) +class EngineClientConfig: + """Deployment facts the setup gates validate. + + The caller (the serving layer) fills this from its own configuration; + the client rejects anything outside the GPU-validated V0 envelope with + a typed config error at construction — never mid-run. + """ + + backend: Optional[str] = None + transport: str = "ipc_proxy" + num_postprocess_workers: int = 0 + post_processor_hook_set: bool = False + speculative_config_set: bool = False + early_first_token_mode: bool = False + world_size: int = 1 + tokenizer_trust_remote_code: bool = False + flag_enabled: Optional[bool] = None # None -> read the env var + + +class FrameStream: + """Async single-consumer stream of ``OutputFrame`` for one request. + + Frames buffer from submit time, so nothing emitted before the first + ``__anext__`` is lost. The stream ends after the request-ending frame. + ``aclose()`` (or ``async with``) triggers the abort-if-incomplete + obligation explicitly — never rely on ``GeneratorExit``/GC timing. + """ + + def __init__(self, client: "LocalProcessEngineClient", binding: RequestBinding): + self._client = client + self._binding = binding + self._finished = False + self._closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._finished or self._closed: + raise StopAsyncIteration + loop = asyncio.get_running_loop() + while True: + frame, ready = self._binding.delivery.pop_nowait() + if ready: + if isinstance(frame, (RequestComplete, ErrorFrame)): + self._finished = True + # Consumed to the ending frame: retire the delivery so + # the request id becomes reusable and router state + # returns to its bounds. + self._client._router.retire_delivery( + self._binding.request_id) + return frame + if self._binding.delivery.closed: + self._finished = True + raise StopAsyncIteration + future = loop.create_future() + if not self._binding.delivery.register_waiter(loop, future): + continue # data arrived while registering + await future + + def pop_ready(self) -> list: + """Pop all immediately-available frames without awaiting. + + Used by consumers that drain-then-render so a final ``TokenDelta`` + and its ``Terminal`` (enqueued together) are processed as one batch. + """ + frames = [] + if self._finished or self._closed: + return frames + while True: + frame, ready = self._binding.delivery.pop_nowait() + if not ready: + return frames + frames.append(frame) + if isinstance(frame, (RequestComplete, ErrorFrame)): + self._finished = True + self._client._router.retire_delivery(self._binding.request_id) + return frames + + async def aclose(self) -> None: + """Close the stream; aborts the runtime request if incomplete.""" + if self._closed: + return + self._closed = True + if not self._finished: + try: + self._client._abort_binding(self._binding) + finally: + self._binding.delivery.close() + # An explicit close retires the delivery either way (the abort's + # eventual ending is absorbed; the id becomes reusable once ended). + self._client._router.retire_delivery(self._binding.request_id) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + await self.aclose() + + +class LocalProcessEngineClient: + """V0 in-process engine client over the IPC-proxy executor.""" + + def __init__(self, executor, config: EngineClientConfig): + self._validate_config(config) + if getattr(executor, "attach_engine_service", None) is None: + raise EngineClientConfigError( + "transport: executor does not expose the engine-service hooks " + "(IPC GenerationExecutorProxy required)") + self._executor = executor + self._config = config + self._capabilities = EngineCapabilities(features=V0_CAPABILITY_FEATURES) + # The service owns contract execution (authoritative routing: contract + # responses never reach legacy delivery); this client is the thin + # in-process transport adapter in front of it. + from .service import EngineService + self._service = EngineService(executor) + self._router = self._service.router + self._closed = False + executor.attach_engine_service(self._service) + + @staticmethod + def _validate_config(config: EngineClientConfig) -> None: + # The environment wins in both directions, even over an explicit + # flag_enabled value (see resolve_engine_client_flag). + flag = resolve_engine_client_flag(bool(config.flag_enabled)) + if not flag: + raise EngineClientConfigError( + f"flag: set {ENGINE_CLIENT_FLAG_ENV}=1 to enable the " + "experimental engine client") + # Gate on the explicit backend name: `_is_pytorch_backend`-style + # checks are also true for the transitional _autodeploy backend, + # which is explicitly rejected here. + if config.backend != "pytorch": + raise EngineClientConfigError( + f"backend: {config.backend!r} is not supported (pytorch only; " + "the transitional _autodeploy backend is explicitly rejected)") + if config.transport != "ipc_proxy": + raise EngineClientConfigError( + f"transport: {config.transport!r} is not supported (IPC proxy only)") + if config.num_postprocess_workers != 0: + raise EngineClientConfigError( + "postproc_workers: num_postprocess_workers must be 0") + if config.post_processor_hook_set: + raise EngineClientConfigError( + "post_processor_hook: the global output hook is not supported") + if config.speculative_config_set: + raise EngineClientConfigError( + "speculative_config: speculative decoding is not validated") + if config.early_first_token_mode: + raise EngineClientConfigError( + "early_first_token: early first-token response mode is not supported") + if config.world_size != 1: + raise EngineClientConfigError( + f"topology: world_size={config.world_size} exceeds the GPU-validated " + "set (TP1 over the IPC proxy)") + if config.tokenizer_trust_remote_code: + raise EngineClientConfigError( + "trust_remote_code_tokenizer: tokenizer provenance cannot be pinned") + + # ------------------------------------------------------------------ # + # EngineClient interface + # ------------------------------------------------------------------ # + + def capabilities(self) -> EngineCapabilities: + return self._capabilities + + def submit(self, engine_request: EngineRequest, + output_config: Optional[FrontendOutputConfig] = None) -> str: + """Submit an eligible request; returns its request id. + + Pre-submit checks (typed, before the engine sees anything): + protocol version, re-derived ``required_features`` vs the caller's, + capability subset, duplicate id. ``output_config`` never crosses the + contract; it only supplies the ordered stop-reason association the + router uses to resolve ``Terminal.stop_reason``. + """ + if self._closed: + raise RequestRejectedError("client is closed") + if self._router.fatal_error is not None: + raise RequestRejectedError( + f"engine failed: {self._router.fatal_error}") + if not isinstance(engine_request, EngineRequest): + raise RequestRejectedError( + f"expected EngineRequest, got {type(engine_request).__name__}") + if engine_request.protocol_version > ENGINE_CONTRACT_VERSION: + raise RequestRejectedError( + f"protocol_version {engine_request.protocol_version} is newer than " + f"supported {ENGINE_CONTRACT_VERSION}") + derived = derive_required_features(engine_request) + if tuple(engine_request.required_features) != derived: + raise RequestRejectedError( + f"required_features {engine_request.required_features!r} do not match " + f"the request's own fields (derived {derived!r})") + missing = set(derived) - set(self._capabilities.features) + if missing: + raise RequestRejectedError( + f"capabilities not supported by this engine: {sorted(missing)}") + + stop_reasons = (output_config.stop_sequence_reasons + if output_config is not None else ()) + try: + self._executor.submit_contract(engine_request, + stop_reasons=stop_reasons) + except (RequestIneligibleError, ConversionError): + raise + except RouterError as e: + raise RequestRejectedError(str(e)) from e + return engine_request.request_id + + def stream(self, request_id: str) -> FrameStream: + """Open the single consumer stream for a submitted request. + + Frames buffer from submit time, so a delayed open loses nothing; + opening twice is a typed error (single consumer). + """ + try: + binding = self._router.open_stream_binding(request_id) + except RouterError as e: + raise RequestRejectedError(str(e)) from e + return FrameStream(self, binding) + + def abort(self, request_id: str) -> None: + self._router.abort(request_id) + + def get_stats(self, timeout: float = 2.0) -> IterationStatsBatch: + return self._service.get_stats(timeout=timeout) + + def get_kv_events(self, timeout: float = 2.0) -> KvCacheEventsBatch: + return self._service.get_kv_events(timeout=timeout) + + def health(self) -> EngineHealth: + return self._service.health() + + def close_client(self) -> None: + """Detach from the engine: poison this client's streams AND abort its + in-flight engine work (best effort). Does NOT shut the shared engine + down; see ``shutdown_engine``. + """ + if self._closed: + return + self._closed = True + self._service.close_client() + + def shutdown_engine(self) -> None: + """Privileged: shut the engine itself down (V0 in-process owner only).""" + self.close_client() + self._executor.shutdown() + + # ------------------------------------------------------------------ # + # Internal + # ------------------------------------------------------------------ # + + def _abort_binding(self, binding: RequestBinding) -> None: + try: + self._router.abort(binding.request_id) + except RouterError: + pass # already ended diff --git a/tensorrt_llm/executor/engine_client/router.py b/tensorrt_llm/executor/engine_client/router.py new file mode 100644 index 000000000000..885cefa5dee0 --- /dev/null +++ b/tensorrt_llm/executor/engine_client/router.py @@ -0,0 +1,793 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Per-request exactly-once terminal state machine over the proxy tap. + +The proxy natively pops each request's result on the final response and +drops late/duplicate/unknown-id frames; abort is asynchronous best-effort. +This router builds the opposite guarantees on top of that behavior: every +runtime-started sequence gets exactly one ``Terminal``, every bound request +gets exactly one request-ending frame (``RequestComplete``, or a standalone +``ErrorFrame`` iff nothing runtime-started), nothing is emitted after the +ending frame, and a stream never hangs. + +Concurrency specification (normative for this module): + +- **State owner.** All binding, tombstone, and per-request lifecycle state + is owned by ``EngineFrameRouter`` and mutated only under its single + ``_lock``. Threads involved: the proxy submit thread (``observe_submit``, + ``on_submit_enqueue_failed``), the proxy dispatch thread + (``on_response``), consumer event-loop threads (``open_stream`` / + ``abort`` / ``FrameStream.aclose``), and the proxy shutdown path + (``fail_all``). Any of these may run concurrently. +- **Lock hierarchy.** ``EngineFrameRouter._lock`` is the outermost and only + router lock; ``_Delivery`` has its own inner lock and never calls back + into the router while holding it. Never acquire ``_lock`` while holding a + delivery lock. +- **Reentrancy.** Frame emission never recursively invokes abort/failure + paths: runtime aborts triggered by router logic (overflow, early close, + router-internal failure) are collected as deferred actions under the lock + and executed after it is released, so a failure while aborting cannot + corrupt state transitions. +- **Atomic bind→enqueue.** Binding happens on the submit thread after + client-id assignment and strictly before the request is enqueued + (``observe_submit`` is called by the proxy between those points); + ``on_submit_enqueue_failed`` rolls the binding back into a standalone + ``ErrorFrame`` ending, so a request can never be in-flight without a + binding or bound without an ending path. +- **Abort-vs-final races.** ``abort`` marks the state under the lock and + issues the runtime abort outside it. Whatever arrives afterwards — a + ``CANCELLED`` final, a normal final that won the race, or nothing until + ``fail_all`` — resolves through the same single transition point + (``_finish_from_envelope`` / ``_fail_locked``), which is idempotent per + request: the first ending wins, later signals are absorbed and counted. +- **State lifetimes.** Router execution state is pruned on the ending + frame; the delivery queue survives until consumed, closed, or expired by + the client's policy; ended request ids move to a bounded tombstone map + (``abort`` on a tombstone is a no-op, on an unknown id a typed error). +""" + +import threading +from collections import OrderedDict, deque +from typing import Callable, Dict, List, Optional, Tuple + +from tensorrt_llm.logger import logger + +from .contract import (CACHED_TOKENS_METRIC_KEY, ContractError, ErrorFrame, + RequestComplete, Terminal, TokenDelta) +from .envelope import (RuntimeResponseEnvelope, normalize_response, + resolve_held_prompt_logprobs) + +__all__ = [ + "RouterError", + "UnknownRequestError", + "RequestBinding", + "EngineFrameRouter", + "DEFAULT_DELIVERY_LIMIT", + "DEFAULT_TOMBSTONE_LIMIT", +] + +DEFAULT_DELIVERY_LIMIT = 8192 +DEFAULT_TOMBSTONE_LIMIT = 4096 + + +class RouterError(ContractError): + """Router-level typed error.""" + + +class UnknownRequestError(RouterError): + """The request id is neither active nor recently ended.""" + + +class _Delivery: + """Bounded, thread-safe frame buffer with an out-of-band ending slot. + + Frames buffer from submit time so nothing is lost before the first + consumer arrives. Ending frames (``Terminal`` / ``RequestComplete`` / + ``ErrorFrame``) go to a dedicated slot so typed termination never + depends on space in the bounded data buffer. + """ + + def __init__(self, limit: int = DEFAULT_DELIVERY_LIMIT): + self._lock = threading.Lock() + self._buffer = deque() + self._end_frames: List = [] + self._limit = limit + self._waiter = None # (loop, future) of the single blocked consumer + self._closed = False + self.overflowed = False + self.dropped_frames = 0 + + def _wake_locked(self): + if self._waiter is not None: + loop, future = self._waiter + self._waiter = None + + def _set(): + if not future.done(): + future.set_result(None) + + try: + loop.call_soon_threadsafe(_set) + except RuntimeError: + # The consumer's event loop is already closed; the frames + # stay buffered for a later pop, nothing to wake. + pass + + def put(self, frame) -> bool: + """Queue a data frame. Returns False on overflow (frame dropped).""" + with self._lock: + if self._closed: + return False + if len(self._buffer) >= self._limit: + self.overflowed = True + self.dropped_frames += 1 + return False + self._buffer.append(frame) + self._wake_locked() + return True + + def put_ending(self, frames) -> None: + """Queue ending frames via the out-of-band slot (never blocked).""" + with self._lock: + if self._closed: + return + self._end_frames.extend(frames) + self._wake_locked() + + def close(self) -> None: + with self._lock: + self._closed = True + self._buffer.clear() + self._wake_locked() + + def pop_nowait(self): + """Pop one frame if available: data first, then ending frames. + + Returns (frame, True) or (None, False). + """ + with self._lock: + if self._buffer: + return self._buffer.popleft(), True + if self._end_frames: + return self._end_frames.pop(0), True + return None, False + + def register_waiter(self, loop, future) -> bool: + """Register the single consumer waiter; False if data already ready.""" + with self._lock: + if self._buffer or self._end_frames or self._closed: + return False + self._waiter = (loop, future) + return True + + @property + def closed(self) -> bool: + return self._closed + + +class RequestBinding: + """Per-request execution state (owned by the router, guarded by its lock).""" + + __slots__ = ("request_id", "client_id", "prompt_token_ids", "stop_reasons", + "delivery", "runtime_started", "terminal_emitted", "ended", + "abort_requested", "event_seq", "completion_tokens", + "cached_tokens", "prompt_logprobs_sent", "held_prompt_logprobs", + "held_raw_prompt_entries", + "recent_tokens", "pending_final_metrics", "final_status", + "stream_opened", "retire_when_ended") + + def __init__(self, request_id: str, prompt_token_ids: Tuple[int, ...], + stop_reasons: Tuple, delivery_limit: int): + self.request_id = request_id + self.client_id: Optional[int] = None + self.prompt_token_ids = tuple(prompt_token_ids) + # Ordered (stop_token_sequence, user_visible_reason) association; + # configuration order, first match wins. Validated primitive-only: + # this crosses into engine-side state and must stay encodable. + validated = [] + for pair in stop_reasons: + sequence, reason = pair + sequence = tuple(sequence) + if not sequence or not all( + isinstance(t, int) and not isinstance(t, bool) + for t in sequence): + raise RouterError( + f"stop_reasons: invalid stop sequence {sequence!r}") + if isinstance(reason, bool) or not isinstance(reason, (int, str)): + raise RouterError( + f"stop_reasons: reason must be int or str, got " + f"{type(reason).__name__}") + validated.append((sequence, reason)) + self.stop_reasons = tuple(validated) + self.delivery = _Delivery(delivery_limit) + self.runtime_started = False + self.terminal_emitted = False + self.ended = False + self.abort_requested = False + self.event_seq = 0 + self.completion_tokens = 0 + self.cached_tokens: Optional[int] = None + self.prompt_logprobs_sent = False + self.held_prompt_logprobs: Optional[tuple] = None + self.held_raw_prompt_entries: Optional[tuple] = None + tail = max((len(seq) for seq, _ in self.stop_reasons), default=0) + self.recent_tokens = deque(maxlen=max(tail, 1)) + self.pending_final_metrics: Optional[dict] = None + self.final_status: Optional[str] = None + self.stream_opened = False + # Set when the consumer retires (closes) the delivery before the + # request has ended: the ending response must complete the + # retirement so the request id does not stay retained until + # tombstone eviction. + self.retire_when_ended = False + + def next_seq(self) -> int: + seq = self.event_seq + self.event_seq += 1 + return seq + + +class EngineFrameRouter: + """Proxy-attached fork target producing the typed engine frame stream. + + ``abort_fn(client_id)`` issues the runtime abort (typically + ``executor.abort_request``); it is always invoked outside the router + lock (deferred-action rule of the concurrency specification). + """ + + def __init__(self, + abort_fn: Optional[Callable[[int], None]] = None, + delivery_limit: int = DEFAULT_DELIVERY_LIMIT, + tombstone_limit: int = DEFAULT_TOMBSTONE_LIMIT): + self._lock = threading.Lock() + self._abort_fn = abort_fn + self._delivery_limit = delivery_limit + self._tombstone_limit = tombstone_limit + # Pending registrations keyed by GenerationRequest object identity, + # matched by observe_submit on the submit thread. + self._pending: Dict[int, RequestBinding] = {} + self._by_client: Dict[int, RequestBinding] = {} + self._by_request: Dict[str, RequestBinding] = {} + # Delivery retention: a request's buffered frames stay reachable for + # a (delayed) stream open even after the request ends, until evicted + # alongside its tombstone. + self._delivery_index: "OrderedDict[str, RequestBinding]" = OrderedDict() + self._tombstones: "OrderedDict[str, str]" = OrderedDict() + # BOUNDED recently-retired proxy client ids, used only for + # late/duplicate accounting. Safety does not depend on this memory: + # proxy client ids are never reused within a proxy lifetime (the + # allocator is monotonic and wrap/collision is rejected at submit), + # so a very late frame whose id aged out of this structure falls + # through to the legacy lookup, finds nothing, and is dropped exactly + # as legacy late frames are dropped today. Frontend request ids + # become REUSABLE once the prior request ended and its delivery is + # retired (consumed to the ending frame, explicitly closed, or + # evicted by the bounded tombstone policy) — per plan AC-4. + self._retired_client_ids: "OrderedDict[int, bool]" = OrderedDict() + self._retired_limit = tombstone_limit + self._fatal_error: Optional[str] = None + # Observability counters (read by the serving layer). + self.counters = { + "late_or_duplicate_absorbed": 0, + "router_failures": 0, + "synthesized_terminals": 0, + "overflow_aborts": 0, + } + + # ------------------------------------------------------------------ # + # Registration (client side, submit thread) + # ------------------------------------------------------------------ # + + def register_pending(self, generation_request, request_id: str, + prompt_token_ids, stop_reasons) -> RequestBinding: + """Register a contract request about to be submitted to the proxy. + + Called by the client immediately before ``proxy.submit``; the proxy's + ``observe_submit`` hook completes the binding with the assigned + client id before the request is enqueued (same thread, race-free). + """ + binding = RequestBinding(request_id, prompt_token_ids, stop_reasons, + self._delivery_limit) + with self._lock: + self._check_request_id_free_locked(request_id) + self._pending[id(generation_request)] = binding + self._by_request[request_id] = binding + self._delivery_index[request_id] = binding + return binding + + def _check_request_id_free_locked(self, request_id: str) -> None: + """Reject a duplicate id while the previous use is still live. + + A frontend request id is reusable once its prior request ended AND + its delivery is retired; it is a typed duplicate only while the + previous binding is active or its delivery is still retained. + """ + if request_id in self._by_request: + raise RouterError(f"duplicate request_id {request_id!r} " + "(previous request still active)") + if request_id in self._delivery_index: + raise RouterError(f"duplicate request_id {request_id!r} " + "(previous delivery not yet retired)") + + def bind(self, request_id: str, client_id: int, prompt_token_ids, + stop_reasons) -> RequestBinding: + """Directly bind a contract request to its allocated client id. + + The contract-native submit path (no pending/object-identity interval): + the caller allocates the client id and enqueues only after this + returns, so a worker response can never beat the binding. + """ + binding = RequestBinding(request_id, prompt_token_ids, stop_reasons, + self._delivery_limit) + binding.client_id = client_id + with self._lock: + self._check_request_id_free_locked(request_id) + if client_id in self._by_client or client_id in self._retired_client_ids: + raise RouterError(f"client id {client_id} already contract-owned") + self._by_request[request_id] = binding + self._by_client[client_id] = binding + self._delivery_index[request_id] = binding + return binding + + def observe_submit(self, generation_request) -> None: + """Proxy hook: bind the assigned client id (before enqueue).""" + with self._lock: + binding = self._pending.pop(id(generation_request), None) + if binding is None: + return # legacy (non-contract) submission + binding.client_id = generation_request.id + self._by_client[generation_request.id] = binding + + def on_submit_enqueue_failed(self, client_id: int) -> None: + """Proxy hook: the request was bound but never reached the worker.""" + with self._lock: + binding = self._by_client.get(client_id) + if binding is None or binding.ended: + return + # Nothing runtime-started: standalone ErrorFrame ending. + self._end_locked(binding, [ + ErrorFrame(request_id=binding.request_id, error_code="enqueue_failed", + message="request could not be enqueued to the engine", + event_seq=binding.next_seq()) + ]) + + def discard_pending(self, generation_request) -> None: + """Roll back a registration whose proxy.submit never ran.""" + with self._lock: + binding = self._pending.pop(id(generation_request), None) + if binding is not None: + self._by_request.pop(binding.request_id, None) + self._delivery_index.pop(binding.request_id, None) + binding.delivery.close() + + # ------------------------------------------------------------------ # + # Response path (proxy dispatch thread) + # ------------------------------------------------------------------ # + + def route_response(self, raw) -> bool: + """Exclusively claim and consume a contract-owned response. + + Returns True when the client id belongs to the contract population + (the response was consumed or deliberately absorbed) — the caller + must then skip all legacy delivery. False means never contract-owned. + """ + client_id = getattr(raw, "client_id", None) + with self._lock: + if (client_id not in self._by_client + and client_id not in self._retired_client_ids): + # Never (or no longer) contract-tracked. Falling through is + # safe: proxy ids are never reused, so the legacy lookup + # finds nothing and drops the frame like any legacy late + # frame. + return False + self.on_response(raw) + return True + + def on_response(self, raw) -> None: + """Proxy hook: observe one raw dispatch item (never raises).""" + deferred: List[Callable[[], None]] = [] + try: + client_id = getattr(raw, "client_id", None) + with self._lock: + binding = self._by_client.get(client_id) + if binding is None: + if client_id in self._retired_client_ids: + # Recently retired contract id: absorb and count. + self.counters["late_or_duplicate_absorbed"] += 1 + return # unbound legacy traffic + if binding.ended: + self.counters["late_or_duplicate_absorbed"] += 1 + return + try: + envelope = normalize_response(raw, binding.prompt_token_ids) + except Exception as e: + self.counters["router_failures"] += 1 + self._fail_request(binding, "router_error", + f"response normalization failed: {e}", deferred) + logger.error(f"engine-frame router: normalization failure for " + f"client_id={client_id}: {e!r}") + return + with self._lock: + if binding.ended: + self.counters["late_or_duplicate_absorbed"] += 1 + return + self._process_envelope_locked(binding, envelope) + except Exception as e: # never propagate into the dispatch thread + self.counters["router_failures"] += 1 + logger.error(f"engine-frame router: internal error: {e!r}") + try: + if 'binding' in locals() and binding is not None: + self._fail_request(binding, "router_error", str(e), deferred) + except Exception: + pass + finally: + for action in deferred: + try: + action() + except Exception as e: + logger.error(f"engine-frame router: deferred action failed: {e!r}") + + def _process_envelope_locked(self, binding: RequestBinding, + envelope: RuntimeResponseEnvelope) -> None: + if envelope.error_msg is not None: + self._end_with_error_locked(binding, "request_error", envelope.error_msg) + return + if envelope.sequence_index != 0: + # V0 lifecycle state is sequence-0 only (n=1); a multi-sequence + # response must fail the request rather than mix sequence state. + self._end_with_error_locked( + binding, "router_error", + f"sequence_index {envelope.sequence_index} unsupported in V0") + return + + binding.runtime_started = True + + if not envelope.new_token_ids and not binding.prompt_logprobs_sent: + # Hold for the next token-carrying delta (empty TokenDeltas do + # not exist by construction). Map-shaped entries arrive raw: the + # first generated token supplies the last lookup key. + if envelope.prompt_logprobs is not None: + binding.held_prompt_logprobs = envelope.prompt_logprobs + elif envelope.raw_prompt_logprob_entries is not None: + binding.held_raw_prompt_entries = envelope.raw_prompt_logprob_entries + # Prompt logprobs re-attached AFTER delivery (the worker caches and + # re-sends them, including on a terminal-only final) are dropped + # quietly: they were already emitted exactly once. + + if envelope.new_token_ids: + prompt_logprobs = None + if not binding.prompt_logprobs_sent: + prompt_logprobs = envelope.prompt_logprobs or binding.held_prompt_logprobs + if prompt_logprobs is None and binding.held_raw_prompt_entries is not None: + try: + prompt_logprobs = resolve_held_prompt_logprobs( + binding.held_raw_prompt_entries, + binding.prompt_token_ids, + envelope.new_token_ids[0], binding.client_id) + except Exception as e: + self._end_with_error_locked( + binding, "router_error", + f"held prompt-logprob resolution failed: {e}") + return + if prompt_logprobs is not None: + binding.prompt_logprobs_sent = True + binding.held_prompt_logprobs = None + binding.held_raw_prompt_entries = None + metrics = dict(envelope.metrics) if envelope.metrics else {} + if envelope.cached_tokens is not None: + metrics[CACHED_TOKENS_METRIC_KEY] = float(envelope.cached_tokens) + binding.cached_tokens = envelope.cached_tokens + binding.completion_tokens += len(envelope.new_token_ids) + binding.recent_tokens.extend(envelope.new_token_ids) + delta = TokenDelta(request_id=binding.request_id, + sequence_id=envelope.sequence_index, + new_token_ids=envelope.new_token_ids, + logprobs=envelope.logprobs, + prompt_logprobs=prompt_logprobs, + metrics=metrics or None, + event_seq=binding.next_seq()) + if not binding.delivery.put(delta): + if binding.delivery.closed: + # The consumer closed the stream (abort already issued by + # the close path); absorb the delta — but fall through so + # a token-carrying FINAL still ends (and, with retirement + # pending, retires) the binding instead of leaking it. + self.counters["late_or_duplicate_absorbed"] += 1 + else: + self._overflow_locked(binding) + return + else: + if envelope.is_final and envelope.metrics: + binding.pending_final_metrics = dict(envelope.metrics) + if envelope.cached_tokens is not None and binding.cached_tokens is None: + # A zero-token request may report cached tokens only on its + # final response (contract divergence note 5). When a delta + # already recorded a value, leave it for the finish-time + # consistency check instead of overwriting. + binding.cached_tokens = envelope.cached_tokens + + if envelope.is_final: + if (binding.held_prompt_logprobs is not None + or binding.held_raw_prompt_entries is not None): + logger.warning( + f"engine-frame router: dropping prompt logprobs for " + f"request {binding.request_id!r}: the request ended before a " + f"token-carrying delta arrived") + binding.held_prompt_logprobs = None + binding.held_raw_prompt_entries = None + self._finish_from_envelope_locked(binding, envelope) + + def _resolve_stop_reason(self, binding: RequestBinding): + tail = tuple(binding.recent_tokens) + for sequence, reason in binding.stop_reasons: + if len(sequence) <= len(tail) and tail[-len(sequence):] == sequence: + return reason + return None + + def _finish_from_envelope_locked(self, binding: RequestBinding, + envelope: RuntimeResponseEnvelope) -> None: + name = envelope.finish_reason_name + stop_reason = None + if name == "END_ID": + finish_reason = "stop" + elif name == "STOP_WORDS": + finish_reason = "stop" + stop_reason = self._resolve_stop_reason(binding) + elif name == "LENGTH": + finish_reason = "length" + elif name == "CANCELLED": + finish_reason = "abort" + elif name == "TIMED_OUT": + finish_reason = "error" + stop_reason = "timeout" + elif name is None or name == "NOT_FINISHED": + if binding.abort_requested: + finish_reason = "abort" + else: + finish_reason = "error" + stop_reason = "not_finished" + else: + finish_reason = "error" + stop_reason = name + status = {"stop": "ok", "length": "ok", "abort": "aborted", + "error": "failed"}[finish_reason] + if (envelope.cached_tokens is not None + and binding.cached_tokens is not None + and envelope.cached_tokens != binding.cached_tokens): + # Consistency rule (divergence note 5): the final value must + # equal the last per-delta value — a mismatch is a typed failure, + # not a warning. + self._end_with_error_locked( + binding, "cached_tokens_mismatch", + f"final cached_tokens {envelope.cached_tokens} != last delta " + f"{binding.cached_tokens}") + return + if envelope.cached_tokens is not None: + binding.cached_tokens = envelope.cached_tokens + frames = [] + if not binding.terminal_emitted: + binding.terminal_emitted = True + frames.append(Terminal(request_id=binding.request_id, + sequence_id=envelope.sequence_index, + finish_reason=finish_reason, + stop_reason=stop_reason, + event_seq=binding.next_seq())) + metrics = binding.pending_final_metrics + if envelope.metrics and not envelope.new_token_ids: + metrics = dict(envelope.metrics) + frames.append( + RequestComplete(request_id=binding.request_id, status=status, + prompt_tokens=len(binding.prompt_token_ids), + completion_tokens=binding.completion_tokens, + cached_tokens=binding.cached_tokens, + metrics=metrics, + event_seq=binding.next_seq())) + self._end_locked(binding, frames) + + def _end_with_error_locked(self, binding: RequestBinding, error_code: str, + message: str) -> None: + frames = [] + if binding.runtime_started: + if not binding.terminal_emitted: + binding.terminal_emitted = True + self.counters["synthesized_terminals"] += 1 + frames.append(Terminal(request_id=binding.request_id, sequence_id=0, + finish_reason="error", stop_reason=None, + event_seq=binding.next_seq())) + frames.append( + RequestComplete(request_id=binding.request_id, status="failed", + prompt_tokens=len(binding.prompt_token_ids), + completion_tokens=binding.completion_tokens, + cached_tokens=binding.cached_tokens, + metrics=binding.pending_final_metrics, + event_seq=binding.next_seq())) + else: + frames.append(ErrorFrame(request_id=binding.request_id, + error_code=error_code, message=message, + event_seq=binding.next_seq())) + self._end_locked(binding, frames) + + def _end_locked(self, binding: RequestBinding, frames) -> None: + binding.ended = True + binding.delivery.put_ending(frames) + # The client-id mapping is retained (until delivery retirement) so + # late/duplicate frames for an ended request are recognized and + # absorbed rather than mistaken for unbound legacy traffic. + self._by_request.pop(binding.request_id, None) + self._tombstones[binding.request_id] = binding.final_status or "ended" + while len(self._tombstones) > self._tombstone_limit: + evicted_id, _ = self._tombstones.popitem(last=False) + evicted = self._delivery_index.pop(evicted_id, None) + if evicted is not None: + self._retire_binding_locked(evicted) + if binding.retire_when_ended: + # The consumer already closed the stream; complete the deferred + # retirement now that the ending response has arrived. + self._retire_binding_locked(binding) + + def _retire_binding_locked(self, binding: RequestBinding) -> None: + """Retire a binding: close its delivery, free its request id for + reuse, and move its client id to the bounded recently-retired set.""" + binding.delivery.close() + self._delivery_index.pop(binding.request_id, None) + if binding.client_id is not None: + self._by_client.pop(binding.client_id, None) + self._retired_client_ids[binding.client_id] = True + while len(self._retired_client_ids) > self._retired_limit: + self._retired_client_ids.popitem(last=False) + + def retire_delivery(self, request_id: str) -> None: + """Client hook: the stream was consumed to its ending frame or + explicitly closed — retire the binding so the request id becomes + reusable and per-request state returns to its bounds. + + On a close BEFORE the request ended (early ``aclose``), retirement + is recorded and deferred to the ending response: the client id must + stay bound until then so the in-flight cancellation final is + recognized and absorbed rather than treated as unbound traffic. + """ + with self._lock: + binding = self._delivery_index.get(request_id) + if binding is None: + return + if not binding.ended: + binding.retire_when_ended = True + return + self._retire_binding_locked(binding) + + def _overflow_locked(self, binding: RequestBinding) -> None: + self.counters["overflow_aborts"] += 1 + binding.abort_requested = True + # Synthesize the typed ending via the out-of-band slot; the runtime + # abort is issued by the client-facing abort path (slow consumer). + if not binding.terminal_emitted: + binding.terminal_emitted = True + self.counters["synthesized_terminals"] += 1 + self._end_locked(binding, [ + Terminal(request_id=binding.request_id, sequence_id=0, + finish_reason="error", stop_reason="delivery_overflow", + event_seq=binding.next_seq()), + RequestComplete(request_id=binding.request_id, status="failed", + prompt_tokens=len(binding.prompt_token_ids), + completion_tokens=binding.completion_tokens, + cached_tokens=binding.cached_tokens, + metrics=None, + event_seq=binding.next_seq()), + ]) + client_id = binding.client_id + if self._abort_fn is not None and client_id is not None: + # Deferred-action rule: schedule outside the lock via a thread — + # on_response's finally block cannot see this, so use a plain + # daemon thread to stay non-reentrant. + threading.Thread(target=self._safe_abort, args=(client_id, ), + daemon=True).start() + + def _safe_abort(self, client_id: int) -> None: + try: + self._abort_fn(client_id) + except Exception as e: + logger.error(f"engine-frame router: runtime abort failed for " + f"client_id={client_id}: {e!r}") + + def _fail_request(self, binding: RequestBinding, error_code: str, message: str, + deferred: List[Callable[[], None]]) -> None: + with self._lock: + if binding.ended: + return + self._end_with_error_locked(binding, error_code, message) + client_id = binding.client_id + if self._abort_fn is not None and client_id is not None: + deferred.append(lambda: self._safe_abort(client_id)) + + # ------------------------------------------------------------------ # + # Client-facing operations + # ------------------------------------------------------------------ # + + def abort(self, request_id: str) -> None: + """Request abort: idempotent on ended requests, typed error on unknown.""" + with self._lock: + binding = self._by_request.get(request_id) + if binding is None: + if request_id in self._tombstones: + return # no-op on a completed/ended request + raise UnknownRequestError(f"unknown request_id {request_id!r}") + if binding.ended or binding.abort_requested: + return + binding.abort_requested = True + client_id = binding.client_id + if self._abort_fn is not None and client_id is not None: + self._safe_abort(client_id) + + def get_binding(self, request_id: str) -> Optional[RequestBinding]: + with self._lock: + return self._by_request.get(request_id) + + def open_stream_binding(self, request_id: str) -> RequestBinding: + """Claim the single consumer stream for a request (typed errors).""" + with self._lock: + binding = self._delivery_index.get(request_id) + if binding is None: + raise UnknownRequestError(f"unknown request_id {request_id!r}") + if binding.stream_opened: + raise RouterError(f"stream for {request_id!r} was already opened") + binding.stream_opened = True + return binding + + def is_ended(self, request_id: str) -> bool: + with self._lock: + return request_id in self._tombstones + + @property + def fatal_error(self) -> Optional[str]: + return self._fatal_error + + def fail_all(self, reason: str) -> None: + """Latch a fatal condition and end every in-flight request typed. + + Wired from the proxy's shutdown/fatal-error path. Non-consuming: it + reads no queues; it only transitions router state. + """ + self.fail_all_and_collect(reason) + + def fail_all_and_collect(self, reason: str) -> List[int]: + """``fail_all`` that also snapshots the active contract client ids. + + The snapshot is taken before the bindings are pruned, so the caller + can issue best-effort runtime cancellation for every request that + was still in flight. Idempotent: already-ended bindings contribute + neither frames nor ids. + """ + with self._lock: + self._fatal_error = self._fatal_error or str(reason) + active_client_ids = [] + bindings = list(self._by_request.values()) + for binding in bindings: + if not binding.ended: + if binding.client_id is not None: + active_client_ids.append(binding.client_id) + self._end_with_error_locked(binding, "executor_failed", + str(reason)) + for binding in self._pending.values(): + if not binding.ended: + self._end_with_error_locked(binding, "executor_failed", + str(reason)) + self._pending.clear() + return active_client_ids + + def active_request_count(self) -> int: + with self._lock: + # Pending bindings are already indexed by request id; count each + # non-ended binding exactly once. + return len(self._by_request) diff --git a/tensorrt_llm/executor/engine_client/service.py b/tensorrt_llm/executor/engine_client/service.py new file mode 100644 index 000000000000..138c4fc4ae74 --- /dev/null +++ b/tensorrt_llm/executor/engine_client/service.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Server-side engine service: the contract execution endpoint. + +``EngineService`` owns everything a contract request needs on the engine +side and nothing the frontend owns: it validates/translates +``EngineRequest`` into the worker-facing request, submits through the +proxy's shared client-id allocator and request queue **without** a rank-0 +``GenerationResult``, exclusively routes contract responses off the +dispatch loop, and implements abort / client-detach / ``fail_all`` with +best-effort runtime cancellation. It holds no tokenizer, HTTP request, +formatter callback, or event loop. + +A later remote detach puts a transport terminator in front of this class +and changes nothing inside: every operation already takes/returns wire +types or primitives (see ``E5B_CUTOVER_DESIGN.md``). +""" + +import json +from typing import Optional, Tuple + +from tensorrt_llm.logger import logger + +from .contract import (EngineHealth, EngineRequest, IterationStatsBatch, + KvCacheEventsBatch) +from .conversion import engine_request_to_generation_request +from .router import EngineFrameRouter, RequestBinding, RouterError + +__all__ = ["EngineService", "EngineServiceError"] + + +class EngineServiceError(RouterError): + """Contract submission failed at the service (engine-side) boundary.""" + + +class EngineService: + """Contract-native execution endpoint over the IPC proxy runtime. + + The proxy supplies the narrow runtime surface this service depends on: + ``_start_dispatch_threads``, ``_get_next_client_id``, ``request_queue``, + ``abort_request``, ``_handle_background_error``, ``get_stats`` / + ``get_kv_events`` / ``check_health``, ``doing_shutdown``, and the + submission/lifecycle lock (``_submission_lock``). It never touches ZMQ + socket construction or MPI process management. + """ + + def __init__(self, proxy): + self._proxy = proxy + self.router = EngineFrameRouter(abort_fn=proxy.abort_request) + self._closed = False + # The shared proxy allocator is monotonic within a proxy lifetime; + # wrap or collision must fail submission rather than reuse an id. + self._max_client_id: int = 0 + + # ------------------------------------------------------------------ # + # Submission + # ------------------------------------------------------------------ # + + def submit_contract(self, engine_request: EngineRequest, + stop_reasons: Tuple = ()) -> int: + """Submit a contract request; returns the allocated proxy client id. + + No rank-0 ``GenerationResult`` is constructed or registered: the + router binding IS the rank-0 request state. Order (per the E5b + design note): translate first (a rejection leaves no engine-side + state), then — under the proxy submission/lifecycle lock — allocate + the shared client id, bind, and enqueue, so shutdown can never + interleave between binding and enqueue and a response can never + beat the binding. + """ + # 1. Translate before any allocation: conversion failures are + # state-free. + generation_request = engine_request_to_generation_request(engine_request) + proxy = self._proxy + proxy._start_dispatch_threads() + with proxy._submission_lock: + if proxy.doing_shutdown or getattr(proxy, "_fatal_error", None) \ + or self.router.fatal_error is not None or self._closed: + raise EngineServiceError( + "engine is shutting down or failed; submission rejected") + client_id = proxy._get_next_client_id() + if client_id in proxy._results: + raise EngineServiceError( + f"client id {client_id} collides with an active legacy " + "request; refusing to reuse an id in this proxy lifetime") + self.observe_legacy_allocation(client_id) + generation_request.set_id(client_id) + binding = self.router.bind(engine_request.request_id, client_id, + engine_request.prompt_token_ids, + stop_reasons) + try: + proxy.request_queue.put(generation_request) + except Exception: + self.router.on_submit_enqueue_failed(client_id) + raise + proxy._handle_background_error() + return client_id + + def observe_legacy_allocation(self, client_id: int) -> None: + """Monotonic-allocation guard shared by BOTH populations. + + Called (submission lock held) for every id the shared proxy + allocator hands out — contract submissions above and legacy + submissions via ``GenerationExecutorProxy.submit``. The raw + allocator wraps modulo 2**64; while contract routing is exclusive, + a wrapped or regressed id must fail the submission it was handed + to, whichever population received it, before it can collide with a + contract-owned or recently retired id. + """ + if client_id <= self._max_client_id: + raise EngineServiceError( + f"client id allocator wrapped or regressed ({client_id} <= " + f"{self._max_client_id}); refusing to reuse an id") + self._max_client_id = client_id + + # ------------------------------------------------------------------ # + # Response routing (dispatch thread) + # ------------------------------------------------------------------ # + + def route_response(self, raw) -> bool: + """Exclusively claim a contract-owned response; see the router.""" + return self.router.route_response(raw) + + # ------------------------------------------------------------------ # + # Abort / detach / failure + # ------------------------------------------------------------------ # + + def abort(self, request_id: str) -> None: + self.router.abort(request_id) + + def fail_all(self, reason: str) -> None: + """Poison every contract stream typed AND cancel runtime work. + + The active client ids are snapshotted before the bindings are + pruned; runtime cancellation is best effort (the worker may already + be gone), poisoning is guaranteed. Idempotent. + """ + # Latch under the proxy submission lock so a concurrent + # submit_contract cannot pass its closed/fatal check and bind after + # the poisoning sweep (the lock is reentrant for shutdown callers). + with self._proxy._submission_lock: + active_client_ids = self.router.fail_all_and_collect(reason) + for client_id in active_client_ids: + try: + self._proxy.abort_request(client_id) + except Exception as e: + logger.debug(f"engine-service: best-effort runtime cancel of " + f"client_id={client_id} failed: {e!r}") + + def close_client(self, reason: str = "client closed") -> None: + """Detach this client: poison its streams and abort its engine work. + + Does NOT shut the shared engine down. + """ + with self._proxy._submission_lock: + if self._closed: + return + self._closed = True + self.fail_all(reason) + + # ------------------------------------------------------------------ # + # Typed control plane + # ------------------------------------------------------------------ # + + def get_stats(self, timeout: float = 2.0) -> IterationStatsBatch: + entries = self._proxy.get_stats(timeout=timeout) + return IterationStatsBatch(entries=tuple( + entry if isinstance(entry, str) else json.dumps(entry) + for entry in entries)) + + def get_kv_events(self, timeout: float = 2.0) -> KvCacheEventsBatch: + entries = self._proxy.get_kv_events(timeout=timeout) + return KvCacheEventsBatch(entries=tuple( + entry if isinstance(entry, str) else json.dumps(entry) + for entry in entries)) + + def health(self) -> EngineHealth: + if self.router.fatal_error is not None: + return EngineHealth(healthy=False, detail=self.router.fatal_error) + try: + healthy = bool(self._proxy.check_health()) + except Exception as e: + return EngineHealth(healthy=False, detail=repr(e)) + return EngineHealth(healthy=healthy, + detail="" if healthy else "executor unhealthy") + + def get_binding(self, request_id: str) -> Optional[RequestBinding]: + return self.router.get_binding(request_id) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 5599cc8df2a8..566d6fbcbb6d 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -14,9 +14,11 @@ # limitations under the License. import atexit import concurrent.futures +import contextlib import json import os import threading +import traceback import weakref from queue import Empty from typing import Dict, List, Optional, Union @@ -296,6 +298,23 @@ def _mark_engine_dead(self, error: Optional[BaseException] = None) -> None: result.queue.put(dead_error) except Exception: # noqa: BLE001 - a full/closed queue must not stop the sweep pass + # The contract population has no rank-0 GenerationResult, so the sweep + # above cannot reach it: poison those streams typed with the same + # immediacy (idempotent; pre_shutdown repeats it later). Without this, + # legacy requests fail fast on worker death while contract streams + # wait for the monitor poll tick. + service = self._engine_service + if service is not None: + try: + service.fail_all("engine dead") + except Exception as e: # noqa: BLE001 - the sweep must complete + logger.error(f"engine service fail_all on engine death: {e!r}") + frame_router = self._engine_frame_router + if frame_router is not None: + try: + frame_router.fail_all("engine dead") + except Exception as e: # noqa: BLE001 - the sweep must complete + logger.error(f"frame router fail_all on engine death: {e!r}") def _handle_worker_death(self, error: BaseException) -> None: """Event-driven worker-death handler. @@ -411,6 +430,50 @@ def _setup_queues(self) -> WorkerCommIpcAddrs: def resource_governor_queue(self): return self._resource_governor_queue + # --- experimental engine-client hooks (inert unless a router is attached) --- + _engine_frame_router = None + _engine_service = None + + def attach_engine_frame_router(self, router) -> None: + """Attach the experimental engine-frame router (TLLM_EXPERIMENTAL_ENGINE_CLIENT). + + The router observes submissions and raw responses to produce the + typed engine-contract frame stream (shadow mode: the legacy path + stays authoritative). With no router attached every hook in this + class is inert and legacy behavior is unchanged. + """ + self._engine_frame_router = router + + @property + def _submission_lock(self): + """Submission/lifecycle lock serializing id allocation + enqueue with + shutdown, shared by the legacy and contract submit paths. Created on + first use; reentrant because a submit-side background-error check can + reach ``pre_shutdown`` while the lock is held.""" + # dict.setdefault is atomic under the GIL, so concurrent first access + # cannot observe two different locks. + return self.__dict__.setdefault("_submission_lock_instance", + threading.RLock()) + + def attach_engine_service(self, service) -> None: + """Attach the engine service (authoritative contract routing). + + Unlike the shadow router, an attached service claims contract + responses exclusively: they never reach legacy delivery. Only one + service may attach per proxy lifetime. + """ + if self._engine_service is not None: + raise RuntimeError( + "an engine service is already attached to this executor") + self._engine_service = service + + def submit_contract(self, engine_request, stop_reasons=()) -> int: + """Contract-native submission (no rank-0 GenerationResult).""" + if self._engine_service is None: + raise RuntimeError("no engine service attached") + return self._engine_service.submit_contract(engine_request, + stop_reasons=stop_reasons) + def abort_request(self, request_id: int) -> None: """Abort a request by sending a cancelling request to the request queue. @@ -462,10 +525,33 @@ def process_res(res): res = res if isinstance(res, list) else [res] + # Contract routing runs BEFORE any GenerationResult lookup: an + # attached engine service exclusively claims contract-owned ids (they + # never reach legacy delivery or the pop-on-final logic). The shadow + # router, when attached instead, only observes — legacy delivery + # still runs. Failures in either must never break legacy dispatch. + service = self._engine_service + router = self._engine_frame_router + for i in res: global_tracer().log_instant("IPC.get") if i is None: return False + if service is not None: + try: + if service.route_response(i): + continue + except Exception as e: + logger.error( + f"engine service raised while routing a response; " + f"continuing legacy dispatch: {e!r}\n{traceback.format_exc()}") + if router is not None: + try: + router.on_response(i) + except Exception as e: + logger.error( + f"engine-frame router raised while observing a response; " + f"continuing legacy dispatch: {e!r}\n{traceback.format_exc()}") process_res(i) if async_queues: @@ -566,15 +652,38 @@ def pre_shutdown(self): return logger_debug('Proxy.pre_shutdown...\n', "yellow") - if self.doing_shutdown: - return - else: + # Serialize the shutdown transition with in-flight submissions when + # the engine service is attached (reentrant: a submit-side + # background-error check can reach here while holding the lock). + lock_ctx = (self._submission_lock if self._engine_service is not None + else contextlib.nullcontext()) + with lock_ctx: + if self.doing_shutdown: + return self.doing_shutdown = True # Wake the error monitor thread immediately so it exits cleanly if hasattr(self, '_shutdown_event'): self._shutdown_event.set() + service = self._engine_service + if service is not None: + try: + # Poisons every contract stream typed and issues best-effort + # runtime cancellation — the contract-population counterpart + # of _abort_all_requests below. + service.fail_all("executor shutdown") + except Exception as e: + logger.error( + f"engine service raised in fail_all during shutdown: {e!r}") + router = self._engine_frame_router + if router is not None: + try: + router.fail_all("executor shutdown") + except Exception as e: + logger.error( + f"engine-frame router raised in fail_all during shutdown: {e!r}") + self._abort_all_requests() # Notify surviving workers to quit. Send the sentinel when: @@ -661,28 +770,54 @@ def submit(self, request: GenerationRequest) -> GenerationResult: self._start_dispatch_threads() - request.set_id(self._get_next_client_id()) - logprob_params = self._get_logprob_params(request) - - result = GenerationResult( - request, - background_error_handler=self._handle_background_error, - executor=self, - disaggregated_params=request.disaggregated_params, - logprob_params=logprob_params) - self._results[request.id] = result - - # Close the submit-vs-death race: _mark_engine_dead() runs on the - # error-monitor thread and its one-shot sweep snapshots _results, so a - # result registered just after that snapshot would never be unblocked - # and would hang forever on a dead worker. Re-check after registering - # and fail fast if the engine died in that window. - if self._engine_dead or self._fatal_error is not None: - self._results.pop(request.id, None) - raise EngineDeadError(self._fatal_error) - - with nvtx_range_debug("request_queue.put"): - self.request_queue.put(request) + # With an engine service attached, the whole allocate+register+enqueue + # sequence is serialized with contract submission and shutdown (shared + # allocator; the worker sentinel must never interleave between + # registration and enqueue). Without one, legacy behavior and cost are + # unchanged. + lock_ctx = (self._submission_lock if self._engine_service is not None + else contextlib.nullcontext()) + with lock_ctx: + request.set_id(self._get_next_client_id()) + if self._engine_service is not None: + # The shared allocator must stay monotonic across BOTH + # populations while contract routing is exclusive: a wrapped + # or regressed id fails this submission before it can collide + # with a contract-owned or recently retired id. + self._engine_service.observe_legacy_allocation(request.id) + router = self._engine_frame_router + if router is not None: + # Bind on the submit thread, after client-id assignment and + # before the enqueue below (race-free by construction). + router.observe_submit(request) + logprob_params = self._get_logprob_params(request) + + result = GenerationResult( + request, + background_error_handler=self._handle_background_error, + executor=self, + disaggregated_params=request.disaggregated_params, + logprob_params=logprob_params) + self._results[request.id] = result + + # Close the submit-vs-death race: _mark_engine_dead() runs on the + # error-monitor thread and its one-shot sweep snapshots _results, + # so a result registered just after that snapshot would never be + # unblocked and would hang forever on a dead worker. Re-check after + # registering and fail fast if the engine died in that window. + if self._engine_dead or self._fatal_error is not None: + self._results.pop(request.id, None) + if router is not None: + router.on_submit_enqueue_failed(request.id) + raise EngineDeadError(self._fatal_error) + + with nvtx_range_debug("request_queue.put"): + try: + self.request_queue.put(request) + except Exception: + if router is not None: + router.on_submit_enqueue_failed(request.id) + raise self._handle_background_error() diff --git a/tensorrt_llm/executor/request.py b/tensorrt_llm/executor/request.py index a003b387f45b..ec099316f207 100644 --- a/tensorrt_llm/executor/request.py +++ b/tensorrt_llm/executor/request.py @@ -114,6 +114,7 @@ def __init__( encoder_input_token_ids: Optional[Union[torch.Tensor, np.ndarray, list]] = None, priority: float = DEFAULT_REQUEST_PRIORITY, + stop_token_sequences: Optional[List[List[int]]] = None, ): if isinstance(prompt_token_ids, list): self.prompt_token_ids = prompt_token_ids @@ -163,6 +164,19 @@ def __init__( raise ValueError( f"priority must be a float in [0.0, 1.0], got {priority}") self.priority = priority + # Pre-tokenized stop sequences forwarded verbatim as runtime stop + # words, in addition to whatever sampling_params derives. Callers that + # tokenize stop strings themselves use this instead of writing + # sampling_params' private tokenizer-derived state. + if stop_token_sequences is not None: + for seq in stop_token_sequences: + if not isinstance(seq, list) or not seq or not all( + isinstance(t, int) and not isinstance(t, bool) + for t in seq): + raise ValueError( + "stop_token_sequences must be a list of non-empty " + f"lists of ints, got {stop_token_sequences!r}") + self.stop_token_sequences = stop_token_sequences @staticmethod def _normalize_optional_token_ids(token_ids: Optional[Union[torch.Tensor, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index faaeea64bc1c..383a1a51a83c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5043,6 +5043,15 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: description="If true, force dynamic quantization. Defaults to False.", status="prototype") + experimental_engine_client: bool = Field( + default=False, + description= + "Enable the experimental engine-client serving path (typed " + "iteration-level engine contract) for eligible streaming requests. " + "The TLLM_EXPERIMENTAL_ENGINE_CLIENT environment variable overrides " + "this setting in both directions ('1' forces on, '0' forces off).", + status="prototype") + allreduce_strategy: Optional[Literal[ 'AUTO', 'NCCL', 'UB', 'MINLATENCY', 'ONESHOT', 'TWOSHOT', 'LOWPRECISION', 'MNNVL', diff --git a/tensorrt_llm/serve/engine_client_serving.py b/tensorrt_llm/serve/engine_client_serving.py new file mode 100644 index 000000000000..d6fa38f8d581 --- /dev/null +++ b/tensorrt_llm/serve/engine_client_serving.py @@ -0,0 +1,530 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Experimental engine-client serving path (TLLM_EXPERIMENTAL_ENGINE_CLIENT). + +With the flag on, ELIGIBLE OpenAI streaming requests route through +``conversion`` → ``LocalProcessEngineClient`` → ``FrontendResponseAssembler`` +and the existing OpenAI stream formatters, never calling +``LLM.generate_async``, never constructing ``PostprocParams``, and never +consuming a ``GenerationResult``. Ineligible requests transparently fall +back to the legacy path, counted per eligibility axis. + +Model/tokenizer information on this path is read only through a +locally-built, data-only ``FrontendModelContext``: the assembler and the +formatters use a tokenizer **reloaded from the context's TokenizerSpec**, +not the server's live tokenizer object (the reload-parity test guards +token-for-token equivalence). The context is built locally in V0; a later +remote detach swaps only its delivery. +""" + +import asyncio +import hashlib +import os +import uuid +from typing import List, Optional + +from tensorrt_llm.executor.engine_client.assembler import FrontendResponseAssembler +from tensorrt_llm.executor.engine_client.contract import (EngineCapabilities, + FrontendModelContext, + TokenizerSpec) +from tensorrt_llm.executor.engine_client.conversion import ( + ConversionError, RequestIneligibleError, convert_request, + prepare_sampling_params_from_context) +from tensorrt_llm.executor.engine_client.local_client import ( + ENGINE_CLIENT_FLAG_ENV, EngineClientConfig, EngineClientConfigError, + LocalProcessEngineClient, RequestRejectedError, + resolve_engine_client_flag) +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.logger import logger + +__all__ = ["ContractStreamView", "EngineClientServing", "build_model_context"] + +_TOKENIZER_FILES = ("tokenizer.json", "tokenizer_config.json", "vocab.json", + "merges.txt", "special_tokens_map.json", + "generation_config.json") + + +class _ViewOutput: + """Duck-typed CompletionOutput view consumed by the stream formatters.""" + + def __init__(self): + self.index = 0 + self.token_ids_diff: List[int] = [] + self.text_diff: str = "" + self.logprobs_diff: List[float] = [] + self.finish_reason: Optional[str] = None + self.stop_reason = None + self.length: int = 0 + self.disaggregated_params = None + self.prompt_logprobs = None + self.token_ids: List[int] = [] + self.text: str = "" + self.logprobs: List[float] = [] + self.request_perf_metrics = None + # Character length of the accumulated text BEFORE the current batch + # (the completions logprob formatter uses it as the text offset). + self._last_text_len: int = 0 + + +class ContractStreamView: + """Duck-typed ``GenerationResultBase`` view fed from assembler updates. + + NOT a ``GenerationResult``: no queue, no executor, no runtime state — + only the read surface the existing OpenAI stream formatters consume. + """ + + def __init__(self, view_id: int, prompt_token_ids): + self.id = view_id + self.prompt_token_ids = list(prompt_token_ids) + self._outputs = [_ViewOutput()] + self._done = False + self.cached_tokens = 0 + self.avg_decoded_tokens_per_iter = None + self.metrics_dict = {} + self.time_breakdown_metrics = None + self._disaggregated_params = None + + @property + def outputs(self): + return self._outputs + + @property + def disaggregated_params(self): + return self._disaggregated_params + + def apply_updates(self, assembler: FrontendResponseAssembler, updates) -> bool: + """Fold one batch of assembly updates into the view. + + Returns True when the request reached its end (complete or error). + """ + output = self._outputs[0] + output.token_ids_diff = [] + output.text_diff = "" + output.logprobs_diff = [] + output._last_text_len = len(output.text) + ended = False + for update in updates: + if update.kind == "delta": + output.token_ids_diff.extend(update.new_token_ids) + output.text_diff += update.text_diff + if update.logprobs: + output.logprobs_diff.extend(update.logprobs) + if update.finish_reason is not None: + output.finish_reason = update.finish_reason + output.stop_reason = update.stop_reason + if update.metrics_cached_tokens is not None: + self.cached_tokens = update.metrics_cached_tokens + elif update.kind == "finish": + output.finish_reason = update.finish_reason + output.stop_reason = update.stop_reason + elif update.kind == "complete": + self._done = True + ended = True + if update.usage and "cached_tokens" in update.usage: + self.cached_tokens = update.usage["cached_tokens"] + elif update.kind == "error": + self._done = True + ended = True + output.finish_reason = output.finish_reason or "error" + raise _ContractStreamError(update.error_code or "error", + update.error_message or "") + output.token_ids = assembler.token_ids + output.text = assembler.text + output.logprobs = assembler.logprobs + output.length = len(assembler.token_ids) + return ended + + +class _ContractStreamError(RuntimeError): + + def __init__(self, code: str, message: str): + super().__init__(f"[{code}] {message}") + self.code = code + + +def _file_manifest(model_dir: str): + manifest = [] + for name in _TOKENIZER_FILES: + path = os.path.join(model_dir, name) + if os.path.isfile(path): + digest = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + digest.update(chunk) + manifest.append((name, digest.hexdigest())) + return tuple(manifest) + + +def _dump_json_or_none(value) -> Optional[str]: + if value is None: + return None + try: + import json + return json.dumps(value, sort_keys=True, default=str) + except (TypeError, ValueError): + return None + + +def build_model_context(llm, capabilities: EngineCapabilities) -> FrontendModelContext: + """Build the data-only model context locally from the in-process engine. + + This is the ONLY point where the contract path touches live model + objects; everything downstream (normalization, detokenization, + formatting) consumes the frozen context plus the spec-reloaded + tokenizer. The remote detach later swaps the delivery (ready-handshake) + without changing any consumer. + """ + import json + model_dir = str(getattr(llm, "_hf_model_dir", None) or "") + tokenizer = llm.tokenizer + chat_template = None + inner = getattr(tokenizer, "tokenizer", None) + added_tokens_json = special_tokens_json = None + normalizer_json = pre_tokenizer_json = None + if inner is not None: + chat_template = getattr(inner, "chat_template", None) + try: + added_tokens_json = _dump_json_or_none(inner.get_added_vocab()) + except Exception: + pass + special_tokens_json = _dump_json_or_none( + getattr(inner, "special_tokens_map", None)) + backend = getattr(inner, "backend_tokenizer", None) + if backend is not None: + try: + backend_config = json.loads(backend.to_str()) + normalizer_json = _dump_json_or_none( + backend_config.get("normalizer")) + pre_tokenizer_json = _dump_json_or_none( + backend_config.get("pre_tokenizer")) + except Exception: + pass + hf_config = getattr(llm, "_hf_model_config", None) + revision = getattr(hf_config, "_commit_hash", None) + spec = TokenizerSpec( + uri=model_dir or "unknown", + files_manifest=_file_manifest(model_dir) if model_dir else (), + revision=revision if isinstance(revision, str) and revision else None, + fast=bool(getattr(tokenizer, "is_fast", True)), + trust_remote_code=bool(getattr(llm.args, "trust_remote_code", False)), + added_tokens_json=added_tokens_json, + special_tokens_json=special_tokens_json, + normalizer_json=normalizer_json, + pre_tokenizer_json=pre_tokenizer_json, + ) + generation_config = getattr(llm, "_generation_config", None) + generation_stop_ids = getattr(generation_config, "eos_token_id", None) + if isinstance(generation_stop_ids, int): + generation_stop_ids = [generation_stop_ids] + chat_template_version = None + if chat_template: + chat_template_version = hashlib.sha256( + chat_template.encode()).hexdigest()[:12] + return FrontendModelContext( + tokenizer=spec, + capabilities=capabilities, + chat_template_source=chat_template, + chat_template_version=chat_template_version, + eos_id=getattr(tokenizer, "eos_token_id", None), + pad_id=getattr(tokenizer, "pad_token_id", None), + max_context_length=getattr(llm.args, "max_seq_len", None), + generation_stop_token_ids=tuple( + int(t) for t in (generation_stop_ids or ())), + model_type=getattr(hf_config, "model_type", None) or None, + ) + + +def load_tokenizer_from_spec(spec: TokenizerSpec): + """Reload the frontend tokenizer purely from the spec (no live handle). + + Provenance is ENFORCED, not advisory: every manifest entry must exist + on disk with a matching content hash before the reload (the manifest is + the content pin — it subsumes a revision string), and the reloaded + tokenizer's added/special-token configuration must match the recorded + source. Every failure is a typed config error (fail closed). + """ + import hashlib + from tensorrt_llm.tokenizer.tokenizer import TransformersTokenizer + if spec.trust_remote_code: + raise EngineClientConfigError( + "trust_remote_code_tokenizer: cannot reload a trust_remote_code " + "tokenizer from a data-only spec") + for rel_path, expected_hash in spec.files_manifest: + candidate = os.path.join(spec.uri, rel_path) + try: + with open(candidate, "rb") as f: + digest = hashlib.sha256(f.read()).hexdigest() + except OSError as e: + raise EngineClientConfigError( + f"tokenizer_manifest: {rel_path!r} unreadable during spec " + f"reload: {e}") + if digest != expected_hash: + raise EngineClientConfigError( + f"tokenizer_manifest: {rel_path!r} content does not match " + "the recorded manifest hash; refusing a divergent reload") + tokenizer = TransformersTokenizer.from_pretrained(spec.uri, + use_fast=spec.fast) + inner = getattr(tokenizer, "tokenizer", None) + if inner is not None: + if spec.added_tokens_json is not None: + try: + reloaded = _dump_json_or_none(inner.get_added_vocab()) + except Exception: + reloaded = None + if reloaded != spec.added_tokens_json: + raise EngineClientConfigError( + "tokenizer_manifest: reloaded added-token vocabulary " + "does not match the recorded spec") + if spec.special_tokens_json is not None: + reloaded = _dump_json_or_none( + getattr(inner, "special_tokens_map", None)) + if reloaded != spec.special_tokens_json: + raise EngineClientConfigError( + "tokenizer_manifest: reloaded special-token map does " + "not match the recorded spec") + return tokenizer + + +async def wrap_sse_with_done(generator, raw_request=None): + """Wrap a contract SSE generator with first-token timing and ``[DONE]``. + + Mirrors the legacy stream wrappers: the first piece stamps + ``server_first_token_time`` on the request state, and the stream always + ends with the ``[DONE]`` sentinel (even when empty). + """ + try: + first = await generator.__anext__() + except StopAsyncIteration: + yield "data: [DONE]\n\n" + return + if raw_request is not None: + from tensorrt_llm.serve.responses_utils import \ + get_steady_clock_now_in_seconds + raw_request.state.server_first_token_time = \ + get_steady_clock_now_in_seconds() + yield first + async for piece in generator: + yield piece + yield "data: [DONE]\n\n" + + +class EngineClientServing: + """Serving-side owner of the experimental engine-client path.""" + + def __init__(self, llm): + config = EngineClientConfig( + backend=llm.args.backend or "pytorch", + num_postprocess_workers=getattr(llm.args, "num_postprocess_workers", + 0) or 0, + post_processor_hook_set=getattr(llm.args, "post_processor_hook", None) + is not None, + speculative_config_set=getattr(llm.args, "speculative_config", None) + is not None, + early_first_token_mode=bool( + getattr(llm.args, "enable_early_first_token_response", False)), + world_size=getattr(getattr(llm.args, "parallel_config", None), + "world_size", 1) or 1, + tokenizer_trust_remote_code=bool( + getattr(llm.args, "trust_remote_code", False)), + flag_enabled=bool( + getattr(llm.args, "experimental_engine_client", False)), + ) + if not isinstance(llm._executor, GenerationExecutorProxy): + raise EngineClientConfigError( + "transport: the executor is not the IPC GenerationExecutorProxy") + self.client = LocalProcessEngineClient(llm._executor, config) + self.model_context = build_model_context(llm, self.client.capabilities()) + # The contract path reads model/tokenizer info only through the + # context: the tokenizer below is reloaded from the spec, not the + # server's live handle. + self.tokenizer = load_tokenizer_from_spec(self.model_context.tokenizer) + self.counters = {"contract_requests": 0, "capability_rejections": 0} + self._stream_interval = getattr(llm.args, "stream_interval", 1) or 1 + # Plain config scalar (not a model reach), captured once. + self._force_return_perf_metrics = bool( + getattr(llm.args, "return_perf_metrics", False)) + + @classmethod + def create_if_enabled(cls, llm) -> Optional["EngineClientServing"]: + """Build the serving path when the resolved flag is on; None when off. + + Raises typed errors (fail closed) for malformed flag values and for + config-gate violations while enablement was requested. + """ + # Fail CLOSED: a malformed flag value or a config-gate violation + # while the resolved flag requests enablement is a typed startup + # error, never a silent fall back to legacy. Only a flag that + # resolves OFF returns None. + enabled = resolve_engine_client_flag( + bool(getattr(llm.args, "experimental_engine_client", False))) + if not enabled: + return None + serving = cls(llm) + logger.info( + f"{ENGINE_CLIENT_FLAG_ENV}: experimental engine-client serving " + "path enabled for eligible streaming requests") + return serving + + # ------------------------------------------------------------------ # + + def _count_fallback(self, axis: str) -> None: + key = f"fallback:{axis}" + self.counters[key] = self.counters.get(key, 0) + 1 + + def get_counters(self) -> dict: + merged = dict(self.counters) + merged.update(self.client._router.counters) + merged["active_requests"] = self.client._router.active_request_count() + return merged + + @staticmethod + def _normalize_scheduling(scheduling_params): + if scheduling_params is None: + return None + if getattr(scheduling_params, "agent_hierarchy", None) is None: + return None # an all-default SchedulingParams carries nothing + return scheduling_params + + def try_stream(self, *, preprocessed, sampling_params, post_processor, + postproc_args, raw_request=None, lora_request=None, + prompt_adapter_request=None, disaggregated_params=None, + conversation_params=None, scheduling_params=None, + cache_salt=None, trace_headers=None, + kv_cache_retention_config=None, priority=None, + echo=False, streaming=True): + """Attempt the contract path; returns an async SSE generator or None. + + None means "fall back to legacy" (the eligibility axis was counted). + Never raises for eligibility reasons. + + ``sampling_params`` MUST be pristine (never live-prepared): the + endpoints capture a copy before any live ``LLM.preprocess`` call so + that EVERY contract request is normalized here, through the + context-only boundary, and never by live model-config reaches. + """ + if getattr(preprocessed, "prompt_token_ids", None) is None: + self._count_fallback("no_preprocessed_inputs") + return None + if trace_headers is not None and not trace_headers: + # A tracing-enabled server extracts {} when no trace header is + # present; only requests actually CARRYING trace headers are + # ineligible. + trace_headers = None + try: + # Context-only normalization — unconditional, so a request that + # arrived live-prepared can never smuggle live-derived sampling + # state past this boundary. + prepare_sampling_params_from_context( + sampling_params, + context=self.model_context, + tokenizer=self.tokenizer, + stream_interval=self._stream_interval, + force_return_perf_metrics=self._force_return_perf_metrics) + except RequestIneligibleError as e: + self._count_fallback(e.axis) + return None + except ConversionError as e: + self._count_fallback("conversion_error") + logger.warning( + f"engine-client normalization error; falling back: {e}") + return None + try: + engine_request, output_config = convert_request( + f"chatcmpl-ec-{uuid.uuid4().hex}", + preprocessed.prompt_token_ids, + sampling_params, + streaming=streaming, + multimodal_params=getattr(preprocessed, "multimodal_params", None), + lora_request=lora_request, + prompt_adapter_request=prompt_adapter_request, + disaggregated_params=disaggregated_params, + scheduling_params=self._normalize_scheduling(scheduling_params), + conversation_params=conversation_params, + trace_headers=trace_headers, + cache_salt=cache_salt, + query_token_ids=getattr(preprocessed, "query_token_ids", None), + encoder_input_token_ids=getattr(preprocessed, + "encoder_input_token_ids", None), + kv_cache_retention_config=kv_cache_retention_config, + priority=priority, + echo=echo, + ) + except RequestIneligibleError as e: + self._count_fallback(e.axis) + return None + except ConversionError as e: + self._count_fallback("conversion_error") + logger.warning(f"engine-client conversion error; falling back: {e}") + return None + + try: + request_id = self.client.submit(engine_request, + output_config=output_config) + except RequestRejectedError as e: + self.counters["capability_rejections"] += 1 + logger.warning(f"engine-client pre-submit rejection; falling back: {e}") + return None + + self.counters["contract_requests"] += 1 + assembler = FrontendResponseAssembler(request_id, output_config, + tokenizer=self.tokenizer, + abort_callback=self.client.abort, + stream_interval=self._stream_interval) + view = ContractStreamView(abs(hash(request_id)) % (1 << 31), + engine_request.prompt_token_ids) + postproc_args.tokenizer = self.tokenizer + postproc_args.num_prompt_tokens = len(engine_request.prompt_token_ids) + return self._sse_generator(request_id, assembler, view, post_processor, + postproc_args, raw_request) + + async def _watch_disconnect(self, raw_request, request_id: str) -> None: + try: + while True: + await asyncio.sleep(1.0) + if await raw_request.is_disconnected(): + # Client went away: end background generation via the + # same abort control edge as an explicit cancel. + try: + self.client.abort(request_id) + except Exception: + pass + return + except asyncio.CancelledError: + pass + + async def _sse_generator(self, request_id, assembler, view, post_processor, + postproc_args, raw_request): + stream = self.client.stream(request_id) + watcher = None + if raw_request is not None: + watcher = asyncio.create_task( + self._watch_disconnect(raw_request, request_id)) + try: + ended = False + while not ended: + frame = await stream.__anext__() + batch = [frame] + stream.pop_ready() + updates = assembler.process_frames(batch) + try: + ended = view.apply_updates(assembler, updates) + except _ContractStreamError as e: + logger.error(f"engine-client stream {request_id} failed: {e}") + raise + for piece in post_processor(view, postproc_args): + yield piece + finally: + if watcher is not None: + watcher.cancel() + await stream.aclose() diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index fd5fb55c25c0..a051b8227639 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2,6 +2,7 @@ import array import asyncio import base64 +import copy import functools import json import os @@ -17,6 +18,7 @@ from datetime import datetime from http import HTTPStatus from pathlib import Path +from types import SimpleNamespace from typing import (Annotated, Any, AsyncGenerator, AsyncIterator, List, Optional, Union) @@ -271,6 +273,20 @@ def __init__( self._embedding_max_queue_delay = embedding_max_queue_delay self._embedding_max_queue_size = embedding_max_queue_size self.embedding_batcher: Optional[EncodeBatcher] = None + # Experimental engine-client serving path; None (fully legacy) + # unless the resolved flag (LlmArgs experimental_engine_client, + # overridden in both directions by TLLM_EXPERIMENTAL_ENGINE_CLIENT) + # is on and the config passes the setup gates. The import is lazy so + # the flag-off import graph is unchanged. + self._engine_client_serving = None + if not self._is_visual_gen and ( + os.environ.get("TLLM_EXPERIMENTAL_ENGINE_CLIENT") is not None + or getattr(getattr(generator, "args", None), + "experimental_engine_client", False)): + from tensorrt_llm.serve.engine_client_serving import \ + EngineClientServing + self._engine_client_serving = EngineClientServing.create_if_enabled( + generator) self.tool_parser = tool_parser self.metadata_server = create_metadata_server(metadata_server_cfg) self.disagg_cluster_config = disagg_cluster_config @@ -1584,12 +1600,6 @@ async def chat_stream_generator( if conversation and conversation[-1].get( "content") and conversation[-1].get("role") == get_role(): postproc_args.last_message_content = conversation[-1]["content"] - postproc_params = PostprocParams( - post_processor=chat_stream_post_processor - if request.stream else chat_response_post_processor, - postproc_args=postproc_args, - ) - trace_headers = (None if raw_request is None else tracing.extract_trace_headers(raw_request.headers)) @@ -1601,6 +1611,13 @@ async def chat_stream_generator( agent_hierarchy=request.agent_hierarchy) generate_inputs = prompt + # The contract path normalizes sampling params EXCLUSIVELY at the + # context-only boundary: capture a pristine copy before the live + # preprocess below can prepare the original in place. + contract_sampling_params = ( + copy.deepcopy(sampling_params) + if request.stream and self._engine_client_serving is not None + else None) preprocess_fn = getattr(self.generator, "preprocess", None) if preprocess_fn is not None: loop = asyncio.get_event_loop() @@ -1609,6 +1626,35 @@ async def chat_stream_generator( functools.partial(preprocess_fn, prompt, sampling_params, disaggregated_params)) + if request.stream and self._engine_client_serving is not None: + contract_generator = self._engine_client_serving.try_stream( + preprocessed=generate_inputs, + sampling_params=contract_sampling_params, + post_processor=chat_stream_post_processor, + postproc_args=postproc_args, + raw_request=raw_request, + lora_request=request.lora_request, + disaggregated_params=disaggregated_params, + conversation_params=conversation_params, + scheduling_params=scheduling_params, + cache_salt=request.cache_salt, + trace_headers=trace_headers, + echo=bool(getattr(request, "echo", False)), + ) + if contract_generator is not None: + from tensorrt_llm.serve.engine_client_serving import \ + wrap_sse_with_done + return StreamingResponse( + content=wrap_sse_with_done(contract_generator, + raw_request), + media_type="text/event-stream") + + postproc_params = PostprocParams( + post_processor=chat_stream_post_processor + if request.stream else chat_response_post_processor, + postproc_args=postproc_args, + ) + promise = self.generator.generate_async( inputs=generate_inputs, sampling_params=sampling_params, @@ -1915,16 +1961,17 @@ async def generator_wrapper(generator: AsyncIterator[Any]): postproc_args.stream_created = stream_created if request.echo: postproc_args.prompt = prompt - postproc_params = PostprocParams( - post_processor=completion_stream_post_processor - if request.stream else completion_response_post_processor, - postproc_args=postproc_args, - ) trace_headers = (None if raw_request is None else tracing.extract_trace_headers( raw_request.headers)) prompt = prompt_inputs(prompt) + # Pristine copy for the contract path, captured before the + # live input processor can prepare the shared original. + contract_sampling_params = ( + copy.deepcopy(sampling_params) + if request.stream and len(prompts) == 1 + and self._engine_client_serving is not None else None) if prompt.get("prompt") is not None: loop = asyncio.get_event_loop() prompt_token_ids, extra_processed_inputs = await loop.run_in_executor( @@ -1939,6 +1986,40 @@ async def generator_wrapper(generator: AsyncIterator[Any]): else: tokens_prompt = prompt + if (request.stream and len(prompts) == 1 + and self._engine_client_serving is not None): + tokens_view = tokens_prompt if isinstance( + tokens_prompt, dict) else {} + preprocessed_shim = SimpleNamespace( + prompt_token_ids=tokens_view.get("prompt_token_ids"), + query_token_ids=tokens_view.get("query_token_ids"), + multimodal_params=tokens_view.get("multi_modal_data") + or tokens_view.get("multi_modal_embeddings"), + encoder_input_token_ids=None) + contract_generator = self._engine_client_serving.try_stream( + preprocessed=preprocessed_shim, + sampling_params=contract_sampling_params, + post_processor=completion_stream_post_processor, + postproc_args=postproc_args, + raw_request=raw_request, + lora_request=request.lora_request, + disaggregated_params=disaggregated_params, + conversation_params=conversation_params, + trace_headers=trace_headers, + echo=bool(getattr(request, "echo", False)), + ) + if contract_generator is not None: + return StreamingResponse( + content=generator_wrapper(contract_generator), + media_type="text/event-stream") + + # Legacy path only: the contract path above never constructs + # PostprocParams (the formatter callable stays frontend-side). + postproc_params = PostprocParams( + post_processor=completion_stream_post_processor + if request.stream else completion_response_post_processor, + postproc_args=postproc_args, + ) promise = self.generator.generate_async( inputs=tokens_prompt, sampling_params=sampling_params, diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 95e5bf391667..52520bda42c7 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -480,6 +480,13 @@ "kind": "value", "path": "encoder_max_num_tokens" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "experimental_engine_client" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 0db864d9e540..59bb61deb5b2 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -35,6 +35,10 @@ methods: annotation: bool default: False status: prototype + experimental_engine_client: + annotation: bool + default: False + status: prototype # Bindings and mirrored configs peft_cache_config: annotation: Optional[tensorrt_llm.llmapi.llm_args.PeftCacheConfig] diff --git a/tests/unittest/executor/engine_client/conftest.py b/tests/unittest/executor/engine_client/conftest.py new file mode 100644 index 000000000000..7c303fa802c1 --- /dev/null +++ b/tests/unittest/executor/engine_client/conftest.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Record per-test pytest verdicts into the GPU run artifact. + +Only active for ``test_gpu_e2e.py`` runs with ``TLLM_ENGINE_CLIENT_GPU_REPORT`` +set: every test's outcome is merged into the artifact so the report carries +the actual pytest verdicts, not just case payloads. +""" + +import json +import os + +import pytest + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + outcome = yield + report = outcome.get_result() + if report.when != "call" or "test_gpu_e2e" not in item.nodeid: + return + path = os.environ.get("TLLM_ENGINE_CLIENT_GPU_REPORT") + if not path: + return + try: + data = {} + if os.path.exists(path): + with open(path) as f: + data = json.load(f) + verdicts = data.setdefault("verdicts", {}) + verdicts[item.nodeid.split("::", 1)[-1]] = report.outcome + with open(path, "w") as f: + json.dump(data, f, indent=2) + except Exception: + pass diff --git a/tests/unittest/executor/engine_client/golden/engine_request_full.msgpack b/tests/unittest/executor/engine_client/golden/engine_request_full.msgpack new file mode 100644 index 000000000000..0af4ad5ea48b Binary files /dev/null and b/tests/unittest/executor/engine_client/golden/engine_request_full.msgpack differ diff --git a/tests/unittest/executor/engine_client/golden/error_frame.msgpack b/tests/unittest/executor/engine_client/golden/error_frame.msgpack new file mode 100644 index 000000000000..51b3eea87875 Binary files /dev/null and b/tests/unittest/executor/engine_client/golden/error_frame.msgpack differ diff --git a/tests/unittest/executor/engine_client/golden/generate_golden.py b/tests/unittest/executor/engine_client/golden/generate_golden.py new file mode 100644 index 000000000000..3cb3dd348584 --- /dev/null +++ b/tests/unittest/executor/engine_client/golden/generate_golden.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Regenerates the checked-in binary golden fixtures (*.msgpack). + +Run from the repo root: + python3 tests/unittest/executor/engine_client/golden/generate_golden.py + +A diff in any .msgpack file is a wire-format change of the same protocol +version and must be treated as such (see ENGINE_CONTRACT.md). +""" + +import pathlib +import sys + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[5] +sys.path.insert(0, str(REPO_ROOT)) + +from tensorrt_llm.executor.engine_client.codec import encode # noqa: E402 + +HERE = pathlib.Path(__file__).parent + + +def golden_objects(): + """The canonical fixture payloads (shared with test_codec.py).""" + sys.path.insert(0, str(HERE.parent)) + from test_codec import GOLDEN_OBJECTS + return GOLDEN_OBJECTS + + +def main(): + for name, factory in golden_objects(): + data = encode(factory()) + path = HERE / f"{name}.msgpack" + path.write_bytes(data) + print(f"wrote {path.name}: {len(data)} bytes") + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/executor/engine_client/golden/minimal_delta.msgpack b/tests/unittest/executor/engine_client/golden/minimal_delta.msgpack new file mode 100644 index 000000000000..df94b7ffaa59 Binary files /dev/null and b/tests/unittest/executor/engine_client/golden/minimal_delta.msgpack differ diff --git a/tests/unittest/executor/engine_client/golden/model_context.msgpack b/tests/unittest/executor/engine_client/golden/model_context.msgpack new file mode 100644 index 000000000000..60af14a32007 Binary files /dev/null and b/tests/unittest/executor/engine_client/golden/model_context.msgpack differ diff --git a/tests/unittest/executor/engine_client/golden/request_complete.msgpack b/tests/unittest/executor/engine_client/golden/request_complete.msgpack new file mode 100644 index 000000000000..981d8833ed6b Binary files /dev/null and b/tests/unittest/executor/engine_client/golden/request_complete.msgpack differ diff --git a/tests/unittest/executor/engine_client/golden/terminal_stop.msgpack b/tests/unittest/executor/engine_client/golden/terminal_stop.msgpack new file mode 100644 index 000000000000..72ee21f73cad Binary files /dev/null and b/tests/unittest/executor/engine_client/golden/terminal_stop.msgpack differ diff --git a/tests/unittest/executor/engine_client/golden/token_delta_full.msgpack b/tests/unittest/executor/engine_client/golden/token_delta_full.msgpack new file mode 100644 index 000000000000..9a651bb33892 Binary files /dev/null and b/tests/unittest/executor/engine_client/golden/token_delta_full.msgpack differ diff --git a/tests/unittest/executor/engine_client/test_assembler.py b/tests/unittest/executor/engine_client/test_assembler.py new file mode 100644 index 000000000000..948dfc61edd7 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_assembler.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Assembler tests: incremental detok, ordered stops, presentation maps.""" + +import pytest + +from tensorrt_llm.executor.engine_client.assembler import FrontendResponseAssembler +from tensorrt_llm.executor.engine_client.contract import (ErrorFrame, + FrontendOutputConfig, + RequestComplete, + Terminal, TokenDelta) + + +class CharTokenizer: + """Deterministic fake: token id N decodes to chr(97 + N % 26).""" + + def decode(self, ids, **kwargs): + return "".join(chr(97 + i % 26) for i in ids) + + def decode_incrementally(self, token_ids, prev_text=None, states=None, *, + flush=False, stream_interval=1, **kwargs): + prev_text = prev_text or "" + return prev_text + self.decode(token_ids), states or {} + + +class HoldingTokenizer(CharTokenizer): + """Holds the last token pending until flush (exercises terminal flush).""" + + def decode_incrementally(self, token_ids, prev_text=None, states=None, *, + flush=False, stream_interval=1, **kwargs): + prev_text = prev_text or "" + states = dict(states or {"held": []}) + pending = states["held"] + list(token_ids) + if flush: + return prev_text + self.decode(pending), {"held": []} + if pending: + states["held"] = pending[-1:] + return prev_text + self.decode(pending[:-1]), states + return prev_text, states + + +def ids_for(text: str) -> list: + return [ord(ch) - 97 for ch in text] + + +def make_assembler(stop_strings=(), stop_sequence_reasons=(), include_stop=False, + tokenizer=None, abort_calls=None): + config = FrontendOutputConfig(stop_strings=tuple(stop_strings), + stop_sequence_reasons=tuple(stop_sequence_reasons), + include_stop_str_in_output=include_stop) + callback = abort_calls.append if abort_calls is not None else None + return FrontendResponseAssembler("req-1", config, + tokenizer=tokenizer or CharTokenizer(), + abort_callback=callback) + + +def delta(tokens, event_seq=0, logprobs=None): + return TokenDelta(request_id="req-1", sequence_id=0, + new_token_ids=tuple(tokens), logprobs=logprobs, + event_seq=event_seq) + + +def terminal(finish_reason, stop_reason=None, event_seq=5): + return Terminal(request_id="req-1", sequence_id=0, finish_reason=finish_reason, + stop_reason=stop_reason, event_seq=event_seq) + + +def complete(status="ok", prompt=3, completion=2, cached=None, event_seq=6): + return RequestComplete(request_id="req-1", status=status, prompt_tokens=prompt, + completion_tokens=completion, cached_tokens=cached, + event_seq=event_seq) + + +class CountingTokenizer(CharTokenizer): + """Records incremental-decode calls to guard the O(total tokens) bound.""" + + def __init__(self): + self.incremental_calls = [] + self.full_decode_calls = 0 + + def decode(self, ids, **kwargs): + self.full_decode_calls += 1 + return super().decode(ids, **kwargs) + + def decode_incrementally(self, token_ids, prev_text=None, states=None, *, + flush=False, stream_interval=1, **kwargs): + self.incremental_calls.append(list(token_ids)) + prev_text = prev_text or "" + piece = "".join(chr(97 + i % 26) for i in token_ids) + return prev_text + piece, states or {} + + +class TestDetokenizationComplexity: + + def test_non_incremental_tokenizer_rejected_at_setup(self): + from tensorrt_llm.executor.engine_client.assembler import AssemblerError + + class DecodeOnly: + + def decode(self, ids, **kwargs): + return "" + + with pytest.raises(AssemblerError): + make_assembler(tokenizer=DecodeOnly()) + + def test_streaming_cost_is_linear_in_tokens(self): + tokenizer = CountingTokenizer() + assembler = make_assembler(tokenizer=tokenizer) + delta_count = 50 + for i in range(delta_count): + assembler.process_frames([delta(ids_for("ab"), event_seq=i)]) + # One incremental call per delta, each fed ONLY the new ids — never + # the accumulated sequence, and never a full re-decode. + assert len(tokenizer.incremental_calls) == delta_count + assert all(len(call) == 2 for call in tokenizer.incremental_calls) + assert tokenizer.full_decode_calls == 0 + + +class TestPlainAssembly: + + def test_text_diffs_accumulate(self): + assembler = make_assembler() + updates = assembler.process_frames([delta(ids_for("hel"))]) + assert updates[0].text_diff == "hel" + updates = assembler.process_frames([delta(ids_for("lo"), 1)]) + assert updates[0].text_diff == "lo" + assert assembler.text == "hello" + + def test_finish_and_usage(self): + assembler = make_assembler() + assembler.process_frames([delta(ids_for("hi"))]) + updates = assembler.process_frames( + [terminal("stop"), complete(cached=2)]) + assert [u.kind for u in updates] == ["finish", "complete"] + assert updates[0].finish_reason == "stop" + assert updates[1].usage == { + "prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5, + "cached_tokens": 2 + } + + def test_logprobs_accumulate(self): + assembler = make_assembler() + assembler.process_frames([delta(ids_for("ab"), logprobs=(-0.1, -0.2))]) + assert assembler.logprobs == [-0.1, -0.2] + + def test_terminal_flushes_held_detok_state(self): + assembler = make_assembler(tokenizer=HoldingTokenizer()) + updates = assembler.process_frames([delta(ids_for("hi"))]) + assert updates[0].text_diff == "h" # last token held + updates = assembler.process_frames([terminal("length"), complete()]) + finish = next(u for u in updates if u.kind == "finish") + deltas = [u for u in updates if u.kind == "delta"] + assert deltas and deltas[0].text_diff == "i" + assert finish.finish_reason == "length" + assert assembler.text == "hi" + + +class TestStopStrings: + + def test_cross_delta_stop_string_trims_and_aborts(self): + abort_calls = [] + assembler = make_assembler(stop_strings=("xy", ), abort_calls=abort_calls) + first = assembler.process_frames([delta(ids_for("ax"))]) + assert first[0].text_diff == "ax" + second = assembler.process_frames([delta(ids_for("yb"), 1)]) + update = second[0] + assert update.finish_reason == "stop" + assert update.stop_reason == "xy" + # Trimmed before the stop string: only nothing new streams out. + assert assembler.text == "a" + assert update.text_diff == "" + assert abort_calls == ["req-1"] + + def test_include_stop_str_keeps_string(self): + assembler = make_assembler(stop_strings=("xy", ), include_stop=True) + assembler.process_frames([delta(ids_for("a"))]) + updates = assembler.process_frames([delta(ids_for("xy"), 1)]) + assert assembler.text == "axy" + assert updates[0].text_diff == "xy" + assert updates[0].finish_reason == "stop" + + def test_post_stop_frames_absorbed(self): + assembler = make_assembler(stop_strings=("x", )) + assembler.process_frames([delta(ids_for("x"))]) + updates = assembler.process_frames( + [delta(ids_for("zz"), 1), terminal("abort"), complete("aborted")]) + kinds = [u.kind for u in updates] + # The extra delta and the abort terminal are absorbed; usage still lands. + assert kinds == ["complete"] + assert assembler.finish_reason == "stop" + + def test_config_order_wins_over_position(self): + assembler = make_assembler(stop_strings=("q", "a")) + # "a" appears earlier in the text, but "q" is configured first and + # both are present in the same scan window: config order wins, + # mirroring the legacy first-match-in-config-order loop. + updates = assembler.process_frames([delta(ids_for("aq"))]) + assert updates[0].stop_reason == "q" + + +class TestEngineStopTrim: + + def test_stop_word_tokens_trimmed_in_same_batch(self): + sequence = tuple(ids_for("xy")) + assembler = make_assembler( + stop_sequence_reasons=((sequence, "xy"), )) + updates = assembler.process_frames([ + delta(ids_for("abxy")), + terminal("stop", stop_reason="xy"), + complete(completion=4), + ]) + by_kind = {u.kind: u for u in updates} + assert assembler.token_ids == ids_for("ab") + assert assembler.text == "ab" + assert by_kind["delta"].text_diff == "ab" + assert by_kind["finish"].finish_reason == "stop" + # User-visible usage counts assembled (post-trim) tokens. + assert by_kind["complete"].usage["completion_tokens"] == 2 + + def test_include_stop_str_keeps_engine_stop_tokens(self): + sequence = tuple(ids_for("xy")) + assembler = make_assembler(stop_sequence_reasons=((sequence, "xy"), ), + include_stop=True) + assembler.process_frames([ + delta(ids_for("abxy")), + terminal("stop", stop_reason="xy"), + ]) + assert assembler.text == "abxy" + + +class TestPresentationMap: + + def test_abort_renders_cancelled(self): + assembler = make_assembler() + updates = assembler.process_frames([terminal("abort")]) + assert updates[0].finish_reason == "cancelled" + + def test_timeout_renders_timeout(self): + assembler = make_assembler() + updates = assembler.process_frames( + [terminal("error", stop_reason="timeout")]) + assert updates[0].kind == "finish" + assert updates[0].finish_reason == "timeout" + + def test_other_errors_render_error_update(self): + assembler = make_assembler() + updates = assembler.process_frames( + [terminal("error", stop_reason="not_finished")]) + assert updates[0].kind == "error" + assert updates[0].error_message == "not_finished" + + def test_error_frame(self): + assembler = make_assembler() + updates = assembler.process_frames( + [ErrorFrame(request_id="req-1", error_code="worker_died", message="x")]) + assert updates[0].kind == "error" + assert updates[0].error_code == "worker_died" + + def test_failed_complete_without_finish_is_error(self): + assembler = make_assembler() + assembler.process_frames([delta(ids_for("a"))]) + updates = assembler.process_frames([complete(status="failed")]) + assert updates[0].kind == "error" diff --git a/tests/unittest/executor/engine_client/test_audit_remediations.py b/tests/unittest/executor/engine_client/test_audit_remediations.py new file mode 100644 index 000000000000..0babd56b2402 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_audit_remediations.py @@ -0,0 +1,351 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Regression tests for the final conformance-audit remediations.""" + +from types import SimpleNamespace + +import msgpack +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.engine_client.assembler import FrontendResponseAssembler +from tensorrt_llm.executor.engine_client.codec import EncodeError, encode +from tensorrt_llm.executor.engine_client.contract import (FrontendOutputConfig, + RequestComplete, + Terminal, TokenDelta) +from tensorrt_llm.executor.engine_client.envelope import (EnvelopeError, + normalize_response) +from tensorrt_llm.executor.engine_client.router import (EngineFrameRouter, + RouterError) +from tensorrt_llm.executor.result import ResponseWrapper + + +def make_result(**overrides): + fields = dict(is_final=False, output_token_ids=[[5]], + finish_reasons=[tllm.FinishReason.NOT_FINISHED], log_probs=None, + cum_log_probs=None, sequence_index=0) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def response_for(client_id, **overrides): + return LlmResponse(request_id=client_id, result=make_result(**overrides), + client_id=client_id) + + +def bind_request(router, request_id="req-1", client_id=41, stop_reasons=()): + binding = router.bind(request_id, client_id, (1, 2), stop_reasons) + binding.stream_opened = True # tests drain the delivery directly + return binding + + +def drain(binding): + frames = [] + while True: + frame, ready = binding.delivery.pop_nowait() + if not ready: + return frames + frames.append(frame) + + +class TestRouterRemediations: + + def test_nonzero_sequence_index_fails_request(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router) + router.on_response(response_for(41, sequence_index=1)) + frames = drain(binding) + assert frames, "expected a typed ending" + assert frames[-1].error_code == "router_error" or isinstance( + frames[-1], RequestComplete) + assert binding.ended + + def test_zero_token_final_carries_cached_tokens(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router) + router.on_response( + response_for(41, is_final=True, output_token_ids=[[]], + finish_reasons=[tllm.FinishReason.LENGTH], + cached_tokens=7)) + frames = drain(binding) + complete = frames[-1] + assert isinstance(complete, RequestComplete) + assert complete.cached_tokens == 7 + + def test_closed_delivery_is_not_recorded_as_overflow(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router) + binding.delivery.close() + router.on_response(response_for(41)) + assert router.counters["overflow_aborts"] == 0 + assert router.counters["late_or_duplicate_absorbed"] == 1 + + def test_malformed_stop_reasons_rejected_at_bind(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + with pytest.raises(RouterError): + bind_request(router, stop_reasons=(((1, 2), lambda: None), )) + with pytest.raises(RouterError): + bind_request(router, request_id="req-2", + stop_reasons=((("a", ), "x"), )) + + +class TestPromptLogprobLifecycle: + """Cached/map-shaped prompt logprobs across multiple responses (R2).""" + + @staticmethod + def map_wrapper(client_id, tokens, prompt_entries): + from tensorrt_llm.executor.result import LogProbsResult + return ResponseWrapper( + response_for(client_id, output_token_ids=[list(tokens)]), + logprobs=LogProbsResult(prompt=prompt_entries, generation=None)) + + @staticmethod + def map_entries(): + from tensorrt_llm.executor.result import Logprob + # prompt (1, 2): expected lookup keys are prompt[1:] + first token. + return [{2: Logprob(-1.0, 1)}, {5: Logprob(-2.0, 1)}] + + def test_reattached_prompt_map_on_terminal_only_final_is_benign(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router) + # First token-carrying response delivers the prompt logprobs. + router.on_response(self.map_wrapper(41, [5], self.map_entries())) + # The worker re-attaches the cached map on the terminal-only final. + final = self.map_wrapper(41, [], self.map_entries()) + final._response.result.is_final = True + final._response.result.finish_reasons = [tllm.FinishReason.END_ID] + router.on_response(final) + frames = drain(binding) + assert isinstance(frames[-1], RequestComplete) + assert frames[-1].status == "ok" + deltas = [f for f in frames if isinstance(f, TokenDelta)] + assert deltas[0].prompt_logprobs == (-1.0, -2.0) + + def test_map_prompt_logprobs_held_from_no_token_response(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router) + # Map-shaped prompt logprobs arrive BEFORE any token: held raw. + router.on_response(self.map_wrapper(41, [], self.map_entries())) + router.on_response(response_for(41, output_token_ids=[[5]])) + router.on_response( + response_for(41, is_final=True, output_token_ids=[[]], + finish_reasons=[tllm.FinishReason.END_ID])) + frames = drain(binding) + deltas = [f for f in frames if isinstance(f, TokenDelta)] + assert deltas[0].prompt_logprobs == (-1.0, -2.0) + assert isinstance(frames[-1], RequestComplete) + assert frames[-1].status == "ok" + + +class TestCachedTokenConsistency: + """Absent / zero / match / mismatch cases (R2).""" + + def run_case(self, delta_cached, final_cached): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router) + router.on_response(response_for(41, cached_tokens=delta_cached)) + router.on_response( + response_for(41, is_final=True, output_token_ids=[[]], + finish_reasons=[tllm.FinishReason.END_ID], + cached_tokens=final_cached)) + return drain(binding) + + def test_absent(self): + frames = self.run_case(None, None) + assert frames[-1].status == "ok" and frames[-1].cached_tokens is None + + def test_zero(self): + frames = self.run_case(0, 0) + assert frames[-1].status == "ok" and frames[-1].cached_tokens == 0 + + def test_match(self): + frames = self.run_case(3, 3) + assert frames[-1].status == "ok" and frames[-1].cached_tokens == 3 + + def test_mismatch_is_typed_failure(self): + frames = self.run_case(3, 9) + assert isinstance(frames[-1], RequestComplete) + assert frames[-1].status == "failed" + + def test_final_only_value_accepted(self): + frames = self.run_case(None, 7) + assert frames[-1].status == "ok" and frames[-1].cached_tokens == 7 + + +class TestLifecycleRework: + """Plan AC-4 lifecycle: completed-id reuse + bounded state (R1).""" + + @staticmethod + def finish_and_retire(router, binding, client_id): + router.on_response( + response_for(client_id, is_final=True, output_token_ids=[[5]], + finish_reasons=[tllm.FinishReason.END_ID])) + drain(binding) + router.retire_delivery(binding.request_id) + + def test_request_id_reusable_after_delivery_retired(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router, "req-reuse", client_id=41) + # While active or un-retired: duplicate is rejected. + with pytest.raises(RouterError): + bind_request(router, "req-reuse", client_id=42) + self.finish_and_retire(router, binding, 41) + # After consumption + retirement: the id is reusable (positive case). + rebound = bind_request(router, "req-reuse", client_id=43) + self.finish_and_retire(router, rebound, 43) + + def test_state_returns_to_bounds_after_stress(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + for i in range(60): + binding = bind_request(router, f"req-{i}", client_id=100 + i) + self.finish_and_retire(router, binding, 100 + i) + assert router.active_request_count() == 0 + assert len(router._by_request) == 0 + assert len(router._by_client) == 0 + assert len(router._delivery_index) == 0 + assert len(router._tombstones) <= router._tombstone_limit + assert len(router._retired_client_ids) <= router._retired_limit + + def test_close_before_end_retires_on_cancellation_final(self): + """AC-4: an early close records retirement intent; the in-flight + (token-carrying) cancellation final both ENDS the binding and + completes the retirement, so the request id is immediately + reusable rather than retained until tombstone eviction.""" + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router, "req-close", client_id=61) + router.on_response(response_for(61)) # mid-stream delta + # The consumer closes early (FrameStream.aclose: abort issued, + # delivery closed, retirement recorded). + binding.delivery.close() + router.retire_delivery("req-close") + assert not binding.ended + assert binding.retire_when_ended + # The in-flight cancellation final arrives afterwards — token- + # carrying, against the closed delivery. + router.on_response( + response_for(61, is_final=True, output_token_ids=[[5]], + finish_reasons=[tllm.FinishReason.CANCELLED])) + assert binding.ended + assert 61 not in router._by_client + assert "req-close" not in router._delivery_index + # Retirement completed: the id is immediately reusable. + rebound = bind_request(router, "req-close", client_id=62) + self.finish_and_retire(router, rebound, 62) + + def test_repeated_close_stress_returns_to_bounds(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + for i in range(60): + binding = bind_request(router, f"req-c{i}", client_id=300 + i) + router.on_response(response_for(300 + i)) + binding.delivery.close() + router.retire_delivery(f"req-c{i}") + router.on_response( + response_for(300 + i, is_final=True, output_token_ids=[[7]], + finish_reasons=[tllm.FinishReason.CANCELLED])) + assert len(router._by_request) == 0 + assert len(router._by_client) == 0 + assert len(router._delivery_index) == 0 + assert len(router._tombstones) <= router._tombstone_limit + assert len(router._retired_client_ids) <= router._retired_limit + + def test_retired_late_frame_absorbed_and_ancient_falls_through(self): + router = EngineFrameRouter(abort_fn=lambda cid: None, + tombstone_limit=2) + router._retired_limit = 2 + for i in range(5): + binding = bind_request(router, f"req-{i}", client_id=200 + i) + self.finish_and_retire(router, binding, 200 + i) + # Recently retired: claimed and absorbed. + assert router.route_response(response_for(204)) is True + # Ancient (aged out of the bounded set): falls through — safe because + # proxy ids are never reused, so legacy lookup drops it. + assert router.route_response(response_for(200)) is False + + +class TestCounterIncrements: + + def test_router_failures_increments_on_normalization_failure(self): + router = EngineFrameRouter(abort_fn=lambda cid: None) + binding = bind_request(router) + before = router.counters["router_failures"] + # Mismatched logprob count fails normalization typed. + router.on_response( + response_for(41, output_token_ids=[[5, 6]], log_probs=[[-0.5]])) + assert router.counters["router_failures"] == before + 1 + frames = drain(binding) + assert frames and binding.ended + + +class TestCodecRemediation: + + def test_mixed_mapping_keys_typed_error(self): + delta = TokenDelta(request_id="r", sequence_id=0, new_token_ids=(1, )) + object.__setattr__(delta, "metrics", {"a": 1.0, 3: 2.0}) + with pytest.raises(EncodeError) as excinfo: + encode(delta) + assert excinfo.value.reason == "non_primitive" + + +class TestEnvelopeRemediation: + + def test_non_finite_metric_typed_error(self): + wrapper = ResponseWrapper(response_for(41), + request_perf_metrics={"t": float("nan")}) + with pytest.raises(EnvelopeError) as excinfo: + normalize_response(wrapper) + assert excinfo.value.reason == "metrics_mismatch" + + +class CharTokenizer: + + def decode(self, ids, **kwargs): + return "".join(chr(97 + i % 26) for i in ids) + + def decode_incrementally(self, token_ids, prev_text=None, states=None, *, + flush=False, stream_interval=1, **kwargs): + prev_text = prev_text or "" + return prev_text + self.decode(token_ids), states or {} + + +def ids_for(text): + return [ord(ch) - 97 for ch in text] + + +class TestAssemblerRemediation: + + def test_engine_stop_trim_spans_multiple_deltas(self): + sequence = tuple(ids_for("xy")) + config = FrontendOutputConfig(stop_sequence_reasons=((sequence, "xy"), )) + assembler = FrontendResponseAssembler("req-1", config, + tokenizer=CharTokenizer()) + # The stop sequence spans two buffered deltas in one batch. + updates = assembler.process_frames([ + TokenDelta(request_id="req-1", sequence_id=0, + new_token_ids=tuple(ids_for("abx")), event_seq=0), + TokenDelta(request_id="req-1", sequence_id=0, + new_token_ids=tuple(ids_for("y")), event_seq=1), + Terminal(request_id="req-1", sequence_id=0, finish_reason="stop", + stop_reason="xy", event_seq=2), + RequestComplete(request_id="req-1", status="ok", prompt_tokens=2, + completion_tokens=4, event_seq=3), + ]) + assert assembler.token_ids == ids_for("ab") + assert assembler.text == "ab" + rendered = "".join(u.text_diff for u in updates if u.kind == "delta") + assert rendered == "ab" + delta_tokens = [t for u in updates if u.kind == "delta" + for t in u.new_token_ids] + assert delta_tokens == ids_for("ab") diff --git a/tests/unittest/executor/engine_client/test_codec.py b/tests/unittest/executor/engine_client/test_codec.py new file mode 100644 index 000000000000..673ffbe20395 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_codec.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Codec round-trip, golden-fixture stability, and strict rejection tests.""" + +import os + +import msgpack +import pytest + +from tensorrt_llm.executor.engine_client.codec import (MAX_MESSAGE_BYTES, + DecodeError, + EncodeError, decode, + encode) +from tensorrt_llm.executor.engine_client.contract import ( + ENGINE_CONTRACT_VERSION, EngineCapabilities, EngineHealth, EngineRequest, + EngineSamplingConfig, ErrorFrame, FrontendModelContext, + FrontendOutputConfig, GuidedDecodingSpec, IterationStatsBatch, + KvCacheEventsBatch, RequestComplete, Terminal, TokenDelta, TokenizerSpec) + + +def full_engine_request() -> EngineRequest: + return EngineRequest( + request_id="req-Ω-你好", + prompt_token_ids=(1, 2, 3, 2**62, 0), + sampling=EngineSamplingConfig( + max_new_tokens=128, end_id=2, pad_id=0, + stop_token_ids=(13, 14), stop_token_sequences=((7, 8), (9, )), + min_tokens=1, temperature=0.5, top_p=0.9, top_k=40, seed=-12345, + repetition_penalty=1.1, presence_penalty=0.0, frequency_penalty=-0.5, + num_logprobs=1, num_prompt_logprobs=1), + guided_decoding=GuidedDecodingSpec(mode="json_schema", payload='{"type":"object"}'), + required_features=("guided_decoding", )) + + +def full_token_delta() -> TokenDelta: + return TokenDelta(request_id="r1", sequence_id=0, new_token_ids=(5, 6, 7), + logprobs=(-0.25, -1.5, -0.001), prompt_logprobs=(-2.0, ), + metrics={"cached_tokens": 3.0, "arrival_time": 12.5}, event_seq=4) + + +def full_model_context() -> FrontendModelContext: + return FrontendModelContext( + tokenizer=TokenizerSpec(uri="/models/llama", + files_manifest=(("tokenizer.json", "ab" * 32), ), + revision="main", fast=True, trust_remote_code=False, + added_tokens_json="[]"), + capabilities=EngineCapabilities(features=("streaming", "logprobs")), + chat_template_source="{{ messages }}", chat_template_version="v1", + eos_id=2, pad_id=0, max_context_length=8192) + + +ROUND_TRIP_CASES = [ + full_engine_request(), + EngineRequest(request_id="minimal", prompt_token_ids=(1, ), + sampling=EngineSamplingConfig(max_new_tokens=1)), + full_token_delta(), + TokenDelta(request_id="m", sequence_id=1, new_token_ids=(42, )), + Terminal(request_id="r1", sequence_id=0, finish_reason="stop", stop_reason="", + event_seq=9), + Terminal(request_id="r1", sequence_id=0, finish_reason="length"), + Terminal(request_id="r1", sequence_id=0, finish_reason="error", stop_reason="timeout"), + RequestComplete(request_id="r1", status="ok", prompt_tokens=5, completion_tokens=3, + cached_tokens=3, metrics={"e2e_time": 1.25}, event_seq=10), + RequestComplete(request_id="r1", status="failed", prompt_tokens=0, completion_tokens=0), + ErrorFrame(request_id="r2", error_code="worker_died", + message="crashed before first response ☠"), + EngineCapabilities(features=("streaming", )), + EngineSamplingConfig(max_new_tokens=4), + GuidedDecodingSpec(mode="json_object"), + FrontendOutputConfig(stop_strings=("stop", ), + stop_sequence_reasons=(((1, 2), "stop"), ((3, ), 13)), end_id=2, + num_logprobs=1), + TokenizerSpec(uri="/m"), + full_model_context(), + EngineHealth(healthy=True, detail="ok"), + IterationStatsBatch(entries=('{"iter": 1}', )), + KvCacheEventsBatch(entries=()), +] + +# Checked-in binary golden fixtures live in golden/.msgpack; regenerate +# with golden/generate_golden.py. A diff in any .msgpack file is a wire-format +# change of the same protocol version and must be treated as such. +GOLDEN_OBJECTS = [ + ("engine_request_full", full_engine_request), + ("token_delta_full", full_token_delta), + ("terminal_stop", + lambda: Terminal(request_id="r1", sequence_id=0, finish_reason="stop", stop_reason="", + event_seq=9)), + ("request_complete", + lambda: RequestComplete(request_id="r1", status="ok", prompt_tokens=5, completion_tokens=3, + cached_tokens=3, metrics={"e2e_time": 1.25}, event_seq=10)), + ("error_frame", + lambda: ErrorFrame(request_id="r2", error_code="worker_died", + message="crashed before first response \u2620")), + ("model_context", full_model_context), + ("minimal_delta", lambda: TokenDelta(request_id="m", sequence_id=1, new_token_ids=(42, ))), +] + +GOLDEN_DIR = os.path.join(os.path.dirname(__file__), "golden") + + + +class TestRoundTrip: + + @pytest.mark.parametrize("obj", ROUND_TRIP_CASES, ids=lambda o: type(o).__name__) + def test_round_trip(self, obj): + data = encode(obj) + assert decode(data) == obj + assert encode(decode(data)) == data + + def test_canonical_metric_key_order(self): + a = TokenDelta(request_id="r", sequence_id=0, new_token_ids=(1, ), + metrics={"a": 1.0, "b": 2.0}) + b = TokenDelta(request_id="r", sequence_id=0, new_token_ids=(1, ), + metrics={"b": 2.0, "a": 1.0}) + assert encode(a) == encode(b) + + +class TestGoldenFixtures: + + @pytest.mark.parametrize("name,factory", GOLDEN_OBJECTS, ids=lambda c: c + if isinstance(c, str) else "") + def test_golden(self, name, factory): + expected = factory() + with open(os.path.join(GOLDEN_DIR, f"{name}.msgpack"), "rb") as f: + data = f.read() + assert decode(data) == expected, f"golden fixture {name} decode mismatch" + assert encode(expected) == data, f"golden fixture {name} encode not byte-stable" + + +def _envelope(kind, fields) -> bytes: + return msgpack.packb({"k": kind, "d": fields}, use_bin_type=True) + + +def _terminal_fields(**overrides): + fields = {"request_id": "r", "sequence_id": 0, "finish_reason": "stop"} + fields.update(overrides) + return fields + + +class TestDecodeRejects: + + def test_unknown_kind(self): + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("mystery_frame", {})) + assert excinfo.value.reason == "unknown_kind" + + def test_unknown_field(self): + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("terminal", _terminal_fields(surprise=1))) + assert excinfo.value.reason == "unknown_field" + + def test_missing_required_field(self): + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("terminal", {"request_id": "r"})) + assert excinfo.value.reason == "missing_field" + + def test_wrong_item_type(self): + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("terminal", _terminal_fields(sequence_id="zero"))) + assert excinfo.value.reason == "invalid_content" + + def test_bool_is_not_int(self): + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("terminal", _terminal_fields(sequence_id=True))) + assert excinfo.value.reason == "invalid_content" + + def test_newer_protocol_version(self): + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("terminal", + _terminal_fields(protocol_version=ENGINE_CONTRACT_VERSION + 1))) + assert excinfo.value.reason == "version_unsupported" + + def test_nested_version_check(self): + fields = { + "request_id": "r", + "prompt_token_ids": [1], + "sampling": {"max_new_tokens": 1, + "protocol_version": ENGINE_CONTRACT_VERSION + 1}, + } + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("engine_request", fields)) + assert excinfo.value.reason == "version_unsupported" + + def test_nan_in_payload(self): + with pytest.raises(DecodeError) as excinfo: + decode(_envelope("token_delta", + {"request_id": "r", "sequence_id": 0, "new_token_ids": [1], + "logprobs": [float("nan")]})) + assert excinfo.value.reason == "invalid_content" + + def test_extension_type_rejected(self): + data = msgpack.packb({"k": "terminal", "d": msgpack.ExtType(4, b"x")}) + with pytest.raises(DecodeError): + decode(data) + + def test_bytes_payload_rejected(self): + data = msgpack.packb({"k": "terminal", "d": _terminal_fields(stop_reason=b"raw")}, + use_bin_type=True) + with pytest.raises(DecodeError): + decode(data) + + def test_malformed_bytes(self): + with pytest.raises(DecodeError): + decode(b"\xc1\x00\x00") + + def test_wrong_envelope_shape(self): + with pytest.raises(DecodeError) as excinfo: + decode(msgpack.packb({"kind": "terminal"})) + assert excinfo.value.reason == "malformed" + + def test_oversized_message(self): + with pytest.raises(DecodeError) as excinfo: + decode(b"\x00" * (MAX_MESSAGE_BYTES + 1)) + assert excinfo.value.reason == "limit_exceeded" + + def test_nesting_depth_limit(self): + nested = 1 + for _ in range(40): + nested = [nested] + with pytest.raises(DecodeError): + decode(msgpack.packb({"k": "terminal", "d": _terminal_fields(stop_reason=nested)})) + + def test_non_bytes_input(self): + with pytest.raises(DecodeError): + decode("not-bytes") + + +class TestEncodeRejects: + + def test_unregistered_type(self): + with pytest.raises(EncodeError) as excinfo: + encode({"not": "a wire type"}) + assert excinfo.value.reason == "unknown_type" + + def test_smuggled_nan_rejected(self): + delta = TokenDelta(request_id="r", sequence_id=0, new_token_ids=(1, ), logprobs=(-0.5, )) + object.__setattr__(delta, "logprobs", (float("nan"), )) + with pytest.raises(EncodeError) as excinfo: + encode(delta) + assert excinfo.value.reason == "not_finite" + + def test_smuggled_object_rejected(self): + delta = TokenDelta(request_id="r", sequence_id=0, new_token_ids=(1, )) + object.__setattr__(delta, "metrics", {"cb": print}) + with pytest.raises(EncodeError) as excinfo: + encode(delta) + assert excinfo.value.reason == "non_primitive" + + def test_smuggled_bytes_rejected(self): + frame = ErrorFrame(request_id="r", error_code="x") + object.__setattr__(frame, "message", b"raw") + with pytest.raises(EncodeError) as excinfo: + encode(frame) + assert excinfo.value.reason == "non_primitive" diff --git a/tests/unittest/executor/engine_client/test_contract.py b/tests/unittest/executor/engine_client/test_contract.py new file mode 100644 index 000000000000..850351d656d9 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_contract.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Construction-time validation tests for the engine contract wire types.""" + +import dataclasses +from types import MappingProxyType + +import pytest + +from tensorrt_llm.executor.engine_client.contract import ( + ENGINE_CONTRACT_VERSION, INT64_MAX, MAX_STOP_SEQUENCES, + ContractConstructionError, EngineCapabilities, EngineRequest, + EngineSamplingConfig, ErrorFrame, FrontendModelContext, + FrontendOutputConfig, GuidedDecodingSpec, RequestComplete, Terminal, + TokenDelta, TokenizerSpec, validate_no_callables) + + +def make_sampling(**overrides) -> EngineSamplingConfig: + kwargs = dict(max_new_tokens=16) + kwargs.update(overrides) + return EngineSamplingConfig(**kwargs) + + +def make_request(**overrides) -> EngineRequest: + kwargs = dict(request_id="req-1", prompt_token_ids=(1, 2, 3), sampling=make_sampling()) + kwargs.update(overrides) + return EngineRequest(**kwargs) + + +class TestConstructionAccepts: + + def test_full_request(self): + request = make_request( + sampling=make_sampling(end_id=2, pad_id=0, stop_token_ids=[13], + stop_token_sequences=[[7, 8]], temperature=0.5, seed=-1, + num_logprobs=1, num_prompt_logprobs=1), + guided_decoding=GuidedDecodingSpec(mode="json_object"), + required_features=["guided_decoding"]) + assert request.protocol_version == ENGINE_CONTRACT_VERSION + # Sequences normalize to tuples so the value graph is immutable. + assert request.sampling.stop_token_ids == (13, ) + assert request.sampling.stop_token_sequences == ((7, 8), ) + assert request.required_features == ("guided_decoding", ) + + def test_metrics_become_immutable_copies(self): + source = {"cached_tokens": 3.0} + delta = TokenDelta(request_id="r", sequence_id=0, new_token_ids=[1], metrics=source) + assert isinstance(delta.metrics, MappingProxyType) + with pytest.raises(TypeError): + delta.metrics["cached_tokens"] = 9.0 + source["cached_tokens"] = 9.0 + assert delta.metrics["cached_tokens"] == 3.0 + + def test_frames_are_frozen(self): + terminal = Terminal(request_id="r", sequence_id=0, finish_reason="stop") + with pytest.raises(dataclasses.FrozenInstanceError): + terminal.finish_reason = "length" + + def test_ordered_stop_sequence_reasons(self): + config = FrontendOutputConfig(stop_strings=["a", "b"], + stop_sequence_reasons=[[[1, 2], "a"], [[1, 2], "b"]]) + # Order is preserved verbatim: first match wins downstream. + assert config.stop_sequence_reasons == (((1, 2), "a"), ((1, 2), "b")) + + def test_model_context(self): + context = FrontendModelContext( + tokenizer=TokenizerSpec(uri="/models/m", files_manifest=[["tokenizer.json", "00"]]), + capabilities=EngineCapabilities(features=["streaming"]), + eos_id=2, max_context_length=4096) + assert context.tokenizer.files_manifest == (("tokenizer.json", "00"), ) + + def test_int_float_field_accepts_int(self): + assert make_sampling(temperature=1).temperature == 1.0 + + +class TestConstructionRejects: + + @pytest.mark.parametrize("kwargs", [ + dict(max_new_tokens=True), + dict(max_new_tokens=0), + dict(end_id=1.5), + dict(stop_token_ids=[True]), + dict(stop_token_sequences=[[]]), + dict(stop_token_sequences=[[1]] * (MAX_STOP_SEQUENCES + 1)), + dict(temperature=float("nan")), + dict(top_p=float("inf")), + dict(seed=INT64_MAX + 1), + dict(num_logprobs=-1), + dict(protocol_version=0), + ]) + def test_sampling_config_rejects(self, kwargs): + with pytest.raises(ContractConstructionError): + make_sampling(**kwargs) + + @pytest.mark.parametrize("kwargs", [ + dict(request_id=""), + dict(request_id=7), + dict(prompt_token_ids=()), + dict(prompt_token_ids=(1, "2")), + dict(sampling=None), + dict(guided_decoding="json"), + dict(required_features=(1, )), + ]) + def test_request_rejects(self, kwargs): + with pytest.raises(ContractConstructionError): + make_request(**kwargs) + + @pytest.mark.parametrize("kwargs", [ + dict(new_token_ids=()), + dict(new_token_ids=(1, ), logprobs=(-0.5, -0.5)), + dict(new_token_ids=(True, )), + dict(sequence_id=-1), + dict(event_seq=-1), + dict(metrics={"k": "not-a-float"}), + dict(metrics={1: 2.0}), + dict(logprobs=(float("nan"), )), + ]) + def test_token_delta_rejects(self, kwargs): + base = dict(request_id="r", sequence_id=0, new_token_ids=(1, )) + base.update(kwargs) + with pytest.raises(ContractConstructionError): + TokenDelta(**base) + + def test_finish_reason_must_be_known(self): + with pytest.raises(ContractConstructionError): + Terminal(request_id="r", sequence_id=0, finish_reason="timeout") + + def test_stop_reason_rejects_bool_and_float(self): + for bad in (True, 1.5): + with pytest.raises(ContractConstructionError): + Terminal(request_id="r", sequence_id=0, finish_reason="stop", stop_reason=bad) + + def test_request_complete_status_must_be_known(self): + with pytest.raises(ContractConstructionError): + RequestComplete(request_id="r", status="done", prompt_tokens=1, completion_tokens=1) + + def test_request_complete_negative_counts(self): + with pytest.raises(ContractConstructionError): + RequestComplete(request_id="r", status="ok", prompt_tokens=-1, completion_tokens=0) + + def test_error_frame_requires_code(self): + with pytest.raises(ContractConstructionError): + ErrorFrame(request_id="r", error_code="") + + def test_guided_decoding_payload_required_unless_json_object(self): + with pytest.raises(ContractConstructionError): + GuidedDecodingSpec(mode="regex") + assert GuidedDecodingSpec(mode="json_object").payload is None + + def test_guided_decoding_unknown_mode(self): + with pytest.raises(ContractConstructionError): + GuidedDecodingSpec(mode="yaml", payload="x") + + def test_output_config_rejects_unordered_or_malformed_pairs(self): + with pytest.raises(ContractConstructionError): + FrontendOutputConfig(stop_sequence_reasons=[[[1, 2]]]) + with pytest.raises(ContractConstructionError): + FrontendOutputConfig(stop_sequence_reasons={(1, 2): "a"}) + with pytest.raises(ContractConstructionError): + FrontendOutputConfig(stop_strings=[""]) + + def test_bool_flags_reject_int(self): + with pytest.raises(ContractConstructionError): + FrontendOutputConfig(detokenize=1) + + def test_tokenizer_spec_manifest_pairs(self): + with pytest.raises(ContractConstructionError): + TokenizerSpec(uri="/m", files_manifest=[["only-path"]]) + + +class TestNoCallables: + + def test_callable_rejected_anywhere(self): + with pytest.raises(ContractConstructionError): + validate_no_callables({"hook": lambda x: x}) + with pytest.raises(ContractConstructionError): + validate_no_callables([1, 2, print]) + + def test_bytes_rejected(self): + with pytest.raises(ContractConstructionError): + validate_no_callables(b"raw") + + def test_arbitrary_object_rejected(self): + class Payload: + pass + + with pytest.raises(ContractConstructionError): + validate_no_callables({"obj": Payload()}) + + def test_valid_graph_passes(self): + validate_no_callables(make_request()) diff --git a/tests/unittest/executor/engine_client/test_conversion.py b/tests/unittest/executor/engine_client/test_conversion.py new file mode 100644 index 000000000000..1c98a88fd465 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_conversion.py @@ -0,0 +1,372 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion, normalization-boundary, and eligibility-matrix tests.""" + +import dataclasses + +import pytest +import torch + +from tensorrt_llm.executor.base_worker import _request_stop_words +from tensorrt_llm.executor.engine_client.codec import decode, encode +from tensorrt_llm.executor.engine_client.contract import EngineRequest +from tensorrt_llm.executor.engine_client.conversion import ( + ELIGIBILITY_MATRIX, ConversionError, RequestIneligibleError, + convert_request, derive_required_features, + engine_request_to_generation_request, prepare_sampling_params) +from tensorrt_llm.sampling_params import (GuidedDecodingParams, LogprobMode, + SamplingParams) + + +class FakeTokenizer: + """Minimal tokenizer surface used by SamplingParams._setup.""" + + eos_token_id = 2 + pad_token_id = 0 + + def encode(self, text, add_special_tokens=False): + # Deterministic non-trivial encoding: one token per character. + return [100 + (ord(ch) % 50) for ch in text] + + +class FakeGenerationConfig: + eos_token_id = [2, 32000] + forced_bos_token_id = None + forced_eos_token_id = None + + +def prepared_params(**kwargs) -> SamplingParams: + params = SamplingParams(max_tokens=8, **kwargs) + return prepare_sampling_params(params, tokenizer=FakeTokenizer(), + generation_config=FakeGenerationConfig()) + + +def convert(params: SamplingParams, **kwargs): + return convert_request("req-1", [1, 2, 3], params, streaming=True, **kwargs) + + +class TestNormalizationBoundary: + + def test_eos_pad_defaults_from_tokenizer(self): + params = prepared_params() + assert params.end_id == 2 + assert params.pad_id == 0 + + def test_generation_config_stop_ids_added_once(self): + params = prepared_params() + assert params.stop_token_ids == [32000] + + def test_stop_strings_tokenized_exactly_once(self): + params = prepared_params(stop=["DONE"]) + first = params._stop_word_ids + assert first is not None + # Conversion must not re-tokenize: same object flows through. + engine_request, output_config = convert(params) + assert engine_request.sampling.stop_token_sequences == (tuple(first[0]), ) + assert params._stop_word_ids is first + + def test_explicit_end_id_skips_setup(self): + params = SamplingParams(max_tokens=8, end_id=7, pad_id=7) + prepare_sampling_params(params, tokenizer=None) + assert params.end_id == 7 + assert params.stop_token_ids is None + + def test_bart_is_ineligible(self): + config = FakeGenerationConfig() + config.forced_bos_token_id = 0 + with pytest.raises(RequestIneligibleError) as excinfo: + prepare_sampling_params(SamplingParams(max_tokens=8), + tokenizer=FakeTokenizer(), + generation_config=config, model_type="bart") + assert excinfo.value.axis == "bart_forced_tokens" + + def test_thinking_budget_is_ineligible(self): + with pytest.raises(RequestIneligibleError) as excinfo: + prepare_sampling_params(SamplingParams(max_tokens=8, thinking_token_budget=100), + tokenizer=FakeTokenizer()) + assert excinfo.value.axis == "thinking_token_budget" + + def test_stream_interval_and_perf_metrics_defaults(self): + params = prepare_sampling_params(SamplingParams(max_tokens=8, end_id=2), + tokenizer=None, stream_interval=4, + force_return_perf_metrics=True) + assert params._stream_interval == 4 + assert params.return_perf_metrics + + +class TestConvertRequest: + + def test_plain_text_conversion(self): + params = prepared_params(temperature=0.5, top_p=0.9, top_k=40, seed=11, + stop_token_ids=[13], stop=["STOP"], logprobs=1) + engine_request, output_config = convert(params) + sampling = engine_request.sampling + assert sampling.max_new_tokens == 8 + assert sampling.end_id == 2 and sampling.pad_id == 0 + # stop_token_ids keeps user ids + generation-config additions. + assert set(sampling.stop_token_ids) == {13, 32000} + assert len(sampling.stop_token_sequences) == 1 + assert output_config.stop_strings == ("STOP", ) + # Ordered association: token-id reasons first, then string reasons. + reasons = output_config.stop_sequence_reasons + assert reasons[-1][1] == "STOP" + assert reasons[-1][0] == sampling.stop_token_sequences[0] + assert all(pair[1] == pair[0][0] for pair in reasons[:-1]) + assert output_config.num_logprobs == 1 + assert engine_request.required_features == () + + def test_guided_decoding_maps_to_spec_and_feature(self): + params = prepared_params(guided_decoding=GuidedDecodingParams(json={"type": "object"})) + engine_request, _ = convert(params) + assert engine_request.guided_decoding.mode == "json_schema" + assert "object" in engine_request.guided_decoding.payload + assert engine_request.required_features == ("guided_decoding", ) + assert derive_required_features(engine_request) == ("guided_decoding", ) + + def test_unprepared_params_rejected(self): + with pytest.raises(ConversionError): + convert(SamplingParams(max_tokens=8)) + + def test_untokenized_stop_strings_rejected(self): + params = SamplingParams(max_tokens=8, end_id=2, stop=["STOP"]) + with pytest.raises(ConversionError): + convert(params) + + +REJECTION_CASES = [ + ("non_streaming", {}, dict(streaming=False)), + ("echo", {}, dict(echo=True)), + ("n_gt_1", dict(n=2, top_p=0.9), {}), + ("n_gt_1", dict(best_of=2, top_p=0.9), {}), + ("beam_search", dict(use_beam_search=True, best_of=1, n=1), {}), + ("beam_search", dict(length_penalty=1.0), {}), + ("top_logprobs", dict(logprobs=5), {}), + ("prompt_top_logprobs", dict(prompt_logprobs=5), {}), + ("logits_processor", dict(apply_batched_logits_processor=True), {}), + ("embedding_bias", dict(embedding_bias=torch.zeros(4)), {}), + ("bad_words", dict(bad=["x"]), {}), + ("bad_words", dict(bad_token_ids=[3]), {}), + ("ignore_eos", dict(ignore_eos=True), {}), + ("min_p", dict(min_p=0.2), {}), + ("top_p_extras", dict(top_p_min=0.5), {}), + ("no_repeat_ngram", dict(no_repeat_ngram_size=2), {}), + ("prompt_ignore_length", dict(prompt_ignore_length=1), {}), + ("return_logits", dict(return_context_logits=True), {}), + ("return_logits", dict(return_generation_logits=True), {}), + ("return_logits", dict(additional_model_outputs=["hidden"]), {}), + ("exclude_input_from_output", dict(exclude_input_from_output=False), {}), + ("truncate_prompt_tokens", dict(truncate_prompt_tokens=4), {}), + ("multimodal", {}, dict(multimodal_params=object())), + ("lora", {}, dict(lora_request=object())), + ("prompt_adapter", {}, dict(prompt_adapter_request=object())), + ("disaggregated", {}, dict(disaggregated_params=object())), + ("scheduling_params", {}, dict(scheduling_params=object())), + ("conversation_params", {}, dict(conversation_params=object())), + ("postproc_params", {}, dict(postproc_params=object())), + ("trace_headers", {}, dict(trace_headers={"traceparent": "x"})), + ("cache_salt", {}, dict(cache_salt="salt")), + ("query_token_ids", {}, dict(query_token_ids=[1])), + ("encoder_input", {}, dict(encoder_input_token_ids=[1])), + ("kv_cache_retention", {}, dict(kv_cache_retention_config=object())), + ("priority", {}, dict(priority=0.9)), +] + + +class TestRejections: + + @pytest.mark.parametrize("axis,param_kwargs,convert_kwargs", REJECTION_CASES) + def test_rejected_pre_submit(self, axis, param_kwargs, convert_kwargs): + params = prepared_params(**param_kwargs) + streaming = convert_kwargs.pop("streaming", True) + with pytest.raises(RequestIneligibleError) as excinfo: + convert_request("req-1", [1, 2, 3], params, streaming=streaming, + **convert_kwargs) + assert excinfo.value.axis == axis + + def test_logprobs_mode_rejected(self): + non_raw = [m for m in LogprobMode if getattr(m, "value", m) not in ("raw", "RAW")] + if not non_raw: + pytest.skip("no non-RAW logprob mode in this build") + params = prepared_params(logprobs=1) + params.logprobs_mode = non_raw[0] + with pytest.raises(RequestIneligibleError) as excinfo: + convert(params) + assert excinfo.value.axis == "logprobs_mode" + + def test_every_ineligible_axis_has_a_rejection_case(self): + tested = {case[0] for case in REJECTION_CASES} + tested |= {"logprobs_mode", "bart_forced_tokens", "thinking_token_budget", + "lookahead_config"} + matrix_axes = {rule.axis for rule in ELIGIBILITY_MATRIX + if rule.classification == "ineligible"} + assert matrix_axes <= tested, f"untested axes: {matrix_axes - tested}" + + def test_lookahead_rejected(self): + params = prepared_params() + params.lookahead_config = object() + with pytest.raises(RequestIneligibleError) as excinfo: + convert(params) + assert excinfo.value.axis == "lookahead_config" + + +class TestOpenAiFieldInventoryCompleteness: + """Machine-readable completeness over the live OpenAI request surface.""" + + def test_every_openai_request_field_dispositioned(self): + from tensorrt_llm.executor.engine_client.conversion import \ + OPENAI_REQUEST_FIELD_DISPOSITION + from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, + CompletionRequest) + surface = set(ChatCompletionRequest.model_fields) | set( + CompletionRequest.model_fields) + missing = surface - set(OPENAI_REQUEST_FIELD_DISPOSITION) + assert not missing, f"OpenAI fields without a V0 disposition: {sorted(missing)}" + stale = set(OPENAI_REQUEST_FIELD_DISPOSITION) - surface + assert not stale, f"dispositions for unknown OpenAI fields: {sorted(stale)}" + + def test_axis_references_exist_in_matrix(self): + from tensorrt_llm.executor.engine_client.conversion import \ + OPENAI_REQUEST_FIELD_DISPOSITION + known = {"preprocessing", "supported", "frontend", "capability_gated"} + matrix_axes = {rule.axis for rule in ELIGIBILITY_MATRIX} + for field, disposition in OPENAI_REQUEST_FIELD_DISPOSITION.items(): + if disposition in known: + continue + assert disposition.startswith("axis:"), (field, disposition) + axis = disposition[len("axis:"):] + assert axis in matrix_axes, ( + f"{field} references unknown matrix axis {axis!r}") + + # The normative dispositions the contract's scope matrix promises: + # completeness alone would let a wrong classification (e.g. `echo` + # marked frontend-supported) pass silently. + EXPECTED_DISPOSITIONS = { + "echo": "axis:echo", + "n": "axis:n_gt_1", + "best_of": "axis:n_gt_1", + "logit_bias": "axis:embedding_bias", + "top_logprobs": "axis:top_logprobs", + "stream": "supported", + "stream_options": "frontend", + "suffix": "frontend", + "max_tokens": "supported", + "temperature": "supported", + "response_format": "capability_gated", + "lora_request": "axis:lora", + "disaggregated_params": "axis:disaggregated", + "cache_salt": "axis:cache_salt", + "messages": "preprocessing", + "prompt": "preprocessing", + "model": "preprocessing", + } + + def test_expected_dispositions(self): + from tensorrt_llm.executor.engine_client.conversion import \ + OPENAI_REQUEST_FIELD_DISPOSITION + for field, expected in self.EXPECTED_DISPOSITIONS.items(): + assert OPENAI_REQUEST_FIELD_DISPOSITION.get(field) == expected, ( + f"{field}: expected disposition {expected!r}, got " + f"{OPENAI_REQUEST_FIELD_DISPOSITION.get(field)!r}") + + +class TestEligibilityMatrixCompleteness: + + def test_every_sampling_params_field_classified(self): + covered = set() + for rule in ELIGIBILITY_MATRIX: + covered.update(rule.sources) + missing = [] + for field in dataclasses.fields(SamplingParams): + if field.name in ("_stop_word_ids", "_bad_word_ids"): + continue # derived storage of stop/bad, covered via their axes + if field.name not in covered: + missing.append(field.name) + assert not missing, f"SamplingParams fields not in ELIGIBILITY_MATRIX: {missing}" + + def test_classifications_are_known(self): + for rule in ELIGIBILITY_MATRIX: + assert rule.classification in ("supported", "normalized", "ineligible", + "config_gate") + + +class TestCompositionalTranslation: + + def make_engine_request(self) -> EngineRequest: + params = prepared_params(temperature=0.25, top_p=0.8, top_k=16, seed=99, + repetition_penalty=1.05, presence_penalty=0.1, + frequency_penalty=0.2, min_tokens=2, + stop_token_ids=[13], stop=["END"], logprobs=1, + prompt_logprobs=0) + engine_request, _ = convert(params) + return engine_request + + def test_translation_consumes_only_the_encoded_request(self): + engine_request = self.make_engine_request() + decoded = decode(encode(engine_request)) + direct = engine_request_to_generation_request(engine_request) + via_wire = engine_request_to_generation_request(decoded) + assert via_wire.prompt_token_ids == direct.prompt_token_ids + assert via_wire.stop_token_sequences == direct.stop_token_sequences + assert via_wire.sampling_params == direct.sampling_params + + def test_every_wire_field_reaches_worker_consumption(self): + engine_request = self.make_engine_request() + generation_request = engine_request_to_generation_request(engine_request) + sp = generation_request.sampling_params + sampling = engine_request.sampling + + # Fields the worker reads directly off SamplingParams at enqueue. + assert sp.max_tokens == sampling.max_new_tokens + assert sp.end_id == sampling.end_id + assert sp.pad_id == sampling.pad_id + assert sp.logprobs == sampling.num_logprobs + assert sp.prompt_logprobs == sampling.num_prompt_logprobs + + # Fields consumed through _get_sampling_config into the runtime config. + config = sp._get_sampling_config() + assert config.temperature == pytest.approx(sampling.temperature) + assert config.top_p == pytest.approx(sampling.top_p) + assert config.top_k == sampling.top_k + assert config.seed == sampling.seed + assert config.repetition_penalty == pytest.approx(sampling.repetition_penalty) + assert config.presence_penalty == pytest.approx(sampling.presence_penalty) + assert config.frequency_penalty == pytest.approx(sampling.frequency_penalty) + assert config.min_tokens == sampling.min_tokens + assert config.beam_width == 1 + assert config.num_return_sequences == 1 + + # Stop handling: ids through SamplingParams, sequences through the + # carrier, merged by the worker's stop-word derivation. + stop_words = _request_stop_words(generation_request) + expected = [[token_id] for token_id in sampling.stop_token_ids] + expected += [list(seq) for seq in sampling.stop_token_sequences] + assert stop_words == expected + + def test_translation_refuses_guided_decoding(self): + params = prepared_params(guided_decoding=GuidedDecodingParams(json_object=True)) + engine_request, _ = convert(params) + with pytest.raises(ConversionError): + engine_request_to_generation_request(engine_request) + + def test_no_setup_needed_downstream(self): + engine_request = self.make_engine_request() + generation_request = engine_request_to_generation_request(engine_request) + # The synthetic params carry no stop strings, so nothing downstream + # can trigger (or need) tokenizer-based re-setup. + assert generation_request.sampling_params.stop is None + assert generation_request.sampling_params._stop_word_ids is None + assert generation_request.stop_token_sequences == [ + list(seq) for seq in engine_request.sampling.stop_token_sequences + ] diff --git a/tests/unittest/executor/engine_client/test_e5b_cutover.py b/tests/unittest/executor/engine_client/test_e5b_cutover.py new file mode 100644 index 000000000000..c8c6b21bb6da --- /dev/null +++ b/tests/unittest/executor/engine_client/test_e5b_cutover.py @@ -0,0 +1,375 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""E5b cut-over tests on a skeleton of the REAL GenerationExecutorProxy. + +Gate ①c: contract requests get no rank-0 GenerationResult, contract +responses are exclusively routed (never legacy-delivered), both request +populations are enumerated symmetrically on shutdown, contract client-id +ownership survives tombstone eviction, and frontend request-id reuse is +forbidden for the service session. +""" + +import threading +from queue import Queue +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.engine_client.contract import (EngineRequest, + EngineSamplingConfig, + ErrorFrame, + RequestComplete) +from tensorrt_llm.executor.engine_client.local_client import ( + EngineClientConfigError, resolve_engine_client_flag) +from tensorrt_llm.executor.engine_client.router import RouterError +from tensorrt_llm.executor.engine_client.service import (EngineService, + EngineServiceError) +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.utils import ErrorResponse + + +class FakeRequestQueue: + + def __init__(self): + self.items = [] + self.raise_on_put = False + self.noblock = [] + + def put(self, item): + if self.raise_on_put: + raise RuntimeError("enqueue failed") + self.items.append(item) + + def put_noblock(self, item, retry=0): + self.noblock.append(item) + + +class FakeLegacyResult: + + def __init__(self): + self.queue = Queue() + self.aborted = False + + def abort(self): + self.aborted = True + + +def make_proxy() -> GenerationExecutorProxy: + proxy = object.__new__(GenerationExecutorProxy) + proxy._results = {} + proxy.doing_shutdown = False + proxy._fatal_error = None + proxy._last_client_id = 0 + proxy.request_queue = FakeRequestQueue() + proxy.workers_started = True + proxy.mpi_futures = [] + proxy._shutdown_event = threading.Event() + stub = SimpleNamespace() + + class _Stub: + + def handle_background_error(self, *args, **kwargs): + return None + + def start_dispatch_threads(self): + return None + + stub_obj = _Stub() + proxy._stub = stub_obj + proxy._handle_background_error = stub_obj.handle_background_error + proxy._start_dispatch_threads = stub_obj.start_dispatch_threads + return proxy + + +def make_engine_request(request_id="req-1") -> EngineRequest: + return EngineRequest(request_id=request_id, prompt_token_ids=(1, 2, 3), + sampling=EngineSamplingConfig(max_new_tokens=8, end_id=2, + pad_id=0)) + + +def make_service(proxy) -> EngineService: + service = EngineService(proxy) + proxy.attach_engine_service(service) + return service + + +def contract_final(client_id): + result = SimpleNamespace(is_final=True, output_token_ids=[[5]], + finish_reasons=[tllm.FinishReason.END_ID], + log_probs=None, cum_log_probs=None, sequence_index=0) + return LlmResponse(request_id=client_id, result=result, client_id=client_id) + + +def drain_binding(binding): + frames = [] + while True: + frame, ready = binding.delivery.pop_nowait() + if not ready: + return frames + frames.append(frame) + + +class TestContractNativeSubmit: + + def test_no_rank0_generation_result(self): + proxy = make_proxy() + service = make_service(proxy) + client_id = service.submit_contract(make_engine_request()) + assert proxy._results == {} # the binding replaces GenerationResult + assert proxy.request_queue.items[0].id == client_id + assert service.router.get_binding("req-1") is not None + + def test_submission_rejected_during_shutdown(self): + proxy = make_proxy() + service = make_service(proxy) + proxy.doing_shutdown = True + with pytest.raises(EngineServiceError): + service.submit_contract(make_engine_request()) + assert proxy.request_queue.items == [] + + def test_enqueue_failure_yields_typed_ending_and_no_state_leak(self): + proxy = make_proxy() + service = make_service(proxy) + proxy.request_queue.raise_on_put = True + with pytest.raises(RuntimeError): + service.submit_contract(make_engine_request()) + binding = service.router.open_stream_binding("req-1") + frames = drain_binding(binding) + assert [type(f).__name__ for f in frames] == ["ErrorFrame"] + assert frames[0].error_code == "enqueue_failed" + assert proxy._results == {} + # The request was never enqueued, so no runtime abort was sent. + assert proxy.aborted == [] if hasattr(proxy, "aborted") else True + + def test_second_service_attachment_rejected(self): + proxy = make_proxy() + make_service(proxy) + with pytest.raises(RuntimeError): + proxy.attach_engine_service(EngineService(proxy)) + + def test_request_id_reuse_rejected_while_delivery_retained(self): + proxy = make_proxy() + service = make_service(proxy) + client_id = service.submit_contract(make_engine_request()) + service.router.route_response(contract_final(client_id)) + # Ended but not consumed/retired: the id is still retained. + with pytest.raises(RouterError): + service.submit_contract(make_engine_request()) + + def test_allocator_wrap_rejected_for_contract_population(self): + proxy = make_proxy() + service = make_service(proxy) + proxy._last_client_id = (1 << 64) - 1 # next allocation wraps to 0 + with pytest.raises(EngineServiceError): + service.submit_contract(make_engine_request()) + assert proxy.request_queue.items == [] + assert service.router.get_binding("req-1") is None + + def test_allocator_wrap_rejected_for_legacy_population(self): + from tensorrt_llm.executor.request import GenerationRequest + from tensorrt_llm.sampling_params import SamplingParams + proxy = make_proxy() + service = make_service(proxy) + # Interleaved monotonic allocations across populations are fine... + contract_id = service.submit_contract(make_engine_request()) + service.observe_legacy_allocation(contract_id + 1) + # ...but a wrapped id handed to a LEGACY submission must fail that + # submission before it can collide with contract-owned state. + proxy._last_client_id = (1 << 64) - 1 + request = GenerationRequest(prompt_token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=4), + streaming=True) + with pytest.raises(EngineServiceError): + proxy.submit(request) + assert proxy._results == {} + assert all(item.id == contract_id + for item in proxy.request_queue.items) + + +class TestExclusiveRouting: + + def test_contract_claimed_legacy_delivered(self): + proxy = make_proxy() + service = make_service(proxy) + contract_id = service.submit_contract(make_engine_request()) + legacy_result = FakeLegacyResult() + legacy_id = 999 + proxy._results[legacy_id] = legacy_result + + contract_response = contract_final(contract_id) + legacy_response = ErrorResponse(client_id=legacy_id, error_msg="x", + request_id=1) + proxy.result_queue = SimpleNamespace( + get=lambda items=[[contract_response, legacy_response]]: items.pop(0) + if items else None) + assert proxy.dispatch_result_task() is True + + # Contract: consumed by the service (frames emitted), absent legacy. + binding = service.router.open_stream_binding("req-1") + frames = drain_binding(binding) + assert isinstance(frames[-1], RequestComplete) + # Legacy: delivered and popped exactly as before. + assert legacy_result.queue.get_nowait() is legacy_response + assert legacy_id not in proxy._results + + def test_contract_response_never_reaches_legacy_even_with_result_entry(self): + proxy = make_proxy() + service = make_service(proxy) + contract_id = service.submit_contract(make_engine_request()) + # Hostile setup: even if a legacy result somehow existed under the + # same id, the claim runs first and legacy never sees the response. + shadow_result = FakeLegacyResult() + proxy._results[contract_id] = shadow_result + response = contract_final(contract_id) + proxy.result_queue = SimpleNamespace( + get=lambda items=[[response]]: items.pop(0) if items else None) + proxy.dispatch_result_task() + assert shadow_result.queue.empty() + assert contract_id in proxy._results # untouched by the claim + + def test_late_frame_after_tombstone_eviction_is_absorbed(self): + proxy = make_proxy() + service = make_service(proxy) + service.router._tombstone_limit = 1 + first_id = service.submit_contract(make_engine_request("req-a")) + service.router.route_response(contract_final(first_id)) + second_id = service.submit_contract(make_engine_request("req-b")) + service.router.route_response(contract_final(second_id)) + # req-a's tombstone is evicted now; a very late frame for it must + # still be classified contract-owned and absorbed, never legacy. + before = service.router.counters["late_or_duplicate_absorbed"] + assert service.route_response(contract_final(first_id)) is True + assert service.router.counters["late_or_duplicate_absorbed"] == before + 1 + + +class TestShutdownSymmetry: + + def test_pre_shutdown_covers_both_populations(self): + proxy = make_proxy() + aborted = [] + proxy.abort_request = aborted.append + service = make_service(proxy) + # Re-point the router abort at the spy (service was built before). + service.router._abort_fn = aborted.append + contract_id = service.submit_contract(make_engine_request()) + legacy_result = FakeLegacyResult() + proxy._results[7] = legacy_result + + proxy.pre_shutdown() + + binding = service.router.open_stream_binding("req-1") + frames = drain_binding(binding) + assert isinstance(frames[-1], (RequestComplete, ErrorFrame)) + assert legacy_result.aborted + assert contract_id in aborted # best-effort runtime cancellation + assert proxy.request_queue.noblock == [None] # worker sentinel sent + + def test_close_client_aborts_in_flight_work(self): + proxy = make_proxy() + aborted = [] + proxy.abort_request = aborted.append + service = make_service(proxy) + service.router._abort_fn = aborted.append + contract_id = service.submit_contract(make_engine_request()) + service.close_client() + assert contract_id in aborted + binding = service.router.open_stream_binding("req-1") + frames = drain_binding(binding) + assert isinstance(frames[-1], (RequestComplete, ErrorFrame)) + + +class TestFlagResolver: + + @pytest.mark.parametrize("env,args_value,expected", [ + (None, False, False), + (None, True, True), + ("0", False, False), + ("0", True, False), + ("1", False, True), + ("1", True, True), + ]) + def test_precedence_matrix(self, monkeypatch, env, args_value, expected): + if env is None: + monkeypatch.delenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", raising=False) + else: + monkeypatch.setenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", env) + assert resolve_engine_client_flag(args_value) is expected + + def test_invalid_env_fails_closed(self, monkeypatch): + monkeypatch.setenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", "yes") + with pytest.raises(EngineClientConfigError): + resolve_engine_client_flag(True) + + def test_llm_args_field_exists(self): + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + field = TorchLlmArgs.model_fields["experimental_engine_client"] + assert field.default is False + + +class TestConcurrentMixedTraffic: + + def test_contract_and_legacy_share_the_proxy_without_misrouting(self): + proxy = make_proxy() + service = make_service(proxy) + contract_count = 12 + legacy_count = 12 + + contract_ids = {} + + def submit_contracts(): + for i in range(contract_count): + cid = service.submit_contract(make_engine_request(f"req-{i}")) + contract_ids[cid] = f"req-{i}" + + legacy_results = {} + + def submit_legacy(): + for i in range(legacy_count): + with proxy._submission_lock: + legacy_id = proxy._get_next_client_id() + result = FakeLegacyResult() + proxy._results[legacy_id] = result + legacy_results[legacy_id] = result + + t1 = threading.Thread(target=submit_contracts) + t2 = threading.Thread(target=submit_legacy) + t1.start() + t2.start() + t1.join() + t2.join() + + # Ids never collide across populations. + assert not set(contract_ids) & set(legacy_results) + + # Dispatch a mixed batch. + responses = [contract_final(cid) for cid in contract_ids] + responses += [ + ErrorResponse(client_id=lid, error_msg="bye", request_id=1) + for lid in legacy_results + ] + proxy.result_queue = SimpleNamespace( + get=lambda items=[responses]: items.pop(0) if items else None) + proxy.dispatch_result_task() + + for cid, request_id in contract_ids.items(): + binding = service.router.open_stream_binding(request_id) + frames = drain_binding(binding) + assert isinstance(frames[-1], RequestComplete) + for lid, result in legacy_results.items(): + assert not result.queue.empty() + assert lid not in proxy._results + assert service.router.active_request_count() == 0 diff --git a/tests/unittest/executor/engine_client/test_envelope.py b/tests/unittest/executor/engine_client/test_envelope.py new file mode 100644 index 000000000000..63a2e54848ff --- /dev/null +++ b/tests/unittest/executor/engine_client/test_envelope.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Envelope normalization tests over the real PyTorch response shapes.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.engine_client.envelope import (EnvelopeError, + normalize_response) +from tensorrt_llm.executor.postproc_worker import PostprocWorker +from tensorrt_llm.executor.result import Logprob, LogProbsResult, ResponseWrapper +from tensorrt_llm.executor.utils import ErrorResponse +from tensorrt_llm.metrics.enums import RequestEventTiming + + +def make_result(**overrides) -> SimpleNamespace: + fields = dict(is_final=False, output_token_ids=[[5, 6]], finish_reasons=None, + log_probs=None, cum_log_probs=None, sequence_index=0) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def make_response(client_id=7, **result_overrides) -> LlmResponse: + return LlmResponse(request_id=1, result=make_result(**result_overrides), + client_id=client_id) + + +class TestTokenDeltaNormalization: + + def test_plain_delta(self): + envelope = normalize_response(make_response()) + assert envelope.client_id == 7 + assert envelope.new_token_ids == (5, 6) + assert not envelope.is_final + assert envelope.finish_reason_name is None + assert envelope.logprobs is None and envelope.prompt_logprobs is None + + def test_final_with_finish_reason(self): + envelope = normalize_response( + make_response(is_final=True, + finish_reasons=[tllm.FinishReason.END_ID])) + assert envelope.is_final + assert envelope.finish_reason_name == "END_ID" + + def test_no_token_final(self): + envelope = normalize_response( + make_response(is_final=True, output_token_ids=[[]], + finish_reasons=[tllm.FinishReason.CANCELLED])) + assert envelope.new_token_ids == () + assert envelope.finish_reason_name == "CANCELLED" + + def test_snapshot_is_immutable_copy(self): + source_tokens = [5, 6] + response = make_response(output_token_ids=[source_tokens]) + envelope = normalize_response(response) + source_tokens.append(999) + assert envelope.new_token_ids == (5, 6) + + def test_cached_tokens_sourced(self): + envelope = normalize_response(make_response(cached_tokens=3)) + assert envelope.cached_tokens == 3 + assert normalize_response(make_response()).cached_tokens is None + + def test_serialized_result_is_deserialized(self): + inner = make_result() + + class SerializedResult: + + def __init__(self): + self._result = b"opaque" + + def deserialize(self): + self._result = inner + for name, value in vars(inner).items(): + setattr(self, name, value) + + response = LlmResponse(request_id=1, result=SerializedResult(), client_id=7) + envelope = normalize_response(response) + assert envelope.new_token_ids == (5, 6) + + +class TestLogprobNormalization: + + def test_float_list_passes_through(self): + envelope = normalize_response( + make_response(log_probs=[[-0.5, -1.5]])) + assert envelope.logprobs == (-0.5, -1.5) + + def test_map_shape_exact_lookup(self): + entries = [{5: Logprob(logprob=-0.5, rank=1)}, + {6: Logprob(logprob=-1.5, rank=2), 9: Logprob(logprob=-0.1, rank=1)}] + envelope = normalize_response(make_response(log_probs=[entries])) + assert envelope.logprobs == (-0.5, -1.5) + + def test_map_missing_key_is_typed_error(self): + entries = [{9: Logprob(logprob=-0.1, rank=1)}, {6: Logprob(logprob=-1.5, rank=1)}] + with pytest.raises(EnvelopeError) as excinfo: + normalize_response(make_response(log_probs=[entries])) + assert excinfo.value.reason == "logprob_mismatch" + + def test_length_mismatch_is_typed_error(self): + with pytest.raises(EnvelopeError) as excinfo: + normalize_response(make_response(log_probs=[[-0.5]])) + assert excinfo.value.reason == "logprob_mismatch" + + def test_prompt_logprobs_float_passthrough(self): + wrapper = ResponseWrapper( + make_response(), + logprobs=LogProbsResult(prompt=[-2.0, -3.0], generation=None)) + envelope = normalize_response(wrapper) + assert envelope.prompt_logprobs == (-2.0, -3.0) + + def test_prompt_logprobs_map_shape_needs_binding(self): + prompt_entries = [{2: Logprob(-1.0, 1)}, {3: Logprob(-2.0, 1)}, {5: Logprob(-3.0, 1)}] + wrapper = ResponseWrapper( + make_response(), + logprobs=LogProbsResult(prompt=prompt_entries, generation=None)) + with pytest.raises(EnvelopeError): + normalize_response(wrapper) + # With the bound prompt ids (offset by one, plus the first generated + # token) the exact lookup succeeds. + envelope = normalize_response(wrapper, prompt_token_ids=(1, 2, 3)) + assert envelope.prompt_logprobs == (-1.0, -2.0, -3.0) + + +class TestMetricsNormalization: + + def test_enum_keys_convert(self): + wrapper = ResponseWrapper( + make_response(), + request_perf_metrics={RequestEventTiming.ARRIVAL_TIME: 1.5}) + envelope = normalize_response(wrapper) + assert envelope.metrics == {RequestEventTiming.ARRIVAL_TIME.value: 1.5} + + def test_non_numeric_metric_is_typed_error(self): + wrapper = ResponseWrapper(make_response(), + request_perf_metrics={"bad": "value"}) + with pytest.raises(EnvelopeError) as excinfo: + normalize_response(wrapper) + assert excinfo.value.reason == "metrics_mismatch" + + +class TestErrorShapes: + + def test_error_response(self): + envelope = normalize_response(ErrorResponse(client_id=3, error_msg="boom", + request_id=1)) + assert envelope.client_id == 3 + assert envelope.is_final + assert envelope.error_msg == "boom" + + def test_llm_response_error(self): + response = LlmResponse(request_id=1, error_msg="worker error", client_id=4) + envelope = normalize_response(response) + assert envelope.error_msg == "worker error" + assert envelope.is_final + + +class TestRejectedShapes: + + def test_postproc_output_rejected(self): + output = PostprocWorker.Output(client_id=1, res="text", is_final=True) + with pytest.raises(EnvelopeError) as excinfo: + normalize_response(output) + assert excinfo.value.reason == "postproc_parallel_shape" + + def test_cpp_engine_shape_rejected(self): + + class FakeCppResponse: + + client_id = 1 + + def has_error(self): + return False + + with pytest.raises(EnvelopeError) as excinfo: + normalize_response(FakeCppResponse()) + assert excinfo.value.reason == "cpp_engine_shape" + + def test_unknown_shape_rejected(self): + with pytest.raises(EnvelopeError) as excinfo: + normalize_response(object()) + assert excinfo.value.reason == "unknown_shape" diff --git a/tests/unittest/executor/engine_client/test_gpu_e2e.py b/tests/unittest/executor/engine_client/test_gpu_e2e.py new file mode 100644 index 000000000000..417ebcf8a881 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_gpu_e2e.py @@ -0,0 +1,826 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Real-GPU end-to-end suite for the engine client (the ①b gate). + +Runs the in-process client against the live PyTorch runtime over the IPC +proxy: pinned model (``LLM_MODELS_ROOT``-relative or the documented local +fallback), greedy sampling (``top_k=1``), TP1. Every stream has a hard +deadline so a hang fails rather than blocks. Both invocations merge into ONE +JSON artifact (commit, model + revision, GPU, backend, every case) at +``TLLM_ENGINE_CLIENT_GPU_REPORT``; the rank-0 measurement additionally +generates ``rank0_measurement.md`` next to it from the recorded values. + +Reproducible command (two invocations: the worker-crash test is destructive +and must not share a process with the main suite; TestShadowOverhead must +run before any client-fixture test, which file order guarantees): + TLLM_ENGINE_CLIENT_GPU_REPORT=/path/report.json \ + python3 -m pytest "tests/unittest/executor/engine_client/test_gpu_e2e.py::TestShadowOverhead" \ + "tests/unittest/executor/engine_client/test_gpu_e2e.py::TestGpuEndToEnd" \ + "tests/unittest/executor/engine_client/test_gpu_e2e.py::TestRank0Measurement" -v --timeout=1200 + TLLM_ENGINE_CLIENT_GPU_REPORT=/path/report.json \ + timeout 900 python3 -m pytest \ + "tests/unittest/executor/engine_client/test_gpu_e2e.py::TestWorkerCrash" -v --timeout=600 +The outer ``timeout`` on the crash invocation is REQUIRED: after the test +passes and the artifact is written, the interpreter can hang at exit on the +SIGKILLed MPI runtime's teardown; the reaper exit code (124/143) is expected +there — the verdict is pytest's summary line and the artifact, not the exit +code. +""" + +import asyncio +import json +import os +import subprocess +import time +import uuid +from copy import deepcopy +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm import LLM +from tensorrt_llm.executor.engine_client.assembler import FrontendResponseAssembler +from tensorrt_llm.executor.engine_client.contract import (ErrorFrame, + RequestComplete, + Terminal, TokenDelta) +from tensorrt_llm.executor.engine_client.conversion import ( + convert_request, prepare_sampling_params) +from tensorrt_llm.executor.engine_client.invariants import InvariantCheckingStream +from tensorrt_llm.executor.engine_client.local_client import ( + EngineClientConfig, LocalProcessEngineClient) +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.sampling_params import SamplingParams + +# The module-scoped LLM fixture keeps the proxy dispatch thread alive across +# tests by design; disable the per-test thread-leak check accordingly. +pytestmark = pytest.mark.threadleak(enabled=False) + +STREAM_DEADLINE = 120.0 +FAILURE_DEADLINE = 30.0 + +REPORT_PATH = os.environ.get("TLLM_ENGINE_CLIENT_GPU_REPORT", + "/tmp/engine_client_gpu_report.json") +REPORT = {"cases": {}} + + +def merge_report(update: dict) -> None: + """Merge results into ONE durable artifact across pytest invocations + (the destructive crash test runs in its own process).""" + existing = {} + try: + with open(REPORT_PATH) as f: + existing = json.load(f) + except (OSError, ValueError): + pass + cases = existing.get("cases", {}) + cases.update(update.pop("cases", {})) + existing.update(update) + existing["cases"] = cases + try: + with open(REPORT_PATH, "w") as f: + json.dump(existing, f, indent=2) + print(f"\nengine-client GPU report updated at {REPORT_PATH}") + except OSError: + pass + + +def model_revision(path: str) -> str: + """Pin the model revision: commit hash when present, else a content + hash over the config + tokenizer manifest.""" + import hashlib + digest = hashlib.sha256() + for name in ("config.json", "tokenizer_config.json", "tokenizer.json"): + candidate = os.path.join(path, name) + if os.path.isfile(candidate): + with open(candidate, "rb") as f: + digest.update(f.read()) + return f"sha256:{digest.hexdigest()[:16]}" + + +def model_path() -> str: + root = os.environ.get("LLM_MODELS_ROOT") + candidates = [] + if root: + candidates.append(os.path.join(root, "Qwen2.5-0.5B-Instruct")) + candidates.append("/workspace/models/Qwen2.5-0.5B-Instruct") + for candidate in candidates: + if os.path.isdir(candidate): + return candidate + pytest.skip(f"pinned model not found (looked in {candidates})") + + +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), + reason="needs a GPU") + + +def greedy_params(**overrides) -> dict: + kwargs = dict(max_tokens=32, top_k=1) + kwargs.update(overrides) + return kwargs + + +@pytest.fixture(scope="module") +def llm(): + # Block reuse off: with (partial) block reuse, a repeated prompt gets + # truncated context logits and therefore truncated prompt logprobs — a + # known runtime limitation (llm_return_logprobs_test_harness disables + # reuse for the same reason). The contract path turns that truncation + # into a typed logprob_mismatch failure instead of delivering silently + # short values, so the suite pins the supported configuration. + instance = LLM(model=model_path(), + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.3, + enable_block_reuse=False), + max_batch_size=8, + max_seq_len=1024) + if not isinstance(instance._executor, GenerationExecutorProxy): + instance.shutdown() + pytest.skip("executor is not the IPC proxy; V0 validates the IPC path only") + yield instance + instance.shutdown() + _write_report(instance) + + +def _write_report(instance) -> None: + try: + commit = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, + text=True, cwd=Path(__file__).parent).stdout.strip() + except Exception: + commit = "unknown" + path = model_path() + REPORT.update({ + "commit": commit, + "model": path, + "model_revision": model_revision(path), + "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none", + "execution_backend": instance.args.backend or "pytorch", + "attention_backend": getattr(instance.args, "attn_backend", "default"), + "topology": "TP1 over IPC proxy", + "sampling": "greedy (top_k=1)", + "flag": "TLLM_EXPERIMENTAL_ENGINE_CLIENT (client constructed with " + "flag_enabled=True)", + }) + merge_report(dict(REPORT)) + + +@pytest.fixture(scope="module") +def client(llm): + config = EngineClientConfig( + backend=llm.args.backend or "pytorch", + num_postprocess_workers=getattr(llm.args, "num_postprocess_workers", 0), + post_processor_hook_set=getattr(llm.args, "post_processor_hook", None) + is not None, + speculative_config_set=getattr(llm.args, "speculative_config", None) + is not None, + world_size=1, + flag_enabled=True) + engine_client = LocalProcessEngineClient(llm._executor, config) + yield engine_client + engine_client.close_client() + + +def prepared(llm, **overrides) -> SamplingParams: + params = SamplingParams(**greedy_params(**overrides)) + return prepare_sampling_params(params, tokenizer=llm.tokenizer, + hf_model_config=llm._hf_model_config, + generation_config=llm._generation_config) + + +def prompt_ids(llm, text="The capital of France is") -> list: + return llm.tokenizer.encode(text) + + +async def consume_contract(client, llm, engine_request, output_config, + stop_after=None, close_early=False): + request_id = client.submit(engine_request, output_config=output_config) + assembler = FrontendResponseAssembler(request_id, output_config, + tokenizer=llm.tokenizer, + abort_callback=client.abort) + stream = client.stream(request_id) + checked = InvariantCheckingStream(stream, request_id) + frames = [] + batch = [] + deltas_seen = 0 + while True: + try: + frame = await asyncio.wait_for(checked.__anext__(), STREAM_DEADLINE) + except StopAsyncIteration: + break + frames.append(frame) + batch.append(frame) + if isinstance(frame, TokenDelta): + deltas_seen += 1 + if stop_after is not None and deltas_seen >= stop_after: + client.abort(request_id) + stop_after = None + if close_early: + assembler.process_frames(batch) + await stream.aclose() + return frames, assembler + if isinstance(frame, (Terminal, RequestComplete, ErrorFrame)): + assembler.process_frames(batch) + batch = [] + if isinstance(frame, (RequestComplete, ErrorFrame)): + break + if batch: + assembler.process_frames(batch) + return frames, assembler + + +def run_contract(client, llm, params, prompt, output_extras=None, **consume_kwargs): + engine_request, output_config = convert_request( + f"gpu-{uuid.uuid4().hex[:8]}", prompt, params, streaming=True) + started = time.monotonic() + frames, assembler = asyncio.run( + consume_contract(client, llm, engine_request, output_config, + **consume_kwargs)) + elapsed = time.monotonic() - started + return frames, assembler, elapsed + + +def run_legacy(llm, prompt, params): + started = time.monotonic() + result = llm.generate(prompt, sampling_params=params) + elapsed = time.monotonic() - started + output = result.outputs[0] + return output, elapsed + + +def legacy_cpu_run(llm, prompts, max_tokens): + """One legacy streaming rep at len(prompts) concurrency; returns + (rank-0 process-CPU seconds, streamed tokens).""" + + async def one(prompt): + promise = llm.generate_async( + prompt, + sampling_params=SamplingParams(**greedy_params( + max_tokens=max_tokens)), + streaming=True) + tokens = 0 + async for output in promise: + tokens = len(output.outputs[0].token_ids) + return tokens + + async def run(): + return sum(await asyncio.gather(*(one(p) for p in prompts))) + + cpu_before = time.process_time() + tokens = asyncio.run(run()) + return time.process_time() - cpu_before, tokens + + +@requires_gpu +class TestShadowOverhead: + """Cost the attached contract machinery imposes on LEGACY traffic + (report-only per DEC-5). + + Post-cutover there is no dual-consumption shadow tap any more — routing + is exclusive — so the shadow-overhead number that still exists is the + per-response contract-population check every legacy response now passes + through (``service.route_response`` -> router membership test). Measured + as legacy per-token rank-0 CPU before the engine client/service is + attached vs after. MUST run before any test that uses the ``client`` + fixture (it attaches irreversibly), hence this class is first in the + file; the two arms are sequential for the same reason. + """ + + CONCURRENCY = 8 + WARMUPS = 2 + REPS = 10 + MAX_TOKENS = 64 + + def _arm(self, llm, prompts): + for _ in range(self.WARMUPS): + legacy_cpu_run(llm, prompts, self.MAX_TOKENS) + costs = [] + for _ in range(self.REPS): + cpu, tokens = legacy_cpu_run(llm, prompts, self.MAX_TOKENS) + costs.append(cpu / max(tokens, 1)) + return costs + + def test_legacy_cost_of_attached_machinery(self, llm, request): + service = getattr(llm._executor, "_engine_service", None) + if service is not None: + pytest.skip("service already attached; baseline arm impossible " + "(run this class first)") + prompts = [ + prompt_ids(llm, f"Write a story number {i} about a river:") + for i in range(self.CONCURRENCY) + ] + baseline = self._arm(llm, prompts) + # Attaching the module client wires the EngineService into the + # proxy dispatch loop — the machinery whose cost we measure. + request.getfixturevalue("client") + attached = self._arm(llm, prompts) + baseline_median = _median(baseline) + attached_median = _median(attached) + REPORT["cases"]["shadow_overhead"] = { + "concurrency": self.CONCURRENCY, + "warmups": self.WARMUPS, + "reps": self.REPS, + "unattached_legacy_cpu_s_per_token": baseline, + "attached_legacy_cpu_s_per_token": attached, + "unattached_median_s_per_token": baseline_median, + "attached_median_s_per_token": attached_median, + "delta_pct": 100.0 * (attached_median - baseline_median) / + baseline_median, + "note": "report-only per DEC-5; arms sequential because attach " + "is irreversible (no dual-consumption tap post-cutover)", + } + assert baseline_median > 0 and attached_median > 0 + + +@requires_gpu +class TestGpuEndToEnd: + + def test_plain_streaming_parity(self, llm, client): + prompt = prompt_ids(llm) + legacy_output, legacy_time = run_legacy(llm, prompt, + SamplingParams(**greedy_params())) + frames, assembler, contract_time = run_contract( + client, llm, prepared(llm), prompt) + assert assembler.token_ids == list(legacy_output.token_ids) + assert assembler.text == legacy_output.text + assert assembler.finish_reason == legacy_output.finish_reason + assert assembler.usage["completion_tokens"] == len(legacy_output.token_ids) + assert assembler.usage["prompt_tokens"] == len(prompt) + REPORT["cases"]["plain_streaming"] = { + "legacy_s": legacy_time, "contract_s": contract_time, + "tokens": len(assembler.token_ids) + } + + def test_stop_token(self, llm, client): + prompt = prompt_ids(llm) + reference, _ = run_legacy(llm, prompt, SamplingParams(**greedy_params())) + if len(reference.token_ids) < 4: + pytest.skip("reference generation too short to derive a stop token") + stop_token = reference.token_ids[3] + legacy_output, _ = run_legacy( + llm, prompt, + SamplingParams(**greedy_params(stop_token_ids=[stop_token]))) + frames, assembler, _ = run_contract( + client, llm, prepared(llm, stop_token_ids=[stop_token]), prompt) + assert assembler.token_ids == list(legacy_output.token_ids) + assert assembler.text == legacy_output.text + assert assembler.finish_reason == legacy_output.finish_reason == "stop" + terminal = next(f for f in frames if isinstance(f, Terminal)) + assert terminal.stop_reason == stop_token + REPORT["cases"]["stop_token"] = { + "stop_token": stop_token, "tokens": len(assembler.token_ids) + } + + def test_stop_string_crossing_token_boundary(self, llm, client): + prompt = prompt_ids(llm) + reference, _ = run_legacy(llm, prompt, SamplingParams(**greedy_params())) + pieces = [llm.tokenizer.decode([t]) for t in reference.token_ids[:6]] + if len(pieces) < 3 or not pieces[1].strip() or not pieces[2].strip(): + pytest.skip("reference tokens unsuitable for a boundary-crossing stop") + # A stop string spanning the boundary between tokens 1 and 2. + stop_string = (pieces[1] + pieces[2]).strip() + if not stop_string or stop_string not in reference.text: + pytest.skip("derived stop string not present in reference text") + legacy_output, _ = run_legacy( + llm, prompt, SamplingParams(**greedy_params(stop=[stop_string]))) + frames, assembler, _ = run_contract( + client, llm, prepared(llm, stop=[stop_string]), prompt) + assert assembler.text == legacy_output.text + assert assembler.finish_reason == legacy_output.finish_reason + REPORT["cases"]["stop_string_boundary"] = {"stop": stop_string} + + def test_logprobs(self, llm, client): + prompt = prompt_ids(llm) + legacy_output, _ = run_legacy( + llm, prompt, SamplingParams(**greedy_params(logprobs=0))) + frames, assembler, _ = run_contract( + client, llm, prepared(llm, logprobs=0), prompt) + legacy_values = [ + entry if not isinstance(entry, dict) else + entry[token].logprob for entry, token in zip( + legacy_output.logprobs, legacy_output.token_ids) + ] + assert len(assembler.logprobs) == len(legacy_values) + for ours, theirs in zip(assembler.logprobs, legacy_values): + assert ours == pytest.approx(theirs, abs=1e-5) + REPORT["cases"]["logprobs"] = {"positions": len(assembler.logprobs)} + + def test_abort_mid_stream(self, llm, client): + prompt = prompt_ids(llm, "Write a very long story about a dragon:") + params = prepared(llm, max_tokens=256) + frames, assembler, elapsed = run_contract(client, llm, params, prompt, + stop_after=2) + terminal = next(f for f in frames if isinstance(f, Terminal)) + complete = frames[-1] + assert terminal.finish_reason == "abort" + assert isinstance(complete, RequestComplete) + assert complete.status == "aborted" + assert assembler.finish_reason == "cancelled" + REPORT["cases"]["abort_mid_stream"] = {"elapsed_s": elapsed} + + def test_stream_close_issues_runtime_abort(self, llm, client): + prompt = prompt_ids(llm, "Write a very long story about a robot:") + params = prepared(llm, max_tokens=256) + # Observe the cancellation itself: spy on the router's abort hook. + aborted = [] + router = client._router + original_abort = router._abort_fn + + def spy(client_id): + aborted.append(client_id) + original_abort(client_id) + + router._abort_fn = spy + try: + frames, assembler, _ = run_contract(client, llm, params, prompt, + close_early=True) + finally: + router._abort_fn = original_abort + assert any(isinstance(f, TokenDelta) for f in frames) + assert aborted, "early close did not issue a runtime cancellation" + # The engine stays healthy and serviceable after the early close. + assert client.health().healthy + legacy_output, _ = run_legacy(llm, prompt_ids(llm), + SamplingParams(**greedy_params(max_tokens=8))) + assert legacy_output.token_ids + REPORT["cases"]["stream_close"] = {"runtime_abort_observed": True} + + def test_prompt_logprobs(self, llm, client): + prompt = prompt_ids(llm) + legacy_output, _ = run_legacy( + llm, prompt, + SamplingParams(**greedy_params(logprobs=0, prompt_logprobs=0))) + frames, assembler, _ = run_contract( + client, llm, prepared(llm, logprobs=0, prompt_logprobs=0), prompt) + legacy_prompt_logprobs = legacy_output.prompt_logprobs + assert assembler.prompt_logprobs, "no prompt logprobs delivered" + assert legacy_prompt_logprobs, "legacy produced no prompt logprobs" + assert len(assembler.prompt_logprobs) == len(legacy_prompt_logprobs) + expected_keys = list(prompt[1:]) + [legacy_output.token_ids[0]] + for position, (ours, theirs) in enumerate( + zip(assembler.prompt_logprobs, legacy_prompt_logprobs)): + if isinstance(theirs, dict): + theirs = theirs[expected_keys[position]].logprob + assert ours == pytest.approx(theirs, abs=1e-4) + REPORT["cases"]["prompt_logprobs"] = { + "positions": len(assembler.prompt_logprobs) + } + + def test_usage_and_report(self, llm, client): + prompt = prompt_ids(llm) + frames, assembler, _ = run_contract(client, llm, + prepared(llm, max_tokens=8), prompt) + complete = frames[-1] + assert isinstance(complete, RequestComplete) + assert complete.prompt_tokens == len(prompt) + assert complete.completion_tokens == len(assembler.token_ids) + REPORT["cases"]["usage"] = { + "prompt_tokens": complete.prompt_tokens, + "completion_tokens": complete.completion_tokens, + } + + +def _median(values): + ordered = sorted(values) + mid = len(ordered) // 2 + if len(ordered) % 2: + return ordered[mid] + return 0.5 * (ordered[mid - 1] + ordered[mid]) + + +def _bootstrap_ci_of_median(deltas, resamples=2000, seed=20260713): + """Percentile-bootstrap 95% CI of the median paired delta + (E5B_CUTOVER_DESIGN.md section 7).""" + import random + rng = random.Random(seed) + n = len(deltas) + medians = sorted( + _median([deltas[rng.randrange(n)] for _ in range(n)]) + for _ in range(resamples)) + lower = medians[int(0.025 * resamples)] + upper = medians[min(int(0.975 * resamples), resamples - 1)] + return lower, upper + + +def _write_measurement_table(rows) -> str: + """Generate the exit-report table FROM the artifact values so the report + cannot diverge from the measurement (design doc section 7).""" + lines = [ + "# Rank-0 CPU per streamed token: legacy vs contract", + "", + "Generated from the values recorded in the GPU run artifact " + "(`rank0_measurement` case). Units: microseconds of rank-0 process " + "CPU per streamed token.", + "", + "| concurrency | pairs | legacy median | contract median | Δ median " + "| Δ% | 95% CI of Δ (bootstrap) |", + "|---|---|---|---|---|---|---|", + ] + for row in rows: + scale = 1e6 + lines.append( + f"| {row['concurrency']} | {row['pairs']} " + f"| {row['legacy_median_s_per_token'] * scale:.2f} " + f"| {row['contract_median_s_per_token'] * scale:.2f} " + f"| {row['delta_median_s_per_token'] * scale:+.2f} " + f"| {row['delta_pct']:+.2f}% " + f"| [{row['delta_ci95_s_per_token'][0] * scale:+.2f}, " + f"{row['delta_ci95_s_per_token'][1] * scale:+.2f}] |") + lines.append("") + path = os.path.join(os.path.dirname(REPORT_PATH) or ".", + "rank0_measurement.md") + with open(path, "w") as f: + f.write("\n".join(lines)) + return path + + +@requires_gpu +class TestRank0Measurement: + """Paired A/B rank-0 CPU measurement: legacy vs contract streaming. + + Executes the pre-registered protocol of E5B_CUTOVER_DESIGN.md section 7: + concurrencies {1, 8, 32}, 2 discarded warm-ups per arm per concurrency, + 10 alternating legacy/contract pairs, paired median delta with a + percentile-bootstrap 95% CI (2000 resamples, fixed seed). Records raw + per-rep values into the artifact and generates the report table from + those recorded values. The gate rule ("no regression or an explicitly + accepted, written delta", DEC-3) is applied at the exit report against + this table. + """ + + CONCURRENCIES = (1, 8, 32) + WARMUPS = 2 + # The pre-registered protocol requires AT LEAST 10 pairs; this run uses + # 30 for a tighter CI (the first 10-pair run left the c=8 upper bound + # +0.31% above zero — both runs are disclosed in the exit report). + PAIRS = 30 + MAX_TOKENS = 64 + + def _run_legacy(self, llm, prompts): + return legacy_cpu_run(llm, prompts, self.MAX_TOKENS) + + def _run_contract(self, llm, client, prompts): + async def one(index, prompt): + params = prepared(llm, max_tokens=self.MAX_TOKENS) + engine_request, output_config = convert_request( + f"gpu-m-{uuid.uuid4().hex[:8]}", prompt, params, streaming=True) + request_id = client.submit(engine_request, + output_config=output_config) + stream = client.stream(request_id) + tokens = 0 + async for frame in stream: + if isinstance(frame, TokenDelta): + tokens += len(frame.new_token_ids) + if isinstance(frame, (RequestComplete, ErrorFrame)): + break + return tokens + + async def run(): + return sum(await asyncio.gather( + *(one(i, p) for i, p in enumerate(prompts)))) + + cpu_before = time.process_time() + tokens = asyncio.run(run()) + return time.process_time() - cpu_before, tokens + + def _measure_concurrency(self, llm, client, concurrency): + prompts = [ + prompt_ids(llm, f"Write a story number {i} about a ship:") + for i in range(concurrency) + ] + for _ in range(self.WARMUPS): + self._run_legacy(llm, prompts) + self._run_contract(llm, client, prompts) + + legacy_costs, contract_costs = [], [] + for pair in range(self.PAIRS): + order = ("legacy", "contract") if pair % 2 == 0 else ("contract", + "legacy") + for mode in order: + if mode == "legacy": + cpu, tokens = self._run_legacy(llm, prompts) + legacy_costs.append(cpu / max(tokens, 1)) + else: + cpu, tokens = self._run_contract(llm, client, prompts) + contract_costs.append(cpu / max(tokens, 1)) + + deltas = [c - l for c, l in zip(contract_costs, legacy_costs)] + legacy_median = _median(legacy_costs) + ci_low, ci_high = _bootstrap_ci_of_median(deltas) + return { + "concurrency": concurrency, + "pairs": self.PAIRS, + "warmups": self.WARMUPS, + "max_tokens": self.MAX_TOKENS, + "legacy_cpu_s_per_token": legacy_costs, + "contract_cpu_s_per_token": contract_costs, + "legacy_median_s_per_token": legacy_median, + "contract_median_s_per_token": _median(contract_costs), + "delta_median_s_per_token": _median(deltas), + "delta_pct": 100.0 * _median(deltas) / legacy_median, + "delta_ci95_s_per_token": [ci_low, ci_high], + "bootstrap": {"resamples": 2000, "seed": 20260713}, + } + + def test_rank0_cpu_per_token(self, llm, client): + rows = [ + self._measure_concurrency(llm, client, concurrency) + for concurrency in self.CONCURRENCIES + ] + table_path = _write_measurement_table(rows) + REPORT["cases"]["rank0_measurement"] = { + "protocol": "E5B_CUTOVER_DESIGN.md section 7 (pre-registered)", + "table": table_path, + "rows": rows, + } + for row in rows: + assert len(row["legacy_cpu_s_per_token"]) == self.PAIRS + assert len(row["contract_cpu_s_per_token"]) == self.PAIRS + assert row["legacy_median_s_per_token"] > 0 + assert row["contract_median_s_per_token"] > 0 + # Record-only for the delta itself; the acceptance decision (no + # regression or an explicit written acceptance, DEC-3) is applied at + # the exit report against the generated table. + + +@requires_gpu +class TestWorkerCrash: + """Kills a worker mid-stream; uses its own LLM instance (destructive). + + Verifies the REAL crash detector first: after SIGKILL the test waits up + to ``NATIVE_DETECTION_DEADLINE`` for the executor's own machinery + (mpi_done_callback -> _error_queue -> monitor -> pre_shutdown -> typed + poisoning) to end the stream without any test interference. Only if the + MPI runtime never surfaces the death in this environment does the test + fall back to injecting the error into ``_error_queue`` — and it records + which mode ran, with the observed native-signal evidence, in the + artifact so the exit report carries the justification. + """ + + NATIVE_DETECTION_DEADLINE = 120.0 + + def test_worker_crash_yields_typed_ending_within_bound(self): + psutil = pytest.importorskip("psutil") + me = psutil.Process() + before = {p.pid for p in me.children(recursive=True)} + instance = LLM(model=model_path(), + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.3, + enable_block_reuse=False), + max_batch_size=8, + max_seq_len=1024) + try: + if not isinstance(instance._executor, GenerationExecutorProxy): + pytest.skip("executor is not the IPC proxy") + workers = [ + p for p in me.children(recursive=True) if p.pid not in before + ] + assert workers, "no worker processes found to crash" + config = EngineClientConfig(backend=instance.args.backend or "pytorch", + world_size=1, flag_enabled=True) + engine_client = LocalProcessEngineClient(instance._executor, config) + params = prepare_sampling_params( + SamplingParams(**greedy_params(max_tokens=512)), + tokenizer=instance.tokenizer, + hf_model_config=instance._hf_model_config, + generation_config=instance._generation_config) + engine_request, output_config = convert_request( + "gpu-crash", prompt_ids(instance, "Tell me a very long story:"), + params, streaming=True) + executor = instance._executor + + def native_signal_seen() -> bool: + futures = getattr(executor, "mpi_futures", None) or [] + if any(fut.done() for fut in futures): + return True + queue = getattr(executor, "_error_queue", None) + try: + return queue is not None and not queue.empty() + except Exception: + return False + + async def drain_to_ending(stream, per_frame_deadline): + tail = [] + while True: + try: + frame = await asyncio.wait_for(stream.__anext__(), + per_frame_deadline) + except StopAsyncIteration: + break + tail.append(frame) + if isinstance(frame, (RequestComplete, ErrorFrame)): + break + return tail + + async def scenario(): + engine_client.submit(engine_request, output_config=output_config) + stream = engine_client.stream("gpu-crash") + first = await asyncio.wait_for(stream.__anext__(), STREAM_DEADLINE) + assert isinstance(first, TokenDelta) + for worker in workers: + try: + worker.kill() + except psutil.NoSuchProcess: + pass + killed_at = time.monotonic() + + # Phase 1: give the executor's own detector the full window. + signal_at = None + frames = [first] + ended = False + while time.monotonic() - killed_at < self.NATIVE_DETECTION_DEADLINE: + if signal_at is None and native_signal_seen(): + signal_at = time.monotonic() + try: + frame = await asyncio.wait_for(stream.__anext__(), 1.0) + except asyncio.TimeoutError: + continue + except StopAsyncIteration: + break + frames.append(frame) + if isinstance(frame, (RequestComplete, ErrorFrame)): + ended = True + break + if signal_at is None and native_signal_seen(): + signal_at = time.monotonic() + + case = { + "native_signal_observed": signal_at is not None, + "native_window_s": self.NATIVE_DETECTION_DEADLINE, + } + if ended: + case["mode"] = "native" + case["detection_s"] = time.monotonic() - killed_at + if signal_at is not None: + case["signal_to_typed_end_s"] = ( + time.monotonic() - signal_at) + elif signal_at is not None: + # The detector saw the death but poisoning never reached + # the stream: that is a genuine bound violation. + pytest.fail( + "crash signal reached the error queue/futures but the " + f"stream stayed open past the " + f"{self.NATIVE_DETECTION_DEADLINE:.0f}s window") + else: + # The MPI runtime never surfaced the SIGKILL here. Record + # that evidence, then verify the ownership machinery + # (error-queue -> monitor -> pre_shutdown -> typed + # poisoning) at the frozen bound via injection. + case["mode"] = "injected_fallback" + case["justification"] = ( + "no mpi_futures completion and empty _error_queue for " + f"{self.NATIVE_DETECTION_DEADLINE:.0f}s after SIGKILL " + "in this environment (confirmed by a dedicated 300s " + "probe: worker observed dead at ~1s, futures never " + "complete — an mpi4py-layer propagation gap that " + "affects the legacy path identically); injected into " + "_error_queue to verify the executor-owned failure " + "path at the frozen DEC-4 bound") + executor._error_queue.put_nowait( + RuntimeError("worker killed by crash test")) + injected_at = time.monotonic() + frames.extend(await drain_to_ending(stream, + FAILURE_DEADLINE)) + ended = isinstance(frames[-1], (RequestComplete, ErrorFrame)) + case["detection_s"] = time.monotonic() - injected_at + + assert ended and isinstance( + frames[-1], (RequestComplete, ErrorFrame)), ( + f"stream did not end typed after worker crash: {frames[-1]!r}") + if isinstance(frames[-1], RequestComplete): + assert frames[-1].status == "failed" + if case["mode"] == "injected_fallback": + assert case["detection_s"] <= FAILURE_DEADLINE + merge_report({"cases": {"worker_crash": case}}) + + asyncio.run(scenario()) + finally: + # The workers are dead by design; a graceful shutdown can block + # forever on the MPI session. Bound it, then reap leftovers. + import threading as _threading + + def _shutdown(): + try: + instance.shutdown() + except Exception: + pass + + shutdown_thread = _threading.Thread(target=_shutdown, daemon=True) + shutdown_thread.start() + shutdown_thread.join(30) + for process in me.children(recursive=True): + if process.pid not in before: + try: + process.kill() + except Exception: + pass diff --git a/tests/unittest/executor/engine_client/test_proxy_hooks.py b/tests/unittest/executor/engine_client/test_proxy_hooks.py new file mode 100644 index 000000000000..911576d93472 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_proxy_hooks.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Proxy hook and worker prompt-logprob-guard tests. + +These tests exercise the engine-client hooks on a skeleton +GenerationExecutorProxy (no MPI workers) to prove the hooks are inert with +no router attached, observe every raw item before the legacy drop logic, +and contain router failures. +""" + +from queue import Queue +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.executor.base_worker import _compute_pytorch_prompt_logprobs +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.utils import ErrorResponse + + +class RecordingRouter: + + def __init__(self, raise_on_response=False): + self.observed_submits = [] + self.observed_responses = [] + self.failed_submits = [] + self.fail_all_calls = [] + self.raise_on_response = raise_on_response + + def observe_submit(self, request): + self.observed_submits.append(request) + + def on_response(self, response): + self.observed_responses.append(response) + if self.raise_on_response: + raise RuntimeError("router bug") + + def on_submit_enqueue_failed(self, client_id): + self.failed_submits.append(client_id) + + def fail_all(self, reason): + self.fail_all_calls.append(reason) + + +class FakeResultQueue: + + def __init__(self, items): + self._items = list(items) + + def get(self): + return self._items.pop(0) if self._items else None + + +def make_skeleton_proxy() -> GenerationExecutorProxy: + proxy = object.__new__(GenerationExecutorProxy) + proxy._results = {} + return proxy + + +def make_legacy_result() -> SimpleNamespace: + return SimpleNamespace(queue=Queue()) + + +class TestDispatchHook: + + def run_dispatch(self, proxy, items): + proxy.result_queue = FakeResultQueue([items]) + return proxy.dispatch_result_task() + + def test_inert_without_router(self): + proxy = make_skeleton_proxy() + result = make_legacy_result() + proxy._results = {3: result} + response = ErrorResponse(client_id=3, error_msg="x", request_id=1) + assert self.run_dispatch(proxy, [response]) is True + # Legacy behavior: delivered then popped on final/error. + assert result.queue.get_nowait() is response + assert 3 not in proxy._results + + def test_router_sees_items_before_legacy_drop(self): + proxy = make_skeleton_proxy() + router = RecordingRouter() + proxy.attach_engine_frame_router(router) + # Unknown client id: the legacy path drops it silently, the router + # must still observe it (late/duplicate frame absorption). + unknown = ErrorResponse(client_id=99, error_msg="late", request_id=1) + assert self.run_dispatch(proxy, [unknown]) is True + assert router.observed_responses == [unknown] + + def test_router_sees_batched_list_items(self): + proxy = make_skeleton_proxy() + router = RecordingRouter() + proxy.attach_engine_frame_router(router) + items = [ErrorResponse(client_id=1, error_msg="a", request_id=1), + ErrorResponse(client_id=2, error_msg="b", request_id=2)] + assert self.run_dispatch(proxy, items) is True + assert router.observed_responses == items + + def test_router_exception_does_not_break_dispatch(self): + proxy = make_skeleton_proxy() + router = RecordingRouter(raise_on_response=True) + proxy.attach_engine_frame_router(router) + result = make_legacy_result() + proxy._results = {3: result} + response = ErrorResponse(client_id=3, error_msg="x", request_id=1) + assert self.run_dispatch(proxy, [response]) is True + # Legacy delivery still happened despite the router raising. + assert result.queue.get_nowait() is response + + +class TestSubmitHook: + + class _Stub: + """Provides bound methods (WeakMethod-compatible) for the skeleton.""" + + def handle_background_error(self, *args, **kwargs): + return None + + def start_dispatch_threads(self): + return None + + def get_logprob_params(self, request): + return None + + def make_submit_proxy(self, put_raises=False): + proxy = make_skeleton_proxy() + proxy._last_client_id = 0 + proxy.dispatch_result_thread = object() # skip thread start + proxy._error_queue = Queue() + proxy.doing_shutdown = False + stub = self._Stub() + proxy._stub_ref = stub # keep alive for WeakMethod + proxy._handle_background_error = stub.handle_background_error + proxy._start_dispatch_threads = stub.start_dispatch_threads + proxy._get_logprob_params = stub.get_logprob_params + + class FakeRequestQueue: + + def __init__(self): + self.items = [] + + def put(self, item): + if put_raises: + raise RuntimeError("zmq down") + self.items.append(item) + + proxy.request_queue = FakeRequestQueue() + return proxy + + def make_request(self): + from tensorrt_llm.executor.request import GenerationRequest + from tensorrt_llm.sampling_params import SamplingParams + return GenerationRequest(prompt_token_ids=[1, 2], + sampling_params=SamplingParams(max_tokens=4, end_id=2), + streaming=True) + + def test_observe_submit_binds_before_enqueue(self): + proxy = self.make_submit_proxy() + router = RecordingRouter() + proxy.attach_engine_frame_router(router) + request = self.make_request() + proxy.submit(request) + assert router.observed_submits == [request] + assert request.id is not None + assert proxy.request_queue.items == [request] + + def test_enqueue_failure_rolls_back(self): + proxy = self.make_submit_proxy(put_raises=True) + router = RecordingRouter() + proxy.attach_engine_frame_router(router) + request = self.make_request() + with pytest.raises(RuntimeError): + proxy.submit(request) + assert router.failed_submits == [request.id] + + def test_submit_without_router_unchanged(self): + proxy = self.make_submit_proxy() + request = self.make_request() + result = proxy.submit(request) + assert proxy.request_queue.items == [request] + assert proxy._results[request.id] is result + + +class TestPromptLogprobGuard: + + def make_generation_result(self, prompt_ids=(1, 2, 3)): + return SimpleNamespace( + _logprob_params=SimpleNamespace(prompt_logprobs=0, + prompt_logprobs_simple_format=True), + _streaming=True, + _generation_request=SimpleNamespace(prompt_token_ids=list(prompt_ids))) + + def make_response(self, output_token_ids, context_logits): + result = SimpleNamespace(output_token_ids=output_token_ids) + inner = SimpleNamespace(context_logits=context_logits, + get_result=lambda: result) + return SimpleNamespace(result=inner, client_id=5) + + def test_no_token_final_degrades_gracefully(self): + generation_result = self.make_generation_result() + response = self.make_response(output_token_ids=[[]], context_logits=None) + assert _compute_pytorch_prompt_logprobs(generation_result, response) is None + + def test_missing_context_logits_degrades_gracefully(self): + generation_result = self.make_generation_result() + response = self.make_response(output_token_ids=[[5]], context_logits=None) + assert _compute_pytorch_prompt_logprobs(generation_result, response) is None + + def test_cached_prompt_logprobs_short_circuit_still_works(self): + generation_result = self.make_generation_result() + generation_result._cached_prompt_logprobs = [-1.0] + response = self.make_response(output_token_ids=[[]], context_logits=None) + logprobs_result = _compute_pytorch_prompt_logprobs(generation_result, response) + assert logprobs_result is not None + assert logprobs_result.prompt == [-1.0] diff --git a/tests/unittest/executor/engine_client/test_replay_oracle.py b/tests/unittest/executor/engine_client/test_replay_oracle.py new file mode 100644 index 000000000000..41c3345bd82f --- /dev/null +++ b/tests/unittest/executor/engine_client/test_replay_oracle.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Golden replay oracle: fixed response streams through both pipelines. + +The legacy side runs the REAL legacy consumer code path — +``DetokenizedGenerationResultBase._handle_response`` with the real +detokenize/trim/stop semantics — never a re-implementation in test code. +The new side runs envelope → router → frames → assembler. Parity fields: +token ids, text, finish reason, stop reason, logprob values, and usage +counts, including the presentation renderings (``timeout``, ``cancelled``). +A deliberately perturbed stream must fail parity, proving the oracle +discriminates. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.engine_client.assembler import FrontendResponseAssembler +from tensorrt_llm.executor.engine_client.contract import ( + ContractConstructionError, FrontendOutputConfig) +from tensorrt_llm.executor.engine_client.router import EngineFrameRouter +from tensorrt_llm.executor.result import (DetokenizedGenerationResultBase, + Logprob) +from tensorrt_llm.sampling_params import SamplingParams + +PROMPT_IDS = (1, 2, 3) + + +class CharTokenizer: + """Deterministic shared tokenizer: id N ↔ chr(97 + N % 26); id 900+ maps + to unicode fragments so multi-char/multi-byte boundaries are exercised.""" + + UNICODE = {900: "é", 901: "☂", 902: "é☂"} + + def _piece(self, token_id: int) -> str: + if token_id in self.UNICODE: + return self.UNICODE[token_id] + return chr(97 + token_id % 26) + + def decode(self, ids, **kwargs): + return "".join(self._piece(i) for i in ids) + + def decode_incrementally(self, token_ids, prev_text=None, states=None, *, + flush=False, stream_interval=1, **kwargs): + prev_text = prev_text or "" + return prev_text + self.decode(token_ids), states or {} + + +TOKENIZER = CharTokenizer() + + +def ids_for(text: str) -> list: + special = {v: k for k, v in CharTokenizer.UNICODE.items()} + out = [] + for ch in text: + out.append(special.get(ch, ord(ch) - 97)) + return out + + +def make_result(**overrides) -> SimpleNamespace: + fields = dict(is_final=False, output_token_ids=[[]], finish_reasons=None, + log_probs=None, cum_log_probs=None, sequence_index=0, + context_phase_params=None, decoding_iter=1, + avg_decoded_tokens_per_iter=None, request_perf_metrics=None, + generation_logits=None, context_logits=None) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def delta(tokens, log_probs=None): + return LlmResponse(request_id=1, client_id=7, + result=make_result( + output_token_ids=[list(tokens)], + finish_reasons=[tllm.FinishReason.NOT_FINISHED], + log_probs=log_probs)) + + +def final(reason, tokens=(), log_probs=None): + return LlmResponse(request_id=1, client_id=7, + result=make_result(is_final=True, + output_token_ids=[list(tokens)], + finish_reasons=[reason], + log_probs=log_probs)) + + +def make_sampling_params(stop=None, stop_token_ids=None) -> SamplingParams: + params = SamplingParams(max_tokens=64, end_id=2, pad_id=0, stop=stop, + stop_token_ids=stop_token_ids) + if stop is not None: + strings = [stop] if isinstance(stop, str) else list(stop) + params._stop_word_ids = [ids_for(s) for s in strings] + params._stream_interval = 1 + return params + + +def stop_association(params: SamplingParams): + return tuple( + (tuple(ids), reason) + for reason, ids in params._get_stop_reasons_and_words()) + + +def run_legacy(responses, params: SamplingParams) -> dict: + result = DetokenizedGenerationResultBase(id=7, sampling_params=params, + tokenizer=TOKENIZER, streaming=True) + result._streaming = True + for response in responses: + result._handle_response(response) + if result._done: + break + output = result._outputs[0] + logprob_values = [ + entry if not isinstance(entry, dict) else + max(entry.values(), key=lambda lp: lp.rank is not None and -lp.rank + if lp.rank else 0).logprob if False else None + for entry in output.logprobs + ] + # Chosen-token projection for dict-shaped entries: look up by token id. + projected = [] + for position, entry in enumerate(output.logprobs): + if isinstance(entry, dict): + token_id = output.token_ids[position] + projected.append(entry[token_id].logprob) + else: + projected.append(entry) + return { + "token_ids": list(output.token_ids), + "text": output.text, + "finish_reason": output.finish_reason, + "stop_reason": output.stop_reason, + "logprobs": projected, + "usage_completion": len(output.token_ids), + } + + +def run_new(responses, params: SamplingParams) -> dict: + association = stop_association(params) + strings = ([params.stop] if isinstance(params.stop, str) else + list(params.stop or ())) + config = FrontendOutputConfig( + stop_strings=tuple(strings), + include_stop_str_in_output=params.include_stop_str_in_output, + stop_sequence_reasons=association, end_id=params.end_id, + num_logprobs=params.logprobs) + aborted = [] + router = EngineFrameRouter(abort_fn=aborted.append) + request = SimpleNamespace(id=None) + request.set_id = lambda v: setattr(request, "id", v) or request + binding = router.register_pending(request, "req-oracle", PROMPT_IDS, + association) + request.set_id(7) + router.observe_submit(request) + assembler = FrontendResponseAssembler("req-oracle", config, + tokenizer=TOKENIZER, + abort_callback=lambda rid: aborted.append(rid)) + + for response in responses: + router.on_response(response) + # Drain the frames this response produced and process them as one + # batch (the SSE layer drains-then-renders the same way). + batch = [] + while True: + frame, ready = binding.delivery.pop_nowait() + if not ready: + break + batch.append(frame) + assembler.process_frames(batch) + if assembler._stopped_by_string: + break + return { + "token_ids": list(assembler.token_ids), + "text": assembler.text, + "finish_reason": assembler.finish_reason, + "stop_reason": assembler.stop_reason, + "logprobs": list(assembler.logprobs), + "usage_completion": (assembler.usage or {}).get( + "completion_tokens", len(assembler.token_ids)), + } + + +PARITY_FIELDS = ("token_ids", "text", "finish_reason", "stop_reason", "logprobs", + "usage_completion") + + +def assert_parity(responses, params, fields=PARITY_FIELDS): + legacy = run_legacy(responses, params) + new = run_new(responses, params) + for field in fields: + assert new[field] == legacy[field], ( + f"parity mismatch on {field!r}: new={new[field]!r} " + f"legacy={legacy[field]!r}") + return legacy, new + + +class TestParityCorpus: + + def test_plain_streaming(self): + responses = [delta(ids_for("hel")), delta(ids_for("lo")), + final(tllm.FinishReason.END_ID)] + assert_parity(responses, make_sampling_params()) + + def test_final_carrying_tokens(self): + responses = [delta(ids_for("hi")), + final(tllm.FinishReason.LENGTH, tokens=ids_for("!"[0:0] or "z"))] + assert_parity(responses, make_sampling_params()) + + def test_stop_token_id(self): + params = make_sampling_params(stop_token_ids=[13]) + responses = [delta(ids_for("ab")), + final(tllm.FinishReason.STOP_WORDS, tokens=[13])] + assert_parity(responses, params) + + def test_engine_stop_string_cross_token_boundary(self): + # Stop string "xy" tokenizes to two tokens; the engine stops on the + # sequence and both pipelines trim it. + params = make_sampling_params(stop=["xy"]) + responses = [delta(ids_for("ab")), + final(tllm.FinishReason.STOP_WORDS, tokens=ids_for("xy"))] + assert_parity(responses, params) + + def test_frontend_stop_string_detection(self): + # The engine does not stop; the frontend detects the string in the + # decoded text mid-stream (legacy scans cumulative text each step). + params = make_sampling_params(stop=["xy"]) + responses = [delta(ids_for("ax")), delta(ids_for("yz")), + final(tllm.FinishReason.LENGTH)] + legacy, new = assert_parity(responses, params) + assert legacy["finish_reason"] == "stop" + assert legacy["stop_reason"] == "xy" + + def test_duplicate_and_colliding_stop_strings(self): + params = make_sampling_params(stop=["xy", "xy", "y"]) + responses = [delta(ids_for("axyb")), final(tllm.FinishReason.LENGTH)] + assert_parity(responses, params) + + def test_unicode_stop_string(self): + params = make_sampling_params(stop=["é☂"]) + responses = [delta([900]), delta([901]), + final(tllm.FinishReason.LENGTH)] + assert_parity(responses, params) + + def test_float_logprobs(self): + responses = [delta(ids_for("ab"), log_probs=[[-0.5, -1.0]]), + delta(ids_for("c"), log_probs=[[-0.25]]), + final(tllm.FinishReason.END_ID)] + assert_parity(responses, make_sampling_params()) + + def test_dict_logprobs_project_to_same_values(self): + entries1 = [{ids_for("a")[0]: Logprob(-0.5, 1)}, + {ids_for("b")[0]: Logprob(-1.0, 2)}] + responses = [delta(ids_for("ab"), log_probs=[entries1]), + final(tllm.FinishReason.END_ID)] + assert_parity(responses, make_sampling_params()) + + def test_timeout_rendering(self): + responses = [delta(ids_for("a")), final(tllm.FinishReason.TIMED_OUT)] + legacy, new = assert_parity(responses, make_sampling_params(), + fields=("token_ids", "text", "finish_reason")) + assert legacy["finish_reason"] == "timeout" + + def test_cancelled_rendering(self): + responses = [delta(ids_for("a")), final(tllm.FinishReason.CANCELLED)] + legacy, new = assert_parity(responses, make_sampling_params(), + fields=("token_ids", "text", "finish_reason")) + assert legacy["finish_reason"] == "cancelled" + + def test_empty_stop_string_rejected_at_construction(self): + with pytest.raises(ContractConstructionError): + FrontendOutputConfig(stop_strings=("", )) + + +class TestOracleDiscriminates: + + def test_perturbed_stream_fails_parity(self): + params = make_sampling_params() + responses = [delta(ids_for("hel")), delta(ids_for("lo")), + final(tllm.FinishReason.END_ID)] + legacy = run_legacy(responses, params) + perturbed = [delta(ids_for("hel")), delta(ids_for("lQ".lower())), + final(tllm.FinishReason.END_ID)] + new = run_new(perturbed, params) + assert new["text"] != legacy["text"] + + def test_dropped_terminal_detected(self): + params = make_sampling_params() + complete_responses = [delta(ids_for("hi")), final(tllm.FinishReason.END_ID)] + truncated = [delta(ids_for("hi"))] + legacy = run_legacy(complete_responses, params) + new = run_new(truncated, params) + assert new["finish_reason"] != legacy["finish_reason"] diff --git a/tests/unittest/executor/engine_client/test_router_client.py b/tests/unittest/executor/engine_client/test_router_client.py new file mode 100644 index 000000000000..4bc60f85ecdb --- /dev/null +++ b/tests/unittest/executor/engine_client/test_router_client.py @@ -0,0 +1,715 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Replay-matrix tests for the exactly-once router and the local client.""" + +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.engine_client.contract import (EngineRequest, + EngineSamplingConfig, + ErrorFrame, + FrontendOutputConfig, + GuidedDecodingSpec, + RequestComplete, + Terminal, TokenDelta) +from tensorrt_llm.executor.engine_client.invariants import ( + InvariantCheckingStream, StreamInvariantViolation) +from tensorrt_llm.executor.engine_client.local_client import ( + EngineClientConfig, EngineClientConfigError, LocalProcessEngineClient, + RequestRejectedError) +from tensorrt_llm.executor.engine_client.router import (DEFAULT_TOMBSTONE_LIMIT, + EngineFrameRouter, + UnknownRequestError) +from tensorrt_llm.executor.result import LogProbsResult, ResponseWrapper +from tensorrt_llm.executor.utils import ErrorResponse + +STREAM_DEADLINE = 5.0 + + +class FakeRequestQueue: + + def __init__(self): + self.items = [] + self.raise_on_put = False + + def put(self, item): + if self.raise_on_put: + raise RuntimeError("enqueue failed after binding") + self.items.append(item) + + +class FakeExecutor: + """Mimics the proxy's contract-native surface for the client/service.""" + + def __init__(self): + self.service = None + self.router = None + self.aborted = [] + self.request_queue = FakeRequestQueue() + self._results = {} + self.doing_shutdown = False + self._fatal_error = None + self._submission_lock = threading.RLock() + self._next_id = 0 + self.shutdown_called = False + self.raise_on_submit = False # fails before any binding exists + + @property + def submitted(self): + return self.request_queue.items + + @property + def raise_after_observe(self): + return self.request_queue.raise_on_put + + @raise_after_observe.setter + def raise_after_observe(self, value): + self.request_queue.raise_on_put = value + + def attach_engine_service(self, service): + if self.service is not None: + raise RuntimeError("service already attached") + self.service = service + self.router = service.router + + def attach_engine_frame_router(self, router): + self.router = router + + def submit_contract(self, engine_request, stop_reasons=()): + return self.service.submit_contract(engine_request, + stop_reasons=stop_reasons) + + def _start_dispatch_threads(self): + if self.raise_on_submit: + raise RuntimeError("submit failed before binding") + + def _get_next_client_id(self): + self._next_id += 1 + return self._next_id + + def _handle_background_error(self): + pass + + def abort_request(self, client_id): + self.aborted.append(client_id) + + def get_stats(self, timeout): + return [{"iteration": 1}] + + def get_kv_events(self, timeout): + return ['{"event": 1}'] + + def check_health(self): + return True + + def shutdown(self): + self.shutdown_called = True + + +def make_config(**overrides) -> EngineClientConfig: + kwargs = dict(backend="pytorch", flag_enabled=True) + kwargs.update(overrides) + return EngineClientConfig(**kwargs) + + +def make_client(executor=None) -> LocalProcessEngineClient: + return LocalProcessEngineClient(executor or FakeExecutor(), make_config()) + + +def make_engine_request(request_id="req-1", **sampling_overrides) -> EngineRequest: + sampling_kwargs = dict(max_new_tokens=16, end_id=2, pad_id=0) + sampling_kwargs.update(sampling_overrides) + return EngineRequest(request_id=request_id, prompt_token_ids=(1, 2, 3), + sampling=EngineSamplingConfig(**sampling_kwargs)) + + +def make_result(**overrides) -> SimpleNamespace: + fields = dict(is_final=False, output_token_ids=[[5, 6]], finish_reasons=None, + log_probs=None, cum_log_probs=None, sequence_index=0) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def response_for(client_id, **result_overrides) -> LlmResponse: + return LlmResponse(request_id=client_id, result=make_result(**result_overrides), + client_id=client_id) + + +def delta_response(client_id, tokens, **overrides): + return response_for(client_id, output_token_ids=[list(tokens)], **overrides) + + +def final_response(client_id, reason, tokens=(), **overrides): + return response_for(client_id, is_final=True, output_token_ids=[list(tokens)], + finish_reasons=[reason], **overrides) + + +async def collect_frames(client, request_id, check_invariants=True): + stream = client.stream(request_id) + if check_invariants: + stream = InvariantCheckingStream(stream, request_id) + frames = [] + while True: + try: + frame = await asyncio.wait_for(stream.__anext__(), STREAM_DEADLINE) + except StopAsyncIteration: + break + frames.append(frame) + if isinstance(frame, (RequestComplete, ErrorFrame)): + break + return frames + + +def submitted_client_id(executor) -> int: + return executor.submitted[-1].id + + +class TestPlainStreaming: + + def test_deltas_terminal_complete(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + router = executor.router + router.on_response(delta_response(client_id, [5, 6])) + router.on_response(delta_response(client_id, [7])) + router.on_response(final_response(client_id, tllm.FinishReason.END_ID)) + + frames = asyncio.run(collect_frames(client, "req-1")) + kinds = [type(f).__name__ for f in frames] + assert kinds == ["TokenDelta", "TokenDelta", "Terminal", "RequestComplete"] + assert frames[0].new_token_ids == (5, 6) + assert frames[2].finish_reason == "stop" and frames[2].stop_reason is None + complete = frames[3] + assert complete.status == "ok" + assert complete.prompt_tokens == 3 + assert complete.completion_tokens == 3 + assert [f.event_seq for f in frames] == [0, 1, 2, 3] + + def test_final_with_last_tokens(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + executor.router.on_response( + final_response(client_id, tllm.FinishReason.LENGTH, tokens=[9])) + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames + ] == ["TokenDelta", "Terminal", "RequestComplete"] + assert frames[1].finish_reason == "length" + assert frames[2].completion_tokens == 1 + + +class TestFinishReasonTable: + + def run_final(self, reason, abort_first=False, stop_reasons=(), tokens=()): + executor = FakeExecutor() + client = make_client(executor) + request = make_engine_request() + output_config = FrontendOutputConfig(stop_sequence_reasons=stop_reasons) + client.submit(request, output_config=output_config) + client_id = submitted_client_id(executor) + if abort_first: + client.abort("req-1") + executor.router.on_response(final_response(client_id, reason, tokens=tokens)) + frames = asyncio.run(collect_frames(client, "req-1")) + terminal = next(f for f in frames if isinstance(f, Terminal)) + complete = next(f for f in frames if isinstance(f, RequestComplete)) + return terminal, complete + + def test_end_id(self): + terminal, complete = self.run_final(tllm.FinishReason.END_ID) + assert (terminal.finish_reason, complete.status) == ("stop", "ok") + + def test_stop_words_resolves_ordered_reason(self): + stop_reasons = (((13, ), 13), ((7, 8), "STOP")) + terminal, complete = self.run_final(tllm.FinishReason.STOP_WORDS, + stop_reasons=stop_reasons, + tokens=[5, 7, 8]) + assert terminal.finish_reason == "stop" + assert terminal.stop_reason == "STOP" + assert complete.status == "ok" + + def test_length(self): + terminal, complete = self.run_final(tllm.FinishReason.LENGTH) + assert (terminal.finish_reason, complete.status) == ("length", "ok") + + def test_cancelled(self): + terminal, complete = self.run_final(tllm.FinishReason.CANCELLED) + assert (terminal.finish_reason, complete.status) == ("abort", "aborted") + + def test_timed_out_maps_to_error_timeout(self): + terminal, complete = self.run_final(tllm.FinishReason.TIMED_OUT) + assert terminal.finish_reason == "error" + assert terminal.stop_reason == "timeout" + assert complete.status == "failed" + + def test_not_finished_without_abort(self): + terminal, complete = self.run_final(tllm.FinishReason.NOT_FINISHED) + assert terminal.finish_reason == "error" + assert terminal.stop_reason == "not_finished" + + def test_not_finished_with_requested_abort(self): + terminal, complete = self.run_final(tllm.FinishReason.NOT_FINISHED, + abort_first=True) + assert (terminal.finish_reason, complete.status) == ("abort", "aborted") + + +class TestAbort: + + def test_abort_before_first_token(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + client.abort("req-1") + assert executor.aborted == [client_id] + executor.router.on_response( + final_response(client_id, tllm.FinishReason.CANCELLED)) + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames] == ["Terminal", "RequestComplete"] + assert frames[0].finish_reason == "abort" + assert frames[1].status == "aborted" + assert frames[1].completion_tokens == 0 + + def test_abort_after_completion_is_noop(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + executor.router.on_response( + final_response(client_id, tllm.FinishReason.END_ID)) + client.abort("req-1") # no-op, no error + assert executor.aborted == [] + + def test_abort_unknown_is_typed_error(self): + client = make_client() + with pytest.raises(UnknownRequestError): + client.abort("nope") + + def test_stream_close_aborts_incomplete_request(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + executor.router.on_response(delta_response(client_id, [5])) + + async def scenario(): + stream = client.stream("req-1") + frame = await asyncio.wait_for(stream.__anext__(), STREAM_DEADLINE) + assert isinstance(frame, TokenDelta) + await stream.aclose() + + asyncio.run(scenario()) + assert executor.aborted == [client_id] + + +class TestLateAndDuplicate: + + def test_late_duplicate_finals_absorbed(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + router = executor.router + router.on_response(final_response(client_id, tllm.FinishReason.END_ID)) + router.on_response(final_response(client_id, tllm.FinishReason.END_ID)) + router.on_response(delta_response(client_id, [9])) + assert router.counters["late_or_duplicate_absorbed"] == 2 + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames] == ["Terminal", "RequestComplete"] + + def test_unbound_legacy_traffic_ignored(self): + executor = FakeExecutor() + make_client(executor) + executor.router.on_response( + ErrorResponse(client_id=999, error_msg="legacy", request_id=1)) + assert executor.router.active_request_count() == 0 + + +class TestFailurePaths: + + def test_error_response_before_start_is_standalone_error_frame(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + executor.router.on_response( + ErrorResponse(client_id=client_id, error_msg="admission failed", + request_id=1)) + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames] == ["ErrorFrame"] + assert frames[0].error_code == "request_error" + + def test_error_after_start_ends_terminal_failed(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + router = executor.router + router.on_response(delta_response(client_id, [5])) + router.on_response( + ErrorResponse(client_id=client_id, error_msg="boom", request_id=1)) + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames + ] == ["TokenDelta", "Terminal", "RequestComplete"] + assert frames[1].finish_reason == "error" + assert frames[2].status == "failed" + + def test_fail_all_nothing_started_is_error_frame(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + executor.router.fail_all("worker crashed") + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames] == ["ErrorFrame"] + assert frames[0].error_code == "executor_failed" + + def test_fail_all_started_synthesizes_terminal(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + executor.router.on_response(delta_response(client_id, [5])) + executor.router.fail_all("worker crashed") + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames + ] == ["TokenDelta", "Terminal", "RequestComplete"] + assert frames[2].status == "failed" + assert executor.router.counters["synthesized_terminals"] == 1 + + def test_fatal_error_latch_blocks_new_submits(self): + executor = FakeExecutor() + client = make_client(executor) + executor.router.fail_all("fatal") + assert not client.health().healthy + with pytest.raises(RequestRejectedError): + client.submit(make_engine_request(request_id="req-2")) + + def test_enqueue_failure_after_binding(self): + executor = FakeExecutor() + executor.raise_after_observe = True + client = make_client(executor) + with pytest.raises(RuntimeError): + client.submit(make_engine_request()) + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames] == ["ErrorFrame"] + assert frames[0].error_code == "enqueue_failed" + + def test_submit_failure_before_binding_rolls_back(self): + executor = FakeExecutor() + executor.raise_on_submit = True + client = make_client(executor) + with pytest.raises(RuntimeError): + client.submit(make_engine_request()) + assert executor.router.active_request_count() == 0 + # The id is reusable: the registration was rolled back entirely. + executor.raise_on_submit = False + client.submit(make_engine_request()) + + +class TestStreamLifecycle: + + def test_submit_without_stream_then_delayed_open(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + router = executor.router + router.on_response(delta_response(client_id, [5])) + router.on_response(final_response(client_id, tllm.FinishReason.END_ID)) + # Open only after the request fully ended: nothing is lost. + frames = asyncio.run(collect_frames(client, "req-1")) + assert [type(f).__name__ for f in frames + ] == ["TokenDelta", "Terminal", "RequestComplete"] + + def test_double_open_is_typed_error(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client.stream("req-1") + with pytest.raises(RequestRejectedError): + client.stream("req-1") + + def test_unknown_stream_is_typed_error(self): + client = make_client() + with pytest.raises(RequestRejectedError): + client.stream("ghost") + + def test_request_id_reuse_after_completion_rejected(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + executor.router.on_response( + final_response(client_id, tllm.FinishReason.END_ID)) + with pytest.raises(RequestRejectedError): + client.submit(make_engine_request()) + + def test_duplicate_active_id_rejected(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + with pytest.raises(RequestRejectedError): + client.submit(make_engine_request()) + + +class TestOverflow: + + def test_overflow_ends_typed_via_out_of_band_path(self): + executor = FakeExecutor() + router = EngineFrameRouter(abort_fn=executor.abort_request, + delivery_limit=2) + executor.attach_engine_frame_router(router) + generation_request = SimpleNamespace(id=None) + + def set_id(value): + generation_request.id = value + return generation_request + + generation_request.set_id = set_id + binding = router.register_pending(generation_request, "req-o", (1, 2), ()) + generation_request.set_id(41) + router.observe_submit(generation_request) + for i in range(4): + router.on_response(delta_response(41, [i])) + assert binding.delivery.overflowed or binding.ended + assert router.counters["overflow_aborts"] == 1 + + async def drain(): + frames = [] + stream_binding = router.open_stream_binding("req-o") + while True: + frame, ready = stream_binding.delivery.pop_nowait() + if not ready: + break + frames.append(frame) + return frames + + frames = asyncio.run(drain()) + assert isinstance(frames[-1], RequestComplete) + assert frames[-1].status == "failed" + terminal = next(f for f in frames if isinstance(f, Terminal)) + assert terminal.stop_reason == "delivery_overflow" + + +class TestPromptLogprobsAndMetrics: + + def test_held_prompt_logprobs_attach_to_next_delta(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + router = executor.router + # A no-token response carrying prompt logprobs is held... + wrapper = ResponseWrapper( + response_for(client_id, output_token_ids=[[]]), + logprobs=LogProbsResult(prompt=[-1.0, -2.0], generation=None)) + router.on_response(wrapper) + # ...and attached to the next token-carrying delta, exactly once. + router.on_response(delta_response(client_id, [5])) + router.on_response(delta_response(client_id, [6])) + router.on_response(final_response(client_id, tllm.FinishReason.END_ID)) + frames = asyncio.run(collect_frames(client, "req-1")) + deltas = [f for f in frames if isinstance(f, TokenDelta)] + assert deltas[0].prompt_logprobs == (-1.0, -2.0) + assert deltas[1].prompt_logprobs is None + + def test_cached_tokens_flow(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client_id = submitted_client_id(executor) + router = executor.router + router.on_response(delta_response(client_id, [5], cached_tokens=2)) + router.on_response(final_response(client_id, tllm.FinishReason.END_ID)) + frames = asyncio.run(collect_frames(client, "req-1")) + delta = frames[0] + assert delta.metrics["cached_tokens"] == 2.0 + assert frames[-1].cached_tokens == 2 + + +class TestSetupGates: + + @pytest.mark.parametrize("overrides,fragment", [ + (dict(flag_enabled=False), "flag"), + (dict(backend="_autodeploy"), "backend"), + (dict(backend="tensorrt"), "backend"), + (dict(transport="rpc"), "transport"), + (dict(num_postprocess_workers=2), "postproc"), + (dict(post_processor_hook_set=True), "post_processor_hook"), + (dict(speculative_config_set=True), "speculative"), + (dict(early_first_token_mode=True), "early_first_token"), + (dict(world_size=2), "topology"), + (dict(tokenizer_trust_remote_code=True), "trust_remote_code"), + ]) + def test_config_gate(self, overrides, fragment): + with pytest.raises(EngineClientConfigError) as excinfo: + LocalProcessEngineClient(FakeExecutor(), make_config(**overrides)) + assert fragment in str(excinfo.value) + + def test_flag_from_env(self, monkeypatch): + monkeypatch.delenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", raising=False) + with pytest.raises(EngineClientConfigError): + LocalProcessEngineClient(FakeExecutor(), make_config(flag_enabled=None)) + monkeypatch.setenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", "1") + LocalProcessEngineClient(FakeExecutor(), make_config(flag_enabled=None)) + + def test_executor_without_hooks_rejected(self): + with pytest.raises(EngineClientConfigError): + LocalProcessEngineClient(SimpleNamespace(), make_config()) + + +class TestPreSubmitChecks: + + def test_newer_protocol_rejected(self): + client = make_client() + request = make_engine_request() + object.__setattr__(request, "protocol_version", 99) + with pytest.raises(RequestRejectedError): + client.submit(request) + + def test_required_features_rederived_not_trusted(self): + client = make_client() + request = make_engine_request() + object.__setattr__(request, "required_features", ("guided_decoding", )) + with pytest.raises(RequestRejectedError): + client.submit(request) + + def test_capability_gated_feature_rejected_pre_submit(self): + client = make_client() + request = EngineRequest( + request_id="req-g", prompt_token_ids=(1, ), + sampling=EngineSamplingConfig(max_new_tokens=4, end_id=2), + guided_decoding=GuidedDecodingSpec(mode="json_object"), + required_features=("guided_decoding", )) + with pytest.raises(RequestRejectedError) as excinfo: + client.submit(request) + assert "guided_decoding" in str(excinfo.value) + + +class TestControlPlane: + + def test_typed_delegation(self): + client = make_client() + stats = client.get_stats() + assert stats.entries and stats.entries[0].startswith("{") + events = client.get_kv_events() + assert events.entries == ('{"event": 1}', ) + assert client.health().healthy + + def test_close_client_vs_shutdown_engine(self): + executor = FakeExecutor() + client = make_client(executor) + client.submit(make_engine_request()) + client.close_client() + assert not executor.shutdown_called + frames = asyncio.run(collect_frames(client, "req-1")) + assert isinstance(frames[-1], ErrorFrame) + executor2 = FakeExecutor() + client2 = make_client(executor2) + client2.shutdown_engine() + assert executor2.shutdown_called + + +class TestInvariantChecker: + + def test_detects_missing_ending(self): + + async def scenario(): + async def frames(): + yield TokenDelta(request_id="r", sequence_id=0, new_token_ids=(1, ), + event_seq=0) + + stream = InvariantCheckingStream(frames(), "r") + with pytest.raises(StreamInvariantViolation): + async for _ in stream: + pass + + asyncio.run(scenario()) + + def test_detects_frames_after_ending(self): + + async def scenario(): + async def frames(): + yield ErrorFrame(request_id="r", error_code="x", event_seq=0) + yield TokenDelta(request_id="r", sequence_id=0, new_token_ids=(1, ), + event_seq=1) + + stream = InvariantCheckingStream(frames(), "r") + with pytest.raises(StreamInvariantViolation): + async for _ in stream: + pass + + asyncio.run(scenario()) + + +class TestStressBaseline: + + def test_concurrent_requests_return_to_baseline(self): + executor = FakeExecutor() + client = make_client(executor) + router = executor.router + request_count = 40 + + for i in range(request_count): + client.submit(make_engine_request(request_id=f"req-{i}")) + + def respond(request): + client_id = request.id + router.on_response(delta_response(client_id, [5, 6])) + router.on_response(final_response(client_id, tllm.FinishReason.END_ID)) + + threads = [ + threading.Thread(target=respond, args=(request, )) + for request in executor.submitted + ] + for t in threads: + t.start() + for t in threads: + t.join() + + async def consume_all(): + for i in range(request_count): + frames = await collect_frames(client, f"req-{i}") + assert isinstance(frames[-1], RequestComplete) + + asyncio.run(consume_all()) + assert router.active_request_count() == 0 + + def test_tombstone_eviction_closes_deliveries(self): + executor = FakeExecutor() + router = EngineFrameRouter(abort_fn=executor.abort_request, + tombstone_limit=4) + executor.attach_engine_frame_router(router) + bindings = [] + for i in range(8): + request = SimpleNamespace(id=None) + request.set_id = lambda v, r=request: setattr(r, "id", v) or r + binding = router.register_pending(request, f"req-{i}", (1, ), ()) + request.set_id(100 + i) + router.observe_submit(request) + router.on_response(final_response(100 + i, tllm.FinishReason.END_ID)) + bindings.append(binding) + # Oldest deliveries were evicted alongside their tombstones. + assert bindings[0].delivery.closed + assert not bindings[-1].delivery.closed + assert DEFAULT_TOMBSTONE_LIMIT >= 4 diff --git a/tests/unittest/executor/engine_client/test_serving_e5a.py b/tests/unittest/executor/engine_client/test_serving_e5a.py new file mode 100644 index 000000000000..1fc7e46f55ca --- /dev/null +++ b/tests/unittest/executor/engine_client/test_serving_e5a.py @@ -0,0 +1,687 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""E5a serving-path tests: SSE semantic parity, fallback counters, model +context/tokenizer reload, event-loop responsiveness, and contract-path +guards. All CPU-only: the engine is faked at the proxy seam.""" + +import asyncio +import json +import os +import time +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.engine_client import conversion as conversion_module +from tensorrt_llm.executor.engine_client.assembler import FrontendResponseAssembler +from tensorrt_llm.executor.engine_client.contract import TokenizerSpec +from tensorrt_llm.executor.engine_client.local_client import ( + EngineClientConfig, LocalProcessEngineClient) +from tensorrt_llm.executor.result import DetokenizedGenerationResultBase +from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.serve.engine_client_serving import (ContractStreamView, + EngineClientServing) +from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, + StreamOptions) +from tensorrt_llm.serve.postprocess_handlers import (ChatPostprocArgs, + chat_stream_post_processor) + +PROMPT_IDS = (1, 2, 3) +MODEL_DIR = "/workspace/models/Qwen2.5-0.5B-Instruct" + + +class CharTokenizer: + + def encode(self, text, add_special_tokens=False): + return [ord(ch) - 97 for ch in text] + + def decode(self, ids, **kwargs): + if isinstance(ids, int): + ids = [ids] + return "".join(chr(97 + i % 26) for i in ids) + + def decode_incrementally(self, token_ids, prev_text=None, states=None, *, + flush=False, stream_interval=1, **kwargs): + prev_text = prev_text or "" + return prev_text + self.decode(token_ids), states or {} + + +def ids_for(text): + return [ord(ch) - 97 for ch in text] + + +class FakeExecutor: + + def __init__(self): + import threading + self.service = None + self.router = None + self.aborted = [] + self.submitted = [] + self._results = {} + self.doing_shutdown = False + self._fatal_error = None + self._submission_lock = threading.RLock() + self._next_id = 0 + + def attach_engine_service(self, service): + self.service = service + self.router = service.router + + def submit_contract(self, engine_request, stop_reasons=()): + return self.service.submit_contract(engine_request, + stop_reasons=stop_reasons) + + class _Queue: + + def __init__(self, outer): + self.outer = outer + + def put(self, item): + self.outer.submitted.append(item) + + @property + def request_queue(self): + return self._Queue(self) + + def _start_dispatch_threads(self): + pass + + def _get_next_client_id(self): + self._next_id += 1 + return self._next_id + + def _handle_background_error(self): + pass + + def abort_request(self, client_id): + self.aborted.append(client_id) + + def check_health(self): + return True + + def get_stats(self, timeout): + return [] + + def get_kv_events(self, timeout): + return [] + + def shutdown(self): + pass + + +def make_serving(tokenizer=None, model_type=None, + generation_stop_token_ids=()) -> EngineClientServing: + """Build the serving glue over a fake engine (bypasses the LLM ctor).""" + from tensorrt_llm.executor.engine_client.contract import ( + EngineCapabilities, FrontendModelContext) + serving = object.__new__(EngineClientServing) + serving.client = LocalProcessEngineClient( + FakeExecutor(), EngineClientConfig(backend="pytorch", flag_enabled=True)) + serving.tokenizer = tokenizer or CharTokenizer() + serving.counters = {"contract_requests": 0, "capability_rejections": 0} + serving._stream_interval = 1 + serving._force_return_perf_metrics = False + serving.model_context = FrontendModelContext( + tokenizer=TokenizerSpec(uri="/fake"), + capabilities=EngineCapabilities(features=("streaming", )), + eos_id=2, pad_id=0, + generation_stop_token_ids=tuple(generation_stop_token_ids), + model_type=model_type) + return serving + + +def make_result(**overrides): + fields = dict(is_final=False, output_token_ids=[[]], + finish_reasons=[tllm.FinishReason.NOT_FINISHED], log_probs=None, + cum_log_probs=None, sequence_index=0, context_phase_params=None, + decoding_iter=1, avg_decoded_tokens_per_iter=None, + request_perf_metrics=None, generation_logits=None, + context_logits=None) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def delta_response(client_id, tokens, **overrides): + return LlmResponse(request_id=1, client_id=client_id, + result=make_result(output_token_ids=[list(tokens)], + **overrides)) + + +def final_response(client_id, reason, tokens=(), **overrides): + return LlmResponse(request_id=1, client_id=client_id, + result=make_result(is_final=True, + output_token_ids=[list(tokens)], + finish_reasons=[reason], + **overrides)) + + +def make_chat_request(**overrides): + kwargs = dict(model="test-model", + messages=[{"role": "user", "content": "hi"}], stream=True, + max_tokens=32) + kwargs.update(overrides) + return ChatCompletionRequest(**kwargs) + + +def make_postproc_args(request, tokenizer, num_prompt_tokens=3) -> ChatPostprocArgs: + args = ChatPostprocArgs.from_request(request) + args.tokenizer = tokenizer + args.num_prompt_tokens = num_prompt_tokens + return args + + +def prepared_sampling_params(**kwargs) -> SamplingParams: + params = SamplingParams(max_tokens=32, end_id=2, pad_id=0, **kwargs) + if params.stop is not None: + strings = [params.stop] if isinstance(params.stop, str) else params.stop + params._stop_word_ids = [ids_for(s) for s in strings] + params._stream_interval = 1 + return params + + +def parse_sse(pieces): + events = [] + for piece in pieces: + piece = piece.strip() + if not piece.startswith("data: "): + continue + body = piece[len("data: "):].strip() + if body == "[DONE]": + events.append("DONE") + continue + events.append(json.loads(body)) + return events + + +def normalize_events(events): + """Semantic normalization: drop ids/timestamps, coalesce content deltas.""" + role = None + content = "" + finish_reason = None + stop_reason = None + usage_snapshots = [] + for event in events: + if event == "DONE": + continue + usage = event.get("usage") + if usage: + usage_snapshots.append({ + "prompt_tokens": usage.get("prompt_tokens"), + "completion_tokens": usage.get("completion_tokens"), + "total_tokens": usage.get("total_tokens"), + "cached_tokens": (usage.get("prompt_tokens_details") or {}).get( + "cached_tokens"), + }) + for choice in event.get("choices", []): + delta = choice.get("delta") or {} + if delta.get("role"): + role = delta["role"] + if delta.get("content"): + content += delta["content"] + if choice.get("text"): # completions-endpoint chunks + content += choice["text"] + if choice.get("finish_reason"): + finish_reason = choice["finish_reason"] + if choice.get("stop_reason") is not None: + stop_reason = choice["stop_reason"] + return { + "role": role, + "content": content, + "finish_reason": finish_reason, + "stop_reason": stop_reason, + # The ORDERED usage sequence is semantic (continuous usage chunks + # must match); chunk-boundary coalescing may merge token deltas but + # never reorders usage snapshots. + "usage_sequence": usage_snapshots, + "final_usage": usage_snapshots[-1] if usage_snapshots else None, + } + + +def run_legacy_sse(responses, sampling_params, request, tokenizer): + result = DetokenizedGenerationResultBase(id=7, sampling_params=sampling_params, + tokenizer=tokenizer, streaming=True) + args = make_postproc_args(request, tokenizer) + pieces = [] + for response in responses: + result._handle_response(response) + pieces.extend(chat_stream_post_processor(result, args)) + if result._done: + break + return pieces + + +def run_contract_sse(responses, sampling_params, request, serving=None): + serving = serving or make_serving() + args = make_postproc_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=sampling_params, + post_processor=chat_stream_post_processor, + postproc_args=args) + assert generator is not None, "expected the request to be eligible" + executor = serving.client._executor + + async def collect(): + # Feed responses incrementally so chunk boundaries (and therefore + # the ordered continuous-usage sequence) match the legacy + # per-response handler calls one-to-one. + pieces = [] + + async def puller(): + async for piece in generator: + pieces.append(piece) + + pull_task = asyncio.ensure_future(puller()) + for response in responses: + executor.router.on_response(response) + for _ in range(5): + await asyncio.sleep(0) + await asyncio.wait_for(pull_task, 10) + return pieces + + return asyncio.run(collect()), serving + + +class TestSseSemanticParity: + + def assert_sse_parity(self, responses, request=None, sampling_kwargs=None): + request = request or make_chat_request() + tokenizer = CharTokenizer() + legacy_params = prepared_sampling_params(**(sampling_kwargs or {})) + contract_params = prepared_sampling_params(**(sampling_kwargs or {})) + legacy = normalize_events( + parse_sse(run_legacy_sse(responses, legacy_params, request, tokenizer))) + contract_pieces, _ = run_contract_sse(responses, contract_params, request) + contract = normalize_events(parse_sse(contract_pieces)) + assert contract == legacy + return contract + + def test_plain_stream(self): + responses = [delta_response(1, ids_for("hel")), + delta_response(1, ids_for("lo")), + final_response(1, tllm.FinishReason.END_ID)] + result = self.assert_sse_parity(responses) + assert result["role"] == "assistant" + assert result["content"] == "hello" + assert result["finish_reason"] == "stop" + + def test_final_usage_and_continuous_usage(self): + request = make_chat_request(stream_options=StreamOptions( + include_usage=True, continuous_usage_stats=True)) + responses = [delta_response(1, ids_for("hi")), + final_response(1, tllm.FinishReason.LENGTH)] + result = self.assert_sse_parity(responses, request=request) + assert result["final_usage"] == { + "prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5, + "cached_tokens": 0, + } + assert result["finish_reason"] == "length" + + def test_stop_string_parity(self): + request = make_chat_request() + responses = [delta_response(1, ids_for("ax")), + delta_response(1, ids_for("yb")), + final_response(1, tllm.FinishReason.LENGTH)] + result = self.assert_sse_parity(responses, request=request, + sampling_kwargs=self.stop_kwargs("xy")) + # The already-streamed "x" prefix leaks before the stop trims — the + # known legacy cross-chunk behavior, deliberately preserved in V0 + # (a semantic fix is future work). Parity, not improvement, is the gate. + assert result["content"] == "ax" + assert result["finish_reason"] == "stop" + assert result["stop_reason"] == "xy" + + @staticmethod + def stop_kwargs(stop_string): + params_kwargs = dict(stop=[stop_string]) + return params_kwargs + + def test_timeout_rendering_parity(self): + responses = [delta_response(1, ids_for("a")), + final_response(1, tllm.FinishReason.TIMED_OUT)] + result = self.assert_sse_parity(responses) + assert result["finish_reason"] == "timeout" + + def test_cancelled_rendering_parity(self): + responses = [delta_response(1, ids_for("a")), + final_response(1, tllm.FinishReason.CANCELLED)] + result = self.assert_sse_parity(responses) + assert result["finish_reason"] == "cancelled" + + +class TestFallbackCounters: + + def test_ineligible_axes_fall_back_with_counters(self): + serving = make_serving() + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + cases = [ + (dict(n=2, best_of=2, top_p=0.9), "n_gt_1"), + (dict(logprobs=5), "top_logprobs"), + (dict(ignore_eos=True), "ignore_eos"), + ] + for sampling_kwargs, axis in cases: + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=[1, 2], + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(**sampling_kwargs), + post_processor=chat_stream_post_processor, postproc_args=args) + assert generator is None + assert serving.counters[f"fallback:{axis}"] == 1 + assert serving.counters["contract_requests"] == 0 + + def test_request_level_fallbacks(self): + serving = make_serving() + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=[1], + query_token_ids=None, + multimodal_params=object(), + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=chat_stream_post_processor, postproc_args=args) + assert generator is None + assert serving.counters["fallback:multimodal"] == 1 + # Missing token ids + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=chat_stream_post_processor, postproc_args=args) + assert generator is None + assert serving.counters["fallback:no_preprocessed_inputs"] == 1 + + def test_all_default_scheduling_params_normalized(self): + serving = make_serving() + normalized = serving._normalize_scheduling( + SimpleNamespace(agent_hierarchy=None)) + assert normalized is None + kept = serving._normalize_scheduling(SimpleNamespace(agent_hierarchy="x")) + assert kept is not None + + def test_counter_snapshot_includes_router(self): + serving = make_serving() + counters = serving.get_counters() + for key in ("late_or_duplicate_absorbed", "router_failures", + "synthesized_terminals", "overflow_aborts", + "active_requests", "contract_requests", + "capability_rejections"): + assert key in counters + + +class TestContractPathGuards: + + def test_contract_path_never_reenters_legacy(self, monkeypatch): + """The contract path must not call _setup, build PostprocParams, or + construct a GenerationResult.""" + from tensorrt_llm.executor import result as result_module + from tensorrt_llm.executor.postproc_worker import PostprocParams + + def forbidden(name): + def _raise(*args, **kwargs): + raise AssertionError(f"contract path invoked forbidden {name}") + return _raise + + monkeypatch.setattr(SamplingParams, "_setup", forbidden("_setup")) + monkeypatch.setattr(PostprocParams, "__init__", + forbidden("PostprocParams")) + monkeypatch.setattr(result_module.GenerationResult, "__init__", + forbidden("GenerationResult")) + + responses = [delta_response(1, ids_for("hi")), + final_response(1, tllm.FinishReason.END_ID)] + request = make_chat_request() + pieces, serving = run_contract_sse(responses, prepared_sampling_params(), + request) + assert pieces + assert serving.counters["contract_requests"] == 1 + + +class TestEventLoopResponsiveness: + + def test_many_concurrent_streams_keep_loop_responsive(self): + serving = make_serving() + request = make_chat_request() + stream_count = 32 + generators = [] + executor = serving.client._executor + client_ids = [] + for i in range(stream_count): + args = make_postproc_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=chat_stream_post_processor, postproc_args=args) + assert generator is not None + generators.append(generator) + client_ids.append(executor.submitted[-1].id) + + async def scenario(): + max_gap = 0.0 + + async def heartbeat(stop_event): + nonlocal max_gap + last = time.monotonic() + while not stop_event.is_set(): + await asyncio.sleep(0.005) + now = time.monotonic() + max_gap = max(max_gap, now - last) + last = now + + stop_event = asyncio.Event() + beat = asyncio.create_task(heartbeat(stop_event)) + + def feed(): + for step in range(20): + for client_id in client_ids: + executor.router.on_response( + delta_response(client_id, [step % 26])) + for client_id in client_ids: + executor.router.on_response( + final_response(client_id, tllm.FinishReason.END_ID)) + + feeder = asyncio.get_running_loop().run_in_executor(None, feed) + + async def drain(generator): + async for _ in generator: + pass + + await asyncio.gather(*(drain(g) for g in generators)) + await feeder + stop_event.set() + await beat + # The event loop must never stall meaningfully while 32 streams + # and a background feeder are active. + assert max_gap < 0.5, f"event loop stalled for {max_gap:.3f}s" + + asyncio.run(scenario()) + + +class TestContextOnlyBoundary: + """The contract path consumes only the data-only model context (R3).""" + + def test_no_live_model_objects_retained(self): + serving = make_serving() + # The glue holds no live LLM/model-config/generation-config handles: + # its normalization inputs are the frozen context + spec tokenizer. + assert not hasattr(serving, "_normalization") + forbidden = ("_hf_model_config", "_generation_config", "_llm") + for name in forbidden: + assert name not in vars(serving) + + def unprepared_params(self, **kwargs): + params = SamplingParams(max_tokens=8, **kwargs) # end_id unset + return params + + def try_unprepared(self, serving, params): + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + return serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=params, + post_processor=chat_stream_post_processor, postproc_args=args) + + def test_context_defaults_end_and_pad_ids(self): + serving = make_serving() + params = self.unprepared_params() + generator = self.try_unprepared(serving, params) + assert generator is not None + submitted = serving.client._executor.submitted[-1] + assert submitted.sampling_params.end_id == 2 + assert submitted.sampling_params.pad_id == 0 + + def test_stop_strings_tokenized_with_spec_tokenizer(self): + serving = make_serving() + params = self.unprepared_params(stop=["xy"]) + generator = self.try_unprepared(serving, params) + assert generator is not None + submitted = serving.client._executor.submitted[-1] + assert submitted.stop_token_sequences == [ids_for("xy")] + + def test_generation_stop_ids_merged_from_context(self): + serving = make_serving(generation_stop_token_ids=(32000, )) + params = self.unprepared_params() + generator = self.try_unprepared(serving, params) + assert generator is not None + submitted = serving.client._executor.submitted[-1] + assert 32000 in submitted.sampling_params.stop_token_ids + + def test_bart_model_type_ineligible_via_context(self): + serving = make_serving(model_type="bart") + generator = self.try_unprepared(serving, self.unprepared_params()) + assert generator is None + assert serving.counters["fallback:bart_forced_tokens"] == 1 + + def test_empty_trace_headers_are_eligible(self): + serving = make_serving() + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=chat_stream_post_processor, postproc_args=args, + trace_headers={}) + assert generator is not None # {} means "tracing on, no header" + generator2 = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=chat_stream_post_processor, postproc_args=args, + trace_headers={"traceparent": "00-abc"}) + assert generator2 is None + assert serving.counters["fallback:trace_headers"] == 1 + + def test_capability_rejection_counts_and_falls_back(self): + from tensorrt_llm.sampling_params import GuidedDecodingParams + serving = make_serving() + params = prepared_sampling_params( + guided_decoding=GuidedDecodingParams(json_object=True)) + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=params, + post_processor=chat_stream_post_processor, postproc_args=args) + assert generator is None + assert serving.counters["capability_rejections"] == 1 + + +class TestFailClosedStartup: + + class FakeArgsLLM(SimpleNamespace): + pass + + def make_llm(self, **arg_overrides): + args = SimpleNamespace(experimental_engine_client=True, + backend="tensorrt", trust_remote_code=False) + for key, value in arg_overrides.items(): + setattr(args, key, value) + return SimpleNamespace(args=args, _executor=SimpleNamespace(), + tokenizer=None) + + def test_flag_off_returns_none(self, monkeypatch): + monkeypatch.delenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", raising=False) + llm = self.make_llm(experimental_engine_client=False) + assert EngineClientServing.create_if_enabled(llm) is None + + def test_malformed_env_fails_closed(self, monkeypatch): + from tensorrt_llm.executor.engine_client.local_client import \ + EngineClientConfigError + monkeypatch.setenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", "yes") + with pytest.raises(EngineClientConfigError): + EngineClientServing.create_if_enabled(self.make_llm()) + + def test_config_gate_violation_fails_closed(self, monkeypatch): + from tensorrt_llm.executor.engine_client.local_client import \ + EngineClientConfigError + monkeypatch.setenv("TLLM_EXPERIMENTAL_ENGINE_CLIENT", "1") + # Wrong backend while enablement is requested: typed startup error, + # never a silent legacy fallback. + with pytest.raises(EngineClientConfigError): + EngineClientServing.create_if_enabled( + self.make_llm(backend="tensorrt")) + + +@pytest.mark.skipif(not os.path.isdir(MODEL_DIR), + reason="local tokenizer artifacts unavailable") +class TestModelContextReload: + + def test_tokenizer_reload_parity(self): + from tensorrt_llm.serve.engine_client_serving import \ + load_tokenizer_from_spec + from tensorrt_llm.tokenizer.tokenizer import TransformersTokenizer + spec = TokenizerSpec(uri=MODEL_DIR) + reloaded = load_tokenizer_from_spec(spec) + reference = TransformersTokenizer.from_pretrained(MODEL_DIR) + corpus = [ + "The capital of France is Paris.", + "def main():\n return 0", + "Unicode: héllo ☂ 你好 — em-dash", + " leading spaces and\ttabs\n", + "", + ] + for text in corpus: + assert reloaded.encode(text) == reference.encode(text) + sample_ids = reference.encode("Hello world, streaming détok!") + assert reloaded.decode(sample_ids) == reference.decode(sample_ids) + + def test_model_context_manifest(self): + from tensorrt_llm.serve.engine_client_serving import _file_manifest + manifest = _file_manifest(MODEL_DIR) + names = [name for name, _ in manifest] + assert "tokenizer_config.json" in names + for _, digest in manifest: + assert len(digest) == 64 diff --git a/tests/unittest/executor/engine_client/test_serving_sse_full.py b/tests/unittest/executor/engine_client/test_serving_sse_full.py new file mode 100644 index 000000000000..83ad07c82062 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_serving_sse_full.py @@ -0,0 +1,569 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Full SSE oracle surface: completions endpoint, ordered continuous usage, +SSE logprobs, assembly errors, disconnect abort, [DONE], and every +request-level fallback axis through the real facade.""" + +import asyncio +import json +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.engine_client.conversion import ELIGIBILITY_MATRIX +from tensorrt_llm.executor.result import DetokenizedGenerationResultBase +from tensorrt_llm.bindings import executor as tllme +from tensorrt_llm.sampling_params import LogprobMode, SamplingParams +from tensorrt_llm.serve.engine_client_serving import (_ContractStreamError, + wrap_sse_with_done) +from tensorrt_llm.serve.openai_protocol import CompletionRequest, StreamOptions +from tensorrt_llm.serve.postprocess_handlers import ( + CompletionPostprocArgs, chat_stream_post_processor, + completion_stream_post_processor) +from tensorrt_llm.executor.utils import ErrorResponse + +from test_serving_e5a import (PROMPT_IDS, CharTokenizer, delta_response, + final_response, ids_for, make_chat_request, + make_postproc_args, make_serving, normalize_events, + parse_sse, prepared_sampling_params, + run_contract_sse, run_legacy_sse) + + +def make_completion_request(**overrides): + kwargs = dict(model="test-model", prompt="hi", stream=True, max_tokens=32) + kwargs.update(overrides) + return CompletionRequest(**kwargs) + + +def make_completion_args(request, tokenizer, num_prompt_tokens=3): + args = CompletionPostprocArgs.from_request(request) + args.prompt_idx = 0 + args.tokenizer = tokenizer + args.num_prompt_tokens = num_prompt_tokens + return args + + +def run_legacy_completion_sse(responses, sampling_params, request, tokenizer): + result = DetokenizedGenerationResultBase(id=7, sampling_params=sampling_params, + tokenizer=tokenizer, streaming=True) + args = make_completion_args(request, tokenizer) + pieces = [] + for response in responses: + result._handle_response(response) + pieces.extend(completion_stream_post_processor(result, args)) + if result._done: + break + return pieces + + +def run_contract_completion_sse(responses, sampling_params, request, + serving=None): + serving = serving or make_serving() + args = make_completion_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=sampling_params, + post_processor=completion_stream_post_processor, postproc_args=args) + assert generator is not None + executor = serving.client._executor + + async def collect(): + pieces = [] + + async def puller(): + async for piece in generator: + pieces.append(piece) + + pull_task = asyncio.ensure_future(puller()) + for response in responses: + executor.router.on_response(response) + for _ in range(5): + await asyncio.sleep(0) + await asyncio.wait_for(pull_task, 10) + return pieces + + return asyncio.run(collect()), serving + + +class TestCompletionsSseParity: + + def assert_parity(self, responses, request=None, sampling_kwargs=None): + request = request or make_completion_request() + tokenizer = CharTokenizer() + legacy = normalize_events(parse_sse( + run_legacy_completion_sse( + responses, prepared_sampling_params(**(sampling_kwargs or {})), + request, tokenizer))) + pieces, _ = run_contract_completion_sse( + responses, prepared_sampling_params(**(sampling_kwargs or {})), + request) + contract = normalize_events(parse_sse(pieces)) + assert contract == legacy + return contract + + def test_plain_stream(self): + responses = [delta_response(1, ids_for("hel")), + delta_response(1, ids_for("lo")), + final_response(1, tllm.FinishReason.END_ID)] + result = self.assert_parity(responses) + assert result["content"] == "hello" + assert result["finish_reason"] == "stop" + + def test_continuous_usage_sequence(self): + request = make_completion_request(stream_options=StreamOptions( + include_usage=True, continuous_usage_stats=True)) + responses = [delta_response(1, ids_for("ab")), + delta_response(1, ids_for("cd")), + final_response(1, tllm.FinishReason.LENGTH)] + result = self.assert_parity(responses, request=request) + assert len(result["usage_sequence"]) >= 2 + assert result["final_usage"]["completion_tokens"] == 4 + + def test_stop_string(self): + responses = [delta_response(1, ids_for("ax")), + delta_response(1, ids_for("yb")), + final_response(1, tllm.FinishReason.LENGTH)] + params_kwargs = dict(stop=["xy"]) + result = self.assert_parity(responses, sampling_kwargs=params_kwargs) + assert result["finish_reason"] == "stop" + assert result["stop_reason"] == "xy" + + def test_cancelled(self): + responses = [delta_response(1, ids_for("a")), + final_response(1, tllm.FinishReason.CANCELLED)] + result = self.assert_parity(responses) + assert result["finish_reason"] == "cancelled" + + +def extract_logprob_values(pieces): + values = [] + for event in parse_sse(pieces): + if event == "DONE": + continue + for choice in event.get("choices", []): + logprobs = choice.get("logprobs") + if not logprobs: + continue + content = logprobs.get("content") or [] + for entry in content: + values.append(round(entry.get("logprob", 0.0), 6)) + token_logprobs = logprobs.get("token_logprobs") or [] + values.extend(round(v, 6) for v in token_logprobs if v is not None) + return values + + +class TestSseLogprobs: + + def test_chat_logprob_chunks_match_legacy(self): + request = make_chat_request(logprobs=True) + tokenizer = CharTokenizer() + # Real legacy streams fuse the last delta with the final response; + # a logprob-carrying stream therefore ends with a token+logprob + # final, which both pipelines must render identically. + responses = [ + delta_response(1, ids_for("ab"), log_probs=[[-0.5, -1.0]]), + final_response(1, tllm.FinishReason.END_ID, tokens=ids_for("c"), + log_probs=[[-0.25]]), + ] + legacy_pieces = run_legacy_sse( + responses, prepared_sampling_params(logprobs=0), request, tokenizer) + contract_pieces, _ = run_contract_sse( + responses, prepared_sampling_params(logprobs=0), request) + assert extract_logprob_values(contract_pieces) == \ + extract_logprob_values(legacy_pieces) + assert extract_logprob_values(contract_pieces) # non-empty + + +class TestAssemblyErrors: + + def test_mid_stream_error_raises_typed(self): + serving = make_serving() + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=chat_stream_post_processor, postproc_args=args) + executor = serving.client._executor + client_id = executor.submitted[-1].id + executor.router.on_response(delta_response(client_id, ids_for("a"))) + executor.router.on_response( + ErrorResponse(client_id=client_id, error_msg="boom", request_id=1)) + + async def consume(): + async for _ in generator: + pass + + with pytest.raises(_ContractStreamError): + asyncio.run(consume()) + + +class TestDisconnectAbort: + + def test_client_disconnect_triggers_engine_abort(self): + serving = make_serving() + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + + class FakeRawRequest: + + def __init__(self): + self.disconnected = False + self.state = SimpleNamespace() + + async def is_disconnected(self): + return self.disconnected + + raw_request = FakeRawRequest() + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=chat_stream_post_processor, postproc_args=args, + raw_request=raw_request) + executor = serving.client._executor + client_id = executor.submitted[-1].id + + async def scenario(): + executor.router.on_response(delta_response(client_id, ids_for("a"))) + task = asyncio.ensure_future( + self._collect_with_timeout(generator)) + await asyncio.sleep(0.05) + raw_request.disconnected = True + # The watcher polls every second; the abort lands, then the fake + # engine answers with a CANCELLED final and the stream ends. + for _ in range(30): + await asyncio.sleep(0.1) + if executor.aborted: + break + assert executor.aborted, "disconnect did not trigger an engine abort" + executor.router.on_response( + final_response(client_id, tllm.FinishReason.CANCELLED)) + await asyncio.wait_for(task, 10) + + asyncio.run(scenario()) + + @staticmethod + async def _collect_with_timeout(generator): + async for _ in generator: + pass + + +class TestDoneWrapper: + + def test_done_appended_and_first_token_time_set(self): + + async def pieces(): + yield "data: {}\n\n" + yield "data: {}\n\n" + + raw_request = SimpleNamespace(state=SimpleNamespace()) + + async def collect(): + return [p async for p in wrap_sse_with_done(pieces(), raw_request)] + + out = asyncio.run(collect()) + assert out[-1] == "data: [DONE]\n\n" + assert len(out) == 3 + assert hasattr(raw_request.state, "server_first_token_time") + + def test_empty_stream_still_emits_done(self): + + async def empty(): + return + yield # pragma: no cover + + async def collect(): + return [p async for p in wrap_sse_with_done(empty())] + + out = asyncio.run(collect()) + assert out == ["data: [DONE]\n\n"] + + +class TestCompletionsSseLogprobs: + + def test_completion_logprob_chunks_match_legacy(self): + request = make_completion_request(logprobs=1) + tokenizer = CharTokenizer() + # Legacy-faithful fused final: the last delta rides the final + # response, carrying both tokens and their logprobs. + responses = [ + delta_response(1, ids_for("ab"), log_probs=[[-0.5, -1.0]]), + final_response(1, tllm.FinishReason.END_ID, tokens=ids_for("c"), + log_probs=[[-0.25]]), + ] + legacy_pieces = run_legacy_completion_sse( + responses, prepared_sampling_params(logprobs=1), request, tokenizer) + contract_pieces, _ = run_contract_completion_sse( + responses, prepared_sampling_params(logprobs=1), request) + assert extract_logprob_values(contract_pieces) == \ + extract_logprob_values(legacy_pieces) + assert extract_logprob_values(contract_pieces) # non-empty + + +class TestCompletionsAssemblyErrors: + + def test_mid_stream_error_raises_typed(self): + serving = make_serving() + request = make_completion_request() + args = make_completion_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=completion_stream_post_processor, + postproc_args=args) + executor = serving.client._executor + client_id = executor.submitted[-1].id + executor.router.on_response(delta_response(client_id, ids_for("a"))) + executor.router.on_response( + ErrorResponse(client_id=client_id, error_msg="boom", request_id=1)) + + async def consume(): + async for _ in generator: + pass + + with pytest.raises(_ContractStreamError): + asyncio.run(consume()) + + +class TestCompletionsDisconnectAbort: + + def test_client_disconnect_triggers_engine_abort(self): + serving = make_serving() + request = make_completion_request() + args = make_completion_args(request, serving.tokenizer) + + class FakeRawRequest: + + def __init__(self): + self.disconnected = False + self.state = SimpleNamespace() + + async def is_disconnected(self): + return self.disconnected + + raw_request = FakeRawRequest() + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=completion_stream_post_processor, + postproc_args=args, raw_request=raw_request) + executor = serving.client._executor + client_id = executor.submitted[-1].id + + async def scenario(): + executor.router.on_response(delta_response(client_id, ids_for("a"))) + + async def consume(): + async for _ in generator: + pass + + task = asyncio.ensure_future(consume()) + await asyncio.sleep(0.05) + raw_request.disconnected = True + for _ in range(30): + await asyncio.sleep(0.1) + if executor.aborted: + break + assert executor.aborted, "disconnect did not trigger an engine abort" + executor.router.on_response( + final_response(client_id, tllm.FinishReason.CANCELLED)) + await asyncio.wait_for(task, 10) + + asyncio.run(scenario()) + + +class TestActiveRequestCounter: + + def test_active_count_transitions_zero_one_zero(self): + serving = make_serving() + assert serving.get_counters()["active_requests"] == 0 + request = make_completion_request() + responses = [delta_response(1, ids_for("ab")), + final_response(1, tllm.FinishReason.END_ID)] + # Mid-stream observation: submit, deliver one delta, snapshot. + args = make_completion_args(request, serving.tokenizer) + generator = serving.try_stream( + preprocessed=SimpleNamespace(prompt_token_ids=list(PROMPT_IDS), + query_token_ids=None, + multimodal_params=None, + encoder_input_token_ids=None), + sampling_params=prepared_sampling_params(), + post_processor=completion_stream_post_processor, + postproc_args=args) + executor = serving.client._executor + client_id = executor.submitted[-1].id + executor.router.on_response(responses[0]) + assert serving.get_counters()["active_requests"] == 1 + + async def finish(): + pieces = [] + + async def puller(): + async for piece in generator: + pieces.append(piece) + + task = asyncio.ensure_future(puller()) + executor.router.on_response(responses[1]) + await asyncio.wait_for(task, 10) + return pieces + + asyncio.run(finish()) + assert serving.get_counters()["active_requests"] == 0 + assert serving.counters["contract_requests"] == 1 + + +class TestNoLiveNormalizationReaches: + """The contract path must normalize ONLY via the context-only boundary: + neither the legacy-input normalizer nor any live ``LLM`` sampling + preparation may run for a contract request.""" + + def test_contract_stream_with_live_normalizers_poisoned(self, monkeypatch): + import tensorrt_llm.serve.engine_client_serving as serving_module + from tensorrt_llm.llmapi.llm import LLM + + def poison(*args, **kwargs): + raise AssertionError( + "live sampling normalization reached on the contract path") + + monkeypatch.setattr(serving_module, "prepare_sampling_params", poison, + raising=False) + monkeypatch.setattr(LLM, "_prepare_sampling_params", poison, + raising=False) + request = make_completion_request() + responses = [delta_response(1, ids_for("hi")), + final_response(1, tllm.FinishReason.END_ID)] + pieces, _ = run_contract_completion_sse( + responses, prepared_sampling_params(), request) + assert pieces, "contract stream produced no output" + + +class TestTokenizerSpecProvenance: + + def test_manifest_hash_mismatch_fails_typed(self, tmp_path): + from tensorrt_llm.executor.engine_client.contract import TokenizerSpec + from tensorrt_llm.serve.engine_client_serving import ( + EngineClientConfigError, load_tokenizer_from_spec) + (tmp_path / "tokenizer.json").write_bytes(b"{}") + spec = TokenizerSpec(uri=str(tmp_path), + files_manifest=(("tokenizer.json", "0" * 64), )) + with pytest.raises(EngineClientConfigError, match="manifest"): + load_tokenizer_from_spec(spec) + + def test_missing_manifest_file_fails_typed(self, tmp_path): + from tensorrt_llm.executor.engine_client.contract import TokenizerSpec + from tensorrt_llm.serve.engine_client_serving import ( + EngineClientConfigError, load_tokenizer_from_spec) + spec = TokenizerSpec(uri=str(tmp_path), + files_manifest=(("tokenizer.json", "0" * 64), )) + with pytest.raises(EngineClientConfigError, match="unreadable"): + load_tokenizer_from_spec(spec) + + +AXIS_RECIPES = [ + ("echo", {}, dict(echo=True)), + ("non_streaming", {}, dict(streaming=False)), + ("logprobs_mode", dict(logprobs=0, logprobs_mode=LogprobMode.PROCESSED), {}), + ("lookahead_config", + dict(lookahead_config=tllme.LookaheadDecodingConfig(2, 2, 2)), {}), + ("n_gt_1", dict(n=2, best_of=2, top_p=0.9), {}), + ("beam_search", dict(use_beam_search=True, n=1, best_of=1), {}), + ("top_logprobs", dict(logprobs=5), {}), + ("prompt_top_logprobs", dict(prompt_logprobs=5), {}), + ("logits_processor", dict(apply_batched_logits_processor=True), {}), + ("embedding_bias", dict(embedding_bias=torch.zeros(4)), {}), + ("bad_words", dict(bad_token_ids=[3]), {}), + ("ignore_eos", dict(ignore_eos=True), {}), + ("min_p", dict(min_p=0.2), {}), + ("top_p_extras", dict(top_p_min=0.5), {}), + ("no_repeat_ngram", dict(no_repeat_ngram_size=2), {}), + ("prompt_ignore_length", dict(prompt_ignore_length=1), {}), + ("return_logits", dict(return_context_logits=True), {}), + ("exclude_input_from_output", dict(exclude_input_from_output=False), {}), + ("truncate_prompt_tokens", dict(truncate_prompt_tokens=4), {}), + ("thinking_token_budget", dict(thinking_token_budget=10), {}), + ("multimodal", {}, dict(shim_multimodal=True)), + ("lora", {}, dict(lora_request=object())), + ("prompt_adapter", {}, dict(prompt_adapter_request=object())), + ("disaggregated", {}, dict(disaggregated_params=object())), + ("scheduling_params", {}, dict(scheduling_params=SimpleNamespace( + agent_hierarchy="x"))), + ("conversation_params", {}, dict(conversation_params=object())), + ("trace_headers", {}, dict(trace_headers={"traceparent": "00-abc"})), + ("cache_salt", {}, dict(cache_salt="salt")), + ("query_token_ids", {}, dict(shim_query=True)), + ("encoder_input", {}, dict(shim_encoder=True)), + ("kv_cache_retention", {}, dict(kv_cache_retention_config=object())), + ("priority", {}, dict(priority=0.9)), +] + + +class TestEveryFallbackAxisThroughFacade: + + @pytest.mark.parametrize("axis,param_kwargs,stream_kwargs", AXIS_RECIPES) + def test_axis_falls_back_with_counter(self, axis, param_kwargs, stream_kwargs): + serving = make_serving() + request = make_chat_request() + args = make_postproc_args(request, serving.tokenizer) + params = prepared_sampling_params(**param_kwargs) + stream_kwargs = dict(stream_kwargs) + shim = SimpleNamespace( + prompt_token_ids=list(PROMPT_IDS), + query_token_ids=[1] if stream_kwargs.pop("shim_query", False) else None, + multimodal_params=object() + if stream_kwargs.pop("shim_multimodal", False) else None, + encoder_input_token_ids=[1] + if stream_kwargs.pop("shim_encoder", False) else None) + generator = serving.try_stream( + preprocessed=shim, sampling_params=params, + post_processor=chat_stream_post_processor, postproc_args=args, + **stream_kwargs) + assert generator is None + assert serving.counters[f"fallback:{axis}"] == 1 + assert serving.counters["contract_requests"] == 0 + + def test_every_request_level_axis_is_driven_through_the_facade(self): + tested = {axis for axis, _, _ in AXIS_RECIPES} + # The only two axes not drivable via AXIS_RECIPES, each verified by a + # dedicated facade/enforcement-point test rather than silently waived: + # - bart_forced_tokens fires from the model context (needs a serving + # fixture with model_type="bart") — driven through the real facade + # with its counter by TestContextOnlyBoundary. + # - postproc_params cannot reach the facade at all (both endpoints + # construct PostprocParams strictly after the contract branch, and + # try_stream exposes no such parameter); its defense-in-depth + # rejection at convert_request is exercised by + # test_conversion.REJECTION_CASES. + exempt = {"bart_forced_tokens", "postproc_params"} + matrix_axes = {rule.axis for rule in ELIGIBILITY_MATRIX + if rule.classification == "ineligible"} + missing = matrix_axes - tested - exempt + assert not missing, f"axes not driven through the facade: {missing}" diff --git a/tests/unittest/executor/engine_client/test_stop_sequence_carrier.py b/tests/unittest/executor/engine_client/test_stop_sequence_carrier.py new file mode 100644 index 000000000000..d1ded6f80f95 --- /dev/null +++ b/tests/unittest/executor/engine_client/test_stop_sequence_carrier.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the pre-tokenized stop-sequence carrier on GenerationRequest.""" + +import pickle + +import pytest + +from tensorrt_llm.executor.base_worker import _request_stop_words +from tensorrt_llm.executor.request import GenerationRequest +from tensorrt_llm.sampling_params import SamplingParams + + +def make_request(**kwargs) -> GenerationRequest: + defaults = dict(prompt_token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=8), + streaming=True) + defaults.update(kwargs) + return GenerationRequest(**defaults) + + +class TestCarrierField: + + def test_default_is_none(self): + assert make_request().stop_token_sequences is None + + def test_accepts_sequences(self): + request = make_request(stop_token_sequences=[[7, 8], [9]]) + assert request.stop_token_sequences == [[7, 8], [9]] + + @pytest.mark.parametrize("bad", [ + [[]], + [[1, "x"]], + [[True]], + ["ab"], + [7], + ]) + def test_rejects_malformed(self, bad): + with pytest.raises((ValueError, TypeError)): + make_request(stop_token_sequences=bad) + + def test_survives_pickle(self): + request = make_request(stop_token_sequences=[[7, 8]]) + clone = pickle.loads(pickle.dumps(request)) + assert clone.stop_token_sequences == [[7, 8]] + assert clone.prompt_token_ids == [1, 2, 3] + + +class TestRequestStopWords: + + def test_carrier_merges_with_sampling_stops(self): + request = make_request( + sampling_params=SamplingParams(max_tokens=8, stop_token_ids=[13]), + stop_token_sequences=[[7, 8], [9]]) + assert _request_stop_words(request) == [[13], [7, 8], [9]] + + def test_carrier_alone(self): + request = make_request(stop_token_sequences=[[7, 8]]) + assert _request_stop_words(request) == [[7, 8]] + + def test_no_stops(self): + assert _request_stop_words(make_request()) == [] + + def test_ignore_eos_suppresses_all(self): + request = make_request( + sampling_params=SamplingParams(max_tokens=8, stop_token_ids=[13], + ignore_eos=True), + stop_token_sequences=[[7, 8]]) + assert _request_stop_words(request) == [] + + def test_legacy_requests_without_field(self): + request = make_request(sampling_params=SamplingParams(max_tokens=8, + stop_token_ids=[13])) + # Requests unpickled from an older writer may lack the attribute. + del request.stop_token_sequences + assert _request_stop_words(request) == [[13]]