From b903166ac56415ab2f5dfcc397fbd7f99d5a4163 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 15:34:43 -0700 Subject: [PATCH 01/20] feat(trtllm): add OpenEngine sidecar Signed-off-by: Connor Carpenter --- Cargo.lock | 66 +- Cargo.toml | 1 + examples/backends/trtllm/deploy/README.md | 52 + .../trtllm/deploy/v1beta1/openengine-agg.yaml | 97 + .../deploy/v1beta1/openengine-audio.yaml | 110 ++ .../deploy/v1beta1/openengine-disagg.yaml | 179 ++ .../deploy/v1beta1/openengine-lora.yaml | 111 ++ .../v1beta1/openengine-multimodal-agg.yaml | 101 ++ .../v1beta1/openengine-multimodal-disagg.yaml | 166 ++ .../phi-4-multimodal-instruct/audio.yaml | 46 + .../engine_configs/qwen3/decode-lora.yaml | 42 + .../trtllm/engine_configs/qwen3/lora.yaml | 40 + .../engine_configs/qwen3/prefill-lora.yaml | 42 + examples/backends/trtllm/launch/OPENENGINE.md | 42 + .../backends/trtllm/launch/openengine_agg.sh | 56 + .../trtllm/launch/openengine_audio.sh | 39 + .../trtllm/launch/openengine_common.sh | 95 + .../trtllm/launch/openengine_disagg.sh | 74 + .../backends/trtllm/launch/openengine_lora.sh | 27 + .../trtllm/launch/openengine_multimodal.sh | 26 + .../launch/openengine_multimodal_disagg.sh | 31 + lib/backend-common/src/adapter.rs | 93 +- lib/backend-common/src/engine.rs | 39 +- lib/backend-common/src/lib.rs | 4 +- lib/backend-common/src/validate.rs | 89 +- lib/backend-common/src/worker.rs | 242 ++- lib/openengine-sidecar/Cargo.toml | 46 + lib/openengine-sidecar/README.md | 17 + lib/openengine-sidecar/build.rs | 54 + lib/openengine-sidecar/src/args.rs | 89 + lib/openengine-sidecar/src/client.rs | 313 ++++ lib/openengine-sidecar/src/convert.rs | 652 +++++++ lib/openengine-sidecar/src/engine.rs | 1583 +++++++++++++++++ lib/openengine-sidecar/src/kv.rs | 504 ++++++ lib/openengine-sidecar/src/lib.rs | 23 + lib/openengine-sidecar/src/main.rs | 9 + lib/openengine-sidecar/src/tests.rs | 1345 ++++++++++++++ 37 files changed, 6459 insertions(+), 86 deletions(-) create mode 100644 examples/backends/trtllm/deploy/v1beta1/openengine-agg.yaml create mode 100644 examples/backends/trtllm/deploy/v1beta1/openengine-audio.yaml create mode 100644 examples/backends/trtllm/deploy/v1beta1/openengine-disagg.yaml create mode 100644 examples/backends/trtllm/deploy/v1beta1/openengine-lora.yaml create mode 100644 examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-agg.yaml create mode 100644 examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-disagg.yaml create mode 100644 examples/backends/trtllm/engine_configs/phi-4-multimodal-instruct/audio.yaml create mode 100644 examples/backends/trtllm/engine_configs/qwen3/decode-lora.yaml create mode 100644 examples/backends/trtllm/engine_configs/qwen3/lora.yaml create mode 100644 examples/backends/trtllm/engine_configs/qwen3/prefill-lora.yaml create mode 100644 examples/backends/trtllm/launch/OPENENGINE.md create mode 100755 examples/backends/trtllm/launch/openengine_agg.sh create mode 100755 examples/backends/trtllm/launch/openengine_audio.sh create mode 100755 examples/backends/trtllm/launch/openengine_common.sh create mode 100755 examples/backends/trtllm/launch/openengine_disagg.sh create mode 100755 examples/backends/trtllm/launch/openengine_lora.sh create mode 100755 examples/backends/trtllm/launch/openengine_multimodal.sh create mode 100755 examples/backends/trtllm/launch/openengine_multimodal_disagg.sh create mode 100644 lib/openengine-sidecar/Cargo.toml create mode 100644 lib/openengine-sidecar/README.md create mode 100644 lib/openengine-sidecar/build.rs create mode 100644 lib/openengine-sidecar/src/args.rs create mode 100644 lib/openengine-sidecar/src/client.rs create mode 100644 lib/openengine-sidecar/src/convert.rs create mode 100644 lib/openengine-sidecar/src/engine.rs create mode 100644 lib/openengine-sidecar/src/kv.rs create mode 100644 lib/openengine-sidecar/src/lib.rs create mode 100644 lib/openengine-sidecar/src/main.rs create mode 100644 lib/openengine-sidecar/src/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 92fe2dfea309..264d39967d93 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2717,6 +2717,32 @@ dependencies = [ "uuid", ] +[[package]] +name = "dynamo-openengine-sidecar" +version = "1.3.0" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "base64 0.22.1", + "clap", + "dynamo-backend-common", + "dynamo-kv-router", + "dynamo-llm", + "dynamo-protocols", + "dynamo-runtime", + "futures", + "openengine-proto", + "parking_lot", + "prost-types 0.14.4", + "serde_json", + "tokio", + "tokio-stream", + "tokio-util", + "tonic 0.14.6", + "tracing", +] + [[package]] name = "dynamo-parsers" version = "4.1.1" @@ -3152,7 +3178,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8acfe553027cd07fc5fafa81a84f19a7a87eaffaccd2162b6db05e8d6ce98084" dependencies = [ "http 1.4.1", - "prost 0.14.3", + "prost 0.14.4", "tokio", "tokio-stream", "tonic 0.14.6", @@ -6236,6 +6262,16 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "openengine-proto" +version = "0.2.0" +dependencies = [ + "prost 0.14.4", + "prost-types 0.14.4", + "tonic 0.14.6", + "tonic-prost", +] + [[package]] name = "openssl-probe" version = "0.1.6" @@ -6298,7 +6334,7 @@ dependencies = [ "opentelemetry-http", "opentelemetry-proto", "opentelemetry_sdk", - "prost 0.14.3", + "prost 0.14.4", "reqwest 0.12.28", "thiserror 2.0.18", "tokio", @@ -6314,7 +6350,7 @@ checksum = "a7175df06de5eaee9909d4805a3d07e28bb752c34cab57fa9cff549da596b30f" dependencies = [ "opentelemetry", "opentelemetry_sdk", - "prost 0.14.3", + "prost 0.14.4", "tonic 0.14.6", "tonic-prost", ] @@ -6965,12 +7001,12 @@ dependencies = [ [[package]] name = "prost" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" dependencies = [ "bytes", - "prost-derive 0.14.3", + "prost-derive 0.14.4", ] [[package]] @@ -7005,8 +7041,8 @@ dependencies = [ "multimap", "petgraph 0.8.3", "prettyplease", - "prost 0.14.3", - "prost-types 0.14.3", + "prost 0.14.4", + "prost-types 0.14.4", "pulldown-cmark", "pulldown-cmark-to-cmark", "regex", @@ -7029,9 +7065,9 @@ dependencies = [ [[package]] name = "prost-derive" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", "itertools 0.14.0", @@ -7051,11 +7087,11 @@ dependencies = [ [[package]] name = "prost-types" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" dependencies = [ - "prost 0.14.3", + "prost 0.14.4", ] [[package]] @@ -9583,7 +9619,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ "bytes", - "prost 0.14.3", + "prost 0.14.4", "tonic 0.14.6", ] @@ -9596,7 +9632,7 @@ dependencies = [ "prettyplease", "proc-macro2", "prost-build 0.14.3", - "prost-types 0.14.3", + "prost-types 0.14.4", "quote", "syn 2.0.117", "tempfile", diff --git a/Cargo.toml b/Cargo.toml index 922bb3474bb5..58126e863d7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ members = [ "lib/backend-common/examples/mocker", "lib/vllm-rs-backend", "lib/sglang-sidecar", + "lib/openengine-sidecar", "lib/bindings/c", "lib/bindings/python/codegen", "deploy/inference-gateway/ext-proc", diff --git a/examples/backends/trtllm/deploy/README.md b/examples/backends/trtllm/deploy/README.md index 8582a69b1580..d1accb8035c4 100644 --- a/examples/backends/trtllm/deploy/README.md +++ b/examples/backends/trtllm/deploy/README.md @@ -1,3 +1,8 @@ + + # TensorRT-LLM Kubernetes Deployment Configurations This directory contains Kubernetes Custom Resource Definition (CRD) templates for deploying TensorRT-LLM inference graphs using the **DynamoGraphDeployment** resource. @@ -68,6 +73,53 @@ TensorRT-LLM worker. This example uses the conservative engine configuration in `experimental.checkpoint.enabled: true` and `startupPolicy: WaitForCheckpoint` +### 8. **OpenEngine Sidecar Examples** (`v1beta1/openengine-*.yaml`) + +The OpenEngine examples keep TRT-LLM in the GPU container and run +`dynamo-openengine-sidecar` as the CPU-only `main` container in the same pod. +Loopback OpenEngine gRPC connects the pair; TRT-LLM's HTTP listener remains +available and uses the same `LLM` instance. + +- `openengine-agg.yaml`: aggregate text +- `openengine-disagg.yaml`: context-first 1P1D text and LoRA +- `openengine-multimodal-agg.yaml`: aggregate image/video +- `openengine-multimodal-disagg.yaml`: image/video context-first 1P1D +- `openengine-audio.yaml`: aggregate Phi-4 audio +- `openengine-lora.yaml`: aggregate lazy LoRA lifecycle + +Replace both placeholder images with builds containing the sidecar and the +local TRT-LLM fork. The TRT-LLM image must also contain this repository's +`/workspace/examples/backends/trtllm/engine_configs` directory. Every engine is +started with the PyTorch backend and immutable OpenEngine schema release +`cea19cb06acf03c911b84d5c147e519b60dd92a6`. + +LoRA manifests expect a PVC named `trtllm-openengine-lora-cache`. For 1P1D it +must support `ReadWriteMany`, because context and generation are different pods +and both mount `/lora-cache`. Create it before the DGD, for example: + +```yaml +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: trtllm-openengine-lora-cache +spec: + accessModes: [ReadWriteMany] + resources: + requests: + storage: 20Gi +``` + +OpenEngine workers publish model cards with no frontend media decoder. This is +the URL/data passthrough mode: ordered media stays encoded until TRT-LLM fetches +and decodes it. Do not add a frontend media decoder to these deployments; +decoded/RDMA media is intentionally rejected by the sidecar. + +The Phi-4 audio manifest additionally expects a PVC named +`trtllm-openengine-phi4-model`. Populate its root with a raw git clone of +`microsoft/Phi-4-multimodal-instruct` before creating the DGD. The mounted +directory must contain the bundled `speech-lora/` subdirectory; TRT-LLM's +Phi-4-MM loader does not support the Hugging Face snapshot layout. + ## CRD Structure All templates use the **DynamoGraphDeployment** CRD: diff --git a/examples/backends/trtllm/deploy/v1beta1/openengine-agg.yaml b/examples/backends/trtllm/deploy/v1beta1/openengine-agg.yaml new file mode 100644 index 000000000000..8d22974ff973 --- /dev/null +++ b/examples/backends/trtllm/deploy/v1beta1/openengine-agg.yaml @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: trtllm-openengine-agg +spec: + backendFramework: trtllm + components: + - name: Frontend + podTemplate: + spec: + containers: + - args: + - python3 -m dynamo.frontend --router-mode round-robin --http-port 8000 + command: [/bin/sh, -c] + env: + - name: DYN_NAMESPACE + value: dynamo + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + replicas: 1 + type: frontend + - name: OpenEngineWorker + podTemplate: + spec: + containers: + - args: + - --openengine-endpoint + - http://127.0.0.1:50051 + - --expected-engine + - tensorrt_llm + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8081" + - name: DYN_LORA_ENABLED + value: "false" + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + - args: + - Qwen/Qwen3-0.6B + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8001" + - --config + - /workspace/examples/backends/trtllm/engine_configs/qwen3/agg.yaml + - --openengine-host + - 127.0.0.1 + - --openengine-port + - "50051" + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /model-cache + name: model-cache + volumes: + - emptyDir: {} + name: model-cache + replicas: 1 + type: worker diff --git a/examples/backends/trtllm/deploy/v1beta1/openengine-audio.yaml b/examples/backends/trtllm/deploy/v1beta1/openengine-audio.yaml new file mode 100644 index 000000000000..15a0e12ee9e4 --- /dev/null +++ b/examples/backends/trtllm/deploy/v1beta1/openengine-audio.yaml @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Phi-4 multimodal audio is aggregate-only for this milestone. The engine config +# enables the model's bundled speech-lora and its exact TRT/HF module mapping. +# The sidecar leaves audio URLs/data URIs encoded for TRT-LLM to fetch/decode. +# Prerequisite: populate a PVC named `trtllm-openengine-phi4-model` with a raw +# git clone of microsoft/Phi-4-multimodal-instruct at the claim root. TRT-LLM's +# Phi-4-MM loader does not support the Hugging Face snapshot layout. + +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: trtllm-openengine-audio +spec: + backendFramework: trtllm + components: + - name: Frontend + podTemplate: + spec: + containers: + - args: + - python3 -m dynamo.frontend --router-mode round-robin --http-port 8000 + command: [/bin/sh, -c] + env: + - name: DYN_NAMESPACE + value: dynamo + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + replicas: 1 + type: frontend + - name: OpenEngineWorker + podTemplate: + spec: + containers: + - args: + - --openengine-endpoint + - http://127.0.0.1:50051 + - --expected-engine + - tensorrt_llm + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8081" + - name: DYN_LORA_ENABLED + value: "false" + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + - args: + - /models/phi-4-multimodal-instruct + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8001" + - --config + - /workspace/examples/backends/trtllm/engine_configs/phi-4-multimodal-instruct/audio.yaml + - --openengine-host + - 127.0.0.1 + - --openengine-port + - "50051" + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /models/phi-4-multimodal-instruct + name: phi4-model + readOnly: true + - mountPath: /model-cache + name: model-cache + volumes: + - name: phi4-model + persistentVolumeClaim: + claimName: trtllm-openengine-phi4-model + - emptyDir: {} + name: model-cache + replicas: 1 + type: worker diff --git a/examples/backends/trtllm/deploy/v1beta1/openengine-disagg.yaml b/examples/backends/trtllm/deploy/v1beta1/openengine-disagg.yaml new file mode 100644 index 000000000000..a9d8a2ddef08 --- /dev/null +++ b/examples/backends/trtllm/deploy/v1beta1/openengine-disagg.yaml @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Prerequisite: create a ReadWriteMany PVC named +# `trtllm-openengine-lora-cache` in the deployment namespace. Both pods mount +# the same claim so an adapter downloaded by either CPU sidecar is immediately +# visible to both TRT-LLM roles. + +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: trtllm-openengine-disagg +spec: + backendFramework: trtllm + components: + - name: Frontend + podTemplate: + spec: + containers: + - args: + - python3 -m dynamo.frontend --router-mode kv --http-port 8000 + command: [/bin/sh, -c] + env: + - name: DYN_NAMESPACE + value: dynamo + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + replicas: 1 + type: frontend + - name: ContextWorker + podTemplate: + spec: + containers: + - args: [--openengine-endpoint, http://127.0.0.1:50051, --expected-engine, tensorrt_llm] + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8081" + - name: DYN_LORA_ENABLED + value: "true" + - name: DYN_LORA_PATH + value: /lora-cache + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + volumeMounts: + - mountPath: /lora-cache + name: lora-cache + - args: + - Qwen/Qwen3-0.6B + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8001" + - --config + - /workspace/examples/backends/trtllm/engine_configs/qwen3/prefill-lora.yaml + - --server_role + - CONTEXT + - --openengine-port + - "50051" + - --openengine-host + - 127.0.0.1 + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /lora-cache + name: lora-cache + - mountPath: /model-cache + name: model-cache + volumes: + - name: lora-cache + persistentVolumeClaim: + claimName: trtllm-openengine-lora-cache + - emptyDir: {} + name: model-cache + replicas: 1 + type: prefill + - name: GenerationWorker + podTemplate: + spec: + containers: + - args: [--openengine-endpoint, http://127.0.0.1:50052, --expected-engine, tensorrt_llm] + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8082" + - name: DYN_LORA_ENABLED + value: "true" + - name: DYN_LORA_PATH + value: /lora-cache + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + volumeMounts: + - mountPath: /lora-cache + name: lora-cache + - args: + - Qwen/Qwen3-0.6B + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8002" + - --config + - /workspace/examples/backends/trtllm/engine_configs/qwen3/decode-lora.yaml + - --server_role + - GENERATION + - --openengine-port + - "50052" + - --openengine-host + - 127.0.0.1 + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /lora-cache + name: lora-cache + - mountPath: /model-cache + name: model-cache + volumes: + - name: lora-cache + persistentVolumeClaim: + claimName: trtllm-openengine-lora-cache + - emptyDir: {} + name: model-cache + replicas: 1 + type: decode diff --git a/examples/backends/trtllm/deploy/v1beta1/openengine-lora.yaml b/examples/backends/trtllm/deploy/v1beta1/openengine-lora.yaml new file mode 100644 index 000000000000..d9edcbe14c9e --- /dev/null +++ b/examples/backends/trtllm/deploy/v1beta1/openengine-lora.yaml @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Prerequisite: create a PVC named `trtllm-openengine-lora-cache` in the +# deployment namespace. The CPU sidecar downloads adapters into the claim and +# the GPU engine consumes the exact same path. + +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: trtllm-openengine-lora +spec: + backendFramework: trtllm + components: + - name: Frontend + podTemplate: + spec: + containers: + - args: + - python3 -m dynamo.frontend --router-mode round-robin --http-port 8000 + command: [/bin/sh, -c] + env: + - name: DYN_NAMESPACE + value: dynamo + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + replicas: 1 + type: frontend + - name: OpenEngineWorker + podTemplate: + spec: + containers: + - args: + - --openengine-endpoint + - http://127.0.0.1:50051 + - --expected-engine + - tensorrt_llm + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8081" + - name: DYN_LORA_ENABLED + value: "true" + - name: DYN_LORA_PATH + value: /lora-cache + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + volumeMounts: + - mountPath: /lora-cache + name: lora-cache + - args: + - Qwen/Qwen3-0.6B + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8001" + - --config + - /workspace/examples/backends/trtllm/engine_configs/qwen3/lora.yaml + - --openengine-host + - 127.0.0.1 + - --openengine-port + - "50051" + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /lora-cache + name: lora-cache + - mountPath: /model-cache + name: model-cache + volumes: + - name: lora-cache + persistentVolumeClaim: + claimName: trtllm-openengine-lora-cache + - emptyDir: {} + name: model-cache + replicas: 1 + type: worker diff --git a/examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-agg.yaml b/examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-agg.yaml new file mode 100644 index 000000000000..3250baabd55a --- /dev/null +++ b/examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-agg.yaml @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# The sidecar registers a model card without a media decoder. The frontend +# therefore preserves ordered image/video http(s) URLs and data URIs for the +# TRT-LLM process in the same pod to fetch and decode. + +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: trtllm-openengine-multimodal-agg +spec: + backendFramework: trtllm + components: + - name: Frontend + podTemplate: + spec: + containers: + - args: + - python3 -m dynamo.frontend --router-mode round-robin --http-port 8000 + command: [/bin/sh, -c] + env: + - name: DYN_NAMESPACE + value: dynamo + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + replicas: 1 + type: frontend + - name: OpenEngineWorker + podTemplate: + spec: + containers: + - args: + - --openengine-endpoint + - http://127.0.0.1:50051 + - --expected-engine + - tensorrt_llm + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8081" + - name: DYN_LORA_ENABLED + value: "false" + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + - args: + - Qwen/Qwen3-VL-2B-Instruct + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8001" + - --config + - /workspace/examples/backends/trtllm/engine_configs/qwen3-vl-2b-instruct/agg.yaml + - --openengine-host + - 127.0.0.1 + - --openengine-port + - "50051" + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /model-cache + name: model-cache + volumes: + - emptyDir: {} + name: model-cache + replicas: 1 + type: worker diff --git a/examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-disagg.yaml b/examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-disagg.yaml new file mode 100644 index 000000000000..155eaddf74e1 --- /dev/null +++ b/examples/backends/trtllm/deploy/v1beta1/openengine-multimodal-disagg.yaml @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Image/video P/D deliberately sends the ordered raw media to both roles. The +# decode engine needs it to reconstruct local mRoPE metadata while consuming the +# KV handoff produced by the context engine. + +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: trtllm-openengine-multimodal-disagg +spec: + backendFramework: trtllm + components: + - name: Frontend + podTemplate: + spec: + containers: + - args: + - python3 -m dynamo.frontend --router-mode kv --http-port 8000 + command: [/bin/sh, -c] + env: + - name: DYN_NAMESPACE + value: dynamo + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + replicas: 1 + type: frontend + - name: ContextWorker + podTemplate: + spec: + containers: + - args: + - --openengine-endpoint + - http://127.0.0.1:50051 + - --expected-engine + - tensorrt_llm + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8081" + - name: DYN_LORA_ENABLED + value: "false" + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + - args: + - Qwen/Qwen3-VL-2B-Instruct + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8001" + - --config + - /workspace/examples/backends/trtllm/engine_configs/qwen3-vl-2b-instruct/prefill.yaml + - --server_role + - CONTEXT + - --openengine-host + - 127.0.0.1 + - --openengine-port + - "50051" + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /model-cache + name: model-cache + volumes: + - emptyDir: {} + name: model-cache + replicas: 1 + type: prefill + - name: GenerationWorker + podTemplate: + spec: + containers: + - args: + - --openengine-endpoint + - http://127.0.0.1:50052 + - --expected-engine + - tensorrt_llm + command: [dynamo-openengine-sidecar] + env: + - name: DYN_NAMESPACE + value: dynamo + - name: DYN_SYSTEM_PORT + value: "8082" + - name: DYN_LORA_ENABLED + value: "false" + image: my-registry/dynamo-openengine-sidecar:my-tag + name: main + resources: + limits: + cpu: "2" + memory: 4Gi + requests: + cpu: "1" + memory: 2Gi + - args: + - Qwen/Qwen3-VL-2B-Instruct + - --backend + - pytorch + - --host + - 127.0.0.1 + - --port + - "8002" + - --config + - /workspace/examples/backends/trtllm/engine_configs/qwen3-vl-2b-instruct/decode.yaml + - --server_role + - GENERATION + - --openengine-host + - 127.0.0.1 + - --openengine-port + - "50052" + command: [trtllm-serve] + env: + - name: HF_HOME + value: /model-cache + - name: OPENENGINE_SCHEMA_RELEASE + value: cea19cb06acf03c911b84d5c147e519b60dd92a6 + image: my-registry/tensorrtllm-openengine:my-tag + name: engine + resources: + limits: + nvidia.com/gpu: "1" + requests: + nvidia.com/gpu: "1" + volumeMounts: + - mountPath: /model-cache + name: model-cache + volumes: + - emptyDir: {} + name: model-cache + replicas: 1 + type: decode diff --git a/examples/backends/trtllm/engine_configs/phi-4-multimodal-instruct/audio.yaml b/examples/backends/trtllm/engine_configs/phi-4-multimodal-instruct/audio.yaml new file mode 100644 index 000000000000..5be3d76be576 --- /dev/null +++ b/examples/backends/trtllm/engine_configs/phi-4-multimodal-instruct/audio.yaml @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Phi-4 multimodal implements audio through the model's bundled speech-lora. +# These mappings match Phi4MMForCausalLM.lora_config() in the local TRT-LLM +# fork. Keep the model's speech-lora directory beside the base checkpoint. +tensor_parallel_size: 1 +enable_attention_dp: false +max_num_tokens: 1024 +max_batch_size: 4 +max_seq_len: 4096 +trust_remote_code: true +backend: pytorch +enable_chunked_prefill: true + +kv_cache_config: + free_gpu_memory_fraction: 0.60 + +lora_config: + lora_ckpt_source: hf + swap_gate_up_proj_lora_b_weight: false + max_loras: 2 + max_cpu_loras: 2 + max_lora_rank: 320 + lora_target_modules: + - attn_qkv + - attn_dense + - mlp_gate_up + - mlp_4h_to_h + trtllm_modules_to_hf_modules: + attn_qkv: qkv_proj + attn_dense: o_proj + mlp_gate_up: gate_up_proj + mlp_4h_to_h: down_proj diff --git a/examples/backends/trtllm/engine_configs/qwen3/decode-lora.yaml b/examples/backends/trtllm/engine_configs/qwen3/decode-lora.yaml new file mode 100644 index 000000000000..f0570e2c99e6 --- /dev/null +++ b/examples/backends/trtllm/engine_configs/qwen3/decode-lora.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +tensor_parallel_size: 1 +moe_expert_parallel_size: 1 +enable_attention_dp: false +max_num_tokens: 8192 +trust_remote_code: true +backend: pytorch +enable_chunked_prefill: true +disable_overlap_scheduler: false +cuda_graph_config: + max_batch_size: 16 +kv_cache_config: + free_gpu_memory_fraction: 0.85 +cache_transceiver_config: + backend: DEFAULT +lora_config: + lora_ckpt_source: hf + max_lora_rank: 64 + max_loras: 2 + max_cpu_loras: 2 + lora_target_modules: + - attn_q + - attn_k + - attn_v + - attn_dense + - mlp_h_to_4h + - mlp_gate + - mlp_4h_to_h diff --git a/examples/backends/trtllm/engine_configs/qwen3/lora.yaml b/examples/backends/trtllm/engine_configs/qwen3/lora.yaml new file mode 100644 index 000000000000..dcc1c92d6cfd --- /dev/null +++ b/examples/backends/trtllm/engine_configs/qwen3/lora.yaml @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +tensor_parallel_size: 1 +moe_expert_parallel_size: 1 +enable_attention_dp: false +max_num_tokens: 8192 +max_batch_size: 16 +trust_remote_code: true +backend: pytorch +enable_chunked_prefill: true +kv_cache_config: + free_gpu_memory_fraction: 0.85 +cuda_graph_config: + max_batch_size: 16 +lora_config: + lora_ckpt_source: hf + max_lora_rank: 64 + max_loras: 2 + max_cpu_loras: 2 + lora_target_modules: + - attn_q + - attn_k + - attn_v + - attn_dense + - mlp_h_to_4h + - mlp_gate + - mlp_4h_to_h diff --git a/examples/backends/trtllm/engine_configs/qwen3/prefill-lora.yaml b/examples/backends/trtllm/engine_configs/qwen3/prefill-lora.yaml new file mode 100644 index 000000000000..c1c3fd23eb1d --- /dev/null +++ b/examples/backends/trtllm/engine_configs/qwen3/prefill-lora.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +tensor_parallel_size: 1 +moe_expert_parallel_size: 1 +enable_attention_dp: false +max_num_tokens: 8192 +trust_remote_code: true +backend: pytorch +enable_chunked_prefill: true +disable_overlap_scheduler: true +cuda_graph_config: + max_batch_size: 16 +kv_cache_config: + free_gpu_memory_fraction: 0.85 +cache_transceiver_config: + backend: DEFAULT +lora_config: + lora_ckpt_source: hf + max_lora_rank: 64 + max_loras: 2 + max_cpu_loras: 2 + lora_target_modules: + - attn_q + - attn_k + - attn_v + - attn_dense + - mlp_h_to_4h + - mlp_gate + - mlp_4h_to_h diff --git a/examples/backends/trtllm/launch/OPENENGINE.md b/examples/backends/trtllm/launch/OPENENGINE.md new file mode 100644 index 000000000000..8465db07c4e9 --- /dev/null +++ b/examples/backends/trtllm/launch/OPENENGINE.md @@ -0,0 +1,42 @@ + + +# TRT-LLM OpenEngine sidecar launches + +These scripts run the TRT-LLM HTTP server and its optional OpenEngine sibling +against the same `LLM`, then register a CPU-only `dynamo-openengine-sidecar` +worker with Dynamo. + +The examples use the OpenEngine schema release +`cea19cb06acf03c911b84d5c147e519b60dd92a6`. Install the matching local +OpenEngine Python package in the TRT-LLM environment and build the Rust sidecar +against the same local checkout before launching them. + +| Script | Model/default | GPUs | Coverage | +| --- | --- | ---: | --- | +| `openengine_agg.sh` | `Qwen/Qwen3-0.6B` | 1 | Aggregate text | +| `openengine_disagg.sh` | `Qwen/Qwen3-0.6B` | 2 | Context-first 1P1D text and LoRA | +| `openengine_multimodal.sh` | `Qwen/Qwen3-VL-2B-Instruct` | 1 | Aggregate image and video | +| `openengine_multimodal_disagg.sh` | `Qwen/Qwen3-VL-2B-Instruct` | 2 | Image/video context-first 1P1D | +| `openengine_audio.sh` | required raw Phi-4-MM clone | 1 | Aggregate audio with bundled `speech-lora` | +| `openengine_lora.sh` | `Qwen/Qwen3-0.6B` | 1 | Aggregate lazy LoRA lifecycle | + +The one-node scripts reserve frontend HTTP `8000`, TRT-LLM HTTP `8001`, +OpenEngine gRPC `50051`, and the Dynamo worker system port `8081`. The 1P1D +scripts add TRT-LLM HTTP `8002`, OpenEngine gRPC `50052`, and system port +`8082`. Every value has an environment-variable override in the script. + +The sidecar intentionally advertises no frontend media decoder. As a result, +the frontend keeps ordered `http(s)` URLs and `data:` URIs intact and the local +TRT-LLM process owns media fetch/decode. Decoded/RDMA media is rejected rather +than silently copied. + +For LoRA, set `DYN_LORA_PATH` to a directory visible at the same absolute path +to the sidecar and TRT-LLM. The provided local scripts share one filesystem; +the Kubernetes examples use a shared PVC. The Qwen configs bound the GPU and +CPU adapter caches to two entries. Phi-4's checkpoint must retain its bundled +`speech-lora/` directory. The TRT-LLM Phi-4-MM loader does not support the +Hugging Face snapshot layout, so `openengine_audio.sh` requires `MODEL_PATH` to +point to a raw git clone of the model repository. diff --git a/examples/backends/trtllm/launch/openengine_agg.sh b/examples/backends/trtllm/launch/openengine_agg.sh new file mode 100755 index 000000000000..b30237cd48e8 --- /dev/null +++ b/examples/backends/trtllm/launch/openengine_agg.sh @@ -0,0 +1,56 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +set -euo pipefail +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" +source "$SCRIPT_DIR/openengine_common.sh" +trap dynamo_exit_trap EXIT + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-0.6B} +HTTP_PORT=${DYN_HTTP_PORT:-${FRONTEND_PORT:-8000}} +export DYN_HTTP_PORT="$HTTP_PORT" +TRTLLM_HTTP_PORT=${TRTLLM_HTTP_PORT:-8001} +OPENENGINE_PORT=${OPENENGINE_PORT:-50051} +SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} +export DYNAMO_HOME=${DYNAMO_HOME:-"$(readlink -f "$SCRIPT_DIR/../../../..")"} +TRTLLM_CONFIG=${TRTLLM_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/qwen3/agg.yaml} +OPENENGINE_LAUNCH_MODE=${OPENENGINE_LAUNCH_MODE:-text} +export DYN_LORA_ENABLED=${DYN_LORA_ENABLED:-false} +export DYN_LORA_PATH=${DYN_LORA_PATH:-/tmp/dynamo_loras} +readonly OPENENGINE_SCHEMA_RELEASE=cea19cb06acf03c911b84d5c147e519b60dd92a6 +export OPENENGINE_SCHEMA_RELEASE + +warn_trtllm_serve_profile_overrides "$TRTLLM_CONFIG" +print_openengine_aggregate_banner "$OPENENGINE_LAUNCH_MODE" "$MODEL_PATH" \ + "$HTTP_PORT" "$TRTLLM_HTTP_PORT" "$OPENENGINE_PORT" "$SYSTEM_PORT" + +# The sidecar advertises no frontend media decoder. This intentionally keeps +# http(s) URLs and data URIs intact until TRT-LLM fetches/decodes them. +python3 -m dynamo.frontend --router-mode round-robin & +trtllm-serve "$MODEL_PATH" \ + --backend pytorch \ + --config "$TRTLLM_CONFIG" \ + --host 127.0.0.1 \ + --port "$TRTLLM_HTTP_PORT" \ + --openengine-host 127.0.0.1 \ + --openengine-port "$OPENENGINE_PORT" & +DYN_SYSTEM_PORT="$SYSTEM_PORT" dynamo-openengine-sidecar \ + --openengine-endpoint "http://127.0.0.1:$OPENENGINE_PORT" \ + --expected-engine tensorrt_llm & + +wait_any_exit diff --git a/examples/backends/trtllm/launch/openengine_audio.sh b/examples/backends/trtllm/launch/openengine_audio.sh new file mode 100755 index 000000000000..dce4ec3efefa --- /dev/null +++ b/examples/backends/trtllm/launch/openengine_audio.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Phi-4 multimodal audio requires TRT-LLM's built-in speech-lora. The engine +# config enables the exact LoRA module mapping and cache bounds for that model. +# Audio P/D is deliberately not advertised by the server in this milestone. +set -euo pipefail +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" +export DYNAMO_HOME=${DYNAMO_HOME:-"$(readlink -f "$SCRIPT_DIR/../../../..")"} + +if [[ -z ${MODEL_PATH:-} ]]; then + echo "MODEL_PATH must point to a raw git clone of microsoft/Phi-4-multimodal-instruct." >&2 + echo "The Hugging Face snapshot layout is not supported by TRT-LLM's Phi-4-MM loader." >&2 + exit 2 +fi +if [[ ! -d "$MODEL_PATH/speech-lora" ]]; then + echo "MODEL_PATH must contain the bundled speech-lora/ directory: $MODEL_PATH" >&2 + exit 2 +fi +export MODEL_PATH +export TRTLLM_CONFIG=${TRTLLM_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/phi-4-multimodal-instruct/audio.yaml} +export DYN_LORA_ENABLED=${DYN_LORA_ENABLED:-false} +export OPENENGINE_LAUNCH_MODE=audio +exec "$SCRIPT_DIR/openengine_agg.sh" "$@" diff --git a/examples/backends/trtllm/launch/openengine_common.sh b/examples/backends/trtllm/launch/openengine_common.sh new file mode 100755 index 000000000000..309858168bd1 --- /dev/null +++ b/examples/backends/trtllm/launch/openengine_common.sh @@ -0,0 +1,95 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Shared helpers for the direct trtllm-serve OpenEngine launch examples. + +warn_trtllm_serve_profile_overrides() { + local configs="$1" + if [[ -z ${_PROFILE_OVERRIDE_TRTLLM_MAX_TOTAL_TOKENS:-} && + -z ${_PROFILE_OVERRIDE_TRTLLM_MAX_GPU_TOTAL_BYTES:-} ]]; then + return + fi + + # Narrow exception: build_trtllm_override_args_with_mem emits JSON for the + # in-process Dynamo worker's --override-engine-args option. trtllm-serve has + # no equivalent JSON merge option, and mapping max_tokens/max_gpu_total_bytes + # to its superficially similar CLI flags would change their semantics. Keep + # the direct-server examples honest: profile KV limits belong in the YAML. + echo "WARNING: TRT-LLM profile overrides are not forwarded to trtllm-serve." >&2 + echo "Set kv_cache_config.max_tokens or max_gpu_total_bytes in: $configs" >&2 +} + +print_openengine_aggregate_banner() { + local mode="$1" model="$2" http_port="$3" + local trtllm_http_port="$4" openengine_port="$5" system_port="$6" + local details=( + "TRT-LLM HTTP: http://127.0.0.1:$trtllm_http_port" + "OpenEngine: http://127.0.0.1:$openengine_port" + "System API: http://127.0.0.1:$system_port" + ) + case "$mode" in + audio) + print_launch_banner --no-curl \ + "Launching TRT-LLM OpenEngine Audio Serving (1 GPU)" \ + "$model" "$http_port" "${details[@]}" \ + "Media: audio URL/data-URI passthrough" + ;; + lora) + print_launch_banner \ + "Launching TRT-LLM OpenEngine Aggregated Serving + LoRA (1 GPU)" \ + "$model" "$http_port" "${details[@]}" + ;; + multimodal) + print_launch_banner --multimodal \ + "Launching TRT-LLM OpenEngine Multimodal Serving (1 GPU)" \ + "$model" "$http_port" "${details[@]}" + ;; + text) + print_launch_banner \ + "Launching TRT-LLM OpenEngine Aggregated Serving (1 GPU)" \ + "$model" "$http_port" "${details[@]}" + ;; + *) + echo "Unknown OpenEngine aggregate launch mode: $mode" >&2 + return 2 + ;; + esac +} + +print_openengine_disagg_banner() { + local mode="$1" model="$2" http_port="$3" + local prefill_port="$4" decode_port="$5" + local details=( + "Prefill OpenEngine: http://127.0.0.1:$prefill_port" + "Decode OpenEngine: http://127.0.0.1:$decode_port" + ) + case "$mode" in + multimodal) + print_launch_banner --multimodal \ + "Launching TRT-LLM OpenEngine Multimodal Context-First 1P1D (2 GPUs)" \ + "$model" "$http_port" "${details[@]}" + ;; + text) + print_launch_banner \ + "Launching TRT-LLM OpenEngine Context-First 1P1D + LoRA (2 GPUs)" \ + "$model" "$http_port" "${details[@]}" + ;; + *) + echo "Unknown OpenEngine disaggregated launch mode: $mode" >&2 + return 2 + ;; + esac +} diff --git a/examples/backends/trtllm/launch/openengine_disagg.sh b/examples/backends/trtllm/launch/openengine_disagg.sh new file mode 100755 index 000000000000..dbd0f2f09497 --- /dev/null +++ b/examples/backends/trtllm/launch/openengine_disagg.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +set -euo pipefail +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" +source "$SCRIPT_DIR/openengine_common.sh" +trap dynamo_exit_trap EXIT + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-0.6B} +HTTP_PORT=${DYN_HTTP_PORT:-${FRONTEND_PORT:-8000}} +export DYN_HTTP_PORT="$HTTP_PORT" +PREFILL_HTTP_PORT=${PREFILL_HTTP_PORT:-8001} +DECODE_HTTP_PORT=${DECODE_HTTP_PORT:-8002} +PREFILL_PORT=${PREFILL_PORT:-50051} +DECODE_PORT=${DECODE_PORT:-50052} +PREFILL_SYSTEM_PORT=${DYN_SYSTEM_PORT1:-8081} +DECODE_SYSTEM_PORT=${DYN_SYSTEM_PORT2:-8082} +export DYNAMO_HOME=${DYNAMO_HOME:-"$(readlink -f "$SCRIPT_DIR/../../../..")"} +PREFILL_CONFIG=${PREFILL_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/qwen3/prefill-lora.yaml} +DECODE_CONFIG=${DECODE_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/qwen3/decode-lora.yaml} +OPENENGINE_LAUNCH_MODE=${OPENENGINE_LAUNCH_MODE:-text} +export DYN_LORA_ENABLED=${DYN_LORA_ENABLED:-true} +export DYN_LORA_PATH=${DYN_LORA_PATH:-/tmp/dynamo_loras} +readonly OPENENGINE_SCHEMA_RELEASE=cea19cb06acf03c911b84d5c147e519b60dd92a6 +export OPENENGINE_SCHEMA_RELEASE + +warn_trtllm_serve_profile_overrides "$PREFILL_CONFIG and $DECODE_CONFIG" +print_openengine_disagg_banner "$OPENENGINE_LAUNCH_MODE" "$MODEL_PATH" \ + "$HTTP_PORT" "$PREFILL_PORT" "$DECODE_PORT" + +# The sidecars advertise no frontend media decoder. URL/data media is therefore +# preserved for both context and generation workers, which is required for +# Qwen3-VL context-first P/D mRoPE reconstruction. +python3 -m dynamo.frontend --router-mode kv & +CUDA_VISIBLE_DEVICES=${PREFILL_CUDA_VISIBLE_DEVICES:-0} \ + trtllm-serve "$MODEL_PATH" --server_role CONTEXT \ + --backend pytorch \ + --config "$PREFILL_CONFIG" \ + --host 127.0.0.1 \ + --port "$PREFILL_HTTP_PORT" \ + --openengine-host 127.0.0.1 \ + --openengine-port "$PREFILL_PORT" & +CUDA_VISIBLE_DEVICES=${DECODE_CUDA_VISIBLE_DEVICES:-1} \ + trtllm-serve "$MODEL_PATH" --server_role GENERATION \ + --backend pytorch \ + --config "$DECODE_CONFIG" \ + --host 127.0.0.1 \ + --port "$DECODE_HTTP_PORT" \ + --openengine-host 127.0.0.1 \ + --openengine-port "$DECODE_PORT" & + +DYN_SYSTEM_PORT="$PREFILL_SYSTEM_PORT" dynamo-openengine-sidecar \ + --openengine-endpoint "http://127.0.0.1:$PREFILL_PORT" \ + --expected-engine tensorrt_llm & +DYN_SYSTEM_PORT="$DECODE_SYSTEM_PORT" dynamo-openengine-sidecar \ + --openengine-endpoint "http://127.0.0.1:$DECODE_PORT" \ + --expected-engine tensorrt_llm & + +wait_any_exit diff --git a/examples/backends/trtllm/launch/openengine_lora.sh b/examples/backends/trtllm/launch/openengine_lora.sh new file mode 100755 index 000000000000..4f984a08c331 --- /dev/null +++ b/examples/backends/trtllm/launch/openengine_lora.sh @@ -0,0 +1,27 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +set -euo pipefail +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" +export DYNAMO_HOME=${DYNAMO_HOME:-"$(readlink -f "$SCRIPT_DIR/../../../..")"} + +export DYN_LORA_ENABLED=true +export DYN_LORA_PATH=${DYN_LORA_PATH:-/tmp/dynamo_loras} +export TRTLLM_CONFIG=${TRTLLM_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/qwen3/lora.yaml} +export OPENENGINE_LAUNCH_MODE=lora +exec "$SCRIPT_DIR/openengine_agg.sh" "$@" diff --git a/examples/backends/trtllm/launch/openengine_multimodal.sh b/examples/backends/trtllm/launch/openengine_multimodal.sh new file mode 100755 index 000000000000..d2f804136a83 --- /dev/null +++ b/examples/backends/trtllm/launch/openengine_multimodal.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +set -euo pipefail +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" +export DYNAMO_HOME=${DYNAMO_HOME:-"$(readlink -f "$SCRIPT_DIR/../../../..")"} + +export MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-VL-2B-Instruct} +export TRTLLM_CONFIG=${TRTLLM_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/qwen3-vl-2b-instruct/agg.yaml} +export OPENENGINE_LAUNCH_MODE=multimodal +exec "$SCRIPT_DIR/openengine_agg.sh" "$@" diff --git a/examples/backends/trtllm/launch/openengine_multimodal_disagg.sh b/examples/backends/trtllm/launch/openengine_multimodal_disagg.sh new file mode 100755 index 000000000000..c5d98f839b98 --- /dev/null +++ b/examples/backends/trtllm/launch/openengine_multimodal_disagg.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# Qwen3-VL supports ordered image/video URLs and data URIs in both aggregate +# and context-first P/D modes. Decode intentionally receives the original media +# again so TRT-LLM can reconstruct its local multimodal position metadata. +set -euo pipefail +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" +export DYNAMO_HOME=${DYNAMO_HOME:-"$(readlink -f "$SCRIPT_DIR/../../../..")"} + +export MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-VL-2B-Instruct} +export PREFILL_CONFIG=${PREFILL_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/qwen3-vl-2b-instruct/prefill.yaml} +export DECODE_CONFIG=${DECODE_CONFIG:-$DYNAMO_HOME/examples/backends/trtllm/engine_configs/qwen3-vl-2b-instruct/decode.yaml} +export DYN_LORA_ENABLED=${DYN_LORA_ENABLED:-false} +export OPENENGINE_LAUNCH_MODE=multimodal +exec "$SCRIPT_DIR/openengine_disagg.sh" "$@" diff --git a/lib/backend-common/src/adapter.rs b/lib/backend-common/src/adapter.rs index a3be68fb5d61..31469c7f2333 100644 --- a/lib/backend-common/src/adapter.rs +++ b/lib/backend-common/src/adapter.rs @@ -13,6 +13,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Instant; use async_trait::async_trait; +use dynamo_llm::protocols::common::FinishReason; use dynamo_llm::protocols::common::llm_backend::LLMEngineOutput; use dynamo_llm::protocols::common::preprocessor::PreprocessedRequest; use dynamo_runtime::engine::AsyncEngineContext; @@ -30,6 +31,7 @@ use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::disagg::DisaggregationMode; use crate::engine::{GenerateContext, LLMEngine, RawEngine}; +use crate::error::{BackendError, DynamoError, ErrorType}; /// Test-only override count. Compiled out of release builds — tests acquire /// an `OtlpExportOverride` RAII guard to force-enable the recording @@ -75,6 +77,10 @@ impl StreamSpanFinalizer { fn mark_completed(&self) { self.completed.set(true); } + + fn completed(&self) -> bool { + self.completed.get() + } } impl Drop for StreamSpanFinalizer { @@ -203,6 +209,7 @@ impl AsyncEngine, ManyOut Result>, Error> { let (request, handle) = input.into_parts(); let ctx: Arc = handle.context(); + let expected_outputs = request.sampling_options.n.unwrap_or(1).max(1) as usize; // Per-request worker-side span. Nests under `handle_payload` (set up // by the runtime's NATS ingress) so the trace tree has a contiguous @@ -371,6 +378,7 @@ impl AsyncEngine, ManyOut, ManyOut { + if usize::try_from(index).map_or(true, |index| index >= expected_outputs) { + yield Annotated::from_err(DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::EngineShutdown)) + .message(format!( + "engine returned terminal output index {index}, but request asked for {expected_outputs} output(s)" + )) + .build()); + finalizer.mark_completed(); + break; + } + if !finished_indexes.insert(index) { + yield Annotated::from_err(DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::EngineShutdown)) + .message(format!( + "engine returned more than one terminal for output index {index}" + )) + .build()); + finalizer.mark_completed(); + break; + } + } + None if expected_outputs > 1 && !globally_terminal => { + yield Annotated::from_err(DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::EngineShutdown)) + .message(format!( + "engine returned an unindexed terminal for a {expected_outputs}-output request" + )) + .build()); + finalizer.mark_completed(); + break; + } + None => {} + } + let all_outputs_finished = globally_terminal + || (chunk.index.is_none() && expected_outputs == 1) + || finished_indexes.len() == expected_outputs; + yield Annotated::from_data(chunk); + if all_outputs_finished { + finalizer.mark_completed(); + break; + } + } else { + yield Annotated::from_data(chunk); } } Err(dynamo_err) => { @@ -468,6 +522,19 @@ impl AsyncEngine, ManyOut>().await; + + assert_eq!(collected.len(), 2); + assert_eq!(collected[0].data.as_ref().unwrap().index, Some(0)); + assert_eq!(collected[1].data.as_ref().unwrap().index, Some(1)); + assert_eq!(abort_ct.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn adapter_cancellation_triggers_engine_abort() { let engine = Arc::new(MockEngine { diff --git a/lib/backend-common/src/engine.rs b/lib/backend-common/src/engine.rs index c7368738db29..ef36b58a493f 100644 --- a/lib/backend-common/src/engine.rs +++ b/lib/backend-common/src/engine.rs @@ -20,10 +20,13 @@ use futures::stream::BoxStream; use tokio::sync::watch; use crate::error::DynamoError; +use dynamo_llm::local_model::LocalModel; +use dynamo_llm::model_type::ModelType; +use dynamo_llm::worker_type::WorkerType; pub use dynamo_llm::kv_router::publisher::KvEventPublisher; pub use dynamo_llm::protocols::common::llm_backend::{ - LLMEngineOutput, LogProbs, TopLogprob, TopLogprobs, + LLMEngineOutput, LogProbs, PromptLogprobEntry, PromptLogprobs, TopLogprob, TopLogprobs, }; pub use dynamo_llm::protocols::common::preprocessor::{ BootstrapInfo, MultimodalData, MultimodalDataMap, PrefillResult, PreprocessedRequest, @@ -200,9 +203,10 @@ pub trait LLMEngine: Send + Sync + 'static { /// emit a terminal `Ok(chunk)` with `FinishReason::Cancelled`. /// /// Stream item: `Result`. - /// * `Ok(chunk)` carries normal output. Exactly one terminal `Ok` - /// chunk (one with `finish_reason` set) must be the last item - /// yielded, and no items may follow it. + /// * `Ok(chunk)` carries normal output. Single-output streams have one + /// terminal `Ok` chunk (one with `finish_reason` set). Multi-output + /// streams have one terminal per distinct `index`; no later chunk may + /// target an already-finished index. /// * `Err(dynamo_err)` carries a typed mid-stream failure (e.g. /// `BackendError::InvalidArgument`). It is itself terminal — the /// framework forwards it as `Annotated::error` and stops polling @@ -258,6 +262,19 @@ pub trait LLMEngine: Send + Sync + 'static { Ok(None) } + /// Long-lived remote-engine liveness watch. Returning is fatal to the + /// serving worker and causes its discovery registration to be removed. + /// In-process engines normally keep the default pending future. + async fn watch(&self) -> Result<(), DynamoError> { + std::future::pending().await + } + + /// Ask an out-of-process engine to stop admission and wait for terminal + /// drain completion. Called before quiescence polling and cleanup. + async fn drain(&self) -> Result<(), DynamoError> { + Ok(()) + } + /// Release all engine resources. Called exactly once. /// /// `Worker` guarantees: @@ -384,6 +401,20 @@ pub trait LLMEngine: Send + Sync + 'static { ) -> Result<(), DynamoError> { Ok(()) } + + /// Hand the engine a clone of the fully populated base-model card before + /// it is attached. Dynamic-LoRA engines use this to publish adapter model + /// cards with identical tokenizer/runtime/routing metadata. + async fn on_model_ready( + &self, + _endpoint: dynamo_runtime::component::Endpoint, + _base_model: LocalModel, + _model_type: ModelType, + _worker_type: WorkerType, + _needs: Vec>, + ) -> Result<(), DynamoError> { + Ok(()) + } } /// Raw media-generation engine trait — the non-token sibling of [`LLMEngine`]. diff --git a/lib/backend-common/src/lib.rs b/lib/backend-common/src/lib.rs index bf9a6ed42c1d..1cd55badf32e 100644 --- a/lib/backend-common/src/lib.rs +++ b/lib/backend-common/src/lib.rs @@ -36,8 +36,8 @@ pub use engine::{ FinishReason, GenerateContext, GuidedDecodingOptions, HEALTH_CHECK_KEY, KvEventPublisher, KvEventSource, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, LlmRegistration, LogProbs, Metrics, MetricsBindings, MetricsCtx, OnPublisherReady, OnSnapshotPublisherReady, - OutputOptions, PrefillResult, PreprocessedRequest, RawEngine, SamplingOptions, StopConditions, - StopReason, TopLogprob, TopLogprobs, chunk, usage, + OutputOptions, PrefillResult, PreprocessedRequest, PromptLogprobEntry, PromptLogprobs, + RawEngine, SamplingOptions, StopConditions, StopReason, TopLogprob, TopLogprobs, chunk, usage, }; pub use error::{BackendError, DynamoError, ErrorType}; pub use metrics::{ComponentGauges, EngineMetrics, LifecycleGauges}; diff --git a/lib/backend-common/src/validate.rs b/lib/backend-common/src/validate.rs index 8900a57262a9..85f4e054a23c 100644 --- a/lib/backend-common/src/validate.rs +++ b/lib/backend-common/src/validate.rs @@ -4,7 +4,8 @@ //! Debug-build stream validator. //! //! Wraps the engine's returned stream and panics on contract violations: -//! - a chunk yielded after a terminal chunk (one carrying `finish_reason`) +//! - a chunk yielded after an unindexed/global terminal, or another chunk for +//! an output index that already emitted its terminal //! - an Encode-mode non-cancelled terminal that lacks an //! `encoder_result: Some(Value::Object(_))` payload (engines that build //! the terminal via `LLMEngineOutput::stop()` instead of @@ -24,52 +25,65 @@ use dynamo_llm::protocols::common::FinishReason; use dynamo_llm::protocols::common::llm_backend::LLMEngineOutput; use futures::StreamExt; use futures::stream::BoxStream; +use std::collections::HashSet; pub(crate) fn wrap( stream: BoxStream<'static, Result>, mode: DisaggregationMode, ) -> BoxStream<'static, Result> { - let mut terminal_seen = false; + let mut global_terminal_seen = false; + let mut finished_indexes = HashSet::new(); Box::pin(async_stream::stream! { let mut inner = stream; while let Some(item) = inner.next().await { assert!( - !terminal_seen, + !global_terminal_seen, "LLMEngine contract violation: item yielded after terminal item \ - (a chunk with finish_reason set, or an Err, must be the last item)" + (an unindexed terminal chunk or Err must be the last item)" ); match &item { - Ok(chunk) if chunk.finish_reason.is_some() => { - // Encode-mode terminal rule: successful terminals MUST - // carry an object-shaped encoder_result. Cancelled is - // exempt because cancellation can land before the encoder - // produces a payload; Error terminals are exempt because a - // failure path (LLMEngineOutput::error) legitimately has no - // encoder_result. - if mode.is_encode() - && !matches!( - chunk.finish_reason, - Some(FinishReason::Cancelled | FinishReason::Error(_)) - ) - { + Ok(chunk) => { + if let Some(index) = chunk.index { assert!( - matches!( - chunk.encoder_result.as_ref(), - Some(v) if v.is_object() - ), - "Encode-mode contract violation: non-cancelled terminal \ - chunk must carry encoder_result: Some(Value::Object(_)). \ - Use LLMEngineOutput::encode_terminal(map) or \ - .with_encoder_result(map) instead of LLMEngineOutput::stop(). \ - Got finish_reason={:?}, encoder_result={:?}", - chunk.finish_reason, - chunk.encoder_result, + !finished_indexes.contains(&index), + "LLMEngine contract violation: item yielded after terminal item for output index {index}" ); } - terminal_seen = true; + if chunk.finish_reason.is_some() { + // Encode-mode terminal rule: successful terminals MUST + // carry an object-shaped encoder_result. Cancelled is + // exempt because cancellation can land before the encoder + // produces a payload; Error terminals are exempt because a + // failure path (LLMEngineOutput::error) legitimately has no + // encoder_result. + if mode.is_encode() + && !matches!( + chunk.finish_reason, + Some(FinishReason::Cancelled | FinishReason::Error(_)) + ) + { + assert!( + matches!( + chunk.encoder_result.as_ref(), + Some(v) if v.is_object() + ), + "Encode-mode contract violation: non-cancelled terminal \ + chunk must carry encoder_result: Some(Value::Object(_)). \ + Use LLMEngineOutput::encode_terminal(map) or \ + .with_encoder_result(map) instead of LLMEngineOutput::stop(). \ + Got finish_reason={:?}, encoder_result={:?}", + chunk.finish_reason, + chunk.encoder_result, + ); + } + if let Some(index) = chunk.index { + finished_indexes.insert(index); + } else { + global_terminal_seen = true; + } + } } - Err(_) => terminal_seen = true, - _ => {} + Err(_) => global_terminal_seen = true, } yield item; } @@ -122,6 +136,19 @@ mod tests { )); } + #[tokio::test] + async fn distinct_output_indexes_may_each_emit_a_terminal() { + let mut first = LLMEngineOutput::length(); + first.index = Some(0); + let mut second = LLMEngineOutput::stop(); + second.index = Some(1); + let wrapped = wrap( + to_stream(vec![first, second]), + DisaggregationMode::Aggregated, + ); + assert_eq!(wrapped.collect::>().await.len(), 2); + } + #[tokio::test] #[should_panic(expected = "item yielded after terminal item")] async fn panics_on_chunk_after_terminal() { diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 360283ddffb4..d6931298b517 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -281,6 +281,20 @@ impl EngineKind { } } + async fn watch(&self) -> Result<(), DynamoError> { + match self { + EngineKind::Llm(e) => e.watch().await, + EngineKind::Raw(_) => std::future::pending().await, + } + } + + async fn drain(&self) -> Result<(), DynamoError> { + match self { + EngineKind::Llm(e) => e.drain().await, + EngineKind::Raw(_) => Ok(()), + } + } + async fn setup_metrics(&self, ctx: MetricsCtx<'_>) -> Result { match self { EngineKind::Llm(e) => e.setup_metrics(ctx).await, @@ -358,6 +372,23 @@ impl EngineKind { } } + async fn on_model_ready( + &self, + endpoint: dynamo_runtime::component::Endpoint, + base_model: LocalModel, + model_type: ModelType, + worker_type: WorkerType, + needs: Vec>, + ) -> Result<(), DynamoError> { + match self { + EngineKind::Llm(e) => { + e.on_model_ready(endpoint, base_model, model_type, worker_type, needs) + .await + } + EngineKind::Raw(_) => Ok(()), + } + } + /// Raw media engines (image/video/audio) register name-only — the engine /// loads the model itself and the model has no LLM artifacts (tokenizer / /// chat template / config.json) for Dynamo to fetch. @@ -743,12 +774,26 @@ impl Worker { /// grace period → engine drain → cleanup. Shared by every shutdown path — /// pre-serve (mid-start signal) and the serve loop's signal arm. async fn orchestrator_steps(&mut self, endpoint: &dynamo_runtime::component::Endpoint) { - if let Err(e) = endpoint.unregister_endpoint_instance().await { - tracing::warn!(error = %e, "discovery unregister failed"); - } else { - tracing::info!("Endpoint unregistered from discovery"); + let _ = self.finalize_serve(Some(endpoint), None).await; + } + + /// Common terminal path for normal shutdown and engine-watch failure. + /// Keeping the watch error until after unregister/drain/cleanup prevents a + /// dead remote from leaving a stale discovery record or skipping cleanup. + async fn finalize_serve( + &mut self, + endpoint: Option<&dynamo_runtime::component::Endpoint>, + terminal_error: Option, + ) -> Result<(), DynamoError> { + if let Some(endpoint) = endpoint { + if let Err(e) = endpoint.unregister_endpoint_instance().await { + tracing::warn!(error = %e, "discovery unregister failed"); + } else { + tracing::info!("Endpoint unregistered from discovery"); + } } self.run_engine_shutdown_steps().await; + terminal_error.map_or(Ok(()), Err) } /// Start the engine exactly once. `Worker::run` consumes `self`, so all @@ -838,6 +883,16 @@ impl Worker { // fire before the engine has the endpoint. self.engine.on_endpoint_ready(endpoint.clone()).await?; + self.engine + .on_model_ready( + endpoint.clone(), + local_model.clone(), + model_type, + worker_type, + needs.clone(), + ) + .await?; + local_model .attach( &endpoint, @@ -962,8 +1017,11 @@ impl Worker { } let serve_fut = builder.start(); tokio::pin!(serve_fut); + let watched_engine = self.engine.clone(); + let watch_fut = watched_engine.watch(); + tokio::pin!(watch_fut); - tokio::select! { + let watch_error = tokio::select! { biased; result = &mut serve_fut => { match result { @@ -974,6 +1032,7 @@ impl Worker { tracing::info!( "Endpoint completed gracefully; running shutdown orchestration" ); + None } // Serve errored; cleanup_once in run() is the safety net. Err(e) => { @@ -986,11 +1045,17 @@ impl Worker { } _ = shutdown.cancelled() => { tracing::info!("Received shutdown signal; running graceful orchestration"); + None } - } + result = &mut watch_fut => { + Some(result.err().unwrap_or_else(|| err( + ErrorType::Backend(BackendError::EngineShutdown), + "engine watch ended unexpectedly", + ))) + } + }; - self.orchestrator_steps(&endpoint).await; - Ok(()) + self.finalize_serve(Some(&endpoint), watch_error).await } /// Engine-facing shutdown sequence: grace period sleep → drain loop on @@ -1013,7 +1078,22 @@ impl Worker { } let drain_start = std::time::Instant::now(); - self.drain_until_idle_or_deadline().await; + let configured = drain_timeout_secs(); + let cap = (graceful_shutdown_timeout().as_secs_f64() - CLEANUP_RESERVE_S).max(0.0); + let budget = Duration::from_secs_f64(configured.min(cap)); + let drain_deadline = tokio::time::Instant::now() + budget; + match tokio::time::timeout(budget, self.engine.drain()).await { + Ok(Ok(())) => {} + Ok(Err(e)) => { + tracing::warn!(error = %e, "engine drain request failed; continuing shutdown") + } + Err(_) => tracing::warn!( + timeout_s = budget.as_secs_f64(), + "engine drain request timed out; continuing shutdown" + ), + } + let remaining = drain_deadline.saturating_duration_since(tokio::time::Instant::now()); + self.drain_until_idle_or_deadline(remaining).await; let drain_elapsed = drain_start.elapsed().as_secs_f64(); if let Some(lifecycle) = self.lifecycle.as_ref() { lifecycle.observe_drain_time(drain_elapsed); @@ -1030,21 +1110,19 @@ impl Worker { /// `DRAIN_POLL_INTERVAL_S`, exiting on `Some(true)` or when the budget /// expires. Budget = `DYN_PREFILL_DRAIN_TIMEOUT_S` capped at /// `graceful_shutdown_timeout - CLEANUP_RESERVE_S`. - async fn drain_until_idle_or_deadline(&self) { + async fn drain_until_idle_or_deadline(&self, budget: Duration) { if !self.config.disaggregation_mode.is_prefill() { return; } - let configured = drain_timeout_secs(); - let cap = (graceful_shutdown_timeout().as_secs_f64() - CLEANUP_RESERVE_S).max(0.0); - let budget = configured.min(cap); - let deadline = std::time::Instant::now() + Duration::from_secs_f64(budget); + let deadline = std::time::Instant::now() + budget; let start = std::time::Instant::now(); let mut last_heartbeat = start; let mut announced = false; loop { - match self.engine.is_quiescent().await { + let remaining = deadline.saturating_duration_since(std::time::Instant::now()); + match tokio::time::timeout(remaining, self.engine.is_quiescent()).await { // Quiescent: in-flight transfers done, safe to exit drain. - Ok(Some(true)) => { + Ok(Ok(Some(true))) => { if announced { tracing::info!( "drain: exited (quiescent, elapsed={:.1}s)", @@ -1054,16 +1132,22 @@ impl Worker { return; } // Busy (Some(false)) or no introspection (None): keep polling. - Ok(Some(false)) | Ok(None) => {} - Err(e) => { + Ok(Ok(Some(false))) | Ok(Ok(None)) => {} + Ok(Err(e)) => { tracing::debug!(error = %e, "is_quiescent raised; treating as not quiescent") } + Err(_) => { + tracing::warn!( + "is_quiescent timed out at drain deadline; proceeding with cleanup" + ); + return; + } } if !announced { // First non-quiescent poll: announce once that we're waiting. tracing::info!( "drain: waiting for prefill to quiesce; polling is_quiescent (timeout={:.1}s)", - budget + budget.as_secs_f64() ); announced = true; } @@ -2477,6 +2561,23 @@ mod tests { assert_eq!(worker.state, LifecycleState::Stopped); } + #[tokio::test] + async fn watch_failure_still_runs_cleanup_before_returning_error() { + let (engine, cleanup_calls) = StateMockEngine::new(false); + let mut worker = worker_with(engine); + worker.start_engine(0).await.unwrap(); + let watch_error = err( + ErrorType::Backend(BackendError::EngineShutdown), + "synthetic watch failure", + ); + + let result = worker.finalize_serve(None, Some(watch_error)).await; + + assert!(result.is_err(), "watch failure must remain fatal"); + assert_eq!(cleanup_calls.load(Ordering::SeqCst), 1); + assert_eq!(worker.state, LifecycleState::Stopped); + } + // The pre-start shutdown path is handled in `run_inner` via a // `CancellationToken` cancellation check before `start_engine` is // called — not by flipping state to `Stopped` first. There is no @@ -2545,6 +2646,11 @@ mod tests { } } + async fn drain(&self) -> Result<(), DynamoError> { + self.log.lock().unwrap().push("drain"); + Ok(()) + } + async fn cleanup(&self) -> Result<(), DynamoError> { self.log.lock().unwrap().push("cleanup"); Ok(()) @@ -2564,8 +2670,8 @@ mod tests { let recorded = log.lock().unwrap().clone(); assert_eq!( recorded, - vec!["start", "is_quiescent", "cleanup"], - "is_quiescent (drain) must run before cleanup" + vec!["start", "drain", "is_quiescent", "cleanup"], + "remote drain and is_quiescent must run before cleanup" ); } @@ -2581,7 +2687,7 @@ mod tests { let _guard = ENV_LOCK.lock().unwrap(); let saved = std::env::var(DRAIN_TIMEOUT_ENV).ok(); // SAFETY: tests in this mod serialize env access via ENV_LOCK. - unsafe { std::env::set_var(DRAIN_TIMEOUT_ENV, "0") }; + unsafe { std::env::set_var(DRAIN_TIMEOUT_ENV, "0.01") }; let (engine, log) = OrderingMockEngine::new(true); // is_quiescent fails let mut worker = worker_with_prefill(engine); @@ -2591,7 +2697,7 @@ mod tests { // is_quiescent ran at least once (and errored), then cleanup ran. let recorded = log.lock().unwrap().clone(); - assert!(recorded.starts_with(&["start", "is_quiescent"])); + assert!(recorded.starts_with(&["start", "drain", "is_quiescent"])); assert_eq!(recorded.last().copied(), Some("cleanup")); assert_eq!(worker.state, LifecycleState::Stopped); @@ -2605,11 +2711,9 @@ mod tests { } #[tokio::test] - async fn shutdown_steps_skip_drain_for_non_prefill() { - // Drain is prefill-only: an aggregated (or decode) worker must go - // straight to cleanup without ever polling is_quiescent, regardless of - // what the engine would report. Guards the mode-gate invariant that - // makes the `Ok(None)` default safe (only prefill workers drain). + async fn shutdown_steps_skip_quiescence_poll_for_non_prefill() { + // Process-wide engine drain applies to every role. Only prefill needs + // the additional KV-transfer quiescence polling phase. let (engine, log) = OrderingMockEngine::new(false); let mut worker = worker_with(engine); // WorkerConfig::default() => Aggregated worker.start_engine(0).await.unwrap(); @@ -2619,11 +2723,87 @@ mod tests { let recorded = log.lock().unwrap().clone(); assert_eq!( recorded, - vec!["start", "cleanup"], - "non-prefill workers must not poll is_quiescent (no drain)" + vec!["start", "drain", "cleanup"], + "non-prefill workers must drain remotely without polling is_quiescent" ); } + struct HangingDrainEngine { + cleanup_calls: Arc, + } + + #[async_trait] + impl LLMEngine for HangingDrainEngine { + async fn start(&self, _worker_id: u64) -> Result { + Ok(EngineConfig { + model: "mock".to_string(), + ..Default::default() + }) + } + + async fn generate( + &self, + _request: PreprocessedRequest, + _ctx: crate::engine::GenerateContext, + ) -> Result< + BoxStream<'static, Result>, + DynamoError, + > { + unreachable!("not used in shutdown test") + } + + async fn drain(&self) -> Result<(), DynamoError> { + std::future::pending().await + } + + async fn cleanup(&self) -> Result<(), DynamoError> { + self.cleanup_calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + #[allow(clippy::await_holding_lock)] + #[tokio::test] + async fn hanging_remote_drain_is_bounded_and_cleanup_still_runs() { + use dynamo_runtime::config::environment_names::worker as env_worker; + + let _guard = ENV_LOCK.lock().unwrap(); + let graceful_key = env_worker::DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT; + let saved_drain = std::env::var(DRAIN_TIMEOUT_ENV).ok(); + let saved_graceful = std::env::var(graceful_key).ok(); + // SAFETY: env-mutating shutdown tests are serialized by ENV_LOCK. + unsafe { + std::env::set_var(DRAIN_TIMEOUT_ENV, "0.01"); + std::env::set_var(graceful_key, "6"); + } + + let cleanup_calls = Arc::new(AtomicUsize::new(0)); + let engine = Arc::new(HangingDrainEngine { + cleanup_calls: cleanup_calls.clone(), + }); + let mut worker = worker_with(engine); + worker.start_engine(0).await.unwrap(); + tokio::time::timeout( + Duration::from_millis(500), + worker.run_engine_shutdown_steps_with_grace(0.0), + ) + .await + .expect("hanging remote drain must be bounded"); + assert_eq!(cleanup_calls.load(Ordering::SeqCst), 1); + + // SAFETY: see serialized env access above. + unsafe { + match saved_drain { + Some(value) => std::env::set_var(DRAIN_TIMEOUT_ENV, value), + None => std::env::remove_var(DRAIN_TIMEOUT_ENV), + } + match saved_graceful { + Some(value) => std::env::set_var(graceful_key, value), + None => std::env::remove_var(graceful_key), + } + } + } + // The "drain skipped when engine never started" scenario isn't // reachable through the public `Worker::run` flow — pre-start // shutdown returns from `run_inner` before `serve_with_orchestrator` @@ -2835,7 +3015,7 @@ mod tests { ); let recorded = log.lock().unwrap().clone(); - assert_eq!(recorded, vec!["start", "is_quiescent", "cleanup"]); + assert_eq!(recorded, vec!["start", "drain", "is_quiescent", "cleanup"]); } // ------------------------------------------------------------------- diff --git a/lib/openengine-sidecar/Cargo.toml b/lib/openengine-sidecar/Cargo.toml new file mode 100644 index 000000000000..69c48a105c30 --- /dev/null +++ b/lib/openengine-sidecar/Cargo.toml @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "dynamo-openengine-sidecar" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +description = "Engine-neutral Dynamo backend for out-of-process OpenEngine servers." +build = "build.rs" + +[package.metadata.openengine] +commit = "cea19cb06acf03c911b84d5c147e519b60dd92a6" + +[[bin]] +name = "dynamo-openengine-sidecar" +path = "src/main.rs" + +[dependencies] +dynamo-backend-common = { workspace = true } +dynamo-kv-router = { workspace = true } +dynamo-llm = { workspace = true } +dynamo-protocols = { workspace = true } +dynamo-runtime = { workspace = true } +openengine-proto = { version = "0.2.0", path = "../../../openengine-trtllm/packages/rust/openengine-proto" } + +anyhow = { workspace = true } +async-stream = { workspace = true } +async-trait = { workspace = true } +base64 = "0.22" +clap = { version = "4", features = ["derive", "env"] } +futures = { workspace = true } +parking_lot = { workspace = true } +prost-types = "0.14.4" +serde_json = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tonic = "0.14.6" +tracing = { workspace = true } + +[dev-dependencies] +dynamo-backend-common = { workspace = true, features = ["testing"] } +tokio-stream = { workspace = true, features = ["net"] } diff --git a/lib/openengine-sidecar/README.md b/lib/openengine-sidecar/README.md new file mode 100644 index 000000000000..7ae93bf6f70b --- /dev/null +++ b/lib/openengine-sidecar/README.md @@ -0,0 +1,17 @@ +# Dynamo OpenEngine sidecar + +This crate is developed against the generated `openengine-proto` 0.2.0 crate +from the sibling OpenEngine worktree. The local checkout layout is fixed: + +```text +sidecar/ +├── dynamo-trtllm-sidecar/ +└── openengine-trtllm/ + └── packages/rust/openengine-proto/ +``` + +`build.rs` verifies that the sibling OpenEngine worktree is exactly commit +`cea19cb06acf03c911b84d5c147e519b60dd92a6`; a different HEAD fails the build. +The Cargo dependency also requires package version 0.2.0. Keep the path +dependency until that exact OpenEngine commit is published at an immutable Git +or registry location, then replace it with the corresponding immutable pin. diff --git a/lib/openengine-sidecar/build.rs b/lib/openengine-sidecar/build.rs new file mode 100644 index 000000000000..02838b9f55b0 --- /dev/null +++ b/lib/openengine-sidecar/build.rs @@ -0,0 +1,54 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::path::PathBuf; +use std::process::Command; + +const OPENENGINE_COMMIT: &str = "cea19cb06acf03c911b84d5c147e519b60dd92a6"; + +fn main() { + let manifest = PathBuf::from(std::env::var_os("CARGO_MANIFEST_DIR").unwrap()); + let repository = manifest.join("../../../openengine-trtllm"); + let package = repository.join("packages/rust/openengine-proto"); + let proto = repository.join("proto"); + + println!("cargo:rerun-if-changed={}", package.display()); + println!("cargo:rerun-if-changed={}", proto.display()); + println!("cargo:rustc-env=OPENENGINE_PROTO_COMMIT={OPENENGINE_COMMIT}"); + + let output = Command::new("git") + .args(["-C"]) + .arg(&repository) + .args(["rev-parse", "HEAD"]) + .output() + .unwrap_or_else(|error| panic!("failed to inspect local OpenEngine dependency: {error}")); + assert!( + output.status.success(), + "local OpenEngine dependency is not a Git worktree: {}", + String::from_utf8_lossy(&output.stderr).trim() + ); + let actual = String::from_utf8_lossy(&output.stdout); + assert_eq!( + actual.trim(), + OPENENGINE_COMMIT, + "local openengine-proto dependency drifted; fetch and checkout the locked OpenEngine commit" + ); + + let status = Command::new("git") + .args(["-C"]) + .arg(&repository) + .args(["status", "--porcelain", "--untracked-files=all", "--"]) + .args(["proto", "packages/rust/openengine-proto", "packages/python"]) + .output() + .unwrap_or_else(|error| panic!("failed to inspect local OpenEngine dependency: {error}")); + assert!( + status.status.success(), + "failed to inspect local OpenEngine dependency status: {}", + String::from_utf8_lossy(&status.stderr).trim() + ); + assert!( + status.stdout.is_empty(), + "local OpenEngine protocol/package sources are dirty; pinning HEAD would not identify the consumed contract:\n{}", + String::from_utf8_lossy(&status.stdout) + ); +} diff --git a/lib/openengine-sidecar/src/args.rs b/lib/openengine-sidecar/src/args.rs new file mode 100644 index 000000000000..9c867ba72dce --- /dev/null +++ b/lib/openengine-sidecar/src/args.rs @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::path::PathBuf; +use std::time::Duration; + +#[derive(clap::Parser, Debug, Clone)] +#[command( + name = "dynamo-openengine-sidecar", + about = "Engine-neutral Dynamo worker for an out-of-process OpenEngine server" +)] +pub struct Args { + /// OpenEngine gRPC endpoint. A bare host:port is interpreted as plaintext HTTP/2. + #[arg(long, env = "OPENENGINE_ENDPOINT")] + pub openengine_endpoint: String, + + /// Optional engine identity assertion (for example `tensorrt_llm`). + #[arg(long, env = "OPENENGINE_EXPECTED_ENGINE")] + pub expected_engine: Option, + + /// Model to select when the OpenEngine server advertises more than one model. + #[arg(long, env = "OPENENGINE_MODEL")] + pub model: Option, + + /// Number of independent HTTP/2 connections used for Generate streams. + #[arg(long, env = "OPENENGINE_CONNECTIONS", default_value_t = 8)] + pub openengine_connections: usize, + + #[arg(long, env = "DYN_NAMESPACE", default_value = "dynamo")] + pub namespace: String, + + #[arg(long, env = "DYN_ENDPOINT", default_value = "generate")] + pub endpoint: String, + + #[arg(long, env = "DYN_ENDPOINT_TYPES", default_value = "chat,completions")] + pub endpoint_types: String, + + #[arg(long, env = "DYN_CUSTOM_JINJA_TEMPLATE")] + pub custom_jinja_template: Option, + + #[arg(long, default_value_t = 30)] + pub connect_timeout_secs: u64, + + #[arg(long, default_value_t = 2)] + pub health_poll_interval_secs: u64, + + #[arg(long, default_value_t = 300)] + pub health_deadline_secs: u64, + + /// Maximum time to wait for OpenEngine's process-wide drain stream. + #[arg(long, env = "OPENENGINE_DRAIN_TIMEOUT_SECS", default_value_t = 25)] + pub drain_timeout_secs: u64, + + /// Polling interval for OpenEngine load snapshots. + #[arg(long, default_value_t = 1)] + pub load_poll_interval_secs: u64, +} + +impl Args { + pub fn transport(&self) -> TransportConfig { + TransportConfig { + connect_timeout: Duration::from_secs(self.connect_timeout_secs), + poll_interval: Duration::from_secs(self.health_poll_interval_secs.max(1)), + deadline: Duration::from_secs(self.health_deadline_secs), + drain_timeout: Duration::from_secs(self.drain_timeout_secs.max(1)), + load_poll_interval: Duration::from_secs(self.load_poll_interval_secs.max(1)), + connections: self.openengine_connections.max(1), + } + } +} + +#[derive(Debug, Clone)] +pub struct TransportConfig { + pub connect_timeout: Duration, + pub poll_interval: Duration, + pub deadline: Duration, + pub drain_timeout: Duration, + pub load_poll_interval: Duration, + pub connections: usize, +} + +pub fn normalize_endpoint(raw: &str) -> String { + let value = raw.trim(); + if value.contains("://") { + value.to_string() + } else { + format!("http://{value}") + } +} diff --git a/lib/openengine-sidecar/src/client.rs b/lib/openengine-sidecar/src/client.rs new file mode 100644 index 000000000000..2bde65e89a5d --- /dev/null +++ b/lib/openengine-sidecar/src/client.rs @@ -0,0 +1,313 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use dynamo_backend_common::{BackendError, DynamoError, ErrorType}; +use tokio::time::Instant; +use tonic::transport::{Channel, Endpoint}; + +use crate::args::TransportConfig; +use crate::proto as pb; +use crate::proto::open_engine_client::OpenEngineClient; + +pub type Client = OpenEngineClient; + +#[derive(Clone, Debug)] +pub struct Discovery { + pub engine: pb::EngineInfo, + pub model: pb::ModelInfo, + pub selected_model: String, +} + +pub struct Pool { + channels: Vec, + next: AtomicUsize, +} + +impl Pool { + pub async fn connect( + uri: &str, + config: &TransportConfig, + size: usize, + ) -> Result { + let mut channels = Vec::with_capacity(size.max(1)); + for _ in 0..size.max(1) { + channels.push(connect_channel(uri, config).await?); + } + Ok(Self { + channels, + next: AtomicUsize::new(0), + }) + } + + pub fn len(&self) -> usize { + self.channels.len() + } + + pub fn stream_client(&self) -> Client { + let index = self.next.fetch_add(1, Ordering::Relaxed) % self.channels.len(); + OpenEngineClient::new(self.channels[index].clone()) + } + + pub fn control_client(&self) -> Client { + OpenEngineClient::new(self.channels[0].clone()) + } + + pub fn channel(&self) -> Channel { + self.channels[0].clone() + } +} + +pub async fn connect(uri: &str, config: &TransportConfig) -> Result { + Ok(OpenEngineClient::new(connect_channel(uri, config).await?)) +} + +async fn connect_channel(uri: &str, config: &TransportConfig) -> Result { + let deadline = Instant::now() + config.deadline; + loop { + let result = async { + Endpoint::from_shared(uri.to_string()) + .map_err(|error| error.to_string())? + .connect_timeout(config.connect_timeout) + .connect() + .await + .map_err(|error| error.to_string()) + } + .await; + match result { + Ok(channel) => return Ok(channel), + Err(error) if Instant::now() < deadline => { + tracing::debug!(%error, %uri, "OpenEngine connection attempt failed"); + tokio::time::sleep(config.poll_interval).await; + } + Err(error) => { + return Err(cannot_connect(format!( + "could not reach OpenEngine at {uri} within {:?}: {error}", + config.deadline + ))); + } + } + } +} + +pub async fn discover( + client: &mut Client, + requested_model: Option<&str>, + expected_engine: Option<&str>, + deadline: Duration, +) -> Result { + let deadline_at = Instant::now() + deadline; + let mut engine = timeout_discovery_rpc( + "GetEngineInfo", + deadline_at, + client.get_engine_info(pb::GetEngineInfoRequest {}), + ) + .await? + .map_err(|status| status_to_dynamo("GetEngineInfo", status))? + .into_inner(); + + validate_schema(&engine)?; + if let Some(expected) = expected_engine + && engine.engine_name != expected + { + return Err(invalid_arg(format!( + "OpenEngine engine mismatch: expected `{expected}`, discovered `{}`", + engine.engine_name + ))); + } + + let selected_model = select_model(&engine.supported_models, requested_model)?; + // Model enumeration is part of GetEngineInfo in revision 2; there is no + // separate ListModels RPC. Selection therefore shares GetEngineInfo's + // bounded deadline above. + let model = timeout_discovery_rpc( + "GetModelInfo", + deadline_at, + client.get_model_info(pb::GetModelInfoRequest { + model: selected_model.clone(), + }), + ) + .await? + .map_err(|status| status_to_dynamo("GetModelInfo", status))? + .into_inner(); + if model.model_id != selected_model { + return Err(invalid_arg(format!( + "OpenEngine GetModelInfo returned model_id `{}` for selected model `{selected_model}`", + model.model_id + ))); + } + let connector = timeout_discovery_rpc( + "GetKvConnectorInfo", + deadline_at, + client.get_kv_connector_info(pb::GetKvConnectorInfoRequest {}), + ) + .await? + .map_err(|status| status_to_dynamo("GetKvConnectorInfo", status))? + .into_inner(); + engine.kv_connector = Some(connector); + Ok(Discovery { + engine, + model, + selected_model, + }) +} + +async fn timeout_discovery_rpc( + rpc: &str, + deadline_at: Instant, + future: F, +) -> Result +where + F: std::future::Future, +{ + let remaining = deadline_at.saturating_duration_since(Instant::now()); + tokio::time::timeout(remaining, future).await.map_err(|_| { + cannot_connect(format!( + "OpenEngine bootstrap {rpc} exceeded the configured discovery deadline" + )) + }) +} + +fn select_model(models: &[String], requested: Option<&str>) -> Result { + if let Some(requested) = requested { + if models.iter().any(|model| model == requested) { + return Ok(requested.to_string()); + } + return Err(invalid_arg(format!( + "requested model `{requested}` is not advertised by OpenEngine (available: {})", + models.join(", ") + ))); + } + match models { + [only] => Ok(only.clone()), + [] => Err(invalid_arg( + "OpenEngine advertised no models; pass --model only after fixing server discovery", + )), + _ => Err(invalid_arg(format!( + "OpenEngine advertises multiple models ({}); select one with --model", + models.join(", ") + ))), + } +} + +fn validate_schema(engine: &pb::EngineInfo) -> Result<(), DynamoError> { + let client_revision = openengine_proto::SCHEMA_REVISION; + if engine.schema_revision == 0 { + return Err(invalid_arg("OpenEngine reported invalid schema revision 0")); + } + if engine.minimum_client_revision > client_revision { + return Err(invalid_arg(format!( + "OpenEngine requires client schema revision {}, but this sidecar supports revision {client_revision}", + engine.minimum_client_revision + ))); + } + if engine.schema_revision > client_revision { + return Err(invalid_arg(format!( + "OpenEngine schema revision {} is newer than this sidecar's tested revision {client_revision}", + engine.schema_revision, + ))); + } + const COMPILED_OPENENGINE_COMMIT: &str = env!("OPENENGINE_PROTO_COMMIT"); + if engine.schema_release != COMPILED_OPENENGINE_COMMIT { + return Err(invalid_arg(format!( + "OpenEngine schema_release `{}` does not match the sidecar's pinned contract commit `{COMPILED_OPENENGINE_COMMIT}`", + engine.schema_release + ))); + } + Ok(()) +} + +fn backend(kind: BackendError, message: impl Into) -> DynamoError { + DynamoError::builder() + .error_type(ErrorType::Backend(kind)) + .message(message) + .build() +} + +pub fn invalid_arg(message: impl Into) -> DynamoError { + backend(BackendError::InvalidArgument, message) +} + +pub fn engine_shutdown(message: impl Into) -> DynamoError { + backend(BackendError::EngineShutdown, message) +} + +pub fn cannot_connect(message: impl Into) -> DynamoError { + backend(BackendError::CannotConnect, message) +} + +pub fn status_to_dynamo(rpc: &str, status: tonic::Status) -> DynamoError { + let kind = match status.code() { + tonic::Code::InvalidArgument + | tonic::Code::NotFound + | tonic::Code::OutOfRange + | tonic::Code::FailedPrecondition + | tonic::Code::AlreadyExists + | tonic::Code::Unimplemented => BackendError::InvalidArgument, + tonic::Code::Unavailable => BackendError::CannotConnect, + tonic::Code::Cancelled => BackendError::Cancelled, + tonic::Code::DeadlineExceeded => BackendError::ConnectionTimeout, + _ => BackendError::Unknown, + }; + backend( + kind, + format!("{rpc}: {} ({:?})", status.message(), status.code()), + ) +} + +pub fn engine_error_to_dynamo(error: &pb::EngineError) -> DynamoError { + let code = pb::ErrorCode::try_from(error.code).unwrap_or(pb::ErrorCode::Unspecified); + let kind = match code { + pb::ErrorCode::InvalidArgument + | pb::ErrorCode::UnsupportedFeature + | pb::ErrorCode::RoleMismatch + | pb::ErrorCode::ModelNotFound + | pb::ErrorCode::KvSessionNotFound + | pb::ErrorCode::RequestNotFound + | pb::ErrorCode::DuplicateRequest => BackendError::InvalidArgument, + pb::ErrorCode::Cancelled => BackendError::Cancelled, + pb::ErrorCode::Draining => BackendError::EngineShutdown, + pb::ErrorCode::KvTransferFailed => BackendError::Disconnected, + pb::ErrorCode::Overloaded | pb::ErrorCode::Internal | pb::ErrorCode::Unspecified => { + BackendError::Unknown + } + }; + backend(kind, format!("OpenEngine [{code:?}]: {}", error.message)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn requires_model_for_multi_model_server() { + let models = vec!["a".to_string(), "b".to_string()]; + assert!(select_model(&models, None).is_err()); + assert_eq!(select_model(&models, Some("b")).unwrap(), "b"); + } + + #[test] + fn rejects_nonimmutable_schema_release() { + let mut engine = pb::EngineInfo { + schema_revision: openengine_proto::SCHEMA_REVISION, + minimum_client_revision: 1, + schema_release: "unreleased".to_string(), + ..Default::default() + }; + assert!(validate_schema(&engine).is_err()); + engine.schema_release = String::new(); + assert!(validate_schema(&engine).is_err()); + engine.schema_release = "main".to_string(); + assert!(validate_schema(&engine).is_err()); + engine.schema_release = "cea19cb".to_string(); + assert!(validate_schema(&engine).is_err()); + engine.schema_release = env!("OPENENGINE_PROTO_COMMIT").to_string(); + assert!(validate_schema(&engine).is_ok()); + engine.schema_release = "0123456789abcdef0123456789abcdef01234567".to_string(); + assert!(validate_schema(&engine).is_err()); + engine.schema_release = "v0.2.0".to_string(); + assert!(validate_schema(&engine).is_err()); + } +} diff --git a/lib/openengine-sidecar/src/convert.rs b/lib/openengine-sidecar/src/convert.rs new file mode 100644 index 000000000000..920eddf09127 --- /dev/null +++ b/lib/openengine-sidecar/src/convert.rs @@ -0,0 +1,652 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeMap; +use std::collections::{HashMap, VecDeque}; + +use dynamo_backend_common::{ + DynamoError, GuidedDecodingOptions, LLMEngineOutput, PreprocessedRequest, TopLogprob, +}; +use dynamo_llm::protocols::common::preprocessor::MultimodalData; + +use crate::client; +use crate::proto as pb; + +pub fn merge_context_metadata( + request: &mut pb::GenerateRequest, + metadata: &BTreeMap, +) { + request.metadata.extend(metadata.clone()); +} + +pub fn build_generate_request( + request: &PreprocessedRequest, + request_id: &str, + model: &str, + is_prefill: bool, +) -> Result { + if request.prompt_embeds.is_some() { + return Err(client::invalid_arg( + "OpenEngine sidecar cannot forward prompt_embeds; use text/token input", + )); + } + if request.encoder_result.is_some() { + return Err(client::invalid_arg( + "OpenEngine sidecar does not support encoder-stage handoffs", + )); + } + + // Context-first multimodal decode still needs the original ordered media + // and processor options (for example, TRT-LLM recomputes mRoPE metadata + // locally). Capability validation in the engine rejects unsupported P/D + // modalities before this conversion runs. + let media = build_media(request)?; + let media_options = build_media_options(request, &media)?; + let sampling = &request.sampling_options; + let stopping = &request.stop_conditions; + let conditions = stopping + .stop + .iter() + .flatten() + .cloned() + .map(|value| pb::StopCondition { + condition: Some(pb::stop_condition::Condition::StopText(value)), + }) + .chain( + stopping + .stop_token_ids + .iter() + .flatten() + .copied() + .map(|value| pb::StopCondition { + condition: Some(pb::stop_condition::Condition::StopTokenId(value)), + }), + ) + .chain( + stopping + .stop_token_ids_hidden + .iter() + .flatten() + .copied() + .map(|value| pb::StopCondition { + condition: Some(pb::stop_condition::Condition::StopTokenId(value)), + }), + ) + .collect(); + + let kv_session = request + .prefill_result + .as_ref() + .map(|result| disagg_json_to_kv_session(&result.disaggregated_params)) + .transpose()?; + let routing = request.routing.as_ref(); + let output = &request.output_options; + + Ok(pb::GenerateRequest { + request_id: request_id.to_string(), + model: model.to_string(), + input: Some(pb::generate_request::Input::TokenIds(pb::TokenIds { + ids: request.token_ids.clone(), + })), + sampling: Some(pb::SamplingParams { + temperature: sampling.temperature.map(f64::from), + top_p: sampling.top_p.map(f64::from), + top_k: sampling.top_k, + 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: sampling.seed.and_then(|seed| u64::try_from(seed).ok()), + num_sequences: sampling.n.map(u32::from), + }), + stopping: Some(pb::StoppingOptions { + max_tokens: if is_prefill { + Some(1) + } else { + stopping.max_tokens + }, + min_tokens: stopping.min_tokens, + conditions, + ignore_eos: stopping.ignore_eos, + include_stop_in_output: sampling.include_stop_str_in_output, + }), + response: Some(pb::ResponseOptions { + return_prompt_logprobs: output.prompt_logprobs.map(|_| true), + prompt_candidates: output.prompt_logprobs.map(top_n_candidates), + return_output_logprobs: output.logprobs.map(|_| true), + output_candidates: output.logprobs.map(top_n_candidates), + prompt_logprob_start: None, + }), + kv: Some(pb::KvOptions { + session: kv_session, + data_parallel_rank: routing.and_then(|value| { + if is_prefill { + value.prefill_dp_rank + } else { + value.dp_rank + } + }), + bypass_prefix_cache: prefix_cache_bypass(request), + cache_salt: routing.and_then(|value| value.cache_namespace.clone()), + }), + guided: sampling + .guided_decoding + .as_ref() + .map(build_guided) + .transpose()?, + media, + lora_name: routing + .and_then(|value| value.lora_name.clone()) + .unwrap_or_default(), + priority: routing.and_then(|value| value.priority), + metadata: request + .annotations + .iter() + .filter_map(|annotation| annotation.split_once(':')) + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect(), + media_options, + }) +} + +fn top_n_candidates(top_n: u32) -> pb::CandidateTokenSelection { + pb::CandidateTokenSelection { + selection: Some(pb::candidate_token_selection::Selection::TopN(top_n)), + } +} + +fn prefix_cache_bypass(request: &PreprocessedRequest) -> Option { + request.extra_args.as_ref().and_then(|args| { + args.get("bypass_prefix_cache") + .and_then(serde_json::Value::as_bool) + .or_else(|| { + args.get("disable_prefix_cache") + .and_then(serde_json::Value::as_bool) + }) + }) +} + +fn build_guided(value: &GuidedDecodingOptions) -> Result { + let guide = if let Some(json) = value.json.as_ref() { + Some(if json.is_null() { + pb::guided_decoding::Guide::JsonObject(pb::JsonObjectConstraint {}) + } else { + pb::guided_decoding::Guide::JsonSchema( + serde_json::to_string(json) + .map_err(|error| client::invalid_arg(format!("guided JSON: {error}")))?, + ) + }) + } else if let Some(regex) = value.regex.as_ref() { + Some(pb::guided_decoding::Guide::Regex(regex.clone())) + } else if let Some(grammar) = value.grammar.as_ref() { + Some(pb::guided_decoding::Guide::EbnfGrammar(grammar.clone())) + } else if let Some(tag) = value.structural_tag.as_ref() { + Some(pb::guided_decoding::Guide::StructuralTag( + serde_json::to_string(tag) + .map_err(|error| client::invalid_arg(format!("structural tag: {error}")))?, + )) + } else { + value.choice.as_ref().map(|choices| { + pb::guided_decoding::Guide::Choice(pb::ChoiceConstraint { + choices: choices.clone(), + }) + }) + }; + Ok(pb::GuidedDecoding { + guide, + backend: value.backend.clone().unwrap_or_default(), + }) +} + +fn modality_for_key(key: &str) -> Option { + match key { + "image_url" => Some(pb::Modality::Image), + "video_url" => Some(pb::Modality::Video), + "audio_url" => Some(pb::Modality::Audio), + _ => None, + } +} + +fn key_for_content_part(value: &serde_json::Value) -> Option<&'static str> { + match value.get("type")?.as_str()? { + "image_url" | "input_image" => Some("image_url"), + "video_url" | "input_video" => Some("video_url"), + "audio_url" | "input_audio" => Some("audio_url"), + _ => None, + } +} + +/// Recover cross-modality order from the original message content preserved by +/// Dynamo's preprocessor, while taking media values from the authoritative map. +fn build_media(request: &PreprocessedRequest) -> Result, DynamoError> { + let Some(map) = request.multi_modal_data.as_ref() else { + return Ok(Vec::new()); + }; + let mut queues: HashMap<&str, VecDeque<&MultimodalData>> = map + .iter() + .filter_map(|(key, values)| { + modality_for_key(key).map(|_| (key.as_str(), values.iter().collect())) + }) + .collect(); + let mut order = Vec::new(); + if let Some(messages) = request + .extra_args + .as_ref() + .and_then(|args| args.get("messages")) + .and_then(serde_json::Value::as_array) + { + for part in messages + .iter() + .filter_map(|message| message.get("content")) + .filter_map(serde_json::Value::as_array) + .flatten() + { + if let Some(key) = key_for_content_part(part) + && queues.get(key).is_some_and(|values| !values.is_empty()) + { + order.push(key); + } + } + } + // Older/pre-tokenized clients may not carry original messages. Retain a + // deterministic fallback and append any entries not represented there. + for key in ["image_url", "video_url", "audio_url"] { + let remaining = queues.get(key).map_or(0, VecDeque::len); + order.extend(std::iter::repeat_n(key, remaining)); + } + + let mut result = Vec::with_capacity(order.len()); + for key in order { + let Some(value) = queues.get_mut(key).and_then(VecDeque::pop_front) else { + continue; + }; + let source = match value { + MultimodalData::Url(value) => media_source(value.as_str()), + MultimodalData::RawUrl(value) => media_source(value), + MultimodalData::Decoded(_) => { + return Err(client::invalid_arg( + "OpenEngine sidecar cannot dereference decoded/RDMA media; configure the frontend for URL/data passthrough", + )); + } + }; + result.push(pb::MediaItem { + modality: modality_for_key(key).expect("filtered modality") as i32, + source: Some(source), + mime_type: String::new(), + uuid: String::new(), + }); + } + Ok(result) +} + +fn media_source(value: &str) -> pb::media_item::Source { + if value.starts_with("data:") { + pb::media_item::Source::DataUri(value.to_string()) + } else { + pb::media_item::Source::Url(value.to_string()) + } +} + +fn build_media_options( + request: &PreprocessedRequest, + media: &[pb::MediaItem], +) -> Result, DynamoError> { + let Some(options) = request.mm_processor_kwargs.as_ref() else { + return Ok(None); + }; + if !options.is_object() { + return Err(client::invalid_arg("mm_processor_kwargs must be an object")); + } + let already_keyed = ["image", "video", "audio"] + .iter() + .any(|key| options.get(key).is_some()); + let value = if already_keyed { + options.clone() + } else { + let mut keyed = serde_json::Map::new(); + for modality in media + .iter() + .filter_map(|item| pb::Modality::try_from(item.modality).ok()) + { + let key = match modality { + pb::Modality::Image | pb::Modality::Unspecified => "image", + pb::Modality::Video => "video", + pb::Modality::Audio => "audio", + }; + keyed + .entry(key.to_string()) + .or_insert_with(|| options.clone()); + } + serde_json::Value::Object(keyed) + }; + json_to_prost_struct(&value) + .map(Some) + .ok_or_else(|| client::invalid_arg("media options must be an object")) +} + +pub fn token_output(value: pb::TokenOutput) -> LLMEngineOutput { + let token_ids = value.tokens.iter().map(|token| token.token_id).collect(); + let tokens = Some( + value + .tokens + .iter() + .map(|token| (!token.token.is_empty()).then(|| token.token.clone())) + .collect(), + ); + let has_logprobs = value.tokens.iter().any(|token| token.logprob.is_some()); + let log_probs = has_logprobs.then(|| { + value + .tokens + .iter() + .map(|token| token.logprob.unwrap_or(f64::NEG_INFINITY)) + .collect() + }); + let top_logprobs = has_logprobs.then(|| { + value + .tokens + .iter() + .map(|token| { + token + .candidates + .iter() + .map(|candidate| TopLogprob { + rank: candidate.rank.unwrap_or(0), + token_id: candidate.token_id, + token: (!candidate.token.is_empty()).then(|| candidate.token.clone()), + logprob: candidate.logprob, + bytes: None, + }) + .collect() + }) + .collect() + }); + LLMEngineOutput { + token_ids, + tokens, + text: (!value.text.is_empty()).then_some(value.text), + log_probs, + top_logprobs, + index: value.output_index, + ..Default::default() + } +} + +pub fn kv_session_to_disagg_json(value: pb::KvSessionRef) -> serde_json::Value { + serde_json::json!({ + "session_id": value.session_id, + "transfer_backend": value.transfer_backend, + "endpoints": value.endpoints.into_iter().map(|endpoint| serde_json::json!({ + "host": endpoint.host, + "port": endpoint.port, + "protocol": endpoint.protocol, + })).collect::>(), + "dp_rank": value.dp_rank, + "attributes_struct": value.attributes_struct.as_ref().map(prost_struct_to_json), + }) +} + +pub fn disagg_json_to_kv_session( + value: &serde_json::Value, +) -> Result { + const PROFILE: &str = "tensorrt_llm.disaggregated_params.v1"; + let object = value.as_object().ok_or_else(|| { + client::invalid_arg("prefill_result.disaggregated_params must be an object") + })?; + let required_string = |key: &str| -> Result { + object + .get(key) + .and_then(serde_json::Value::as_str) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + .ok_or_else(|| { + client::invalid_arg(format!("handoff `{key}` must be a non-empty string")) + }) + }; + let session_id = required_string("session_id")?; + let transfer_backend = required_string("transfer_backend")?; + let dp_rank = object + .get("dp_rank") + .and_then(serde_json::Value::as_u64) + .and_then(|value| u32::try_from(value).ok()) + .ok_or_else(|| client::invalid_arg("handoff `dp_rank` must be a uint32"))?; + let endpoints_value = object + .get("endpoints") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| client::invalid_arg("handoff `endpoints` must be an array"))?; + let mut endpoints = Vec::with_capacity(endpoints_value.len()); + for endpoint in endpoints_value { + let host = endpoint + .get("host") + .and_then(serde_json::Value::as_str) + .filter(|value| !value.is_empty() && value.trim() == *value) + .ok_or_else(|| client::invalid_arg("handoff endpoint host must be non-empty"))?; + let port = endpoint + .get("port") + .and_then(serde_json::Value::as_u64) + .and_then(|value| u16::try_from(value).ok()) + .filter(|value| *value > 0) + .ok_or_else(|| client::invalid_arg("handoff endpoint port must be in 1..=65535"))?; + let protocol = endpoint + .get("protocol") + .and_then(serde_json::Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or_else(|| client::invalid_arg("handoff endpoint protocol must be non-empty"))?; + endpoints.push(pb::KvEndpoint { + host: host.to_owned(), + port: u32::from(port), + protocol: protocol.to_owned(), + }); + } + let attributes = object + .get("attributes_struct") + .and_then(serde_json::Value::as_object) + .ok_or_else(|| client::invalid_arg("handoff `attributes_struct` must be an object"))?; + let encoded_profile = attributes + .get(PROFILE) + .and_then(serde_json::Value::as_str) + .ok_or_else(|| { + client::invalid_arg(format!("handoff is missing string profile `{PROFILE}`")) + })?; + validate_trt_handoff_profile(encoded_profile, dp_rank)?; + + Ok(pb::KvSessionRef { + session_id, + transfer_backend, + endpoints, + dp_rank, + attributes_struct: json_to_prost_struct(&serde_json::Value::Object(attributes.clone())), + }) +} + +fn validate_trt_handoff_profile(encoded: &str, session_dp_rank: u32) -> Result<(), DynamoError> { + let payload: serde_json::Value = serde_json::from_str(encoded) + .map_err(|error| client::invalid_arg(format!("invalid TRT-LLM handoff JSON: {error}")))?; + let payload = payload + .as_object() + .ok_or_else(|| client::invalid_arg("TRT-LLM handoff profile must be a JSON object"))?; + for key in [ + "first_gen_tokens", + "first_gen_log_probs", + "ctx_request_id", + "disagg_request_id", + "ctx_dp_rank", + "ctx_info_endpoint", + "draft_tokens", + "ctx_usage", + "conversation_id", + "schedule_style", + "requires_decode_media", + "opaque_state", + ] { + if !payload.contains_key(key) { + return Err(client::invalid_arg(format!( + "TRT-LLM handoff profile is missing `{key}`" + ))); + } + } + const PROFILE_KEYS: &[&str] = &[ + "first_gen_tokens", + "first_gen_log_probs", + "ctx_request_id", + "disagg_request_id", + "ctx_dp_rank", + "ctx_info_endpoint", + "draft_tokens", + "ctx_usage", + "conversation_id", + "schedule_style", + "requires_decode_media", + "opaque_state", + ]; + if let Some(unknown) = payload + .keys() + .find(|key| !PROFILE_KEYS.contains(&key.as_str())) + { + return Err(client::invalid_arg(format!( + "TRT-LLM handoff contains unsupported field `{unknown}`" + ))); + } + if payload + .get("schedule_style") + .and_then(serde_json::Value::as_str) + != Some("context_first") + { + return Err(client::invalid_arg( + "only context-first TRT-LLM handoffs are supported", + )); + } + for key in ["ctx_request_id", "disagg_request_id"] { + let value = &payload[key]; + if !value.is_null() && !value.as_str().is_some_and(canonical_u64) { + return Err(client::invalid_arg(format!( + "TRT-LLM handoff `{key}` must be null or a decimal string" + ))); + } + } + let ctx_dp_rank = payload["ctx_dp_rank"] + .as_u64() + .and_then(|value| u32::try_from(value).ok()); + if !payload["ctx_dp_rank"].is_null() && ctx_dp_rank.is_none() { + return Err(client::invalid_arg( + "TRT-LLM handoff `ctx_dp_rank` must be null or uint32", + )); + } + if let Some(ctx_dp_rank) = ctx_dp_rank + && ctx_dp_rank != session_dp_rank + { + return Err(client::invalid_arg( + "TRT-LLM handoff ctx_dp_rank does not match KV session dp_rank", + )); + } + for key in ["first_gen_tokens", "first_gen_log_probs", "draft_tokens"] { + if !payload[key].is_null() && !payload[key].is_array() { + return Err(client::invalid_arg(format!( + "TRT-LLM handoff `{key}` must be null or an array" + ))); + } + } + if !payload["ctx_usage"].is_null() && !payload["ctx_usage"].is_object() { + return Err(client::invalid_arg( + "TRT-LLM handoff `ctx_usage` must be null or an object", + )); + } + for key in ["ctx_info_endpoint", "conversation_id"] { + if !payload[key].is_null() && !payload[key].is_string() { + return Err(client::invalid_arg(format!( + "TRT-LLM handoff `{key}` must be null or a string" + ))); + } + } + if !payload["requires_decode_media"].is_boolean() { + return Err(client::invalid_arg( + "TRT-LLM handoff `requires_decode_media` must be boolean", + )); + } + if !payload["opaque_state"].is_null() { + let encoded = payload["opaque_state"].as_str().ok_or_else(|| { + client::invalid_arg("TRT-LLM handoff `opaque_state` must be null or base64") + })?; + let decoded = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, encoded) + .map_err(|_| { + client::invalid_arg("TRT-LLM handoff `opaque_state` must be canonical base64") + })?; + if base64::Engine::encode(&base64::engine::general_purpose::STANDARD, decoded) != encoded { + return Err(client::invalid_arg( + "TRT-LLM handoff `opaque_state` must be canonical base64", + )); + } + } + Ok(()) +} + +fn canonical_u64(value: &str) -> bool { + value + .parse::() + .is_ok_and(|parsed| parsed.to_string() == value) +} + +pub fn json_to_prost_struct(value: &serde_json::Value) -> Option { + let serde_json::Value::Object(fields) = value else { + return None; + }; + Some(prost_types::Struct { + fields: fields + .iter() + .map(|(key, value)| (key.clone(), json_to_prost_value(value))) + .collect(), + }) +} + +fn json_to_prost_value(value: &serde_json::Value) -> prost_types::Value { + use prost_types::value::Kind; + let kind = match value { + serde_json::Value::Null => Kind::NullValue(0), + serde_json::Value::Bool(value) => Kind::BoolValue(*value), + serde_json::Value::Number(value) => Kind::NumberValue(value.as_f64().unwrap_or_default()), + serde_json::Value::String(value) => Kind::StringValue(value.clone()), + serde_json::Value::Array(values) => Kind::ListValue(prost_types::ListValue { + values: values.iter().map(json_to_prost_value).collect(), + }), + serde_json::Value::Object(fields) => Kind::StructValue(prost_types::Struct { + fields: fields + .iter() + .map(|(key, value)| (key.clone(), json_to_prost_value(value))) + .collect(), + }), + }; + prost_types::Value { kind: Some(kind) } +} + +pub fn prost_struct_to_json(value: &prost_types::Struct) -> serde_json::Value { + serde_json::Value::Object( + value + .fields + .iter() + .map(|(key, value)| (key.clone(), prost_value_to_json(value))) + .collect(), + ) +} + +fn prost_value_to_json(value: &prost_types::Value) -> serde_json::Value { + use prost_types::value::Kind; + match value.kind.as_ref() { + None | Some(Kind::NullValue(_)) => serde_json::Value::Null, + Some(Kind::BoolValue(value)) => serde_json::Value::Bool(*value), + Some(Kind::NumberValue(value)) + if value.is_finite() + && value.fract() == 0.0 + && *value >= i64::MIN as f64 + && *value <= i64::MAX as f64 => + { + serde_json::Value::Number((*value as i64).into()) + } + Some(Kind::NumberValue(value)) => serde_json::Number::from_f64(*value) + .map(serde_json::Value::Number) + .unwrap_or(serde_json::Value::Null), + Some(Kind::StringValue(value)) => serde_json::Value::String(value.clone()), + Some(Kind::ListValue(value)) => { + serde_json::Value::Array(value.values.iter().map(prost_value_to_json).collect()) + } + Some(Kind::StructValue(value)) => prost_struct_to_json(value), + } +} diff --git a/lib/openengine-sidecar/src/engine.rs b/lib/openengine-sidecar/src/engine.rs new file mode 100644 index 000000000000..0c2504a44d43 --- /dev/null +++ b/lib/openengine-sidecar/src/engine.rs @@ -0,0 +1,1583 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use async_trait::async_trait; +use dynamo_backend_common::{ + AsyncEngineContext, CompletionUsage, ComponentSnapshot, DisaggregationMode, DynamoError, + EngineConfig, GenerateContext, KvEventSource, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, + LlmRegistration, MetricsBindings, MetricsCtx, ModelInput, PreprocessedRequest, + PromptLogprobEntry, PromptLogprobs, SnapshotPublisher, WorkerConfig, usage, +}; +use dynamo_llm::local_model::LocalModel; +use dynamo_llm::lora::{LoRACache, LoRADownloader, LoRASource, LocalLoRASource, S3LoRASource}; +use dynamo_llm::model_card::LoraInfo; +use dynamo_llm::model_type::ModelType; +use dynamo_llm::protocols::common::preprocessor::MultimodalData; +use dynamo_llm::utils::lora_name_to_id; +use dynamo_llm::worker_type::WorkerType; +use futures::stream::BoxStream; +use parking_lot::Mutex; +use tokio::sync::{OnceCell, watch}; +use tokio::task::JoinHandle; +use tokio::time::{Instant, MissedTickBehavior}; +use tokio_util::sync::CancellationToken; + +use crate::args::{Args, TransportConfig, normalize_endpoint}; +use crate::client::{self, Discovery, Pool}; +use crate::convert; +use crate::kv; +use crate::proto as pb; + +pub struct OpenEngineSidecar { + endpoint: String, + expected_engine: Option, + requested_model: Option, + transport: TransportConfig, + bootstrap: Discovery, + disaggregation_mode: DisaggregationMode, + pool: OnceCell, + discovery: OnceCell, + cancel: CancellationToken, + fatal: watch::Sender>, + background_tasks: Arc>>>, + lora_downloader: Option, + lora_discovery: OnceCell, + loaded_loras: tokio::sync::Mutex>, +} + +struct LoraDiscovery { + endpoint: dynamo_runtime::component::Endpoint, + base_model: LocalModel, + model_type: ModelType, + worker_type: WorkerType, + needs: Vec>, + models: tokio::sync::Mutex>, +} + +impl LoraDiscovery { + async fn attach(&self, adapter: &pb::LoraAdapter) -> Result<(), DynamoError> { + let mut model = self.base_model.clone(); + model + .attach( + &self.endpoint, + self.model_type, + ModelInput::Tokens, + Some(LoraInfo { + name: adapter.lora_name.clone(), + max_gpu_lora_count: None, + }), + Some(self.worker_type), + self.needs.clone(), + ) + .await + .map_err(|error| { + client::engine_shutdown(format!("register LoRA model card: {error}")) + })?; + self.models + .lock() + .await + .insert(adapter.lora_name.clone(), model); + Ok(()) + } + + async fn detach(&self, name: &str) -> Result<(), DynamoError> { + self.models.lock().await.remove(name); + LocalModel::detach_from_endpoint(&self.endpoint, Some(name)) + .await + .map_err(|error| { + client::engine_shutdown(format!("unregister LoRA model card: {error}")) + })?; + Ok(()) + } +} + +impl OpenEngineSidecar { + pub fn from_args(argv: Option>) -> Result<(Self, WorkerConfig), DynamoError> { + let args = match argv { + Some(argv) => ::try_parse_from(argv), + None => ::try_parse(), + } + .map_err(|error| client::invalid_arg(error.to_string()))?; + let endpoint = normalize_endpoint(&args.openengine_endpoint); + let transport = args.transport(); + let bootstrap = bootstrap_discover( + &endpoint, + &transport, + args.model.as_deref(), + args.expected_engine.as_deref(), + )?; + validate_model(&bootstrap)?; + let role = engine_role(&bootstrap)?; + let disaggregation_mode = role_to_mode(role); + let model_name = bootstrap.model.model_id.clone(); + let served_model_name = (!bootstrap.model.served_model_name.is_empty()) + .then(|| bootstrap.model.served_model_name.clone()); + let (fatal, _) = watch::channel(None); + let worker = WorkerConfig { + namespace: args.namespace, + component: component_for_role(role).to_string(), + endpoint: args.endpoint, + endpoint_types: args.endpoint_types, + custom_jinja_template: args.custom_jinja_template, + disaggregation_mode, + model_name, + served_model_name, + model_input: ModelInput::Tokens, + reasoning_parser: nonempty(&bootstrap.model.reasoning_parser), + tool_call_parser: nonempty(&bootstrap.model.tool_call_parser), + // OpenEngine owns fetch/decode/preprocessing; the CPU sidecar must + // receive dereferenceable URL/data inputs. + media_decoder: None, + ..Default::default() + }; + tracing::info!( + %endpoint, + engine = %bootstrap.engine.engine_name, + model = %bootstrap.model.model_id, + ?role, + "OpenEngine sidecar discovery complete" + ); + Ok(( + Self { + endpoint, + expected_engine: args.expected_engine, + requested_model: Some(bootstrap.selected_model.clone()), + transport, + bootstrap, + disaggregation_mode, + pool: OnceCell::new(), + discovery: OnceCell::new(), + cancel: CancellationToken::new(), + fatal, + background_tasks: Arc::new(Mutex::new(Vec::new())), + lora_downloader: build_lora_downloader(), + lora_discovery: OnceCell::new(), + loaded_loras: tokio::sync::Mutex::new(HashMap::new()), + }, + worker, + )) + } + + async fn await_ready(&self, client: &mut client::Client) -> Result<(), DynamoError> { + let deadline = Instant::now() + self.transport.deadline; + loop { + match tokio::time::timeout( + self.transport.connect_timeout, + client.health(health_request()), + ) + .await + { + Ok(Ok(response)) => { + let state = pb::HealthState::try_from(response.into_inner().state) + .unwrap_or(pb::HealthState::Unspecified); + match state { + pb::HealthState::Ready => return Ok(()), + pb::HealthState::Draining => { + return Err(client::engine_shutdown("OpenEngine is draining")); + } + _ => {} + } + } + Ok(Err(error)) => tracing::debug!(%error, "OpenEngine health poll failed"), + Err(_) => tracing::debug!("OpenEngine health poll timed out"), + } + if Instant::now() >= deadline { + return Err(client::engine_shutdown(format!( + "OpenEngine did not become READY within {:?}", + self.transport.deadline + ))); + } + tokio::time::sleep(self.transport.poll_interval).await; + } + } + + #[cfg(test)] + pub(crate) fn enable_local_lora_for_test(&mut self) { + self.lora_downloader = Some(LoRADownloader::new( + vec![Arc::new(LocalLoRASource::new())], + LoRACache::new(std::env::temp_dir().join("dynamo-openengine-lora-tests")), + )); + } + + #[cfg(test)] + pub(crate) async fn lora_card_count(&self) -> usize { + match self.lora_discovery.get() { + Some(discovery) => discovery.models.lock().await.len(), + None => 0, + } + } +} + +#[async_trait] +impl LLMEngine for OpenEngineSidecar { + async fn start(&self, _worker_id: u64) -> Result { + if self.pool.initialized() { + return Err(client::engine_shutdown( + "OpenEngine sidecar already started", + )); + } + let pool = + Pool::connect(&self.endpoint, &self.transport, self.transport.connections).await?; + let mut control = pool.control_client(); + self.await_ready(&mut control).await?; + let discovery = tokio::time::timeout( + self.transport.deadline, + client::discover( + &mut control, + self.requested_model.as_deref(), + self.expected_engine.as_deref(), + self.transport.deadline, + ), + ) + .await + .map_err(|_| client::engine_shutdown("OpenEngine discovery timed out"))??; + validate_model(&discovery)?; + if role_to_mode(engine_role(&discovery)?) != self.disaggregation_mode { + return Err(client::invalid_arg( + "OpenEngine role changed between bootstrap discovery and startup", + )); + } + if discovery.model.model_id != self.bootstrap.model.model_id { + return Err(client::invalid_arg( + "OpenEngine selected model changed between bootstrap discovery and startup", + )); + } + let config = build_engine_config(&discovery); + let connections = pool.len(); + self.pool + .set(pool) + .map_err(|_| client::engine_shutdown("OpenEngine sidecar already started"))?; + self.discovery + .set(discovery) + .map_err(|_| client::engine_shutdown("OpenEngine discovery already initialized"))?; + tracing::info!(connections, model = %config.model, "OpenEngine sidecar started"); + Ok(config) + } + + async fn generate( + &self, + request: PreprocessedRequest, + ctx: GenerateContext, + ) -> Result>, DynamoError> { + let pool = self + .pool + .get() + .ok_or_else(|| client::engine_shutdown("generate called before start"))?; + let discovery = self + .discovery + .get() + .ok_or_else(|| client::engine_shutdown("generate called before discovery"))?; + if self.disaggregation_mode.is_decode() && request.prefill_result.is_none() { + return Err(client::invalid_arg( + "decode worker requires a context-first prefill_result", + )); + } + validate_request_capabilities(discovery, &request, self.disaggregation_mode)?; + + let request_id = ctx.id().to_string(); + let prompt_tokens = request.token_ids.len() as u32; + let requested_outputs = request.sampling_options.n.unwrap_or(1).max(1) as usize; + let is_prefill = self.disaggregation_mode.is_prefill(); + let remote_model = if discovery.model.served_model_name.is_empty() { + discovery.selected_model.as_str() + } else { + discovery.model.served_model_name.as_str() + }; + let mut grpc_request = + convert::build_generate_request(&request, &request_id, remote_model, is_prefill)?; + validate_decode_handoff(discovery, &grpc_request, self.disaggregation_mode)?; + convert::merge_context_metadata(&mut grpc_request, ctx.metadata()); + let mut grpc_client = pool.stream_client(); + let abort_client = pool.control_client(); + let cancel = self.cancel.clone(); + let fatal = self.fatal.clone(); + let current_failure = self.fatal.borrow().clone(); + let response_discovery = discovery.clone(); + + Ok(Box::pin(async_stream::stream! { + let mut abort_guard = AbortOnDrop::new(request_id.clone(), abort_client); + if let Some(reason) = current_failure { + yield Err(client::engine_shutdown(reason)); + return; + } + if ctx.is_stopped() || cancel.is_cancelled() { + yield Ok(LLMEngineOutput::cancelled().with_usage(usage(prompt_tokens, 0))); + return; + } + let opened = tokio::select! { + biased; + _ = ctx.stopped() => None, + _ = cancel.cancelled() => None, + response = grpc_client.generate(grpc_request) => Some(response), + }; + let Some(opened) = opened else { + yield Ok(LLMEngineOutput::cancelled().with_usage(usage(prompt_tokens, 0))); + return; + }; + let mut stream = match opened { + Ok(response) => response.into_inner(), + Err(status) => { + let error = client::status_to_dynamo("Generate", status); + yield Err(error); + return; + } + }; + let mut generated = 0u32; + let mut effective_prompt_tokens = prompt_tokens; + let mut finished_outputs = HashSet::new(); + let mut prompt_logprobs: Option = None; + let mut prompt_seen = false; + loop { + let message = tokio::select! { + biased; + _ = ctx.stopped() => { + yield Ok(LLMEngineOutput::cancelled().with_usage(usage(effective_prompt_tokens, generated))); + break; + } + _ = cancel.cancelled() => { + yield Ok(LLMEngineOutput::cancelled().with_usage(usage(effective_prompt_tokens, generated))); + break; + } + message = stream.message() => message, + }; + match message { + Ok(Some(response)) => { + if response.request_id != request_id { + let error = client::engine_shutdown(format!( + "OpenEngine returned request_id `{}` on stream `{request_id}`", + response.request_id + )); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let terminal_usage = response.usage.as_ref().map(openengine_usage); + if let Some(value) = response.usage.as_ref() { + effective_prompt_tokens = value.prompt_tokens; + } + match response.event { + Some(pb::generate_response::Event::Prompt(prompt)) => { + if response.usage.is_some() || prompt_seen { + let error = client::engine_shutdown("OpenEngine returned duplicate/misplaced PromptOutput"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + prompt_seen = true; + prompt_logprobs = Some(prompt_logprobs_from_openengine(prompt.tokens)); + } + Some(pb::generate_response::Event::Token(token)) => { + if is_prefill || response.usage.is_some() { + let error = client::engine_shutdown("OpenEngine returned TokenOutput for a prefill role or with terminal usage"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let Some(output_index) = token.output_index else { + let error = client::engine_shutdown("OpenEngine TokenOutput omitted required output_index"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + }; + if output_index as usize >= requested_outputs || finished_outputs.contains(&output_index) { + let error = client::engine_shutdown(format!("OpenEngine returned invalid TokenOutput index {output_index}")); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + generated = generated.saturating_add(token.tokens.len() as u32); + yield Ok(convert::token_output(token)); + } + Some(pb::generate_response::Event::PrefillReady(prefill)) => { + if !is_prefill { + let error = client::engine_shutdown("OpenEngine returned PrefillReady for a non-prefill role"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let Some(session) = prefill.kv_session else { + let error = client::engine_shutdown("OpenEngine PrefillReady omitted required kv_session"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + }; + let Some(usage) = terminal_usage else { + let error = client::engine_shutdown("OpenEngine PrefillReady omitted required terminal usage"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + }; + let params = convert::kv_session_to_disagg_json(session); + let validated_session = match convert::disagg_json_to_kv_session(¶ms) { + Ok(session) => session, + Err(error) => { + let error = client::engine_shutdown(format!("OpenEngine returned an invalid PrefillReady handoff: {error}")); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + }; + if let Err(error) = validate_prefill_handoff(&response_discovery, &validated_session) { + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let mut output = LLMEngineOutput::stop(); + output.completion_usage = Some(usage); + output.disaggregated_params = Some(params); + attach_prompt_logprobs(&mut output, prompt_logprobs.take()); + abort_guard.complete(); + yield Ok(output); + break; + } + Some(pb::generate_response::Event::Finished(finished)) => { + if is_prefill { + let error = client::engine_shutdown("OpenEngine returned GenerationFinished for a prefill role"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let Some(output_index) = finished.output_index else { + let error = client::engine_shutdown("OpenEngine GenerationFinished omitted required output_index"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + }; + if output_index as usize >= requested_outputs || !finished_outputs.insert(output_index) { + let error = client::engine_shutdown(format!("OpenEngine returned invalid/duplicate terminal index {output_index}")); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let reason = pb::FinishReason::try_from(finished.reason).unwrap_or(pb::FinishReason::Unspecified); + if reason == pb::FinishReason::Unspecified { + let error = client::engine_shutdown("OpenEngine GenerationFinished used unspecified finish reason"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let all_outputs_finished = finished_outputs.len() >= requested_outputs; + if response.usage.is_some() != all_outputs_finished { + let error = client::engine_shutdown("OpenEngine usage must appear only on the final output terminal"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + let mut output = finish_output(reason, finished.stop_match); + output.completion_usage = terminal_usage; + output.index = finished.output_index; + if all_outputs_finished { + attach_prompt_logprobs(&mut output, prompt_logprobs.take()); + abort_guard.complete(); + } + yield Ok(output); + if all_outputs_finished { break; } + } + Some(pb::generate_response::Event::Error(error)) => { + abort_guard.complete(); + yield Err(client::engine_error_to_dynamo(&error)); + break; + } + None => {} + } + } + Ok(None) => { + let error = client::engine_shutdown("OpenEngine closed Generate before a terminal event"); + let _ = fatal.send(Some(error.message().to_string())); + yield Err(error); + break; + } + Err(status) => { + let error = client::status_to_dynamo("Generate stream", status); + yield Err(error); + break; + } + } + } + })) + } + + async fn abort(&self, ctx: Arc) { + if let Some(pool) = self.pool.get() { + abort_request(pool.control_client(), ctx.id().to_string()).await; + } + } + + async fn cleanup(&self) -> Result<(), DynamoError> { + self.cancel.cancel(); + if let Some(discovery) = self.lora_discovery.get() { + let names = self + .loaded_loras + .lock() + .await + .keys() + .cloned() + .collect::>(); + for name in names { + if let Err(error) = discovery.detach(&name).await { + tracing::warn!(%error, %name, "failed to detach LoRA card during cleanup"); + } + } + } + let tasks = std::mem::take(&mut *self.background_tasks.lock()); + for task in tasks { + task.abort(); + } + Ok(()) + } + + async fn watch(&self) -> Result<(), DynamoError> { + let pool = self + .pool + .get() + .ok_or_else(|| client::engine_shutdown("watch called before start"))?; + let mut grpc_client = pool.control_client(); + let mut fatal = self.fatal.subscribe(); + let mut interval = tokio::time::interval(self.transport.poll_interval); + interval.set_missed_tick_behavior(MissedTickBehavior::Skip); + loop { + tokio::select! { + changed = fatal.changed() => { + if changed.is_err() { + return Err(client::engine_shutdown("fatal watcher closed")); + } + if let Some(reason) = fatal.borrow().clone() { + return Err(client::engine_shutdown(reason)); + } + } + _ = interval.tick() => { + let response = tokio::time::timeout( + self.transport.connect_timeout, + grpc_client.health(health_request()), + ).await.map_err(|_| client::engine_shutdown("OpenEngine Health timed out"))? + .map_err(|status| client::status_to_dynamo("Health", status))?; + let state = pb::HealthState::try_from(response.into_inner().state) + .unwrap_or(pb::HealthState::Unspecified); + if state != pb::HealthState::Ready { + return Err(client::engine_shutdown(format!("OpenEngine health changed to {state:?}"))); + } + } + } + } + } + + async fn is_quiescent(&self) -> Result, DynamoError> { + let Some(pool) = self.pool.get() else { + return Ok(None); + }; + let load = tokio::time::timeout( + self.transport.connect_timeout, + pool.control_client().get_load(pb::GetLoadRequest { + include_per_rank: false, + }), + ) + .await + .map_err(|_| client::engine_shutdown("OpenEngine GetLoad timed out"))? + .map_err(|status| client::status_to_dynamo("GetLoad", status))? + .into_inner(); + Ok(load.running_requests.map(|running| running == 0)) + } + + async fn kv_event_sources(&self) -> Result, DynamoError> { + let pool = self + .pool + .get() + .ok_or_else(|| client::engine_shutdown("kv_event_sources called before start"))?; + let discovery = self + .discovery + .get() + .ok_or_else(|| client::engine_shutdown("kv_event_sources called before discovery"))?; + kv::discover_sources( + pool.channel(), + pool.control_client(), + advertised_dp_ranks(discovery)?, + self.cancel.clone(), + self.background_tasks.clone(), + self.fatal.clone(), + ) + .await + } + + async fn setup_metrics(&self, _ctx: MetricsCtx<'_>) -> Result { + let pool = self + .pool + .get() + .ok_or_else(|| client::engine_shutdown("setup_metrics called before start"))?; + let discovery = self + .discovery + .get() + .ok_or_else(|| client::engine_shutdown("setup_metrics called before discovery"))?; + let parallelism = discovery.engine.parallelism.unwrap_or_default(); + let start = parallelism.data_parallel_start_rank.unwrap_or(0); + let count = parallelism.data_parallel_size.unwrap_or(1).max(1); + let ranks = (start..start + count).collect::>(); + let channel = pool.channel(); + let cancel = self.cancel.clone(); + let interval = self.transport.load_poll_interval; + let tasks = self.background_tasks.clone(); + Ok(MetricsBindings { + dp_ranks: ranks.clone(), + on_publisher_ready: Some(Box::new(move |publisher| { + let task = tokio::spawn(load_loop(channel, publisher, ranks, interval, cancel)); + tasks.lock().push(task); + Ok(()) + })), + }) + } + + async fn drain(&self) -> Result<(), DynamoError> { + let Some(pool) = self.pool.get() else { + return Ok(()); + }; + let deadline = Instant::now() + self.transport.drain_timeout; + let deadline_ms = self + .transport + .drain_timeout + .as_millis() + .min(u32::MAX as u128) as u32; + let mut stream = tokio::time::timeout( + self.transport.drain_timeout, + pool.control_client().drain(pb::DrainRequest { + stop_accepting_new_requests: true, + deadline_ms: Some(deadline_ms), + abort_after_deadline: true, + }), + ) + .await + .map_err(|_| client::engine_shutdown("OpenEngine Drain startup timed out"))? + .map_err(|status| client::status_to_dynamo("Drain", status))? + .into_inner(); + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(client::engine_shutdown("OpenEngine Drain timed out")); + } + let response = tokio::time::timeout(remaining, stream.message()) + .await + .map_err(|_| client::engine_shutdown("OpenEngine Drain stream timed out"))? + .map_err(|status| client::status_to_dynamo("Drain stream", status))?; + let Some(response) = response else { + break; + }; + match response.event { + Some(pb::drain_response::Event::State(value)) + if pb::DrainState::try_from(value).ok() == Some(pb::DrainState::Complete) => + { + return Ok(()); + } + Some(pb::drain_response::Event::Error(error)) => { + return Err(client::engine_error_to_dynamo(&error)); + } + _ => {} + } + } + Err(client::engine_shutdown( + "OpenEngine Drain stream ended without COMPLETE", + )) + } + + async fn supported_updates(&self) -> Result, DynamoError> { + if self.bootstrap.model.supports_lora == Some(true) { + Ok(vec!["load_lora", "unload_lora", "list_loras"] + .into_iter() + .map(str::to_string) + .collect()) + } else { + Ok(Vec::new()) + } + } + + async fn engine_update( + &self, + update: String, + body: serde_json::Value, + ) -> Result { + let pool = self + .pool + .get() + .ok_or_else(|| client::engine_shutdown("engine_update called before start"))?; + let mut grpc_client = pool.control_client(); + match update.as_str() { + "load_lora" => { + let name = required_string(&body, "lora_name")?; + let uri = body + .get("source") + .and_then(|source| source.get("uri")) + .and_then(serde_json::Value::as_str) + .or_else(|| body.get("source_path").and_then(serde_json::Value::as_str)) + .ok_or_else(|| client::invalid_arg("load_lora requires source.uri"))?; + let downloader = self.lora_downloader.as_ref().ok_or_else(|| { + client::invalid_arg( + "LoRA downloading is disabled; set DYN_LORA_ENABLED=true and mount the shared cache", + ) + })?; + let path = downloader.download_if_needed(uri).await.map_err(|error| { + client::invalid_arg(format!("download LoRA `{name}`: {error}")) + })?; + validate_lora_directory(&path)?; + let path = std::path::absolute(&path) + .map_err(|error| client::invalid_arg(format!("resolve LoRA path: {error}")))?; + let id = i64::from(lora_name_to_id(name)); + let requested = pb::LoraAdapter { + lora_id: id, + lora_name: name.to_string(), + source_path: path.to_string_lossy().into_owned(), + }; + let mut loaded = self.loaded_loras.lock().await; + if let Some(existing) = loaded.get(name) { + if existing == &requested { + return Ok(lora_response(Some(existing.clone()), Some(true))); + } + return Err(client::invalid_arg(format!( + "LoRA adapter `{name}` conflicts with the loaded ID/path" + ))); + } + let response = grpc_client + .load_lora(pb::LoadLoraRequest { + adapter: Some(requested.clone()), + }) + .await + .map_err(|status| client::status_to_dynamo("LoadLora", status))? + .into_inner(); + let adapter = response.adapter.unwrap_or(requested); + let discovery = self.lora_discovery.get().ok_or_else(|| { + client::engine_shutdown("LoRA model-card discovery is not initialized") + })?; + if let Err(error) = discovery.attach(&adapter).await { + let _ = grpc_client + .unload_lora(pb::UnloadLoraRequest { + lora_name: adapter.lora_name.clone(), + }) + .await; + return Err(error); + } + loaded.insert(name.to_string(), adapter.clone()); + Ok(lora_response(Some(adapter), Some(response.already_loaded))) + } + "unload_lora" => { + let name = required_string(&body, "lora_name")?; + let mut loaded = self.loaded_loras.lock().await; + if loaded.remove(name).is_none() { + return Err(client::invalid_arg(format!( + "LoRA adapter `{name}` is not loaded" + ))); + } + drop(loaded); + let remote = grpc_client + .unload_lora(pb::UnloadLoraRequest { + lora_name: name.to_string(), + }) + .await; + let detach = self + .lora_discovery + .get() + .ok_or_else(|| { + client::engine_shutdown("LoRA model-card discovery is not initialized") + })? + .detach(name) + .await; + let response = + remote.map_err(|status| client::status_to_dynamo("UnloadLora", status)); + match (response, detach) { + (Ok(response), Ok(())) => { + Ok(lora_response(response.into_inner().adapter, None)) + } + (Err(remote), Ok(())) => Err(remote), + (Ok(_), Err(detach)) => Err(detach), + (Err(remote), Err(detach)) => Err(client::engine_shutdown(format!( + "remote LoRA unload failed: {remote}; local model-card detach also failed: {detach}" + ))), + } + } + "list_loras" => { + let response = grpc_client + .list_loras(pb::ListLorasRequest {}) + .await + .map_err(|status| client::status_to_dynamo("ListLoras", status))? + .into_inner(); + Ok(serde_json::json!({ + "status": "ok", + "adapters": response.adapters.into_iter().map(lora_json).collect::>() + })) + } + _ => Ok( + serde_json::json!({"status": "error", "message": format!("unsupported engine update: {update}")}), + ), + } + } + + async fn health_check_payload(&self) -> Result, DynamoError> { + // A decode canary cannot synthesize a valid engine-owned handoff. + if self.disaggregation_mode.is_decode() || self.disaggregation_mode.is_prefill() { + return Ok(None); + } + Ok(Some(serde_json::json!({ + "token_ids": [1], + "stop_conditions": {"max_tokens": 1, "ignore_eos": true}, + "sampling_options": {"temperature": 0.0} + }))) + } + + async fn on_model_ready( + &self, + endpoint: dynamo_runtime::component::Endpoint, + base_model: LocalModel, + model_type: ModelType, + worker_type: WorkerType, + needs: Vec>, + ) -> Result<(), DynamoError> { + self.lora_discovery + .set(LoraDiscovery { + endpoint, + base_model, + model_type, + worker_type, + needs, + models: tokio::sync::Mutex::new(HashMap::new()), + }) + .map_err(|_| client::engine_shutdown("LoRA discovery already initialized")) + } +} + +async fn load_loop( + channel: tonic::transport::Channel, + publisher: Arc, + ranks: Vec, + interval: std::time::Duration, + cancel: CancellationToken, +) { + let mut grpc_client = pb::open_engine_client::OpenEngineClient::new(channel); + let mut ticker = tokio::time::interval(interval); + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = cancel.cancelled() => break, + _ = ticker.tick() => { + let Ok(Ok(response)) = tokio::time::timeout( + interval, + grpc_client.get_load(pb::GetLoadRequest { include_per_rank: true }), + ).await else { + continue; + }; + let load = response.into_inner(); + if load.ranks.is_empty() { + if ranks.len() == 1 && let Some(&rank) = ranks.first() { + publish_load(&publisher, rank, load.used_kv_blocks, load.total_kv_blocks); + } + } else { + for rank in load.ranks { + let Some(dp_rank) = rank.data_parallel_rank else { continue; }; + if !ranks.contains(&dp_rank) { continue; } + publish_load( + &publisher, + dp_rank, + rank.used_kv_blocks, + rank.total_kv_blocks, + ); + } + } + } + } + } +} + +fn publish_load(publisher: &SnapshotPublisher, rank: u32, used: Option, total: Option) { + let (Some(used), Some(total)) = (used, total) else { + return; + }; + if total == 0 || used > total { + return; + } + publisher.publish( + rank, + ComponentSnapshot { + kv_used_blocks: used, + kv_total_blocks: total, + gpu_cache_usage: used as f32 / total as f32, + kv_cache_hit_rate: None, + dp_rank: rank, + }, + ); +} + +struct AbortOnDrop { + request_id: Option, + client: Option, +} + +impl AbortOnDrop { + fn new(request_id: String, client: client::Client) -> Self { + Self { + request_id: Some(request_id), + client: Some(client), + } + } + + fn complete(&mut self) { + self.request_id = None; + self.client = None; + } +} + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + let (Some(request_id), Some(client)) = (self.request_id.take(), self.client.take()) else { + return; + }; + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(abort_request(client, request_id)); + } + } +} + +async fn abort_request(mut client: client::Client, request_id: String) { + if let Err(error) = client + .abort(pb::AbortRequest { + target: Some(pb::abort_request::Target::RequestId(request_id.clone())), + }) + .await + { + tracing::debug!(%error, %request_id, "OpenEngine Abort failed"); + } +} + +fn bootstrap_discover( + endpoint: &str, + transport: &TransportConfig, + model: Option<&str>, + expected_engine: Option<&str>, +) -> Result { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|error| client::engine_shutdown(format!("bootstrap runtime: {error}")))? + .block_on(async { + let mut grpc_client = client::connect(endpoint, transport).await?; + client::discover(&mut grpc_client, model, expected_engine, transport.deadline).await + }) +} + +fn engine_role(discovery: &Discovery) -> Result { + match pb::EngineRole::try_from(discovery.engine.role) { + Ok( + role @ (pb::EngineRole::Aggregated | pb::EngineRole::Prefill | pb::EngineRole::Decode), + ) => Ok(role), + Ok(pb::EngineRole::Unspecified) | Err(_) => Err(client::invalid_arg(format!( + "OpenEngine advertised unsupported role value {}", + discovery.engine.role + ))), + } +} + +fn role_to_mode(role: pb::EngineRole) -> DisaggregationMode { + match role { + pb::EngineRole::Prefill => DisaggregationMode::Prefill, + pb::EngineRole::Decode => DisaggregationMode::Decode, + _ => DisaggregationMode::Aggregated, + } +} + +fn component_for_role(role: pb::EngineRole) -> &'static str { + if role == pb::EngineRole::Prefill { + "prefill" + } else { + "backend" + } +} + +fn health_request() -> pb::HealthRequest { + pb::HealthRequest { + include_inference_probe: false, + model: String::new(), + role: pb::EngineRole::Unspecified as i32, + } +} + +fn nonempty(value: &str) -> Option { + (!value.is_empty()).then(|| value.to_string()) +} + +fn validate_model(discovery: &Discovery) -> Result<(), DynamoError> { + if discovery.model.supports_token_ids_input != Some(true) { + return Err(client::invalid_arg( + "OpenEngine model must explicitly advertise token-id input required by the Dynamo worker pipeline", + )); + } + validate_role_connector(discovery)?; + Ok(()) +} + +fn validate_role_connector(discovery: &Discovery) -> Result<(), DynamoError> { + let role = engine_role(discovery)?; + if role == pb::EngineRole::Aggregated { + return Ok(()); + } + let connector = discovery.engine.kv_connector.as_ref().ok_or_else(|| { + client::invalid_arg("disaggregated OpenEngine role omitted KV connector discovery") + })?; + if connector.enabled != Some(true) + || connector.supports_abort_cleanup != Some(true) + || connector.supports_drain != Some(true) + || connector.transfer_backend.is_empty() + || connector.supported_protocols.is_empty() + { + return Err(client::invalid_arg( + "disaggregated OpenEngine role requires an enabled, drainable, abort-cleanable KV connector with a backend and protocol", + )); + } + if role == pb::EngineRole::Prefill && connector.supports_remote_prefill != Some(true) { + return Err(client::invalid_arg( + "OpenEngine prefill role does not advertise remote-prefill KV support", + )); + } + if role == pb::EngineRole::Decode && connector.supports_decode_pull != Some(true) { + return Err(client::invalid_arg( + "OpenEngine decode role does not advertise decode-pull KV support", + )); + } + Ok(()) +} + +fn validate_decode_handoff( + discovery: &Discovery, + request: &pb::GenerateRequest, + mode: DisaggregationMode, +) -> Result<(), DynamoError> { + if !mode.is_decode() { + return Ok(()); + } + let session = request + .kv + .as_ref() + .and_then(|kv| kv.session.as_ref()) + .ok_or_else(|| client::invalid_arg("decode request omitted its KV session"))?; + let connector = discovery + .engine + .kv_connector + .as_ref() + .ok_or_else(|| client::invalid_arg("decode engine omitted KV connector discovery"))?; + if session.transfer_backend != connector.transfer_backend { + return Err(client::invalid_arg(format!( + "handoff transfer backend `{}` is incompatible with decode backend `{}`", + session.transfer_backend, connector.transfer_backend + ))); + } + for endpoint in &session.endpoints { + if !connector + .supported_protocols + .iter() + .any(|protocol| protocol == &endpoint.protocol) + { + return Err(client::invalid_arg(format!( + "handoff protocol `{}` is not supported by the decode KV connector", + endpoint.protocol + ))); + } + } + Ok(()) +} + +fn validate_prefill_handoff( + discovery: &Discovery, + session: &pb::KvSessionRef, +) -> Result<(), DynamoError> { + let connector = + discovery.engine.kv_connector.as_ref().ok_or_else(|| { + client::engine_shutdown("prefill engine omitted KV connector discovery") + })?; + if session.transfer_backend != connector.transfer_backend { + return Err(client::engine_shutdown(format!( + "OpenEngine PrefillReady transfer backend `{}` does not match discovered backend `{}`", + session.transfer_backend, connector.transfer_backend + ))); + } + for endpoint in &session.endpoints { + if !connector + .supported_protocols + .iter() + .any(|protocol| protocol == &endpoint.protocol) + { + return Err(client::engine_shutdown(format!( + "OpenEngine PrefillReady protocol `{}` is not supported by the discovered KV connector", + endpoint.protocol + ))); + } + } + if !advertised_dp_ranks(discovery)?.contains(&session.dp_rank) { + return Err(client::engine_shutdown(format!( + "OpenEngine PrefillReady dp_rank {} is outside the advertised parallelism range", + session.dp_rank + ))); + } + Ok(()) +} + +fn advertised_dp_ranks(discovery: &Discovery) -> Result, DynamoError> { + let parallelism = discovery.engine.parallelism.as_ref(); + let start = parallelism + .and_then(|parallelism| parallelism.data_parallel_start_rank) + .unwrap_or(0); + let count = parallelism + .and_then(|parallelism| parallelism.data_parallel_size) + .unwrap_or(1) + .max(1); + let end = start.checked_add(count).ok_or_else(|| { + client::invalid_arg("OpenEngine data-parallel rank range overflows uint32") + })?; + Ok((start..end).collect()) +} + +fn validate_request_capabilities( + discovery: &Discovery, + request: &PreprocessedRequest, + mode: DisaggregationMode, +) -> Result<(), DynamoError> { + let generation = discovery.model.generation.as_ref(); + let sampling = &request.sampling_options; + let stopping = &request.stop_conditions; + let output = &request.output_options; + let n = u32::from(sampling.n.unwrap_or(1).max(1)); + if n > 1 + && generation + .and_then(|capabilities| capabilities.max_num_sequences) + .is_none_or(|maximum| n > maximum) + { + return Err(client::invalid_arg(format!( + "OpenEngine model does not advertise support for {n} output sequences" + ))); + } + if let Some(max_tokens) = stopping.max_tokens + && discovery + .model + .max_output_tokens + .is_some_and(|maximum| max_tokens > maximum) + { + return Err(client::invalid_arg(format!( + "requested max_tokens {max_tokens} exceeds OpenEngine maximum {}", + discovery.model.max_output_tokens.unwrap_or_default() + ))); + } + if sampling + .best_of + .is_some_and(|best_of| best_of != sampling.n.unwrap_or(1)) + || sampling.use_beam_search == Some(true) + || sampling.length_penalty.is_some() + { + return Err(client::invalid_arg( + "OpenEngine revision 2 does not support best_of or beam-search semantics", + )); + } + if sampling.seed.is_some_and(|seed| seed < 0) { + return Err(client::invalid_arg( + "OpenEngine sampling seed must be non-negative", + )); + } + if stopping + .stop_token_ids_visible + .as_ref() + .is_some_and(|ids| !ids.is_empty()) + || stopping.max_thinking_tokens.is_some() + { + return Err(client::invalid_arg( + "OpenEngine revision 2 cannot preserve visible stop-token or thinking-budget semantics", + )); + } + if output.skip_special_tokens == Some(false) + || output.formatted_prompt == Some(true) + || output.return_tokens_as_token_ids == Some(true) + { + return Err(client::invalid_arg( + "OpenEngine revision 2 cannot preserve the requested output formatting semantics", + )); + } + if sampling.include_stop_str_in_output == Some(true) + && generation.and_then(|value| value.supports_stop_in_output) != Some(true) + { + return Err(client::invalid_arg( + "OpenEngine model does not support including stop matches in output", + )); + } + if request.require_reasoning { + return Err(client::invalid_arg( + "OpenEngine revision 2 does not support require_reasoning", + )); + } + if request.bootstrap_info.is_some() { + return Err(client::invalid_arg( + "OpenEngine sidecar does not support Dynamo bootstrap handoffs", + )); + } + if let Some(extra) = request.extra_args.as_ref() { + let object = extra + .as_object() + .ok_or_else(|| client::invalid_arg("extra_args must be an object"))?; + for key in object.keys() { + if !matches!( + key.as_str(), + "messages" | "bypass_prefix_cache" | "disable_prefix_cache" + ) { + return Err(client::invalid_arg(format!( + "OpenEngine sidecar cannot preserve extra_args field `{key}`" + ))); + } + } + } + let routing = request.routing.as_ref(); + if routing.and_then(|routing| routing.priority).is_some() + && generation.and_then(|value| value.supports_priority) != Some(true) + { + return Err(client::invalid_arg( + "OpenEngine model does not support per-request priority", + )); + } + if routing + .and_then(|routing| routing.cache_namespace.as_deref()) + .is_some_and(|namespace| !namespace.is_empty()) + && generation.and_then(|value| value.supports_cache_salt) != Some(true) + { + return Err(client::invalid_arg( + "OpenEngine model does not support cache-salt isolation", + )); + } + if routing + .and_then(|routing| routing.lora_name.as_deref()) + .is_some_and(|name| !name.is_empty()) + && discovery.model.supports_lora != Some(true) + { + return Err(client::invalid_arg( + "OpenEngine model does not support LoRA selection", + )); + } + validate_logprob_capability( + "prompt", + output.prompt_logprobs, + generation.and_then(|value| value.prompt_logprobs.as_ref()), + )?; + validate_logprob_capability( + "output", + output.logprobs, + generation.and_then(|value| value.output_logprobs.as_ref()), + )?; + validate_guided_capability(sampling.guided_decoding.as_ref(), generation)?; + let requests_prefix_cache_bypass = request.extra_args.as_ref().is_some_and(|args| { + args.get("bypass_prefix_cache") + .and_then(serde_json::Value::as_bool) + .or_else(|| { + args.get("disable_prefix_cache") + .and_then(serde_json::Value::as_bool) + }) + == Some(true) + }); + if requests_prefix_cache_bypass + && discovery + .model + .generation + .as_ref() + .and_then(|generation| generation.supports_prefix_cache_bypass) + != Some(true) + { + return Err(client::invalid_arg( + "OpenEngine model does not support per-request prefix-cache bypass", + )); + } + let Some(media) = request.multi_modal_data.as_ref() else { + if request.mm_processor_kwargs.is_some() { + return Err(client::invalid_arg("media options require media inputs")); + } + return Ok(()); + }; + let Some(capabilities) = discovery.model.multimodal_capabilities.as_ref() else { + return Err(client::invalid_arg( + "OpenEngine model does not advertise multimodal support", + )); + }; + let supported = if mode.is_prefill() || mode.is_decode() { + &capabilities.prefill_decode_modalities + } else { + &capabilities.aggregate_modalities + }; + for (key, values) in media { + let modality = match key.as_str() { + "image_url" => pb::Modality::Image, + "video_url" => pb::Modality::Video, + "audio_url" => pb::Modality::Audio, + _ => { + return Err(client::invalid_arg(format!( + "unsupported media key `{key}`" + ))); + } + }; + if !supported.contains(&(modality as i32)) { + return Err(client::invalid_arg(format!( + "OpenEngine model does not support {modality:?} in {mode:?} mode" + ))); + } + for value in values { + let source_type = match value { + MultimodalData::Url(url) if url.as_str().starts_with("data:") => { + pb::MediaSourceType::DataUri + } + MultimodalData::RawUrl(url) if url.starts_with("data:") => { + pb::MediaSourceType::DataUri + } + MultimodalData::Url(_) | MultimodalData::RawUrl(_) => pb::MediaSourceType::Url, + MultimodalData::Decoded(_) => { + return Err(client::invalid_arg( + "OpenEngine sidecar cannot dereference decoded/RDMA media; configure URL/data passthrough", + )); + } + }; + if !capabilities.source_types.contains(&(source_type as i32)) { + return Err(client::invalid_arg(format!( + "OpenEngine model does not support {source_type:?} media sources" + ))); + } + } + } + if request.mm_processor_kwargs.is_some() + && capabilities.supports_per_request_media_options != Some(true) + { + return Err(client::invalid_arg( + "OpenEngine model does not support per-request media options", + )); + } + Ok(()) +} + +fn validate_logprob_capability( + kind: &str, + requested: Option, + capabilities: Option<&pb::LogprobCapabilities>, +) -> Result<(), DynamoError> { + let Some(requested) = requested else { + return Ok(()); + }; + let supported = capabilities.is_some_and(|capabilities| { + capabilities.supported == Some(true) + && capabilities + .candidate_selection_modes + .contains(&(pb::CandidateTokenSelectionMode::TopN as i32)) + && capabilities + .max_top_n + .is_none_or(|maximum| requested <= maximum) + }); + if !supported { + return Err(client::invalid_arg(format!( + "OpenEngine model does not support requested {kind} top-{requested} logprobs" + ))); + } + Ok(()) +} + +fn validate_guided_capability( + requested: Option<&dynamo_backend_common::GuidedDecodingOptions>, + generation: Option<&pb::GenerationCapabilities>, +) -> Result<(), DynamoError> { + let Some(requested) = requested else { + return Ok(()); + }; + if requested + .backend + .as_ref() + .is_some_and(|backend| !backend.is_empty()) + || requested.whitespace_pattern.is_some() + { + return Err(client::invalid_arg( + "OpenEngine does not advertise guided backend overrides or whitespace patterns", + )); + } + let mode = if requested + .json + .as_ref() + .is_some_and(serde_json::Value::is_null) + { + pb::GuidedDecodingMode::JsonObject + } else if requested.json.is_some() { + pb::GuidedDecodingMode::JsonSchema + } else if requested.regex.is_some() { + pb::GuidedDecodingMode::Regex + } else if requested.grammar.is_some() { + pb::GuidedDecodingMode::EbnfGrammar + } else if requested.structural_tag.is_some() { + pb::GuidedDecodingMode::StructuralTag + } else if requested + .choice + .as_ref() + .is_some_and(|choice| !choice.is_empty()) + { + pb::GuidedDecodingMode::Choice + } else { + return Err(client::invalid_arg( + "guided decoding request has no constraint", + )); + }; + let capabilities = generation.and_then(|value| value.guided_decoding.as_ref()); + if capabilities.is_none_or(|capabilities| { + capabilities.supported != Some(true) || !capabilities.modes.contains(&(mode as i32)) + }) { + return Err(client::invalid_arg(format!( + "OpenEngine model does not support guided decoding mode {mode:?}" + ))); + } + Ok(()) +} + +fn build_engine_config(discovery: &Discovery) -> EngineConfig { + let model = &discovery.model; + let parallelism = discovery.engine.parallelism.unwrap_or_default(); + EngineConfig { + model: model.model_id.clone(), + served_model_name: nonempty(&model.served_model_name), + runtime_data: [ + ( + "openengine_engine".to_string(), + serde_json::json!(discovery.engine.engine_name), + ), + ( + "openengine_schema_release".to_string(), + serde_json::json!(discovery.engine.schema_release), + ), + ] + .into_iter() + .collect(), + llm: Some(LlmRegistration { + context_length: model.max_context_length, + kv_cache_block_size: model.kv_block_size, + total_kv_blocks: model.total_kv_blocks, + max_num_seqs: model.max_running_requests, + max_num_batched_tokens: model.max_batched_tokens, + data_parallel_size: parallelism.data_parallel_size, + data_parallel_start_rank: parallelism.data_parallel_start_rank, + bootstrap_host: None, + bootstrap_port: None, + }), + } +} + +fn prompt_logprobs_from_openengine(tokens: Vec) -> PromptLogprobs { + tokens + .into_iter() + .map(|token| { + let mut candidates = HashMap::new(); + if let Some(logprob) = token.logprob { + candidates.insert( + token.token_id, + PromptLogprobEntry { + logprob: logprob as f32, + rank: token.rank, + decoded_token: (!token.token.is_empty()).then_some(token.token), + }, + ); + } + for candidate in token.candidates { + candidates.insert( + candidate.token_id, + PromptLogprobEntry { + logprob: candidate.logprob as f32, + rank: candidate.rank, + decoded_token: (!candidate.token.is_empty()).then_some(candidate.token), + }, + ); + } + (!candidates.is_empty()).then_some(candidates) + }) + .collect() +} + +fn attach_prompt_logprobs(output: &mut LLMEngineOutput, prompt_logprobs: Option) { + if let Some(prompt_logprobs) = prompt_logprobs { + output.engine_data = Some(serde_json::json!({"prompt_logprobs": prompt_logprobs})); + } +} + +fn finish_output(reason: pb::FinishReason, stop_match: Option) -> LLMEngineOutput { + let mut output = match reason { + pb::FinishReason::Length => LLMEngineOutput::length(), + pb::FinishReason::Cancelled => LLMEngineOutput::cancelled(), + _ => LLMEngineOutput::stop(), + }; + output.stop_reason = stop_match.and_then(|stop_match| match stop_match.r#match? { + pb::stop_match::Match::StopTokenId(token_id) + | pb::stop_match::Match::EosTokenId(token_id) => { + Some(dynamo_backend_common::StopReason::Int(i64::from(token_id))) + } + pb::stop_match::Match::StopText(text) => { + Some(dynamo_backend_common::StopReason::String(text)) + } + }); + output +} + +fn openengine_usage(value: &pb::Usage) -> CompletionUsage { + use dynamo_protocols::types::{CompletionTokensDetails, PromptTokensDetails}; + CompletionUsage { + prompt_tokens: value.prompt_tokens, + completion_tokens: value.completion_tokens, + total_tokens: value.total_tokens, + prompt_tokens_details: value.cached_prompt_tokens.map(|cached_tokens| { + PromptTokensDetails { + audio_tokens: None, + cached_tokens: Some(cached_tokens), + } + }), + completion_tokens_details: value.reasoning_tokens.map(|reasoning_tokens| { + CompletionTokensDetails { + reasoning_tokens: Some(reasoning_tokens), + ..Default::default() + } + }), + } +} + +fn required_string<'a>(body: &'a serde_json::Value, key: &str) -> Result<&'a str, DynamoError> { + body.get(key) + .and_then(serde_json::Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or_else(|| client::invalid_arg(format!("`{key}` is required"))) +} + +fn build_lora_downloader() -> Option { + if !dynamo_llm::lora::lora_serving_enabled() { + return None; + } + let cache = LoRACache::from_env() + .map_err(|error| tracing::warn!(%error, "failed to configure LoRA cache")) + .ok()?; + let mut sources: Vec> = vec![Arc::new(LocalLoRASource::new())]; + if let Ok(source) = S3LoRASource::from_env() { + sources.push(Arc::new(source)); + } + Some(LoRADownloader::new(sources, cache)) +} + +fn validate_lora_directory(path: &std::path::Path) -> Result<(), DynamoError> { + if !path.is_dir() || !path.join("adapter_config.json").is_file() { + return Err(client::invalid_arg(format!( + "LoRA path `{}` is not a PEFT adapter directory (adapter_config.json missing)", + path.display() + ))); + } + Ok(()) +} + +fn lora_json(value: pb::LoraAdapter) -> serde_json::Value { + serde_json::json!({ + "lora_id": value.lora_id, + "lora_name": value.lora_name, + "source_path": value.source_path, + }) +} + +fn lora_response( + value: Option, + already_loaded: Option, +) -> serde_json::Value { + serde_json::json!({ + "status": "ok", + "adapter": value.map(lora_json), + "already_loaded": already_loaded, + }) +} diff --git a/lib/openengine-sidecar/src/kv.rs b/lib/openengine-sidecar/src/kv.rs new file mode 100644 index 000000000000..13d6e2179235 --- /dev/null +++ b/lib/openengine-sidecar/src/kv.rs @@ -0,0 +1,504 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::atomic::AtomicU32; + +use dynamo_backend_common::{DynamoError, KvEventPublisher, KvEventSource}; +use dynamo_kv_router::protocols::{ + BlockExtraInfo, BlockMmObjectInfo, ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, + KvCacheRemoveData, KvCacheStoreData, +}; +use dynamo_kv_router::zmq_wire::create_stored_blocks; +use parking_lot::Mutex; +use tokio::sync::watch; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; +use tonic::transport::Channel; + +use crate::client::{self, Client}; +use crate::proto as pb; + +pub async fn discover_sources( + channel: Channel, + mut client: Client, + expected_ranks: HashSet, + cancel: CancellationToken, + tasks: Arc>>>, + fatal: watch::Sender>, +) -> Result, DynamoError> { + let response = client + .get_kv_event_sources(pb::GetKvEventSourcesRequest { + data_parallel_ranks: Vec::new(), + }) + .await + .map_err(|status| client::status_to_dynamo("GetKvEventSources", status))? + .into_inner(); + let mut result = Vec::new(); + let mut seen_ranks = HashSet::new(); + for source in response.sources { + let rank = source.data_parallel_rank.ok_or_else(|| { + client::invalid_arg("OpenEngine KV event source omitted data_parallel_rank") + })?; + if !seen_ranks.insert(rank) { + return Err(client::invalid_arg(format!( + "OpenEngine advertised duplicate KV event source for rank {rank}" + ))); + } + if !expected_ranks.contains(&rank) { + return Err(client::invalid_arg(format!( + "OpenEngine advertised KV event source rank {rank} outside its data-parallel range" + ))); + } + match source.transport.as_str() { + "zmq" => { + let Some(endpoint) = source.endpoint_addr else { + return Err(client::invalid_arg("ZMQ KV source omitted endpoint_addr")); + }; + let protocol = if endpoint.protocol.is_empty() { + "tcp" + } else { + endpoint.protocol.as_str() + }; + result.push(KvEventSource::Zmq { + endpoint: format!("{protocol}://{}:{}", endpoint.host, endpoint.port), + topic: source.topic, + dp_rank: rank, + }); + } + "grpc" => { + let channel = channel.clone(); + let cancel = cancel.clone(); + let tasks = tasks.clone(); + let fatal = fatal.clone(); + result.push(KvEventSource::Push { + dp_rank: rank, + on_ready: Box::new(move |publisher| { + let task = + tokio::spawn(subscribe_loop(channel, rank, publisher, cancel, fatal)); + tasks.lock().push(task); + Ok(()) + }), + }); + } + unsupported => { + tracing::warn!(unsupported, rank, "ignoring unsupported KV event transport") + } + } + } + Ok(result) +} + +async fn subscribe_loop( + channel: Channel, + rank: u32, + publisher: Arc, + cancel: CancellationToken, + fatal: watch::Sender>, +) { + let mut client = pb::open_engine_client::OpenEngineClient::new(channel); + let response = client.subscribe_kv_events(subscription_request(rank)).await; + let mut stream = match response { + Ok(response) => response.into_inner(), + Err(error) => { + tracing::error!(%error, rank, "OpenEngine KV subscription failed"); + let _ = fatal.send(Some(format!( + "OpenEngine KV subscription for rank {rank} failed: {error}" + ))); + return; + } + }; + let warnings = Arc::new(AtomicU32::new(0)); + let mut last_sequence = None; + loop { + let message = tokio::select! { + _ = cancel.cancelled() => break, + message = stream.message() => message, + }; + match message { + Ok(Some(response)) => match response.event { + Some(pb::subscribe_kv_events_response::Event::Batch(batch)) => { + if batch.data_parallel_rank != rank { + let message = format!( + "OpenEngine KV stream for rank {rank} returned batch rank {}", + batch.data_parallel_rank + ); + tracing::error!(%message); + let _ = fatal.send(Some(message)); + return; + } + if !accept_sequence(&mut last_sequence, batch.sequence_number) { + tracing::error!( + rank, + sequence_number = batch.sequence_number, + ?last_sequence, + "OpenEngine KV stream is non-monotonic" + ); + let _ = fatal.send(Some(format!( + "OpenEngine KV stream for rank {rank} is non-monotonic" + ))); + return; + } + if batch.events.is_empty() { + continue; + } + if batch.events.len() != 1 { + let message = format!( + "OpenEngine KV batch {} for rank {rank} contained {} events; revision-2 TRT contract requires at most one", + batch.sequence_number, + batch.events.len() + ); + let _ = fatal.send(Some(message)); + return; + } + let event = batch.events.into_iter().next().expect("checked one event"); + let event = match convert_event(event, rank, batch.sequence_number, &warnings) { + Ok(event) => event, + Err(message) => { + let message = format!( + "invalid OpenEngine KV batch {} for rank {rank}: {message}", + batch.sequence_number + ); + let _ = fatal.send(Some(message)); + return; + } + }; + if let Some(event) = event + && let Err(error) = publisher.publish(event) + { + tracing::debug!(%error, rank, "Dynamo KV publisher closed"); + return; + } + } + Some(pb::subscribe_kv_events_response::Event::Error(error)) => { + tracing::error!(?error.code, %error.message, rank, "OpenEngine KV stream error"); + let _ = fatal.send(Some(format!( + "OpenEngine KV stream error for rank {rank}: {}", + error.message + ))); + return; + } + None => {} + }, + Ok(None) => { + tracing::warn!(rank, "OpenEngine KV subscription ended"); + if !cancel.is_cancelled() { + let _ = fatal.send(Some(format!( + "OpenEngine KV subscription for rank {rank} ended" + ))); + } + return; + } + Err(error) => { + tracing::warn!(%error, rank, "OpenEngine KV subscription transport failed"); + if !cancel.is_cancelled() { + let _ = fatal.send(Some(format!( + "OpenEngine KV subscription for rank {rank} failed: {error}" + ))); + } + return; + } + } + } +} + +fn subscription_request(rank: u32) -> pb::SubscribeKvEventsRequest { + pb::SubscribeKvEventsRequest { + data_parallel_ranks: vec![rank], + // Revision 2 does not advertise snapshot/replay capability. Request + // only the live stream until discovery can negotiate it. + include_snapshot: false, + start_sequence_number: 0, + } +} + +fn convert_event( + value: pb::KvEvent, + rank: u32, + sequence_number: u64, + warnings: &Arc, +) -> Result, String> { + let Some(wire_event) = value.event else { + return Err("KvEvent omitted its event payload".to_string()); + }; + let data = match wire_event { + pb::kv_event::Event::BlockStored(stored) => { + if !is_gpu_medium(stored.medium) { + return Ok(None); + } + if stored.block_size == 0 { + return Err("BlockStored block_size must be non-zero".to_string()); + } + if stored.block_hashes.is_empty() && stored.token_ids.is_empty() { + return Ok(None); + } + if !stored.kv_cache_spec_kind.is_empty() || stored.kv_cache_spec_sliding_window != 0 { + return Err("unsupported KV cache group/spec metadata".to_string()); + } + if stored.lora_id != 0 && stored.lora_name.is_empty() { + return Err("BlockStored with a non-zero LoRA ID omitted its name".to_string()); + } + let expected_tokens = usize::try_from(stored.block_size) + .ok() + .and_then(|size| size.checked_mul(stored.block_hashes.len())) + .ok_or_else(|| "BlockStored token cardinality overflow".to_string())?; + if stored.token_ids.len() != expected_tokens { + return Err(format!( + "BlockStored has {} tokens for {} block(s) of size {}", + stored.token_ids.len(), + stored.block_hashes.len(), + stored.block_size + )); + } + let hashes = stored + .block_hashes + .iter() + .map(block_hash) + .collect::, _>>()?; + let block_mm_infos = trt_mm_infos(&stored.extra_keys, hashes.len())?; + let num_block_tokens = vec![u64::from(stored.block_size); hashes.len()]; + let blocks = create_stored_blocks( + stored.block_size, + &stored.token_ids, + &num_block_tokens, + &hashes, + (!stored.lora_name.is_empty()).then_some(stored.lora_name.as_str()), + None, + warnings, + block_mm_infos.as_deref(), + None, + None, + ); + if blocks.len() != hashes.len() { + return Err( + "BlockStored could not be normalized without dropping blocks".to_string(), + ); + } + KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: stored + .parent_block_hash + .as_ref() + .map(block_hash) + .transpose()? + .map(ExternalSequenceBlockHash), + start_position: None, + blocks, + }) + } + pb::kv_event::Event::BlockRemoved(removed) => { + if !is_gpu_medium(removed.medium) { + return Ok(None); + } + let block_hashes = removed + .block_hashes + .iter() + .map(block_hash) + .collect::, _>>()? + .into_iter() + .map(ExternalSequenceBlockHash) + .collect::>(); + if block_hashes.is_empty() { + return Ok(None); + } + KvCacheEventData::Removed(KvCacheRemoveData { block_hashes }) + } + pb::kv_event::Event::AllBlocksCleared(_) => KvCacheEventData::Cleared, + }; + Ok(Some(KvCacheEvent { + // Keep the producer sequence intact so queue drops remain visible to + // Dynamo's indexer rather than being hidden by a local counter. + event_id: sequence_number, + data, + dp_rank: rank, + })) +} + +fn accept_sequence(last: &mut Option, current: u64) -> bool { + if last.is_some_and(|last| current <= last) { + return false; + } + *last = Some(current); + true +} + +fn is_gpu_medium(value: i32) -> bool { + matches!( + pb::StorageMedium::try_from(value), + Ok(pb::StorageMedium::Gpu | pb::StorageMedium::Unspecified) | Err(_) + ) +} + +fn block_hash(value: &pb::KvBlockHash) -> Result { + if value.encoding != "decimal_int64" { + return Err(format!( + "KV block hash encoding `{}` is not canonical decimal_int64", + value.encoding + )); + } + let text = std::str::from_utf8(&value.value) + .map_err(|_| "KV block hash is not ASCII decimal".to_string())?; + let parsed = text + .parse::() + .map_err(|_| "KV block hash is not a signed int64 decimal".to_string())?; + if parsed.to_string() != text { + return Err("KV block hash is not canonically encoded".to_string()); + } + Ok(parsed as u64) +} + +fn trt_mm_infos( + tuples: &[pb::OpaqueKeyTuple], + block_count: usize, +) -> Result>>, String> { + if tuples.is_empty() { + return Ok(None); + } + let mut infos = vec![None; block_count]; + for tuple in tuples { + let [tag, block_index, mm_hash, start_offset] = tuple.values.as_slice() else { + return Err("TRT multimodal tuple must contain exactly four values".to_string()); + }; + if tag != "trt_mm_v1" { + return Err(format!("unknown KV extra-key tag `{tag}`")); + } + let block_index = parse_canonical_decimal::(block_index, "MM block index")?; + if block_index >= block_count { + return Err("MM block index exceeds BlockStored cardinality".to_string()); + } + let mm_hash = parse_canonical_decimal::(mm_hash, "MM hash")?; + let _start_offset = parse_canonical_decimal::(start_offset, "MM start offset")?; + infos[block_index] + .get_or_insert_with(|| BlockExtraInfo { + mm_objects: Vec::new(), + }) + .mm_objects + .push(BlockMmObjectInfo { + mm_hash, + // TRT supplies a logical-media-item offset, while this field + // represents block-relative ranges. Do not mislabel it; + // Dynamo's KV hash depends only on mm_hash. + offsets: Vec::new(), + }); + } + Ok(Some(infos)) +} + +fn parse_canonical_decimal(value: &str, name: &str) -> Result +where + T: std::str::FromStr + ToString, +{ + let parsed = value + .parse::() + .map_err(|_| format!("{name} is not decimal"))?; + if parsed.to_string() != value { + return Err(format!("{name} is not canonically encoded")); + } + Ok(parsed) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_only_canonical_decimal_int64_hashes() { + assert_eq!( + block_hash(&pb::KvBlockHash { + value: b"-1".to_vec(), + encoding: "decimal_int64".into() + }), + Ok(u64::MAX) + ); + assert!( + block_hash(&pb::KvBlockHash { + value: b"01".to_vec(), + encoding: "decimal_int64".into() + }) + .is_err() + ); + assert!( + block_hash(&pb::KvBlockHash { + value: 42u64.to_le_bytes().to_vec(), + encoding: "bytes".into() + }) + .is_err() + ); + } + + #[test] + fn preserves_gaps_and_rejects_non_monotonic_sequences() { + let warnings = Arc::new(AtomicU32::new(0)); + let cleared = || pb::KvEvent { + event: Some(pb::kv_event::Event::AllBlocksCleared( + pb::AllBlocksCleared {}, + )), + ..Default::default() + }; + assert_eq!( + convert_event(cleared(), 2, 7, &warnings) + .unwrap() + .unwrap() + .event_id, + 7 + ); + assert_eq!( + convert_event(cleared(), 2, 11, &warnings) + .unwrap() + .unwrap() + .event_id, + 11 + ); + let mut last = None; + assert!(accept_sequence(&mut last, 7)); + assert!(accept_sequence(&mut last, 11)); + assert!(!accept_sequence(&mut last, 11)); + assert!(!accept_sequence(&mut last, 10)); + } + + #[test] + fn live_subscription_does_not_request_unsupported_snapshot_or_replay() { + let request = subscription_request(3); + assert_eq!(request.data_parallel_ranks, vec![3]); + assert!(!request.include_snapshot); + assert_eq!(request.start_sequence_number, 0); + } + + #[test] + fn accepts_filtered_nonzero_group_and_named_id_zero_lora_with_mm_metadata() { + let warnings = Arc::new(AtomicU32::new(0)); + let event = pb::KvEvent { + event: Some(pb::kv_event::Event::BlockStored(pb::BlockStored { + block_hashes: vec![pb::KvBlockHash { + value: b"9".to_vec(), + encoding: "decimal_int64".into(), + }], + token_ids: vec![1, 2], + block_size: 2, + lora_id: 0, + lora_name: "vision-lora".into(), + medium: pb::StorageMedium::Gpu as i32, + extra_keys: vec![pb::OpaqueKeyTuple { + values: vec!["trt_mm_v1".into(), "0".into(), "42".into(), "99".into()], + }], + group_idx: 1, + ..Default::default() + })), + ..Default::default() + }; + let converted = convert_event(event, 3, 7, &warnings).unwrap().unwrap(); + let KvCacheEventData::Stored(stored) = converted.data else { + panic!("expected stored event"); + }; + assert_eq!(stored.blocks.len(), 1); + assert_eq!( + stored.blocks[0].mm_extra_info.as_ref().unwrap().mm_objects[0].mm_hash, + 42 + ); + assert_eq!( + stored.blocks[0].mm_extra_info.as_ref().unwrap().mm_objects[0].offsets, + Vec::<(usize, usize)>::new() + ); + assert_eq!(converted.dp_rank, 3); + } +} diff --git a/lib/openengine-sidecar/src/lib.rs b/lib/openengine-sidecar/src/lib.rs new file mode 100644 index 000000000000..3f8751812e12 --- /dev/null +++ b/lib/openengine-sidecar/src/lib.rs @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Engine-neutral Dynamo sidecar for the OpenEngine gRPC contract. + +mod args; +mod client; +mod convert; +mod engine; +mod kv; + +pub use engine::OpenEngineSidecar; + +/// Immutable OpenEngine source commit checked by build.rs for the local +/// development dependency. +pub const OPENENGINE_PROTO_COMMIT: &str = env!("OPENENGINE_PROTO_COMMIT"); + +pub mod proto { + pub use openengine_proto::openengine::v1::*; +} + +#[cfg(test)] +mod tests; diff --git a/lib/openengine-sidecar/src/main.rs b/lib/openengine-sidecar/src/main.rs new file mode 100644 index 000000000000..c1435f774cba --- /dev/null +++ b/lib/openengine-sidecar/src/main.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +fn main() -> anyhow::Result<()> { + let (engine, config) = dynamo_openengine_sidecar::OpenEngineSidecar::from_args(None)?; + dynamo_backend_common::run(Arc::new(engine), config) +} diff --git a/lib/openengine-sidecar/src/tests.rs b/lib/openengine-sidecar/src/tests.rs new file mode 100644 index 000000000000..0424f177a44d --- /dev/null +++ b/lib/openengine-sidecar/src/tests.rs @@ -0,0 +1,1345 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicI32, AtomicU8, AtomicU64, Ordering}; + +use dynamo_backend_common::{GenerateContext, LLMEngine, PreprocessedRequest}; +use dynamo_llm::protocols::common::preprocessor::{MultimodalData, RoutingHints}; +use parking_lot::Mutex; + +use crate::convert; +use crate::proto as pb; + +type GrpcStream = Pin> + Send>>; + +fn request() -> PreprocessedRequest { + PreprocessedRequest::builder() + .model("model".to_string()) + .token_ids(vec![1, 2, 3]) + .stop_conditions(Default::default()) + .sampling_options(Default::default()) + .output_options(Default::default()) + .build() + .unwrap() +} + +fn trt_handoff( + dp_rank: u32, + ctx_request_id: &str, + opaque_state: Option<&str>, +) -> serde_json::Value { + let profile = serde_json::json!({ + "first_gen_tokens": null, + "first_gen_log_probs": null, + "ctx_request_id": ctx_request_id, + "disagg_request_id": "9223372036854775808", + "ctx_dp_rank": dp_rank, + "ctx_info_endpoint": null, + "draft_tokens": null, + "ctx_usage": null, + "conversation_id": null, + "schedule_style": "context_first", + "requires_decode_media": true, + "opaque_state": opaque_state, + }); + serde_json::json!({ + "session_id": "ctx", + "transfer_backend": "tensorrt_llm", + "endpoints": [{"host": "127.0.0.1", "port": 1234, "protocol": "nixl"}], + "dp_rank": dp_rank, + "attributes_struct": { + "tensorrt_llm.disaggregated_params.v1": serde_json::to_string(&profile).unwrap() + } + }) +} + +#[test] +fn request_conversion_forwards_sampling_routing_and_stopping() { + let mut request = request(); + request.sampling_options.temperature = Some(0.25); + request.sampling_options.top_k = Some(7); + request.sampling_options.n = Some(2); + request.stop_conditions.max_tokens = Some(16); + request.stop_conditions.stop = Some(vec!["done".to_string()]); + request.routing = Some(RoutingHints { + dp_rank: Some(3), + priority: Some(9), + lora_name: Some("adapter".to_string()), + cache_namespace: Some("tenant".to_string()), + ..Default::default() + }); + let converted = convert::build_generate_request(&request, "r1", "served", false).unwrap(); + assert_eq!(converted.model, "served"); + assert_eq!(converted.priority, Some(9)); + assert_eq!(converted.lora_name, "adapter"); + assert_eq!(converted.sampling.unwrap().num_sequences, Some(2)); + assert_eq!(converted.stopping.unwrap().max_tokens, Some(16)); + let kv = converted.kv.unwrap(); + assert_eq!(kv.data_parallel_rank, Some(3)); + assert_eq!(kv.cache_salt.as_deref(), Some("tenant")); +} + +#[test] +fn multimodal_order_and_media_options_follow_original_messages() { + let mut request = request(); + request.multi_modal_data = Some(HashMap::from([ + ( + "image_url".to_string(), + vec![MultimodalData::RawUrl("https://host/image.png".to_string())], + ), + ( + "audio_url".to_string(), + vec![MultimodalData::RawUrl( + "data:audio/wav;base64,AAAA".to_string(), + )], + ), + ])); + request.extra_args = Some(serde_json::json!({ + "messages": [{"content": [ + {"type": "audio_url", "audio_url": {"url": "ignored"}}, + {"type": "image_url", "image_url": {"url": "ignored"}} + ]}] + })); + request.mm_processor_kwargs = Some(serde_json::json!({"num_frames": 8})); + + let converted = convert::build_generate_request(&request, "r1", "served", false).unwrap(); + assert_eq!( + pb::Modality::try_from(converted.media[0].modality).unwrap(), + pb::Modality::Audio + ); + assert_eq!( + pb::Modality::try_from(converted.media[1].modality).unwrap(), + pb::Modality::Image + ); + assert!(matches!( + converted.media[0].source, + Some(pb::media_item::Source::DataUri(_)) + )); + let options = convert::prost_struct_to_json(converted.media_options.as_ref().unwrap()); + assert_eq!(options["audio"]["num_frames"], 8); + assert_eq!(options["image"]["num_frames"], 8); +} + +#[test] +fn decoded_media_fails_closed() { + let mut request = request(); + let descriptor = serde_json::from_value(serde_json::json!({ + "nixl_metadata": "opaque-agent-metadata", + "nixl_descriptor": { + "addr": 0, + "size": 1, + "mem_type": "Dram", + "device_id": 0 + }, + "shape": [1], + "dtype": "UINT8", + "metadata": null + })) + .expect("construct a real serialized RDMA media descriptor"); + request.multi_modal_data = Some(HashMap::from([( + "image_url".to_string(), + vec![MultimodalData::Decoded(descriptor)], + )])); + assert!(convert::build_generate_request(&request, "r1", "m", false).is_err()); +} + +#[test] +fn decode_handoff_preserves_media_and_media_options_for_mrope() { + let mut request = request(); + request.multi_modal_data = Some(HashMap::from([( + "image_url".to_string(), + vec![MultimodalData::RawUrl("https://host/image.png".to_string())], + )])); + request.mm_processor_kwargs = Some(serde_json::json!({"min_pixels": 64})); + request.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: trt_handoff(2, "42", None), + prompt_tokens_details: None, + }); + let converted = convert::build_generate_request(&request, "r1", "served", false).unwrap(); + assert_eq!(converted.media.len(), 1); + assert!(converted.media_options.is_some()); + assert_eq!(converted.kv.unwrap().session.unwrap().session_id, "ctx"); +} + +#[test] +fn trt_handoff_preserves_large_ids_and_binary_as_strings() { + let large_id = "18446744073709551615"; + let opaque = "AP8QICo="; + let json = trt_handoff(2, large_id, Some(opaque)); + let restored = convert::disagg_json_to_kv_session(&json).unwrap(); + let restored_attributes = + convert::prost_struct_to_json(restored.attributes_struct.as_ref().unwrap()); + let profile: serde_json::Value = serde_json::from_str( + restored_attributes["tensorrt_llm.disaggregated_params.v1"] + .as_str() + .unwrap(), + ) + .unwrap(); + assert_eq!(profile["ctx_request_id"], large_id); + assert_eq!(profile["opaque_state"], opaque); + assert_eq!(restored.endpoints[0].port, 1234); +} + +#[test] +fn malformed_or_generation_first_trt_handoffs_fail_closed() { + let mut missing_profile = trt_handoff(0, "42", None); + missing_profile["attributes_struct"] = serde_json::json!({}); + assert!(convert::disagg_json_to_kv_session(&missing_profile).is_err()); + + let mut generation_first = trt_handoff(0, "42", None); + let encoded = generation_first["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] + .as_str() + .unwrap(); + let mut profile: serde_json::Value = serde_json::from_str(encoded).unwrap(); + profile["schedule_style"] = serde_json::json!("generation_first"); + generation_first["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] = + serde_json::json!(serde_json::to_string(&profile).unwrap()); + assert!(convert::disagg_json_to_kv_session(&generation_first).is_err()); + + let mut numeric_id = trt_handoff(0, "42", None); + let encoded = numeric_id["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] + .as_str() + .unwrap(); + let mut profile: serde_json::Value = serde_json::from_str(encoded).unwrap(); + profile["ctx_request_id"] = serde_json::json!(u64::MAX); + numeric_id["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] = + serde_json::json!(serde_json::to_string(&profile).unwrap()); + assert!(convert::disagg_json_to_kv_session(&numeric_id).is_err()); + + let mut bad_base64 = trt_handoff(0, "42", None); + let encoded = bad_base64["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] + .as_str() + .unwrap(); + let mut profile: serde_json::Value = serde_json::from_str(encoded).unwrap(); + profile["opaque_state"] = serde_json::json!("not-base64"); + bad_base64["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] = + serde_json::json!(serde_json::to_string(&profile).unwrap()); + assert!(convert::disagg_json_to_kv_session(&bad_base64).is_err()); + + let mut unsupported_logits = trt_handoff(0, "42", None); + let encoded = unsupported_logits["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] + .as_str() + .unwrap(); + let mut profile: serde_json::Value = serde_json::from_str(encoded).unwrap(); + profile["first_gen_logits"] = serde_json::json!([0.5]); + unsupported_logits["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] = + serde_json::json!(serde_json::to_string(&profile).unwrap()); + assert!(convert::disagg_json_to_kv_session(&unsupported_logits).is_err()); + + let mut incoherent_rank = trt_handoff(0, "42", None); + incoherent_rank["dp_rank"] = serde_json::json!(1); + assert!(convert::disagg_json_to_kv_session(&incoherent_rank).is_err()); + + let mut bad_endpoint = trt_handoff(0, "42", None); + bad_endpoint["endpoints"][0]["port"] = serde_json::json!(70000); + assert!(convert::disagg_json_to_kv_session(&bad_endpoint).is_err()); + + let mut empty_endpoints = trt_handoff(0, "42", None); + empty_endpoints["endpoints"] = serde_json::json!([]); + assert!( + convert::disagg_json_to_kv_session(&empty_endpoints).is_ok(), + "endpoint list is optional; entries are validated only when present" + ); +} + +#[test] +fn token_delta_preserves_text_logprobs_and_output_index() { + let output = convert::token_output(pb::TokenOutput { + output_index: Some(4), + tokens: vec![pb::TokenInfo { + token_id: 9, + token: "hello".into(), + logprob: Some(-0.25), + rank: Some(1), + candidates: vec![pb::LogProb { + token_id: 10, + logprob: -1.5, + token: "world".into(), + rank: Some(2), + }], + }], + text: "hello".into(), + }); + assert_eq!(output.index, Some(4)); + assert_eq!(output.token_ids, vec![9]); + assert_eq!(output.text.as_deref(), Some("hello")); + assert_eq!(output.log_probs.as_ref().unwrap(), &vec![-0.25]); + assert_eq!(output.top_logprobs.as_ref().unwrap()[0][0].token_id, 10); +} + +#[test] +fn runtime_trace_metadata_is_merged_into_generate_request() { + let mut converted = convert::build_generate_request(&request(), "r1", "served", false).unwrap(); + convert::merge_context_metadata( + &mut converted, + &BTreeMap::from([ + ("traceparent".to_string(), "00-abcd-1234-01".to_string()), + ("tracestate".to_string(), "vendor=value".to_string()), + ]), + ); + assert_eq!(converted.metadata["traceparent"], "00-abcd-1234-01"); + assert_eq!(converted.metadata["tracestate"], "vendor=value"); +} + +const AGGREGATE: u8 = 0; +const PREFILL: u8 = 1; +const PENDING: u8 = 2; +const MULTI_OUTPUT: u8 = 3; +const PROMPT_LOGPROBS: u8 = 4; +const PREFILL_NO_USAGE: u8 = 5; +const PREFILL_BAD_HANDOFF: u8 = 6; + +struct FakeState { + engine: Mutex, + model: Mutex, + health: AtomicI32, + behavior: AtomicU8, + requests: Mutex>, + aborts: Mutex>, + load: Mutex, + loras: Mutex>, + subscriptions: Mutex>, + discovery_delay_ms: AtomicU64, +} + +impl Default for FakeState { + fn default() -> Self { + Self { + engine: Mutex::new(pb::EngineInfo { + engine_name: "tensorrt_llm".into(), + role: pb::EngineRole::Aggregated as i32, + supported_models: vec!["model".into()], + schema_revision: openengine_proto::SCHEMA_REVISION, + minimum_client_revision: 1, + schema_release: crate::OPENENGINE_PROTO_COMMIT.into(), + parallelism: Some(pb::ParallelismInfo { + data_parallel_size: Some(2), + data_parallel_start_rank: Some(0), + ..Default::default() + }), + kv_connector: Some(pb::KvConnectorInfo { + enabled: Some(true), + transfer_backend: "tensorrt_llm".into(), + supported_protocols: vec!["tcp".into()], + supports_remote_prefill: Some(true), + supports_decode_pull: Some(true), + supports_abort_cleanup: Some(true), + supports_drain: Some(true), + schema_version: Some(1), + ..Default::default() + }), + ..Default::default() + }), + model: Mutex::new(pb::ModelInfo { + model_id: "model".into(), + served_model_name: "model".into(), + supports_token_ids_input: Some(true), + supports_lora: Some(true), + kv_block_size: Some(16), + generation: Some(pb::GenerationCapabilities { + prompt_logprobs: Some(pb::LogprobCapabilities { + supported: Some(true), + candidate_selection_modes: vec![ + pb::CandidateTokenSelectionMode::TopN as i32, + ], + max_top_n: Some(4), + }), + max_num_sequences: Some(4), + ..Default::default() + }), + multimodal_capabilities: Some(pb::MultimodalCapabilities { + aggregate_modalities: vec![pb::Modality::Image as i32], + prefill_decode_modalities: vec![pb::Modality::Image as i32], + source_types: vec![ + pb::MediaSourceType::Url as i32, + pb::MediaSourceType::DataUri as i32, + ], + supports_per_request_media_options: Some(true), + }), + ..Default::default() + }), + health: AtomicI32::new(pb::HealthState::Ready as i32), + behavior: AtomicU8::new(AGGREGATE), + requests: Mutex::new(Vec::new()), + aborts: Mutex::new(Vec::new()), + load: Mutex::new(pb::LoadInfo { + running_requests: Some(0), + used_kv_blocks: Some(1), + total_kv_blocks: Some(10), + ..Default::default() + }), + loras: Mutex::new(HashMap::new()), + subscriptions: Mutex::new(Vec::new()), + discovery_delay_ms: AtomicU64::new(0), + } + } +} + +#[derive(Clone)] +struct FakeOpenEngine(Arc); + +#[tonic::async_trait] +impl pb::open_engine_server::OpenEngine for FakeOpenEngine { + type GenerateStream = GrpcStream; + + async fn generate( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let request = request.into_inner(); + let request_id = request.request_id.clone(); + self.0.requests.lock().push(request); + let behavior = self.0.behavior.load(Ordering::SeqCst); + if matches!(behavior, PREFILL | PREFILL_NO_USAGE | PREFILL_BAD_HANDOFF) { + let mut handoff = trt_handoff(0, "18446744073709551615", Some("AP8QICo=")); + if behavior == PREFILL_BAD_HANDOFF { + let encoded = handoff["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] + .as_str() + .unwrap(); + let mut profile: serde_json::Value = serde_json::from_str(encoded).unwrap(); + profile["schedule_style"] = serde_json::json!("generation_first"); + handoff["attributes_struct"]["tensorrt_llm.disaggregated_params.v1"] = + serde_json::json!(serde_json::to_string(&profile).unwrap()); + } + let attributes = handoff["attributes_struct"].clone(); + let usage = (behavior != PREFILL_NO_USAGE).then_some(pb::Usage { + prompt_tokens: 3, + total_tokens: 3, + ..Default::default() + }); + return Ok(tonic::Response::new(Box::pin(futures::stream::once( + async move { + Ok(pb::GenerateResponse { + request_id: request_id.clone(), + event: Some(pb::generate_response::Event::PrefillReady( + pb::PrefillReady { + kv_session: Some(pb::KvSessionRef { + session_id: request_id, + transfer_backend: "tensorrt_llm".into(), + endpoints: vec![pb::KvEndpoint { + host: "127.0.0.1".into(), + port: 9000, + protocol: "tcp".into(), + }], + dp_rank: 0, + attributes_struct: convert::json_to_prost_struct(&attributes), + }), + }, + )), + usage, + }) + }, + )))); + } + if self.0.behavior.load(Ordering::SeqCst) == PENDING { + return Ok(tonic::Response::new(Box::pin(async_stream::try_stream! { + yield pb::GenerateResponse { + request_id, + event: Some(pb::generate_response::Event::Token(pb::TokenOutput { + output_index: Some(0), + tokens: vec![pb::TokenInfo { token_id: 42, token: "x".into(), ..Default::default() }], + text: "x".into(), + })), + usage: None, + }; + std::future::pending::<()>().await; + }))); + } + if self.0.behavior.load(Ordering::SeqCst) == MULTI_OUTPUT { + return Ok(tonic::Response::new(Box::pin(futures::stream::iter([ + Ok(pb::GenerateResponse { + request_id: request_id.clone(), + event: Some(pb::generate_response::Event::Finished( + pb::GenerationFinished { + output_index: Some(0), + reason: pb::FinishReason::Length as i32, + ..Default::default() + }, + )), + usage: None, + }), + Ok(pb::GenerateResponse { + request_id, + event: Some(pb::generate_response::Event::Finished( + pb::GenerationFinished { + output_index: Some(1), + reason: pb::FinishReason::Stop as i32, + ..Default::default() + }, + )), + usage: Some(pb::Usage { + prompt_tokens: 5, + completion_tokens: 7, + total_tokens: 99, + ..Default::default() + }), + }), + ])))); + } + if self.0.behavior.load(Ordering::SeqCst) == PROMPT_LOGPROBS { + return Ok(tonic::Response::new(Box::pin(futures::stream::iter([ + Ok(pb::GenerateResponse { + request_id: request_id.clone(), + event: Some(pb::generate_response::Event::Prompt(pb::PromptOutput { + tokens: vec![ + pb::TokenInfo { + token_id: 1, + token: "".into(), + ..Default::default() + }, + pb::TokenInfo { + token_id: 2, + token: "hello".into(), + logprob: Some(-0.25), + rank: Some(1), + candidates: vec![pb::LogProb { + token_id: 3, + token: "world".into(), + logprob: -1.5, + rank: Some(2), + }], + }, + ], + })), + usage: None, + }), + Ok(pb::GenerateResponse { + request_id, + 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: 2, + total_tokens: 2, + ..Default::default() + }), + }), + ])))); + } + let responses = vec![ + Ok(pb::GenerateResponse { + request_id: request_id.clone(), + event: Some(pb::generate_response::Event::Token(pb::TokenOutput { + output_index: Some(0), + tokens: vec![pb::TokenInfo { + token_id: 42, + token: "x".into(), + ..Default::default() + }], + text: "x".into(), + })), + usage: None, + }), + Ok(pb::GenerateResponse { + request_id, + 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: 3, + completion_tokens: 1, + total_tokens: 4, + ..Default::default() + }), + }), + ]; + Ok(tonic::Response::new(Box::pin(futures::stream::iter( + responses, + )))) + } + + async fn embed( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Err(tonic::Status::unimplemented("embed")) + } + + async fn classify( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Err(tonic::Status::unimplemented("classify")) + } + + async fn score( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Err(tonic::Status::unimplemented("score")) + } + + async fn get_engine_info( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + let delay_ms = self.0.discovery_delay_ms.load(Ordering::SeqCst); + if delay_ms > 0 { + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + } + Ok(tonic::Response::new(self.0.engine.lock().clone())) + } + + async fn get_model_info( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Ok(tonic::Response::new(self.0.model.lock().clone())) + } + + async fn get_load( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Ok(tonic::Response::new(self.0.load.lock().clone())) + } + + async fn health( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Ok(tonic::Response::new(pb::HealthResponse { + state: self.0.health.load(Ordering::SeqCst), + checks: Vec::new(), + })) + } + + async fn abort( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + if let Some(pb::abort_request::Target::RequestId(request_id)) = request.into_inner().target + { + self.0.aborts.lock().push(request_id); + } + Ok(tonic::Response::new(pb::AbortResponse { + status: pb::AbortStatus::Aborted as i32, + message: String::new(), + })) + } + + type DrainStream = GrpcStream; + + async fn drain( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + if self.0.behavior.load(Ordering::SeqCst) == PENDING { + return Ok(tonic::Response::new(Box::pin(async_stream::try_stream! { + yield pb::DrainResponse { + event: Some(pb::drain_response::Event::State(pb::DrainState::Started as i32)), + ..Default::default() + }; + std::future::pending::<()>().await; + }))); + } + Ok(tonic::Response::new(Box::pin(futures::stream::once( + async { + Ok(pb::DrainResponse { + event: Some(pb::drain_response::Event::State( + pb::DrainState::Complete as i32, + )), + ..Default::default() + }) + }, + )))) + } + + async fn load_lora( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let adapter = request + .into_inner() + .adapter + .ok_or_else(|| tonic::Status::invalid_argument("adapter required"))?; + let mut loras = self.0.loras.lock(); + let already_loaded = loras.get(&adapter.lora_name) == Some(&adapter); + if let Some(existing) = loras.get(&adapter.lora_name) + && existing != &adapter + { + return Err(tonic::Status::already_exists("conflicting adapter")); + } + loras.insert(adapter.lora_name.clone(), adapter.clone()); + Ok(tonic::Response::new(pb::LoadLoraResponse { + adapter: Some(adapter), + already_loaded, + })) + } + + async fn unload_lora( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let adapter = self.0.loras.lock().remove(&request.into_inner().lora_name); + Ok(tonic::Response::new(pb::UnloadLoraResponse { adapter })) + } + + async fn list_loras( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Ok(tonic::Response::new(pb::ListLorasResponse { + adapters: self.0.loras.lock().values().cloned().collect(), + })) + } + + async fn get_kv_connector_info( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Ok(tonic::Response::new( + self.0 + .engine + .lock() + .kv_connector + .clone() + .unwrap_or_default(), + )) + } + + async fn get_kv_event_sources( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Ok(tonic::Response::new(pb::GetKvEventSourcesResponse { + sources: vec![ + pb::KvEventSource { + transport: "zmq".into(), + endpoint_addr: Some(pb::KvEndpoint { + host: "127.0.0.1".into(), + port: 5555, + protocol: "tcp".into(), + }), + topic: "kv".into(), + data_parallel_rank: Some(0), + ..Default::default() + }, + pb::KvEventSource { + transport: "grpc".into(), + data_parallel_rank: Some(1), + ..Default::default() + }, + ], + })) + } + + type SubscribeKvEventsStream = GrpcStream; + + async fn subscribe_kv_events( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + self.0.subscriptions.lock().push(request.into_inner()); + Ok(tonic::Response::new(Box::pin(futures::stream::once( + async { + Ok(pb::SubscribeKvEventsResponse { + event: Some(pb::subscribe_kv_events_response::Event::Batch( + pb::KvEventBatch { + data_parallel_rank: 1, + events: vec![pb::KvEvent { + event: Some(pb::kv_event::Event::AllBlocksCleared( + pb::AllBlocksCleared {}, + )), + ..Default::default() + }], + ..Default::default() + }, + )), + }) + }, + )))) + } + + type SubscribeRuntimeEventsStream = GrpcStream; + + async fn subscribe_runtime_events( + &self, + _: tonic::Request, + ) -> Result, tonic::Status> { + Err(tonic::Status::unimplemented("runtime events")) + } +} + +struct FakeServer { + address: std::net::SocketAddr, + shutdown: tokio::sync::oneshot::Sender<()>, + task: tokio::task::JoinHandle<()>, +} + +impl FakeServer { + async fn start(state: Arc) -> Self { + use tokio_stream::wrappers::TcpListenerStream; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind fake tonic server; these integration tests must never be skipped"); + let address = listener.local_addr().unwrap(); + let (shutdown, shutdown_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(pb::open_engine_server::OpenEngineServer::new( + FakeOpenEngine(state), + )) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }) + .await + .unwrap(); + }); + Self { + address, + shutdown, + task, + } + } + + async fn stop(self) { + let _ = self.shutdown.send(()); + self.task.await.unwrap(); + } +} + +async fn build_sidecar( + address: std::net::SocketAddr, + expected_engine: &str, +) -> Result< + ( + crate::OpenEngineSidecar, + dynamo_backend_common::WorkerConfig, + ), + dynamo_backend_common::DynamoError, +> { + let expected_engine = expected_engine.to_string(); + tokio::task::spawn_blocking(move || { + crate::OpenEngineSidecar::from_args(Some(vec![ + "dynamo-openengine-sidecar".to_string(), + "--openengine-endpoint".to_string(), + address.to_string(), + "--expected-engine".to_string(), + expected_engine, + "--openengine-connections".to_string(), + "1".to_string(), + "--health-poll-interval-secs".to_string(), + "1".to_string(), + "--health-deadline-secs".to_string(), + "5".to_string(), + "--drain-timeout-secs".to_string(), + "1".to_string(), + ])) + }) + .await + .unwrap() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_server_discovery_and_aggregate_stream() { + use futures::StreamExt; + + let state = Arc::new(FakeState::default()); + let server = FakeServer::start(state).await; + let (engine, config) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + assert_eq!(config.model_name, "model"); + let started = engine.start(1).await.unwrap(); + assert_eq!(started.model, "model"); + let outputs = engine + .generate( + request(), + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .unwrap() + .collect::>() + .await; + assert_eq!(outputs.len(), 2); + assert_eq!(outputs[0].as_ref().unwrap().token_ids, vec![42]); + let terminal = outputs[1].as_ref().unwrap(); + assert!(terminal.finish_reason.is_some()); + assert_eq!(terminal.completion_usage.as_ref().unwrap().total_tokens, 4); + let mut bypass = request(); + bypass.extra_args = Some(serde_json::json!({"bypass_prefix_cache": true})); + assert!( + engine + .generate( + bypass, + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .is_err() + ); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_rejects_unadvertised_request_semantics_before_scheduling() { + let state = Arc::new(FakeState::default()); + state.model.lock().supports_lora = Some(false); + let server = FakeServer::start(state.clone()).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + + let mut requests = Vec::new(); + let mut priority = request(); + priority.routing = Some(RoutingHints { + priority: Some(1), + ..Default::default() + }); + requests.push(priority); + let mut salted = request(); + salted.routing = Some(RoutingHints { + cache_namespace: Some("tenant".into()), + ..Default::default() + }); + requests.push(salted); + let mut lora = request(); + lora.routing = Some(RoutingHints { + lora_name: Some("adapter".into()), + ..Default::default() + }); + requests.push(lora); + let mut guided = request(); + guided.sampling_options.guided_decoding = Some(dynamo_backend_common::GuidedDecodingOptions { + regex: Some("a+".into()), + ..Default::default() + }); + requests.push(guided); + let mut too_many = request(); + too_many.sampling_options.n = Some(5); + requests.push(too_many); + let mut too_many_logprobs = request(); + too_many_logprobs.output_options.prompt_logprobs = Some(5); + requests.push(too_many_logprobs); + + for request in requests { + assert!( + engine + .generate( + request, + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .is_err() + ); + } + assert!( + state.requests.lock().is_empty(), + "rejected requests reached Generate" + ); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_rejects_schema_engine_and_role_mismatches() { + let state = Arc::new(FakeState::default()); + let server = FakeServer::start(state.clone()).await; + assert!(build_sidecar(server.address, "vllm").await.is_err()); + state.engine.lock().schema_release = "main".into(); + assert!(build_sidecar(server.address, "tensorrt_llm").await.is_err()); + state.engine.lock().schema_release = crate::OPENENGINE_PROTO_COMMIT.into(); + state.engine.lock().role = pb::EngineRole::Unspecified as i32; + assert!(build_sidecar(server.address, "tensorrt_llm").await.is_err()); + state.engine.lock().role = pb::EngineRole::Aggregated as i32; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + state.engine.lock().role = pb::EngineRole::Decode as i32; + assert!(engine.start(1).await.is_err()); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_prefill_decode_preserves_media_options_and_handoff() { + use futures::StreamExt; + + let state = Arc::new(FakeState::default()); + state.engine.lock().role = pb::EngineRole::Prefill as i32; + state.behavior.store(PREFILL, Ordering::SeqCst); + let server = FakeServer::start(state.clone()).await; + let mut media_request = request(); + media_request.multi_modal_data = Some(HashMap::from([( + "image_url".to_string(), + vec![MultimodalData::RawUrl("https://host/image.png".into())], + )])); + media_request.mm_processor_kwargs = Some(serde_json::json!({"image": {"min_pixels": 64}})); + let (prefill, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + prefill.start(1).await.unwrap(); + let mut output = prefill + .generate( + media_request.clone(), + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .unwrap() + .collect::>() + .await; + let handoff = output.remove(0).unwrap().disaggregated_params.unwrap(); + prefill.cleanup().await.unwrap(); + + state.engine.lock().role = pb::EngineRole::Decode as i32; + state.behavior.store(AGGREGATE, Ordering::SeqCst); + state.requests.lock().clear(); + media_request.prefill_result = Some(dynamo_backend_common::PrefillResult { + disaggregated_params: handoff, + prompt_tokens_details: None, + }); + let (decode, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + decode.start(2).await.unwrap(); + let _ = decode + .generate( + media_request, + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .unwrap() + .collect::>() + .await; + let decoded = state.requests.lock().last().cloned().unwrap(); + assert!(decoded.kv.unwrap().session.is_some()); + assert_eq!(decoded.media.len(), 1); + assert!(decoded.media_options.is_some()); + decode.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_prefill_rejects_missing_usage_and_malformed_server_handoff() { + use futures::StreamExt; + + for behavior in [PREFILL_NO_USAGE, PREFILL_BAD_HANDOFF] { + let state = Arc::new(FakeState::default()); + state.engine.lock().role = pb::EngineRole::Prefill as i32; + state.behavior.store(behavior, Ordering::SeqCst); + let server = FakeServer::start(state).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + let outputs = engine + .generate( + request(), + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .unwrap() + .collect::>() + .await; + assert_eq!(outputs.len(), 1); + assert!(outputs[0].is_err()); + engine.cleanup().await.unwrap(); + server.stop().await; + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_health_load_kv_discovery_and_watch_failure() { + let state = Arc::new(FakeState::default()); + let server = FakeServer::start(state.clone()).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + assert_eq!(engine.is_quiescent().await.unwrap(), Some(true)); + state.load.lock().running_requests = Some(2); + assert_eq!(engine.is_quiescent().await.unwrap(), Some(false)); + let sources = engine.kv_event_sources().await.unwrap(); + assert_eq!(sources.len(), 2); + assert!(matches!( + sources[0], + dynamo_backend_common::KvEventSource::Zmq { dp_rank: 0, .. } + )); + assert!(matches!( + sources[1], + dynamo_backend_common::KvEventSource::Push { dp_rank: 1, .. } + )); + state + .health + .store(pb::HealthState::Draining as i32, Ordering::SeqCst); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(3), engine.watch()) + .await + .unwrap() + .is_err() + ); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_rejects_kv_source_rank_outside_discovered_parallelism() { + let state = Arc::new(FakeState::default()); + state + .engine + .lock() + .parallelism + .as_mut() + .unwrap() + .data_parallel_size = Some(1); + let server = FakeServer::start(state).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + let error = engine + .kv_event_sources() + .await + .err() + .expect("out-of-range KV source must fail discovery"); + assert!( + error + .to_string() + .contains("outside its data-parallel range") + ); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn bootstrap_discovery_rpc_timeout_is_typed_and_bounded() { + let state = Arc::new(FakeState::default()); + state.discovery_delay_ms.store(100, Ordering::SeqCst); + let server = FakeServer::start(state).await; + let transport = crate::args::TransportConfig { + connect_timeout: std::time::Duration::from_secs(1), + poll_interval: std::time::Duration::from_millis(10), + deadline: std::time::Duration::from_secs(1), + drain_timeout: std::time::Duration::from_secs(1), + load_poll_interval: std::time::Duration::from_secs(1), + connections: 1, + }; + let mut grpc_client = crate::client::connect(&format!("http://{}", server.address), &transport) + .await + .unwrap(); + let error = crate::client::discover( + &mut grpc_client, + None, + Some("tensorrt_llm"), + std::time::Duration::from_millis(20), + ) + .await + .unwrap_err(); + assert_eq!( + error.error_type(), + dynamo_backend_common::ErrorType::Backend( + dynamo_backend_common::BackendError::CannotConnect + ) + ); + assert!(error.to_string().contains("GetEngineInfo")); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_blackholed_drain_stream_times_out() { + let state = Arc::new(FakeState::default()); + state.behavior.store(PENDING, Ordering::SeqCst); + let server = FakeServer::start(state).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + let result = tokio::time::timeout(std::time::Duration::from_secs(2), engine.drain()) + .await + .expect("sidecar drain must have its own timeout"); + assert!(result.is_err()); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_dropped_stream_sends_abort() { + use futures::StreamExt; + + let state = Arc::new(FakeState::default()); + state.behavior.store(PENDING, Ordering::SeqCst); + let server = FakeServer::start(state.clone()).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + let mut stream = engine + .generate( + request(), + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .unwrap(); + assert!(stream.next().await.is_some()); + drop(stream); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + if !state.aborts.lock().is_empty() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_forwards_every_output_terminal_and_exact_final_usage() { + use dynamo_backend_common::FinishReason; + use futures::StreamExt; + + let state = Arc::new(FakeState::default()); + state.behavior.store(MULTI_OUTPUT, Ordering::SeqCst); + let server = FakeServer::start(state).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + let mut generate_request = request(); + generate_request.sampling_options.n = Some(2); + let outputs = engine + .generate( + generate_request, + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .unwrap() + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!(outputs.len(), 2); + assert_eq!(outputs[0].index, Some(0)); + assert_eq!(outputs[0].finish_reason, Some(FinishReason::Length)); + assert!(outputs[0].completion_usage.is_none()); + assert_eq!(outputs[1].index, Some(1)); + assert_eq!(outputs[1].finish_reason, Some(FinishReason::Stop)); + let usage = outputs[1].completion_usage.as_ref().unwrap(); + assert_eq!(usage.prompt_tokens, 5); + assert_eq!(usage.completion_tokens, 7); + assert_eq!(usage.total_tokens, 99); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_attaches_exact_prompt_logprobs_to_terminal_engine_data() { + use dynamo_backend_common::PromptLogprobs; + use futures::StreamExt; + + let state = Arc::new(FakeState::default()); + state.behavior.store(PROMPT_LOGPROBS, Ordering::SeqCst); + let server = FakeServer::start(state).await; + let (engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.start(1).await.unwrap(); + let mut generate_request = request(); + generate_request.output_options.prompt_logprobs = Some(2); + let outputs = engine + .generate( + generate_request, + GenerateContext::new(dynamo_backend_common::testing::mock_context(), None), + ) + .await + .unwrap() + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!( + outputs.len(), + 1, + "PromptOutput is metadata, not a Dynamo chunk" + ); + let engine_data = outputs[0].engine_data.as_ref().unwrap(); + assert!(outputs[0].extra_args.is_none()); + let prompt: PromptLogprobs = + serde_json::from_value(engine_data["prompt_logprobs"].clone()).unwrap(); + assert_eq!(prompt.len(), 2); + assert!(prompt[0].is_none()); + let second = prompt[1].as_ref().unwrap(); + assert_eq!(second[&2].decoded_token.as_deref(), Some("hello")); + assert_eq!(second[&3].rank, Some(2)); + engine.cleanup().await.unwrap(); + server.stop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fake_tonic_lora_lifecycle_publishes_and_removes_model_cards() { + use dynamo_llm::local_model::LocalModelBuilder; + use dynamo_llm::model_type::ModelType; + use dynamo_llm::worker_type::WorkerType; + use dynamo_runtime::distributed::DistributedConfig; + use dynamo_runtime::{DistributedRuntime, Runtime}; + + let state = Arc::new(FakeState::default()); + let server = FakeServer::start(state).await; + let (mut engine, _) = build_sidecar(server.address, "tensorrt_llm").await.unwrap(); + engine.enable_local_lora_for_test(); + engine.start(1).await.unwrap(); + + let runtime = Runtime::from_current().unwrap(); + let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) + .await + .unwrap(); + let endpoint = drt + .namespace("lora-test") + .unwrap() + .component("backend") + .unwrap() + .endpoint("generate"); + let base_model = LocalModelBuilder::default() + .model_name(Some("model".into())) + .build() + .await + .unwrap(); + engine + .on_model_ready( + endpoint, + base_model, + ModelType::Chat | ModelType::Completions, + WorkerType::Aggregated, + Vec::new(), + ) + .await + .unwrap(); + + let adapter_dir = std::env::temp_dir().join(format!( + "dynamo-openengine-lora-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&adapter_dir).unwrap(); + std::fs::write(adapter_dir.join("adapter_config.json"), "{}").unwrap(); + let load = serde_json::json!({ + "lora_name": "adapter", + "source": {"uri": format!("file://{}", adapter_dir.display())} + }); + let first = engine + .engine_update("load_lora".into(), load.clone()) + .await + .unwrap(); + assert_eq!(first["status"], "ok"); + assert_eq!(engine.lora_card_count().await, 1); + let repeated = engine + .engine_update("load_lora".into(), load) + .await + .unwrap(); + assert_eq!(repeated["already_loaded"], true); + let listed = engine + .engine_update("list_loras".into(), serde_json::Value::Null) + .await + .unwrap(); + assert_eq!(listed["adapters"].as_array().unwrap().len(), 1); + engine + .engine_update( + "unload_lora".into(), + serde_json::json!({"lora_name": "adapter"}), + ) + .await + .unwrap(); + assert_eq!(engine.lora_card_count().await, 0); + engine.cleanup().await.unwrap(); + std::fs::remove_dir_all(adapter_dir).unwrap(); + server.stop().await; +} From a73cbc2642ee22677c42f5db06cceeb66fdeaf14 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 23:05:15 -0700 Subject: [PATCH 02/20] fix(llm): preserve adapter and media semantics Signed-off-by: Connor Carpenter --- lib/llm/src/local_model.rs | 8 ++ lib/llm/src/preprocessor.rs | 209 +++++++++++++++++++++++++++++++++++- 2 files changed, 216 insertions(+), 1 deletion(-) diff --git a/lib/llm/src/local_model.rs b/lib/llm/src/local_model.rs index f88103f8dc60..cee4ee03f0f3 100644 --- a/lib/llm/src/local_model.rs +++ b/lib/llm/src/local_model.rs @@ -575,6 +575,14 @@ impl LocalModel { worker_type: Option, needs: Vec>, ) -> anyhow::Result<()> { + // LoRA adapters are independent public model names even though their + // metadata and tokenizer are inherited from the base model. Callers + // may therefore clone a base LocalModel before attaching an adapter; + // make sure discovery does not merge that card back into the base + // model's entry. + if let Some(info) = &lora_info { + self.card.set_name(&info.name); + } self.card.model_type = model_type; self.card.model_input = model_input; self.card.worker_type = worker_type; diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index b6764fa68754..035583fbd645 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -1129,12 +1129,103 @@ impl OpenAIPreprocessor { } else { self.formatter.render(request)? }; - Ok(Some(formatted_prompt)) + Ok(Some(Self::repair_numbered_media_placeholders( + formatted_prompt, + request.typed_messages(), + ))) } else { Ok(None) } } + /// Repair string-content templates that expose a single numbered image + /// placeholder for every media item. The renderer flattens audio/video + /// parts through that image template because it does not retain + /// per-modality placeholder maps on this path. Rewrite only when the full + /// expected numbered sequence is present; otherwise preserve the rendered + /// prompt byte-for-byte. + fn repair_numbered_media_placeholders( + prompt: String, + messages: Option<&[ChatCompletionRequestMessage]>, + ) -> String { + let Some(messages) = messages else { + return prompt; + }; + let mut replacements = Vec::new(); + let mut expected_sources = HashMap::::new(); + let mut literal_text = Vec::new(); + let mut needs_repair = false; + let (mut image, mut video, mut audio) = (0_u32, 0_u32, 0_u32); + for message in messages { + let ChatCompletionRequestMessage::User(user) = message else { + continue; + }; + let ChatCompletionRequestUserMessageContent::Array(parts) = &user.content else { + continue; + }; + // String-content renderers restart their synthetic image slot for + // each message, while TRT-LLM indexes each target modality over + // the complete conversation. + let mut slot = 0_u32; + for part in parts { + let target = match part { + ChatCompletionRequestUserMessageContentPart::ImageUrl(_) => { + slot += 1; + image += 1; + format!("<|image_{image}|>") + } + ChatCompletionRequestUserMessageContentPart::VideoUrl(_) => { + slot += 1; + video += 1; + needs_repair = true; + format!("<|video_{video}|>") + } + ChatCompletionRequestUserMessageContentPart::AudioUrl(_) => { + slot += 1; + audio += 1; + needs_repair = true; + format!("<|audio_{audio}|>") + } + ChatCompletionRequestUserMessageContentPart::Text(text) => { + literal_text.push(text.text.as_str()); + continue; + } + _ => continue, + }; + let source = format!("<|image_{slot}|>"); + *expected_sources.entry(source.clone()).or_default() += 1; + replacements.push((source, target)); + } + } + if !needs_repair || replacements.is_empty() { + return prompt; + } + + // A user may include placeholder-looking text of their own. Without + // renderer provenance, an extra occurrence is ambiguous; preserve the + // prompt instead of rewriting user content or a mismatched slot. + if expected_sources.iter().any(|(source, expected)| { + prompt.match_indices(source).count() != *expected + || literal_text.iter().any(|text| text.contains(source)) + }) { + return prompt; + } + + let mut repaired = String::with_capacity(prompt.len()); + let mut cursor = 0; + for (source, target) in replacements { + let Some(relative) = prompt[cursor..].find(&source) else { + return prompt; + }; + let start = cursor + relative; + repaired.push_str(&prompt[cursor..start]); + repaired.push_str(&target); + cursor = start + source.len(); + } + repaired.push_str(&prompt[cursor..]); + repaired + } + /// Replace inline `data:` URLs with empty strings in message content parts. /// Preserves HTTP(S) URLs, text content, and overall message structure. fn strip_inline_data_urls(messages: &mut serde_json::Value) { @@ -3563,6 +3654,7 @@ impl #[cfg(test)] mod strip_tests { use super::OpenAIPreprocessor; + use dynamo_protocols::types::ChatCompletionRequestMessage; #[test] fn test_strip_inline_data_urls_replaces_data_urls() { @@ -3616,6 +3708,121 @@ mod strip_tests { OpenAIPreprocessor::strip_inline_data_urls(&mut messages); assert_eq!(messages, serde_json::json!([])); } + + #[test] + fn test_repair_numbered_media_placeholders_preserves_modality_indexes() { + let messages: Vec = + serde_json::from_value(serde_json::json!([{ + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,AAAA"}}, + {"type": "text", "text": "compare"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + {"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,BBBB"}} + ] + }])) + .unwrap(); + let rendered = "<|user|><|image_1|>compare<|image_2|><|image_3|><|end|>"; + assert_eq!( + OpenAIPreprocessor::repair_numbered_media_placeholders( + rendered.to_string(), + Some(&messages), + ), + "<|user|><|audio_1|>compare<|image_1|><|audio_2|><|end|>" + ); + } + + #[test] + fn test_repair_numbered_media_placeholders_fails_closed() { + let messages: Vec = + serde_json::from_value(serde_json::json!([{ + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,AAAA"}} + ] + }])) + .unwrap(); + let rendered = "<|user|>