diff --git a/.github/filters.yaml b/.github/filters.yaml index 3c7f3867ab43..8538fc85d790 100644 --- a/.github/filters.yaml +++ b/.github/filters.yaml @@ -358,6 +358,7 @@ sidecar: - 'tests/router/helper.py' - 'tests/router/router_process.py' - 'tests/utils/gpu_args.py' + - 'tests/utils/payloads.py' - 'tests/utils/prometheus.py' - 'tests/utils/router_logs.py' - 'tests/utils/router_nvext.py' diff --git a/.github/scripts/test-filters.js b/.github/scripts/test-filters.js index c9484a9dda10..266a9c51f38a 100755 --- a/.github/scripts/test-filters.js +++ b/.github/scripts/test-filters.js @@ -127,7 +127,7 @@ const testCases = [ desc: 'common sidecar source avoids unrelated build and E2E filters' }, { - file: 'lib/sidecar/trtllm/proto/trtllm_service.proto', + file: 'lib/sidecar/trtllm/proto/openengine/v1/generation.proto', expect: { sidecar: true, rust: true, core: false, frontend: false, vllm: false, sglang: false, trtllm: false }, desc: 'sidecar proto contracts trigger Rust checks without backend E2E' }, diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cd932c6295be..8eec0aad0c5b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,7 +14,7 @@ # limitations under the License. default_install_hook_types: [pre-commit, commit-msg] -exclude: ^(src/grpc_generated|.*\.patch$|.*/connect/.*\.py|components/src/dynamo/planner/plugins/proto/v1/plugin_pb2(_grpc)?\.pyi?$|lib/sidecar/trtllm/proto/trtllm_service\.proto$) +exclude: ^(src/grpc_generated|.*\.patch$|.*/connect/.*\.py|components/src/dynamo/planner/plugins/proto/v1/plugin_pb2(_grpc)?\.pyi?$|lib/sidecar/trtllm/proto/openengine/v1/.*\.proto$) repos: - repo: https://github.com/timothycrosley/isort rev: 5.12.0 diff --git a/Cargo.lock b/Cargo.lock index c3c3d47524d1..cb3d75dc4f20 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3101,6 +3101,9 @@ dependencies = [ "clap", "dynamo-backend-common", "futures", + "prost-types 0.13.5", + "prost-types 0.14.3", + "serde_json", "tokio", "tonic 0.13.1", "tonic 0.14.6", @@ -3161,6 +3164,8 @@ dependencies = [ "dynamo-sidecar-common", "futures", "prost 0.13.5", + "prost-types 0.13.5", + "serde", "serde_json", "tokio", "tokio-stream", diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/tensorrt-llm/sidecar.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/tensorrt-llm/sidecar.md index ee434c48e4bc..2462a645d0c5 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/tensorrt-llm/sidecar.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/tensorrt-llm/sidecar.md @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 title: TensorRT-LLM Sidecar -subtitle: Run Dynamo beside a stock TensorRT-LLM engine through native gRPC. +subtitle: Run Dynamo beside a TensorRT-LLM engine through its OpenEngine gRPC API. --- > [!WARNING] @@ -10,9 +10,10 @@ subtitle: Run Dynamo beside a stock TensorRT-LLM engine through native gRPC. > coverage can change without notice. `dynamo-trtllm-sidecar` is a CPU-only Dynamo worker that connects to -TensorRT-LLM's native gRPC service. It preserves the upstream engine process -and argument surface while using Dynamo for request handling and distributed -serving. See the +TensorRT-LLM's OpenEngine gRPC API (`openengine.v1`), served by +`trtllm-serve --grpc --grpc-protocol openengine`. It preserves the upstream +engine process and argument surface while using Dynamo for request handling and +distributed serving. See the [Sidecar Backends](../../../concepts/system-architecture/sidecar-backends.md) page for the common architecture. @@ -20,51 +21,53 @@ architecture. | Deployment path | Aggregated | Disaggregated | |---|---|---| -| Local launcher | Validated on one GPU | Not supported | -| Kubernetes example | Validated | Not supported | +| Local launcher | Validated on one GPU | Validated on one GPU, both engines co-located; the launcher defaults to two | +| Kubernetes example | Validated | Validated, prefill and decode on separate pods | This table covers launch topology only. The [TensorRT-LLM feature matrix](overview.md#feature-support-matrix) describes the -in-process backend; sidecar feature parity is still under evaluation. The -current native gRPC contract does not provide the prefill/decode handoff needed -for disaggregated serving. See the +in-process backend; sidecar feature parity is still under evaluation. +Disaggregated prefill/decode is supported over the OpenEngine contract: a +prefill worker marks its request `context_only` and returns the `PrefillReady` +KV handoff that a decode worker replays. Running it needs an engine with a KV +cache transceiver configured on both legs. See the [TensorRT-LLM sidecar README](https://github.com/ai-dynamo/dynamo/blob/main/lib/sidecar/trtllm/README.md) for other protocol limitations. ## Launch Locally -From a Dynamo source checkout, build or install Dynamo so -`dynamo-trtllm-sidecar` is on `PATH`. Install a TensorRT-LLM release that -provides `tensorrt_llm.commands.serve --grpc`. - -Start Dynamo's local discovery services, then run the aggregated launcher: +Build or install Dynamo from a source checkout so `dynamo-trtllm-sidecar` is on +`PATH`. The engine must serve `--grpc --grpc-protocol openengine`, which means +`nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc27.dev202609170000` or newer; the +launchers install the pinned OpenEngine Python bindings, which those releases do +not ship. ```bash -docker compose -f dev/docker-compose.yml up -d +docker compose -f dev/docker-compose.yml up -d # local discovery services ./lib/sidecar/trtllm/launch/agg.sh --model Qwen/Qwen3-0.6B ``` -The launcher starts the Dynamo frontend, TensorRT-LLM engine, and sidecar. It -binds TensorRT-LLM's native gRPC endpoint to loopback. - -Verify the frontend: +Each launcher starts the Dynamo frontend, the engine, and the sidecar, and binds +the engine's gRPC endpoint to loopback — it is unauthenticated and plaintext. +`launch/disagg.sh` is the same on two GPUs, with a NIXL cache transceiver on both +engines for the KV handoff. Either way the frontend serves the usual endpoint: ```bash -curl localhost:8000/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "Qwen/Qwen3-0.6B", - "messages": [{"role": "user", "content": "Hello"}], - "max_tokens": 32 - }' +curl localhost:8000/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"model":"Qwen/Qwen3-0.6B","messages":[{"role":"user","content":"Hello"}],"max_tokens":32}' ``` ## Deploy on Kubernetes -No published sidecar image is available yet. Follow the -[Kubernetes quick start](https://github.com/ai-dynamo/dynamo/blob/main/lib/sidecar/trtllm/README.md#deploy-on-kubernetes-quick-start) -to build `dynamo-sidecar`, which contains all three engine-specific sidecar -executables. The TensorRT-LLM manifest runs `dynamo-trtllm-sidecar` as the -container command and pairs it with a stock upstream TensorRT-LLM image. The -source tree includes an -[aggregated deployment manifest](https://github.com/ai-dynamo/dynamo/blob/main/lib/sidecar/trtllm/deploy/agg.yaml). +The source tree ships two +[deployment manifests](https://github.com/ai-dynamo/dynamo/tree/main/lib/sidecar/trtllm/deploy): +`agg.yaml`, and `disagg.yaml`, which runs prefill and decode as separate worker +pods. Both need two images you build yourself: `dynamo-sidecar`, which carries +all three engine-specific executables, and a TensorRT-LLM image with the pinned +OpenEngine bindings layered on — the release ships the servicer but not the +bindings that it and the manifests' health probes import. + +Read the disaggregated manifest's header before applying it: it runs the engines +over TCP/CUDA-IPC and requests no `rdma/ib`, which you add on a fabric that +provides it. The +[README](https://github.com/ai-dynamo/dynamo/blob/main/lib/sidecar/trtllm/README.md#deploy-on-kubernetes) has the full walkthrough. diff --git a/lib/backend-common/src/engine.rs b/lib/backend-common/src/engine.rs index 2439dcc0a94f..2c5daf1643e9 100644 --- a/lib/backend-common/src/engine.rs +++ b/lib/backend-common/src/engine.rs @@ -33,7 +33,7 @@ pub use dynamo_llm::protocols::common::preprocessor::{ pub use dynamo_llm::protocols::common::{ FinishReason, GuidedDecodingOptions, OutputOptions, SamplingOptions, StopConditions, }; -pub use dynamo_protocols::types::{CompletionUsage, StopReason}; +pub use dynamo_protocols::types::{CompletionUsage, PromptTokensDetails, StopReason}; pub use dynamo_runtime::engine::AsyncEngineContext; /// Per-request handle wrapping the runtime context. `Deref`s to diff --git a/lib/backend-common/src/lib.rs b/lib/backend-common/src/lib.rs index 197b80e4fff0..fd749de7e1ab 100644 --- a/lib/backend-common/src/lib.rs +++ b/lib/backend-common/src/lib.rs @@ -39,8 +39,9 @@ pub use engine::{ KV_HINT_TRANSFER_CAPABILITY_KEY, KvEventPublisher, KvEventSource, KvHint, KvHintAction, KvSourceLocationsPayload, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, LlmRegistration, LogProbs, Metrics, MetricsBindings, MetricsCtx, MultimodalData, OnPublisherReady, - OnSnapshotPublisherReady, OutputOptions, PrefillResult, PreprocessedRequest, RawEngine, - SamplingOptions, StopConditions, StopReason, TopLogprob, TopLogprobs, chunk, usage, + OnSnapshotPublisherReady, OutputOptions, PrefillResult, PreprocessedRequest, + PromptTokensDetails, RawEngine, SamplingOptions, StopConditions, StopReason, TopLogprob, + TopLogprobs, chunk, usage, }; pub use error::{BackendError, DynamoError, ErrorType}; pub use metrics::{ComponentGauges, EngineMetrics, LifecycleGauges}; diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 81dbd880e125..d5fb2fa516b2 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2590,6 +2590,9 @@ dependencies = [ "clap", "dynamo-backend-common", "futures", + "prost-types 0.13.5", + "prost-types 0.14.4", + "serde_json", "tokio", "tonic 0.13.1", "tonic 0.14.6", @@ -2647,6 +2650,8 @@ dependencies = [ "dynamo-sidecar-common", "futures", "prost 0.13.5", + "prost-types 0.13.5", + "serde", "serde_json", "tokio", "tokio-util", diff --git a/lib/sidecar/Dockerfile b/lib/sidecar/Dockerfile index d253f5b9a237..dae9c342d696 100644 --- a/lib/sidecar/Dockerfile +++ b/lib/sidecar/Dockerfile @@ -74,6 +74,13 @@ COPY examples/router/custom-policy-example/ examples/router/custom-policy-exampl COPY deploy/inference-gateway/ext-proc/ deploy/inference-gateway/ext-proc/ COPY deploy/inference-gateway/sidecar/ deploy/inference-gateway/sidecar/ +# The `/src/target` cache below is keyed by a name that changes whenever a +# build-script *input* is renamed or removed. The TensorRT-LLM sidecar's +# vendored protos moved from `proto/trtllm_service.proto` to +# `proto/openengine/v1/`, and a cached build-script binary compiled against the +# old layout keeps asking protoc for a file this tree no longer has. Bumping the +# id costs one cold build and makes that impossible. + # ---- vLLM builder ----------------------------------------------------------- FROM builder-base AS vllm-builder @@ -89,7 +96,7 @@ ARG TARGETARCH # build do not serialize against each other. RUN --mount=type=cache,id=dynamo-sidecar-cargo-registry-${TARGETARCH},target=/usr/local/cargo/registry,sharing=locked \ --mount=type=cache,id=dynamo-sidecar-cargo-git-${TARGETARCH},target=/usr/local/cargo/git,sharing=locked \ - --mount=type=cache,id=dynamo-sidecar-target-${TARGETARCH},target=/src/target,sharing=locked \ + --mount=type=cache,id=dynamo-sidecar-target-${TARGETARCH}-openengine,target=/src/target,sharing=locked \ cargo build --release --locked -p dynamo-vllm-sidecar \ && strip target/release/dynamo-vllm-sidecar \ && install -D target/release/dynamo-vllm-sidecar /out/dynamo-vllm-sidecar @@ -101,7 +108,7 @@ ARG TARGETARCH RUN --mount=type=cache,id=dynamo-sidecar-cargo-registry-${TARGETARCH},target=/usr/local/cargo/registry,sharing=locked \ --mount=type=cache,id=dynamo-sidecar-cargo-git-${TARGETARCH},target=/usr/local/cargo/git,sharing=locked \ - --mount=type=cache,id=dynamo-sidecar-target-${TARGETARCH},target=/src/target,sharing=locked \ + --mount=type=cache,id=dynamo-sidecar-target-${TARGETARCH}-openengine,target=/src/target,sharing=locked \ cargo build --release --locked -p dynamo-sglang-sidecar \ && strip target/release/dynamo-sglang-sidecar \ && install -D target/release/dynamo-sglang-sidecar /out/dynamo-sglang-sidecar @@ -113,7 +120,7 @@ ARG TARGETARCH RUN --mount=type=cache,id=dynamo-sidecar-cargo-registry-${TARGETARCH},target=/usr/local/cargo/registry,sharing=locked \ --mount=type=cache,id=dynamo-sidecar-cargo-git-${TARGETARCH},target=/usr/local/cargo/git,sharing=locked \ - --mount=type=cache,id=dynamo-sidecar-target-${TARGETARCH},target=/src/target,sharing=locked \ + --mount=type=cache,id=dynamo-sidecar-target-${TARGETARCH}-openengine,target=/src/target,sharing=locked \ cargo build --release --locked -p dynamo-trtllm-sidecar \ && strip target/release/dynamo-trtllm-sidecar \ && install -D target/release/dynamo-trtllm-sidecar /out/dynamo-trtllm-sidecar diff --git a/lib/sidecar/common/Cargo.toml b/lib/sidecar/common/Cargo.toml index 428d93591070..b95740ef6580 100644 --- a/lib/sidecar/common/Cargo.toml +++ b/lib/sidecar/common/Cargo.toml @@ -11,13 +11,25 @@ homepage.workspace = true repository.workspace = true description = "Shared gRPC transport and CLI infrastructure for Dynamo sidecars" +[package.metadata.cargo-machete] +# Referenced only as `::prost_types_v14` inside the `json_codec!` macro body, +# which cargo-machete's source scan does not expand. +ignored = ["prost-types-v14"] + [dependencies] dynamo-backend-common = { workspace = true } clap = { version = "4", features = ["derive", "env"] } futures = { workspace = true } +prost-types = { workspace = true } +prost-types-v14 = { workspace = true, optional = true } +serde_json = { workspace = true } tokio = { workspace = true } tonic = { workspace = true } tonic-v14 = { workspace = true, optional = true } tracing = { workspace = true } url = { workspace = true } + +[features] +# The v14 protobuf/tonic surface, for sidecars generated against prost 0.14. +tonic-v14 = ["dep:tonic-v14", "dep:prost-types-v14"] diff --git a/lib/sidecar/common/src/error.rs b/lib/sidecar/common/src/error.rs index 27173d48c208..063cb9b9bd24 100644 --- a/lib/sidecar/common/src/error.rs +++ b/lib/sidecar/common/src/error.rs @@ -80,6 +80,17 @@ pub fn connection_timeout(message: impl Into) -> DynamoError { backend(BackendError::ConnectionTimeout, message) } +fn typed(kind: ErrorType, message: impl Into) -> DynamoError { + DynamoError::builder() + .error_type(kind) + .message(message) + .build() +} + +pub fn cancelled(message: impl Into) -> DynamoError { + typed(ErrorType::Cancelled, message) +} + pub fn status_to_dynamo(rpc: &str, status: tonic::Status) -> DynamoError { status_to_dynamo_parts(rpc, status.message(), status.code()) } @@ -125,6 +136,12 @@ mod tests { BackendError::ConnectionTimeout, ), (tonic::Code::Internal, BackendError::Unknown), + // Not WorkerOverloaded: that type is migratable, and a shared + // helper cannot tell an overloaded engine from a per-request + // rejection such as an oversized message, which every worker would + // reject identically. A backend that knows its server means + // backpressure maps this itself. + (tonic::Code::ResourceExhausted, BackendError::Unknown), ] { let error = status_to_dynamo("Test", tonic::Status::new(code, "failure")); assert_eq!(error.error_type(), ErrorType::Backend(expected)); diff --git a/lib/sidecar/common/src/json.rs b/lib/sidecar/common/src/json.rs new file mode 100644 index 000000000000..ab16866f0b89 --- /dev/null +++ b/lib/sidecar/common/src/json.rs @@ -0,0 +1,224 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `google.protobuf.Struct` <-> `serde_json::Value` conversion. +//! +//! Engine gRPC contracts carry their opaque, engine-specific payloads (vLLM's +//! `kv_transfer_params`, the OpenEngine `KvSessionRef.attributes_struct`) as a +//! `Struct`, while Dynamo carries them as JSON. Sidecars share this codec so a +//! fix to the number or null handling lands for all of them at once. +//! +//! `what` names the payload in error messages; `peer` names the engine whose +//! response failed to decode. +//! +//! The sidecars do not agree on a protobuf version -- vLLM generates against +//! prost 0.14 while TensorRT-LLM and SGLang are on 0.13 -- and the two +//! `prost_types::Struct` types are unrelated. Rather than keep a copy of the +//! codec per version, [`json_codec`] emits it once per `prost_types` path, the +//! same shape `error::status_to_dynamo_v14` uses for the equivalent problem. + +/// Emits the `Struct` <-> JSON codec against one `prost_types` crate. +macro_rules! json_codec { + ($prost_types:path) => { + use $prost_types as prost_types; + + use dynamo_backend_common::DynamoError; + use prost_types::value::Kind; + + use $crate::error::{invalid_argument, protocol_error}; + + /// A `Struct` number is an IEEE-754 double, so integers above 2^53 + /// cannot round-trip exactly. Reject them rather than silently + /// truncate. + const MAX_EXACT_INTEGER: u64 = 1_u64 << 53; + + /// Encodes a JSON object as a protobuf `Struct`. + pub fn json_to_struct( + value: serde_json::Value, + what: &str, + ) -> Result { + let serde_json::Value::Object(fields) = value else { + return Err(invalid_argument(format!("{what} must be a JSON object"))); + }; + Ok(prost_types::Struct { + fields: fields + .into_iter() + .map(|(key, value)| Ok((key, json_to_value(value, what)?))) + .collect::>()?, + }) + } + + fn json_to_value( + value: serde_json::Value, + what: &str, + ) -> Result { + let kind = match value { + serde_json::Value::Null => { + Kind::NullValue(prost_types::NullValue::NullValue as i32) + } + serde_json::Value::Bool(value) => Kind::BoolValue(value), + serde_json::Value::String(value) => Kind::StringValue(value), + serde_json::Value::Number(value) => Kind::NumberValue(number_to_f64(&value, what)?), + serde_json::Value::Array(values) => Kind::ListValue(prost_types::ListValue { + values: values + .into_iter() + .map(|value| json_to_value(value, what)) + .collect::>()?, + }), + serde_json::Value::Object(values) => Kind::StructValue(prost_types::Struct { + fields: values + .into_iter() + .map(|(key, value)| Ok((key, json_to_value(value, what)?))) + .collect::>()?, + }), + }; + Ok(prost_types::Value { kind: Some(kind) }) + } + + fn number_to_f64(value: &serde_json::Number, what: &str) -> Result { + if let Some(value) = value.as_u64() + && value > MAX_EXACT_INTEGER + { + return Err(invalid_argument(format!( + "{what} integer {value} cannot be represented exactly by protobuf Struct" + ))); + } + if let Some(value) = value.as_i64() + && value.unsigned_abs() > MAX_EXACT_INTEGER + { + return Err(invalid_argument(format!( + "{what} integer {value} cannot be represented exactly by protobuf Struct" + ))); + } + value.as_f64().ok_or_else(|| { + invalid_argument(format!( + "{what} number {value} cannot be represented by protobuf Struct" + )) + }) + } + + /// Decodes a protobuf `Struct` from `peer` into a JSON object. + pub fn struct_to_json( + value: prost_types::Struct, + peer: &str, + what: &str, + ) -> Result { + Ok(serde_json::Value::Object( + value + .fields + .into_iter() + .map(|(key, value)| Ok((key, value_to_json(value, peer, what)?))) + .collect::>()?, + )) + } + + fn value_to_json( + value: prost_types::Value, + peer: &str, + what: &str, + ) -> Result { + match value.kind { + None | Some(Kind::NullValue(_)) => Ok(serde_json::Value::Null), + Some(Kind::BoolValue(value)) => Ok(serde_json::Value::Bool(value)), + Some(Kind::StringValue(value)) => Ok(serde_json::Value::String(value)), + Some(Kind::NumberValue(value)) => { + if !value.is_finite() { + return Err(protocol_error( + peer, + format!("{what} contains NaN or infinity"), + )); + } + // Render whole doubles as JSON integers so a round-trip + // does not turn `dp_rank: 0` into `0.0`. + let number = if value.fract() == 0.0 && value.abs() <= MAX_EXACT_INTEGER as f64 + { + if value.is_sign_negative() { + serde_json::Number::from(value as i64) + } else { + serde_json::Number::from(value as u64) + } + } else { + serde_json::Number::from_f64(value).expect("finite f64 is a JSON number") + }; + Ok(serde_json::Value::Number(number)) + } + Some(Kind::ListValue(values)) => Ok(serde_json::Value::Array( + values + .values + .into_iter() + .map(|value| value_to_json(value, peer, what)) + .collect::>()?, + )), + Some(Kind::StructValue(value)) => struct_to_json(value, peer, what), + } + } + }; +} + +/// The codec for sidecars generated against prost 0.13. +mod v13 { + json_codec!(::prost_types); +} + +/// The codec for sidecars generated against prost 0.14. +#[cfg(feature = "tonic-v14")] +mod v14 { + json_codec!(::prost_types_v14); +} + +pub use v13::{json_to_struct, struct_to_json}; + +#[cfg(feature = "tonic-v14")] +pub use v14::{json_to_struct as json_to_struct_v14, struct_to_json as struct_to_json_v14}; + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{json_to_struct, struct_to_json}; + + #[test] + fn nested_payload_round_trips_without_shape_changes() { + let payload = json!({ + "string": "value", + "bool": true, + "number": 42, + "null": null, + "list": [1, "two", false, {"nested": 3.5}], + }); + let encoded = json_to_struct(payload.clone(), "payload").expect("encode"); + assert_eq!( + struct_to_json(encoded, "test-engine", "payload").expect("decode"), + payload + ); + } + + #[test] + fn rejects_non_objects_and_inexact_integers() { + assert!(json_to_struct(json!([1, 2]), "payload").is_err()); + assert!(json_to_struct(json!({"value": 9_007_199_254_740_993_u64}), "payload").is_err()); + } + + /// The two instantiations are one implementation, so a payload encoded by + /// either decodes to the same JSON. Without this, a divergence between the + /// protobuf versions would only show up in a sidecar's own tests. + #[cfg(feature = "tonic-v14")] + #[test] + fn both_protobuf_versions_agree() { + let payload = json!({"dp_rank": 0, "ttft_ms": 12.5, "tags": ["a", null]}); + let v13 = struct_to_json( + json_to_struct(payload.clone(), "payload").expect("encode"), + "test-engine", + "payload", + ) + .expect("decode"); + let v14 = super::struct_to_json_v14( + super::json_to_struct_v14(payload.clone(), "payload").expect("encode"), + "test-engine", + "payload", + ) + .expect("decode"); + assert_eq!(v13, payload); + assert_eq!(v14, payload); + } +} diff --git a/lib/sidecar/common/src/lib.rs b/lib/sidecar/common/src/lib.rs index d6e455dce8e6..296ec195f518 100644 --- a/lib/sidecar/common/src/lib.rs +++ b/lib/sidecar/common/src/lib.rs @@ -6,6 +6,7 @@ mod args; mod endpoint; mod error; +mod json; mod transport; #[cfg(feature = "tonic-v14")] @@ -21,7 +22,12 @@ pub mod v14 { pub use args::{GrpcTransportArgs, GrpcTransportConfig, SidecarArgs}; pub use endpoint::{GrpcEndpoint, HttpEndpoint}; pub use error::{ - SidecarStartupError, cannot_connect, connection_timeout, engine_shutdown, invalid_argument, - protocol_error, status_to_dynamo, + SidecarStartupError, cancelled, cannot_connect, connection_timeout, engine_shutdown, + invalid_argument, protocol_error, status_to_dynamo, +}; +pub use json::{json_to_struct, struct_to_json}; +#[cfg(feature = "tonic-v14")] +pub use json::{json_to_struct_v14, struct_to_json_v14}; +pub use transport::{ + DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPool, format_error_chain, startup_deadline, }; -pub use transport::{DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPool, format_error_chain}; diff --git a/lib/sidecar/common/src/transport.rs b/lib/sidecar/common/src/transport.rs index 9ef2f046e425..162df4e411c7 100644 --- a/lib/sidecar/common/src/transport.rs +++ b/lib/sidecar/common/src/transport.rs @@ -203,6 +203,15 @@ async fn connect_until_ready( } } +/// The instant by which startup work -- connecting, waiting for the engine to +/// answer, discovering the model -- must be done. Every sidecar bounds its whole +/// startup path by this one deadline rather than by per-RPC constants, so +/// `--grpc-startup-deadline-secs` actually governs how long a slow-loading engine is +/// given. +pub fn startup_deadline(duration: Duration) -> Result { + checked_instant_add(Instant::now(), duration, "gRPC startup deadline") +} + fn checked_instant_add( instant: Instant, duration: Duration, diff --git a/lib/sidecar/trtllm/Cargo.toml b/lib/sidecar/trtllm/Cargo.toml index c308e4b4a923..0c3aa851fc1d 100644 --- a/lib/sidecar/trtllm/Cargo.toml +++ b/lib/sidecar/trtllm/Cargo.toml @@ -9,11 +9,11 @@ authors.workspace = true license.workspace = true homepage.workspace = true repository.workspace = true -description = "Dynamo TensorRT-LLM sidecar backend — drives an out-of-process TensorRT-LLM engine over its native gRPC service through the backend-common LLMEngine trait." +description = "Dynamo TensorRT-LLM sidecar backend — drives an out-of-process TensorRT-LLM engine over its OpenEngine gRPC service through the backend-common LLMEngine trait." [package.metadata.cargo-machete] -# Referenced by the protobuf client generated into OUT_DIR at build time. -ignored = ["prost"] +# Referenced by protobuf code generated into OUT_DIR. +ignored = ["prost", "prost-types"] [[bin]] name = "dynamo-trtllm-sidecar" @@ -28,18 +28,18 @@ async-stream = { workspace = true } async-trait = { workspace = true } clap = { version = "4", features = ["derive", "env"] } futures = { workspace = true } +serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } tokio = { workspace = true } tokio-util = { workspace = true } tracing = { workspace = true } -# tonic / prost are pinned per-crate (not workspace deps): tonic 0.13.1 -# depends on prost 0.13.5, matching `lib/llm` and the sglang sidecar. -prost = { version = "0.13.5" } -tonic = { version = "0.13.1" } +prost = { workspace = true } +prost-types = { workspace = true } +tonic = { workspace = true } [build-dependencies] -tonic-build = { version = "0.13.1" } +tonic-build = { workspace = true } [dev-dependencies] dynamo-backend-common = { workspace = true, features = ["testing"] } diff --git a/lib/sidecar/trtllm/README.md b/lib/sidecar/trtllm/README.md index 6a594aa5989d..bbee38c63a0f 100644 --- a/lib/sidecar/trtllm/README.md +++ b/lib/sidecar/trtllm/README.md @@ -6,44 +6,105 @@ SPDX-License-Identifier: Apache-2.0 # TensorRT-LLM sidecar > [!WARNING] -> **Experimental.** This sidecar and its deployment example are experimental. -> The Python launcher ships in the `ai-dynamo` and `ai-dynamo-runtime` wheel -> pair, but the container image is not yet packaged for distribution. The -> manifest, flags, and behavior may change without notice. - -`dynamo-trtllm-sidecar` connects a Dynamo worker to TensorRT-LLM's native -`trtllm.TrtllmService` gRPC `Generate` service. It is a standalone Rust -executable composed with `dynamo_backend_common::run` and is also compiled into -`ai-dynamo-runtime` for the importable `dynamo.trtllm.sidecar` launcher. +> **Experimental.** This sidecar and its deployment example are experimental and +> not yet packaged for distribution (see [Packaging](#packaging)). The manifest, +> flags, and behavior may change without notice. + +`dynamo-trtllm-sidecar` connects a Dynamo worker to TensorRT-LLM's OpenEngine +(`openengine.v1`) gRPC server — the `Inference.Generate` streaming RPC. It is a +standalone Rust executable composed with `dynamo_backend_common::run`, and is +also compiled into `ai-dynamo-runtime` for the importable +`dynamo.trtllm.sidecar` launcher: TensorRT-LLM runs as its own process while the sidecar owns Dynamo worker registration, request conversion, transport, cancellation, and abort. ## Supported - Aggregated generation +- Disaggregated (prefill/decode) serving — see [Disaggregation](#disaggregation) - Token requests through Dynamo preprocessing - Sampling, stop conditions, structured output (JSON schema / regex / grammar / structural tag), and logprobs - Streaming delta tokens with a terminal usage/finish summary -- `Abort` on cancellation - -The initial protocol does **not** support disaggregated (prefill/decode) -serving, multimodal input, LoRA, KV-aware routing, encode workers, beam search, -or `n > 1`. Disaggregation is excluded because the `Generate` response contract -carries no context-phase handoff. +- Cancellation via `Control.Abort` and by closing the gRPC stream + +The integration does **not** support multimodal input, LoRA, encode workers, +beam search, or `n > 1`. + +Data-parallel rank targeting is rejected: the server answers +`openengine-target-dp-rank` with `UNIMPLEMENTED`, so a request carrying a rank +hint is refused up front instead of failing in the engine. +`KvSessionRef.dp_rank` still carries a disaggregated session's KV affinity, +inside the request body. + +> [!NOTE] +> `Control.GetModelInfo` supplies the registered context length (and the default +> `max_tokens` for requests that omit one) unless `--context-length` supplies it +> instead. `Control.Abort` cancels an in-flight request; closing the `Generate` +> stream also aborts it, so cancellation is covered either way. +> Unclaimed prefill KV handoffs require separate cleanup; see [Known issues](#known-issues). +> +> `Control`'s LoRA RPCs (`LoadLora`, `UnloadLora`, `ListLoras`) and KV-event +> RPCs (`GetKvEventSources`, `SubscribeKvEvents`) return `UNIMPLEMENTED`: the +> LLM API has no runtime adapter load/unload entry point, and KV events are +> published out of band. The sidecar uses neither. + +## Protocol + +The gRPC types are vendored, like the vLLM and SGLang sidecars': `proto/` +carries the `openengine.v1` contract from +[`ai-dynamo/openengine`](https://github.com/ai-dynamo/openengine) `v0.1.0`, +compiled by `build.rs` with `tonic-build`. The pinned revision is the git commit +behind the Buf Schema Registry module commit (`768a93c7b44e`) TensorRT-LLM's +server is generated from, so the two sides agree; `proto/README.md` records the +commit and per-file SHA-256. + +Building needs only `protoc`, which the workspace already requires for `lib/llm` +and the other sidecars — there is no registry to configure and no token to +obtain. + +To bump the protocol, re-copy `proto/openengine/v1/*.proto` from a newer upstream +commit, update the revision, checksums, and `build.rs`'s `PROTOS` list together, +and re-run the tests. ## Run -Start TensorRT-LLM with its native gRPC listener. Its protobuf bindings are an -optional extra, so install them first. Install the package directly rather than -via `tensorrt_llm[grpc-smg]`: the extra makes pip re-resolve TensorRT-LLM's whole -dependency closure, which fails on an image whose site-packages is read-only. +Start TensorRT-LLM with its OpenEngine gRPC server. Published releases carry the +server from `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc27.dev202609170000` +onward; earlier tags have no OpenEngine servicer at all. + +The bindings are a separate install even on those releases. The servicer imports +`openengine.v1` directly, so a stock image refuses to start the listener: + +```text +Error: Failed to import OpenEngine support: No module named 'openengine'. +``` ```bash -pip install "smg-grpc-proto>=0.4.2" -python -m tensorrt_llm.commands.serve --grpc --host 0.0.0.0 --port 50051 +# Install the two packages directly rather than through the +# `tensorrt_llm[openengine]` extra: the extra makes pip re-resolve +# TensorRT-LLM's whole dependency closure, which fails on an image whose +# site-packages is read-only. Both are pinned to BSR module commit +# 768a93c7b44e, the revision `proto/` was generated from. The protobuf package +# is additionally pinned by gencode version -- a gencode newer than the image's +# protobuf runtime fails at import -- so raise it only with the image's +# protobuf. The `-pyi` package is pinned because the protobuf wheel requires it +# with no upper bound and the name is unregistered on PyPI, which under +# `--extra-index-url` would let anyone claiming it there run code here. +python -m pip install --extra-index-url https://buf.build/gen/python \ + "openengine-openengine-grpc-python==1.78.1.1.20260730172104+768a93c7b44e" \ + "openengine-openengine-protocolbuffers-python==33.5.0.1.20260730172104+768a93c7b44e" \ + "openengine-openengine-protocolbuffers-pyi==36.2.0.1.20260730172104+768a93c7b44e" + +python -m tensorrt_llm.commands.serve \ + --grpc --grpc-protocol openengine --host 127.0.0.1 --port 50051 \ + --max_seq_len 4096 ``` +Without `--max_seq_len` the servicer leaves `max_context_length` unset. The +sidecar below passes `--context-length`, so it starts either way; with neither +setting it fails on the engine's first answer, naming both fixes. + This listener is unauthenticated and plaintext. Keep colocated deployments on loopback or a private interface. Remote access requires network controls or a secure proxy. @@ -53,116 +114,123 @@ Start the Dynamo worker: ```bash dynamo-trtllm-sidecar \ --grpc-endpoint 127.0.0.1:50051 \ - --model-path + --model-path \ + --context-length 4096 ``` -After installing `ai-dynamo`, the Python module runs the same native worker: - -```bash -python -m dynamo.trtllm.sidecar \ - --grpc-endpoint 127.0.0.1:50051 \ - --model-path -``` +The context length comes from `--context-length` (or `TRTLLM_CONTEXT_LENGTH`) +when supplied, and from `Control.GetModelInfo` otherwise; a disagreement is +logged at WARN and the configured value wins. Supply it whenever the engine was +started without `--max_seq_len`, because TensorRT-LLM then leaves +`max_context_length` unset. With neither source the sidecar fails as soon as the +engine answers, rather than registering a worker that would reject every request +omitting `max_tokens`. It does not wait out `--grpc-startup-deadline-secs` for +that case: the engine binds its gRPC port only once the model has loaded, so an +engine that answers is not going to start reporting a window. That deadline +covers an engine that is not up yet, where only transport failures are retried. Use `DYN_SIDECAR_GRPC_ENDPOINT` instead of `--grpc-endpoint` when the endpoint is provided through the environment. -## Deploy on Kubernetes (quick start) - -`deploy/agg.yaml` deploys a frontend and one worker pod. The worker runs the -sidecar next to a TensorRT-LLM engine and serves `Qwen/Qwen3-0.6B` on one GPU. +## Disaggregation -There is no published sidecar image yet (see [Packaging](#packaging)), so build -and push the image from `lib/sidecar/Dockerfile`. It contains all three -engine-specific sidecar executables; this manifest runs `dynamo-trtllm-sidecar` -as the container command. +Prefill and decode run as two workers, each with its own TensorRT-LLM engine and +its own sidecar, selected with `--disaggregation-mode`: -### Prerequisites - -- A Kubernetes cluster (**v1.29+**, or v1.28 with the `SidecarContainers` feature - gate) with the Dynamo operator and a GPU node. The engine runs as a native - sidecar (`initContainers` with `restartPolicy: Always`), which requires that - version. -- `kubectl` set to that cluster, and a namespace to deploy into. -- A Hugging Face token for the model. -- A container registry you can push to and the cluster can pull from. +```bash +# Prefill (context) worker — registers under the `prefill` component. +dynamo-trtllm-sidecar --disaggregation-mode prefill \ + --grpc-endpoint 127.0.0.1:50051 --model-path -### 1. Build and push the sidecar image +# Decode (generation) worker. +dynamo-trtllm-sidecar --disaggregation-mode decode \ + --grpc-endpoint 127.0.0.1:50052 --model-path +``` -Build and push the image to a registry your cluster can pull from: +`launch/disagg.sh` brings up the whole topology on two GPUs — frontend, both +engines with a cache transceiver, and both sidecars: ```bash -docker buildx build --platform linux/amd64,linux/arm64 \ - -f lib/sidecar/Dockerfile \ - -t /dynamo-sidecar:1.3.0 --push . +./lib/sidecar/trtllm/launch/disagg.sh --model ``` -See [Build the image](../README.md#build-the-image) for a single-architecture -build. This manifest sets the container `command` to -`dynamo-trtllm-sidecar`. +Run `--help` for the ports, GPU assignment, and transceiver backend it accepts. -### 2. Point the manifest at your image +Both engines must be started with a KV cache transceiver so they can move KV +cache between themselves (`cache_transceiver_config`); without it the engines +cannot complete the handoff. Use the default `NIXL` backend — it picks its own +underlying transport (UCX where there is no RDMA fabric) and is the path Dynamo +uses elsewhere for disaggregation. -In `deploy/agg.yaml`, set the `main` worker image to the one you just pushed. -If your registry is private, add `imagePullSecrets` to the worker pod spec. +OpenEngine has no request-type field, so the phase is carried on the wire like +this: -### 3. Create the Hugging Face token secret +- The prefill worker sets `extra.request_type = "context_only"` and caps + generation at one token. The server answers with a terminal `PrefillReady` + event holding a `KvSessionRef`; there is no `finished` event for a context + request. +- The sidecar encodes that `KvSessionRef` as the opaque JSON Dynamo carries in + `PrefillResult.disaggregated_params`, and emits it on the prefill worker's + terminal chunk. The prefill worker streams no tokens to the client. +- The decode worker decodes that JSON back into `kv.session` on its own + `Generate` request, which the server maps to `generation_only`. It streams the + full completion and reports the authoritative usage. -Read the token from an env var so it stays out of your shell history (or use -`--from-file` / an external secret manager): +The handoff JSON mirrors `KvSessionRef` field-for-field (`session_id`, +`transfer_backend`, `endpoints`, `dp_rank`, `attributes`) and is never +interpreted between the two workers. See `src/disagg.rs`. -```bash -kubectl create secret generic hf-token-secret \ - --from-literal=HF_TOKEN="$HF_TOKEN" -n -``` +## Known issues -### 4. Deploy +In disaggregated serving, decode-side validation can reject a request after prefill has produced a KV handoff, for example when `min_tokens` cannot fit in the decode context window. This rejection happens before the decode `Generate` RPC, so no decode request or KV transfer starts. The unclaimed prefill blocks remain reserved until the engine transfer timeout (60 seconds in the engine version pinned in [Run](#run)), which can delay subsequent valid requests when KV cache capacity is exhausted. -```bash -kubectl apply -f lib/sidecar/trtllm/deploy/agg.yaml -n -``` +Tracked in [#15404: abort abandoned prefill KV sessions after decode rejection](https://github.com/ai-dynamo/dynamo/issues/15404). The intended fix requires coordinated TensorRT-LLM and sidecar support for `Control.Abort(kv_session)`, with the engine releasing blocks only after transfers can no longer access them. That engine version returns `UNIMPLEMENTED` for session abort; aborting the completed prefill request by request ID does not release its pending handoff. Until that follow-up is delivered, cleanup relies on the engine transfer timeout. -Wait for the worker pod to reach `2/2 Running`: +## Deploy on Kubernetes -```bash -kubectl get pods -n -w -``` +`deploy/agg.yaml` runs a frontend and one worker pod serving `Qwen/Qwen3-0.6B` +on one GPU. `deploy/disagg.yaml` runs prefill and decode as separate worker +pods. Read the disaggregated manifest's header before applying it: it runs the +engines over TCP/CUDA-IPC and requests no `rdma/ib`, which you add on a fabric +that provides it. -### 5. Send a request +You need a cluster on **v1.29+** (or v1.28 with the `SidecarContainers` gate) +with the Dynamo operator and a GPU node — the engine runs as a native sidecar — +plus a namespace, a registry, and two images: -Port-forward the frontend and call it: +- **The sidecar image.** Not published yet (see [Packaging](#packaging)), so + build it from `lib/sidecar/Dockerfile`; it carries all three engine-specific + executables. Multi-arch if your nodes are mixed: + `docker buildx build --platform linux/amd64,linux/arm64 -f lib/sidecar/Dockerfile -t /dynamo-sidecar:1.3.0 --push .` +- **The engine image.** `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc27.dev202609170000` + or newer, with the pinned bindings from [Run](#run) layered on: the release + ships the servicer but not the `openengine.v1` package that both it and the + manifests' health probes import. -```bash -kubectl port-forward -n svc/trtllm-sidecar-agg-frontend 8000:8000 & - -curl -s localhost:8000/v1/models | jq . +Set both images in the manifest, add `imagePullSecrets` for a private registry, +then: -curl -s localhost:8000/v1/chat/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"Qwen/Qwen3-0.6B","messages":[{"role":"user","content":"Hello"}],"max_tokens":32}' | jq . +```bash +kubectl create secret generic hf-token-secret --from-literal=HF_TOKEN="$HF_TOKEN" -n +kubectl apply -f lib/sidecar/trtllm/deploy/agg.yaml -n +kubectl get pods -n -w # every worker pod reaches 2/2 ``` -`/v1/models` should list `Qwen/Qwen3-0.6B`, and the chat call returns a reply. +Then port-forward `svc/trtllm-sidecar-agg-frontend` (or +`svc/trtllm-sidecar-disagg-frontend`) on 8000 and call `/v1/models` and +`/v1/chat/completions`. ## Tuning -The engine streams tokens to the sidecar over gRPC. By default it sends one -message per token, and that per-token serialization is the sidecar's main -throughput cost versus an in-process backend. The `trtllm-engine-config` -ConfigMap in `deploy/agg.yaml` sets `stream_interval`, which emits one chunk per -`N` decode steps instead: - -- Higher `N` → fewer, larger gRPC messages → higher throughput under load. -- Trade-off: the client receives tokens in bursts of `N`. - -On a single GB200 (Qwen3-0.6B, 2000-in / 256-out) raising `stream_interval` from -1 to 5 roughly doubled output throughput at high concurrency (~6.3k → ~12k -tok/s) and even lowered TTFT. `5` keeps streaming smooth while capturing nearly -all the gain. +Per-token gRPC serialization is the sidecar's main throughput cost versus an +in-process backend. `stream_interval: N` emits one chunk per `N` decode steps +instead of one per token, trading burstier delivery for throughput. On a single +GB200 (Qwen3-0.6B, 2000-in / 256-out) raising it from 1 to 5 roughly doubled +output throughput at high concurrency (~6.3k → ~12k tok/s) and lowered TTFT. +The manifests set `5`; `launch/disagg.sh` leaves it at the engine default. ## Packaging -There is no published sidecar image yet. The image contains the vLLM, -SGLang, and TensorRT-LLM executables and uses a minimal CPU-only base. Until -official packaging is available, build and push the sidecar image as described -in [Build the image](../README.md#build-the-image). +No published sidecar image yet. Build and push it yourself — see +[Build the image](../README.md#build-the-image). One minimal CPU-only image +carries the vLLM, SGLang, and TensorRT-LLM executables. diff --git a/lib/sidecar/trtllm/build.rs b/lib/sidecar/trtllm/build.rs index 86da65b15fe8..c235a1649c90 100644 --- a/lib/sidecar/trtllm/build.rs +++ b/lib/sidecar/trtllm/build.rs @@ -1,23 +1,41 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Compiles client stubs from the vendored TensorRT-LLM gRPC contract. +//! Compiles client stubs from the vendored OpenEngine (`openengine.v1`) contract. use std::env; use std::path::PathBuf; +/// The `openengine.v1` files, in `proto/openengine/v1/`. Listed explicitly so a +/// new upstream file has to be vendored deliberately rather than picked up by a +/// glob. +const PROTOS: &[&str] = &[ + "error.proto", + "generation.proto", + "kv.proto", + "lifecycle.proto", + "lora.proto", + "model.proto", + "openengine.proto", + "server.proto", +]; + fn main() -> Result<(), Box> { - let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR")?); - let proto_dir = manifest_dir.join("proto"); - let proto_path = proto_dir.join("trtllm_service.proto"); + let proto_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR")?).join("proto"); + let paths: Vec = PROTOS + .iter() + .map(|name| proto_dir.join("openengine/v1").join(name)) + .collect(); // `--experimental_allow_proto3_optional` lets older protoc compile the - // proto3 `optional` fields (e.g. `top_k`, `temperature`). The flag was - // added in protoc 3.12 for exactly this and is a no-op on 3.15+. + // proto3 `optional` fields (e.g. `max_context_length`). The flag was added + // in protoc 3.12 for exactly this and is a no-op on 3.15+. tonic_build::configure() .protoc_arg("--experimental_allow_proto3_optional") - .compile_protos(&[proto_path.as_path()], &[proto_dir.as_path()])?; + .compile_protos(&paths, &[proto_dir.as_path()])?; - println!("cargo:rerun-if-changed={}", proto_path.display()); + for path in &paths { + println!("cargo:rerun-if-changed={}", path.display()); + } Ok(()) } diff --git a/lib/sidecar/trtllm/deploy/agg.yaml b/lib/sidecar/trtllm/deploy/agg.yaml index 71d4f1a116f7..c158fd953d28 100644 --- a/lib/sidecar/trtllm/deploy/agg.yaml +++ b/lib/sidecar/trtllm/deploy/agg.yaml @@ -6,8 +6,8 @@ # Worker pod = two containers: # - main: Dynamo worker (dynamo-trtllm-sidecar), no GPU; talks to the engine # over gRPC on 127.0.0.1. Must be named "main" (operator injects env + probes). -# - trtllm-engine: stock TRT-LLM image, holds the GPU, runs the gRPC server as -# a native sidecar (see initContainers). +# - trtllm-engine: TRT-LLM image with OpenEngine gRPC support, holds the GPU, +# runs the OpenEngine gRPC server as a native sidecar (see initContainers). # # The engine's gRPC has no auth/TLS; safe only because it stays on loopback. @@ -45,16 +45,22 @@ spec: # starts before `main`, stops after. initContainers: - name: trtllm-engine - image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc25 + # Layer the OpenEngine Python bindings onto this release and push the + # result: the release ships the servicer, but both it and the probes + # below import `openengine.v1`, which it does not install. See the + # README's Run section for the pinned versions. + image: /tensorrt-llm-openengine:1.3.0rc27.dev202609170000 restartPolicy: Always - # gRPC HealthCheck via exec (tcpSocket/grpc: dial the pod IP, not the - # loopback engine). startupProbe gates `main`; readinessProbe drops dead workers. + # Readiness via the OpenEngine Control.Health RPC, which reports the + # engine's own state (grpc + model, plus kv_connector when configured) + # rather than just a connectable socket. + # startupProbe gates `main`; readinessProbe drops dead workers. startupProbe: exec: command: - python3 - -c - - import grpc; grpc.insecure_channel('127.0.0.1:50051').unary_unary('/trtllm.TrtllmService/HealthCheck')(b'', timeout=3) + - import grpc; from openengine.v1 import lifecycle_pb2, openengine_pb2_grpc; s=openengine_pb2_grpc.ControlStub(grpc.insecure_channel('127.0.0.1:50051')); r=s.Health(lifecycle_pb2.HealthRequest(), timeout=3); exit(0 if r.state==lifecycle_pb2.HEALTH_STATE_READY else 1) periodSeconds: 5 timeoutSeconds: 10 # must exceed the RPC timeout + interpreter spawn failureThreshold: 60 # ~5 min for model load; raise for big models @@ -63,10 +69,34 @@ spec: command: - python3 - -c - - import grpc; grpc.insecure_channel('127.0.0.1:50051').unary_unary('/trtllm.TrtllmService/HealthCheck')(b'', timeout=3) + - import grpc; from openengine.v1 import lifecycle_pb2, openengine_pb2_grpc; s=openengine_pb2_grpc.ControlStub(grpc.insecure_channel('127.0.0.1:50051')); r=s.Health(lifecycle_pb2.HealthRequest(), timeout=3); exit(0 if r.state==lifecycle_pb2.HEALTH_STATE_READY else 1) periodSeconds: 10 timeoutSeconds: 10 failureThreshold: 3 + env: + # TensorRT-LLM would otherwise start its TP=1 worker by spawning an + # MPI process, and that spawn fails inside a pod: PRTE reaches + # `VM READY` and then terminates its daemons with no diagnostic. + # Measured on release:1.3.0rc27. This worker runs one GPU, which is + # exactly the case the switch covers + # (`enable_worker_single_process_for_tp1`), so the worker runs + # in-process and never spawns. + # Not free: the engine logs "Using single process worker for TP1, + # this may hurt streaming generation performance" + # (`executor/executor.py`), because the worker now shares the server's + # process instead of gathering over IPC. `stream_interval` in the + # ConfigMap above is the lever that buys that back -- see the README's + # Tuning section. Revisit both together if the engine needs TP > 1, + # which this switch does not cover. + - name: TLLM_WORKER_USE_SINGLE_PROCESS + value: "1" + # The release image runs as uid 0 and PRTE refuses to run as root + # unless told twice. Harmless with the switch above, and it keeps the + # engine from dying at startup if the spawn path is ever re-entered. + - name: PRTE_ALLOW_RUN_AS_ROOT + value: "1" + - name: PRTE_ALLOW_RUN_AS_ROOT_CONFIRM + value: "1" envFrom: - secretRef: name: hf-token-secret @@ -76,27 +106,31 @@ spec: requests: # Increase for larger models. ephemeral-storage: 2Gi - # `--grpc` needs `smg-grpc-proto`, which TRT-LLM keeps behind the - # optional `grpc-smg` extra rather than shipping in the image. - # Constraint copied from that extra so we resolve what upstream - # resolves. Needs egress to PyPI; bake it into a derived image if the - # cluster has none. command: - - sh - - -c - - | - set -e - python3 -c "import smg_grpc_proto" 2>/dev/null \ - || pip install --no-cache-dir "smg-grpc-proto>=0.4.2" - exec python3 -m tensorrt_llm.commands.serve "$@" - - trtllm-serve + - python3 + - -m + - tensorrt_llm.commands.serve args: - Qwen/Qwen3-0.6B - --grpc + - --grpc-protocol + - openengine - --host - 127.0.0.1 - --port - "50051" + # Workaround, not a requirement: `GetModelInfo` reads the unresolved + # `llm.args.max_seq_len`, which stays `None` on the PyTorch path even + # though the executor has already resolved the real window from the + # model config. Measured on release:1.3.0rc27.dev202609170000 with + # TinyLlama-1.1B: the executor runs at `max_seq_len=2048` while + # `GetModelInfo` answers UNSET. Stating it here is what stops the + # sidecar polling for its whole startup deadline and then registering + # a worker that rejects every request omitting max_tokens. Drop this + # and the sidecar's --context-length once the servicer reports the + # resolved window. + - --max_seq_len + - "4096" # Engine tuning from the ConfigMap above. - --extra_llm_api_options - /config/config.yaml @@ -117,10 +151,9 @@ spec: - 127.0.0.1:50051 - --model-path - Qwen/Qwen3-0.6B - # Engine reports 0 for context length; set it here or requests that - # omit max_tokens are rejected. + # Matches the engine's --max_seq_len above. - --context-length - - "40960" + - "4096" envFrom: - secretRef: name: hf-token-secret diff --git a/lib/sidecar/trtllm/deploy/disagg.yaml b/lib/sidecar/trtllm/deploy/disagg.yaml new file mode 100644 index 000000000000..01abbf30c507 --- /dev/null +++ b/lib/sidecar/trtllm/deploy/disagg.yaml @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# EXPERIMENTAL. Disaggregated TensorRT-LLM deployment using the Rust sidecar. +# +# Prefill and decode run as separate worker pods. Each pod = two containers: +# - main: Dynamo worker (dynamo-trtllm-sidecar), no GPU; talks to its engine +# over gRPC on 127.0.0.1. Its --disaggregation-mode (prefill|decode) also +# picks the component it registers under. Must be named "main" (operator +# injects env + probes). +# - trtllm-engine: TRT-LLM image with OpenEngine gRPC support, holds the GPU, +# runs the OpenEngine gRPC server as a native sidecar (see initContainers). +# +# The frontend routes to a prefill worker, gets back a `KvSessionRef` handoff, +# then routes to a decode worker carrying it. The KV cache itself moves +# engine-to-engine over the cache transceiver, never through the sidecars, so +# both engines need `cache_transceiver_config` (the ConfigMaps below). +# +# Use the default router, not --router-mode kv. Missing KV events would only +# make KV routing ineffective; the blocker is harder than that. KV mode stamps +# routing.dp_rank on every request -- Some(0) for a single-rank worker -- and +# this sidecar rejects any dp-rank targeting, because the OpenEngine servicer +# answers `openengine-target-dp-rank` with UNIMPLEMENTED. So every KV-routed +# request fails, and adding KV events later will not change that on its own. +# +# The engine's gRPC has no auth/TLS; safe only because it stays on loopback. +# The cache transceiver binds a separate port on the pod IP, which is what the +# peer engine dials. +# +# PREREQUISITES / CAVEATS: +# - The prefill -> decode flow was validated with both engines co-located on +# one node, NIXL over TCP/CUDA-IPC without rdma/ib, which is what this +# manifest asks for. On a real RDMA fabric, add `rdma/ib: "1"` to each +# engine's resource limits and requests. Configure memory-lock limits and +# any additional capabilities only as required by that deployment. +# - No published sidecar image yet — build from lib/sidecar/Dockerfile and set +# . + +apiVersion: v1 +kind: ConfigMap +metadata: + name: trtllm-prefill-engine-config +data: + config.yaml: | + # Without a transceiver the engines cannot complete the handoff. NIXL picks + # its own underlying transport (UCX where there is no RDMA fabric). + cache_transceiver_config: + backend: NIXL +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: trtllm-decode-engine-config +data: + config.yaml: | + cache_transceiver_config: + backend: NIXL + # One gRPC chunk per N decode steps, not per token: N>1 cuts per-token gRPC + # serialization for a throughput win, at the cost of bursty delivery. Only + # the decode engine streams, so only it is tuned here. + stream_interval: 5 +--- +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: trtllm-sidecar-disagg +spec: + components: + - name: Frontend + type: frontend + replicas: 1 + podTemplate: + spec: + containers: + # Dynamo frontend; operator adds the command. + - name: main + # 1.4.0 or newer: the 1.3.0 router fails a disaggregated request + # whose one-token prefill stops on EOS or a stop match. + image: nvcr.io/nvidia/ai-dynamo/dynamo-frontend:1.5.0 + + - name: TRTLLMPrefillWorker + type: prefill + replicas: 1 + podTemplate: + spec: + volumes: + - name: engine-config + configMap: + name: trtllm-prefill-engine-config + # UCX shared-memory transports need more than the default 64 MiB. + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 10Gi + # trtllm-engine as a native sidecar (initContainer + restartPolicy: Always): + # starts before `main`, stops after. + initContainers: + - name: trtllm-engine + # Layer the OpenEngine Python bindings onto this release and push the + # result: the release ships the servicer, but both it and the probes + # below import `openengine.v1`, which it does not install. See the + # README's Run section for the pinned versions. + image: /tensorrt-llm-openengine:1.3.0rc27.dev202609170000 + restartPolicy: Always + # Readiness via the OpenEngine Control.Health RPC, which reports the + # engine's own state rather than just a connectable socket. + # startupProbe gates `main`; readinessProbe drops dead workers. + startupProbe: + exec: + command: + - python3 + - -c + - import grpc; from openengine.v1 import lifecycle_pb2, openengine_pb2_grpc; s=openengine_pb2_grpc.ControlStub(grpc.insecure_channel('127.0.0.1:50051')); r=s.Health(lifecycle_pb2.HealthRequest(), timeout=3); exit(0 if r.state==lifecycle_pb2.HEALTH_STATE_READY else 1) + periodSeconds: 5 + timeoutSeconds: 10 # must exceed the RPC timeout + interpreter spawn + failureThreshold: 60 # ~5 min for model load; raise for big models + readinessProbe: + exec: + command: + - python3 + - -c + - import grpc; from openengine.v1 import lifecycle_pb2, openengine_pb2_grpc; s=openengine_pb2_grpc.ControlStub(grpc.insecure_channel('127.0.0.1:50051')); r=s.Health(lifecycle_pb2.HealthRequest(), timeout=3); exit(0 if r.state==lifecycle_pb2.HEALTH_STATE_READY else 1) + periodSeconds: 10 + timeoutSeconds: 10 + failureThreshold: 3 + env: + # TensorRT-LLM would otherwise start its TP=1 worker by spawning an + # MPI process, and that spawn fails inside a pod: PRTE reaches + # `VM READY` and then terminates its daemons with no diagnostic. + # Measured on release:1.3.0rc27. Both engines run one GPU each, which + # is exactly the case this switch covers + # (`enable_worker_single_process_for_tp1`), so the worker runs + # in-process and never spawns. + # Not free: the engine logs "Using single process worker for TP1, + # this may hurt streaming generation performance" + # (`executor/executor.py`). It costs least here -- a context-only + # request emits one token -- and most on the decode leg below, where + # `stream_interval` is the lever that buys it back. Revisit both + # together if an engine needs TP > 1, which this switch does not cover. + - name: TLLM_WORKER_USE_SINGLE_PROCESS + value: "1" + # The release image runs as uid 0 and PRTE refuses to run as root + # unless told twice. Harmless with the switch above, and it keeps the + # engine from dying at startup if the spawn path is ever re-entered. + - name: PRTE_ALLOW_RUN_AS_ROOT + value: "1" + - name: PRTE_ALLOW_RUN_AS_ROOT_CONFIRM + value: "1" + envFrom: + - secretRef: + name: hf-token-secret + resources: + # The tested flow needs no RDMA fabric: NIXL falls back to + # TCP/CUDA-IPC. Add `rdma/ib: "1"` to both limits and requests only + # where a device plugin advertises it -- without one, both worker + # pods stay Pending. + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + # Increase for larger models. + ephemeral-storage: 2Gi + volumeMounts: + - name: engine-config + mountPath: /config + - name: dshm + mountPath: /dev/shm + command: + - python3 + - -m + - tensorrt_llm.commands.serve + args: + - Qwen/Qwen3-0.6B + - --grpc + - --grpc-protocol + - openengine + - --host + - 127.0.0.1 + - --port + - "50051" + # Workaround, not a requirement: `GetModelInfo` reads the unresolved + # `llm.args.max_seq_len`, which stays `None` on the PyTorch path even + # though the executor has already resolved the real window from the + # model config. Stating it here is what stops the sidecar polling for + # its whole startup deadline and then registering a worker that + # rejects every request omitting max_tokens. Drop this and the + # sidecar's --context-length once the servicer reports the resolved + # window. + - --max_seq_len + - "4096" + - --extra_llm_api_options + - /config/config.yaml + containers: + # Sidecar worker. No published image yet — build from + # lib/sidecar/Dockerfile, set (see README), add + # imagePullSecrets if private. Must be named "main" (operator injects + # discovery env + probes by name). + - name: main + image: /dynamo-sidecar:1.3.0 + command: + - dynamo-trtllm-sidecar + args: + - --disaggregation-mode + - prefill + - --grpc-endpoint + - 127.0.0.1:50051 + - --model-path + - Qwen/Qwen3-0.6B + # Matches the engine's --max_seq_len above. + - --context-length + - "4096" + envFrom: + - secretRef: + name: hf-token-secret + + - name: TRTLLMDecodeWorker + type: decode + replicas: 1 + podTemplate: + spec: + volumes: + - name: engine-config + configMap: + name: trtllm-decode-engine-config + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 10Gi + initContainers: + - name: trtllm-engine + image: /tensorrt-llm-openengine:1.3.0rc27.dev202609170000 + restartPolicy: Always + startupProbe: + exec: + command: + - python3 + - -c + - import grpc; from openengine.v1 import lifecycle_pb2, openengine_pb2_grpc; s=openengine_pb2_grpc.ControlStub(grpc.insecure_channel('127.0.0.1:50051')); r=s.Health(lifecycle_pb2.HealthRequest(), timeout=3); exit(0 if r.state==lifecycle_pb2.HEALTH_STATE_READY else 1) + periodSeconds: 5 + timeoutSeconds: 10 + failureThreshold: 60 + readinessProbe: + exec: + command: + - python3 + - -c + - import grpc; from openengine.v1 import lifecycle_pb2, openengine_pb2_grpc; s=openengine_pb2_grpc.ControlStub(grpc.insecure_channel('127.0.0.1:50051')); r=s.Health(lifecycle_pb2.HealthRequest(), timeout=3); exit(0 if r.state==lifecycle_pb2.HEALTH_STATE_READY else 1) + periodSeconds: 10 + timeoutSeconds: 10 + failureThreshold: 3 + env: + # See the prefill engine above: run the TP=1 worker in-process rather + # than spawning it over MPI, which does not work inside a pod. This is + # the leg that streams, so it is the leg that pays the engine's + # "may hurt streaming generation performance" warning -- the + # `stream_interval: 5` in this engine's ConfigMap is what offsets it. + - name: TLLM_WORKER_USE_SINGLE_PROCESS + value: "1" + - name: PRTE_ALLOW_RUN_AS_ROOT + value: "1" + - name: PRTE_ALLOW_RUN_AS_ROOT_CONFIRM + value: "1" + envFrom: + - secretRef: + name: hf-token-secret + resources: + # The tested flow needs no RDMA fabric: NIXL falls back to + # TCP/CUDA-IPC. Add `rdma/ib: "1"` to both limits and requests only + # where a device plugin advertises it -- without one, both worker + # pods stay Pending. + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + ephemeral-storage: 2Gi + volumeMounts: + - name: engine-config + mountPath: /config + - name: dshm + mountPath: /dev/shm + command: + - python3 + - -m + - tensorrt_llm.commands.serve + args: + - Qwen/Qwen3-0.6B + - --grpc + - --grpc-protocol + - openengine + - --host + - 127.0.0.1 + - --port + - "50051" + - --max_seq_len + - "4096" + - --extra_llm_api_options + - /config/config.yaml + containers: + - name: main + image: /dynamo-sidecar:1.3.0 + command: + - dynamo-trtllm-sidecar + args: + - --disaggregation-mode + - decode + - --grpc-endpoint + - 127.0.0.1:50051 + - --model-path + - Qwen/Qwen3-0.6B + - --context-length + - "4096" + envFrom: + - secretRef: + name: hf-token-secret diff --git a/lib/sidecar/trtllm/launch/agg.sh b/lib/sidecar/trtllm/launch/agg.sh index afc2940704a7..67c1861bc4bd 100755 --- a/lib/sidecar/trtllm/launch/agg.sh +++ b/lib/sidecar/trtllm/launch/agg.sh @@ -2,7 +2,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -# Aggregated serving through TensorRT-LLM's native gRPC server (1 GPU). +# Aggregated serving through TensorRT-LLM's OpenEngine gRPC server (1 GPU). +# +# Run this where `TRTLLM_PYTHON` has TensorRT-LLM installed -- +# `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc27.dev202609170000` or newer, the +# first releases carrying the OpenEngine servicer. The bindings it needs are not +# in that image; the pip step below adds them. set -e @@ -15,6 +20,8 @@ SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" source "$SCRIPT_DIR/../../../../examples/common/gpu_utils.sh" # build_trtllm_override_args_with_mem # shellcheck disable=SC1091 # Resolved relative to this script at runtime. source "$SCRIPT_DIR/../../../../examples/common/launch_utils.sh" # print_launch_banner, wait_any_exit +# shellcheck disable=SC1091 # Resolved relative to this script at runtime. +source "$SCRIPT_DIR/common.sh" # trtllm_ensure_openengine_bindings, trtllm_resolve_context_length MODEL="${MODEL:-Qwen/Qwen3-0.6B}" @@ -68,41 +75,9 @@ TRTLLM_PYTHON="${TRTLLM_PYTHON:-python3}" TRTLLM_GRPC_PORT="${TRTLLM_GRPC_PORT:-50051}" CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" -# Keep the engine and the sidecar on one number. Started without `--max_seq_len`, -# TensorRT-LLM reports its `max_input_len` default instead of a context length -# and the sidecar discards it, so pass the same value to both. When the caller -# supplies `--max_seq_len`, theirs wins and the sidecar adopts the engine's -# report rather than overriding it with a default it was never told about. -TRTLLM_MAX_SEQ_LEN_ARGS=() -TRTLLM_CONTEXT_LENGTH_ARGS=() -trtllm_max_seq_len_supplied=0 -for arg in "${EXTRA_ARGS[@]}"; do - case "$arg" in - --max_seq_len|--max_seq_len=*) trtllm_max_seq_len_supplied=1 ;; - esac -done -if [[ "$trtllm_max_seq_len_supplied" -eq 0 ]]; then - TRTLLM_CONTEXT_LENGTH="${TRTLLM_CONTEXT_LENGTH:-4096}" - TRTLLM_MAX_SEQ_LEN_ARGS=(--max_seq_len "$TRTLLM_CONTEXT_LENGTH") -fi -if [[ -n "$TRTLLM_CONTEXT_LENGTH" ]]; then - TRTLLM_CONTEXT_LENGTH_ARGS=(--context-length "$TRTLLM_CONTEXT_LENGTH") -fi +trtllm_resolve_context_length "${EXTRA_ARGS[@]}" -# `--grpc` needs `smg-grpc-proto`. Pinned to the exact version -# lib/sidecar/trtllm/proto/trtllm_service.proto was vendored from (see -# proto/README.md's checksum) -- 0.4.2 lacks the include_stop_token_in_output -# field (added by 0.4.14) our proto and Rust code both expect, which makes -# every request fail with "'GenerateRequest' object has no attribute -# 'include_stop_token_in_output'". Check the resolved version, not just -# importability: an image whose TRT-LLM install already pulled an older -# smg-grpc-proto (e.g. via its own grpc-smg extra) would otherwise satisfy a -# bare `import` check and skip straight past this pin. sys.exit, not assert: -# `assert` is stripped entirely under `python -O`/`PYTHONOPTIMIZE`, which -# would make this check fail open (exit 0) even with the package missing. -if ! "$TRTLLM_PYTHON" -c "import importlib.metadata as m, sys; sys.exit(0 if m.version('smg-grpc-proto') == '0.4.14' else 1)" >/dev/null 2>&1; then - "$TRTLLM_PYTHON" -m pip install --no-cache-dir "smg-grpc-proto==0.4.14" -fi +trtllm_ensure_openengine_bindings "$TRTLLM_PYTHON" HTTP_PORT="${DYN_HTTP_PORT:-8000}" GPU_MEM_ARGS=$(build_trtllm_override_args_with_mem) @@ -113,16 +88,17 @@ if [[ -n "$GPU_MEM_ARGS" ]]; then TRTLLM_GPU_MEM_ARGS=(--extra_llm_api_options "$TRTLLM_EXTRA_CONFIG") fi -print_launch_banner "Launching TensorRT-LLM Native-gRPC Sidecar (1 GPU)" "$MODEL" "$HTTP_PORT" \ +print_launch_banner "Launching TensorRT-LLM OpenEngine-gRPC Sidecar (1 GPU)" "$MODEL" "$HTTP_PORT" \ "TensorRT-LLM gRPC: 127.0.0.1:${TRTLLM_GRPC_PORT}" \ "Context length: ${TRTLLM_CONTEXT_LENGTH:-from engine report}" python3 -m dynamo.frontend & -# TensorRT-LLM's native gRPC listener is unauthenticated; keep it on loopback. +# TensorRT-LLM's OpenEngine gRPC listener is unauthenticated; keep it on loopback. CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" \ "$TRTLLM_PYTHON" -m tensorrt_llm.commands.serve "$MODEL" \ --grpc \ + --grpc-protocol openengine \ --host 127.0.0.1 \ --port "$TRTLLM_GRPC_PORT" \ "${TRTLLM_MAX_SEQ_LEN_ARGS[@]}" \ diff --git a/lib/sidecar/trtllm/launch/common.sh b/lib/sidecar/trtllm/launch/common.sh new file mode 100755 index 000000000000..21e4a901dacf --- /dev/null +++ b/lib/sidecar/trtllm/launch/common.sh @@ -0,0 +1,80 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Shared by agg.sh and disagg.sh. Both need the same OpenEngine bindings and the +# same context-length derivation; keeping one copy is what stops the two +# topologies drifting onto different contract revisions. +# +# Source it, then call: +# trtllm_ensure_openengine_bindings "$TRTLLM_PYTHON" +# trtllm_resolve_context_length "${EXTRA_ARGS[@]}" +# +# `trtllm_resolve_context_length` sets TRTLLM_MAX_SEQ_LEN_ARGS, +# TRTLLM_CONTEXT_LENGTH_ARGS, and TRTLLM_CONTEXT_LENGTH. + +# `--grpc-protocol openengine` needs the OpenEngine bindings, which resolve only +# from a custom index. Both packages are pinned to BSR module commit +# 768a93c7b44e, the same revision the vendored protos in `proto/` were generated +# from (see `proto/README.md`), so the engine and the sidecar speak the same +# contract revision. +# +# The protobuf package is pinned by *gencode* version as well: buf publishes one +# build per protoc release, and a gencode newer than the runtime in the +# TensorRT-LLM image fails at import with "Detected incompatible Protobuf +# Gencode/Runtime versions". 33.5 matches the protobuf 6.33.x runtime those +# images ship. Raise it only together with the image's protobuf. +TRTLLM_OPENENGINE_PROTOBUF_VERSION="33.5.0.1.20260730172104+768a93c7b44e" +TRTLLM_OPENENGINE_GRPC_VERSION="1.78.1.1.20260730172104+768a93c7b44e" +# Pinned because the protobuf wheel requires it with no upper bound, and the +# name is unregistered on PyPI: under --extra-index-url pip would take the +# highest version from either index, so anyone claiming the name there could +# run code in the engine's environment. A `+local` version cannot come from +# PyPI, so pinning it exactly closes that. TensorRT-LLM's +# requirements-openengine.txt pins it the same way. +TRTLLM_OPENENGINE_PROTOBUF_PYI_VERSION="36.2.0.1.20260730172104+768a93c7b44e" + +# Installs the bindings unless they are already importable. Set +# TRTLLM_SKIP_BINDINGS_INSTALL=1 on an air-gapped host or one whose +# site-packages is read-only, and install them yourself. +trtllm_ensure_openengine_bindings() { + local python="$1" + if "$python" -c "import openengine.v1.openengine_pb2" >/dev/null 2>&1; then + return 0 + fi + if [[ -n "${TRTLLM_SKIP_BINDINGS_INSTALL:-}" ]]; then + echo "OpenEngine bindings are missing and TRTLLM_SKIP_BINDINGS_INSTALL is set." >&2 + echo "Install them into $python before starting the engine." >&2 + return 1 + fi + echo "Installing the pinned OpenEngine bindings into $python..." + "$python" -m pip install --no-cache-dir \ + --extra-index-url https://buf.build/gen/python \ + "openengine-openengine-grpc-python==${TRTLLM_OPENENGINE_GRPC_VERSION}" \ + "openengine-openengine-protocolbuffers-python==${TRTLLM_OPENENGINE_PROTOBUF_VERSION}" \ + "openengine-openengine-protocolbuffers-pyi==${TRTLLM_OPENENGINE_PROTOBUF_PYI_VERSION}" +} + +# Keeps the engine and the sidecar on one number. Started without +# `--max_seq_len`, TensorRT-LLM leaves `max_context_length` unset and the sidecar +# has no window to register, so pass the same value to both. When the caller +# supplies `--max_seq_len`, theirs wins and the sidecar adopts the engine's +# `Control.GetModelInfo` report rather than overriding it with a default it was +# never told about. +trtllm_resolve_context_length() { + TRTLLM_MAX_SEQ_LEN_ARGS=() + TRTLLM_CONTEXT_LENGTH_ARGS=() + local arg supplied=0 + for arg in "$@"; do + case "$arg" in + --max_seq_len|--max_seq_len=*) supplied=1 ;; + esac + done + if [[ "$supplied" -eq 0 ]]; then + TRTLLM_CONTEXT_LENGTH="${TRTLLM_CONTEXT_LENGTH:-4096}" + TRTLLM_MAX_SEQ_LEN_ARGS=(--max_seq_len "$TRTLLM_CONTEXT_LENGTH") + fi + if [[ -n "${TRTLLM_CONTEXT_LENGTH:-}" ]]; then + TRTLLM_CONTEXT_LENGTH_ARGS=(--context-length "$TRTLLM_CONTEXT_LENGTH") + fi +} diff --git a/lib/sidecar/trtllm/launch/disagg.sh b/lib/sidecar/trtllm/launch/disagg.sh new file mode 100755 index 000000000000..4fcc1aaea56e --- /dev/null +++ b/lib/sidecar/trtllm/launch/disagg.sh @@ -0,0 +1,166 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Disaggregated serving through two TensorRT-LLM OpenEngine gRPC servers (2 workers). +# +# Run this where `TRTLLM_PYTHON` has TensorRT-LLM installed -- +# `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc27.dev202609170000` or newer, the +# first releases carrying the OpenEngine servicer. The bindings it needs are not +# in that image; the pip step below adds them. +# +# KV cache moves prefill -> decode engine-to-engine over the cache transceiver, +# not through the sidecars. Dynamo carries only the handoff: the prefill worker +# returns a `KvSessionRef` and the decode worker replays it. + +set -e + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +# Resolved relative to this script, not via $DYNAMO_HOME, for the reason agg.sh +# states: some runtime images bake DYNAMO_HOME to a minimal install path with no +# examples/ directory, and `${DYNAMO_HOME:-...}` lets that baked value win and +# silently break sourcing. +# shellcheck disable=SC1091 # Resolved relative to this script at runtime. +source "$SCRIPT_DIR/../../../../examples/common/gpu_utils.sh" # build_trtllm_override_args_with_mem +# shellcheck disable=SC1091 # Resolved relative to this script at runtime. +source "$SCRIPT_DIR/../../../../examples/common/launch_utils.sh" # print_launch_banner, wait_any_exit +# shellcheck disable=SC1091 # Resolved relative to this script at runtime. +source "$SCRIPT_DIR/common.sh" # trtllm_ensure_openengine_bindings, trtllm_resolve_context_length + +MODEL="${MODEL:-Qwen/Qwen3-0.6B}" + +EXTRA_ARGS=() +while [[ $# -gt 0 ]]; do + case $1 in + --model|--model-path) + if [[ $# -lt 2 || "$2" == -* ]]; then + echo "Missing value for $1" + echo "Use --help for usage information" + exit 1 + fi + MODEL="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [--model|--model-path ] [TensorRT-LLM engine options...]" + echo + echo "Additional options are passed to both TensorRT-LLM engines." + echo + echo "Environment overrides:" + echo " MODEL Model to serve (default: Qwen/Qwen3-0.6B)" + echo " TRTLLM_PYTHON Python with TensorRT-LLM installed (default: python3)" + echo " TRTLLM_PREFILL_GPU Prefill GPU assignment (default: 0)" + echo " TRTLLM_DECODE_GPU Decode GPU assignment (default: 1)" + echo " DYN_HTTP_PORT Dynamo frontend port (default: 8000)" + echo " DYN_SYSTEM_PORT1 Prefill sidecar system port (default: 8081)" + echo " DYN_SYSTEM_PORT2 Decode sidecar system port (default: 8082)" + echo " TRTLLM_PREFILL_GRPC_PORT Prefill TensorRT-LLM gRPC port (default: 50051)" + echo " TRTLLM_DECODE_GRPC_PORT Decode TensorRT-LLM gRPC port (default: 50052)" + echo " TRTLLM_CACHE_TRANSCEIVER_BACKEND KV transfer backend (default: NIXL)" + echo " TRTLLM_CONTEXT_LENGTH Model context length, applied to both engines and both" + echo " sidecars (default: 4096; unset when --max_seq_len is given)" + exit 0 + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +TRTLLM_EXTRA_CONFIG="" +trtllm_exit_trap() { + local rc=$? + if [[ -n "$TRTLLM_EXTRA_CONFIG" ]]; then + rm -f -- "$TRTLLM_EXTRA_CONFIG" + fi + echo "Cleaning up..." + dynamo_reap_and_exit "$rc" +} +trap trtllm_exit_trap EXIT + +TRTLLM_PYTHON="${TRTLLM_PYTHON:-python3}" +TRTLLM_HOST="127.0.0.1" +TRTLLM_PREFILL_GRPC_PORT="${TRTLLM_PREFILL_GRPC_PORT:-50051}" +TRTLLM_DECODE_GRPC_PORT="${TRTLLM_DECODE_GRPC_PORT:-50052}" +TRTLLM_PREFILL_GPU="${TRTLLM_PREFILL_GPU:-0}" +TRTLLM_DECODE_GPU="${TRTLLM_DECODE_GPU:-1}" +TRTLLM_CACHE_TRANSCEIVER_BACKEND="${TRTLLM_CACHE_TRANSCEIVER_BACKEND:-NIXL}" + +# `--extra_llm_api_options` is last-wins, not additive, so a forwarded copy +# would drop the transceiver and leave the prefill worker producing a handoff +# no decode worker can consume -- as an opaque engine-side transfer error, not +# a launcher one. Refuse it rather than silently losing the setting. +for arg in "${EXTRA_ARGS[@]}"; do + case "$arg" in + --extra_llm_api_options|--extra_llm_api_options=*) + echo "Cannot forward --extra_llm_api_options: this launcher needs it for" >&2 + echo "cache_transceiver_config. Merge your settings into that file, or set" >&2 + echo "TRTLLM_CACHE_TRANSCEIVER_BACKEND and run the engines yourself." >&2 + exit 1 + ;; + esac +done + +trtllm_resolve_context_length "${EXTRA_ARGS[@]}" + +trtllm_ensure_openengine_bindings "$TRTLLM_PYTHON" + +# Both engines need a cache transceiver or the handoff has nothing to move the +# KV cache over. NIXL picks its own underlying transport (UCX where there is no +# RDMA fabric) and is the path Dynamo uses elsewhere for disaggregation. +TRTLLM_EXTRA_CONFIG=$(mktemp "${TMPDIR:-/tmp}/dynamo-trtllm-sidecar.XXXXXX.yaml") +build_trtllm_override_args_with_mem \ + --merge-with-json "{\"cache_transceiver_config\": {\"backend\": \"${TRTLLM_CACHE_TRANSCEIVER_BACKEND}\"}}" \ + > "$TRTLLM_EXTRA_CONFIG" + +HTTP_PORT="${DYN_HTTP_PORT:-8000}" + +print_launch_banner "Launching TensorRT-LLM OpenEngine-gRPC Sidecar (Disaggregated, 2 workers)" "$MODEL" "$HTTP_PORT" \ + "Prefill: GPU ${TRTLLM_PREFILL_GPU}, gRPC ${TRTLLM_HOST}:${TRTLLM_PREFILL_GRPC_PORT}" \ + "Decode: GPU ${TRTLLM_DECODE_GPU}, gRPC ${TRTLLM_HOST}:${TRTLLM_DECODE_GRPC_PORT}" \ + "KV transfer: ${TRTLLM_CACHE_TRANSCEIVER_BACKEND} cache transceiver" \ + "Context length: ${TRTLLM_CONTEXT_LENGTH:-from engine report}" + +python3 -m dynamo.frontend & + +# TensorRT-LLM's OpenEngine gRPC listener is unauthenticated; keep it on +# loopback. The cache transceiver binds its own port on a routable address, so +# the engines still reach each other. +CUDA_VISIBLE_DEVICES="$TRTLLM_PREFILL_GPU" \ +"$TRTLLM_PYTHON" -m tensorrt_llm.commands.serve "$MODEL" \ + --grpc \ + --grpc-protocol openengine \ + --host "$TRTLLM_HOST" \ + --port "$TRTLLM_PREFILL_GRPC_PORT" \ + "${TRTLLM_MAX_SEQ_LEN_ARGS[@]}" \ + --extra_llm_api_options "$TRTLLM_EXTRA_CONFIG" \ + "${EXTRA_ARGS[@]}" & + +CUDA_VISIBLE_DEVICES="$TRTLLM_DECODE_GPU" \ +"$TRTLLM_PYTHON" -m tensorrt_llm.commands.serve "$MODEL" \ + --grpc \ + --grpc-protocol openengine \ + --host "$TRTLLM_HOST" \ + --port "$TRTLLM_DECODE_GRPC_PORT" \ + "${TRTLLM_MAX_SEQ_LEN_ARGS[@]}" \ + --extra_llm_api_options "$TRTLLM_EXTRA_CONFIG" \ + "${EXTRA_ARGS[@]}" & + +# `--disaggregation-mode prefill` also registers this worker under the `prefill` +# component, which is what the frontend's prefill router targets. +DYN_SYSTEM_PORT="${DYN_SYSTEM_PORT1:-8081}" \ + dynamo-trtllm-sidecar \ + --disaggregation-mode prefill \ + --grpc-endpoint "${TRTLLM_HOST}:${TRTLLM_PREFILL_GRPC_PORT}" \ + --model-path "$MODEL" \ + "${TRTLLM_CONTEXT_LENGTH_ARGS[@]}" & + +DYN_SYSTEM_PORT="${DYN_SYSTEM_PORT2:-8082}" \ + dynamo-trtllm-sidecar \ + --disaggregation-mode decode \ + --grpc-endpoint "${TRTLLM_HOST}:${TRTLLM_DECODE_GRPC_PORT}" \ + --model-path "$MODEL" \ + "${TRTLLM_CONTEXT_LENGTH_ARGS[@]}" & + +wait_any_exit diff --git a/lib/sidecar/trtllm/proto/README.md b/lib/sidecar/trtllm/proto/README.md index 2019a75c91ac..d77fd88aa732 100644 --- a/lib/sidecar/trtllm/proto/README.md +++ b/lib/sidecar/trtllm/proto/README.md @@ -3,21 +3,27 @@ SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. SPDX-License-Identifier: Apache-2.0 --> -# Vendored TensorRT-LLM gRPC protocol +# Vendored OpenEngine protocol -- Source package: `smg-grpc-proto` -- Version: `0.4.14` -- File: `smg_grpc_proto/proto/trtllm_service.proto` -- Upstream SHA-256: `ec2f93de3af59047e3779283369eba89c2e6981b27849553144671b1fbc8eea5` +- Source: [`proto/openengine/v1`](https://github.com/ai-dynamo/openengine/tree/b5f2bd93721f7b888d3e2440679e0ae7012939d1/proto/openengine/v1) +- Commit: `b5f2bd93721f7b888d3e2440679e0ae7012939d1` (tag `v0.1.0`) -`trtllm_service.proto` here is derived from the upstream file above with a -single, mechanical edit: the `import "common.proto"` line and the two RPCs that -reference `smg.grpc.common` types — `GetTokenizer` and `SubscribeKvEvents` — are -removed. The sidecar drives only `Generate`, so dropping the tokenizer/KV-event -RPCs lets the contract compile standalone without vendoring `common.proto`. -Every message the sidecar exchanges is preserved byte-for-byte. +This is the same revision TensorRT-LLM's OpenEngine server is generated from: it +is the git commit behind Buf Schema Registry module commit `768a93c7b44e` of +`buf.build/openengine/openengine`. -TensorRT-LLM loads the matching generated stubs from `smg-grpc-proto` at -runtime (`tensorrt_llm/grpc/__init__.py`), first shipped in TRT-LLM -`1.3.0rc21`. Update the version, checksum, and this note together when bumping -the protocol. +SHA-256: + +| File | SHA-256 | +|------|---------| +| `openengine/v1/error.proto` | `b1c559a3b7847c55a247145f1493805b56b58491a320fb26e6ebcd7072b62fb3` | +| `openengine/v1/generation.proto` | `3c4e9f11664d50dd26e393cb2b34dca9af680419097493d4f7ceb2654b0f1058` | +| `openengine/v1/kv.proto` | `4f5c46a65a7998ecf34ea37117b325b0f0a5a70149702d6bca919726b40526cd` | +| `openengine/v1/lifecycle.proto` | `0b39bbecc617db3de900333efe6c070dadca2f71f5e372c306d44973229d48c9` | +| `openengine/v1/lora.proto` | `c1aa58fca3ab1b40f7c75dbbd17280b2b52ee295256479b6b0fc74f924047ef9` | +| `openengine/v1/model.proto` | `5166cdcdaa260b1f5de2fa8d170dacbab1595af9413e73f09c6902b25f553284` | +| `openengine/v1/openengine.proto` | `1d1baafee0aaf0dbf51962ae2d293431fbd3be96601b375288fab785fc2d8572` | +| `openengine/v1/server.proto` | `deb379508d5e93c719cc0c5ebe222e66af7d8784384d3cc270e95d0b48274d96` | + +The files are copied without modification. Update the revision and checksums +together, and add or remove entries in `build.rs`'s `PROTOS` list to match. diff --git a/lib/sidecar/trtllm/proto/openengine/v1/error.proto b/lib/sidecar/trtllm/proto/openengine/v1/error.proto new file mode 100644 index 000000000000..de2e13d23b81 --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/error.proto @@ -0,0 +1,31 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Terminal errors for accepted streaming requests. + +syntax = "proto3"; + +package openengine.v1; + +// Accepted failures emit one terminal EngineError and close with OK. +// Validation and transport failures use non-OK gRPC status instead. +message EngineError { + ErrorCode code = 1; + string message = 2; + bool retryable = 3; // Retry may succeed without changing the request. +} + +enum ErrorCode { + ERROR_CODE_UNSPECIFIED = 0; + ERROR_CODE_INVALID_ARGUMENT = 1; + ERROR_CODE_UNSUPPORTED_FEATURE = 2; + ERROR_CODE_ROLE_MISMATCH = 3; + ERROR_CODE_MODEL_NOT_FOUND = 4; + ERROR_CODE_OVERLOADED = 5; + ERROR_CODE_REQUEST_NOT_FOUND = 6; + ERROR_CODE_DUPLICATE_REQUEST = 7; + ERROR_CODE_KV_SESSION_NOT_FOUND = 8; + ERROR_CODE_KV_TRANSFER_FAILED = 9; + ERROR_CODE_CANCELLED = 10; + ERROR_CODE_INTERNAL = 11; +} diff --git a/lib/sidecar/trtllm/proto/openengine/v1/generation.proto b/lib/sidecar/trtllm/proto/openengine/v1/generation.proto new file mode 100644 index 000000000000..c1f81ddac8b1 --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/generation.proto @@ -0,0 +1,223 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Generation inputs, requests, streamed events, and usage. + +syntax = "proto3"; + +package openengine.v1; + +import "google/protobuf/struct.proto"; +import "openengine/v1/error.proto"; +import "openengine/v1/kv.proto"; + +message TokenIds { + repeated uint32 ids = 1; +} + +// Multimodal modality discriminator. UNSPECIFIED means the sender left it unset. +enum Modality { + MODALITY_UNSPECIFIED = 0; + MODALITY_IMAGE = 1; + MODALITY_VIDEO = 2; + MODALITY_AUDIO = 3; +} + +// A single multimodal input. Exactly one `source` should be set. The engine +// owns fetch, decode, and preprocessing, so pre-decoded or RDMA media +// descriptors are not represented here. +message MediaItem { + Modality modality = 1; + oneof source { + string url = 2; // http(s):// -- engine fetches + string data_uri = 3; // data:;base64,<...> -- engine decodes + bytes raw_bytes = 4; // pre-fetched bytes -- engine still preprocesses + } + string mime_type = 5; // optional, hints raw_bytes decode + string uuid = 6; // optional caller id / mm_hash +} + +message GenerateRequest { + string request_id = 1; + string model = 2; + + oneof input { + string prompt = 3; + TokenIds token_ids = 4; + } + + SamplingParams sampling = 5; + StoppingOptions stopping = 6; + ResponseOptions response = 7; + KvOptions kv = 8; + GuidedDecoding guided = 9; + + // Multimodal inputs. Order is significant: the i-th item aligns with the + // i-th (un-expanded) placeholder marker carried in the prompt/token_ids. + // The engine fetches/decodes and preprocesses each item, then expands the + // marker into the model's replacement run. Empty for text-only requests. + repeated MediaItem media = 10; + + // Loaded LoRA adapter name to apply to this request. Empty = base model. + // ModelInfo.supports_lora advertises whether lifecycle and selection are + // available through OpenEngine. + string lora_name = 11; + + // Engine-specific request parameters that have no portable field above (e.g. + // an experimental sampler knob). NOT part of the portable contract: an engine + // MAY ignore keys it does not recognize, and a request MUST remain valid with + // this field empty. Clients treat it as best-effort and never depend on it + // for correctness. Carried as a Struct so JSON types survive the wire. + google.protobuf.Struct extra = 12; +} + +message SamplingParams { + optional double temperature = 1; + optional double top_p = 2; + optional int32 top_k = 3; + optional double min_p = 4; + optional double frequency_penalty = 5; + optional double presence_penalty = 6; + optional double repetition_penalty = 7; + optional uint64 seed = 8; + optional uint32 num_sequences = 9; +} + +message StoppingOptions { + optional uint32 max_tokens = 1; + optional uint32 min_tokens = 2; + repeated StopCondition conditions = 3; + optional bool ignore_eos = 4; + optional bool include_stop_in_output = 5; +} + +message ResponseOptions { + optional bool return_prompt_logprobs = 1; + CandidateTokenSelection prompt_candidates = 2; + optional bool return_output_logprobs = 3; + CandidateTokenSelection output_candidates = 4; + optional uint32 prompt_logprob_start = 5; +} + +message CandidateTokenSelection { + oneof selection { + uint32 top_n = 1; + TokenIds token_ids = 2; + AllCandidates all = 3; + } +} + +message AllCandidates {} + +message KvOptions { + KvSessionRef session = 1; + optional bool bypass_prefix_cache = 2; + optional string cache_salt = 3; +} + +message StopCondition { + oneof condition { + string stop_text = 1; + uint32 stop_token_id = 2; + } +} + +// Constrained / guided decoding spec. At most one of `guide` should be set. +// The engine enforces the constraint during sampling via its grammar backend +// (xgrammar / outlines / llguidance); clients cannot apply it post-hoc. +message GuidedDecoding { + oneof guide { + string json_schema = 1; // output conforms to this JSON schema + string regex = 2; // output matches this regex + string ebnf_grammar = 3; // output follows this EBNF / context-free grammar + string structural_tag = 4; // xgrammar structural-tag constraint (JSON string) + ChoiceConstraint choice = 5; + JsonObjectConstraint json_object = 6; + } + // Optional grammar backend override (e.g. "xgrammar", "outlines", + // "llguidance"). Empty = engine default. + string backend = 7; +} + +message ChoiceConstraint { + repeated string choices = 1; +} + +message JsonObjectConstraint {} + +message GenerateResponse { + string request_id = 1; + + // PromptOutput, PrefillReady, and EngineError are request-scoped. + // TokenOutput and GenerationFinished are output-scoped. + oneof event { + PromptOutput prompt = 2; + TokenOutput token = 3; + PrefillReady prefill_ready = 4; + GenerationFinished finished = 5; + EngineError error = 6; + } + + Usage usage = 10; // Cumulative request usage; only set on the final response. +} + +// Request-scoped prompt token information, emitted at most once. +message PromptOutput { + repeated TokenInfo tokens = 1; +} + +// An incremental output delta. Tokens and text are never cumulative. +message TokenOutput { + optional uint32 output_index = 1; // Required, including for output zero. + repeated TokenInfo tokens = 2; + string text = 3; +} + +message TokenInfo { + uint32 token_id = 1; + string token = 2; + optional double logprob = 3; + optional uint32 rank = 4; + repeated LogProb candidates = 5; +} + +message LogProb { + uint32 token_id = 1; + double logprob = 2; + string token = 3; + optional uint32 rank = 4; +} + +message PrefillReady { + KvSessionRef kv_session = 1; +} + +message GenerationFinished { + optional uint32 output_index = 1; // Required, including for output zero. + FinishReason reason = 2; + string message = 3; + StopMatch stop_match = 4; +} + +message StopMatch { + oneof match { + uint32 stop_token_id = 1; + string stop_text = 2; + uint32 eos_token_id = 3; + } +} + +enum FinishReason { + FINISH_REASON_UNSPECIFIED = 0; + FINISH_REASON_STOP = 1; + FINISH_REASON_LENGTH = 2; + FINISH_REASON_CANCELLED = 3; +} + +message Usage { + uint32 prompt_tokens = 1; + uint32 completion_tokens = 2; + uint32 total_tokens = 3; + optional uint32 cached_prompt_tokens = 4; + optional uint32 reasoning_tokens = 5; +} diff --git a/lib/sidecar/trtllm/proto/openengine/v1/kv.proto b/lib/sidecar/trtllm/proto/openengine/v1/kv.proto new file mode 100644 index 000000000000..d73a63dbd43b --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/kv.proto @@ -0,0 +1,134 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// KV sessions, connector discovery, and cache events. + +syntax = "proto3"; + +package openengine.v1; + +import "google/protobuf/struct.proto"; +import "openengine/v1/error.proto"; + +message KvSessionRef { + string session_id = 1; + string transfer_backend = 2; + repeated KvEndpoint endpoints = 3; + uint32 dp_rank = 4; + // Engine-specific KV-transfer and rendezvous parameters, carried as a Struct + // so numbers, booleans, and arrays survive the wire with their JSON type + // intact. Struct numbers are IEEE-754 doubles (exact only to 2^53); carry + // larger integer values as decimal strings. + google.protobuf.Struct attributes_struct = 5; +} + +message KvEndpoint { + string host = 1; + uint32 port = 2; + string protocol = 3; // grpc, nixl, ucx, tcp, shm, etc. +} + +message KvConnectorInfo { + optional bool enabled = 1; + string transfer_backend = 2; + repeated KvEndpoint local_endpoints = 3; + repeated string supported_protocols = 4; + optional bool supports_remote_prefill = 5; + optional bool supports_decode_pull = 6; + optional bool supports_abort_cleanup = 7; + optional uint32 schema_version = 8; +} + +message GetKvEventSourcesRequest { + repeated uint32 data_parallel_ranks = 1; +} + +message GetKvEventSourcesResponse { + repeated KvEventSource sources = 1; +} + +message KvEventSource { + string transport = 1; // grpc, zmq + // Connectable address. MUST use a routable host, never a bind wildcard. + KvEndpoint endpoint_addr = 2; + string topic = 3; + string replay_endpoint = 4; // optional, for ZMQ replay + optional uint32 data_parallel_rank = 5; + string encoding = 6; // protobuf, msgpack + optional uint32 schema_version = 7; + optional uint32 buffer_steps = 8; + optional uint32 hwm = 9; + optional uint32 max_queue_size = 10; +} + +message SubscribeKvEventsRequest { + repeated uint32 data_parallel_ranks = 1; + bool include_snapshot = 2; + uint64 start_sequence_number = 3; +} + +message SubscribeKvEventsResponse { + oneof event { + KvEventBatch batch = 1; + EngineError error = 2; // Terminal. + } +} + +message KvEventBatch { + uint64 sequence_number = 1; + uint64 timestamp_unix_nanos = 2; + uint32 data_parallel_rank = 3; + repeated KvEvent events = 4; +} + +message KvEvent { + string request_id = 1; + KvSessionRef kv_session = 2; + + oneof event { + BlockStored block_stored = 10; + BlockRemoved block_removed = 11; + AllBlocksCleared all_blocks_cleared = 12; + } +} + +message BlockStored { + repeated KvBlockHash block_hashes = 1; + KvBlockHash parent_block_hash = 2; + repeated uint32 token_ids = 3; + uint32 block_size = 4; + int64 lora_id = 5; + string lora_name = 6; + StorageMedium medium = 7; + + // vLLM-compatible optional metadata for reconstructing block keys. + repeated OpaqueKeyTuple extra_keys = 20; + uint32 group_idx = 21; + string kv_cache_spec_kind = 22; + uint32 kv_cache_spec_sliding_window = 23; +} + +message BlockRemoved { + repeated KvBlockHash block_hashes = 1; + StorageMedium medium = 2; + uint32 group_idx = 3; +} + +message AllBlocksCleared {} + +message KvBlockHash { + bytes value = 1; + string encoding = 2; // int64, string, bytes, engine_specific +} + +message OpaqueKeyTuple { + repeated string values = 1; +} + +enum StorageMedium { + STORAGE_MEDIUM_UNSPECIFIED = 0; + STORAGE_MEDIUM_GPU = 1; + STORAGE_MEDIUM_CPU_PINNED = 2; + STORAGE_MEDIUM_DISK = 3; + STORAGE_MEDIUM_EXTERNAL = 4; +} diff --git a/lib/sidecar/trtllm/proto/openengine/v1/lifecycle.proto b/lib/sidecar/trtllm/proto/openengine/v1/lifecycle.proto new file mode 100644 index 000000000000..c5611b022866 --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/lifecycle.proto @@ -0,0 +1,64 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Health and abort operations. + +syntax = "proto3"; + +package openengine.v1; + +import "openengine/v1/kv.proto"; +import "openengine/v1/server.proto"; + +message HealthRequest { + // False means a lightweight readiness/liveness check. True asks the engine to + // run a role-appropriate minimal inference probe and report it as a check. + bool include_inference_probe = 1; + + // Optional. Used when include_inference_probe is true. Empty means engine + // default served model. + string model = 2; + + // Optional expected role for role-specific inference probes. + EngineRole role = 3; +} + +message HealthResponse { + HealthState state = 1; + repeated HealthCheck checks = 2; +} + +enum HealthState { + HEALTH_STATE_UNSPECIFIED = 0; + HEALTH_STATE_STARTING = 1; + HEALTH_STATE_READY = 2; + HEALTH_STATE_DEGRADED = 3; + HEALTH_STATE_NOT_READY = 4; +} + +message HealthCheck { + string name = 1; // grpc, scheduler, model, kv_connector, role, inference_probe + HealthState state = 2; + string message = 3; +} + +message AbortRequest { + oneof target { + string request_id = 1; + KvSessionRef kv_session = 2; + AllRequests all_requests = 3; + } +} + +message AllRequests {} + +message AbortResponse { + AbortStatus status = 1; + string message = 2; +} + +enum AbortStatus { + ABORT_STATUS_UNSPECIFIED = 0; + ABORT_STATUS_ABORTED = 1; + ABORT_STATUS_ALREADY_FINISHED = 2; +} diff --git a/lib/sidecar/trtllm/proto/openengine/v1/lora.proto b/lib/sidecar/trtllm/proto/openengine/v1/lora.proto new file mode 100644 index 000000000000..e0951ad2bf47 --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/lora.proto @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// LoRA adapter lifecycle messages. + +syntax = "proto3"; + +package openengine.v1; + +message LoraAdapter { + int64 lora_id = 1; + string lora_name = 2; + string source_path = 3; +} + +message LoadLoraRequest { + LoraAdapter adapter = 1; +} + +message LoadLoraResponse { + LoraAdapter adapter = 1; + bool already_loaded = 2; +} + +message UnloadLoraRequest { + string lora_name = 1; +} + +message UnloadLoraResponse { + LoraAdapter adapter = 1; +} + +message ListLorasRequest {} + +message ListLorasResponse { + repeated LoraAdapter adapters = 1; +} diff --git a/lib/sidecar/trtllm/proto/openengine/v1/model.proto b/lib/sidecar/trtllm/proto/openengine/v1/model.proto new file mode 100644 index 000000000000..f8b8c9c8a2d5 --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/model.proto @@ -0,0 +1,80 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Model metadata and portable inference capabilities. + +syntax = "proto3"; + +package openengine.v1; + +import "google/protobuf/struct.proto"; + +message GetModelInfoRequest { + string model = 1; +} + +message ModelInfo { + string model_id = 1; + string served_model_name = 2; + repeated string served_model_aliases = 3; + optional uint32 max_context_length = 4; // Effective context-window limit in this deployment. + optional uint32 max_output_tokens = 5; // Effective generated-token limit in this deployment. + repeated string tokenizer_modes = 10; + + optional bool supports_text_input = 20; + optional bool supports_token_ids_input = 21; + GenerationCapabilities generation = 22; + optional bool supports_lora = 23; + optional bool supports_multimodal = 24; + + // Engine-advertised response parser names. Clients can apply these to a + // model's output stream for tool-call extraction or reasoning separation. + // Empty means no parser is configured for this model. + string reasoning_parser = 25; + string tool_call_parser = 26; + + // Engine-specific model metadata with no portable field (e.g. kv-cache dtype, + // quantization, capabilities not yet standardized). NOT part of the portable + // contract: clients read it opportunistically and never depend on it for + // correctness. Carried as a Struct so JSON types survive the wire. + google.protobuf.Struct extra = 28; +} + +message GenerationCapabilities { + LogprobCapabilities prompt_logprobs = 1; + LogprobCapabilities output_logprobs = 2; + GuidedDecodingCapabilities guided_decoding = 3; + optional uint32 max_num_sequences = 4; + optional bool supports_priority = 5; + optional bool supports_stop_in_output = 6; + optional bool supports_cache_salt = 7; + optional bool supports_prefix_cache_bypass = 8; +} + +message LogprobCapabilities { + optional bool supported = 1; + repeated CandidateTokenSelectionMode candidate_selection_modes = 2; + optional uint32 max_top_n = 3; +} + +enum CandidateTokenSelectionMode { + CANDIDATE_TOKEN_SELECTION_MODE_UNSPECIFIED = 0; + CANDIDATE_TOKEN_SELECTION_MODE_TOP_N = 1; + CANDIDATE_TOKEN_SELECTION_MODE_TOKEN_IDS = 2; + CANDIDATE_TOKEN_SELECTION_MODE_ALL = 3; +} + +message GuidedDecodingCapabilities { + optional bool supported = 1; + repeated GuidedDecodingMode modes = 2; +} + +enum GuidedDecodingMode { + GUIDED_DECODING_MODE_UNSPECIFIED = 0; + GUIDED_DECODING_MODE_JSON_SCHEMA = 1; + GUIDED_DECODING_MODE_REGEX = 2; + GUIDED_DECODING_MODE_EBNF_GRAMMAR = 3; + GUIDED_DECODING_MODE_STRUCTURAL_TAG = 4; + GUIDED_DECODING_MODE_CHOICE = 5; + GUIDED_DECODING_MODE_JSON_OBJECT = 6; +} diff --git a/lib/sidecar/trtllm/proto/openengine/v1/openengine.proto b/lib/sidecar/trtllm/proto/openengine/v1/openengine.proto new file mode 100644 index 000000000000..bd55318141f9 --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/openengine.proto @@ -0,0 +1,61 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// OpenEngine API v1 -- vendor-neutral runtime API between model engines (e.g. +// vLLM, SGLang, TensorRT-LLM) and clients, including direct applications and +// distributed frameworks (e.g. Dynamo). +// +// The files in this package are the canonical contract and single source of +// truth for the wire shape. Implementations should generate client and server +// bindings from the complete openengine/v1 package at a known revision. +// +// Human-readable reference: +// https://github.com/ai-dynamo/openengine/blob/main/docs/api.md + +syntax = "proto3"; + +package openengine.v1; + +import "openengine/v1/generation.proto"; +import "openengine/v1/kv.proto"; +import "openengine/v1/lifecycle.proto"; +import "openengine/v1/lora.proto"; +import "openengine/v1/model.proto"; +import "openengine/v1/server.proto"; + +// Standard ASCII gRPC request metadata: +// - openengine-routing-key: opaque application routing key. +// - openengine-target-dp-rank: base-10 uint32 routing target. This does not +// establish KV affinity; KvSessionRef.dp_rank is authoritative for a session. +// - openengine-priority: base-10 int32 admission priority; higher values have +// higher priority. +// - traceparent and tracestate: W3C Trace Context. +// +// The openengine- prefix is reserved for metadata defined by this protocol. +// Routing and admission keys apply to Inference RPCs. Trace Context applies to +// every RPC and should be propagated across downstream calls. + +service Inference { + // Core inference path. + rpc Generate(GenerateRequest) returns (stream GenerateResponse); +} + +service Control { + // Runtime metadata and scheduling state. + rpc GetServerInfo(GetServerInfoRequest) returns (ServerInfo); + rpc GetModelInfo(GetModelInfoRequest) returns (ModelInfo); + rpc GetLoad(GetLoadRequest) returns (LoadInfo); + + // Health and lifecycle. + rpc Health(HealthRequest) returns (HealthResponse); + rpc Abort(AbortRequest) returns (AbortResponse); + + // LoRA lifecycle. + rpc LoadLora(LoadLoraRequest) returns (LoadLoraResponse); + rpc UnloadLora(UnloadLoraRequest) returns (UnloadLoraResponse); + rpc ListLoras(ListLorasRequest) returns (ListLorasResponse); + + // Disaggregated serving / KV transfer. Connector info: ServerInfo.kv_connector. + rpc GetKvEventSources(GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse); + rpc SubscribeKvEvents(SubscribeKvEventsRequest) returns (stream SubscribeKvEventsResponse); +} diff --git a/lib/sidecar/trtllm/proto/openengine/v1/server.proto b/lib/sidecar/trtllm/proto/openengine/v1/server.proto new file mode 100644 index 000000000000..f039aa1d179b --- /dev/null +++ b/lib/sidecar/trtllm/proto/openengine/v1/server.proto @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Server identity, deployment capacity, roles, parallelism, and load. + +syntax = "proto3"; + +package openengine.v1; + +import "google/protobuf/struct.proto"; +import "openengine/v1/kv.proto"; + +enum EngineRole { + ENGINE_ROLE_UNSPECIFIED = 0; + ENGINE_ROLE_AGGREGATED = 1; + ENGINE_ROLE_PREFILL = 2; + ENGINE_ROLE_DECODE = 3; +} + +message GetServerInfoRequest {} + +message ServerInfo { + string engine_name = 1; // sglang, vllm, tensorrt_llm, etc. + string engine_version = 2; + EngineRole engine_role = 3; + string instance_id = 4; + repeated string supported_models = 5; + ParallelismInfo parallelism = 6; + KvConnectorInfo kv_connector = 7; + uint32 schema_revision = 8; // Contract revision implemented; zero is invalid. + uint32 minimum_client_revision = 9; // Oldest compatible contract revision. + // Immutable BSR module commit; unpublished builds may use a source commit. + string schema_release = 10; + DeploymentCapacity capacity = 11; // Configured capacity for this deployed server. + + // Engine-specific server metadata with no portable field (e.g. attention + // backend, build flags, experimental capabilities). NOT part of the portable + // contract: clients read it opportunistically and never depend on it for + // correctness. Carried as a Struct so JSON types survive the wire. + google.protobuf.Struct extra = 12; +} + +message DeploymentCapacity { + optional uint32 kv_block_size = 1; // Tokens per deployed KV-cache block. + optional uint64 total_kv_blocks = 2; // Allocatable KV blocks in the reporting scope. + optional uint64 max_running_requests = 3; // Concurrent running-request ceiling. + optional uint64 max_batched_tokens = 4; // Scheduler token ceiling per batch. + optional uint32 max_loras = 5; // Maximum simultaneously resident LoRA adapters. +} + +message ParallelismInfo { + optional uint32 tensor_parallel_size = 1; + optional uint32 pipeline_parallel_size = 2; + optional uint32 data_parallel_size = 3; + optional uint32 data_parallel_rank = 4; + optional uint32 data_parallel_start_rank = 5; + optional uint32 decode_context_parallel_size = 6; // Ranks per decode-context group; at least 1. +} + +message GetLoadRequest { + bool include_per_rank = 1; +} + +message LoadInfo { + string instance_id = 1; + optional uint64 timestamp_unix_nanos = 2; + optional uint32 running_requests = 3; + optional uint32 queued_requests = 4; + optional uint32 active_kv_sessions = 5; + optional uint64 used_kv_blocks = 6; + optional uint64 total_kv_blocks = 7; + optional uint64 running_tokens = 8; + optional uint64 waiting_tokens = 9; + optional uint32 prefill_batch_size = 10; + optional uint32 decode_batch_size = 11; + repeated RankLoadInfo ranks = 20; + google.protobuf.Struct attributes = 30; +} + +message RankLoadInfo { + optional uint32 data_parallel_rank = 1; + optional uint32 running_requests = 2; + optional uint32 queued_requests = 3; + optional uint64 used_kv_blocks = 4; + optional uint64 total_kv_blocks = 5; + optional uint32 prefill_batch_size = 6; + optional uint32 decode_batch_size = 7; +} diff --git a/lib/sidecar/trtllm/proto/trtllm_service.proto b/lib/sidecar/trtllm/proto/trtllm_service.proto deleted file mode 100644 index 8e666b3c99e8..000000000000 --- a/lib/sidecar/trtllm/proto/trtllm_service.proto +++ /dev/null @@ -1,517 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-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. - -// Derived from the `smg-grpc-proto` package (trtllm_service.proto). The -// upstream file imports `common.proto` for two optional RPCs — GetTokenizer -// and SubscribeKvEvents — that this sidecar does not use. Those RPCs and the -// `import "common.proto"` are removed so the contract compiles standalone; -// every message the sidecar exchanges is preserved verbatim. See -// `proto/README.md` for the source revision and checksum. - -syntax = "proto3"; -package trtllm; - -// TensorRT-LLM gRPC Service -// -// This service provides high-performance inference for LLMs using TensorRT-LLM. -// It accepts pre-tokenized requests and returns raw token IDs, enabling efficient -// binary communication with external routers (e.g., Shepherd Model Gateway). -// -// Key Design Principles: -// - Pre-tokenized input: Router tokenizes prompt, server tokenizes stop/bad words -// - Streaming delta mode: Chunks contain only new tokens since last chunk -// - Full feature support: All TensorRT-LLM capabilities exposed - -service TrtllmService { - // Generate tokens from pre-tokenized input - // Returns a stream of responses containing token IDs - rpc Generate(GenerateRequest) returns (stream GenerateResponse); - - // Generate embeddings from pre-tokenized input (for embedding models) - rpc Embed(EmbedRequest) returns (EmbedResponse); - - // Health check endpoint - rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse); - - // Abort a running generation request - rpc Abort(AbortRequest) returns (AbortResponse); - - // Get model information (vocab size, max lengths, etc.) - rpc GetModelInfo(GetModelInfoRequest) returns (GetModelInfoResponse); - - // Get server information (version, parallelism, etc.) - rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse); -} - -// ============================================================================ -// Generate Request -// ============================================================================ - -message GenerateRequest { - // Unique request identifier (assigned by client/router) - string request_id = 1; - - // Pre-tokenized input (REQUIRED - router tokenizes) - TokenizedInput tokenized = 2; - - // Sampling configuration - SamplingConfig sampling_config = 3; - - // Output configuration - OutputConfig output_config = 4; - - // Maximum tokens to generate (REQUIRED) - uint32 max_tokens = 5; - - // Enable streaming mode - bool streaming = 6; - - reserved 7, 8; // Removed: end_id, pad_id (resolved by TRT-LLM's _setup()) - reserved 9, 10; // Removed: bad_words, stop_words (replaced by string fields below) - - // Guided decoding parameters (JSON schema, regex, grammar) - optional GuidedDecodingParams guided_decoding = 11; - - // Embedding bias tensor (vocab_size floats, optional) - repeated float embedding_bias = 12; - - // LoRA adapter configuration - optional LoraConfig lora_config = 13; - - // Prompt tuning configuration - optional PromptTuningConfig prompt_tuning_config = 14; - - // Multimodal input (for VLMs) - optional MultimodalInput multimodal_input = 15; - - // KV cache retention configuration - optional KvCacheRetentionConfig kv_cache_retention = 16; - - // Disaggregated inference parameters - optional DisaggregatedParams disaggregated_params = 17; - - // Lookahead decoding configuration - optional LookaheadConfig lookahead_config = 18; - - // Cache salt ID for cache hashing - optional int64 cache_salt_id = 19; - - // Request arrival time (unix timestamp for metrics) - optional double arrival_time = 20; - - // Stop strings (TRT-LLM tokenizes these internally via _setup()) - repeated string stop = 21; - - // Stop token IDs (generation stops when any of these tokens are produced) - repeated uint32 stop_token_ids = 22; - - // Ignore end-of-sequence token (continue generating past EOS) - bool ignore_eos = 23; - - // Bad word strings (TRT-LLM tokenizes these internally via _setup()) - repeated string bad = 24; - - // Bad word token IDs (generation avoids producing these tokens) - repeated uint32 bad_token_ids = 25; - - // Include stop tokens in output token IDs (default: false). - // When true, stop token IDs are retained in output_token_ids instead of - // being stripped. - bool include_stop_token_in_output = 26; -} - -// Tokenized input from router -message TokenizedInput { - // Original text (for debugging/logging only, not used for generation) - string original_text = 1; - - // Pre-tokenized input token IDs (REQUIRED) - repeated uint32 input_token_ids = 2; - - // Query token IDs for VLM star attention (optional) - repeated uint32 query_token_ids = 3; -} - -// ============================================================================ -// Sampling Configuration -// Maps to tensorrt_llm.bindings.executor.SamplingConfig -// ============================================================================ - -message SamplingConfig { - // Beam width (1 for sampling, >1 for beam search) - int32 beam_width = 1; - - // Number of sequences to return - uint32 num_return_sequences = 2; - - // Top-K sampling (0 = disabled, considers all tokens) - optional int32 top_k = 3; - - // Top-P (nucleus) sampling threshold - optional float top_p = 4; - - // Top-P minimum threshold for decay - optional float top_p_min = 5; - - // Top-P reset token IDs - optional int32 top_p_reset_ids = 6; - - // Top-P decay factor - optional float top_p_decay = 7; - - // Random seed for reproducibility - optional uint64 seed = 8; - - // Temperature for sampling (0 = greedy, higher = more random) - optional float temperature = 9; - - // Minimum tokens to generate before stopping - optional uint32 min_tokens = 10; - - // Beam search diversity rate - optional float beam_search_diversity_rate = 11; - - // Repetition penalty (>1 discourages, <1 encourages repetition) - optional float repetition_penalty = 12; - - // Presence penalty (penalizes tokens that have appeared) - optional float presence_penalty = 13; - - // Frequency penalty (penalizes based on frequency of appearance) - optional float frequency_penalty = 14; - - // Number of prompt tokens to ignore for penalties - optional int32 prompt_ignore_length = 15; - - // Length penalty for beam search - optional float length_penalty = 16; - - // Early stopping for beam search - optional int32 early_stopping = 17; - - // No repeat n-gram size - optional int32 no_repeat_ngram_size = 18; - - // Min-P sampling threshold - optional float min_p = 19; - - // Variable beam width array for beam search - repeated int32 beam_width_array = 20; -} - -// ============================================================================ -// Output Configuration -// Maps to tensorrt_llm.bindings.executor.OutputConfig -// ============================================================================ - -message OutputConfig { - // Number of top log probabilities to return per output token - optional int32 logprobs = 1; - - // Number of top log probabilities to return per prompt token - optional int32 prompt_logprobs = 2; - - // Return context logits tensor (large, use with caution) - bool return_context_logits = 3; - - // Return generation logits tensor (large, use with caution) - bool return_generation_logits = 4; - - // Exclude input tokens from output (set to true to enable; TRT-LLM defaults to true internally) - bool exclude_input_from_output = 5; - - // Return encoder output (for encoder-decoder models) - bool return_encoder_output = 6; - - // Return performance metrics - bool return_perf_metrics = 7; -} - -// ============================================================================ -// Guided Decoding -// ============================================================================ - -message GuidedDecodingParams { - // Guide type enumeration - enum GuideType { - GUIDE_TYPE_UNSPECIFIED = 0; - GUIDE_TYPE_JSON = 1; // JSON format (any valid JSON) - GUIDE_TYPE_JSON_SCHEMA = 2; // JSON with schema constraint - GUIDE_TYPE_REGEX = 3; // Regular expression constraint - GUIDE_TYPE_EBNF_GRAMMAR = 4; // EBNF grammar constraint - GUIDE_TYPE_STRUCTURAL_TAG = 5; // Structural tag (xgrammar backend) - } - - GuideType guide_type = 1; - - // Guide content (schema string, regex pattern, or grammar definition) - string guide = 2; -} - -// ============================================================================ -// LoRA Configuration -// ============================================================================ - -message LoraConfig { - // LoRA task/adapter ID - int64 task_id = 1; - - // LoRA weights (serialized tensor, optional if already cached) - optional bytes weights = 2; - - // LoRA config as JSON string - optional string config_json = 3; -} - -// ============================================================================ -// Prompt Tuning Configuration -// ============================================================================ - -message PromptTuningConfig { - // Embedding table (serialized tensor) - bytes embedding_table = 1; -} - -// ============================================================================ -// Multimodal Input (for Vision-Language Models) -// ============================================================================ - -message MultimodalInput { - // Raw image bytes (JPEG/PNG) for each image in the request. - // TRT-LLM's input processor handles hashing, position tracking, - // and vision encoding server-side. - repeated bytes image_data = 1; -} - -// ============================================================================ -// KV Cache Retention -// ============================================================================ - -message KvCacheRetentionConfig { - // Retention policy name - string policy = 1; - - // Additional configuration as JSON string - string config_json = 2; -} - -// ============================================================================ -// Disaggregated Inference -// ============================================================================ - -message DisaggregatedParams { - // Request type for disaggregated inference - enum RequestType { - REQUEST_TYPE_CONTEXT_AND_GENERATION = 0; // Normal full request - REQUEST_TYPE_CONTEXT_ONLY = 1; // Prefill only - REQUEST_TYPE_GENERATION_ONLY = 2; // Decode only - } - - RequestType request_type = 1; - - // Context request ID (links context and generation phases) - string ctx_request_id = 2; - - // Context phase parameters (for generation_only requests) - optional ContextPhaseParams context_phase_params = 3; -} - -message ContextPhaseParams { - // First generated token ID from context phase - uint32 first_gen_token_id = 1; - - // KV cache block pointers (serialized) - bytes kv_cache_blocks = 2; -} - -// ============================================================================ -// Lookahead Decoding -// ============================================================================ - -message LookaheadConfig { - int32 max_window_size = 1; - int32 max_ngram_size = 2; - int32 max_verification_set_size = 3; -} - -// ============================================================================ -// Generate Response -// ============================================================================ - -message GenerateResponse { - // Request ID echo - string request_id = 1; - - reserved 4; // was GenerateError — removed, errors use gRPC status - - // Response type (oneof ensures exactly one is set) - oneof response { - GenerateStreamChunk chunk = 2; // Streaming delta - GenerateComplete complete = 3; // Final response - } -} - -// Streaming chunk containing delta tokens (new tokens since last chunk) -message GenerateStreamChunk { - // NEW token IDs only (delta from previous chunk) - repeated uint32 token_ids = 1; - - // Beam/sequence index (for beam_width > 1 or n > 1) - uint32 sequence_index = 2; - - // Token counts for usage tracking - uint32 prompt_tokens = 3; - uint32 completion_tokens = 4; - uint32 cached_tokens = 5; - - // Log probabilities for this chunk's tokens (if requested) - repeated TokenLogprob logprobs = 6; -} - -// Final/complete response with all output tokens -message GenerateComplete { - // All output token IDs (cumulative, not delta) - repeated uint32 output_token_ids = 1; - - // Beam/sequence index - uint32 sequence_index = 2; - - // Finish reason: "stop", "length", "stop_word" - string finish_reason = 3; - - // Matched stop information (for stop sequences) - oneof matched_stop { - string matched_stop_str = 4; // Stop string that matched (wire-compatible with old stop_reason) - uint32 matched_token_id = 13; // Stop token ID that matched - } - - // Token counts for usage tracking - uint32 prompt_tokens = 5; - uint32 completion_tokens = 6; - uint32 cached_tokens = 7; - - // Generation log probabilities (if requested) - repeated TokenLogprob logprobs = 8; - - // Prompt log probabilities (if requested) - repeated TokenLogprob prompt_logprobs = 9; - - // Performance metrics (if requested) - optional PerfMetrics perf_metrics = 10; - - // Context logits (if requested) - serialized float tensor - optional bytes context_logits = 11; - - // Generation logits (if requested) - serialized float tensor - optional bytes generation_logits = 12; -} - -// Token log probability information -message TokenLogprob { - uint32 token_id = 1; - float logprob = 2; - - // Top alternative tokens and their log probabilities - repeated TopLogprob top_logprobs = 3; -} - -message TopLogprob { - uint32 token_id = 1; - float logprob = 2; -} - -// Performance metrics for request -message PerfMetrics { - double arrival_time = 1; // When request arrived - double first_scheduled_time = 2; // When first scheduled - double first_token_time = 3; // Time to first token (TTFT) - double last_token_time = 4; // When last token generated - double kv_cache_transfer_start = 5; // KV cache transfer start (disagg) - double kv_cache_transfer_end = 6; // KV cache transfer end (disagg) - int64 kv_cache_size = 7; // KV cache size in bytes -} - -// ============================================================================ -// Embed Request/Response (for embedding models) -// ============================================================================ - -message EmbedRequest { - string request_id = 1; - TokenizedInput tokenized = 2; -} - -message EmbedResponse { - string request_id = 1; - repeated float embedding = 2; // Embedding vector - uint32 prompt_tokens = 3; -} - -// ============================================================================ -// Health Check -// ============================================================================ - -message HealthCheckRequest {} - -message HealthCheckResponse { - string status = 1; // "OK" or error description -} - -// ============================================================================ -// Abort Request -// ============================================================================ - -message AbortRequest { - string request_id = 1; -} - -message AbortResponse { - bool success = 1; - string message = 2; -} - -// ============================================================================ -// Model Info -// ============================================================================ - -message GetModelInfoRequest {} - -message GetModelInfoResponse { - string model_id = 1; // Model identifier/path - int32 max_input_len = 2; // Maximum input length - int32 max_seq_len = 3; // Maximum sequence length (input + output) - int32 max_batch_size = 4; // Maximum batch size - int32 vocab_size = 5; // Vocabulary size - int32 hidden_size = 6; // Hidden dimension - int32 num_layers = 7; // Number of transformer layers - int32 num_heads = 8; // Number of attention heads - - // Supported features - repeated string supported_features = 9; // e.g., "lora", "guided_decoding" -} - -// ============================================================================ -// Server Info -// ============================================================================ - -message GetServerInfoRequest {} - -message GetServerInfoResponse { - string version = 1; // TensorRT-LLM version - string backend = 2; // "tensorrt" or "pytorch" - int32 tensor_parallel_size = 3; // TP size - int32 pipeline_parallel_size = 4; // PP size - int32 context_parallel_size = 5; // CP size - int32 world_size = 6; // Total world size -} diff --git a/lib/sidecar/trtllm/src/args.rs b/lib/sidecar/trtllm/src/args.rs index 5d9ce90b36a6..e5832bfb65e0 100644 --- a/lib/sidecar/trtllm/src/args.rs +++ b/lib/sidecar/trtllm/src/args.rs @@ -6,7 +6,7 @@ use dynamo_sidecar_common::SidecarArgs; #[derive(clap::Parser, Clone, Debug)] #[command( name = "dynamo-trtllm-sidecar", - about = "Run a Dynamo worker against TensorRT-LLM's native gRPC TrtllmService" + about = "Run a Dynamo worker against TensorRT-LLM's OpenEngine gRPC server" )] pub(crate) struct Args { #[command(flatten)] @@ -19,13 +19,12 @@ pub(crate) struct Args { /// Model maximum sequence length (input + output). Used to register the /// context length and to derive a default `max_tokens` when a request omits /// one. A value supplied here takes precedence over the context length - /// TensorRT-LLM's `GetModelInfo` gRPC reports; that report is used only when - /// this argument is omitted, and a disagreement is logged at WARN. Supply - /// this whenever the engine was started without `--max_seq_len`, because - /// TensorRT-LLM then reports its `max_input_len` default instead of a real - /// context length and `client::model_info` discards it. With neither - /// source, requests that omit `max_tokens` are rejected. See the note in - /// `convert.rs`. - #[arg(long, env = "TRTLLM_CONTEXT_LENGTH")] + /// `Control.GetModelInfo` reports; that report is used only when this + /// argument is omitted, and a disagreement is logged at WARN. Supply this + /// whenever the engine was started without `--max_seq_len`, because + /// TensorRT-LLM then leaves `max_context_length` unset and the sidecar has + /// no window to register. With neither source, requests that omit + /// `max_tokens` are rejected. + #[arg(long, env = "TRTLLM_CONTEXT_LENGTH", value_parser = clap::value_parser!(u32).range(1..))] pub context_length: Option, } diff --git a/lib/sidecar/trtllm/src/client.rs b/lib/sidecar/trtllm/src/client.rs index 477da7c6aaf9..d1f4f046a014 100644 --- a/lib/sidecar/trtllm/src/client.rs +++ b/lib/sidecar/trtllm/src/client.rs @@ -1,26 +1,45 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Thin client for TensorRT-LLM's native `trtllm.TrtllmService`. +//! Thin client for TensorRT-LLM's OpenEngine gRPC services (`openengine.v1`). use std::time::Duration; -use dynamo_backend_common::DynamoError; +use tokio::time::{Instant, sleep_until, timeout_at}; + +use dynamo_backend_common::{BackendError, DynamoError, ErrorType}; use dynamo_sidecar_common::{ DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPool, GrpcEndpoint, GrpcTransportConfig, connection_timeout, }; use tonic::transport::Channel; -pub(crate) use dynamo_sidecar_common::{engine_shutdown, invalid_argument, status_to_dynamo}; +pub(crate) use dynamo_sidecar_common::{ + cancelled, engine_shutdown, invalid_argument, status_to_dynamo, +}; use crate::proto as pb; -use crate::proto::trtllm_service_client::TrtllmServiceClient; +use crate::proto::control_client::ControlClient; +use crate::proto::inference_client::InferenceClient; -/// Deadline for the one-shot control RPCs issued at startup / on cancel, so a -/// connected-but-unresponsive server cannot hang `start` or `abort`. +/// Deadline for the one-shot `Control.Abort`, so a connected-but-unresponsive +/// server cannot hang cancellation. Startup is bounded by the operator's +/// `--grpc-startup-deadline-secs` instead. const RPC_TIMEOUT: Duration = Duration::from_secs(30); +/// The engine's limits, resolved at startup from `--context-length` and +/// `Control.GetModelInfo`. +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct ModelLimits { + /// Maximum input + output tokens; backs both the registered context window + /// and the default `max_tokens`. Absent when neither source supplied one, + /// which leaves requests that omit `max_tokens` to be rejected. + pub(crate) context_length: Option, + /// Cap on generated tokens, when the engine reports one. A context window + /// alone can imply a larger budget than the engine will accept. + pub(crate) max_output_tokens: Option, +} + pub(crate) struct TrtllmClient { pool: GrpcChannelPool, } @@ -41,80 +60,160 @@ impl TrtllmClient { self.pool.len() } - fn client(&self) -> TrtllmServiceClient { - TrtllmServiceClient::new(self.pool.next_channel()) + fn inference(&self) -> InferenceClient { + InferenceClient::new(self.pool.next_channel()) + .max_decoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) + .max_encoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) + } + + fn control(&self) -> ControlClient { + ControlClient::new(self.pool.next_channel()) .max_decoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) .max_encoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) } + /// No routing metadata is attached: the server rejects both + /// `openengine-target-dp-rank` and `openengine-priority` with UNIMPLEMENTED, + /// so `convert::validate_request` refuses those requests up front rather + /// than letting the engine fail them. `KvSessionRef.dp_rank` still carries a + /// disaggregated session's KV affinity, inside the request body. pub(crate) async fn generate( &self, request: pb::GenerateRequest, ) -> Result, DynamoError> { - self.client() - .generate(request) + self.inference() + .generate(tonic::Request::new(request)) .await .map(tonic::Response::into_inner) .map_err(|status| status_to_dynamo("Generate", status)) } - /// Queries `GetModelInfo` and returns the reported maximum sequence length - /// (input + output) as the registration context length, if it is usable. + /// Reads the server's advertised limits, keeping positive values. /// - /// TensorRT-LLM answers `max_seq_len` with `args.max_seq_len or - /// args.max_input_len`, and on the PyTorch path `max_seq_len` is unset - /// while `max_input_len` holds its 1024 default. An equal pair is that - /// fallback's signature: the number describes the argument defaults rather - /// than the model, so it is discarded in favour of `--context-length` or - /// the context length the frontend reads from the model itself. - pub(crate) async fn model_info(&self) -> Result, DynamoError> { - let info = tokio::time::timeout( - RPC_TIMEOUT, - self.client().get_model_info(pb::GetModelInfoRequest {}), - ) - .await - .map_err(|_| { - connection_timeout(format!( - "GetModelInfo did not respond within {RPC_TIMEOUT:?}" - )) - })? - .map(tonic::Response::into_inner) - .map_err(|status| status_to_dynamo("GetModelInfo", status))?; - let max_seq_len = u32::try_from(info.max_seq_len).ok().filter(|len| *len > 0); - let max_input_len = u32::try_from(info.max_input_len) - .ok() - .filter(|len| *len > 0); - if max_seq_len.is_some() && max_seq_len == max_input_len { - tracing::warn!( - reported_context_length = max_seq_len, - "TensorRT-LLM reported max_seq_len == max_input_len, which is how it \ - answers when the engine was started without --max_seq_len; ignoring \ - the report" - ); - return Ok(None); + /// An engine started without `--max_seq_len` leaves `max_context_length` + /// unset rather than substituting its `max_input_len` default (measured + /// against TensorRT-LLM main at 8bbaf66bd5), so `None` here means "the + /// server did not say", and `--context-length` supplies it instead. + async fn get_model_info(&self, model: &str) -> Result { + let info = self + .control() + .get_model_info(pb::GetModelInfoRequest { + model: model.to_string(), + }) + .await? + .into_inner(); + Ok(ModelLimits { + context_length: info.max_context_length.filter(|len| *len > 0), + max_output_tokens: info.max_output_tokens.filter(|cap| *cap > 0), + }) + } + + /// One `GetModelInfo` call, for when `--context-length` already supplies the + /// window and the engine is consulted only to cross-check it and to learn + /// its output cap. + pub(crate) async fn model_limits(&self, model: &str) -> Result { + tokio::time::timeout(RPC_TIMEOUT, self.get_model_info(model)) + .await + .map_err(|_| { + connection_timeout(format!( + "GetModelInfo did not respond within {RPC_TIMEOUT:?}" + )) + })? + .map_err(|status| status_to_dynamo("GetModelInfo", status)) + } + + /// Polls `GetModelInfo` until the engine reports a usable context length or + /// `deadline` passes. + /// + /// Used only when no `--context-length` was supplied, which makes the engine + /// the sole source. TensorRT-LLM binds its gRPC port only once the model has + /// loaded, so waiting is for the engine not being up yet, and only transport + /// failures are worth retrying. An engine that answers has finished loading: + /// if that answer carries no window it was started without `--max_seq_len` + /// and will say the same thing at the deadline, so fail now with the fix + /// rather than polling for half an hour and then blaming a slow load. An + /// answer that says the request itself is wrong ends the wait for the same + /// reason. + pub(crate) async fn wait_for_model_limits( + &self, + model: &str, + deadline: Instant, + retry_interval: Duration, + ) -> Result { + let mut last = "it never answered".to_string(); + loop { + match timeout_at(deadline, self.get_model_info(model)).await { + Ok(Ok(limits)) if limits.context_length.is_some() => return Ok(limits), + Ok(Ok(_)) => { + return Err(protocol_error(format!( + "TensorRT-LLM answered GetModelInfo for {model} with no \ + max_context_length, which means it was started without \ + --max_seq_len. Start the engine with --max_seq_len, or give the \ + sidecar --context-length." + ))); + } + Ok(Err(status)) if is_request_itself_wrong(&status) => { + return Err(status_to_dynamo("GetModelInfo", status)); + } + Ok(Err(status)) => last = format!("{}: {}", status.code(), status.message()), + Err(_) => break, + } + let next_attempt = Instant::now() + retry_interval; + if next_attempt >= deadline { + break; + } + tracing::info!(model, reason = %last, "waiting for TensorRT-LLM GetModelInfo"); + sleep_until(next_attempt).await; } - Ok(max_seq_len) + Err(connection_timeout(format!( + "TensorRT-LLM never answered GetModelInfo before the gRPC startup deadline \ + ({last}). Raise --grpc-startup-deadline-secs if the engine is still starting, \ + or pin the window with --context-length." + ))) } pub(crate) async fn abort(&self, request_id: String) -> Result<(), DynamoError> { let response = tokio::time::timeout( RPC_TIMEOUT, - self.client().abort(pb::AbortRequest { request_id }), + self.control().abort(pb::AbortRequest { + target: Some(pb::abort_request::Target::RequestId(request_id)), + }), ) .await .map_err(|_| connection_timeout(format!("Abort did not respond within {RPC_TIMEOUT:?}")))? .map(tonic::Response::into_inner) .map_err(|status| status_to_dynamo("Abort", status))?; - if !response.success { - return Err(protocol_error(format!( - "TensorRT-LLM rejected abort: {}", + // A request already gone (finished/never seen) is not an error; only an + // unspecified status is protocol drift. + match pb::AbortStatus::try_from(response.status) { + Ok(pb::AbortStatus::Aborted) | Ok(pb::AbortStatus::AlreadyFinished) => Ok(()), + _ => Err(protocol_error(format!( + "TensorRT-LLM returned an unexpected abort status: {}", response.message - ))); + ))), } - Ok(()) } } +/// Whether the server answered that the request itself is wrong -- a model it +/// does not serve, or no Control service at all. Waiting cannot change any of +/// these, so a startup probe stops rather than burning its whole deadline. +fn is_request_itself_wrong(status: &tonic::Status) -> bool { + matches!( + status.code(), + tonic::Code::InvalidArgument | tonic::Code::NotFound | tonic::Code::Unimplemented + ) +} + pub(crate) fn protocol_error(message: impl Into) -> DynamoError { dynamo_sidecar_common::protocol_error("TensorRT-LLM", message) } + +/// A generation failure the engine reported in-band via an `EngineError` event +/// (as opposed to a transport/validation gRPC status). +pub(crate) fn engine_error(message: impl Into) -> DynamoError { + DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::Unknown)) + .message(message) + .build() +} diff --git a/lib/sidecar/trtllm/src/convert.rs b/lib/sidecar/trtllm/src/convert.rs index 60492a2b1ce5..a36c861d695a 100644 --- a/lib/sidecar/trtllm/src/convert.rs +++ b/lib/sidecar/trtllm/src/convert.rs @@ -1,85 +1,144 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Conversion between Dynamo's `PreprocessedRequest` / `LLMEngineOutput` and -//! the TensorRT-LLM `TrtllmService` protobuf messages. +//! Conversion between Dynamo's `PreprocessedRequest` / `LLMEngineOutput` and the +//! TensorRT-LLM OpenEngine (`openengine.v1`) protobuf messages. //! -//! Scope: aggregated generation. Disaggregation, multimodal, LoRA, beam search, -//! and `n > 1` are rejected before dispatch — the `Generate` response contract -//! carries no disaggregation handoff, and the sidecar streams a single sequence. +//! Scope: aggregated and disaggregated (prefill/decode) generation. Multimodal, +//! LoRA, beam search, and `n > 1` are rejected before dispatch — the sidecar +//! streams a single sequence. +//! +//! Disaggregation is driven by [`DisaggregationMode`]: a prefill worker marks +//! its request `context_only` and returns the `PrefillReady` handoff as its +//! terminal chunk, and a decode worker replays that handoff in `kv.session`. +//! See [`crate::disagg`] for the codec. use std::collections::BTreeSet; use dynamo_backend_common::{ - DynamoError, FinishReason, LLMEngineOutput, PreprocessedRequest, StopReason, TopLogprob, usage, + CompletionUsage, DisaggregationMode, DynamoError, FinishReason, LLMEngineOutput, + PreprocessedRequest, PromptTokensDetails, StopReason, TopLogprob, usage, }; -use crate::client; +use crate::client::{self, ModelLimits}; +use crate::disagg; use crate::proto as pb; -/// Per-chunk logprobs: the selected-token logprob sequence plus the per-token -/// top-k alternatives, both aligned with the chunk's delta tokens. -type MappedLogprobs = (Option>, Option>>); +/// A chunk's delta token IDs, the selected-token logprob sequence, and the +/// per-token top-k alternatives, all aligned with each other. +type MappedTokens = (Vec, Option>, Option>>); pub(crate) fn build_generate_request( request: &PreprocessedRequest, request_id: &str, - context_length: Option, + model: &str, + limits: Option, + mode: DisaggregationMode, ) -> Result { - validate_request(request)?; + validate_request(request, mode)?; let sampling = &request.sampling_options; let stop = &request.stop_conditions; let output = &request.output_options; + // A prefill worker only needs the context phase; TensorRT-LLM still requires + // a positive budget, and one token is what the context phase produces. + // No minimum check here: the minimum applies to the decode worker's budget, + // and discovery pairs the two roles without requiring equal context lengths + // (`discovery/model_manager.rs`). This worker's own window can neither + // prove the request impossible nor prove it servable, so checking it would + // reject requests a wider decode window would serve. Releasing the handoff + // when decode does reject is the sound fix and belongs on that side. + let max_tokens = if mode.is_prefill() { + 1 + } else { + max_tokens(request, limits)? + }; + // The decode worker replays the prefill worker's session; the prefill worker + // marks its request `context_only` through `extra`. + let (kv, extra) = match mode { + DisaggregationMode::Prefill => (None, Some(disagg::context_only_extra())), + // A handoff is the normal case, but conditional disaggregation + // dispatches straight to a decode worker with none, expecting it to run + // the context phase itself. Requiring one would fail every bypassed + // request. + DisaggregationMode::Decode => match request.prefill_result.as_ref() { + Some(handoff) => ( + Some(pb::KvOptions { + session: Some(disagg::session_from_json(&handoff.disaggregated_params)?), + ..Default::default() + }), + None, + ), + None => (None, None), + }, + DisaggregationMode::Aggregated | DisaggregationMode::Encode => (None, None), + }; + Ok(pb::GenerateRequest { request_id: request_id.to_string(), - tokenized: Some(pb::TokenizedInput { - original_text: String::new(), - input_token_ids: request.token_ids.as_ref().clone(), - query_token_ids: Vec::new(), - }), - sampling_config: Some(pb::SamplingConfig { - beam_width: 1, - num_return_sequences: 1, + // The OpenEngine server rejects an empty model; any non-empty name is + // served by the loaded model (single-model server). + model: model.to_string(), + input: Some(pb::generate_request::Input::TokenIds(pb::TokenIds { + ids: request.token_ids.as_ref().clone(), + })), + sampling: Some(pb::SamplingParams { + temperature: sampling.temperature.map(f64::from), + top_p: sampling.top_p.map(f64::from), top_k: normalize_top_k(sampling.top_k)?, - top_p: sampling.top_p, - temperature: sampling.temperature, - min_p: sampling.min_p, + min_p: sampling.min_p.map(f64::from), + frequency_penalty: sampling.frequency_penalty.map(f64::from), + presence_penalty: sampling.presence_penalty.map(f64::from), + repetition_penalty: sampling.repetition_penalty.map(f64::from), seed: normalize_seed(sampling.seed)?, + num_sequences: Some(1), + }), + stopping: Some(pb::StoppingOptions { + max_tokens: Some(max_tokens), + // Forwarded on a context-only request too. The minimum only masks + // EOS until the sequence reaches it, so it cannot extend a context + // phase that `max_tokens` already caps at one token. Dropping it + // would let the prefill worker sample EOS on that token, finish with + // `Stop` instead of `PrefillReady`, and leave the frontend treating + // a request that never met its minimum as complete. min_tokens: stop.min_tokens, - repetition_penalty: sampling.repetition_penalty, - presence_penalty: sampling.presence_penalty, - frequency_penalty: sampling.frequency_penalty, - ..Default::default() + conditions: stop_conditions(request), + ignore_eos: stop.ignore_eos, + // `include_stop_in_output` retains matched stop *strings*; the + // request-level flag is rejected in `validate_request`, so leave it + // unset (the server strips them). + include_stop_in_output: None, }), - output_config: Some(pb::OutputConfig { - logprobs: output.logprobs.map(request_logprob_count).transpose()?, - // prompt_logprobs is intentionally not requested: it is rejected in - // `validate_request` (no `LLMEngineOutput` field to surface it). - // TRT-LLM streams delta tokens; input tokens must not be echoed. - exclude_input_from_output: true, - ..Default::default() + response: Some(pb::ResponseOptions { + // A prefill worker streams no tokens to the client, but it must + // still be asked for logprobs: the context phase produces the first + // generated token, and its logprob only reaches the decode worker + // if the context request computed it (the server carries it as + // `first_gen_log_probs` in the handoff). Suppressing it here makes + // the decode worker report that first token as missing its logprob. + return_output_logprobs: Some(output.logprobs.is_some()), + output_candidates: output.logprobs.map(output_candidates), + // Prompt logprobs are rejected in `validate_request` (no + // `LLMEngineOutput` field to surface them). + return_prompt_logprobs: None, + prompt_candidates: None, + prompt_logprob_start: None, }), - max_tokens: max_tokens(request, context_length)?, - streaming: true, - guided_decoding: guided_decoding(request)?, - stop: stop.stop.clone().unwrap_or_default(), - stop_token_ids: stop_token_ids(request), - ignore_eos: stop.ignore_eos.unwrap_or(false), - // `include_stop_token_in_output` retains matched stop *token IDs*, a - // different axis from Dynamo's `include_stop_str_in_output` (stop - // *strings*). Mapping one to the other would leak hidden stop/EOS tokens, - // so leave it at the default (strip) and reject the request-level flag in - // `validate_request` instead. - ..Default::default() + guided: guided_decoding(request)?, + // Text generation only: no multimodal media or LoRA selection. + media: Vec::new(), + lora_name: String::new(), + kv, + extra, }) } -// Temporary workaround: TensorRT-LLM's gRPC `GenerateRequest.max_tokens` is -// REQUIRED (see proto/trtllm_service.proto), so an omitted `max_tokens` — which -// the Dynamo frontend forwards as `None` for the backend to default — has no -// natural value. We mirror the in-process backend's text-only default, +// Temporary workaround: the OpenEngine contract makes `stopping.max_tokens` +// optional, but an omitted value falls through to TensorRT-LLM's small +// `SamplingParams` default rather than filling the context. The Dynamo frontend +// forwards an omitted `max_tokens` as `None` expecting the backend to default, +// so we mirror the in-process backend's text-only default, // `max(1, context_length - prompt_len)` (components/src/dynamo/trtllm // `_default_max_tokens`); the sidecar rejects multimodal before dispatch, so // `token_ids.len()` is the true prompt length. `context_length` is resolved in @@ -97,25 +156,59 @@ pub(crate) fn build_generate_request( // not govern. fn max_tokens( request: &PreprocessedRequest, - context_length: Option, + limits: Option, ) -> Result { if let Some(max_tokens) = request.stop_conditions.max_tokens { return Ok(max_tokens); } - let context_length = context_length.ok_or_else(|| { - client::invalid_argument( - "TensorRT-LLM requires max_tokens, and no model context length is available to \ - derive a default; specify max_tokens explicitly", - ) - })?; + let context_length = limits + .and_then(|limits| limits.context_length) + .ok_or_else(|| { + client::invalid_argument( + "TensorRT-LLM requires max_tokens, and no model context length is known to \ + derive a default from; specify max_tokens explicitly or start the sidecar \ + with --context-length", + ) + })?; + let limits = limits.unwrap_or_default(); let prompt_len = request.token_ids.len() as u32; - Ok(context_length.saturating_sub(prompt_len).max(1)) + let remaining = context_length.saturating_sub(prompt_len).max(1); + // The window is input + output, so on a short prompt the remainder can + // exceed what the engine will actually generate and it would reject the + // request we derived. + let derived = match limits.max_output_tokens { + Some(cap) => remaining.min(cap), + None => remaining, + }; + reject_impossible_minimum(request, derived, context_length, prompt_len)?; + Ok(derived) +} + +/// Both the output cap and the `.max(1)` floor can land under an explicit +/// minimum. Sending `min_tokens` above `max_tokens` is a request no engine can +/// honour, and TensorRT-LLM does not cross-validate the pair, so say which two +/// values conflict instead of letting it resolve them silently. +fn reject_impossible_minimum( + request: &PreprocessedRequest, + derived: u32, + context_length: u32, + prompt_len: u32, +) -> Result<(), DynamoError> { + if let Some(min_tokens) = request.stop_conditions.min_tokens + && min_tokens > derived + { + return Err(client::invalid_argument(format!( + "min_tokens ({min_tokens}) exceeds the {derived} tokens left for this request; \ + the {context_length}-token window already holds a {prompt_len}-token prompt" + ))); + } + Ok(()) } fn normalize_top_k(top_k: Option) -> Result, DynamoError> { - // Dynamo uses -1/0 (or absence) for "consider all tokens"; TRT-LLM treats an - // unset top_k the same way. Forward only a positive cap; reject other - // negatives rather than silently widening them to "all tokens". + // Dynamo uses -1/0 (or absence) for "consider all tokens"; the OpenEngine + // server treats an unset top_k the same way. Forward only a positive cap; + // reject other negatives rather than silently widening them to "all tokens". match top_k { None | Some(-1) | Some(0) => Ok(None), Some(value) if value > 0 => Ok(Some(value)), @@ -126,7 +219,7 @@ fn normalize_top_k(top_k: Option) -> Result, DynamoError> { } fn normalize_seed(seed: Option) -> Result, DynamoError> { - // TRT-LLM's seed is `uint64`; reject a negative seed rather than silently + // OpenEngine's seed is `uint64`; reject a negative seed rather than silently // dropping it and losing reproducibility. seed.map(|seed| { u64::try_from(seed) @@ -135,14 +228,33 @@ fn normalize_seed(seed: Option) -> Result, DynamoError> { .transpose() } -fn request_logprob_count(count: u32) -> Result { - // TRT-LLM computes the selected-token logprob only when at least one logprob - // is requested, so floor the wire value at 1: `logprobs=0` (selected token, - // no alternatives) still yields the chosen-token logprob. The original count - // is preserved in `ResponseState` to decide whether to surface alternatives. - i32::try_from(count.max(1)).map_err(|_| { - client::invalid_argument(format!("logprobs request must fit in i32; got {count}")) - }) +fn output_candidates(count: u32) -> pb::CandidateTokenSelection { + // TRT-LLM computes the selected-token logprob only when at least one + // candidate is requested, so floor the wire value at 1: `logprobs=0` + // (selected token, no alternatives) still yields the chosen-token logprob. + // The original count is preserved in `ResponseState` to decide whether to + // surface alternatives. + pb::CandidateTokenSelection { + selection: Some(pb::candidate_token_selection::Selection::TopN(count.max(1))), + } +} + +fn stop_conditions(request: &PreprocessedRequest) -> Vec { + let stop = &request.stop_conditions; + let mut conditions = Vec::new(); + if let Some(stop_strings) = stop.stop.as_ref() { + for text in stop_strings { + conditions.push(pb::StopCondition { + condition: Some(pb::stop_condition::Condition::StopText(text.clone())), + }); + } + } + for id in stop_token_ids(request) { + conditions.push(pb::StopCondition { + condition: Some(pb::stop_condition::Condition::StopTokenId(id)), + }); + } + conditions } fn stop_token_ids(request: &PreprocessedRequest) -> Vec { @@ -162,46 +274,46 @@ fn stop_token_ids(request: &PreprocessedRequest) -> Vec { fn guided_decoding( request: &PreprocessedRequest, -) -> Result, DynamoError> { +) -> Result, DynamoError> { let Some(guided) = request.sampling_options.guided_decoding.as_ref() else { return Ok(None); }; if guided.backend.is_some() || guided.whitespace_pattern.is_some() { return Err(client::invalid_argument( - "guided decoding backend and whitespace_pattern are not supported by TensorRT-LLM gRPC", + "guided decoding backend and whitespace_pattern are not supported by the TensorRT-LLM OpenEngine server", )); } - use pb::guided_decoding_params::GuideType; - let mut constraints = Vec::new(); + use pb::guided_decoding::Guide; + let mut guides = Vec::new(); if let Some(json) = &guided.json { - constraints.push((GuideType::JsonSchema, json_guide(json))); + guides.push(Guide::JsonSchema(json_guide(json))); } if let Some(regex) = &guided.regex { - constraints.push((GuideType::Regex, regex.clone())); + guides.push(Guide::Regex(regex.clone())); } if let Some(grammar) = &guided.grammar { - constraints.push((GuideType::EbnfGrammar, grammar.clone())); + guides.push(Guide::EbnfGrammar(grammar.clone())); } if let Some(tag) = &guided.structural_tag { - constraints.push((GuideType::StructuralTag, json_guide(tag))); + guides.push(Guide::StructuralTag(json_guide(tag))); } if guided.choice.is_some() { + // `choice` is in the OpenEngine contract but the TensorRT-LLM server + // rejects it; fail fast with a clear message. return Err(client::invalid_argument( - "guided decoding `choice` is not supported by TensorRT-LLM gRPC", + "guided decoding `choice` is not supported by the TensorRT-LLM OpenEngine server", )); } - if constraints.len() > 1 { + if guides.len() > 1 { return Err(client::invalid_argument( "only one guided decoding constraint may be set", )); } - Ok(constraints - .pop() - .map(|(guide_type, guide)| pb::GuidedDecodingParams { - guide_type: guide_type as i32, - guide, - })) + Ok(guides.pop().map(|guide| pb::GuidedDecoding { + guide: Some(guide), + backend: String::new(), + })) } /// A guide is either a JSON string carried verbatim or a JSON value rendered to @@ -213,7 +325,10 @@ fn json_guide(value: &serde_json::Value) -> String { } } -fn validate_request(request: &PreprocessedRequest) -> Result<(), DynamoError> { +fn validate_request( + request: &PreprocessedRequest, + mode: DisaggregationMode, +) -> Result<(), DynamoError> { if request.token_ids.is_empty() { return Err(client::invalid_argument("token_ids must not be empty")); } @@ -230,15 +345,18 @@ fn validate_request(request: &PreprocessedRequest) -> Result<(), DynamoError> { "multimodal requests are not supported by the TensorRT-LLM sidecar", )); } - if request.prefill_result.is_some() { - return Err(client::invalid_argument( - "disaggregated (prefill/decode) requests are not supported by the TensorRT-LLM sidecar", - )); + // Only a decode worker consumes a prefill handoff. Seeing one anywhere else + // means the frontend routed a decode request to the wrong worker role. + if request.prefill_result.is_some() && !mode.is_decode() { + return Err(client::invalid_argument(format!( + "received a prefill_result on a worker running in '{mode}' mode; \ + disaggregated decode requests must be routed to a decode worker" + ))); } if request.output_options.prompt_logprobs.is_some() { - // TRT-LLM would compute these, but the terminal `complete` carries them - // with no `LLMEngineOutput` field to surface — reject rather than pay for - // and drop them. + // TRT-LLM would compute these, but the terminal carries them with no + // `LLMEngineOutput` field to surface — reject rather than pay for and + // drop them. return Err(client::invalid_argument( "prompt logprobs are not supported by the TensorRT-LLM sidecar", )); @@ -250,8 +368,9 @@ fn validate_request(request: &PreprocessedRequest) -> Result<(), DynamoError> { .is_some_and(|ids| !ids.is_empty()) { // A visible stop token must halt generation *and* stay in the output. - // TRT-LLM's single `include_stop_token_in_output` flag cannot honor that - // per token, so reject rather than silently drop or mis-retain them. + // OpenEngine's stop conditions and single `include_stop_in_output` (stop + // *strings*) cannot honor that per token, so reject rather than silently + // drop or mis-retain them. return Err(client::invalid_argument( "visible stop token IDs are not supported by the TensorRT-LLM sidecar", )); @@ -269,38 +388,46 @@ fn validate_request(request: &PreprocessedRequest) -> Result<(), DynamoError> { if request .routing .as_ref() - .is_some_and(|routing| routing.dp_rank.is_some() || routing.prefill_dp_rank.is_some()) + .and_then(|routing| routing.cache_namespace.as_deref()) + .is_some_and(|namespace| !namespace.is_empty()) { + // `cache_namespace` is the request-scoped KV-cache isolation contract. + // The OpenEngine request carries `kv.cache_salt`, but the sidecar does + // not yet map the namespace onto it, so honoring the request would let + // requests from different namespaces share prefix-cache entries. Reject + // until the mapping is implemented. return Err(client::invalid_argument( - "KV-aware data-parallel routing is not supported by the TensorRT-LLM sidecar", + "cache namespace isolation is not supported by the TensorRT-LLM sidecar", )); } if request .routing .as_ref() - .and_then(|routing| routing.cache_namespace.as_deref()) - .is_some_and(|namespace| !namespace.is_empty()) + .and_then(|routing| routing.priority) + .is_some_and(|priority| priority != 0) { - // `cache_namespace` is the request-scoped KV-cache isolation contract. The - // proto carries `cache_salt_id`, but TRT-LLM 1.3.0rc21 does not plumb it - // through the gRPC servicer, so honoring it would let requests from - // different namespaces share prefix-cache entries. Reject until supported. + // OpenEngine carries priority as request metadata, which the TensorRT-LLM + // server rejects; a nonzero priority would otherwise be silently dropped + // and change queue ordering. Reject until the server honors it. return Err(client::invalid_argument( - "cache namespace isolation is not supported by the TensorRT-LLM sidecar", + "request priority is not supported by the TensorRT-LLM sidecar", )); } if request .routing .as_ref() - .and_then(|routing| routing.priority) - .is_some_and(|priority| priority != 0) + .is_some_and(|routing| routing.dp_rank.is_some() || routing.prefill_dp_rank.is_some()) { - // The shared contract forwards a nonzero priority to the engine, but the - // TRT-LLM gRPC request has no priority field, so it would be silently - // dropped and change queue ordering. Reject until the wire protocol - // carries it. + // The same server branch that rejects `openengine-priority` also rejects + // `openengine-target-dp-rank` (`grpc/openengine/request_mapping.py`, + // `_trace_headers`), and the servicer turns that into UNIMPLEMENTED -- + // measured against TensorRT-LLM main at 8bbaf66bd5, rank 0 included. + // Sending it anyway failed the whole request with a non-migratable + // 5xx; rejecting here names the unsupported feature in a 4xx instead. + // `nvext.dp_rank` and the `x-dynamo-dp-rank` header both reach this + // field, so it is reachable without a KV router. return Err(client::invalid_argument( - "request priority is not supported by the TensorRT-LLM sidecar", + "data-parallel rank targeting is not supported by the TensorRT-LLM sidecar", )); } if request.stop_conditions.max_thinking_tokens.is_some() { @@ -311,10 +438,10 @@ fn validate_request(request: &PreprocessedRequest) -> Result<(), DynamoError> { } let sampling = &request.sampling_options; if sampling.include_stop_str_in_output == Some(true) { - // Retaining stop *strings* has no faithful mapping in the TRT-LLM gRPC - // request (its only control, `include_stop_token_in_output`, governs stop - // *token IDs* and would leak hidden stop/EOS tokens). Reject rather than - // mis-map — consistent with rejecting visible stop token IDs above. + // Retaining stop *strings* maps to `include_stop_in_output`, but the + // sidecar rejects visible stop token IDs above; honoring string retention + // alone would be inconsistent, so reject until both are supported + // together. return Err(client::invalid_argument( "include_stop_str_in_output is not supported by the TensorRT-LLM sidecar", )); @@ -331,22 +458,61 @@ fn validate_request(request: &PreprocessedRequest) -> Result<(), DynamoError> { Ok(()) } -/// Streaming response reducer. TRT-LLM streams delta `chunk`s followed by a -/// single terminal `complete`; this maps each onto an `LLMEngineOutput`. +/// Streaming response reducer. The OpenEngine server streams `token` events +/// followed by a terminal `finished` event carrying authoritative usage; this +/// maps each onto an `LLMEngineOutput`. pub(crate) struct ResponseState { prompt_tokens: u32, completion_tokens: u32, - streamed_tokens: u32, output_logprobs: Option, + phase: Phase, +} + +/// What this worker does with the tokens the engine streams back. Each arm owns +/// the state that only exists in that mode, so a prefill worker cannot reach a +/// decode worker's cache accounting and vice versa. +enum Phase { + /// Aggregated, or a decode leg: tokens stream straight to the client. + Stream { + /// The handoff this leg replays, if any. Cache hits are measured during + /// the context phase, so on a decode leg the count arrives here rather + /// than from the local engine -- which counts the blocks transferred + /// into it and reports a different quantity under the same name. + handoff: Option, + }, + /// A prefill leg: `PrefillReady` is the terminal event, and context tokens + /// are held back because the decode leg replays them. They only reach the + /// client through a context request that ends without a handoff, which has + /// no decode leg to do the replaying. + Prefill { held: Vec }, +} + +/// What a decode leg knows about the context phase that preceded it. +struct Handoff { + /// The context phase's prefix-cache hit count, when it reported one. Absent + /// means the context phase reported none -- not that the decode engine's + /// own count should stand in for it. + cached_tokens: Option, } impl ResponseState { - pub(crate) fn new(request: &PreprocessedRequest) -> Self { + pub(crate) fn new(request: &PreprocessedRequest, mode: DisaggregationMode) -> Self { Self { prompt_tokens: request.token_ids.len() as u32, completion_tokens: 0, - streamed_tokens: 0, output_logprobs: request.output_options.logprobs, + phase: if mode.is_prefill() { + Phase::Prefill { held: Vec::new() } + } else { + Phase::Stream { + handoff: request.prefill_result.as_ref().map(|prefill| Handoff { + cached_tokens: prefill + .prompt_tokens_details + .as_ref() + .and_then(|details| details.cached_tokens), + }), + } + }, } } @@ -362,60 +528,92 @@ impl ResponseState { &mut self, response: pb::GenerateResponse, ) -> Result, DynamoError> { - match response.response { - Some(pb::generate_response::Response::Chunk(chunk)) => self.convert_chunk(chunk), - Some(pb::generate_response::Response::Complete(complete)) => { - self.convert_complete(complete).map(Some) + let pb::GenerateResponse { event, usage, .. } = response; + // A response with no event is protocol drift, not an empty delta. + let Some(event) = event else { + return Err(client::protocol_error("response carried no event")); + }; + match event { + pb::generate_response::Event::Token(token) => self.convert_token(token), + pb::generate_response::Event::Finished(finished) => { + self.convert_finished(finished, usage).map(Some) } - // A response with neither payload is protocol drift, not an - // empty delta — fail rather than silently discard it. - None => Err(client::protocol_error( - "response carried neither a chunk nor a complete payload", + pb::generate_response::Event::PrefillReady(prefill) => { + self.convert_prefill_ready(prefill, usage).map(Some) + } + pb::generate_response::Event::Error(error) => Err(engine_error(error)), + // Prompt logprobs are never requested, so a prompt event is drift. + pb::generate_response::Event::Prompt(_) => Err(client::protocol_error( + "received an unexpected prompt event; prompt logprobs are not requested", )), } } - /// Validates the single-sequence invariant and folds the response's usage - /// counts into the running totals. Shared by chunk and terminal handling. - fn accumulate( + /// Terminal chunk for a prefill worker. + /// + /// `PrefillReady` *is* the terminal event for a `context_only` request — the + /// server suppresses `finished` because the engine reports the sequence as + /// unfinished — so this synthesizes the terminal chunk, carrying the handoff + /// the decode worker will replay. + fn convert_prefill_ready( &mut self, - sequence_index: u32, - prompt_tokens: u32, - completion_tokens: u32, - ) -> Result<(), DynamoError> { - if sequence_index != 0 { - return Err(client::protocol_error(format!( - "received unsupported sequence index {sequence_index}" - ))); - } - if prompt_tokens != 0 { - self.prompt_tokens = prompt_tokens; - } - // TRT-LLM reports cumulative (not per-chunk-delta) usage counts, verified - // against 1.3.0rc21 over multi-chunk streams; `max` tolerates chunks that - // omit the count (report 0) without regressing the running total. - self.completion_tokens = self.completion_tokens.max(completion_tokens); - Ok(()) + prefill: pb::PrefillReady, + reported: Option, + ) -> Result { + let Phase::Prefill { .. } = self.phase else { + return Err(client::protocol_error( + "received a prefill_ready event on a worker that is not running in prefill mode", + )); + }; + let session = prefill + .kv_session + .ok_or_else(|| client::protocol_error("prefill_ready event carried no kv_session"))?; + Ok(LLMEngineOutput { + // The client sees no tokens from prefill; the decode worker emits + // the full completion. + token_ids: Vec::new(), + index: Some(0), + // `Length` is what the frontend's prefill router requires to chain + // into decode: it returns any other terminal reason straight to the + // caller as an already-complete request (see + // `kv_router::prefill_router::admission`). It is also accurate — + // the context request is capped at one token. + finish_reason: Some(FinishReason::Length), + // PrefillReady is the final response for a context request, so it + // carries the engine's usage -- including cached_prompt_tokens, + // which the frontend forwards to the decode leg and which cannot be + // reconstructed client-side. + completion_usage: Some(prefill_usage(reported, self.prompt_tokens)), + disaggregated_params: Some(disagg::session_to_json(session)?), + ..Default::default() + }) } - fn convert_chunk( + fn convert_token( &mut self, - chunk: pb::GenerateStreamChunk, + token: pb::TokenOutput, ) -> Result, DynamoError> { - self.accumulate( - chunk.sequence_index, - chunk.prompt_tokens, - chunk.completion_tokens, - )?; - self.streamed_tokens = self - .streamed_tokens - .saturating_add(chunk.token_ids.len() as u32); - if chunk.token_ids.is_empty() { + check_output_index(token.output_index)?; + if let Phase::Prefill { held } = &mut self.phase { + // Whole `TokenInfo`s, not just the IDs: a context request that ends + // without a handoff returns these to the client, and by then the + // logprobs it asked for are gone if only the IDs were kept. + held.extend(token.tokens); + self.completion_tokens = held.len() as u32; + return Ok(None); + } + let (token_ids, log_probs, top_logprobs) = self.map_tokens(token.tokens)?; + self.completion_tokens = self + .completion_tokens + .saturating_add(token_ids.len() as u32); + if token_ids.is_empty() { + // A text-only delta (e.g. stop-string holdback): the sidecar streams + // token IDs and lets the frontend detokenize, so there is nothing to + // surface yet. return Ok(None); } - let (log_probs, top_logprobs) = self.map_logprobs(&chunk.token_ids, &chunk.logprobs)?; Ok(Some(LLMEngineOutput { - token_ids: chunk.token_ids, + token_ids, log_probs, top_logprobs, index: Some(0), @@ -423,114 +621,243 @@ impl ResponseState { })) } - fn convert_complete( + fn convert_finished( &mut self, - complete: pb::GenerateComplete, + finished: pb::GenerationFinished, + reported: Option, ) -> Result { - self.accumulate( - complete.sequence_index, - complete.prompt_tokens, - complete.completion_tokens, - )?; - - // The terminal echoes the cumulative output token IDs (proto: "cumulative, - // not delta"). When present they must match what we streamed; a mismatch - // is protocol drift. Tolerate an omitted echo (empty) rather than - // false-erroring a server that only carries finish state in the terminal. - if !complete.output_token_ids.is_empty() - && complete.output_token_ids.len() as u32 != self.streamed_tokens - { - return Err(client::protocol_error(format!( - "terminal output_token_ids ({}) do not match {} streamed delta tokens", - complete.output_token_ids.len(), - self.streamed_tokens - ))); + check_output_index(finished.output_index)?; + // A decode engine counts the blocks transferred into it as cache hits, + // a different quantity from the context phase's prefix-cache hit that + // would read as a ~100% hit rate on every disaggregated request. Once + // this leg took a handoff the context phase is the only source, even + // when it reported no hits. A leg without one ran its own context + // phase, so its engine is the source. + let (mut cached_tokens, engine_measured_the_cache) = match &self.phase { + Phase::Stream { + handoff: Some(handoff), + } => (handoff.cached_tokens, false), + Phase::Stream { handoff: None } | Phase::Prefill { .. } => (None, true), + }; + // The final response carries authoritative usage; prefer it over the + // counts accumulated while streaming. + if let Some(reported) = reported { + if reported.prompt_tokens != 0 { + self.prompt_tokens = reported.prompt_tokens; + } + if reported.completion_tokens != 0 { + self.completion_tokens = reported.completion_tokens; + } + if engine_measured_the_cache { + cached_tokens = reported.cached_prompt_tokens; + } } - // The delta tokens were already streamed via chunks; the terminal - // response only carries finish state and usage. + let finish_reason = match pb::FinishReason::try_from(finished.reason).map_err(|_| { + client::protocol_error(format!("unknown finish reason {}", finished.reason)) + })? { + pb::FinishReason::Stop => FinishReason::Stop, + pb::FinishReason::Length => FinishReason::Length, + pb::FinishReason::Cancelled => FinishReason::Cancelled, + // Fail closed on an unspecified reason rather than reporting a clean + // stop for a version-skewed or malformed terminal. + pb::FinishReason::Unspecified => { + return Err(client::protocol_error( + "terminal response has an unspecified finish reason", + )); + } + }; + + // A context request only reaches a `finished` event when it produced no + // handoff: the server suppresses the terminal once it has sent + // `PrefillReady`. `Length` there means the context phase burned its + // one-token budget without ever transmitting KV -- disaggregation is + // not working on this engine -- so returning that single token would + // silently truncate the completion. Any other reason is a genuine stop + // during the context phase, and with no decode leg to run, it is the + // whole answer. + let held = match &mut self.phase { + Phase::Prefill { held } if finish_reason == FinishReason::Length => { + return Err(client::protocol_error(format!( + "context request ended at its token budget without a kv_session handoff \ + after {} token(s); the engine ran the context phase but never transmitted \ + its KV, which usually means disaggregation is not configured on it (check \ + cache_transceiver_config)", + held.len() + ))); + } + Phase::Prefill { held } => std::mem::take(held), + Phase::Stream { .. } => Vec::new(), + }; + + let (token_ids, log_probs, top_logprobs) = if held.is_empty() { + (Vec::new(), None, None) + } else { + self.map_tokens(held)? + }; + let mut terminal = LLMEngineOutput { + token_ids, + log_probs, + top_logprobs, index: Some(0), - finish_reason: Some(finish_reason(&complete.finish_reason)?), - completion_usage: Some(usage(self.prompt_tokens, self.completion_tokens)), + finish_reason: Some(finish_reason), + completion_usage: Some(CompletionUsage { + prompt_tokens_details: cached_prompt_tokens(cached_tokens, self.prompt_tokens), + ..usage(self.prompt_tokens, self.completion_tokens) + }), ..Default::default() }; - terminal.stop_reason = complete.matched_stop.map(|matched| match matched { - pb::generate_complete::MatchedStop::MatchedStopStr(value) => StopReason::String(value), - pb::generate_complete::MatchedStop::MatchedTokenId(id) => { - StopReason::Int(i64::from(id)) - } - }); + terminal.stop_reason = finished + .stop_match + .and_then(|stop_match| stop_match.r#match) + .map(|matched| match matched { + pb::stop_match::Match::StopTokenId(id) | pb::stop_match::Match::EosTokenId(id) => { + StopReason::Int(i64::from(id)) + } + pb::stop_match::Match::StopText(text) => StopReason::String(text), + }); Ok(terminal) } - fn map_logprobs( - &self, - token_ids: &[u32], - logprobs: &[pb::TokenLogprob], - ) -> Result { + /// One pass over the owned deltas. Token IDs, selected-token logprobs, and + /// the candidate lists all come from the same `TokenInfo`s, and this runs on + /// every token of every request. + fn map_tokens(&self, tokens: Vec) -> Result { + let mut token_ids = Vec::with_capacity(tokens.len()); let Some(count) = self.output_logprobs else { - return Ok((None, None)); + token_ids.extend(tokens.into_iter().map(|info| info.token_id)); + return Ok((token_ids, None, None)); }; - // Logprobs were requested, so every delta token must carry one; an empty - // or short slice is a protocol violation, not a silent no-op. - if logprobs.len() != token_ids.len() { - return Err(client::protocol_error(format!( - "logprob count {} does not match {} delta tokens", - logprobs.len(), - token_ids.len() - ))); - } - let log_probs = logprobs.iter().map(|lp| f64::from(lp.logprob)).collect(); - if count == 0 { - // `logprobs=0` keeps the selected-token logprob but omits the top - // alternatives, matching the vLLM sidecar contract. - return Ok((Some(log_probs), None)); - } - let top_logprobs = logprobs - .iter() - .map(|lp| { - if lp.top_logprobs.is_empty() { + let mut log_probs = Vec::with_capacity(tokens.len()); + // `logprobs=0` keeps the selected-token logprob but omits the top + // alternatives, matching the vLLM sidecar contract. + let wants_candidates = count != 0; + let mut top_logprobs = Vec::with_capacity(if wants_candidates { tokens.len() } else { 0 }); + for info in tokens { + let token_id = info.token_id; + let rank = info.rank.unwrap_or(0); + // Logprobs were requested, so every delta token must carry its + // selected-token logprob; a missing value is protocol drift. + let logprob = info.logprob.ok_or_else(|| { + client::protocol_error(format!("token {token_id} is missing its logprob")) + })?; + if wants_candidates { + top_logprobs.push(if info.candidates.is_empty() { vec![TopLogprob { - rank: 0, - token_id: lp.token_id, + rank, + token_id, token: None, - logprob: f64::from(lp.logprob), + logprob, bytes: None, }] } else { - lp.top_logprobs - .iter() - .enumerate() - .map(|(rank, candidate)| TopLogprob { - rank: rank as u32, + info.candidates + .into_iter() + .map(|candidate| TopLogprob { + rank: candidate.rank.unwrap_or(0), token_id: candidate.token_id, token: None, - logprob: f64::from(candidate.logprob), + logprob: candidate.logprob, bytes: None, }) .collect() - } - }) - .collect(); - Ok((Some(log_probs), Some(top_logprobs))) + }); + } + token_ids.push(token_id); + log_probs.push(logprob); + } + Ok(( + token_ids, + Some(log_probs), + wants_candidates.then_some(top_logprobs), + )) } } -fn finish_reason(reason: &str) -> Result { - Ok(match reason.trim().to_ascii_lowercase().as_str() { - "length" => FinishReason::Length, - "cancelled" | "canceled" | "aborted" | "timeout" => FinishReason::Cancelled, - "error" => FinishReason::Error("TensorRT-LLM reported an error finish".to_string()), - // Proto-documented stop reasons ("stop", "stop_word"); an empty string is - // a bare terminal (finished implies stop). - "stop" | "stop_word" | "eos" | "" => FinishReason::Stop, - // Fail closed on an unrecognized reason rather than reporting a clean stop - // for a version-skewed or malformed terminal. - other => { - return Err(client::protocol_error(format!( - "unknown finish reason {other:?}" - ))); - } +/// The OpenEngine contract requires an explicit index even for output zero, and +/// the sidecar streams a single sequence: anything else is drift. +fn check_output_index(index: Option) -> Result<(), DynamoError> { + match index { + Some(0) => Ok(()), + Some(index) => Err(client::protocol_error(format!( + "received unsupported output index {index}" + ))), + None => Err(client::protocol_error( + "response carried no output index; the OpenEngine contract requires one", + )), + } +} + +/// Usage for the prefill terminal. `PrefillReady` *is* the final response for a +/// context request, so it carries the engine's authoritative counts; the +/// client-side prompt length is only a fallback for a server that omits them. +fn prefill_usage(reported: Option, prompt_tokens: u32) -> CompletionUsage { + let Some(reported) = reported else { + return usage(prompt_tokens, 0); + }; + let prompt_tokens = if reported.prompt_tokens != 0 { + reported.prompt_tokens + } else { + prompt_tokens + }; + CompletionUsage { + prompt_tokens_details: cached_prompt_tokens(reported.cached_prompt_tokens, prompt_tokens), + ..usage(prompt_tokens, 0) + } +} + +/// A server that counts cache hits against an expanded prompt would otherwise +/// report more cached tokens than the prompt the client actually sent. +fn cached_prompt_tokens(cached: Option, prompt_tokens: u32) -> Option { + cached.map(|cached_tokens| PromptTokensDetails { + audio_tokens: None, + cached_tokens: Some(cached_tokens.min(prompt_tokens)), }) } + +pub(crate) fn engine_error(error: pb::EngineError) -> DynamoError { + let code = pb::ErrorCode::try_from(error.code).unwrap_or(pb::ErrorCode::Unspecified); + let message = if error.message.trim().is_empty() { + format!("TensorRT-LLM reported engine error {code:?}") + } else { + format!("TensorRT-LLM engine error: {}", error.message) + }; + match code { + pb::ErrorCode::InvalidArgument | pb::ErrorCode::UnsupportedFeature => { + client::invalid_argument(message) + } + // Not `worker_overloaded`, despite the name. The only site in the + // TensorRT-LLM servicer that emits this code is the 30-second + // consumer-stall watchdog (`grpc/openengine/servicer.py`, "response + // consumer stalled"), which fires when *this* sidecar stopped draining + // the stream -- the engine has capacity. Marking it migratable would + // re-dispatch to a second worker and stall there too, spending two + // workers' GPU time and recording local backpressure as worker + // capacity. Revisit if a server starts emitting it for real admission + // pressure. + pb::ErrorCode::Overloaded => client::engine_error(message), + pb::ErrorCode::Cancelled => client::cancelled(message), + // A decode request reached a worker whose engine is not in that role, + // so every request to it fails the same way. That is a deployment + // mistake, and an opaque 500 gives the operator nothing to search for. + pb::ErrorCode::RoleMismatch => client::invalid_argument(format!( + "{message} (the engine rejected this request's disaggregation role: check that the \ + sidecar's --disaggregation-mode matches how its engine was started)" + )), + // The handoff named a context worker this engine could not reach or + // whose session is gone. Migration is linked above the prefill router + // (`entrypoint/input/common.rs:501-505`), so a retry would re-run + // prefill and get a fresh handoff rather than replay this one -- this + // is recoverable in principle. Kept non-migratable for now because the + // common cause is a transceiver both engines lack, where retrying + // across every worker buries the one error that names the fix. + pb::ErrorCode::KvSessionNotFound | pb::ErrorCode::KvTransferFailed => { + client::engine_error(format!( + "{message} (the prefill handoff could not be resolved: check that both engines \ + were started with a cache transceiver and can reach each other)" + )) + } + _ => client::engine_error(message), + } +} diff --git a/lib/sidecar/trtllm/src/disagg.rs b/lib/sidecar/trtllm/src/disagg.rs new file mode 100644 index 000000000000..5350f1d2755a --- /dev/null +++ b/lib/sidecar/trtllm/src/disagg.rs @@ -0,0 +1,158 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Disaggregated prefill/decode handoff between Dynamo and OpenEngine. +//! +//! OpenEngine has no request-type field. A context (prefill) request is marked +//! by `extra.request_type = "context_only"`, and the server answers it with a +//! terminal `PrefillReady` event carrying a [`pb::KvSessionRef`]. A generation +//! (decode) request replays that same session in `kv.session`, which the server +//! decodes back into TensorRT-LLM's context handoff. +//! +//! Dynamo carries the handoff as opaque JSON (`PrefillResult.disaggregated_params`), +//! so this module is the codec between the two. The JSON mirrors `KvSessionRef` +//! field-for-field; it is written by the prefill worker and read by the decode +//! worker, and never interpreted in between. + +use dynamo_backend_common::DynamoError; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use dynamo_sidecar_common::{json_to_struct, struct_to_json}; + +use crate::client; +use crate::proto as pb; + +/// `extra` key the OpenEngine servicer reads to select the disaggregation phase. +pub(crate) const REQUEST_TYPE_KEY: &str = "request_type"; +/// `extra.request_type` value marking a prefill-only request. +pub(crate) const CONTEXT_ONLY: &str = "context_only"; + +const ATTRIBUTES: &str = "prefill handoff attributes"; + +/// The handoff's JSON shape, mirroring [`pb::KvSessionRef`] field for field. +/// +/// Every field is required, `attributes` most of all: the server reads the +/// session's location and rank out of it (`opaque_state`, `ctx_info_endpoint`, +/// `ctx_dp_rank`), treats `endpoints` only as a fallback, and never reads +/// `transfer_backend`. A handoff that lost `attributes` in transit still passes +/// the server's own guard, which accepts an endpoint in place of an +/// `opaque_state`, and then resumes a session with no opaque state and a rank +/// defaulted to 0. Requiring it here turns that into a named failure. +/// +/// Unknown fields are accepted on purpose. The required fields already reject +/// another engine's handoff, so rejecting unknown ones would only reject a +/// *newer* peer's -- during a rolling upgrade, every new-prefill/old-decode +/// request, as a non-migratable 400. +#[derive(Serialize, Deserialize)] +struct Handoff { + session_id: String, + transfer_backend: String, + endpoints: Vec, + dp_rank: u32, + attributes: Value, +} + +// No `deny_unknown_fields` here either: an unknown field inside `endpoints[]` +// is the same rolling-upgrade case the module doc describes for the handoff +// itself, and rejecting it would fail a newer peer's request as a +// non-migratable 400 one level down from where that rule is stated. +#[derive(Serialize, Deserialize)] +struct Endpoint { + host: String, + port: u32, + protocol: String, +} + +/// Encodes the prefill worker's `KvSessionRef` as the opaque JSON Dynamo +/// forwards to the decode worker. +pub(crate) fn session_to_json(session: pb::KvSessionRef) -> Result { + let pb::KvSessionRef { + session_id, + transfer_backend, + endpoints, + dp_rank, + attributes_struct, + } = session; + + if session_id.is_empty() { + return Err(client::protocol_error( + "prefill_ready carried no kv_session.session_id", + )); + } + // The decode leg cannot locate the session without these, so refuse to emit + // a handoff that would fail there instead of here. + let Some(attributes_struct) = attributes_struct else { + return Err(client::protocol_error( + "prefill_ready carried no kv_session.attributes_struct", + )); + }; + + let handoff = Handoff { + session_id, + transfer_backend, + endpoints: endpoints + .into_iter() + .map(|endpoint| Endpoint { + host: endpoint.host, + port: endpoint.port, + protocol: endpoint.protocol, + }) + .collect(), + dp_rank, + attributes: struct_to_json(attributes_struct, "TensorRT-LLM", ATTRIBUTES)?, + }; + serde_json::to_value(handoff).map_err(|error| { + client::protocol_error(format!("prefill handoff could not be encoded: {error}")) + }) +} + +/// Decodes the handoff JSON produced by [`session_to_json`] back into the +/// `KvSessionRef` the decode request replays. +pub(crate) fn session_from_json(value: &Value) -> Result { + // Deserialize from the borrowed tree: `from_value` would deep-clone the whole + // handoff, including `opaque_state` and `first_gen_log_probs`, only to drop + // the copy again. + let handoff = Handoff::deserialize(value).map_err(|error| { + client::invalid_argument(format!( + "decode request prefill_result.disaggregated_params is not a TensorRT-LLM handoff: \ + {error}" + )) + })?; + if handoff.session_id.is_empty() { + return Err(client::invalid_argument( + "decode request prefill handoff has an empty session_id", + )); + } + Ok(pb::KvSessionRef { + session_id: handoff.session_id, + transfer_backend: handoff.transfer_backend, + endpoints: handoff + .endpoints + .into_iter() + .map(|endpoint| pb::KvEndpoint { + host: endpoint.host, + port: endpoint.port, + protocol: endpoint.protocol, + }) + .collect(), + dp_rank: handoff.dp_rank, + attributes_struct: Some(json_to_struct(handoff.attributes, ATTRIBUTES)?), + }) +} + +/// `extra` payload marking a request as prefill-only. +pub(crate) fn context_only_extra() -> prost_types::Struct { + prost_types::Struct { + fields: [( + REQUEST_TYPE_KEY.to_string(), + prost_types::Value { + kind: Some(prost_types::value::Kind::StringValue( + CONTEXT_ONLY.to_string(), + )), + }, + )] + .into_iter() + .collect(), + } +} diff --git a/lib/sidecar/trtllm/src/engine.rs b/lib/sidecar/trtllm/src/engine.rs index 0daf9983cb54..be7ccb0dd15e 100644 --- a/lib/sidecar/trtllm/src/engine.rs +++ b/lib/sidecar/trtllm/src/engine.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Dynamo backend for TensorRT-LLM's native `trtllm.TrtllmService` gRPC server. +//! Dynamo backend for TensorRT-LLM's OpenEngine (`openengine.v1`) gRPC server. use std::sync::Arc; @@ -10,13 +10,15 @@ use dynamo_backend_common::{ AsyncEngineContext, DisaggregationMode, DynamoError, EngineConfig, GenerateContext, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, PreprocessedRequest, WorkerConfig, usage, }; -use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportConfig, SidecarStartupError}; +use dynamo_sidecar_common::{ + GrpcEndpoint, GrpcTransportConfig, SidecarStartupError, startup_deadline, +}; use futures::stream::BoxStream; use tokio::sync::OnceCell; use tokio_util::sync::CancellationToken; use crate::args::Args; -use crate::client::{self, TrtllmClient}; +use crate::client::{self, ModelLimits, TrtllmClient}; use crate::convert::{ResponseState, build_generate_request}; use crate::model::ConfiguredModel; @@ -32,11 +34,14 @@ pub struct TrtllmSidecarEngine { endpoint: GrpcEndpoint, transport: GrpcTransportConfig, model: ConfiguredModel, + /// Disaggregation role this worker plays. Selects the `context_only` / + /// `kv.session` divergence in `convert`. + mode: DisaggregationMode, client: OnceCell, - /// Resolved model context length (`--context-length`, else `GetModelInfo`), - /// cached at `start` so `generate` can derive a default `max_tokens` for - /// requests that omit one. - context_length: OnceCell, + /// Engine limits resolved at `start` from `--context-length` and + /// `Control.GetModelInfo`, so `generate` can derive a default `max_tokens` + /// for requests that omit one. + limits: OnceCell, cancel: CancellationToken, } @@ -45,13 +50,15 @@ impl TrtllmSidecarEngine { endpoint: GrpcEndpoint, transport: GrpcTransportConfig, model: ConfiguredModel, + mode: DisaggregationMode, ) -> Self { Self { endpoint, transport, model, + mode, client: OnceCell::new(), - context_length: OnceCell::new(), + limits: OnceCell::new(), cancel: CancellationToken::new(), } } @@ -77,15 +84,10 @@ impl TrtllmSidecarEngine { if args.model_path.trim().is_empty() { return Err(client::invalid_argument("model-path must not be empty")); } - if args.context_length == Some(0) { + let mode = args.sidecar.common.disaggregation_mode; + if mode.is_encode() { return Err(client::invalid_argument( - "context-length must be greater than zero", - )); - } - if args.sidecar.common.disaggregation_mode != DisaggregationMode::Aggregated { - return Err(client::invalid_argument( - "the TensorRT-LLM sidecar supports aggregated serving only; the Generate \ - response contract carries no disaggregation handoff", + "encode mode is not supported by the TensorRT-LLM sidecar", )); } if args.sidecar.common.route_to_encoder { @@ -98,12 +100,21 @@ impl TrtllmSidecarEngine { let transport = args.sidecar.grpc.config(); let model = ConfiguredModel { source: args.model_path, + // Absent unless `--context-length` supplied one; `start` falls back + // to the server's `Control.GetModelInfo` report. context_length: args.context_length, }; - let engine = Self::new(endpoint, transport, model.clone()); + let engine = Self::new(endpoint, transport, model.clone(), mode); let config = WorkerConfig { namespace: args.sidecar.common.namespace, - component: args.sidecar.common.component, + // Every disaggregated role registers under its own component so + // the frontend can target each separately; only an aggregated + // worker uses the operator-configured one. + component: if mode == DisaggregationMode::Aggregated { + args.sidecar.common.component + } else { + mode.discovery_component().to_string() + }, endpoint: args.sidecar.common.endpoint, endpoint_types: args.sidecar.common.endpoint_types, custom_jinja_template: args.sidecar.common.custom_jinja_template, @@ -116,7 +127,7 @@ impl TrtllmSidecarEngine { .common .exclude_tools_when_tool_choice_none, enable_kv_routing: false, - disaggregation_mode: DisaggregationMode::Aggregated, + disaggregation_mode: mode, route_to_encoder: false, enable_rl: args.sidecar.common.enable_rl, ..Default::default() @@ -136,45 +147,87 @@ impl LLMEngine for TrtllmSidecarEngine { connections = self.transport.connections.get(), "connecting to TensorRT-LLM gRPC" ); + // One deadline for the whole startup path. `GrpcChannelPool::connect` + // derives its own from the same duration, so taking this before + // connecting is what stops the two stages spending a full budget each. + let deadline = startup_deadline(self.transport.startup_deadline)?; let client = TrtllmClient::connect(&self.endpoint, self.transport).await?; let connection_count = client.connection_count(); - // `GetModelInfo` reports the engine's `--max_seq_len`, which is unset by - // default; `client::model_info` discards the value TensorRT-LLM - // substitutes for it. A configured `--context-length` wins over what - // survives that check. + // `--context-length` wins over what the engine reports, and is the + // only source when the engine reports nothing usable. The resolved + // value backs both the registered window and the default-`max_tokens` + // path in `convert::max_tokens`. let mut model = self.model.clone(); - let reported = match client.model_info().await { - Ok(reported) => reported, - Err(error) => { - match model.context_length { - Some(configured) => tracing::warn!( - %error, + let limits = match model.context_length { + // Configured: the engine is consulted once, to cross-check the + // value and to learn its output cap. It may not answer at all -- + // an older server has no Control service -- and that must not stop + // a worker whose window the operator already supplied. + // + // Neither branch waits for the engine to be *ready*, only to answer. + // Readiness belongs to the engine's health service, which is what + // the Kubernetes probes use and what DEP #14897 standardises on + // (`grpc.health.v1` Check/Watch, with registration driven by + // `Watch`). Do not grow a second readiness protocol here. + Some(configured) => { + let reported = match client.model_limits(&model.source).await { + Ok(reported) => reported, + Err(error) => { + // Tolerated because some OpenEngine servicer builds ship + // without Control, which is what --context-length exists + // for. UNIMPLEMENTED also fits an engine left on the + // default --grpc-protocol smg, where Generate will fail + // too -- name both so the log points at the real fix. + tracing::warn!( + %error, + configured_context_length = configured, + "Control.GetModelInfo failed; using the configured \ + --context-length. If this is UNIMPLEMENTED, the engine either \ + ships no Control service or was started without \ + --grpc-protocol openengine" + ); + ModelLimits::default() + } + }; + if let Some(engine_context_length) = reported.context_length + && engine_context_length != configured + { + tracing::warn!( configured_context_length = configured, - "GetModelInfo failed; using the configured --context-length" - ), - None => tracing::warn!( - %error, - "GetModelInfo failed and no --context-length was configured; \ - no context length is available" - ), + engine_context_length, + "--context-length disagrees with the context length TensorRT-LLM \ + reported; using the configured --context-length" + ); } - None + ModelLimits { + context_length: Some(configured), + ..reported + } + } + // Nothing configured: the engine is the only source. It binds its + // gRPC port only after the model has loaded, but the sidecar and the + // engine start independently, so wait for it either way. + None => { + // Failing here beats registering a worker that advertises + // capacity and then rejects, with a non-migratable 4xx, every + // request that omits `max_tokens` -- which is most chat + // traffic. The operator gets one startup error naming the fix + // instead of a worker that looks healthy and serves half. + client + .wait_for_model_limits(&model.source, deadline, self.transport.retry_interval) + .await + .map_err(|error| { + client::invalid_argument(format!( + "no context length is available: {error}. TensorRT-LLM reports one \ + only when it was started with --max_seq_len; otherwise supply \ + --context-length." + )) + })? } }; - match (model.context_length, reported) { - (Some(configured), Some(reported)) if configured != reported => tracing::warn!( - configured_context_length = configured, - engine_context_length = reported, - "--context-length disagrees with the context length TensorRT-LLM reported; \ - using the configured --context-length" - ), - (None, Some(reported)) => model.context_length = Some(reported), - _ => {} - } - if let Some(context_length) = model.context_length { - let _ = self.context_length.set(context_length); - } + model.context_length = limits.context_length; + let _ = self.limits.set(limits); self.client .set(client) @@ -183,7 +236,8 @@ impl LLMEngine for TrtllmSidecarEngine { endpoint = %self.endpoint, connections = connection_count, model = %model.source, - context_length = ?model.context_length, + context_length = ?limits.context_length, + max_output_tokens = ?limits.max_output_tokens, "TensorRT-LLM gRPC is ready" ); Ok(model.engine_config()) @@ -199,15 +253,42 @@ impl LLMEngine for TrtllmSidecarEngine { .get() .ok_or_else(|| client::engine_shutdown("TensorRT-LLM sidecar is not started"))?; let request_id = ctx.id().to_string(); - let proto_request = - build_generate_request(&request, &request_id, self.context_length.get().copied())?; - let mut state = ResponseState::new(&request); - let cancel = self.cancel.clone(); + let proto_request = build_generate_request( + &request, + &request_id, + &self.model.source, + self.limits.get().copied(), + self.mode, + )?; + let mut state = ResponseState::new(&request, self.mode); + // A decode request that took a handoff has KV transferred into it, and + // the transceiver releases those blocks when the engine finishes the + // request -- not when the client goes away. Dropping the stream on + // cancellation would strand the prefill worker's blocks. Defer only + // until the first token proves the transfer landed; deferring past + // that would let a cancelled request generate its whole budget with no + // consumer. A decode-mode request without a handoff ran locally and + // has nothing to strand. + let needs_cancellation_deferral = self.mode.is_decode() && request.prefill_result.is_some(); + let stopped_ctx = ctx.inner_arc(); + // Hoisted: `stopped()` is an async-trait method, so re-creating it per + // streamed chunk costs a boxed future and a waker registration on every + // token. + let mut request_cancellation = Box::pin(async move { stopped_ctx.stopped().await }); + let shutdown = self.cancel.clone(); + let mut shutdown_cancellation = Box::pin(async move { shutdown.cancelled().await }); + // The same deferral applies here, not just to the streaming loop below. + // `generate` sends the request and then awaits response headers, so + // losing this race can drop a request the engine has already accepted + // and begun pulling KV for; and an already-stopped context would skip + // the dispatch entirely, leaving the prefill worker's blocks with no + // decode leg to claim them. Both strand exactly what the deferral + // exists to protect. Shutdown still wins -- the process is going away. let stream = tokio::select! { biased; - _ = ctx.stopped() => None, - _ = cancel.cancelled() => None, + _ = &mut request_cancellation, if !needs_cancellation_deferral => None, + _ = &mut shutdown_cancellation => None, result = client.generate(proto_request) => Some(result?), }; let Some(mut stream) = stream else { @@ -216,14 +297,15 @@ impl LLMEngine for TrtllmSidecarEngine { }; Ok(Box::pin(async_stream::stream! { + let mut is_transfer_settled = false; loop { tokio::select! { biased; - _ = ctx.stopped() => { + _ = &mut request_cancellation, if !needs_cancellation_deferral || is_transfer_settled => { yield Ok(cancelled(&state)); break; } - _ = cancel.cancelled() => { + _ = &mut shutdown_cancellation => { yield Ok(cancelled(&state)); break; } @@ -231,6 +313,7 @@ impl LLMEngine for TrtllmSidecarEngine { match message { Ok(Some(response)) => match state.convert(response) { Ok(Some(output)) => { + is_transfer_settled |= !output.token_ids.is_empty(); let terminal = output.finish_reason.is_some(); yield Ok(output); if terminal { @@ -265,7 +348,14 @@ impl LLMEngine for TrtllmSidecarEngine { return; }; if let Err(error) = client.abort(ctx.id().to_string()).await { - tracing::debug!(request_id = ctx.id(), %error, "TensorRT-LLM Abort RPC failed"); + // Escaped: the message embeds the engine's gRPC status text, and + // this site logs at the default level, so raw newlines from the + // peer would let it forge what look like separate log records. + tracing::warn!( + request_id = ctx.id(), + error = %error.to_string().escape_debug(), + "TensorRT-LLM Control.Abort failed" + ); } } diff --git a/lib/sidecar/trtllm/src/lib.rs b/lib/sidecar/trtllm/src/lib.rs index 2859d60cea0f..47b3024e9e44 100644 --- a/lib/sidecar/trtllm/src/lib.rs +++ b/lib/sidecar/trtllm/src/lib.rs @@ -1,11 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Dynamo sidecar for TensorRT-LLM's native `TrtllmService` gRPC API. +//! Dynamo sidecar for TensorRT-LLM's OpenEngine (`openengine.v1`) gRPC API. mod args; mod client; mod convert; +mod disagg; mod engine; mod model; mod proto; diff --git a/lib/sidecar/trtllm/src/model.rs b/lib/sidecar/trtllm/src/model.rs index d3e50129a9a0..6989b63658d2 100644 --- a/lib/sidecar/trtllm/src/model.rs +++ b/lib/sidecar/trtllm/src/model.rs @@ -21,7 +21,7 @@ impl ConfiguredModel { let mut runtime_data = HashMap::new(); runtime_data.insert( "grpc_service".to_string(), - serde_json::Value::String("trtllm.TrtllmService".to_string()), + serde_json::Value::String("openengine.v1.Inference".to_string()), ); EngineConfig { diff --git a/lib/sidecar/trtllm/src/proto.rs b/lib/sidecar/trtllm/src/proto.rs index 87b71b4e3f2e..5300cb6ade84 100644 --- a/lib/sidecar/trtllm/src/proto.rs +++ b/lib/sidecar/trtllm/src/proto.rs @@ -1,12 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Generated protobuf / gRPC types for the `trtllm` package. +//! Re-exports the generated OpenEngine (`openengine.v1`) gRPC types. //! -//! Codegen is driven by `build.rs` from the vendored, checksum-tracked -//! TensorRT-LLM `trtllm_service.proto`. +//! Generated by `build.rs` from the vendored contract in `proto/`; see +//! `proto/README.md` for the upstream revision. #![allow(clippy::all)] #![allow(missing_docs)] -tonic::include_proto!("trtllm"); +tonic::include_proto!("openengine.v1"); diff --git a/lib/sidecar/trtllm/src/tests.rs b/lib/sidecar/trtllm/src/tests.rs index 50dc426487e8..411aa819176e 100644 --- a/lib/sidecar/trtllm/src/tests.rs +++ b/lib/sidecar/trtllm/src/tests.rs @@ -6,11 +6,12 @@ use std::net::SocketAddr; use std::num::NonZeroUsize; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::Duration; use dynamo_backend_common::{ - FinishReason, GenerateContext, LLMEngine, OutputOptions, PreprocessedRequest, SamplingOptions, - StopConditions, StopReason, + DisaggregationMode, ErrorType, FinishReason, GenerateContext, LLMEngine, OutputOptions, + PreprocessedRequest, SamplingOptions, StopConditions, StopReason, }; use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportConfig}; use futures::{Stream, StreamExt}; @@ -20,14 +21,27 @@ use tokio::sync::{Mutex, oneshot}; use tokio_stream::wrappers::TcpListenerStream; use tonic::{Request, Response, Status}; -use crate::client::TrtllmClient; -use crate::convert::{ResponseState, build_generate_request}; +use crate::client::{ModelLimits, TrtllmClient}; +use crate::convert::{ResponseState, build_generate_request, engine_error}; use crate::engine::TrtllmSidecarEngine; use crate::model::ConfiguredModel; use crate::proto as pb; +/// Most tests exercise aggregated serving; the disaggregation tests name their +/// mode explicitly. +const AGG: DisaggregationMode = DisaggregationMode::Aggregated; +const PREFILL: DisaggregationMode = DisaggregationMode::Prefill; + +// The tests themselves live in `tests/`, grouped by the surface they cover; +// this file holds only the fakes and fixtures they share. +mod convert_request; +mod convert_response; +mod disagg; +mod e2e; +mod engine; + // --------------------------------------------------------------------------- -// Fake TensorRT-LLM gRPC service +// Fake TensorRT-LLM OpenEngine services // --------------------------------------------------------------------------- #[derive(Clone, Default)] @@ -37,20 +51,48 @@ struct FakeTrtllm { peers: Arc>>, reject: Arc, hang: Arc, - /// `(max_seq_len, max_input_len)` for `GetModelInfo`; unset reports 4096 - /// with `max_input_len` absent. - model_info: Arc>>, + /// Simulates a server that has not yet accepted the request, so the RPC + /// itself is still in flight. + hang_before_stream: Arc, + /// Simulates a server whose Control service is not implemented. + no_control: Arc, + /// Simulates a server that answers GetModelInfo without a context length. + empty_model_info: Arc, + /// Answers `UNAVAILABLE`, the shape of an engine that is not serving yet. + unavailable_model_info: Arc, + model_info_calls: Arc, +} + +fn prompt_len(request: &pb::GenerateRequest) -> u32 { + match request.input.as_ref() { + Some(pb::generate_request::Input::TokenIds(tokens)) => tokens.ids.len() as u32, + _ => 0, + } } -impl FakeTrtllm { - async fn reporting(self, max_seq_len: i32, max_input_len: i32) -> Self { - *self.model_info.lock().await = Some((max_seq_len, max_input_len)); - self - } +/// Mirrors the OpenEngine servicer: a request is prefill-only when its `extra` +/// Struct carries `request_type = "context_only"`. +fn is_context_only(request: &pb::GenerateRequest) -> bool { + request + .extra + .as_ref() + .and_then(|extra| extra.fields.get("request_type")) + .and_then(|value| value.kind.as_ref()) + .is_some_and(|kind| { + matches!(kind, prost_types::value::Kind::StringValue(value) if value == "context_only") + }) +} + +fn wants_logprobs(request: &pb::GenerateRequest) -> bool { + request + .response + .as_ref() + .and_then(|response| response.return_output_logprobs) + .unwrap_or(false) } #[tonic::async_trait] -impl pb::trtllm_service_server::TrtllmService for FakeTrtllm { +impl pb::inference_server::Inference for FakeTrtllm { type GenerateStream = Pin> + Send>>; async fn generate( @@ -67,41 +109,62 @@ impl pb::trtllm_service_server::TrtllmService for FakeTrtllm { } let request_id = request.request_id.clone(); - let prompt_tokens = request - .tokenized - .as_ref() - .map(|t| t.input_token_ids.len() as u32) - .unwrap_or(0); - let wants_logprobs = request - .output_config - .as_ref() - .and_then(|o| o.logprobs) - .is_some(); + let prompt_tokens = prompt_len(&request); + let wants_logprobs = wants_logprobs(&request); + let context_only = is_context_only(&request); let hang = self.hang.load(Ordering::SeqCst); + if self.hang_before_stream.load(Ordering::SeqCst) { + std::future::pending::<()>().await; + } let stream = async_stream::try_stream! { - let logprobs = if wants_logprobs { - vec![pb::TokenLogprob { + let tokens = if wants_logprobs { + vec![pb::TokenInfo { token_id: 42, - logprob: -0.25, - top_logprobs: vec![pb::TopLogprob { token_id: 43, logprob: -0.5 }], + token: String::new(), + logprob: Some(-0.25), + rank: Some(1), + candidates: vec![pb::LogProb { + token_id: 43, + logprob: -0.5, + token: String::new(), + rank: Some(2), + }], }] } else { - Vec::new() + vec![pb::TokenInfo { + token_id: 42, + token: String::new(), + logprob: None, + rank: None, + candidates: Vec::new(), + }] }; yield pb::GenerateResponse { request_id: request_id.clone(), - response: Some(pb::generate_response::Response::Chunk(pb::GenerateStreamChunk { - token_ids: vec![42], - sequence_index: 0, - prompt_tokens, - completion_tokens: 1, - cached_tokens: 0, - logprobs, + event: Some(pb::generate_response::Event::Token(pb::TokenOutput { + output_index: Some(0), + tokens, + text: String::new(), })), + usage: None, }; + // A context_only request terminates on PrefillReady: the servicer + // suppresses the `finished` event because the engine reports the + // sequence as unfinished. + if context_only { + yield pb::GenerateResponse { + request_id, + event: Some(pb::generate_response::Event::PrefillReady(pb::PrefillReady { + kv_session: Some(fake_session()), + })), + usage: None, + }; + return; + } + if hang { loop { tokio::time::sleep(std::time::Duration::from_millis(10)).await; @@ -110,34 +173,90 @@ impl pb::trtllm_service_server::TrtllmService for FakeTrtllm { yield pb::GenerateResponse { request_id, - response: Some(pb::generate_response::Response::Complete(pb::GenerateComplete { - output_token_ids: vec![42], - sequence_index: 0, - finish_reason: "stop".to_string(), - matched_stop: Some(pb::generate_complete::MatchedStop::MatchedTokenId(2)), + event: Some(pb::generate_response::Event::Finished(pb::GenerationFinished { + output_index: Some(0), + reason: pb::FinishReason::Stop as i32, + message: String::new(), + stop_match: Some(pb::StopMatch { + r#match: Some(pb::stop_match::Match::StopTokenId(2)), + }), + })), + usage: Some(pb::Usage { prompt_tokens, completion_tokens: 1, - cached_tokens: 0, - ..Default::default() - })), + total_tokens: prompt_tokens + 1, + cached_prompt_tokens: None, + reasoning_tokens: None, + }), }; }; Ok(Response::new(Box::pin(stream))) } +} - async fn embed( - &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("embed is not used")) +/// The handoff a context worker returns, shaped like TensorRT-LLM's: the +/// session id is the context request id and the engine-specific state rides in +/// `attributes_struct`. +fn fake_session() -> pb::KvSessionRef { + pb::KvSessionRef { + session_id: "12345".to_string(), + transfer_backend: "NIXL".to_string(), + endpoints: vec![pb::KvEndpoint { + host: "10.0.0.7".to_string(), + port: 5601, + protocol: "grpc".to_string(), + }], + dp_rank: 0, + attributes_struct: Some(prost_types::Struct { + fields: [ + ( + "opaque_state".to_string(), + prost_types::Value { + kind: Some(prost_types::value::Kind::StringValue( + "c3RhdGU=".to_string(), + )), + }, + ), + ( + "first_gen_tokens".to_string(), + prost_types::Value { + kind: Some(prost_types::value::Kind::ListValue( + prost_types::ListValue { + values: vec![prost_types::Value { + kind: Some(prost_types::value::Kind::NumberValue(42.0)), + }], + }, + )), + }, + ), + ] + .into_iter() + .collect(), + }), } +} - async fn health_check( +#[tonic::async_trait] +impl pb::control_server::Control for FakeTrtllm { + async fn get_model_info( &self, - _request: Request, - ) -> Result, Status> { - Ok(Response::new(pb::HealthCheckResponse { - status: "OK".to_string(), + _request: Request, + ) -> Result, Status> { + self.model_info_calls.fetch_add(1, Ordering::SeqCst); + if self.no_control.load(Ordering::SeqCst) { + return Err(Status::unimplemented("Control is not implemented")); + } + if self.unavailable_model_info.load(Ordering::SeqCst) { + return Err(Status::unavailable("not serving yet")); + } + Ok(Response::new(pb::ModelInfo { + model_id: "fake-model".to_string(), + max_context_length: if self.empty_model_info.load(Ordering::SeqCst) { + None + } else { + Some(4096) + }, + ..Default::default() })) } @@ -145,40 +264,78 @@ impl pb::trtllm_service_server::TrtllmService for FakeTrtllm { &self, request: Request, ) -> Result, Status> { - let request_id = request.into_inner().request_id; + let request_id = match request.into_inner().target { + Some(pb::abort_request::Target::RequestId(id)) => id, + other => { + return Err(Status::invalid_argument(format!( + "unexpected abort target {other:?}" + ))); + } + }; self.aborts.lock().await.push(request_id.clone()); Ok(Response::new(pb::AbortResponse { - success: true, + status: pb::AbortStatus::Aborted as i32, message: format!("aborted {request_id}"), })) } - async fn get_model_info( + async fn get_server_info( &self, - _request: Request, - ) -> Result, Status> { - let (max_seq_len, max_input_len) = self.model_info.lock().await.unwrap_or((4096, 0)); - Ok(Response::new(pb::GetModelInfoResponse { - model_id: "fake-model".to_string(), - max_seq_len, - max_input_len, - vocab_size: 32000, - ..Default::default() - })) + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetServerInfo is not used")) } - async fn get_server_info( + async fn get_load( &self, - _request: Request, - ) -> Result, Status> { - Ok(Response::new(pb::GetServerInfoResponse { - version: "1.3.0rc21".to_string(), - backend: "pytorch".to_string(), - tensor_parallel_size: 1, - pipeline_parallel_size: 1, - context_parallel_size: 1, - world_size: 1, - })) + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetLoad is not used")) + } + + async fn health( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Health is not used")) + } + + async fn load_lora( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("LoadLora is not used")) + } + + async fn unload_lora( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("UnloadLora is not used")) + } + + async fn list_loras( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("ListLoras is not used")) + } + + async fn get_kv_event_sources( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetKvEventSources is not used")) + } + + type SubscribeKvEventsStream = + Pin> + Send>>; + + async fn subscribe_kv_events( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("SubscribeKvEvents is not used")) } } @@ -196,9 +353,10 @@ impl FakeServer { let server_service = service.clone(); tokio::spawn(async move { tonic::transport::Server::builder() - .add_service(pb::trtllm_service_server::TrtllmServiceServer::new( - server_service, + .add_service(pb::inference_server::InferenceServer::new( + server_service.clone(), )) + .add_service(pb::control_server::ControlServer::new(server_service)) .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { let _ = shutdown_rx.await; }) @@ -267,22 +425,54 @@ fn transport(connections: usize) -> GrpcTransportConfig { } } +/// A transport that gives up on startup quickly, for the paths that retry until +/// the deadline rather than failing on the first answer. +fn impatient_transport() -> GrpcTransportConfig { + GrpcTransportConfig { + retry_interval: Duration::from_millis(10), + startup_deadline: Duration::from_millis(200), + ..transport(1) + } +} + +/// Engine limits as `Control.GetModelInfo` would report them, with no +/// separate output cap. +fn limits(context_length: u32) -> Option { + Some(ModelLimits { + context_length: Some(context_length), + max_output_tokens: None, + }) +} + fn engine(endpoint: &str, connections: usize) -> TrtllmSidecarEngine { - engine_with_context_length(endpoint, connections, None) + engine_in_mode(endpoint, connections, AGG) } -fn engine_with_context_length( +fn engine_in_mode( endpoint: &str, connections: usize, + mode: DisaggregationMode, +) -> TrtllmSidecarEngine { + engine_with(endpoint, transport(connections), None, mode) +} + +/// The one place a test engine is built. Everything a test varies -- the +/// transport, whether `--context-length` was supplied, the disaggregation role +/// -- is a parameter here. +fn engine_with( + endpoint: &str, + transport: GrpcTransportConfig, context_length: Option, + mode: DisaggregationMode, ) -> TrtllmSidecarEngine { TrtllmSidecarEngine::new( GrpcEndpoint::parse(endpoint, "--grpc-endpoint").expect("valid test endpoint"), - transport(connections), + transport, ConfiguredModel { source: "model-source".to_string(), context_length, }, + mode, ) } @@ -305,523 +495,32 @@ async fn collect( fn assert_rejected(mutate: impl FnOnce(&mut PreprocessedRequest), expect: &str) { let mut req = request(); mutate(&mut req); - let error = build_generate_request(&req, "req", None).expect_err("request must be rejected"); + let error = build_generate_request(&req, "req", "model", None, AGG) + .expect_err("request must be rejected"); assert!( error.to_string().contains(expect), "error {error:?} should mention {expect:?}" ); } -// --------------------------------------------------------------------------- -// Request-building unit tests -// --------------------------------------------------------------------------- - -#[test] -fn request_maps_sampling_stop_and_output_fields() { - let proto = build_generate_request(&request(), "req-1", None).expect("build request"); - assert_eq!(proto.request_id, "req-1"); - assert_eq!( - proto.tokenized.as_ref().unwrap().input_token_ids, - [11, 22, 33] - ); - assert_eq!(proto.max_tokens, 16); - assert!(proto.streaming); - assert!(proto.ignore_eos); - // Stop-token inclusion is never enabled from `include_stop_str_in_output` - // (different axis; would leak hidden stop tokens). - assert!(!proto.include_stop_token_in_output); - assert_eq!(proto.stop, ["done"]); - assert_eq!(proto.stop_token_ids, [2]); - - let sampling = proto.sampling_config.as_ref().unwrap(); - assert_eq!(sampling.top_k, Some(4)); - assert_eq!(sampling.top_p, Some(0.9)); - assert_eq!(sampling.min_p, Some(0.1)); - assert_eq!(sampling.temperature, Some(0.2)); - assert_eq!(sampling.seed, Some(123)); - assert_eq!(sampling.repetition_penalty, Some(1.1)); - assert_eq!(sampling.min_tokens, Some(1)); - - let output = proto.output_config.as_ref().unwrap(); - assert_eq!(output.logprobs, Some(1)); - assert!(output.exclude_input_from_output); - - let guided = proto.guided_decoding.as_ref().unwrap(); - assert_eq!( - guided.guide_type, - pb::guided_decoding_params::GuideType::JsonSchema as i32 - ); - assert!(guided.guide.contains("object")); -} - -#[test] -fn omitted_max_tokens_without_context_length_is_rejected() { - let mut req = request(); - req.stop_conditions.max_tokens = None; - let error = build_generate_request(&req, "req", None).expect_err("must require max_tokens"); - assert!(error.to_string().contains("max_tokens")); -} - -#[test] -fn omitted_max_tokens_defaults_to_remaining_context() { - let mut req = request(); - req.stop_conditions.max_tokens = None; - // request() carries three prompt tokens ([11, 22, 33]); the default fills the - // remaining context: max(1, context_length - prompt_len). - let proto = build_generate_request(&req, "req", Some(100)).expect("build with fallback"); - assert_eq!(proto.max_tokens, 97); -} - -#[test] -fn omitted_max_tokens_default_is_floored_at_one() { - let mut req = request(); - req.stop_conditions.max_tokens = None; - // Prompt already fills (or exceeds) the context: default must not underflow to 0. - let proto = build_generate_request(&req, "req", Some(2)).expect("build with fallback"); - assert_eq!(proto.max_tokens, 1); -} - -#[test] -fn top_k_all_tokens_is_left_unset() { - let mut req = request(); - req.sampling_options.top_k = Some(-1); - let proto = build_generate_request(&req, "req", None).expect("build"); - assert_eq!(proto.sampling_config.unwrap().top_k, None); -} - -#[test] -fn oversized_logprob_count_is_rejected() { - let mut req = request(); - req.output_options.logprobs = Some(i32::MAX as u32 + 1); - let error = - build_generate_request(&req, "req", None).expect_err("oversized logprobs must fail"); - assert!(error.to_string().contains("must fit in i32")); -} - -#[test] -fn unsupported_request_controls_are_rejected() { - // Controls the native gRPC contract can neither forward nor faithfully honor: - // reject rather than fail open. - assert_rejected( - |r| r.sampling_options.include_stop_str_in_output = Some(true), - "include_stop_str_in_output", - ); - assert_rejected( - |r| r.stop_conditions.max_thinking_tokens = Some(32), - "max_thinking_tokens", - ); - assert_rejected( - |r| { - r.routing = Some(dynamo_backend_common::engine::RoutingHints { - cache_namespace: Some("tenant-a".to_string()), - ..Default::default() - }) - }, - "cache namespace", - ); - assert_rejected( - |r| { - r.routing = Some(dynamo_backend_common::engine::RoutingHints { - priority: Some(5), - ..Default::default() - }) - }, - "priority", - ); - // A negative top_k other than -1/0 (the "all tokens" sentinels) is invalid, - // not a silent widening to "all tokens". - assert_rejected(|r| r.sampling_options.top_k = Some(-5), "top_k must be"); - assert_rejected( - |r| r.sampling_options.seed = Some(-1), - "seed must be non-negative", - ); -} - -#[test] -fn logprobs_zero_keeps_selected_without_alternatives() { - let mut req = request(); - req.output_options.logprobs = Some(0); - // The wire request must still ask TRT-LLM for one logprob so the selected - // token's value is computed. - let proto = build_generate_request(&req, "req", None).expect("build"); - assert_eq!(proto.output_config.unwrap().logprobs, Some(1)); - - let mut state = ResponseState::new(&req); - let chunk = pb::GenerateResponse { - request_id: "r".to_string(), - response: Some(pb::generate_response::Response::Chunk( - pb::GenerateStreamChunk { - token_ids: vec![7], - sequence_index: 0, - prompt_tokens: 3, - completion_tokens: 1, - cached_tokens: 0, - logprobs: vec![pb::TokenLogprob { - token_id: 7, - logprob: -0.1, - top_logprobs: vec![], - }], - }, - )), - }; - let delta = state.convert(chunk).expect("convert").expect("delta"); - assert_eq!(delta.log_probs.as_deref(), Some(&[f64::from(-0.1_f32)][..])); - // logprobs=0 surfaces the selected-token logprob but no top alternatives. - assert!(delta.top_logprobs.is_none()); -} - -// --------------------------------------------------------------------------- -// Response-conversion unit tests -// --------------------------------------------------------------------------- - -#[test] -fn chunk_then_complete_produces_delta_then_terminal_usage() { - let req = request(); - let mut state = ResponseState::new(&req); - - let chunk = pb::GenerateResponse { - request_id: "r".to_string(), - response: Some(pb::generate_response::Response::Chunk( - pb::GenerateStreamChunk { - token_ids: vec![7, 8], - sequence_index: 0, - prompt_tokens: 3, - completion_tokens: 2, - cached_tokens: 0, - logprobs: vec![ - pb::TokenLogprob { - token_id: 7, - logprob: -0.1, - top_logprobs: vec![], - }, - pb::TokenLogprob { - token_id: 8, - logprob: -0.2, - top_logprobs: vec![], - }, - ], - }, - )), - }; - let delta = state.convert(chunk).expect("convert chunk").expect("delta"); - assert_eq!(delta.token_ids, [7, 8]); - assert!(delta.finish_reason.is_none()); - let log_probs = delta.log_probs.as_ref().expect("log_probs"); - assert_eq!(log_probs.len(), 2); - assert!((log_probs[0] - f64::from(-0.1_f32)).abs() < 1e-9); - assert!((log_probs[1] - f64::from(-0.2_f32)).abs() < 1e-9); - assert_eq!(delta.top_logprobs.as_ref().unwrap().len(), 2); - - let complete = pb::GenerateResponse { - request_id: "r".to_string(), - response: Some(pb::generate_response::Response::Complete( - pb::GenerateComplete { - output_token_ids: vec![7, 8], - sequence_index: 0, - finish_reason: "length".to_string(), - prompt_tokens: 3, - completion_tokens: 2, - ..Default::default() - }, - )), - }; - let terminal = state - .convert(complete) - .expect("convert complete") - .expect("terminal"); - assert!(terminal.token_ids.is_empty()); - assert_eq!(terminal.finish_reason, Some(FinishReason::Length)); - let usage = terminal.completion_usage.as_ref().expect("usage"); - assert_eq!((usage.prompt_tokens, usage.completion_tokens), (3, 2)); -} - -#[test] -fn unsupported_sequence_index_is_rejected() { - let req = request(); - let mut state = ResponseState::new(&req); - let chunk = pb::GenerateResponse { +fn token_response(tokens: Vec) -> pb::GenerateResponse { + pb::GenerateResponse { request_id: "r".to_string(), - response: Some(pb::generate_response::Response::Chunk( - pb::GenerateStreamChunk { - token_ids: vec![1], - sequence_index: 1, - ..Default::default() - }, - )), - }; - assert!(state.convert(chunk).is_err()); -} - -#[test] -fn missing_response_payload_is_rejected() { - let req = request(); - let mut state = ResponseState::new(&req); - let empty = pb::GenerateResponse { - request_id: "r".to_string(), - response: None, - }; - assert!(state.convert(empty).is_err()); -} - -#[test] -fn unknown_finish_reason_is_rejected() { - let req = request(); - let mut state = ResponseState::new(&req); - let complete = pb::GenerateResponse { - request_id: "r".to_string(), - response: Some(pb::generate_response::Response::Complete( - pb::GenerateComplete { - finish_reason: "teleported".to_string(), - sequence_index: 0, - ..Default::default() - }, - )), - }; - assert!(state.convert(complete).is_err()); -} - -#[test] -fn terminal_token_count_mismatch_is_rejected() { - let mut req = request(); - req.output_options.logprobs = None; - let mut state = ResponseState::new(&req); - // Stream a single delta token... - let chunk = pb::GenerateResponse { - request_id: "r".to_string(), - response: Some(pb::generate_response::Response::Chunk( - pb::GenerateStreamChunk { - token_ids: vec![7], - sequence_index: 0, - prompt_tokens: 3, - completion_tokens: 1, - cached_tokens: 0, - logprobs: vec![], - }, - )), - }; - state.convert(chunk).expect("convert chunk"); - // ...but the terminal claims two cumulative output tokens. - let complete = pb::GenerateResponse { - request_id: "r".to_string(), - response: Some(pb::generate_response::Response::Complete( - pb::GenerateComplete { - output_token_ids: vec![7, 8], - sequence_index: 0, - finish_reason: "stop".to_string(), - prompt_tokens: 3, - completion_tokens: 2, - ..Default::default() - }, - )), - }; - assert!(state.convert(complete).is_err()); -} - -// --------------------------------------------------------------------------- -// Integration tests against the fake server -// --------------------------------------------------------------------------- - -#[tokio::test] -async fn aggregated_generation_streams_delta_then_terminal() { - let server = FakeServer::start(FakeTrtllm::default()).await; - let engine = engine(&server.endpoint, 2); - let config = engine.start(0).await.expect("start"); - assert_eq!(config.model, "model-source"); - // GetModelInfo reports max_seq_len 4096. - assert_eq!(config.llm.unwrap().context_length, Some(4096)); - - let outputs = collect(&engine, request()).await; - assert_eq!(outputs.len(), 2); - assert_eq!(outputs[0].token_ids, [42]); - assert!(outputs[0].finish_reason.is_none()); - assert_eq!(outputs[0].log_probs.as_deref(), Some(&[-0.25][..])); - - let terminal = &outputs[1]; - assert!(terminal.token_ids.is_empty()); - assert_eq!(terminal.finish_reason, Some(FinishReason::Stop)); - assert_eq!(terminal.stop_reason, Some(StopReason::Int(2))); - let usage = terminal.completion_usage.as_ref().expect("usage"); - assert_eq!((usage.prompt_tokens, usage.completion_tokens), (3, 1)); - - let requests = server.service.requests.lock().await; - let sent = requests.first().expect("recorded request"); - assert!(sent.streaming); - assert_eq!( - sent.tokenized.as_ref().unwrap().input_token_ids, - [11, 22, 33] - ); -} - -#[tokio::test] -async fn configured_context_length_overrides_the_engine_report() { - let server = FakeServer::start(FakeTrtllm::default()).await; - // The fake's GetModelInfo reports 4096, standing in for a release that - // under-reports `max_seq_len`; the operator configured 8192. - let engine = engine_with_context_length(&server.endpoint, 1, Some(8192)); - let config = engine.start(0).await.expect("start"); - assert_eq!(config.llm.unwrap().context_length, Some(8192)); - - let mut omits_max_tokens = request(); - omits_max_tokens.stop_conditions.max_tokens = None; - let outputs = collect(&engine, omits_max_tokens).await; - assert_eq!( - outputs.last().unwrap().finish_reason, - Some(FinishReason::Stop) - ); - - let requests = server.service.requests.lock().await; - let sent = requests.first().expect("recorded request"); - // 8192 configured minus request()'s three prompt tokens; the reported - // 4096 would give 4093. - assert_eq!(sent.max_tokens, 8189); -} - -#[tokio::test] -async fn a_max_seq_len_equal_to_max_input_len_is_not_registered() { - // TensorRT-LLM answers `max_seq_len` with `max_input_len` when the engine - // was started without `--max_seq_len`, so an equal pair carries no model - // information and must not reach registration. - let server = FakeServer::start(FakeTrtllm::default().reporting(1024, 1024).await).await; - let engine = engine(&server.endpoint, 1); - let config = engine.start(0).await.expect("start"); - // Nothing is registered, so the frontend keeps the context length it read - // from the model itself instead of being pinned to 1024. An omitted - // `max_tokens` then has no source to derive from and is rejected, which - // `omitted_max_tokens_without_context_length_is_rejected` covers. - assert_eq!(config.llm.unwrap().context_length, None); -} - -#[tokio::test] -async fn a_distinct_max_seq_len_is_registered() { - // `max_seq_len != max_input_len` means the engine was given an explicit - // `--max_seq_len`, so the report is real and is adopted. - let server = FakeServer::start(FakeTrtllm::default().reporting(2048, 1024).await).await; - let engine = engine(&server.endpoint, 1); - let config = engine.start(0).await.expect("start"); - assert_eq!(config.llm.unwrap().context_length, Some(2048)); -} - -#[tokio::test] -async fn grpc_request_errors_are_propagated() { - let service = FakeTrtllm::default(); - service.reject.store(true, Ordering::SeqCst); - let server = FakeServer::start(service).await; - let engine = engine(&server.endpoint, 1); - engine.start(0).await.expect("start"); - - // TRT-LLM surfaces an invalid-argument on the initial response header, so - // opening the stream fails rather than yielding an error item. - let context = dynamo_backend_common::testing::mock_context(); - let result = engine - .generate(request(), GenerateContext::new(context, None)) - .await; - assert!(result.is_err()); -} - -#[tokio::test] -async fn cancellation_yields_a_cancelled_terminal() { - let service = FakeTrtllm::default(); - service.hang.store(true, Ordering::SeqCst); - let server = FakeServer::start(service).await; - let engine = engine(&server.endpoint, 1); - engine.start(0).await.expect("start"); - - let context = dynamo_backend_common::testing::mock_context(); - let mut stream = engine - .generate(request(), GenerateContext::new(context.clone(), None)) - .await - .expect("generate"); - let first = stream.next().await.unwrap().unwrap(); - assert_eq!(first.token_ids, [42]); - context.stop_generating(); - let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) - .await - .expect("terminal within deadline") - .unwrap() - .unwrap(); - assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); -} - -#[tokio::test] -async fn abort_sends_the_abort_rpc_to_the_server() { - let server = FakeServer::start(FakeTrtllm::default()).await; - let engine = engine(&server.endpoint, 1); - engine.start(0).await.expect("start"); - - let context = dynamo_backend_common::testing::mock_context(); - let request_id = context.id().to_string(); - engine.abort(context).await; - - // The cancelled generation's ID must reach TensorRT-LLM, not just produce a - // local terminal, or the server keeps generating. - assert_eq!(server.service.aborts.lock().await.as_slice(), [request_id]); -} - -#[tokio::test] -async fn unsupported_features_fail_before_rpc_submission() { - let server = FakeServer::start(FakeTrtllm::default()).await; - let engine = engine(&server.endpoint, 1); - engine.start(0).await.expect("start"); - - let mut multiple = request(); - multiple.sampling_options.n = Some(2); - - let mut beam = request(); - beam.sampling_options.use_beam_search = Some(true); - - let mut embeds = request(); - embeds.prompt_embeds = Some("encoded".to_string()); - - let mut prompt_logprobs = request(); - prompt_logprobs.output_options.prompt_logprobs = Some(1); - - let mut visible_stops = request(); - visible_stops.stop_conditions.stop_token_ids_visible = Some(vec![7]); - - for unsupported in [multiple, beam, embeds, prompt_logprobs, visible_stops] { - let context = dynamo_backend_common::testing::mock_context(); - let result = engine - .generate(unsupported, GenerateContext::new(context, None)) - .await; - assert!(result.is_err()); + event: Some(pb::generate_response::Event::Token(pb::TokenOutput { + output_index: Some(0), + tokens, + text: String::new(), + })), + usage: None, } - assert!(server.service.requests.lock().await.is_empty()); } -#[tokio::test] -async fn pool_uses_each_configured_connection() { - let server = FakeServer::start(FakeTrtllm::default()).await; - let endpoint = - GrpcEndpoint::parse(&server.endpoint, "--grpc-endpoint").expect("valid endpoint"); - let client = TrtllmClient::connect(&endpoint, transport(2)) - .await - .expect("connect pool"); - assert_eq!(client.connection_count(), 2); - - for index in 0..4 { - let mut stream = client - .generate(pb::GenerateRequest { - request_id: format!("request-{index}"), - tokenized: Some(pb::TokenizedInput { - input_token_ids: vec![1, 2], - ..Default::default() - }), - max_tokens: 4, - streaming: true, - ..Default::default() - }) - .await - .expect("start stream"); - while stream.message().await.expect("message").is_some() {} +fn logprob_token(token_id: u32, logprob: f64) -> pb::TokenInfo { + pb::TokenInfo { + token_id, + token: String::new(), + logprob: Some(logprob), + rank: Some(1), + candidates: Vec::new(), } - - let ports: BTreeSet<_> = server - .service - .peers - .lock() - .await - .iter() - .map(SocketAddr::port) - .collect(); - assert_eq!(ports.len(), 2); } diff --git a/lib/sidecar/trtllm/src/tests/convert_request.rs b/lib/sidecar/trtllm/src/tests/convert_request.rs new file mode 100644 index 000000000000..4c2d868d83b7 --- /dev/null +++ b/lib/sidecar/trtllm/src/tests/convert_request.rs @@ -0,0 +1,250 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Building an OpenEngine `GenerateRequest` from a Dynamo request. + +use super::*; + +#[test] +fn request_maps_sampling_stop_and_output_fields() { + let proto = build_generate_request(&request(), "req-1", "served-model", None, AGG) + .expect("build request"); + assert_eq!(proto.request_id, "req-1"); + assert_eq!(proto.model, "served-model"); + match proto.input.as_ref().expect("input") { + pb::generate_request::Input::TokenIds(tokens) => assert_eq!(tokens.ids, [11, 22, 33]), + other => panic!("expected token IDs input, got {other:?}"), + } + + let stopping = proto.stopping.as_ref().unwrap(); + assert_eq!(stopping.max_tokens, Some(16)); + assert_eq!(stopping.min_tokens, Some(1)); + assert_eq!(stopping.ignore_eos, Some(true)); + // Stop-string retention is never enabled from `include_stop_str_in_output`. + assert_eq!(stopping.include_stop_in_output, None); + let conditions: Vec<_> = stopping + .conditions + .iter() + .map(|condition| condition.condition.clone().unwrap()) + .collect(); + assert!( + conditions.contains(&pb::stop_condition::Condition::StopText("done".to_string())), + "missing stop text: {conditions:?}" + ); + assert!( + conditions.contains(&pb::stop_condition::Condition::StopTokenId(2)), + "missing stop token id: {conditions:?}" + ); + + let sampling = proto.sampling.as_ref().unwrap(); + assert_eq!(sampling.top_k, Some(4)); + assert_eq!(sampling.top_p, Some(f64::from(0.9_f32))); + assert_eq!(sampling.min_p, Some(f64::from(0.1_f32))); + assert_eq!(sampling.temperature, Some(f64::from(0.2_f32))); + assert_eq!(sampling.seed, Some(123)); + assert_eq!(sampling.repetition_penalty, Some(f64::from(1.1_f32))); + assert_eq!(sampling.num_sequences, Some(1)); + + let response = proto.response.as_ref().unwrap(); + assert_eq!(response.return_output_logprobs, Some(true)); + match response + .output_candidates + .as_ref() + .unwrap() + .selection + .as_ref() + .unwrap() + { + pb::candidate_token_selection::Selection::TopN(n) => assert_eq!(*n, 1), + other => panic!("expected top_n selection, got {other:?}"), + } + + match proto.guided.as_ref().unwrap().guide.as_ref().unwrap() { + pb::guided_decoding::Guide::JsonSchema(guide) => assert!(guide.contains("object")), + other => panic!("expected JSON schema guide, got {other:?}"), + } +} + +#[test] +fn omitted_max_tokens_without_context_length_is_rejected() { + let mut req = request(); + req.stop_conditions.max_tokens = None; + let error = build_generate_request(&req, "req", "model", None, AGG) + .expect_err("must require max_tokens"); + assert!(error.to_string().contains("max_tokens")); +} + +#[test] +fn omitted_max_tokens_defaults_to_remaining_context() { + let mut req = request(); + req.stop_conditions.max_tokens = None; + // request() carries three prompt tokens ([11, 22, 33]); the default fills the + // remaining context: max(1, context_length - prompt_len). + let proto = build_generate_request(&req, "req", "model", limits(100), AGG) + .expect("build with fallback"); + assert_eq!(proto.stopping.unwrap().max_tokens, Some(97)); +} + +#[test] +fn omitted_max_tokens_default_is_floored_at_one() { + let mut req = request(); + req.stop_conditions.max_tokens = None; + // Prompt already fills (or exceeds) the context: default must not underflow to 0. + let proto = + build_generate_request(&req, "req", "model", limits(2), AGG).expect("build with fallback"); + assert_eq!(proto.stopping.unwrap().max_tokens, Some(1)); +} + +#[test] +fn top_k_all_tokens_is_left_unset() { + let mut req = request(); + req.sampling_options.top_k = Some(-1); + let proto = build_generate_request(&req, "req", "model", None, AGG).expect("build"); + assert_eq!(proto.sampling.unwrap().top_k, None); +} + +#[test] +fn unsupported_request_controls_are_rejected() { + // Controls the OpenEngine contract can neither forward nor faithfully honor: + // reject rather than fail open. + assert_rejected( + |r| r.sampling_options.include_stop_str_in_output = Some(true), + "include_stop_str_in_output", + ); + assert_rejected( + |r| r.stop_conditions.max_thinking_tokens = Some(32), + "max_thinking_tokens", + ); + assert_rejected( + |r| { + r.routing = Some(dynamo_backend_common::engine::RoutingHints { + cache_namespace: Some("tenant-a".to_string()), + ..Default::default() + }) + }, + "cache namespace", + ); + assert_rejected( + |r| { + r.routing = Some(dynamo_backend_common::engine::RoutingHints { + priority: Some(5), + ..Default::default() + }) + }, + "priority", + ); + // A negative top_k other than -1/0 (the "all tokens" sentinels) is invalid, + // not a silent widening to "all tokens". + assert_rejected(|r| r.sampling_options.top_k = Some(-5), "top_k must be"); + assert_rejected( + |r| r.sampling_options.seed = Some(-1), + "seed must be non-negative", + ); +} + +#[test] +fn logprobs_zero_keeps_selected_without_alternatives() { + let mut req = request(); + req.output_options.logprobs = Some(0); + // The wire request must still ask TRT-LLM for one candidate so the selected + // token's logprob is computed. + let proto = build_generate_request(&req, "req", "model", None, AGG).expect("build"); + match proto + .response + .unwrap() + .output_candidates + .unwrap() + .selection + .unwrap() + { + pb::candidate_token_selection::Selection::TopN(n) => assert_eq!(n, 1), + other => panic!("expected top_n selection, got {other:?}"), + } + + let mut state = ResponseState::new(&req, AGG); + let delta = state + .convert(token_response(vec![logprob_token(7, -0.1)])) + .expect("convert") + .expect("delta"); + assert_eq!(delta.log_probs.as_deref(), Some(&[-0.1_f64][..])); + // logprobs=0 surfaces the selected-token logprob but no top alternatives. + assert!(delta.top_logprobs.is_none()); +} + +/// The engine's own output cap can be smaller than the context window leaves +/// room for, and a request derived past it is rejected by the engine. +#[test] +fn the_derived_max_tokens_respects_the_engines_output_cap() { + let mut req = request(); + req.stop_conditions.max_tokens = None; + let proto = build_generate_request( + &req, + "req", + "model", + Some(ModelLimits { + context_length: Some(4096), + max_output_tokens: Some(256), + }), + AGG, + ) + .expect("build with an output cap"); + assert_eq!( + proto.stopping.expect("stopping").max_tokens, + Some(256), + "the derived default must not exceed what the engine will generate" + ); +} + +/// A `max_tokens` derived from the context window is clamped by the engine's +/// output cap and floored at 1, so it can land under an explicit `min_tokens`. +/// TensorRT-LLM does not cross-validate the pair, so a request carrying +/// `min_tokens > max_tokens` is one the engine resolves however it likes -- +/// reject it here, naming both values. +#[test] +fn a_derived_budget_below_min_tokens_is_rejected() { + let mut req = request(); + req.stop_conditions.max_tokens = None; + req.stop_conditions.min_tokens = Some(64); + // A 40-token window against this request's prompt leaves fewer than 64. + let prompt_len = req.token_ids.len() as u32; + let window = prompt_len + 8; + + let error = build_generate_request(&req, "req", "model", limits(window), AGG) + .expect_err("a minimum above the remaining budget cannot be served"); + let message = error.to_string(); + assert!( + message.contains("64") && message.contains('8'), + "the error must name both the minimum and what is left: {message}" + ); +} + +/// The prefill worker must not judge the minimum: it applies to the decode +/// worker's budget, and the two roles are paired without requiring equal +/// context lengths. A window too small here says nothing about the window the +/// decode worker will use. +#[test] +fn a_prefill_request_does_not_judge_a_minimum_against_its_own_window() { + let mut req = request(); + req.stop_conditions.max_tokens = None; + req.stop_conditions.min_tokens = Some(64); + let window = req.token_ids.len() as u32 + 8; + + build_generate_request(&req, "req", "model", limits(window), PREFILL) + .expect("only the decode worker's window governs the minimum"); +} + +/// The same request is served when the window leaves room for the minimum -- +/// the guard must reject the contradiction, not every request that has one. +#[test] +fn a_minimum_within_the_derived_budget_is_served() { + let mut req = request(); + req.stop_conditions.max_tokens = None; + req.stop_conditions.min_tokens = Some(8); + let window = req.token_ids.len() as u32 + 64; + + let proto = build_generate_request(&req, "req", "model", limits(window), AGG) + .expect("a minimum that fits is not a conflict"); + let stopping = proto.stopping.expect("stopping options"); + assert_eq!(stopping.min_tokens, Some(8)); + assert_eq!(stopping.max_tokens, Some(64)); +} diff --git a/lib/sidecar/trtllm/src/tests/convert_response.rs b/lib/sidecar/trtllm/src/tests/convert_response.rs new file mode 100644 index 000000000000..d60bebe3e785 --- /dev/null +++ b/lib/sidecar/trtllm/src/tests/convert_response.rs @@ -0,0 +1,304 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Reducing the server's streamed events into `LLMEngineOutput`s. + +use super::*; + +#[test] +fn token_then_finished_produces_delta_then_terminal_usage() { + let req = request(); + let mut state = ResponseState::new(&req, AGG); + + let delta = state + .convert(token_response(vec![ + logprob_token(7, -0.1), + logprob_token(8, -0.2), + ])) + .expect("convert token") + .expect("delta"); + assert_eq!(delta.token_ids, [7, 8]); + assert!(delta.finish_reason.is_none()); + let log_probs = delta.log_probs.as_ref().expect("log_probs"); + assert_eq!(log_probs, &[-0.1, -0.2]); + assert_eq!(delta.top_logprobs.as_ref().unwrap().len(), 2); + + let finished = pb::GenerateResponse { + request_id: "r".to_string(), + event: Some(pb::generate_response::Event::Finished( + pb::GenerationFinished { + output_index: Some(0), + reason: pb::FinishReason::Length as i32, + message: String::new(), + stop_match: None, + }, + )), + usage: Some(pb::Usage { + prompt_tokens: 3, + completion_tokens: 2, + total_tokens: 5, + cached_prompt_tokens: None, + reasoning_tokens: None, + }), + }; + let terminal = state + .convert(finished) + .expect("convert finished") + .expect("terminal"); + assert!(terminal.token_ids.is_empty()); + assert_eq!(terminal.finish_reason, Some(FinishReason::Length)); + let usage = terminal.completion_usage.as_ref().expect("usage"); + assert_eq!((usage.prompt_tokens, usage.completion_tokens), (3, 2)); +} + +#[test] +fn unsupported_output_index_is_rejected() { + let req = request(); + let mut state = ResponseState::new(&req, AGG); + let response = pb::GenerateResponse { + request_id: "r".to_string(), + event: Some(pb::generate_response::Event::Token(pb::TokenOutput { + output_index: Some(1), + tokens: vec![logprob_token(1, -0.1)], + text: String::new(), + })), + usage: None, + }; + assert!(state.convert(response).is_err()); +} + +#[test] +fn missing_event_is_rejected() { + let req = request(); + let mut state = ResponseState::new(&req, AGG); + let empty = pb::GenerateResponse { + request_id: "r".to_string(), + event: None, + usage: None, + }; + assert!(state.convert(empty).is_err()); +} + +#[test] +fn unspecified_finish_reason_is_rejected() { + let req = request(); + let mut state = ResponseState::new(&req, AGG); + let finished = pb::GenerateResponse { + request_id: "r".to_string(), + event: Some(pb::generate_response::Event::Finished( + pb::GenerationFinished { + output_index: Some(0), + reason: pb::FinishReason::Unspecified as i32, + message: String::new(), + stop_match: None, + }, + )), + usage: None, + }; + assert!(state.convert(finished).is_err()); +} + +#[test] +fn engine_error_event_is_surfaced_as_error() { + let req = request(); + let mut state = ResponseState::new(&req, AGG); + let error = pb::GenerateResponse { + request_id: "r".to_string(), + event: Some(pb::generate_response::Event::Error(pb::EngineError { + code: pb::ErrorCode::Internal as i32, + message: "boom".to_string(), + retryable: false, + })), + usage: None, + }; + let error = state.convert(error).expect_err("engine error must surface"); + assert!(error.to_string().contains("boom")); +} + +/// Cache hits are measured during the context phase, so a decode worker can +/// only report them by carrying the handoff's count through to its terminal. +/// An aggregated worker reads its own engine's count off the same terminal. +#[test] +fn cached_prompt_tokens_reach_the_client_on_both_paths() { + let terminal = |mut state: ResponseState, reported: Option| { + let response = pb::GenerateResponse { + request_id: "req".to_string(), + event: Some(pb::generate_response::Event::Finished( + pb::GenerationFinished { + output_index: Some(0), + reason: pb::FinishReason::Stop as i32, + ..Default::default() + }, + )), + usage: Some(pb::Usage { + prompt_tokens: 11, + completion_tokens: 3, + total_tokens: 14, + cached_prompt_tokens: reported, + reasoning_tokens: None, + }), + }; + state + .convert(response) + .expect("finished converts") + .expect("finished yields a terminal") + .completion_usage + .expect("usage is set") + .prompt_tokens_details + .and_then(|details| details.cached_tokens) + }; + + assert_eq!( + terminal(ResponseState::new(&request(), AGG), Some(7)), + Some(7) + ); + + let mut decode_request = request(); + decode_request.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: crate::disagg::session_to_json(fake_session()).expect("handoff"), + prompt_tokens_details: Some(dynamo_backend_common::PromptTokensDetails { + audio_tokens: None, + cached_tokens: Some(5), + }), + }); + let state = ResponseState::new(&decode_request, DisaggregationMode::Decode); + assert_eq!( + terminal(state, None), + Some(5), + "the decode leg reports the prefill leg's cache hits" + ); + + // The decode engine counts the blocks transferred into it as cache hits, + // which would read as a ~100% prefix-cache hit rate on every disaggregated + // request. The handoff's count is the real one. + let state = ResponseState::new(&decode_request, DisaggregationMode::Decode); + assert_eq!( + terminal(state, Some(11)), + Some(5), + "the decode engine's own count must not overwrite the handoff's" + ); + let state = ResponseState::new(&decode_request, DisaggregationMode::Decode); + assert_eq!( + terminal(state, Some(0)), + Some(5), + "a decode engine reporting zero must not erase the handoff's count" + ); + + // A handoff that carried no count means the context phase measured no + // hits. The decode engine's own count is about transferred blocks, so it + // must not fill the gap. + let mut uncounted = decode_request.clone(); + uncounted.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: crate::disagg::session_to_json(fake_session()).expect("handoff"), + prompt_tokens_details: None, + }); + let state = ResponseState::new(&uncounted, DisaggregationMode::Decode); + assert_eq!(terminal(state, Some(11)), None); + + // A decode leg that ran its own context phase (conditional disaggregation + // bypassed the prefill worker) has no handoff, so its engine is the only + // source. + let state = ResponseState::new(&request(), DisaggregationMode::Decode); + assert_eq!(terminal(state, Some(4)), Some(4)); + + // A count measured against an expanded prompt must not exceed the prompt + // the client sent. + assert_eq!( + terminal(ResponseState::new(&request(), AGG), Some(9_999)), + Some(11) + ); +} + +/// `Cancelled` must survive as itself; `Overloaded` must not become +/// `WorkerOverloaded`. The only server site that emits `Overloaded` is the +/// consumer-stall watchdog, so migrating on it would re-dispatch to a second +/// worker that stalls the same way. +#[test] +fn engine_error_codes_the_router_acts_on_are_preserved() { + let overloaded = engine_error(pb::EngineError { + code: pb::ErrorCode::Overloaded as i32, + message: "at capacity".to_string(), + retryable: true, + }); + assert!(!matches!( + overloaded.error_type(), + ErrorType::WorkerOverloaded + )); + + let cancelled = engine_error(pb::EngineError { + code: pb::ErrorCode::Cancelled as i32, + message: "aborted".to_string(), + retryable: false, + }); + assert!(matches!(cancelled.error_type(), ErrorType::Cancelled)); +} + +/// A role mismatch means the frontend routed to a worker whose engine is in +/// the other disaggregation role, so every request to it fails identically. +/// An opaque 500 gives the operator nothing to search for, and the class must +/// stay client-side: retrying elsewhere cannot fix a misconfigured deployment. +#[test] +fn a_role_mismatch_names_the_misconfiguration() { + let error = engine_error(pb::EngineError { + code: pb::ErrorCode::RoleMismatch as i32, + message: "generation_only request on a context engine".to_string(), + retryable: false, + }); + assert!( + matches!( + error.error_type(), + ErrorType::Backend(dynamo_backend_common::BackendError::InvalidArgument) + ), + "a deployment mistake is a client error, not an opaque engine failure" + ); + assert!( + error.to_string().contains("--disaggregation-mode"), + "the error must name the flag the operator has to check: {error}" + ); +} + +/// A handoff the decode engine cannot resolve is not migratable: the retry +/// would replay the same dead session and fail the same way on the next +/// worker. It still has to say what to look at, which is why these two codes +/// no longer fall into the catch-all. +#[test] +fn an_unresolvable_handoff_is_named_but_not_migratable() { + for code in [ + pb::ErrorCode::KvSessionNotFound, + pb::ErrorCode::KvTransferFailed, + ] { + let error = engine_error(pb::EngineError { + code: code as i32, + message: "session 12345 is gone".to_string(), + retryable: false, + }); + assert!( + !matches!(error.error_type(), ErrorType::WorkerOverloaded), + "{code:?} must not be migratable" + ); + assert!( + error.to_string().contains("cache transceiver"), + "{code:?} must point at the transceiver: {error}" + ); + } +} + +/// The OpenEngine contract requires an explicit output index. Defaulting an +/// absent one to zero hides exactly the drift the index check exists to catch. +#[test] +fn a_response_without_an_output_index_is_rejected() { + let mut state = ResponseState::new(&request(), AGG); + let response = pb::GenerateResponse { + request_id: "req".to_string(), + event: Some(pb::generate_response::Event::Token(pb::TokenOutput { + output_index: None, + tokens: vec![pb::TokenInfo { + token_id: 7, + ..Default::default() + }], + ..Default::default() + })), + usage: None, + }; + assert!(state.convert(response).is_err()); +} diff --git a/lib/sidecar/trtllm/src/tests/disagg.rs b/lib/sidecar/trtllm/src/tests/disagg.rs new file mode 100644 index 000000000000..05f7a4639968 --- /dev/null +++ b/lib/sidecar/trtllm/src/tests/disagg.rs @@ -0,0 +1,401 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The prefill/decode legs and the KV handoff codec between them. + +use super::*; + +/// The prefill worker marks its request `context_only` and caps generation at +/// the single token the context phase produces. +#[test] +fn prefill_request_is_marked_context_only() { + let mut req = request(); + req.stop_conditions.max_tokens = Some(128); + req.stop_conditions.min_tokens = Some(8); + req.output_options.logprobs = Some(1); + let proto = build_generate_request(&req, "req", "model", None, DisaggregationMode::Prefill) + .expect("build prefill request"); + + assert!( + is_context_only(&proto), + "prefill must set extra.request_type" + ); + let stopping = proto.stopping.expect("stopping options"); + assert_eq!(stopping.max_tokens, Some(1)); + // The minimum masks EOS rather than extending generation, so it must + // survive: without it the context phase can sample EOS on its single token + // and terminate with `Stop` instead of the `PrefillReady` handoff. + assert_eq!(stopping.min_tokens, Some(8)); + // The prefill worker surfaces no tokens, but it must still compute + // logprobs: the first generated token comes from the context phase and its + // logprob only reaches the decode worker through the handoff. + assert_eq!( + proto + .response + .expect("response options") + .return_output_logprobs, + Some(true) + ); + assert!(proto.kv.is_none(), "prefill carries no session to replay"); +} + +/// `PrefillReady` is the prefill worker's terminal chunk: no tokens, and the +/// handoff the decode worker will replay. +#[tokio::test] +async fn prefill_ready_is_the_terminal_handoff() { + let server = FakeServer::start(FakeTrtllm::default()).await; + let engine = engine_in_mode(&server.endpoint, 1, DisaggregationMode::Prefill); + engine.start(0).await.expect("start"); + + let outputs = collect(&engine, request()).await; + let terminal = outputs + .iter() + .find(|output| output.finish_reason.is_some()) + .expect("a terminal chunk"); + // The frontend's prefill router chains into decode only for `Length`; any + // other terminal reason is returned to the caller as a finished request. + assert_eq!(terminal.finish_reason, Some(FinishReason::Length)); + + assert!( + outputs.iter().all(|output| output.token_ids.is_empty()), + "a prefill worker must not stream tokens to the client" + ); + let handoff = terminal + .disaggregated_params + .as_ref() + .expect("terminal carries the prefill handoff"); + assert_eq!(handoff["session_id"], json!("12345")); + assert_eq!(handoff["transfer_backend"], json!("NIXL")); + assert_eq!(handoff["endpoints"][0]["port"], json!(5601)); + assert_eq!(handoff["attributes"]["first_gen_tokens"], json!([42])); +} + +/// The decode worker replays the prefill handoff verbatim in `kv.session`. +#[tokio::test] +async fn decode_request_replays_the_prefill_session() { + let server = FakeServer::start(FakeTrtllm::default()).await; + + // Phase 1: prefill produces the handoff. + let prefill = engine_in_mode(&server.endpoint, 1, DisaggregationMode::Prefill); + prefill.start(0).await.expect("start prefill"); + let handoff = collect(&prefill, request()) + .await + .into_iter() + .find_map(|output| output.disaggregated_params) + .expect("prefill handoff"); + + // Phase 2: decode replays it. + let decode = engine_in_mode(&server.endpoint, 1, DisaggregationMode::Decode); + decode.start(0).await.expect("start decode"); + let mut req = request(); + req.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: handoff, + prompt_tokens_details: None, + }); + let outputs = collect(&decode, req).await; + assert!( + outputs.iter().any(|output| !output.token_ids.is_empty()), + "the decode worker streams the completion" + ); + + let session = server.service.requests.lock().await[1] + .kv + .as_ref() + .and_then(|kv| kv.session.as_ref()) + .expect("decode request carries kv.session") + .clone(); + assert_eq!( + session, + fake_session(), + "the handoff must round-trip intact" + ); +} + +/// Conditional disaggregation dispatches straight to a decode worker with no +/// handoff, expecting it to run the context phase itself. Rejecting that would +/// fail every bypassed request. +#[test] +fn decode_without_a_prefill_result_runs_the_whole_request() { + let proto = build_generate_request( + &request(), + "req", + "model", + limits(4096), + DisaggregationMode::Decode, + ) + .expect("a decode request without a handoff is the bypass path"); + assert!( + proto.kv.and_then(|kv| kv.session).is_none(), + "no session should be replayed when none was handed off" + ); + assert!( + proto.extra.is_none(), + "a bypassed request is not context_only" + ); +} + +/// A handoff on a non-decode worker means the frontend routed the request to +/// the wrong role - fail loudly rather than silently prefill it again. +#[test] +fn prefill_result_on_a_non_decode_worker_is_rejected() { + let mut req = request(); + req.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: json!({"session_id": "1"}), + prompt_tokens_details: None, + }); + for mode in [DisaggregationMode::Aggregated, DisaggregationMode::Prefill] { + let error = build_generate_request(&req, "req", "model", None, mode) + .expect_err("handoff must be rejected"); + assert!( + error + .to_string() + .contains("must be routed to a decode worker"), + "unexpected error for {mode}: {error}" + ); + } +} + +/// An aggregated worker never asks for a context handoff, so receiving one is +/// protocol drift rather than a silent no-op. +#[test] +fn unexpected_prefill_ready_on_an_aggregated_worker_is_rejected() { + let mut state = ResponseState::new(&request(), AGG); + let error = state + .convert(pb::GenerateResponse { + request_id: "req".to_string(), + event: Some(pb::generate_response::Event::PrefillReady( + pb::PrefillReady { + kv_session: Some(fake_session()), + }, + )), + usage: None, + }) + .expect_err("prefill_ready must be rejected"); + assert!( + error.to_string().contains("not running in prefill mode"), + "unexpected error: {error}" + ); +} + +/// The OpenEngine servicer suppresses `finished` on a context request only +/// when it already sent a `PrefillReady`; otherwise -- a stop condition during +/// the one-token context phase, or a cancellation before transmission -- it +/// deliberately emits the real terminal, because there is no decode leg to +/// answer the caller. Context tokens are held back on the normal path so the +/// decode leg's replay cannot duplicate them, which makes this terminal the +/// only place they can surface. +#[test] +fn a_prefill_terminal_without_a_handoff_carries_its_tokens() { + let mut state = ResponseState::new(&request(), DisaggregationMode::Prefill); + let token = pb::GenerateResponse { + request_id: "req".to_string(), + event: Some(pb::generate_response::Event::Token(pb::TokenOutput { + output_index: Some(0), + tokens: vec![pb::TokenInfo { + token_id: 99, + logprob: Some(-0.5), + rank: Some(1), + candidates: vec![pb::LogProb { + token_id: 100, + logprob: -1.5, + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + })), + usage: None, + }; + assert!( + state.convert(token).expect("token converts").is_none(), + "a context token is held back, not streamed" + ); + + let response = pb::GenerateResponse { + request_id: "req".to_string(), + event: Some(pb::generate_response::Event::Finished( + pb::GenerationFinished { + output_index: Some(0), + reason: pb::FinishReason::Stop as i32, + message: String::new(), + stop_match: None, + }, + )), + usage: None, + }; + let terminal = state + .convert(response) + .expect("a handoff-less prefill terminal is serviceable") + .expect("it yields a terminal"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Stop)); + assert_eq!( + terminal.token_ids, + [99], + "the caller gets what the context phase produced" + ); + // The held token is the whole answer here, so dropping to bare IDs would + // silently strip logprobs a caller asked for and cannot get anywhere else. + assert_eq!(terminal.log_probs.as_deref(), Some(&[-0.5][..])); + let candidates = terminal.top_logprobs.expect("candidates"); + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0][0].token_id, 100); +} + +/// `PrefillReady` is the final response for a context request, so its usage is +/// authoritative -- the frontend forwards `cached_tokens` from it to the decode +/// leg and cannot reconstruct it. +#[test] +fn prefill_ready_keeps_the_engines_reported_usage() { + let mut state = ResponseState::new(&request(), DisaggregationMode::Prefill); + let response = pb::GenerateResponse { + request_id: "req".to_string(), + event: Some(pb::generate_response::Event::PrefillReady( + pb::PrefillReady { + kv_session: Some(fake_session()), + }, + )), + usage: Some(pb::Usage { + prompt_tokens: 11, + completion_tokens: 0, + total_tokens: 11, + cached_prompt_tokens: Some(7), + reasoning_tokens: None, + }), + }; + let output = state + .convert(response) + .expect("prefill_ready converts") + .expect("prefill_ready yields a terminal"); + let usage = output.completion_usage.expect("usage is set"); + assert_eq!(usage.prompt_tokens, 11); + assert_eq!( + usage + .prompt_tokens_details + .and_then(|details| details.cached_tokens), + Some(7), + "cached prompt tokens must survive to the decode leg" + ); +} + +/// A context request that ends at its one-token budget never transmitted its +/// KV, so there is no decode leg to finish the completion. Returning that +/// single token would silently truncate the answer, and passing the `Length` +/// terminal through makes the frontend's prefill router report a missing +/// handoff instead of the real cause. +#[test] +fn a_prefill_terminal_that_ran_out_of_budget_is_rejected() { + let mut state = ResponseState::new(&request(), DisaggregationMode::Prefill); + let response = pb::GenerateResponse { + request_id: "req".to_string(), + event: Some(pb::generate_response::Event::Finished( + pb::GenerationFinished { + output_index: Some(0), + reason: pb::FinishReason::Length as i32, + ..Default::default() + }, + )), + usage: None, + }; + let error = state + .convert(response) + .expect_err("a budget-exhausted context request has no answer to return"); + assert!( + error.to_string().contains("without a kv_session handoff"), + "the error must name the missing handoff: {error}" + ); +} + +/// The handoff codec is symmetric: the decode worker requires exactly what the +/// prefill worker wrote. A field lost in transit must fail by name rather than +/// decode into a plausible session. `attributes` matters most -- the server +/// reads the session's location and rank out of it, and accepts an endpoint in +/// place of an `opaque_state`, so a handoff that lost it would pass the +/// server's own guard and resume a session with no opaque state at rank 0. +#[test] +fn a_handoff_missing_a_field_is_rejected() { + let mut session = fake_session(); + session.dp_rank = 3; + let encoded = crate::disagg::session_to_json(session).expect("encode"); + + let decoded = crate::disagg::session_from_json(&encoded).expect("round trip"); + assert_eq!(decoded.dp_rank, 3, "the rank must survive the round trip"); + + for field in [ + "attributes", + "dp_rank", + "endpoints", + "transfer_backend", + "session_id", + ] { + let mut mangled = encoded.clone(); + mangled + .as_object_mut() + .expect("handoff object") + .remove(field) + .expect("field is present"); + let error = crate::disagg::session_from_json(&mangled) + .expect_err("a handoff missing a field is not usable"); + assert!( + error.to_string().contains(field), + "the error must name {field}: {error}" + ); + } +} + +/// A newer prefill worker may add fields this decode worker has never heard of. +/// Rejecting them would fail every new-prefill/old-decode request during a +/// rolling upgrade, as a non-migratable 400, with no version to negotiate on. +#[test] +fn a_handoff_carrying_an_unknown_field_still_decodes() { + let encoded = crate::disagg::session_to_json(fake_session()).expect("encode"); + let mut newer = encoded.clone(); + newer + .as_object_mut() + .expect("handoff object") + .insert("schedule_style".to_string(), serde_json::json!(2)); + + crate::disagg::session_from_json(&newer) + .expect("a handoff from a newer peer must still decode"); +} + +/// The prefill leg refuses to emit a handoff the decode leg could not resolve, +/// so the failure names the prefill worker rather than surfacing one hop later. +#[test] +fn a_prefill_ready_without_attributes_is_rejected() { + let mut session = fake_session(); + session.attributes_struct = None; + + let error = crate::disagg::session_to_json(session) + .expect_err("a session with no attributes cannot locate the context worker"); + assert!( + error.to_string().contains("attributes_struct"), + "the error must name the missing field: {error}" + ); +} + +/// The same rolling-upgrade rule one level down: a newer peer may add fields +/// inside `endpoints[]`. `deny_unknown_fields` there reversed the module's +/// stated policy and would fail that request as a non-migratable 400. +#[test] +fn an_unknown_field_inside_an_endpoint_still_decodes() { + let encoded = crate::disagg::session_to_json(fake_session()).expect("encode"); + let mut newer = encoded.clone(); + newer + .as_object_mut() + .expect("handoff object") + .get_mut("endpoints") + .expect("endpoints") + .as_array_mut() + .expect("endpoints array") + .first_mut() + .expect("at least one endpoint") + .as_object_mut() + .expect("endpoint object") + .insert( + "future_field".to_string(), + serde_json::json!("from a newer peer"), + ); + + crate::disagg::session_from_json(&newer) + .expect("an endpoint from a newer peer must still decode"); +} diff --git a/lib/sidecar/trtllm/src/tests/e2e.rs b/lib/sidecar/trtllm/src/tests/e2e.rs new file mode 100644 index 000000000000..984391b2e9eb --- /dev/null +++ b/lib/sidecar/trtllm/src/tests/e2e.rs @@ -0,0 +1,204 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Ignored by default: these drive a live TensorRT-LLM OpenEngine server. + +use super::*; + +/// Drives a real disaggregated prefill -> decode handoff against two live +/// OpenEngine servers, each backed by a TensorRT-LLM engine with a KV cache +/// transceiver. Ignored by default; run explicitly with both endpoints: +/// +/// ```text +/// TRTLLM_E2E_PREFILL_ENDPOINT=http://127.0.0.1:50051 \ +/// TRTLLM_E2E_DECODE_ENDPOINT=http://127.0.0.1:50052 \ +/// cargo test -p dynamo-trtllm-sidecar e2e_real_disagg -- --ignored --nocapture +/// ``` +#[tokio::test] +#[ignore = "requires two live TensorRT-LLM OpenEngine servers with a KV transceiver"] +async fn e2e_real_disagg_handoff() { + let prefill_endpoint = std::env::var("TRTLLM_E2E_PREFILL_ENDPOINT") + .expect("set TRTLLM_E2E_PREFILL_ENDPOINT, e.g. http://127.0.0.1:50051"); + let decode_endpoint = std::env::var("TRTLLM_E2E_DECODE_ENDPOINT") + .expect("set TRTLLM_E2E_DECODE_ENDPOINT, e.g. http://127.0.0.1:50052"); + let model = std::env::var("TRTLLM_E2E_MODEL") + .unwrap_or_else(|_| "TinyLlama/TinyLlama-1.1B-Chat-v1.0".to_string()); + + // `start` resolves the context length from `Control.GetModelInfo`. + let configured = |source: String| ConfiguredModel { + source, + context_length: None, + }; + let prefill = TrtllmSidecarEngine::new( + GrpcEndpoint::parse(&prefill_endpoint, "--grpc-endpoint").expect("valid endpoint"), + transport(1), + configured(model.clone()), + DisaggregationMode::Prefill, + ); + let decode = TrtllmSidecarEngine::new( + GrpcEndpoint::parse(&decode_endpoint, "--grpc-endpoint").expect("valid endpoint"), + transport(1), + configured(model.clone()), + DisaggregationMode::Decode, + ); + prefill.start(0).await.expect("start prefill worker"); + decode.start(0).await.expect("start decode worker"); + + // " Hello, my name is" in the Llama tokenizer. + let base = || { + PreprocessedRequest::builder() + .model(model.clone()) + .token_ids(vec![1, 15043, 29892, 590, 1024, 338]) + .stop_conditions(StopConditions { + max_tokens: Some(16), + ..Default::default() + }) + .sampling_options(SamplingOptions { + temperature: Some(0.0), + ..Default::default() + }) + .output_options(OutputOptions::default()) + .build() + .expect("request") + }; + + // Phase 1: prefill. + let prefill_outputs = collect(&prefill, base()).await; + eprintln!( + "[e2e-disagg] prefill produced {} chunk(s)", + prefill_outputs.len() + ); + assert!( + prefill_outputs + .iter() + .all(|output| output.token_ids.is_empty()), + "the prefill worker must not stream tokens" + ); + let handoff = prefill_outputs + .iter() + .find_map(|output| output.disaggregated_params.clone()) + .expect("prefill terminal carries the KV handoff"); + eprintln!("[e2e-disagg] handoff = {handoff}"); + + // Phase 2: decode replays the handoff. + let mut decode_request = base(); + decode_request.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: handoff, + prompt_tokens_details: None, + }); + let decode_outputs = collect(&decode, decode_request).await; + + let mut generated = Vec::new(); + let mut terminal = None; + for output in &decode_outputs { + generated.extend(output.token_ids.iter().copied()); + if output.finish_reason.is_some() { + terminal = Some(output); + } + } + eprintln!("[e2e-disagg] decode generated token IDs: {generated:?}"); + let terminal = terminal.expect("a terminal output carrying a finish_reason"); + eprintln!("[e2e-disagg] finish_reason = {:?}", terminal.finish_reason); + let usage = terminal.completion_usage.as_ref().expect("terminal usage"); + eprintln!( + "[e2e-disagg] usage: prompt={}, completion={}", + usage.prompt_tokens, usage.completion_tokens + ); + + assert!( + !generated.is_empty(), + "the decode worker must generate tokens from the prefill handoff" + ); +} + +/// Drives the real `TrtllmSidecarEngine` against a live OpenEngine gRPC server +/// (a real TensorRT-LLM engine on a GPU). Ignored by default; run explicitly with +/// a reachable endpoint: +/// +/// ```text +/// TRTLLM_E2E_ENDPOINT=http://127.0.0.1:50051 \ +/// cargo test -p dynamo-trtllm-sidecar e2e_real_openengine -- --ignored --nocapture +/// ``` +#[tokio::test] +#[ignore = "requires a live TensorRT-LLM OpenEngine server; set TRTLLM_E2E_ENDPOINT"] +async fn e2e_real_openengine_server() { + let endpoint = std::env::var("TRTLLM_E2E_ENDPOINT") + .expect("set TRTLLM_E2E_ENDPOINT, e.g. http://127.0.0.1:50051"); + let model = std::env::var("TRTLLM_E2E_MODEL") + .unwrap_or_else(|_| "TinyLlama/TinyLlama-1.1B-Chat-v1.0".to_string()); + + let engine = TrtllmSidecarEngine::new( + GrpcEndpoint::parse(&endpoint, "--grpc-endpoint").expect("valid endpoint"), + transport(2), + // `start` resolves the context length from `Control.GetModelInfo`. + ConfiguredModel { + source: model.clone(), + context_length: None, + }, + AGG, + ); + let config = engine + .start(0) + .await + .expect("start against the live server"); + eprintln!("[e2e] connected; registered model = {}", config.model); + + // " Hello, my name is" in the Llama tokenizer; any valid token IDs work. + let request = PreprocessedRequest::builder() + .model(model) + .token_ids(vec![1, 15043, 29892, 590, 1024, 338]) + .stop_conditions(StopConditions { + max_tokens: Some(16), + ..Default::default() + }) + .sampling_options(SamplingOptions { + temperature: Some(0.0), + ..Default::default() + }) + .output_options(OutputOptions { + logprobs: Some(1), + ..Default::default() + }) + .build() + .expect("request"); + + let outputs = collect(&engine, request).await; + eprintln!("[e2e] received {} stream item(s)", outputs.len()); + + let mut generated = Vec::new(); + let mut terminal = None; + let mut saw_logprobs = false; + for output in &outputs { + generated.extend(output.token_ids.iter().copied()); + saw_logprobs |= output.log_probs.is_some(); + if output.finish_reason.is_some() { + terminal = Some(output); + } + } + eprintln!("[e2e] generated token IDs: {generated:?}"); + + let terminal = terminal.expect("a terminal output carrying a finish_reason"); + eprintln!("[e2e] finish_reason = {:?}", terminal.finish_reason); + let usage = terminal.completion_usage.as_ref().expect("terminal usage"); + eprintln!( + "[e2e] usage: prompt={}, completion={}", + usage.prompt_tokens, usage.completion_tokens + ); + + assert!( + !generated.is_empty(), + "expected at least one generated token" + ); + assert_eq!( + usage.prompt_tokens, 6, + "prompt token count should echo the input" + ); + assert!( + usage.completion_tokens > 0, + "expected nonzero completion usage" + ); + assert!( + saw_logprobs, + "logprobs were requested but none were surfaced" + ); +} diff --git a/lib/sidecar/trtllm/src/tests/engine.rs b/lib/sidecar/trtllm/src/tests/engine.rs new file mode 100644 index 000000000000..f5dfac6bc0b3 --- /dev/null +++ b/lib/sidecar/trtllm/src/tests/engine.rs @@ -0,0 +1,542 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The engine driven against an in-process fake OpenEngine server. + +use super::*; + +#[tokio::test] +async fn aggregated_generation_streams_delta_then_terminal() { + let server = FakeServer::start(FakeTrtllm::default()).await; + let engine = engine(&server.endpoint, 2); + let config = engine.start(0).await.expect("start"); + assert_eq!(config.model, "model-source"); + // GetModelInfo reports max_context_length 4096. + assert_eq!(config.llm.unwrap().context_length, Some(4096)); + + let outputs = collect(&engine, request()).await; + assert_eq!(outputs.len(), 2); + assert_eq!(outputs[0].token_ids, [42]); + assert!(outputs[0].finish_reason.is_none()); + assert_eq!(outputs[0].log_probs.as_deref(), Some(&[-0.25][..])); + + let terminal = &outputs[1]; + assert!(terminal.token_ids.is_empty()); + assert_eq!(terminal.finish_reason, Some(FinishReason::Stop)); + assert_eq!(terminal.stop_reason, Some(StopReason::Int(2))); + let usage = terminal.completion_usage.as_ref().expect("usage"); + assert_eq!((usage.prompt_tokens, usage.completion_tokens), (3, 1)); + + let requests = server.service.requests.lock().await; + let sent = requests.first().expect("recorded request"); + assert_eq!(sent.model, "model-source"); + match sent.input.as_ref().expect("input") { + pb::generate_request::Input::TokenIds(tokens) => assert_eq!(tokens.ids, [11, 22, 33]), + other => panic!("expected token IDs input, got {other:?}"), + } +} + +#[tokio::test] +async fn grpc_request_errors_are_propagated() { + let service = FakeTrtllm::default(); + service.reject.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine(&server.endpoint, 1); + engine.start(0).await.expect("start"); + + // TRT-LLM surfaces an invalid-argument on the initial response header, so + // opening the stream fails rather than yielding an error item. + let context = dynamo_backend_common::testing::mock_context(); + let result = engine + .generate(request(), GenerateContext::new(context, None)) + .await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn cancellation_yields_a_cancelled_terminal() { + let service = FakeTrtllm::default(); + service.hang.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine(&server.endpoint, 1); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + let mut stream = engine + .generate(request(), GenerateContext::new(context.clone(), None)) + .await + .expect("generate"); + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(first.token_ids, [42]); + context.stop_generating(); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) + .await + .expect("terminal within deadline") + .unwrap() + .unwrap(); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); +} + +/// A request cancelled before dispatch must not reach the engine. +#[tokio::test] +async fn an_aggregated_request_cancelled_before_dispatch_never_reaches_the_engine() { + let server = FakeServer::start(FakeTrtllm::default()).await; + let engine = engine(&server.endpoint, 1); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + context.stop_generating(); + + let mut stream = engine + .generate(request(), GenerateContext::new(context.clone(), None)) + .await + .expect("generate"); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) + .await + .expect("terminal within deadline") + .unwrap() + .unwrap(); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); + assert!( + server.service.requests.lock().await.is_empty(), + "a request cancelled before dispatch must not be sent to the engine" + ); +} + +/// The server answers `openengine-target-dp-rank` with UNIMPLEMENTED, so a +/// rank hint has to be refused before dispatch: sending it anyway fails the +/// whole request with a non-migratable 5xx. `nvext.dp_rank` and the +/// `x-dynamo-dp-rank` header both land in these fields, so this is reachable +/// without a KV router. +#[tokio::test] +async fn a_data_parallel_rank_hint_is_rejected_before_dispatch() { + let server = FakeServer::start(FakeTrtllm::default()).await; + + for (mode, hints) in [ + ( + DisaggregationMode::Prefill, + dynamo_backend_common::engine::RoutingHints { + prefill_dp_rank: Some(7), + ..Default::default() + }, + ), + ( + DisaggregationMode::Decode, + dynamo_backend_common::engine::RoutingHints { + dp_rank: Some(3), + ..Default::default() + }, + ), + ( + AGG, + dynamo_backend_common::engine::RoutingHints { + dp_rank: Some(3), + ..Default::default() + }, + ), + ] { + let engine = engine_in_mode(&server.endpoint, 1, mode); + engine.start(0).await.expect("start"); + let mut req = request(); + req.routing = Some(hints); + let context = dynamo_backend_common::testing::mock_context(); + let result = engine + .generate(req, GenerateContext::new(context, None)) + .await; + assert!(result.is_err(), "{mode:?} must refuse a rank hint"); + } + assert!( + server.service.requests.lock().await.is_empty(), + "nothing may reach the engine" + ); +} + +/// A decode request holding transferred KV blocks must outlive its client until +/// the first token proves the transfer landed -- dropping the stream earlier +/// strands the prefill worker's blocks. It must not outlive it any longer than +/// that, or a cancelled request generates its whole budget with no consumer. +#[tokio::test] +async fn a_cancelled_decode_request_survives_only_until_the_transfer_lands() { + let service = FakeTrtllm::default(); + service.hang.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine_in_mode(&server.endpoint, 1, DisaggregationMode::Decode); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + let mut req = request(); + req.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: crate::disagg::session_to_json(fake_session()).expect("handoff"), + prompt_tokens_details: None, + }); + let mut stream = engine + .generate(req, GenerateContext::new(context.clone(), None)) + .await + .expect("generate"); + + context.stop_generating(); + // Still deferring: nothing has confirmed the KV transfer yet. + let first = stream.next().await.unwrap().unwrap(); + assert_eq!( + first.finish_reason, None, + "the deferral outlasts the client" + ); + assert!(!first.token_ids.is_empty()); + + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) + .await + .expect("the deferral lifts once a token lands") + .unwrap() + .unwrap(); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); +} + +#[tokio::test] +async fn abort_sends_the_abort_rpc_to_the_server() { + let server = FakeServer::start(FakeTrtllm::default()).await; + let engine = engine(&server.endpoint, 1); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + let request_id = context.id().to_string(); + engine.abort(context).await; + + // The cancelled generation's ID must reach TensorRT-LLM, not just produce a + // local terminal, or the server keeps generating. + assert_eq!(server.service.aborts.lock().await.as_slice(), [request_id]); +} + +#[tokio::test] +async fn unsupported_features_fail_before_rpc_submission() { + let server = FakeServer::start(FakeTrtllm::default()).await; + let engine = engine(&server.endpoint, 1); + engine.start(0).await.expect("start"); + + let mut multiple = request(); + multiple.sampling_options.n = Some(2); + + let mut beam = request(); + beam.sampling_options.use_beam_search = Some(true); + + let mut embeds = request(); + embeds.prompt_embeds = Some("encoded".to_string()); + + let mut prompt_logprobs = request(); + prompt_logprobs.output_options.prompt_logprobs = Some(1); + + let mut visible_stops = request(); + visible_stops.stop_conditions.stop_token_ids_visible = Some(vec![7]); + + for unsupported in [multiple, beam, embeds, prompt_logprobs, visible_stops] { + let context = dynamo_backend_common::testing::mock_context(); + let result = engine + .generate(unsupported, GenerateContext::new(context, None)) + .await; + assert!(result.is_err()); + } + assert!(server.service.requests.lock().await.is_empty()); +} + +#[tokio::test] +async fn pool_uses_each_configured_connection() { + let server = FakeServer::start(FakeTrtllm::default()).await; + let endpoint = + GrpcEndpoint::parse(&server.endpoint, "--grpc-endpoint").expect("valid endpoint"); + let client = TrtllmClient::connect(&endpoint, transport(2)) + .await + .expect("connect pool"); + assert_eq!(client.connection_count(), 2); + + for index in 0..4 { + let mut stream = client + .generate(pb::GenerateRequest { + request_id: format!("request-{index}"), + model: "model-source".to_string(), + input: Some(pb::generate_request::Input::TokenIds(pb::TokenIds { + ids: vec![1, 2], + })), + ..Default::default() + }) + .await + .expect("start stream"); + while stream.message().await.expect("message").is_some() {} + } + + let ports: BTreeSet<_> = server + .service + .peers + .lock() + .await + .iter() + .map(SocketAddr::port) + .collect(); + assert_eq!(ports.len(), 2); +} + +/// A server with no Control service leaves the window unknown. Registering +/// anyway would advertise capacity and then reject, with a non-migratable 4xx, +/// every request that omits `max_tokens` -- most chat traffic. Fail at startup +/// instead, naming the flag that fixes it. +#[tokio::test] +async fn start_without_a_context_length_fails() { + let service = FakeTrtllm::default(); + service.no_control.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine_with(&server.endpoint, impatient_transport(), None, AGG); + + let error = engine + .start(0) + .await + .expect_err("a worker with no window must not register"); + assert!( + error.to_string().contains("--context-length"), + "the error must name the flag that supplies a window: {error}" + ); +} + +/// An engine that answers has finished loading, so an answer carrying no +/// context length means it was started without `--max_seq_len` and will not +/// change. Fail on the first answer rather than polling to the deadline and +/// then blaming a slow load. +#[tokio::test] +async fn start_fails_at_once_when_the_server_answers_without_a_context_length() { + let service = FakeTrtllm::default(); + service.empty_model_info.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine_with(&server.endpoint, impatient_transport(), None, AGG); + + let error = engine + .start(0) + .await + .expect_err("a window is required to register"); + assert!( + error.to_string().contains("--max_seq_len"), + "the error must name the fix: {error}" + ); + assert_eq!( + server.service.model_info_calls.load(Ordering::SeqCst), + 1, + "an answer that cannot change must not be retried" + ); +} + +/// The wait exists for an engine that is not serving yet: it keeps asking +/// through transport failures and picks up the window once the engine answers. +#[tokio::test] +async fn start_waits_for_a_server_that_is_not_serving_yet() { + let service = FakeTrtllm::default(); + service.unavailable_model_info.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let ready = Arc::clone(&server.service.unavailable_model_info); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(40)).await; + ready.store(false, Ordering::SeqCst); + }); + let engine = engine_with(&server.endpoint, impatient_transport(), None, AGG); + + let config = engine.start(0).await.expect("start once the engine serves"); + assert_eq!( + config.llm.expect("llm registration").context_length, + Some(4096) + ); +} + +/// `--context-length` wins over what the engine reports. The engine is still +/// asked, so the disagreement can be logged and its output cap picked up, but +/// its answer does not decide the window. +#[tokio::test] +async fn a_configured_context_length_outranks_the_servers() { + let server = FakeServer::start(FakeTrtllm::default()).await; + let engine = engine_with(&server.endpoint, transport(1), Some(8192), AGG); + + let config = engine + .start(0) + .await + .expect("start with a configured window"); + assert_eq!( + config.llm.expect("llm registration").context_length, + Some(8192), + "the configured window must outrank the server's 4096" + ); + assert_eq!( + server.service.model_info_calls.load(Ordering::SeqCst), + 1, + "the engine is still consulted, to cross-check and to learn its output cap" + ); +} + +/// A configured window also means a server with no Control service is no longer +/// a startup problem: there is nothing left to ask it for. +#[tokio::test] +async fn a_configured_context_length_survives_a_server_without_control() { + let service = FakeTrtllm::default(); + service.no_control.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine_with(&server.endpoint, transport(1), Some(8192), AGG); + + let config = engine + .start(0) + .await + .expect("start with a configured window"); + assert_eq!( + config.llm.expect("llm registration").context_length, + Some(8192) + ); +} + +/// Shutdown has to terminate in-flight requests: the worker cannot drain if a +/// stream waits forever on an engine that will never answer. +#[tokio::test] +async fn cleanup_terminates_an_in_flight_request() { + let service = FakeTrtllm::default(); + service.hang.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = Arc::new(engine(&server.endpoint, 1)); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + let mut stream = engine + .generate(request(), GenerateContext::new(context, None)) + .await + .expect("generate"); + let first = stream.next().await.expect("first item").expect("delta"); + assert_eq!(first.token_ids, [42]); + + engine.cleanup().await.expect("cleanup"); + let terminal = tokio::time::timeout(Duration::from_secs(5), stream.next()) + .await + .expect("shutdown must terminate the stream") + .expect("terminal item") + .expect("terminal"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); +} + +/// A decode leg's cancellation deferral has to cover the dispatch itself, not +/// just the streaming loop. `generate` sends the request and then awaits +/// response headers, so a cancellation that wins that race abandons a request +/// the engine has already accepted and begun pulling KV for -- stranding the +/// prefill worker's blocks with no leg left to claim them. Shutdown still wins. +#[tokio::test] +async fn a_cancelled_decode_dispatch_is_not_abandoned_mid_flight() { + let service = FakeTrtllm::default(); + service.hang_before_stream.store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = Arc::new(engine_in_mode( + &server.endpoint, + 1, + DisaggregationMode::Decode, + )); + engine.start(0).await.expect("start"); + + let mut decode_request = request(); + decode_request.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: crate::disagg::session_to_json(fake_session()).expect("handoff"), + prompt_tokens_details: None, + }); + let context = dynamo_backend_common::testing::mock_context(); + let mut dispatch = tokio::spawn({ + let engine = Arc::clone(&engine); + let context = context.clone(); + async move { + engine + .generate(decode_request, GenerateContext::new(context, None)) + .await + } + }); + + // The fake records the request before it withholds response headers, which + // is exactly the window this test is about: the engine has the request and + // the sidecar does not know it yet. + while server.service.requests.lock().await.is_empty() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + context.stop_generating(); + assert!( + tokio::time::timeout(Duration::from_millis(200), &mut dispatch) + .await + .is_err(), + "the dispatch must outlive the client's cancellation" + ); + + engine.cleanup().await.expect("cleanup"); + let mut stream = tokio::time::timeout(Duration::from_secs(5), dispatch) + .await + .expect("shutdown must release the dispatch") + .expect("dispatch task") + .expect("generate"); + let terminal = stream + .next() + .await + .expect("a terminal item") + .expect("terminal"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); +} + +/// Argument parsing decides where each worker registers. A disaggregated leg +/// that lands on the operator-configured component instead of its role +/// component is invisible to the frontend's prefill router, and no amount of +/// request-level testing would show it. +#[test] +fn parsed_arguments_map_onto_the_worker_registration() { + let parse = |extra: &[&str]| { + let mut argv = vec![ + "dynamo-trtllm-sidecar".to_string(), + "--grpc-endpoint".to_string(), + "127.0.0.1:50051".to_string(), + "--model-path".to_string(), + "model-source".to_string(), + ]; + argv.extend(extra.iter().map(|arg| arg.to_string())); + TrtllmSidecarEngine::from_args(argv) + }; + + let (_, aggregated) = parse(&["--component", "operator-chosen"]).expect("aggregated parses"); + assert_eq!(aggregated.component, "operator-chosen"); + assert_eq!(aggregated.disaggregation_mode, AGG); + assert_eq!(aggregated.model_name, "model-source"); + assert!( + !aggregated.enable_kv_routing, + "the sidecar has no KV events" + ); + + for (mode, expected) in [ + (DisaggregationMode::Prefill, "prefill"), + (DisaggregationMode::Decode, "backend"), + ] { + let (_, config) = parse(&[ + "--component", + "operator-chosen", + "--disaggregation-mode", + &mode.to_string(), + ]) + .expect("a disaggregated leg parses"); + assert_eq!( + config.component, + mode.discovery_component(), + "{mode} must register under its own component, not the configured one" + ); + assert_eq!(config.component, expected); + assert_eq!(config.disaggregation_mode, mode); + } + + for rejected in [ + vec!["--disaggregation-mode", "encode"], + vec!["--route-to-encoder"], + ] { + assert!( + parse(&rejected).is_err(), + "{rejected:?} is not supported by this sidecar" + ); + } + assert!( + TrtllmSidecarEngine::from_args(vec![ + "dynamo-trtllm-sidecar".to_string(), + "--grpc-endpoint".to_string(), + "127.0.0.1:50051".to_string(), + "--model-path".to_string(), + " ".to_string(), + ]) + .is_err(), + "an empty model path has nothing to tokenize with" + ); +} diff --git a/lib/sidecar/vllm/src/client.rs b/lib/sidecar/vllm/src/client.rs index add4c72196d7..797a9605790b 100644 --- a/lib/sidecar/vllm/src/client.rs +++ b/lib/sidecar/vllm/src/client.rs @@ -17,7 +17,7 @@ use tonic_health::pb::health_check_response::ServingStatus; use tonic_health::pb::{HealthCheckRequest, health_client::HealthClient}; pub(crate) use dynamo_sidecar_common::v14::status_to_dynamo; -pub(crate) use dynamo_sidecar_common::{engine_shutdown, invalid_argument}; +pub(crate) use dynamo_sidecar_common::{engine_shutdown, invalid_argument, startup_deadline}; use crate::proto as pb; @@ -227,14 +227,6 @@ impl VllmClient { } } -pub(crate) fn startup_deadline(duration: Duration) -> Result { - Instant::now().checked_add(duration).ok_or_else(|| { - invalid_argument(format!( - "gRPC startup deadline {duration:?} exceeds the supported monotonic clock range" - )) - }) -} - pub(crate) fn protocol_error(message: impl Into) -> DynamoError { dynamo_sidecar_common::protocol_error("vLLM", message) } diff --git a/lib/sidecar/vllm/src/convert.rs b/lib/sidecar/vllm/src/convert.rs index 6921de6f86d6..ab38a9a7f2f0 100644 --- a/lib/sidecar/vllm/src/convert.rs +++ b/lib/sidecar/vllm/src/convert.rs @@ -12,7 +12,12 @@ use dynamo_llm::protocols::common::{preprocessed_mm_identifier, preprocessed_mm_ use serde::{Deserialize, de::DeserializeOwned}; use crate::client; -use crate::json::{json_to_struct, struct_to_json}; +use dynamo_sidecar_common::{json_to_struct_v14, struct_to_json_v14}; + +/// Payload names and peer, used only to name the source in codec errors. +const PEER: &str = "vLLM"; +const KV_TRANSFER_PARAMS: &str = "kv_transfer_params"; +const EC_TRANSFER_PARAMS: &str = "ec_transfer_params"; use crate::proto as pb; const VLLM_LOGPROB_FLOOR: f64 = -9999.0; @@ -1116,8 +1121,12 @@ fn build_kv_parameters( cache_salt: cache_salt .map(|cache_salt| format!("{DYNAMO_CACHE_SALT_PREFIX}{cache_salt}")) .unwrap_or_default(), - kv_transfer_params: kv_transfer_params.map(json_to_struct).transpose()?, - ec_transfer_params: ec_transfer_params.map(json_to_struct).transpose()?, + kv_transfer_params: kv_transfer_params + .map(|value| json_to_struct_v14(value, KV_TRANSFER_PARAMS)) + .transpose()?, + ec_transfer_params: ec_transfer_params + .map(|value| json_to_struct_v14(value, EC_TRANSFER_PARAMS)) + .transpose()?, }) } @@ -1416,7 +1425,7 @@ impl ResponseState { } let params = finish .ec_transfer_params - .map(struct_to_json) + .map(|value| struct_to_json_v14(value, PEER, EC_TRANSFER_PARAMS)) .transpose()? .and_then(|value| value.as_object().cloned()) .ok_or_else(|| { @@ -1424,7 +1433,10 @@ impl ResponseState { })?; return Ok(Some(LLMEngineOutput::encode_terminal(params))); } - mapped.disaggregated_params = finish.kv_transfer_params.map(struct_to_json).transpose()?; + mapped.disaggregated_params = finish + .kv_transfer_params + .map(|value| struct_to_json_v14(value, PEER, KV_TRANSFER_PARAMS)) + .transpose()?; if self.mode.is_prefill() && mapped.disaggregated_params.is_none() { return Err(client::protocol_error( "prefill terminal is missing kv_transfer_params", diff --git a/lib/sidecar/vllm/src/json.rs b/lib/sidecar/vllm/src/json.rs deleted file mode 100644 index 5879cca711a6..000000000000 --- a/lib/sidecar/vllm/src/json.rs +++ /dev/null @@ -1,138 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use prost_types_v14 as prost_types; - -use dynamo_backend_common::DynamoError; -use prost_types::value::Kind; - -use crate::client; - -const MAX_EXACT_INTEGER: u64 = 1_u64 << 53; - -pub(crate) fn json_to_struct(value: serde_json::Value) -> Result { - let serde_json::Value::Object(fields) = value else { - return Err(client::invalid_argument( - "kv_transfer_params must be a JSON object", - )); - }; - Ok(prost_types::Struct { - fields: fields - .into_iter() - .map(|(key, value)| Ok((key, json_to_value(value)?))) - .collect::>()?, - }) -} - -fn json_to_value(value: serde_json::Value) -> Result { - let kind = match value { - serde_json::Value::Null => Kind::NullValue(prost_types::NullValue::NullValue as i32), - serde_json::Value::Bool(value) => Kind::BoolValue(value), - serde_json::Value::String(value) => Kind::StringValue(value), - serde_json::Value::Number(value) => Kind::NumberValue(number_to_f64(&value)?), - serde_json::Value::Array(values) => Kind::ListValue(prost_types::ListValue { - values: values - .into_iter() - .map(json_to_value) - .collect::>()?, - }), - serde_json::Value::Object(values) => Kind::StructValue(prost_types::Struct { - fields: values - .into_iter() - .map(|(key, value)| Ok((key, json_to_value(value)?))) - .collect::>()?, - }), - }; - Ok(prost_types::Value { kind: Some(kind) }) -} - -fn number_to_f64(value: &serde_json::Number) -> Result { - if let Some(value) = value.as_u64() - && value > MAX_EXACT_INTEGER - { - return Err(client::invalid_argument(format!( - "kv_transfer_params integer {value} cannot be represented exactly by protobuf Struct" - ))); - } - if let Some(value) = value.as_i64() - && value.unsigned_abs() > MAX_EXACT_INTEGER - { - return Err(client::invalid_argument(format!( - "kv_transfer_params integer {value} cannot be represented exactly by protobuf Struct" - ))); - } - value.as_f64().ok_or_else(|| { - client::invalid_argument(format!( - "kv_transfer_params number {value} cannot be represented by protobuf Struct" - )) - }) -} - -pub(crate) fn struct_to_json(value: prost_types::Struct) -> Result { - Ok(serde_json::Value::Object( - value - .fields - .into_iter() - .map(|(key, value)| Ok((key, value_to_json(value)?))) - .collect::>()?, - )) -} - -fn value_to_json(value: prost_types::Value) -> Result { - match value.kind { - None | Some(Kind::NullValue(_)) => Ok(serde_json::Value::Null), - Some(Kind::BoolValue(value)) => Ok(serde_json::Value::Bool(value)), - Some(Kind::StringValue(value)) => Ok(serde_json::Value::String(value)), - Some(Kind::NumberValue(value)) => { - if !value.is_finite() { - return Err(client::protocol_error( - "kv_transfer_params contains NaN or infinity", - )); - } - let number = if value.fract() == 0.0 && value.abs() <= MAX_EXACT_INTEGER as f64 { - if value.is_sign_negative() { - serde_json::Number::from(value as i64) - } else { - serde_json::Number::from(value as u64) - } - } else { - serde_json::Number::from_f64(value).expect("finite f64 is a JSON number") - }; - Ok(serde_json::Value::Number(number)) - } - Some(Kind::ListValue(values)) => Ok(serde_json::Value::Array( - values - .values - .into_iter() - .map(value_to_json) - .collect::>()?, - )), - Some(Kind::StructValue(value)) => struct_to_json(value), - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::{json_to_struct, struct_to_json}; - - #[test] - fn nested_payload_round_trips_without_shape_changes() { - let payload = json!({ - "string": "value", - "bool": true, - "number": 42, - "null": null, - "list": [1, "two", false, {"nested": 3.5}], - }); - let encoded = json_to_struct(payload.clone()).expect("encode"); - assert_eq!(struct_to_json(encoded).expect("decode"), payload); - } - - #[test] - fn rejects_non_objects_and_inexact_integers() { - assert!(json_to_struct(json!([1, 2])).is_err()); - assert!(json_to_struct(json!({"value": 9_007_199_254_740_993_u64})).is_err()); - } -} diff --git a/lib/sidecar/vllm/src/lib.rs b/lib/sidecar/vllm/src/lib.rs index 015dfa544c08..995b2ebb7d19 100644 --- a/lib/sidecar/vllm/src/lib.rs +++ b/lib/sidecar/vllm/src/lib.rs @@ -7,7 +7,6 @@ mod args; mod client; mod convert; mod engine; -mod json; mod lora; mod model; diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index a50a89b4dc6e..64ce2b06aa0f 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -36,9 +36,9 @@ use tonic_health::ServingStatus as HealthServingStatus; use crate::client::{CONTROL_SERVICE, INFERENCE_SERVICE, VllmClient}; use crate::convert::{ResponseState, build_generate_request, normalize_response_options}; use crate::engine::VllmSidecarEngine; -use crate::json::{json_to_struct, struct_to_json}; use crate::model::DiscoveredModel; use crate::proto as pb; +use dynamo_sidecar_common::{json_to_struct_v14, struct_to_json_v14}; #[derive(Clone, Default)] struct FakeVllm { @@ -196,7 +196,7 @@ impl pb::inference_server::Inference for FakeVllm { .kv .as_ref() .and_then(|kv| kv.kv_transfer_params.clone()) - .map(struct_to_json) + .map(|value| struct_to_json_v14(value, "vLLM", "kv_transfer_params")) .transpose() .map_err(|error| Status::invalid_argument(error.to_string()))?; let is_prefill = request_kv @@ -274,7 +274,7 @@ impl pb::inference_server::Inference for FakeVllm { } } else if encoder_response { let ec = (!omit_encoder_metadata).then(|| { - json_to_struct(encoder_handoff).expect("encoder handoff") + json_to_struct_v14(encoder_handoff, "ec_transfer_params").expect("encoder handoff") }); yield encode_response(ec); } else if let Some(outputs) = sequence_outputs { @@ -286,7 +286,7 @@ impl pb::inference_server::Inference for FakeVllm { } } else { let kv = is_prefill.then(|| { - json_to_struct(handoff.clone()).expect("encode handoff") + json_to_struct_v14(handoff.clone(), "kv_transfer_params").expect("encode handoff") }); yield sequence_response(true, wants_logprobs, kv); } @@ -790,7 +790,8 @@ fn encode_response(ec_transfer_params: Option) -> pb::Gener #[test] fn encode_response_enforces_terminal_contract() { let request = epd_image_request(); - let ec_transfer_params = || json_to_struct(encoder_handoff()).expect("encoder handoff"); + let ec_transfer_params = + || json_to_struct_v14(encoder_handoff(), "ec_transfer_params").expect("encoder handoff"); let mut length = encode_response(Some(ec_transfer_params())); length @@ -2103,7 +2104,12 @@ async fn aggregated_generation_converts_request_stream_and_usage() { assert!(kv.bypass_prefix_cache); assert_eq!(kv.cache_salt, "dynamo-cache-salt:cache-salt"); assert_eq!( - struct_to_json(kv.kv_transfer_params.clone().unwrap()).unwrap(), + struct_to_json_v14( + kv.kv_transfer_params.clone().unwrap(), + "vLLM", + "kv_transfer_params" + ) + .unwrap(), json!({"connector_data": {"values": [1, true, null]}}) ); } @@ -2686,12 +2692,14 @@ async fn encoder_cache_handoff_is_opaque_for_e_pd_and_e_p_d() { assert_eq!(downstream_wire.media.len(), 2, "{topology}"); assert_eq!(downstream_wire.media[0].uuid, "image-a", "{topology}"); assert_eq!(downstream_wire.media[1].uuid, "image-b", "{topology}"); - let forwarded_ec = struct_to_json( + let forwarded_ec = struct_to_json_v14( downstream_wire .kv .as_ref() .and_then(|kv| kv.ec_transfer_params.clone()) .expect("forwarded EC metadata"), + "vLLM", + "ec_transfer_params", ) .expect("EC metadata JSON"); assert_eq!(forwarded_ec, encoder_handoff(), "{topology}"); @@ -2732,9 +2740,12 @@ async fn encoder_cache_handoff_is_opaque_for_e_pd_and_e_p_d() { assert_eq!(decode_wire.media[1].uuid, "image-b"); let decode_cache = decode_wire.kv.expect("decode cache parameters"); assert!(decode_cache.kv_transfer_params.is_some()); - let decode_ec = - struct_to_json(decode_cache.ec_transfer_params.expect("decode EC metadata")) - .expect("decode EC metadata JSON"); + let decode_ec = struct_to_json_v14( + decode_cache.ec_transfer_params.expect("decode EC metadata"), + "vLLM", + "ec_transfer_params", + ) + .expect("decode EC metadata JSON"); assert_eq!(decode_ec, encoder_handoff()); } } @@ -3603,7 +3614,12 @@ async fn prefill_decode_handoff_is_opaque_and_repeatable() { let requests = server.service.requests.lock().await; let decode_wire = requests.last().unwrap().kv.as_ref().unwrap(); - let decoded = struct_to_json(decode_wire.kv_transfer_params.clone().unwrap()).unwrap(); + let decoded = struct_to_json_v14( + decode_wire.kv_transfer_params.clone().unwrap(), + "vLLM", + "kv_transfer_params", + ) + .unwrap(); // Every field round-trips opaquely except remote_port, which the sidecar // stringifies so vLLM builds a valid NIXL side-channel URL (a protobuf // Struct number would reach the engine as `20097.0`). diff --git a/tests/serve/test_sidecar.py b/tests/serve/test_sidecar.py index 8bc92b515a35..60adc1e2d3c3 100644 --- a/tests/serve/test_sidecar.py +++ b/tests/serve/test_sidecar.py @@ -4,8 +4,10 @@ """E2E coverage for native-gRPC sidecar launch scripts.""" import dataclasses +import importlib.util import json import os +import pathlib import pytest @@ -45,6 +47,38 @@ def _sidecar_worker_gpu_env(backend: str) -> dict[str, str]: return {f"{backend.upper()}_WORKER{index + 1}_GPU": device for index in range(2)} +def _trtllm_serves_openengine() -> bool: + """Is the installed TensorRT-LLM new enough to serve `openengine.v1`? + + `lib/sidecar/trtllm` speaks OpenEngine, which the servicer provides from + release 1.3.0rc27 onward. `container/context.yaml` pins that release, so + this passes in the stock container and the case runs; it guards the + off-container and downgraded-image paths, where + `trtllm-serve --grpc-protocol openengine` does not exist and the launcher + cannot come up at all. + + Probing for the servicer module beats comparing version strings: it is the + thing the sidecar actually needs, and it tracks the pinned image by itself. + `find_spec` on the parent does not import TensorRT-LLM, so this stays cheap + and works on a host without a GPU. + """ + spec = importlib.util.find_spec("tensorrt_llm") + locations = getattr(spec, "submodule_search_locations", None) if spec else None + if not locations: + return False + return any( + (pathlib.Path(root) / "grpc" / "openengine" / "servicer.py").is_file() + for root in locations + ) + + +TRTLLM_OPENENGINE_SKIP_REASON = ( + "the installed TensorRT-LLM has no OpenEngine servicer; the sidecar needs " + "release 1.3.0rc27 or newer, which container/context.yaml pins. Seeing this " + "skip in CI means the pin moved backwards." +) + + def _disaggregated_chat_payload() -> DisaggregatedChatPayload: return DisaggregatedChatPayload( body={ @@ -110,6 +144,10 @@ def _disaggregated_chat_payload() -> DisaggregatedChatPayload: pytest.mark.gpu_1, pytest.mark.timeout(780), pytest.mark.pre_merge, + pytest.mark.skipif( + not _trtllm_serves_openengine(), + reason=TRTLLM_OPENENGINE_SKIP_REASON, + ), ], model="Qwen/Qwen3-0.6B", env={ @@ -123,6 +161,51 @@ def _disaggregated_chat_payload() -> DisaggregatedChatPayload: ], ), # Prefill/decode handoff is a critical native-sidecar path. + "trtllm_disaggregated": EngineConfig( + name="trtllm_disaggregated", + directory=trtllm_sidecar_dir, + script_name="disagg.sh", + marks=[ + pytest.mark.trtllm, + # Both engines share one GPU: the handoff is what this covers, and + # a co-resident pair keeps it on the existing 1-GPU sidecar runner. + # `requested_trtllm_kv_tokens` caps each engine's KV pool so the + # second one has memory left to load. The KV cache still moves over + # the launcher's NIXL default; note that NIXL's shared-memory + # transport needs more than a 64 MiB `/dev/shm` to bring up a + # second agent, as `deploy/disagg.yaml` also calls out. + pytest.mark.gpu_1, + pytest.mark.requested_trtllm_kv_tokens(2048), + # Two engines load serially before either sidecar can register, so + # this needs longer than the single-engine configs above. Raise the + # readiness budget with it: `pytest.mark.timeout` is only the outer + # kill timer, while EngineConfig.timeout is what the harness gives + # the health checks (engine_process.py passes it as their deadline). + # Leaving that at its 600s default would let the health check fail + # at the single-engine budget and then idle until the kill timer. + pytest.mark.timeout(1200), + pytest.mark.pre_merge, + pytest.mark.skipif( + not _trtllm_serves_openengine(), + reason=TRTLLM_OPENENGINE_SKIP_REASON, + ), + ], + model="Qwen/Qwen3-0.6B", + timeout=1000, + health_check_workers=True, + health_check_worker_count=2, + env={ + "TLLM_ALLOW_N_GREEDY_DECODING": "1", + "PYTHONUNBUFFERED": "1", + # The TensorRT-LLM release image runs as uid 0 and PRTE refuses + # to run as root unless told twice; without these the engines' + # worker spawn fails with `MPI_ERR_UNKNOWN` and neither binds its + # gRPC port. Same pair, same reason, as `deploy/disagg.yaml`. + "PRTE_ALLOW_RUN_AS_ROOT": "1", + "PRTE_ALLOW_RUN_AS_ROOT_CONFIRM": "1", + }, + request_payloads=[_disaggregated_chat_payload()], + ), "vllm_disaggregated": EngineConfig( name="vllm_disaggregated", directory=vllm_sidecar_dir, @@ -203,14 +286,18 @@ def test_serve_deployment( "DYN_NAMESPACE": f"sidecar-disagg-{generate_random_suffix()}", "MODEL": config.model, } - num_engine_ports = {"vllm": 4, "sglang": 5}[backend] + num_engine_ports = {"vllm": 4, "sglang": 5, "trtllm": 2}[backend] with reserved_ports( num_engine_ports, start_port=DynamoPortRange.SERVE.value ) as engine_ports: for index, role in enumerate(roles): prefix = f"{backend.upper()}_{role}" - engine_env[f"{prefix}_HTTP_PORT"] = str(engine_ports[index * 2]) - engine_env[f"{prefix}_GRPC_PORT"] = str(engine_ports[index * 2 + 1]) + if backend == "trtllm": + # TensorRT-LLM exposes only a gRPC listener per engine. + engine_env[f"{prefix}_GRPC_PORT"] = str(engine_ports[index]) + else: + engine_env[f"{prefix}_HTTP_PORT"] = str(engine_ports[index * 2]) + engine_env[f"{prefix}_GRPC_PORT"] = str(engine_ports[index * 2 + 1]) if backend == "vllm": engine_env[f"{prefix}_NIXL_SIDE_CHANNEL_PORT"] = str( dynamo_dynamic_ports.nixl_side_channel_ports[index]