From d775d347cae3ad5aa2d404d7b48cc4cc5db2644b Mon Sep 17 00:00:00 2001 From: "Dr. Stefan Schimanski" Date: Tue, 30 Jun 2026 13:07:17 +0200 Subject: [PATCH 001/320] docs(kubernetes): add GAIE quickstart (#10956) Signed-off-by: Dr. Stefan Schimanski --- README.md | 17 +- deploy/inference-gateway/README.md | 2 +- docs/design-docs/architecture.md | 8 +- docs/getting-started/introduction.md | 16 +- docs/getting-started/kubernetes-deployment.md | 3 +- docs/index.yml | 12 +- docs/kubernetes/README.md | 6 +- docs/kubernetes/dgdr.md | 4 +- docs/kubernetes/gateway-api/README.mdx | 284 +++++++ docs/kubernetes/gateway-api/quickstart.mdx | 364 ++++++++ docs/kubernetes/gateway-api/reference.mdx | 281 +++++++ docs/kubernetes/inference-gateway.md | 776 ------------------ fern/docs.yml | 4 +- .../qwen3-0.6b/vllm/agg/gaie/httproute.yaml | 2 - 14 files changed, 973 insertions(+), 806 deletions(-) create mode 100644 docs/kubernetes/gateway-api/README.mdx create mode 100644 docs/kubernetes/gateway-api/quickstart.mdx create mode 100644 docs/kubernetes/gateway-api/reference.mdx delete mode 100644 docs/kubernetes/inference-gateway.md diff --git a/README.md b/README.md index fe75adfc205e..be9711fa4754 100644 --- a/README.md +++ b/README.md @@ -108,18 +108,21 @@ Most inference engines optimize a single GPU or a single node. Dynamo is the **o - **K8s Inference Gateway plugin:** KV-aware routing inside the standard Kubernetes gateway - **Storage-tier KV offload:** S3/Azure blob support + global KV events for cluster-wide cache visibility -## Deployment Modes +## Request Routing Topologies -Dynamo can run in two deployment modes. Both expose an OpenAI-compatible API and support the same backends, disaggregated serving, and KV-aware routing. +Dynamo can expose traffic through two Kubernetes request routing topologies. Both expose an +OpenAI-compatible API and support the same backends, disaggregated serving, and KV-aware routing. -| Mode | What it is | When to use | +| Topology | What it is | When to use | |------|------------|-------------| -| **Standalone** *(default)* | Dynamo's own Frontend serves HTTP and the integrated Dynamo Router makes KV-aware routing decisions. No external gateway required. | Local development, single-cluster deployments, and any environment where you want Dynamo to own the request entry point end to end. | -| **Gateway (GAIE)** | Dynamo runs behind a Kubernetes [Gateway API Inference Extension](https://gateway-api-inference-extension.sigs.k8s.io/) gateway. KV-aware routing is performed at the gateway layer by the Dynamo Endpoint Picker Plugin (EPP); the Frontend runs as a sidecar in `--router-mode direct` and respects the EPP's per-request worker selection. | Production Kubernetes platforms that already standardize on the Inference Gateway, mixed-tenant clusters, or when you need gateway-level policy (auth, rate limiting, observability) co-located with KV-aware routing. | +| **Dynamo-native Frontend routing** | The Dynamo Frontend serves HTTP and the integrated Dynamo Router makes worker-selection decisions. No external gateway is required. | Local development, single-cluster deployments, and environments where Dynamo should own the request entry point end to end. | +| **Gateway API routing with GAIE** | A Kubernetes [Gateway API Inference Extension](https://github.com/kubernetes-sigs/gateway-api-inference-extension) gateway calls the Dynamo Endpoint Picker Plugin (EPP) before forwarding to the selected worker's Frontend sidecar in `--router-mode direct`. | Kubernetes platforms that standardize on Gateway API, or deployments where gateway-level policy, auth, rate limiting, and observability should sit at the cluster edge. | -In **standalone** mode, request flow is `client → Frontend → Router → workers`. In **gateway** mode, request flow is `client → Inference Gateway → EPP (KV-aware routing) → Frontend sidecar (direct) → workers`. +Request flow for the Dynamo-native path is `client → Frontend → Router → workers`. Request flow for +the Gateway API path is `client → Gateway → EPP → Frontend sidecar (direct) → workers`. -See the [Inference Gateway (GAIE) guide](docs/kubernetes/inference-gateway.md) for the full setup, supported features, and configuration of gateway mode. +See the [Gateway API Inference Extension (GAIE) guide](docs/kubernetes/gateway-api/README.mdx) for +the Gateway API setup, supported features, and configuration. ## Quick Start diff --git a/deploy/inference-gateway/README.md b/deploy/inference-gateway/README.md index ee4363129d2d..34c04a19113b 100644 --- a/deploy/inference-gateway/README.md +++ b/deploy/inference-gateway/README.md @@ -5,4 +5,4 @@ SPDX-License-Identifier: Apache-2.0 --> Integrate Dynamo with the Gateway API Inference Extension for intelligent KV-aware request routing at the gateway layer. -See [Inference Gateway documentation](../../docs/kubernetes/inference-gateway.md) for full setup instructions, configuration options, and deployment examples. +See [Gateway API Inference Extension documentation](../../docs/kubernetes/gateway-api/README.mdx) for setup instructions, configuration options, and deployment examples. diff --git a/docs/design-docs/architecture.md b/docs/design-docs/architecture.md index 4bb7cebaa507..331ce5bd7465 100644 --- a/docs/design-docs/architecture.md +++ b/docs/design-docs/architecture.md @@ -132,14 +132,14 @@ In Kubernetes deployments, the same architecture maps to declarative resources: The diagram labels such as `PodClique A/B`, `ScalingGroup "Prefill"`, `ScalingGroup "Decode"`, and `(replicas, min)` represent this grouped scaling model. -## Deployment Modes +## Request Routing Topologies The request plane can be exposed in two ways: -- **Standalone mode** (default) — the Dynamo Frontend is the request entry point and the integrated Dynamo Router selects workers using KV-aware scoring. Used by all local installs and the default Kubernetes deployment. -- **Gateway mode (GAIE)** — Dynamo runs behind a Kubernetes [Gateway API Inference Extension](https://gateway-api-inference-extension.sigs.k8s.io/) gateway. KV-aware routing is performed at the gateway layer by the Dynamo Endpoint Picker Plugin (EPP); the Frontend runs as a sidecar in `--router-mode direct` and respects the EPP's per-request worker selection passed via request headers. +- **Dynamo-native Frontend routing** -- the Dynamo Frontend is the request entry point and the integrated Dynamo Router selects workers using KV-aware scoring. +- **Gateway API routing with GAIE** -- a Kubernetes [Gateway API Inference Extension](https://github.com/kubernetes-sigs/gateway-api-inference-extension) gateway calls the Dynamo Endpoint Picker Plugin (EPP), then forwards to the selected worker's Frontend sidecar in `--router-mode direct`. -Both modes share the same control plane, storage/events plane, and backend integrations — only the request entry point and the location of the routing decision differ. See the [Inference Gateway (GAIE) guide](../kubernetes/inference-gateway.md) for the gateway-mode setup and configuration reference. +Both topologies share the same control plane, storage/events plane, and backend integrations; only the request entry point and Gateway API integration boundary differ. See the [Gateway API Inference Extension (GAIE) guide](../kubernetes/gateway-api/README.mdx) for setup and configuration. ## Fault Tolerance Architecture diff --git a/docs/getting-started/introduction.md b/docs/getting-started/introduction.md index 8c1cec4b9d95..6c0ca54e5127 100644 --- a/docs/getting-started/introduction.md +++ b/docs/getting-started/introduction.md @@ -106,17 +106,17 @@ Kubernetes cluster: | Local or container | You are evaluating, developing, or adopting one component at a time. | OpenAI-compatible frontend, router, workers, file or etcd discovery, Python/Rust APIs, and installable packages. | | Kubernetes | You are deploying shared GPU capacity, multi-node serving, autoscaling, or platform-integrated inference. | Helm install, Dynamo operator, DGD/DCD/DGDR CRDs, Kubernetes-native discovery, Gateway API Inference Extension, Grove/LWS scheduling, ModelExpress, observability, and lifecycle management. | -## Request Routing Modes +## Request Routing Topologies -Dynamo supports two request-routing modes. Both expose the same -OpenAI-compatible API and the same backends; they differ in *where* request -routing happens. +Dynamo supports two Kubernetes request routing topologies. Both expose the same +OpenAI-compatible API and the same backends; they differ in where the request +enters the system and where worker selection is integrated. -- **Standalone mode** (default) -- The Dynamo Frontend serves HTTP requests directly, and the integrated Dynamo Router makes KV-aware routing decisions before dispatching to workers. No external gateway is required. This is the mode used by all local installs and the default Kubernetes deployment. Request flow: `client -> Frontend -> Router -> workers`. +- **Dynamo-native Frontend routing** -- The Dynamo Frontend serves HTTP requests directly, and the integrated Dynamo Router makes KV-aware routing decisions before dispatching to workers. No external gateway is required. Request flow: `client -> Frontend -> Router -> workers`. -- **Gateway mode (GAIE)** -- Dynamo runs behind a Kubernetes [Gateway API Inference Extension](https://gateway-api-inference-extension.sigs.k8s.io/) gateway. KV-aware routing is performed at the gateway layer by the Dynamo Endpoint Picker Plugin (EPP); the Frontend runs as a sidecar in `--router-mode direct` and forwards requests to the worker the EPP selected. Use this mode when your platform standardizes on the Inference Gateway, or when you want gateway-level policy (auth, rate limiting, observability) co-located with KV-aware routing. Request flow: `client -> Inference Gateway -> EPP (KV-aware) -> Frontend sidecar (direct) -> workers`. +- **Gateway API routing with GAIE** -- A Kubernetes [Gateway API Inference Extension](https://github.com/kubernetes-sigs/gateway-api-inference-extension) gateway calls the Dynamo Endpoint Picker Plugin (EPP) before forwarding to the selected worker's Frontend sidecar in `--router-mode direct`. Use this topology when your platform standardizes on Gateway API, or when you want gateway-level policy, auth, rate limiting, and observability at the cluster edge. Request flow: `client -> Gateway -> EPP -> Frontend sidecar (direct) -> workers`. -Both modes support disaggregated serving, multimodal, and the same set of backends (vLLM, SGLang, TensorRT-LLM). For full setup, supported features, and configuration of gateway mode, see the [Inference Gateway (GAIE) guide](../kubernetes/inference-gateway.md). +Both topologies support disaggregated serving, multimodal, and the same set of backends (vLLM, SGLang, TensorRT-LLM). For setup and configuration of the Gateway API path, see the [Gateway API Inference Extension (GAIE) guide](../kubernetes/gateway-api/README.mdx). ## Performance @@ -192,7 +192,7 @@ Explore the following resources to go deeper: - [KV Cache Offloading](../components/kvbm/kvbm-guide.md) -- Set up multi-tier memory management - [Planner](../components/planner/planner-guide.md) -- Configure SLA-based autoscaling - [Kubernetes Deployment](../kubernetes/README.md) -- Deploy at scale with Grove -- [Inference Gateway (GAIE)](../kubernetes/inference-gateway.md) -- Run Dynamo in gateway mode behind the K8s Inference Gateway +- [Gateway API Inference Extension (GAIE)](../kubernetes/gateway-api/README.mdx) -- Run Dynamo behind Kubernetes Gateway API with Dynamo EPP routing - [Overall Architecture](../design-docs/architecture.md) -- Full technical design - [Support Matrix](../reference/support-matrix.md) -- Check hardware and engine compatibility diff --git a/docs/getting-started/kubernetes-deployment.md b/docs/getting-started/kubernetes-deployment.md index 137289b65f7d..1668054d7ac7 100644 --- a/docs/getting-started/kubernetes-deployment.md +++ b/docs/getting-started/kubernetes-deployment.md @@ -24,7 +24,8 @@ Start with the [Kubernetes Quickstart](../kubernetes/README.md) to run one model | Deploy and manage models | [Deployment Overview](../kubernetes/model-deployment-guide.md) | | Load models faster across pods | [Model Caching](../kubernetes/model-caching.md) and [ModelExpress](../kubernetes/modelexpress.md) | | Operate a cluster deployment | [Autoscaling](../kubernetes/autoscaling.md), [Rolling Update](../kubernetes/rolling-update.md), [Disagg Communication](../kubernetes/disagg-communication-guide.md), and [Observability Metrics](../kubernetes/observability/metrics.md) | +| Route traffic kube-natively | [Gateway API Inference Extension (GAIE)](../kubernetes/gateway-api/README.mdx) | | Scale disaggregated serving | [Multinode Deployments](../kubernetes/deployment/multinode-deployment.md), [Grove](../kubernetes/grove.md), and [Topology Aware Scheduling](../kubernetes/topology-aware-scheduling.md) | -| Integrate with Kubernetes serving APIs | [Gateway API Inference Extension (GAIE)](../kubernetes/inference-gateway.md) and [LWS](../kubernetes/lws.md) | +| Scheduler support | [Grove](../kubernetes/grove.md) and [LWS](../kubernetes/lws.md) | If you are still evaluating Dynamo locally, start with the [Quickstart](quickstart.mdx) and [Local Installation](local-installation.md) first. diff --git a/docs/index.yml b/docs/index.yml index bfb0dd9ebaea..67c82520ddc7 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -143,6 +143,16 @@ navigation: path: kubernetes/dgdr.md - page: DGDR Examples path: kubernetes/dgdr-examples.md + - section: Request Routing + contents: + - section: Gateway API Inference Extension + contents: + - page: Overview + path: kubernetes/gateway-api/README.mdx + - page: Quickstart + path: kubernetes/gateway-api/quickstart.mdx + - page: Reference + path: kubernetes/gateway-api/reference.mdx - section: Model Loading contents: - page: Model Caching @@ -475,8 +485,6 @@ navigation: contents: - page: LWS path: kubernetes/lws.md - - page: Gateway API Inference Extension (GAIE) - path: kubernetes/inference-gateway.md # ==================== Design Docs ==================== - section: Design Docs diff --git a/docs/kubernetes/README.md b/docs/kubernetes/README.md index e10ac6fe265e..53d5692b5833 100644 --- a/docs/kubernetes/README.md +++ b/docs/kubernetes/README.md @@ -14,7 +14,11 @@ container guides remain useful for development, but Kubernetes is the canonical path for shared GPU clusters and multi-node serving. > [!NOTE] -> **Deployment modes.** Dynamo supports two deployment modes on Kubernetes. This quickstart uses **standalone mode**, where the Dynamo Frontend serves requests and the integrated Dynamo Router does KV-aware routing. Dynamo can also run in **gateway mode** behind a [Gateway API Inference Extension](https://gateway-api-inference-extension.sigs.k8s.io/) gateway, where KV-aware routing happens in the Dynamo Endpoint Picker Plugin (EPP) at the gateway layer and the Frontend runs as a sidecar in `--router-mode direct`. See the [Inference Gateway (GAIE) guide](inference-gateway.md) to set up gateway mode. +> **Request entry.** This quickstart uses Dynamo-native Frontend routing: the Dynamo Frontend +> receives requests and the integrated Dynamo Router selects workers. Dynamo can also integrate +> Kubernetes-natively with [Gateway API Inference Extension](https://github.com/kubernetes-sigs/gateway-api-inference-extension), +> where Gateway API receives requests and calls the Dynamo EPP for endpoint selection. See the +> [GAIE guide](gateway-api/README.mdx) for the Gateway API path. ## Prerequisites diff --git a/docs/kubernetes/dgdr.md b/docs/kubernetes/dgdr.md index 107391cc32c8..06b8895399db 100644 --- a/docs/kubernetes/dgdr.md +++ b/docs/kubernetes/dgdr.md @@ -316,9 +316,7 @@ so it cannot be used to add a missing `Epp` service to a DGDR-generated deployment. Use a direct DGD manifest or a GAIE recipe for EPP deployments. For manifests, `frontendSidecar` configuration, direct routing, EPP routing variables such as `DYN_USE_KV_EVENTS`, and route setup, see -[Gateway API Inference Extension](inference-gateway.md). The same guide also -documents the optional [Rust EPP](inference-gateway.md#4b-build-rust-epp-image-optional--experimental), -which is currently experimental. +[Gateway API Inference Extension](gateway-api/README.mdx). ### SKU Format diff --git a/docs/kubernetes/gateway-api/README.mdx b/docs/kubernetes/gateway-api/README.mdx new file mode 100644 index 000000000000..463fd41d8e2e --- /dev/null +++ b/docs/kubernetes/gateway-api/README.mdx @@ -0,0 +1,284 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Gateway API Inference Extension (GAIE) +sidebar-title: Overview +subtitle: Expose DynamoGraphDeployments through Kubernetes Gateway API and Dynamo EPP routing. +--- + +Dynamo supports two request routing topologies on Kubernetes: + +- **Dynamo-native Frontend routing.** The Dynamo Frontend receives HTTP requests and the integrated + Dynamo Router selects workers. +- **Gateway API routing with GAIE.** A Kubernetes `Gateway` receives HTTP requests, the + [Gateway API Inference Extension (GAIE)](https://github.com/kubernetes-sigs/gateway-api-inference-extension) + calls the Dynamo Endpoint Picker Plugin (EPP) for endpoint selection, and the selected worker's + Frontend sidecar forwards the request in direct mode. + +This guide covers the Gateway API path for `DynamoGraphDeployment` resources managed by the Dynamo +operator. Use it when your Kubernetes platform wants Gateway API to own traffic entry, policy, and +observability while Dynamo owns the serving graph, discovery, event plane, and routing logic inside +the EPP. + +## Components + +The operator-managed GAIE path combines user-created Gateway API objects with resources created +from the `DynamoGraphDeployment`. + +| Component | Role | Created by | +|---|---|---| +| `Gateway` | Receives external HTTP traffic for the namespace. | User or platform team | +| `HTTPRoute` | Attaches model traffic to the `Gateway` and points at the `InferencePool`. | User | +| `DynamoGraphDeployment` | Describes the serving graph, EPP component, workers, and Frontend sidecars. | User | +| Dynamo operator | Reconciles the DGD into Kubernetes resources. | Dynamo platform | +| `InferencePool` | Connects GAIE endpoint selection to the Dynamo EPP service. | Dynamo operator | +| Dynamo EPP | Scores endpoints and returns the selected worker to the gateway. | Dynamo operator | +| Frontend sidecar | Receives the already-selected request and forwards in direct mode. | Dynamo operator | +| Worker | Runs the model backend. | Dynamo operator | + +## Request Flow + +```mermaid +flowchart LR + Client["Client
/v1/chat/completions"] -->|"sends request"| Gateway["Gateway API
Gateway"] + Gateway -->|"matches"| Route["HTTPRoute"] + Route -->|"targets"| Pool["InferencePool"] + Pool -->|"calls"| EPP["Dynamo EPP
EndpointPicker"] + EPP -->|"x-dynamo-worker-instance-id"| Gateway + Gateway -->|"forward request"| Sidecar["Frontend sidecar
--router-mode direct"] + Sidecar -->|"forwards"| Worker["Dynamo worker"] + Worker -. "publishes KV events" .-> Runtime["Dynamo runtime
NATS/JetStream"] + Runtime -. "updates routing state" .-> EPP +``` + +Gateway API owns the external request path. Dynamo still owns the serving graph: the operator +creates the EPP Service, worker pods, Frontend sidecars, and `InferencePool` that binds the route to +the EPP. The EPP receives Dynamo routing state from the runtime event plane and returns the selected +worker ID to the gateway. The gateway forwards the request to the selected worker's Frontend sidecar, +which runs in direct routing mode. + +In this operator-managed path, the EPP consumes routing state through the Dynamo event plane using +NATS/JetStream. Direct vLLM ZMQ KV-event subscriptions are used by other integration shapes, but not +by this quickstart path. + +## Shared Prerequisites + +- A Kubernetes cluster with GPU nodes. For the baseline Gateway API environment, start with the + upstream [Gateway API getting started guide](https://gateway-api.sigs.k8s.io/guides/getting-started/introduction/) + and the upstream [GAIE guide](https://github.com/kubernetes-sigs/gateway-api-inference-extension/blob/main/site-src/guides/index.md). +- `kubectl`, [Helm](https://helm.sh/docs/intro/install/), and + [jq](https://jqlang.org/download/) configured for the cluster. +- Gateway API and GAIE CRDs installed. The quickstart installs them explicitly from pinned upstream + release manifests. +- A Gateway API implementation that supports GAIE `InferencePool` resources and `endpointPickerRef` + calls. +- Dynamo platform installed with the operator. See the [Kubernetes Quickstart](../README.md) and + [Installation Guide](../installation-guide.md). +- Model credentials and storage needed by the selected model. Hugging Face token secrets are a + Dynamo model-serving prerequisite, not a GAIE-specific resource; see the + [Hugging Face token secret](../README.md#huggingface-token-secret) setup. + +## Gateway Implementation + +GAIE requires a Gateway API implementation that can call an Endpoint Picker Plugin before forwarding +the request to a backend. Dynamo is independent of the Gateway implementation: pick the gateway that +matches your platform, then point its `HTTPRoute` and generated `InferencePool` at the Dynamo EPP. + +The quickstart shows two verified paths: agentgateway and Istio. Other Gateway API +implementations might work when they support the same GAIE `InferencePool` and `endpointPickerRef` +EPP path; check the upstream +[GAIE gateway implementation list](https://github.com/kubernetes-sigs/gateway-api-inference-extension/blob/main/site-src/implementations/gateways.md) +and your controller's documentation before choosing another implementation. + +Istio uses Envoy in its data plane. agentgateway is a Rust-based AI gateway. The requirement for +this guide is not Envoy specifically; it is support for Gateway API plus the GAIE EndpointPicker +flow. + + + + Use agentgateway for a small Gateway API footprint or when the cluster does not already + standardize on a service mesh. Install the agentgateway chart with + `inferenceExtension.enabled=true`; the GatewayClass is `agentgateway`. + + + Use Istio when the cluster already standardizes on Istio for ingress, mesh policy, or telemetry. + Install Istio with `ENABLE_GATEWAY_API_INFERENCE_EXTENSION=true`; the GatewayClass is `istio`. + Configure EPP TLS policy with a `DestinationRule` when mesh policy requires it. + + + +The quickstart walks through the two verified implementation paths shown in this table: + +| | agentgateway | Istio | +|---|---|---| +| Good fit | New clusters or clusters without a mesh standard | Clusters that already standardize on Istio | +| Install footprint | agentgateway CRDs and controller in `agentgateway-system` | Istio control plane in `istio-system` or your chosen namespace | +| GatewayClass | `agentgateway` | `istio` | +| GAIE support | Enable `inferenceExtension.enabled=true` on the chart | Install Istio with `ENABLE_GATEWAY_API_INFERENCE_EXTENSION=true` | +| Mesh interaction | Use `AgentgatewayParameters` to keep `agentgateway-proxy` out of sidecar injection | Configure EPP TLS with a `DestinationRule` when mesh policy applies | + +## Gateway API Concepts + +```mermaid +flowchart TB + GatewayNS["Gateway namespace"] -->|"contains"| Gateway["Gateway
listener and address"] + ModelNS["Model namespace"] -->|"contains"| Route["HTTPRoute"] + ModelNS -->|"contains"| Pool["InferencePool"] + Route -->|"attaches with spec.parentRefs"| Gateway + Route -->|"routes with rules.backendRefs"| Pool + Pool -->|"selects with endpointPickerRef"| EPP["Dynamo EPP Service"] +``` + +`HTTPRoute.spec.parentRefs` attaches a route to a `Gateway`. If the `HTTPRoute` and `Gateway` live +in different namespaces, set `parentRefs[].namespace` to the Gateway namespace. `rules[].backendRefs` +points at the `InferencePool`; the pool points at the EPP service through `endpointPickerRef`. + +For the upstream API model, see the +[Gateway API HTTP routing guide](https://gateway-api.sigs.k8s.io/guides/user-guides/http-routing/) and the +[cross-namespace routing guide](https://gateway-api.sigs.k8s.io/guides/user-guides/multiple-ns/). + +## Configure DynamoGraphDeployments for GAIE + +In GAIE mode, the EPP chooses workers. The worker Frontend sidecar must run in direct routing mode so +it honors the EPP selection instead of choosing a worker again. + +```yaml +frontendSidecar: sidecar-frontend +podTemplate: + spec: + containers: + - name: sidecar-frontend + args: + - -m + - dynamo.frontend + - --router-mode + - direct +``` + +The EPP component is part of the `DynamoGraphDeployment`. The operator creates the EPP Service and +the matching `InferencePool`, so users apply the DGD and the route instead of hand-crafting the pool. + +### EPP Component Configuration + +Start from the recipe EPP component and update the `DynamoGraphDeployment` for your cluster. Change +deployment-level settings such as replicas and resources to fit gateway traffic volume. Change +routing plugin settings only when you want different endpoint-selection behavior, then validate the +result with production-like traffic. + +| Setting | When to change it | Rule | +|---|---|---| +| `replicas` and `podTemplate.spec.containers[].resources` | Scale or reserve capacity for EPP pods. | Keep EPP capacity aligned with gateway request volume. | +| `DYN_MODEL_NAME` | Change the served model. | Match the worker model name. | +| `DYN_KV_CACHE_BLOCK_SIZE` | Change the backend KV block size. | Match the backend `--block-size` value. | +| `DYN_ENFORCE_DISAGG` | Require strict prefill/decode routing separation. | Set `"true"` only for disaggregated deployments that should fail closed when topology labels are missing. | +| `label-filter` parameters | Change worker topology labels or component names. | Keep filter labels and values aligned with worker pod labels. | +| `schedulingProfiles[].plugins[].weight` | Adjust how much each scorer influences endpoint selection. | Tune scorer weights deliberately; keep required filters and the picker in the profile. | +| scorer and picker plugins | Change the routing strategy. | Treat this as advanced EPP tuning and validate with traffic. | + +For upstream EPP configuration semantics, see the GAIE +[EPP YAML configuration guide](https://gateway-api-inference-extension.sigs.k8s.io/guides/epp-configuration/config-text/) +and its +[Scheduling Profiles](https://gateway-api-inference-extension.sigs.k8s.io/guides/epp-configuration/config-text/#scheduling-profiles) +section for plugin weights. For label-based endpoint selection, see the upstream +[InferencePool configuration guide](https://gateway-api-inference-extension.sigs.k8s.io/api-types/inferencepool/#how-to-configure-an-inferencepool). +The `label-filter` plugin shown here is Dynamo-specific; the component role label comes from the +Dynamo [ComponentType](../api-reference.md#componenttype) field. + +The operator reconciles the EPP `Deployment`, EPP `Service`, and generated `InferencePool` from the +DGD. Tune the DGD first; patch generated resources only for short-lived debugging. + +```yaml +- name: Epp + type: epp + replicas: 1 + eppConfig: + config: + plugins: + - type: disagg-profile-handler + - name: decode-filter + type: label-filter + parameters: + label: nvidia.com/dynamo-component-type + validValues: [decode] + allowsNoLabel: true # Aggregated recipes can route unlabeled decode pods. + - name: dyn-decode + type: dyn-decode-scorer + - name: picker + type: max-score-picker + schedulingProfiles: + - name: decode + plugins: + - pluginRef: decode-filter + weight: 1 # Keep topology filters aligned with the worker labels. + - pluginRef: dyn-decode + weight: 1 # Tune scorer weights to change endpoint scoring. + - pluginRef: picker + weight: 1 +``` + +The `DYN_*` environment values are runtime contracts between the EPP router logic and the workers. +Update them when the worker backend changes; do not use them to tune scoring. + +```yaml +- name: Epp + type: epp + podTemplate: + spec: + containers: + - name: main + env: + - name: DYN_MODEL_NAME + value: Qwen/Qwen3-0.6B # Match the worker model name. + - name: DYN_KV_CACHE_BLOCK_SIZE + value: "16" # Match the worker backend's --block-size. + - name: DYN_ENFORCE_DISAGG + value: "false" # Use "true" for disaggregated fail-closed behavior. +``` + +See the complete EPP examples in the source tree: +`recipes/qwen3-0.6b/vllm/agg/gaie/deploy.yaml` for the Qwen 0.6B aggregated recipe manifest and +`examples/backends/vllm/deploy/gaie/disagg.yaml` for the Qwen 0.6B disaggregated example manifest. + +## Routing Behavior + +GAIE does not require one scoring strategy. Choose the routing behavior based on the routing state +available to the EPP. + +| Mode | What the EPP uses | When to use it | +|---|---|---| +| KV cache aware routing | Worker-published KV cache events delivered through the Dynamo event plane. | Default path when workers publish KV events and you want cache locality to influence endpoint selection. | +| Approximate routing | Endpoint availability plus local bookkeeping from tokenized requests and request lifecycle. | Fallback path when precise worker-published KV events are unavailable, disabled, or not yet supported by the chosen backend or deployment shape. | + +With operator-managed GAIE, NATS/JetStream backs routing-state delivery. The EPP can receive startup +state and subsequent updates through the Dynamo event plane instead of rebuilding all state from new +traffic after every EPP restart. + +## Compatibility and Defaults + +The quickstart pins the Gateway API layer so manual setup is repeatable. Keep the Dynamo platform, +EPP, and runtime images on the same Dynamo release line. + +| Component | Default shown here | Notes | +|---|---|---| +| Gateway API CRDs | `v1.5.1` | Installed from the upstream Gateway API release. | +| GAIE CRDs | `v1.2.1` | Installed from the upstream Gateway API Inference Extension release. | +| agentgateway | `v1.0.0` | Installed with `inferenceExtension.enabled=true`. | +| Istio | `1.29.2` | Install with `ENABLE_GATEWAY_API_INFERENCE_EXTENSION=true`. | +| Dynamo images | `1.2.1` | Use one Dynamo release line for the platform chart, EPP image, and runtime images. | + +## Troubleshooting Signals + +| Symptom | Likely cause | Check | +|---|---|---| +| `HTTPRoute` is not accepted | `parentRefs` points at the wrong Gateway name or namespace. | `kubectl describe httproute -n ` and compare `spec.parentRefs` with the Gateway. | +| Requests reach a model but EPP logs stay quiet | The route bypasses the `InferencePool`, or the pool points at the wrong EPP service. | Verify `rules.backendRefs` points at the `InferencePool` and `endpointPickerRef` points at the Dynamo EPP service. | +| EPP does not receive routing state | Dynamo event-plane components are not ready, or image tags do not match. | Check Dynamo platform pods, DGD status, EPP logs, and image tags against the compatibility table. | +| Istio path cannot call the EPP | Istio was installed without GAIE enabled, or mesh TLS policy blocks the EPP call. | Confirm `ENABLE_GATEWAY_API_INFERENCE_EXTENSION=true` and configure the EPP `DestinationRule`. | + +## Next Step + +Run the [GAIE Quickstart](./quickstart.mdx) to deploy a `DynamoGraphDeployment`, expose it through +Gateway API, and verify an end-to-end request through the Dynamo EPP. + +Use [GAIE Reference](./reference.mdx) for resource contracts, routing knobs, and service mesh +settings. diff --git a/docs/kubernetes/gateway-api/quickstart.mdx b/docs/kubernetes/gateway-api/quickstart.mdx new file mode 100644 index 000000000000..e7c6f146c4d6 --- /dev/null +++ b/docs/kubernetes/gateway-api/quickstart.mdx @@ -0,0 +1,364 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: GAIE Quickstart +sidebar-title: Quickstart +subtitle: Deploy a DynamoGraphDeployment behind Gateway API Inference Extension. +--- + +This quickstart deploys a Dynamo operator-managed serving graph behind Gateway API. The `Gateway` +receives requests, GAIE calls the Dynamo EPP for endpoint selection, and the selected worker's +Frontend sidecar forwards the request in direct mode. + +## What This Deploys + +```mermaid +flowchart LR + Client["Client"] -->|"sends request"| Gateway["Gateway API
Gateway"] + Gateway -->|"matches"| Route["HTTPRoute"] + Route -->|"targets"| Pool["InferencePool
generated by operator"] + Pool -->|"calls"| EPP["Dynamo EPP"] + EPP -->|"selects worker"| Sidecar["Frontend sidecar
--router-mode direct"] + Sidecar -->|"forwards"| Worker["Dynamo vLLM worker"] + Worker -. "publishes KV events" .-> Runtime["Dynamo runtime
NATS/JetStream"] + Runtime -. "updates routing state" .-> EPP + + DGD["DynamoGraphDeployment"] -->|"declares graph"| Operator["Dynamo operator"] + Operator -->|"creates"| Pool + Operator -->|"creates"| EPP + Operator -->|"injects"| Sidecar + Operator -->|"creates"| Worker +``` + +## Prerequisites + +- Kubernetes cluster with GPU nodes. +- `kubectl`, Helm, and `jq`. +- Access to `nvcr.io/nvidia/ai-dynamo` images for the Dynamo release you use. +- Hugging Face access to `Qwen/Qwen3-0.6B`. +- Shared RWX storage for the recipe's `model-cache` PVC. + +Set the common variables: + +```bash +export DYNAMO_VERSION=1.2.1 +export NAMESPACE=gaie-dynamo +export DYNAMO_SYSTEM_NAMESPACE=dynamo-system +export AGW_NAMESPACE=agentgateway-system +export ISTIO_NAMESPACE=istio-system + +kubectl create namespace "$NAMESPACE" --dry-run=client -o yaml | kubectl apply -f - +``` + +Clone the Dynamo source tree that contains the recipe manifests used below: + +```bash +git clone https://github.com/ai-dynamo/dynamo.git +cd dynamo +``` + +## Install Dynamo Platform + +Install the Dynamo platform and operator with the [Installation Guide](../installation-guide.md). +Use the same Dynamo release line for the platform chart, EPP image, and runtime images. + +After installation, verify that the Helm release exists in the platform namespace: + +```bash +helm status dynamo-platform --namespace "$DYNAMO_SYSTEM_NAMESPACE" +``` + +## Create Model Credentials + +Create model credentials if the model requires them. This is a Dynamo model-serving prerequisite, +not a GAIE-specific resource. The general Kubernetes quickstart explains the +[Hugging Face token secret](../README.md#huggingface-token-secret) pattern. + +```bash +export HF_TOKEN='your-hf-token' + +kubectl create secret generic hf-token-secret \ + -n "$NAMESPACE" \ + --from-literal=HF_TOKEN="$HF_TOKEN" +``` + +## Install Gateway API and GAIE CRDs + +Install the Gateway API layer explicitly. If your platform team already installed Gateway API, GAIE, +and a compatible Gateway implementation, skip to [Create the Gateway](#create-the-gateway). + +```bash +kubectl apply --server-side --force-conflicts \ + -f https://github.com/kubernetes-sigs/gateway-api/releases/download/v1.5.1/standard-install.yaml + +kubectl apply \ + -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/releases/download/v1.2.1/manifests.yaml +``` + +## Create the Gateway + +Choose the Gateway implementation for this namespace. + + + + ```bash + helm upgrade -i --create-namespace --namespace "$AGW_NAMESPACE" --version v1.0.0 \ + agentgateway-crds oci://cr.agentgateway.dev/charts/agentgateway-crds + + helm upgrade -i --namespace "$AGW_NAMESPACE" --version v1.0.0 \ + agentgateway oci://cr.agentgateway.dev/charts/agentgateway \ + --set inferenceExtension.enabled=true \ + --wait + + kubectl get gatewayclass agentgateway + ``` + + Create the `AgentgatewayParameters` resource in the model namespace. The parameters resource + excludes Istio sidecar injection from `agentgateway-proxy` pods when the namespace has + `istio-injection=enabled`. + + ```bash + kubectl apply --server-side -n "$NAMESPACE" -f - <<'YAML' + apiVersion: agentgateway.dev/v1alpha1 + kind: AgentgatewayParameters + metadata: + name: inference-gateway-params + spec: + deployment: + spec: + template: + metadata: + annotations: + sidecar.istio.io/inject: "false" + YAML + ``` + + Create the `Gateway` that uses those parameters. + + ```bash + kubectl apply -n "$NAMESPACE" -f - <<'YAML' + apiVersion: gateway.networking.k8s.io/v1 + kind: Gateway + metadata: + name: inference-gateway + spec: + gatewayClassName: agentgateway + infrastructure: + parametersRef: + group: agentgateway.dev + kind: AgentgatewayParameters + name: inference-gateway-params + listeners: + - name: http + port: 80 + protocol: HTTP + YAML + ``` + + Wait for the gateway controller to program the gateway. + + ```bash + kubectl wait gateway/inference-gateway -n "$NAMESPACE" \ + --for=condition=Programmed --timeout=180s + ``` + + + ```bash + export ISTIO_VERSION=1.29.2 + + if ! command -v istioctl >/dev/null 2>&1; then + curl -fsSL https://istio.io/downloadIstio | ISTIO_VERSION="$ISTIO_VERSION" sh - + export PATH="$PWD/istio-$ISTIO_VERSION/bin:$PATH" + fi + + istioctl install -y \ + --set values.global.istioNamespace="$ISTIO_NAMESPACE" \ + --set values.pilot.env.ENABLE_GATEWAY_API_INFERENCE_EXTENSION=true + + kubectl wait --for=condition=Available --timeout=180s \ + -n "$ISTIO_NAMESPACE" deployment/istiod + + helm upgrade -i dynamo-platform \ + oci://helm.ngc.nvidia.com/nvidia/ai-dynamo/charts/dynamo-platform \ + --version "$DYNAMO_VERSION" \ + --namespace "$DYNAMO_SYSTEM_NAMESPACE" \ + --reuse-values \ + --set dynamo.serviceMesh.enabled=true \ + --set dynamo.serviceMesh.provider=istio \ + --wait + + kubectl apply -n "$NAMESPACE" -f - <<'YAML' + apiVersion: gateway.networking.k8s.io/v1 + kind: Gateway + metadata: + name: inference-gateway + spec: + gatewayClassName: istio + listeners: + - name: http + port: 80 + protocol: HTTP + YAML + + kubectl wait gateway/inference-gateway -n "$NAMESPACE" \ + --for=condition=Programmed --timeout=180s + + kubectl get gatewayclass istio + ``` + + + +Keeping the `Gateway` and `HTTPRoute` in the same namespace avoids a cross-namespace +`parentRefs[].namespace` field in the route. + +## Prepare the Model Cache + +The Qwen recipe mounts a shared `model-cache` PVC. Edit +`recipes/qwen3-0.6b/model-cache/model-cache.yaml` first and set `storageClassName` to a RWX storage +class available in your cluster. For the general pattern, see [Model Caching](../model-caching.md). + +```bash +kubectl apply -n "$NAMESPACE" -f recipes/qwen3-0.6b/model-cache/ + +kubectl wait --for=condition=Complete job/model-download \ + -n "$NAMESPACE" --timeout=3600s +``` + +## Deploy the Serving Graph + +Deploy the Qwen 0.6B aggregated recipe and its route: + +```bash +kubectl apply -n "$NAMESPACE" \ + -f recipes/qwen3-0.6b/vllm/agg/gaie/deploy.yaml + +kubectl apply -n "$NAMESPACE" \ + -f recipes/qwen3-0.6b/vllm/agg/gaie/httproute.yaml +``` + +Wait for the operator-created resources: + +```bash +export ROUTE_MODEL=Qwen/Qwen3-0.6B + +kubectl wait -n "$NAMESPACE" dynamographdeployment/qwen3-0-6b-agg \ + --for=condition=Ready --timeout=1800s + +kubectl get inferencepool qwen3-0-6b-agg-pool -n "$NAMESPACE" +kubectl get httproute qwen3-0-6b-agg -n "$NAMESPACE" +``` + +## Verify End-to-End + +Use one access mode to set `GATEWAY_URL`, then send a request through the Gateway and EPP. Keep +`ROUTE_MODEL` set to the model name from the route manifest you applied. + + + + ```bash + export GATEWAY_SERVICE=$(kubectl get svc -n "$NAMESPACE" \ + -l gateway.networking.k8s.io/gateway-name=inference-gateway \ + -o jsonpath='{.items[0].metadata.name}') + + kubectl -n "$NAMESPACE" port-forward "svc/$GATEWAY_SERVICE" 8000:80 + ``` + + In another terminal: + + ```bash + export GATEWAY_URL=http://localhost:8000 + ``` + + + ```bash + export GATEWAY_HOST=$(kubectl get gateway inference-gateway -n "$NAMESPACE" \ + -o jsonpath='{.status.addresses[0].value}') + export GATEWAY_URL=http://$GATEWAY_HOST + ``` + + + + +```bash title="List models" +curl --max-time 20 -sS "$GATEWAY_URL/v1/models" \ + -H "X-Gateway-Model-Name: $ROUTE_MODEL" | jq . +``` + +```bash title="Send a chat request" +curl --max-time 180 -sS "$GATEWAY_URL/v1/chat/completions" \ + -H "X-Gateway-Model-Name: $ROUTE_MODEL" \ + -H "content-type: application/json" \ + -d '{ + "model": "'"$ROUTE_MODEL"'", + "messages": [{"role": "user", "content": "Explain KV cache aware routing in one sentence."}], + "max_tokens": 96 + }' | jq . +``` + + +Finish by checking that the EPP path handled the request. A successful smoke test should show the EPP +receiving endpoint-picker traffic and selecting a worker near the time of your request; that proves +the request flowed through Gateway API and the Dynamo EPP before it reached the Frontend sidecar. + +```bash +kubectl logs -n "$NAMESPACE" -l nvidia.com/dynamo-component-type=epp --tail=200 +``` + +If the log output is quiet, run the chat request again while tailing the EPP logs in another +terminal. + +## Troubleshooting + + + + ```bash + kubectl describe gateway inference-gateway -n "$NAMESPACE" + kubectl get pods -n "$AGW_NAMESPACE" + kubectl logs -n "$AGW_NAMESPACE" deployment/agentgateway --tail=50 + kubectl get gatewayclass agentgateway + kubectl get inferencepool -n "$NAMESPACE" + kubectl describe httproute -n "$NAMESPACE" + ``` + + If requests return HTTP 500 and the namespace has `istio-injection=enabled`, verify the + `agentgateway-proxy` pod does not have an `istio-proxy` sidecar: + + ```bash + kubectl get pods -n "$NAMESPACE" \ + -l gateway.networking.k8s.io/gateway-name=inference-gateway \ + -o jsonpath='{.items[*].spec.containers[*].name}' + ``` + + See [GAIE Reference](./reference.mdx#agentgateway-and-istio-injection) for the sidecar + injection contract. + + + ```bash + kubectl describe gateway inference-gateway -n "$NAMESPACE" + kubectl get pods -n "$ISTIO_NAMESPACE" + kubectl logs -n "$ISTIO_NAMESPACE" deployment/istiod --tail=50 + kubectl get gatewayclass istio + kubectl get inferencepool -n "$NAMESPACE" + kubectl describe httproute -n "$NAMESPACE" + ``` + + Confirm Istio was installed with `ENABLE_GATEWAY_API_INFERENCE_EXTENSION=true`, and confirm the + EPP service has a `DestinationRule` when Istio sidecars enforce TLS policy for gateway-to-EPP + traffic. + + See [GAIE Reference](./reference.mdx#service-mesh-integration) for the generated + `DestinationRule` behavior. + + + +If model pods restart while loading, inspect the pod events. When events show startup probe failures +and the model load time is expected, increase `startupProbe.failureThreshold` on the affected DGD +component. This is general Kubernetes probe tuning, not a GAIE-specific setting. + +## Clean Up + +If this namespace is only for the quickstart, delete it: + +```bash +kubectl delete namespace "$NAMESPACE" +``` diff --git a/docs/kubernetes/gateway-api/reference.mdx b/docs/kubernetes/gateway-api/reference.mdx new file mode 100644 index 000000000000..417e931aee21 --- /dev/null +++ b/docs/kubernetes/gateway-api/reference.mdx @@ -0,0 +1,281 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: GAIE Reference +sidebar-title: Reference +subtitle: Runtime contracts, routing controls, and mesh integration for Dynamo with GAIE. +--- + +Use this reference after the [GAIE Quickstart](./quickstart.mdx) when you need to inspect generated +resources, tune routing behavior, or adapt the Gateway API path to a cluster policy. + +This page is user-facing runtime reference. It does not cover building custom EPP images, local +development loops, minikube-specific setup, or full uninstall procedures. + +## Resource Contract + +Operator-managed GAIE connects Gateway API resources to the Dynamo serving graph through an +operator-generated `InferencePool`. + +```mermaid +flowchart LR + Route["HTTPRoute"] -->|"routes to backendRef"| Pool["InferencePool"] + Pool -->|"calls endpointPickerRef"| EPP["Dynamo EPP Service
grpc 9002"] + Pool -->|"selects eligible pods"| Pods["Worker pods
Frontend sidecar :8000"] + EPP -->|"returns selected endpoint"| Gateway["Gateway data plane"] + Gateway -->|"forwards request"| Pods +``` + +| Resource | Contract | +|---|---| +| `Gateway` | Owns listeners, addresses, and Gateway implementation behavior. | +| `HTTPRoute` | Attaches traffic to the `Gateway` through `spec.parentRefs` and points `rules[].backendRefs` at the `InferencePool`. | +| `InferencePool` | Defines the eligible backend pod set and the EPP endpoint the gateway calls before forwarding. | +| EPP `Service` | Exposes the Dynamo EPP on gRPC port `9002`. | +| Frontend sidecar | Receives the selected request on the pod's `http` port and runs with `--router-mode direct`. | + +The Dynamo operator creates the `InferencePool` for a `DynamoGraphDeployment` that contains an EPP +component. The pool name is `-pool`, it lives in the DGD namespace, and its +`endpointPickerRef` points to the generated EPP `Service`. + +The generated selector matches worker pods by Dynamo labels: + +```yaml +spec: + selector: + matchLabels: + nvidia.com/dynamo-component-class: worker + nvidia.com/dynamo-namespace: + endpointPickerRef: + kind: Service + name: -epp + port: + number: 9002 + targetPorts: + - number: 8000 +``` + +Do not hand-edit the generated `InferencePool` unless you also keep its selector aligned with the +operator's worker-pod labels. If the pool selector and Dynamo discovery disagree, the EPP can select +a worker that the gateway data plane refuses to forward to. + +For the upstream Gateway API model, see the +[HTTP routing guide](https://gateway-api.sigs.k8s.io/guides/user-guides/http-routing/) and +[cross-namespace routing guide](https://gateway-api.sigs.k8s.io/guides/user-guides/multiple-ns/). + +## Request Contract + +With GAIE, worker selection happens in the EPP before the request reaches the worker sidecar. The +sidecar must run in direct mode so it honors the EPP decision instead of routing again. + +```yaml +frontendSidecar: sidecar-frontend +podTemplate: + spec: + containers: + - name: sidecar-frontend + args: + - -m + - dynamo.frontend + - --router-mode + - direct +``` + +The EPP sends routing decisions to the selected sidecar through request headers. + +| Header | Meaning | +|---|---| +| `x-dynamo-worker-instance-id` | Decode or aggregated worker selected for the request. | +| `x-dynamo-dp-rank` | Data-parallel rank for the selected decode or aggregated worker. | +| `x-dynamo-routing-mode` | `aggregated` or `disaggregated`. | +| `x-dynamo-prefill-instance-id` | Prefill worker selected for disaggregated requests. | +| `x-dynamo-prefill-dp-rank` | Data-parallel rank for the selected prefill worker, when present. | + +For body-bearing OpenAI requests, the EPP also tokenizes the request and injects token data into the +request body so the sidecar can avoid repeating the same tokenization work. + +## Routing Modes + +The same Dynamo router logic can run behind the Dynamo-native Frontend entry path or inside the +GAIE EPP. In the Gateway API path, the EPP owns endpoint selection and the worker sidecar owns +request forwarding. + +| Mode | EPP input | When to use | +|---|---|---| +| KV cache aware routing | Worker KV cache events plus local request bookkeeping. | Use when workers publish KV events and cache locality should influence endpoint selection. | +| Approximate routing | Tokenized requests, request lifecycle, and local predicted state. | Use when KV events are unavailable, disabled, or not supported by the selected deployment shape. | + +In the operator-managed GAIE path, KV events reach the EPP through the Dynamo event plane using +NATS/JetStream. vLLM can also publish KV events through ZMQ in other integration shapes; the +operator-managed `DynamoGraphDeployment` path does not use ZMQ for the EPP. + +```mermaid +flowchart LR + Worker["Dynamo worker"] -->|"publishes KV events"| NATS["Dynamo runtime
NATS/JetStream"] + NATS -->|"delivers state"| EPP["Dynamo EPP"] + EPP -->|"scores endpoints"| Gateway["Gateway data plane"] + Gateway -->|"forwards selected request"| Sidecar["Frontend sidecar"] +``` + +To use KV cache aware routing: + +1. Enable worker prefix caching and KV event publishing for your backend. +2. Keep EPP KV events enabled. +3. Keep the worker KV block size aligned with the EPP block size. + +Backend examples: + +| Backend | Worker setting | +|---|---| +| vLLM | Pass `--enable-prefix-caching` and `--kv-events-config '{"enable_kv_cache_events":true}'`. | +| SGLang | Pass the backend's supported `--kv-events-config`. | +| TensorRT-LLM | Pass `--publish-events-and-metrics`. | + +Set `DYN_KV_CACHE_BLOCK_SIZE` on the EPP only when discovery does not already provide the backend's +block size. It must match the workers' `--block-size`. A mismatch changes the block hashes used for +prefix overlap and produces incorrect routing scores. + +To use approximate routing, disable worker KV events and set the EPP to predicted local state: + +```yaml +env: + - name: DYN_USE_KV_EVENTS + value: "false" + - name: DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT + value: "0" +``` + +## Router Tuning + +Set these values on the EPP component unless the deployment manifest says otherwise. + +| Setting | Default | Effect | +|---|---:|---| +| `DYN_ENFORCE_DISAGG` | `false` | When `true`, fail requests if prefill routing is unavailable. When `false`, fall back to aggregated routing until prefill workers appear. | +| `DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT` | `1.0` | Controls device-local prefix-overlap credit. Higher values prefer workers with cached prompt prefixes. | +| `DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT_DECAY` | `0.0` | Reduces prefix-overlap credit as active prefill load rises above the least-loaded eligible worker. | +| `DYN_ROUTER_PREFILL_LOAD_SCALE` | `1.0` | Scales prompt-side prefill load after cache-hit credits are applied. | +| `DYN_ROUTER_TEMPERATURE` | `0.0` | `0.0` selects deterministically. Higher values allow more worker exploration through softmax sampling. | +| `DYN_ROUTER_REPLICA_SYNC` | `false` | Publishes and subscribes router state across router replicas. | +| `DYN_ROUTER_TRACK_ACTIVE_BLOCKS` | `true` | Tracks active decode blocks for load balancing. | +| `DYN_ROUTER_TRACK_OUTPUT_BLOCKS` | `false` | Predicts output blocks during generation and decays them by progress toward expected output length. | +| `DYN_ROUTER_TRACK_PREFILL_TOKENS` | `true` | Includes active prompt-side prefill tokens in load accounting. | +| `DYN_ROUTER_PREDICTED_TTL_SECS` | unset | Enables predicted entries in the local indexer for this TTL when KV events are enabled. | +| `DYN_ADMISSION_CONTROL` | `none` | Set to `token-capacity` to skip workers that exceed active decode or prefill thresholds. | +| `DYN_ACTIVE_DECODE_BLOCKS_THRESHOLD` | unset | Decode worker is busy above this active-block fraction. Setting a numeric value enables token-capacity admission. | +| `DYN_ACTIVE_PREFILL_TOKENS_THRESHOLD` | unset | Worker is busy above this absolute active-prefill-token count. | +| `DYN_ACTIVE_PREFILL_TOKENS_THRESHOLD_FRAC` | unset | Worker is busy above this fraction of `max_num_batched_tokens`. | + +For the broader router configuration surface, see +[Router Configuration](../../components/router/router-configuration.md). + +## Service Mesh Integration + +The EPP serves gRPC on port `9002`. When an Istio sidecar mediates traffic from the gateway proxy to +the EPP service, configure mesh TLS explicitly so the proxy connects to the EPP's serving mode. + +Enable operator-managed Istio `DestinationRule` generation when installing or upgrading the Dynamo +platform chart: + +```bash +helm upgrade -i dynamo-platform \ + oci://helm.ngc.nvidia.com/nvidia/ai-dynamo/charts/dynamo-platform \ + --version "$DYNAMO_VERSION" \ + --namespace "$DYNAMO_SYSTEM_NAMESPACE" \ + --reuse-values \ + --set dynamo.serviceMesh.enabled=true \ + --set dynamo.serviceMesh.provider=istio +``` + +The platform values are: + +| Value | Default | Meaning | +|---|---|---| +| `dynamo.serviceMesh.enabled` | `false` | Generate service-mesh resources for EPP services. | +| `dynamo.serviceMesh.provider` | `istio` | Mesh provider. Only Istio is supported. | +| `dynamo.serviceMesh.istio.tlsMode` | `SIMPLE` | TLS mode for generated `DestinationRule` resources. | +| `dynamo.serviceMesh.istio.insecureSkipVerify` | `true` | Skip server certificate verification for the EPP's self-signed certificate. | +| `dynamo.serviceMesh.istio.clientCertificate` | `""` | Client certificate path for `MUTUAL` TLS mode. | +| `dynamo.serviceMesh.istio.privateKey` | `""` | Client private key path for `MUTUAL` TLS mode. | +| `dynamo.serviceMesh.istio.caCertificates` | `""` | CA certificate path for `MUTUAL` TLS mode. | + +When enabled and Istio CRDs are installed, the operator creates a `DestinationRule` for each EPP +service: + +```yaml +apiVersion: networking.istio.io/v1beta1 +kind: DestinationRule +metadata: + name: +spec: + host: ..svc.cluster.local + trafficPolicy: + tls: + mode: SIMPLE + insecureSkipVerify: true +``` + +If you install without the Dynamo operator Helm chart or leave `dynamo.serviceMesh.enabled=false`, +create an equivalent `DestinationRule` for each EPP service used through Istio. + +## agentgateway and Istio Injection + +When namespace-level Istio injection is enabled, the `agentgateway-proxy` pod can receive an Istio +sidecar. That sidecar can intercept the ext_proc gRPC connection from agentgateway to the EPP and +cause HTTP 500 responses from the gateway. + +Use a per-Gateway `AgentgatewayParameters` resource in the same namespace as the `Gateway`: + +```yaml +apiVersion: agentgateway.dev/v1alpha1 +kind: AgentgatewayParameters +metadata: + name: inference-gateway-params +spec: + deployment: + spec: + template: + metadata: + annotations: + sidecar.istio.io/inject: "false" +``` + +Reference that parameters resource from the `Gateway`: + +```yaml +apiVersion: gateway.networking.k8s.io/v1 +kind: Gateway +metadata: + name: inference-gateway +spec: + gatewayClassName: agentgateway + infrastructure: + parametersRef: + group: agentgateway.dev + kind: AgentgatewayParameters + name: inference-gateway-params + listeners: + - name: http + port: 80 + protocol: HTTP +``` + +Verify that the proxy pod does not contain `istio-proxy`: + +```bash +kubectl get pods -n "$NAMESPACE" \ + -l gateway.networking.k8s.io/gateway-name=inference-gateway \ + -o jsonpath='{.items[*].spec.containers[*].name}{"\n"}' +``` + +> [!WARNING] +> Patch the default `AgentgatewayParameters` resource in `agentgateway-system` only as a +> cluster-wide policy decision. Gateways without `spec.infrastructure.parametersRef` inherit that +> default. + +## Developer References + +Image build commands belong with the component source, not in this user reference. Use this source +location when developing or replacing the standard EPP image: + +- [Go EPP source](https://github.com/ai-dynamo/dynamo/tree/main/deploy/inference-gateway/epp) diff --git a/docs/kubernetes/inference-gateway.md b/docs/kubernetes/inference-gateway.md deleted file mode 100644 index b0ad15c56624..000000000000 --- a/docs/kubernetes/inference-gateway.md +++ /dev/null @@ -1,776 +0,0 @@ ---- -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -title: Gateway API Inference Extension (GAIE) ---- - -## Gateway API Inference Extension Setup with Dynamo - -Integrate Dynamo with the Gateway API Inference Extension, also known as Inference Gateway, for intelligent KV-aware request routing at the gateway layer. - -## Features - -- EPP's default kv-routing approach is not token-aware because the prompt is not tokenized. But the Dynamo plugin uses a token-aware KV algorithm. It employs the dynamo router which implements kv routing by running your model's tokenizer inline. The EPP plugin configuration is embedded in the recipe-based GAIE deploy YAMLs under [`recipes/llama-3-70b/vllm/agg/gaie/`](https://github.com/ai-dynamo/dynamo/tree/main/recipes/llama-3-70b/vllm/agg/gaie) and [`recipes/llama-3-70b/vllm/disagg-single-node/gaie/`](https://github.com/ai-dynamo/dynamo/tree/main/recipes/llama-3-70b/vllm/disagg-single-node/gaie), following the GAIE/EPP configuration layout used by this repository. - -- Dynamo Integration with the Inference Gateway supports Aggregated and Disaggregated Serving. A request only exercises disaggregated routing when the EPP config defines a `prefill` profile and prefill workers are available. The recipe examples provide separate aggregated and disaggregated configs under `recipes/llama-3-70b/vllm/agg/gaie/` and `recipes/llama-3-70b/vllm/disagg-single-node/gaie/`. Unless `DYN_ENFORCE_DISAGG=true`, deployments without a `prefill` profile or prefill workers fall back to aggregated serving. - -- GAIE integration supports Data Parallelism. - -- If you want to use LoRA deploy Dynamo without the Inference Gateway. - -- These setups use [agentgateway](https://agentgateway.dev/) as the Inference Gateway implementation. For the Istio Inference Gateway, check out [`recipes/qwen3-0.6b/vllm/agg/gaie`](https://github.com/ai-dynamo/dynamo/tree/main/recipes/qwen3-0.6b/vllm/agg/gaie). - -## Prerequisites - -- Kubernetes cluster with kubectl configured -- NVIDIA GPU drivers installed on worker nodes - -## Installation Steps - -### 1. Install Dynamo Platform ### - -[See Quickstart Guide](./README.md) to install Dynamo Kubernetes Platform. -If you are installing from the source tree rather than a release chart, follow [Advanced: Build from Source](./installation-guide.md#advanced-build-from-source) and run `helm dep build ./platform/` before `helm install` so the vendored subcharts match the local chart contents. - -### 2. Deploy Inference Gateway ### - -First, deploy an inference gateway service. In this example, we'll install agentgateway with the inference extension enabled. - -```bash -cd deploy/inference-gateway -export NAMESPACE=my-model # You can put the inference gateway into another namespace and then adjust your http-route.yaml -./scripts/install_gaie_crd_agentgateway.sh -``` -This script installs the Gateway API CRDs, the GAIE CRDs, agentgateway into `agentgateway-system`, and a `Gateway` named `inference-gateway` into `${NAMESPACE}`. - -#### Verify the Gateway is running - -```bash -kubectl get gateway inference-gateway -n ${NAMESPACE} - -# Sample output -# NAME CLASS ADDRESS PROGRAMMED AGE -# inference-gateway agentgateway True 1m -``` - - -### 2b. Istio Gateway (Alternative) ### - -If you are using Istio as your gateway implementation, -the EPP uses secure serving (TLS) by default. The gateway proxy needs an -Istio `DestinationRule` to talk to the EPP service; without it the Istio -`ext_proc` filter fails with `connection termination` errors. - -The Dynamo operator can create this `DestinationRule` for you. Install or -upgrade the platform Helm chart with `dynamo.serviceMesh.enabled=true` -(see [Service Mesh Integration (Istio)](#service-mesh-integration-istio) -below). When that is set, you can skip the rest of this section. - -If you are not using the operator's Helm chart, or have left -`dynamo.serviceMesh.enabled=false`, apply a `DestinationRule` manually for -each EPP service: - -```yaml -apiVersion: networking.istio.io/v1 -kind: DestinationRule -metadata: - name: -epp -spec: - host: -epp..svc.cluster.local - trafficPolicy: - tls: - insecureSkipVerify: true - mode: SIMPLE -``` - -Replace `` with your DynamoGraphDeployment name and `` with the namespace where the EPP is deployed. See [`recipes/qwen3-0.6b/vllm/agg/gaie/dr.yaml`](../../recipes/qwen3-0.6b/vllm/agg/gaie/dr.yaml) for an example. - -### 3. Setup secrets ### - -Do not forget docker registry secret if needed. - -```bash -kubectl create secret docker-registry docker-imagepullsecret \ - --docker-server=$DOCKER_SERVER \ - --docker-username=$DOCKER_USERNAME \ - --docker-password=$DOCKER_PASSWORD \ - --namespace=$NAMESPACE -``` - -Do not forget to include the HuggingFace token. - -```bash -export HF_TOKEN=your_hf_token -kubectl create secret generic hf-token-secret \ - --from-literal=HF_TOKEN=${HF_TOKEN} \ - -n ${NAMESPACE} -``` - -### 4. Build EPP image (Optional) - -You can either use the provided Dynamo FrontEnd image for the EPP image or you need to build your own Dynamo EPP custom image following the steps below. - -```bash -# export env vars -export DOCKER_SERVER=ghcr.io/nvidia/dynamo # Container registry -export IMAGE_TAG=YOUR-TAG # Or auto from git tag -cd deploy/inference-gateway/epp -make all # Do everything in one command -# or make all-push to also push - - -# Or step-by-step -make dynamo-lib # Build Dynamo library and copy to project -make image-load # Build Docker image and load locally -make image-push # Build and push to registry -make info # Check image tag -``` - -#### All-in-one Targets - -| Target | Description | -|--------|-------------| -| `make dynamo-lib` | Build Dynamo static library and copy to project | -| `make all` | Build Dynamo lib + Docker image + load locally | -| `make all-push` | Build Dynamo lib + Docker image + push to registry | - -### 4b. Build Rust EPP image (Optional — experimental) - -A pure-Rust EPP implementation is available as an alternative to the Go-based EPP. -It replaces the Go EPP + CGO bridge with a single native Rust binary that implements -the Envoy ext_proc gRPC service and uses Dynamo's KV-aware router directly — no FFI -boundary, no Go runtime. - -```bash -cd deploy/inference-gateway/ext-proc - -# Build and load Docker image locally -make image-load -# Creates: dynamo/dynamo-rust-epp: - -# Or build and push to a registry -export DOCKER_SERVER=ghcr.io/nvidia/dynamo -make image-push -``` - -To build the binary locally without Docker: - -```bash -cd deploy/inference-gateway/ext-proc -make build -# Binary at: /target/release/dynamo-ext-proc -``` - -#### Rust EPP Makefile Targets - -| Target | Description | -|--------|-------------| -| `make build` | Build the Rust EPP binary locally via cargo | -| `make image-load` | Build Docker image and load locally | -| `make image-push` | Build and push Docker image to registry | -| `make image-kind` | Build and load into a kind cluster | -| `make image-multiarch-push` | Build and push multi-arch image | -| `make fmt` / `make clippy` / `make test` | Development checks | -| `make info` | Show image tag and build configuration | - -#### Rust EPP Configuration - -The Rust EPP uses the same environment variables as the Go EPP for namespace -resolution and router configuration: - -| Variable | Default | Description | -|----------|---------|-------------| -| `DYN_NAMESPACE_PREFIX` | *(unset)* | Dynamo discovery namespace (highest priority) | -| `DYN_NAMESPACE` | `vllm-agg` | Dynamo discovery namespace (fallback) | -| `DYN_COMPONENT_NAME` | `backend` | Dynamo component name | -| `DYN_ENFORCE_DISAGG` | `false` | Enforce disaggregated prefill/decode routing | -| `DYN_KUBE_DISCOVERY_MODE` | `pod` | Kubernetes discovery identity mode; Rust EPP currently rejects `container` | -| `RUST_LOG` | `info` | Tracing log level filter | - -The gRPC port is hardcoded to `9002` (matching the operator's `EPPGRPCPort` constant). - -Namespace resolution follows the same logic as the Go EPP plugin: -`DYN_NAMESPACE_PREFIX` > `DYN_NAMESPACE` > `"vllm-agg"` (default). - -The Rust EPP also respects the standard Dynamo router environment variables -(`DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT`, `DYN_ROUTER_PREFILL_LOAD_SCALE`, -`DYN_ROUTER_TEMPERATURE`, `DYN_USE_KV_EVENTS`, etc.) documented in the -Configuration section below. The deprecated overlap-weight aliases remain -supported with the same precedence as the Go EPP. - -> [!NOTE] -> The Rust EPP is experimental. It uses Dynamo's native discovery system -> (`DistributedRuntime`) instead of the GAIE Kubernetes controllers, so it -> does not require `InferencePool` or `InferenceModel` CRDs for endpoint -> discovery. It discovers workers through Dynamo's own registration mechanism. - -> [!WARNING] -> The Rust EPP currently supports only pod-level Kubernetes discovery. Deploy -> one Rust EPP replica per pool because request selection and booking are not -> yet atomic across concurrent EPP replicas. After a worker-generation rolling -> update, restart the Rust EPP so it binds to the new generation namespace. -> Exact streamed output-block updates are also not yet wired into the Rust EPP. - -#### `InferencePool` and the data plane (Istio, kGateway, Agentgateway) - -Although the Rust EPP does not consult `InferencePool` for worker discovery, -the CRD is still required by the gateway **data plane**. Gateway -implementations (Istio, kGateway, Agentgateway) read `InferencePool` to: - -1. Attach the `ext_proc` filter pointing at the EPP service. -2. Enable the `override_host` LB policy so the EPP's - `x-gateway-destination-endpoint` header / dynamic-metadata is honored. -3. Scope which pods are eligible to receive traffic — the pool's selector - becomes the `envoy.lb.subset_hint` metadata that the EPP intersects with - its own discovered workers before picking one. - -The Dynamo operator **auto-generates the `InferencePool`** for every -`DynamoGraphDeployment` ([`deploy/operator/internal/dynamo/epp/inference_pool.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/dynamo/epp/inference_pool.go)). -Its `Selector` matches the operator's worker-pod labels and its -`EndpointPickerRef` points at the EPP service on `9002`, so Dynamo's -discovery and the pool's pod set stay in sync automatically — users do not -hand-craft the pool. - -**Using Istio instead of kGateway:** - -- The only Istio-specific step is creating an Istio `Gateway` / `HTTPRoute` - that references the operator-generated `InferencePool` as its `backendRef`. - The DGD, the generated pool, and the Rust EPP image are all unchanged. -- The operator targets the stable `inference.networking.k8s.io/v1` API group, - supported in Istio ≥ 1.27. Older Istio versions used the experimental - `inference.networking.x-k8s.io` group and are not compatible. -- **mTLS to the EPP.** Istio expects mTLS between the gateway and the EPP - service. The Rust EPP serves self-signed TLS on `9002` by default - (`DYN_SECURE_SERVING=true`). See *Service Mesh Integration (Istio)* below - for the `DestinationRule` the Dynamo Helm chart can generate so Istio - terminates the EPP's TLS correctly. - -> [!IMPORTANT] -> Model card discovery, worker liveness, KV-aware routing, and bookkeeping -> remain entirely in Dynamo's control. The `InferencePool` provides the -> data-plane envelope (which pods, which port, which EPP); Dynamo's -> discovery and the Rust EPP provide the routing intelligence inside that -> envelope. Customizing the pool selector by hand is supported but requires -> keeping it consistent with the operator's worker-pod labels — otherwise -> pods discovered by Dynamo will fail subset filtering and the EPP will -> return `RoutingFailed`. - -### 5. Deploy - -We provide an example for the Qwen vLLM below. -You have to deploy the Dynamo Graph and the `HTTPRoute`. -The example `http-route.yaml` resolves the `Gateway` in the same namespace as -the `HTTPRoute`, so the simplest path is to apply the route in the same -namespace where you installed the `Gateway` (i.e. `${NAMESPACE}`). If your -`Gateway` lives in a different namespace, add `parentRefs[].namespace` to point -at it explicitly: -```yaml - parentRefs: - - group: gateway.networking.k8s.io - kind: Gateway - name: inference-gateway - namespace: my-model # only needed if the Gateway is in a different namespace -``` - -```bash -cd -# kubectl get httproutes -n my-model # Make sure you do not have an incompatible HTTPRoute running, delete if so. -# Choose disagg or agg example -kubectl apply -f examples/backends/vllm/deploy/gaie/disagg.yaml -n my-model -# or -kubectl apply -f examples/backends/vllm/deploy/gaie/agg.yaml -n my-model -# make sure to apply the route -kubectl apply -f examples/backends/vllm/deploy/gaie/http-route.yaml -n my-model -``` - -Examples for other models can be found in the recipes folder. - -```bash -# Deploy PVC, having first Update `storageClassName` in recipes/llama-3-70b/model-cache/model-cache.yaml to match your cluster before deploying -kubectl apply -f recipes/llama-3-70b/model-cache/model-cache.yaml -n ${NAMESPACE} -kubectl apply -f recipes/llama-3-70b/model-cache/model-download.yaml -n ${NAMESPACE} -``` -We provide examples for llama-3-70b vLLM under the `recipes/llama-3-70b/vllm/agg/gaie/` for aggregated and `recipes/llama-3-70b/vllm/disagg-single-node/gaie/` for disaggregated serving. -Note for the aggregated serving you need to disable DYN_ENFORCE_DISAGG in epp config. -```bash - - name: DYN_ENFORCE_DISAGG - value: "false" -``` -Use the proper folder in commands below. - -```bash -# Deploy your Dynamo Graph. - -# agg -kubectl apply -f recipes/llama-3-70b/vllm/agg/gaie/deploy.yaml -n ${NAMESPACE} -# Deploy the GAIE http-route CR. The route resolves the Gateway in the same namespace by default; -# if your Gateway is elsewhere, add parentRefs[].namespace before applying. -kubectl apply -f recipes/llama-3-70b/vllm/agg/gaie/http-route.yaml -n ${NAMESPACE} - -# or disagg -kubectl apply -f recipes/llama-3-70b/vllm/disagg-single-node/gaie/deploy.yaml -n ${NAMESPACE} -kubectl apply -f recipes/llama-3-70b/vllm/disagg-single-node/gaie/http-route.yaml -n ${NAMESPACE} -``` - -- When using GAIE the FrontEnd does not choose the workers. The routing is determined in the EPP. -- The FrontEnd must run with `--router-mode direct` so that it respects the EPP's routing decisions passed via request headers. -- In v1beta1 DGD manifests, set the `frontendSidecar` field on a worker - component to the name of a container in that component's pod template. The - operator merges the required Dynamo env vars, probes, and ports into that - sidecar container: - -```yaml -frontendSidecar: sidecar-frontend -podTemplate: - spec: - containers: - - name: main - image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.3.0 - command: - - /bin/sh - - -c - args: - - python3 -m dynamo.vllm --model $MODEL_PATH --served-model-name $SERVED_MODEL_NAME - - name: sidecar-frontend - image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.3.0 - args: - - -m - - dynamo.frontend - - --router-mode - - direct - envFrom: - - secretRef: - name: hf-token-secret -``` - -- The pre-selected workers (decode and prefill in case of disaggregated serving) are passed in request headers and injected into the request routing hints. -- The `--router-mode direct` flag ensures the routing respects this selection. - -**Startup Probe Timeout:** The EPP has a default startup probe timeout of 30 minutes (10s × 180 failures). -If your model takes longer to load, increase the `failureThreshold` in the EPP's `startupProbe`. For example, -to allow 60 minutes for startup: - -```yaml -extraPodSpec: - mainContainer: - startupProbe: - failureThreshold: 360 # 10s × 360 = 60 minutes -``` - -**Gateway Namespace** -The example `http-route.yaml` resolves the `Gateway` in the same namespace as -the route. If you install the `Gateway` in one namespace and apply the route in -another, add `parentRefs[].namespace: ` to `http-route.yaml`. - -Common Vars for Routing Configuration: - -**Enabling KV-Aware Routing (most precise)** - -KV-aware routing uses live KV cache block events from workers so the EPP can route requests to the worker with the best prefix cache overlap. To enable it (default): - -1. **Workers — enable prefix caching and KV event publishing.** Each worker must publish KV cache events to event plane (NATS/ZMQ) so the EPP's router can track per-worker cache state. - - **vLLM:** Pass `--enable-prefix-caching` and `--kv-events-config '{"enable_kv_cache_events":true}'`. - - **SGLang:** Pass `--kv-events-config` with the appropriate endpoint. - - **TRT-LLM:** Pass `--publish-events-and-metrics`. -2. **EPP — leave `DYN_USE_KV_EVENTS` at its default (`true`).** The EPP subscribes to worker KV events via event plane (NATS/ZMQ) and uses them for prefix-overlap scoring. -3. **Block size — must be consistent.** The `--block-size` on all workers must match `DYN_KV_CACHE_BLOCK_SIZE` on the EPP (default: 128). Mismatched block sizes cause incorrect block hash computation. - -**Disabling KV-Aware Routing** - -To disable the EPP from listening for KV events (e.g., when prefix caching is off on workers, or for simpler load-balanced routing): - -1. **EPP:** Set `DYN_USE_KV_EVENTS=false`. The router falls back to approximate mode (routing decisions are tracked locally with TTL decay instead of live KV events from workers). -2. **Workers:** Pass `--no-enable-prefix-caching` to disable prefix caching entirely. Without prefix caching, no KV events are generated regardless of other flags. -3. **Optionally** set `DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT=0` on the EPP to skip prefix-overlap scoring altogether, making the router select workers based on load only. - -- Set `DYN_BUSY_THRESHOLD` to configure the upper bound on how "full" a worker can be (often derived from kv_active_blocks or other load metrics) before the router skips it. If the selected worker exceeds this value, routing falls back to the next best candidate. By default the value is negative meaning this is not enabled. -- Set `DYN_ENFORCE_DISAGG=true` (default: `false`) to control per-request behavior when prefill workers are unavailable: - - **`true` (recommended for disaggregated serving):** Requests fail with an error if prefill workers are not available. Use this when disaggregated serving is required and aggregated fallback is not acceptable. - - **`false` (default):** Requests gracefully fall back to aggregated mode (skip prefill, route directly to decode) when prefill workers are not available. When prefill workers appear later, subsequent requests automatically use disaggregated routing. -- Set `DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT` to control the device-local prefix-overlap credit multiplier, from 0.0 to 1.0. Higher values bias toward reusing workers with similar cached prefixes. (default: 1) -- Set `DYN_ROUTER_PREFILL_LOAD_SCALE` to scale adjusted prompt-side prefill load before decode blocks are added. (default: 1) -- Set `DYN_ROUTER_TEMPERATURE` (default: `0.0`) to soften or sharpen normalized worker sampling. Low temperature makes the router pick the top candidate deterministically; higher temperature lets lower-scoring workers through more often (exploration). -- `DYN_ROUTER_REPLICA_SYNC` — Enable replica synchronization (default: false) -- `DYN_ROUTER_TRACK_ACTIVE_BLOCKS` — Track active blocks (default: true) -- `DYN_ROUTER_TRACK_OUTPUT_BLOCKS` — Track output blocks during generation (default: false) -- `DYN_ROUTER_PREDICTED_TTL_SECS` — Enable predict-on-route entries with this TTL in seconds -- See the [KV cache routing design](../design-docs/router-design.md) for details. - - -**Service Mesh Integration (Istio)** - -When running under a service mesh such as Istio, the mesh sidecar proxy may conflict with the EPP's own TLS serving, causing connection failures (double-TLS). To avoid this, the mesh must be told how to connect to the EPP service via an Istio `DestinationRule`. - -The Dynamo operator can generate this DestinationRule automatically. Enable it by setting the `dynamo.serviceMesh` parameters when installing or upgrading the Dynamo platform Helm chart: - -```bash -helm install dynamo deploy/helm/charts/platform \ - --set dynamo.serviceMesh.enabled=true -``` - -Or equivalently in a custom values file: - -```yaml -dynamo: - serviceMesh: - enabled: true - provider: "istio" - istio: - tlsMode: "SIMPLE" - insecureSkipVerify: true -``` - -**Helm Parameters** - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `dynamo.serviceMesh.enabled` | bool | `false` | Enable automatic DestinationRule generation for EPP services. | -| `dynamo.serviceMesh.provider` | string | `"istio"` | Service mesh provider. Only `"istio"` is supported. | -| `dynamo.serviceMesh.istio.tlsMode` | string | `"SIMPLE"` | TLS mode for the DestinationRule. Supported values: `DISABLE`, `SIMPLE`, `MUTUAL`, `ISTIO_MUTUAL`. | -| `dynamo.serviceMesh.istio.insecureSkipVerify` | bool | `true` | Skip TLS certificate verification. Set to `true` when EPP uses self-signed certificates (the default). | - -> [!NOTE] -> The Istio CRDs (`networking.istio.io`) must be installed on the cluster before enabling this feature. The operator detects Istio availability at startup — if the CRDs are not present, DestinationRule reconciliation is skipped even when `serviceMesh.enabled` is `true`. - -When enabled, the operator produces a `DestinationRule` for each EPP service equivalent to: - -```yaml -apiVersion: networking.istio.io/v1beta1 -kind: DestinationRule -metadata: - name: -spec: - host: ..svc.cluster.local - trafficPolicy: - tls: - mode: SIMPLE - insecureSkipVerify: true -``` - -If you are **not** using the Dynamo operator's Helm chart, you must create this `DestinationRule` manually for each EPP service. Without it, Istio's default mTLS policy will conflict with the EPP's gRPC TLS endpoint. - -**Inference-gateway Istio sidecar exclusion** - -When namespace-level Istio sidecar injection is enabled (`istio-injection=enabled`), the agentgateway-proxy pod also receives an Istio sidecar. This sidecar intercepts the ext_proc gRPC connection from agentgateway-proxy to EPP (port 9002) and routes it through `PassthroughCluster`, which breaks the connection and causes all inference requests to return HTTP 500 with an empty body. - -The fix is to tell agentgateway to stamp `sidecar.istio.io/inject: "false"` on the proxy pod template so the Istio webhook skips that pod. EPP and worker pods still receive sidecars normally. - -You have two options depending on how you set up the gateway: - -***Option A: Per-gateway `AgentgatewayParameters` (recommended)*** - -This is what `install_gaie_crd_agentgateway.sh` does automatically. It only affects the `inference-gateway` proxy pods and leaves any other agentgateway-managed gateways untouched. - -1. Create an `AgentgatewayParameters` resource in **the same namespace as the `inference-gateway` Gateway** (e.g. `dynamo-cloud`). It must be co-located with the `Gateway` because the Gateway API `spec.infrastructure.parametersRef` is a `LocalParametersReference` — it has no `namespace` field. - - ```yaml - apiVersion: agentgateway.dev/v1alpha1 - kind: AgentgatewayParameters - metadata: - name: inference-gateway-params - namespace: dynamo-cloud # same as the Gateway - spec: - deployment: - spec: - template: - metadata: - annotations: - sidecar.istio.io/inject: "false" - ``` - - Apply it with server-side apply (recommended by agentgateway): - - ```bash - kubectl apply --server-side -n dynamo-cloud -f agentgateway-params.yaml - ``` - -2. Wire the existing `Gateway` to use it. If the Gateway already exists, patch it in place: - - ```bash - kubectl patch gateway inference-gateway -n dynamo-cloud --type='merge' -p '{ - "spec": { - "infrastructure": { - "parametersRef": { - "group": "agentgateway.dev", - "kind": "AgentgatewayParameters", - "name": "inference-gateway-params" - } - } - } - }' - ``` - - Or include the `infrastructure` block directly in your `Gateway` manifest: - - ```yaml - apiVersion: gateway.networking.k8s.io/v1 - kind: Gateway - metadata: - name: inference-gateway - namespace: dynamo-cloud - spec: - gatewayClassName: agentgateway - infrastructure: - parametersRef: - group: agentgateway.dev - kind: AgentgatewayParameters - name: inference-gateway-params - listeners: - - name: http - port: 80 - protocol: HTTP - ``` - -3. agentgateway will roll the proxy pod. Verify the new pod no longer has an `istio-proxy` container: - - ```bash - kubectl get pod -l gateway.networking.k8s.io/gateway-name=inference-gateway \ - -n dynamo-cloud \ - -o jsonpath='{.items[0].spec.containers[*].name}{"\n"}' - # Expect: agentgateway (NOT "agentgateway istio-proxy") - ``` - -***Option B: Patch the default `AgentgatewayParameters` CR (cluster-wide)*** - -The agentgateway controller creates a default `AgentgatewayParameters` resource named `agentgateway` in `agentgateway-system`. Any `Gateway` that does not set `spec.infrastructure.parametersRef` inherits this default. Patching it affects **all** agentgateway-managed proxies in the cluster. - -```bash -kubectl patch agentgatewayparameters agentgateway -n agentgateway-system \ - --type='merge' -p '{ - "spec": { - "deployment": { - "spec": { - "template": { - "metadata": { - "annotations": { - "sidecar.istio.io/inject": "false" - } - } - } - } - } - } -}' -``` - -Use Option A instead if you have multiple agentgateway-managed gateways in the cluster and only want the `inference-gateway` proxy to skip injection. - -The annotation is a no-op on clusters where Istio is not installed, so it is safe to set unconditionally. - -> [!NOTE] -> With both the `DestinationRule` (for EPP) and the `AgentgatewayParameters` sidecar exclusion (for agentgateway-proxy) in place, end-to-end GAIE inference works correctly under Istio namespace-level injection. - -### 6. Verify Installation ### - -Check that all resources are properly deployed: - -```bash -kubectl get inferencepool -n ${NAMESPACE} -kubectl get httproute -n ${NAMESPACE} -kubectl get service -n ${NAMESPACE} -kubectl get gateway -n ${NAMESPACE} -``` - -Sample output: - -```bash -# kubectl get inferencepool -NAME AGE -qwen-pool 33m - -# kubectl get httproute -NAME HOSTNAMES AGE -qwen-route 33m -``` - -### 7. Usage ### - -The Inference Gateway provides HTTP endpoints for model inference. - -#### 1: Populate gateway URL for your k8s cluster #### - -a. To test the integration in minikube, proceed as below: -Use minikube tunnel to expose the gateway to the host. This requires `sudo` access to the host machine. Alternatively, you can use port-forward to expose the gateway to the host as shown in alternative (b). - -```bash -# in first terminal -ps aux | grep "minikube tunnel" | grep -v grep # make sure minikube tunnel is not already running. -minikube tunnel # start the tunnel - -# in second terminal where you want to send inference requests -GATEWAY_URL=$(kubectl get svc inference-gateway -n my-model -o jsonpath='{.spec.clusterIP}') && echo $GATEWAY_URL -``` - -b. To test on a cluster use commands below: - -use port-forward to expose the gateway to the host - -```bash -# in first terminal -kubectl port-forward svc/inference-gateway 8000:80 -n ${NAMESPACE} -# for NAMESPACE use the namespace where the Gateway service was created, for example agentgateway-system - -# in second terminal where you want to send inference requests -GATEWAY_URL=http://localhost:8000 -``` - -#### 2: Check models deployed to inference gateway #### - -a. Query models: - -```bash -# in the second terminal where you GATEWAY_URL is set -curl $GATEWAY_URL/v1/models | jq . -# or if you added the host name to http route: -curl -H "Host: llama3-70b-disagg.example.com" $GATEWAY_URL/v1/models | jq . -``` - -Sample output: - -```json -{ - "data": [ - { - "created": 1753768323, - "id": "Qwen/Qwen3-0.6B", - "object": "object", - "owned_by": "nvidia" - } - ], - "object": "list" -} -``` - -b. Send inference request to gateway: - -```bash -MODEL_NAME="Qwen/Qwen3-0.6B" -curl $GATEWAY_URL/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "'"${MODEL_NAME}"'", - "messages": [ - { - "role": "user", - "content": "In the heart of Eldoria, an ancient land of boundless magic and mysterious creatures, lies the long-forgotten city of Aeloria. Once a beacon of knowledge and power, Aeloria was buried beneath the shifting sands of time, lost to the world for centuries. You are an intrepid explorer, known for your unparalleled curiosity and courage, who has stumbled upon an ancient map hinting at ests that Aeloria holds a secret so profound that it has the potential to reshape the very fabric of reality. Your journey will take you through treacherous deserts, enchanted forests, and across perilous mountain ranges. Your Task: Character Background: Develop a detailed background for your character. Describe their motivations for seeking out Aeloria, their skills and weaknesses, and any personal connections to the ancient city or its legends. Are they driven by a quest for knowledge, a search for lost familt clue is hidden." - } - ], - "stream":false, - "max_tokens": 30, - "temperature": 0.0 - }' -``` -or - -```bash -MODEL_NAME="RedHatAI/Llama-3.3-70B-Instruct-FP8-dynamic" -curl -H "Host: llama3-70b-disagg.example.com" http://localhost:8000/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "'"${MODEL_NAME}"'", - "messages": [ - { - "role": "user", - "content": "In the heart of Eldoria, an ancient land of boundless magic and mysterious creatures, lies the long-forgotten city of Aeloria. Once a beacon of knowledge and power, Aeloria was buried beneath the shifting sands of time, lost to the world for centuries. You are an intrepid explorer, known for your unparalleled curiosity and courage, who has stumbled upon an ancient map hinting at ests that Aeloria holds a secret so profound that it has the potential to reshape the very fabric of reality. Your journey will take you through treacherous deserts, enchanted forests, and across perilous mountain ranges. Your Task: Character Background: Develop a detailed background for your character. Describe their motivations for seeking out Aeloria, their skills and weaknesses, and any personal connections to the ancient city or its legends. Are they driven by a quest for knowledge, a search for lost familt clue is hidden." - } - ], - "stream":false, - "max_tokens": 30, - "temperature": 0.0 - }' -``` - -Sample inference output: - -```json -{ - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "audio": null, - "content": "\nOkay, I need to develop a character background for the user's query. Let me start by understanding the requirements. The character is an", - "function_call": null, - "refusal": null, - "role": "assistant", - "tool_calls": null - } - } - ], - "created": 1753768682, - "id": "chatcmpl-772289b8-5998-4f6d-bd61-3659b684b347", - "model": "Qwen/Qwen3-0.6B", - "object": "chat.completion", - "service_tier": null, - "system_fingerprint": null, - "usage": { - "completion_tokens": 29, - "completion_tokens_details": null, - "prompt_tokens": 196, - "prompt_tokens_details": null, - "total_tokens": 225 - } -} -``` - -***If you have more than one HTTPRoute running on the cluster*** -Add the host to your `http-route.yaml` and add the header -`curl -H "Host: llama3-70b-agg.example.com" ...` or `curl -H "Host: llama3-70b-disagg.example.com" http://localhost:8000/v1/models` - -```bash -spec: - hostnames: - - llama3-70b-agg.example.com -``` - -### 8. Deleting the installation ### - -If you need to uninstall run: - -```bash -kubectl delete dynamoGraphDeployment vllm-agg -helm uninstall dynamo-gaie -n my-model - -# To uninstall GAIE -# 1. Delete the inference-gateway -kubectl delete gateway inference-gateway --ignore-not-found - -# 2. Uninstall agentgateway helm releases -helm uninstall agentgateway -n agentgateway-system -helm uninstall agentgateway-crds -n agentgateway-system - -# 3. Delete the agentgateway-system namespace (optional, cleans up everything in it) -kubectl delete namespace agentgateway-system --ignore-not-found - -# 4. Delete the Inference Extension CRDs -IGW_LATEST_RELEASE=v1.5.0-rc.2 -kubectl delete -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/releases/download/${IGW_LATEST_RELEASE}/manifests.yaml --ignore-not-found - -# 5. Delete the Gateway API CRDs -GATEWAY_API_VERSION=v1.5.1 -kubectl delete -f https://github.com/kubernetes-sigs/gateway-api/releases/download/$GATEWAY_API_VERSION/standard-install.yaml --ignore-not-found -``` - -## Gateway API Inference Extension Integration - -This section documents the updated plugin implementation for Gateway API Inference Extension **v1.5.0-rc.2**. - -### Router bookkeeping operations - -EPP performs Dynamo router book keeping operations so the FrontEnd's Router does not have to sync its state. - - -### Header Routing Hints - -Since v1.5.0-rc.1, the EPP uses **headers and body mutations** for communicating routing decisions. -The plugins set HTTP headers for worker targeting and inject pre-computed token IDs -into the request body (`nvext.token_data`) so the frontend sidecar can skip redundant tokenization. - -#### Headers Set by Dynamo Plugins - -| Header | Description | Set By | -|--------|-------------|--------| -| `x-dynamo-worker-instance-id` | Primary worker ID (decode worker in disagg mode) | kv-aware-scorer | -| `x-dynamo-prefill-instance-id` | Prefill worker ID (disaggregated mode only) | kv-aware-scorer | diff --git a/fern/docs.yml b/fern/docs.yml index e4d4afce8fee..dd70dd568815 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -180,7 +180,9 @@ redirects: - source: "/dynamo/dev/kubernetes-deployment/deployment-guide/developing-with-tilt" destination: "/dynamo/dev/kubernetes-deployment/advanced-platform/developing-with-tilt" - source: "/dynamo/dev/kubernetes-deployment/deployment-guide/inference-gateway-gaie" - destination: "/dynamo/dev/integrations/kubernetes-integrations/gateway-api-inference-extension-gaie" + destination: "/dynamo/dev/kubernetes-deployment/request-routing/gateway-api-inference-extension/overview" + - source: "/dynamo/dev/integrations/kubernetes-integrations/gateway-api-inference-extension-gaie" + destination: "/dynamo/dev/kubernetes-deployment/request-routing/gateway-api-inference-extension/overview" # SGLang HiCache doc: canonical URL under Integrations > KV Cache Integrations. - source: "/dynamo/dev/backends/sg-lang/hi-cache" destination: "/dynamo/dev/integrations/kv-cache-integrations/hi-cache" diff --git a/recipes/qwen3-0.6b/vllm/agg/gaie/httproute.yaml b/recipes/qwen3-0.6b/vllm/agg/gaie/httproute.yaml index a3980dd5044c..a810406c0ba8 100644 --- a/recipes/qwen3-0.6b/vllm/agg/gaie/httproute.yaml +++ b/recipes/qwen3-0.6b/vllm/agg/gaie/httproute.yaml @@ -10,8 +10,6 @@ spec: - group: gateway.networking.k8s.io kind: Gateway name: inference-gateway - # Adjust to the namespace where your Gateway is installed. - namespace: kgateway-system rules: - matches: - headers: From 3cabd83e1f5f03783a81a185f1696d1f68cfa03c Mon Sep 17 00:00:00 2001 From: "Dr. Stefan Schimanski" Date: Tue, 30 Jun 2026 13:22:09 +0200 Subject: [PATCH 002/320] docs(inference-gateway): recover EPP developer build docs (#10993) Signed-off-by: Dr. Stefan Schimanski --- deploy/inference-gateway/epp/DEVEL.md | 166 +++++++++++++++++++++ deploy/inference-gateway/ext-proc/DEVEL.md | 145 ++++++++++++++++++ 2 files changed, 311 insertions(+) create mode 100644 deploy/inference-gateway/epp/DEVEL.md create mode 100644 deploy/inference-gateway/ext-proc/DEVEL.md diff --git a/deploy/inference-gateway/epp/DEVEL.md b/deploy/inference-gateway/epp/DEVEL.md new file mode 100644 index 000000000000..dec21c3d2665 --- /dev/null +++ b/deploy/inference-gateway/epp/DEVEL.md @@ -0,0 +1,166 @@ + + +# Dynamo EPP Development + +This directory contains the standard Dynamo Endpoint Picker Plugin (EPP) for Gateway API Inference +Extension (GAIE). The implementation builds a Rust Dynamo routing library and links it into the Go +EPP binary through CGO. + +Use this file for developer build, test, and image workflows. User-facing GAIE setup belongs in the +published Kubernetes Gateway API documentation. + +## Prerequisites + +- Docker with BuildKit/buildx for image builds. +- Go and a C toolchain for host-native Go builds. +- Rust and Cargo for host-native Dynamo library builds. +- `kind` only when using `make image-kind` or `make all-kind`. + +Image builds are self-contained and do not require a host-built Dynamo library. Host-native binary +builds do require the local Dynamo static library artifacts. + +## Image Builds + +From this directory: + +```bash +cd deploy/inference-gateway/epp +``` + +Build and load a local image: + +```bash +make image-load +``` + +Build with a temporary local buildx builder and load the image: + +```bash +make all +``` + +Build and push an image: + +```bash +export DOCKER_SERVER=ghcr.io/nvidia/dynamo +export IMAGE_TAG=ghcr.io/nvidia/dynamo/dynamo-epp: +make image-push +``` + +Build and push through the aggregate target: + +```bash +make all-push +``` + +Build, load, and import into a kind cluster: + +```bash +export KIND_CLUSTER=kind +make image-kind +``` + +Build and import through the aggregate target: + +```bash +make all-kind +``` + +Build and push a multi-architecture image: + +```bash +export DOCKER_SERVER=ghcr.io/nvidia/dynamo +export IMAGE_TAG=ghcr.io/nvidia/dynamo/dynamo-epp: +make image-multiarch-push +``` + +Useful image variables: + +| Variable | Default | Purpose | +|---|---|---| +| `DOCKER_SERVER` | `dynamo` | Registry or registry namespace used to form `IMAGE_REPO`. | +| `IMAGE_TAG` | `$(DOCKER_SERVER)/dynamo-epp:$(git describe ...)` | Full image reference to build. | +| `DYNAMO_DIR` | Repository root auto-detected from this directory | Named Docker build context for the Dynamo workspace. | +| `PLATFORMS` | Host architecture | Platform for local image builds. | +| `MULTIARCH_PLATFORMS` | `linux/amd64,linux/arm64` | Platforms for multi-architecture builds. | +| `DOCKER_PROXY` | unset | Optional image prefix or mirror for base images. | +| `EXTRA_BUILD_ARGS` | unset | Extra arguments passed to `docker buildx build`. | +| `KIND_CLUSTER` | `kind` | kind cluster name for `make image-kind`. | + +Run `make info` to print the resolved build values. + +Common image targets: + +| Target | Purpose | +|---|---| +| `make image-build` | Build the image with the default buildx builder. | +| `make image-load` | Build the image and load it into the local Docker daemon. | +| `make image-push` | Build the image and push it to `IMAGE_TAG`. | +| `make image-kind` | Build, load, and import the image into `KIND_CLUSTER`. | +| `make image-multiarch-push` | Build and push a multi-architecture image. | +| `make image-local-build` | Build with a temporary local buildx builder. | +| `make image-local-load` | Build with a temporary local buildx builder and load locally. | +| `make image-local-push` | Build with a temporary local buildx builder and push. | +| `make all` | Alias for `make image-local-load`. | +| `make all-push` | Alias for `make image-push`. | +| `make all-kind` | Alias for `make image-kind`. | + +## Host-Native Build + +Host-native `go build` needs the Dynamo C API static library and header copied into this project. +Build those artifacts first: + +```bash +make dynamo-lib +``` + +Then build the EPP binary: + +```bash +make build +``` + +The combined target builds the library and binary: + +```bash +make build-with-lib +``` + +The binary is written to: + +```text +deploy/inference-gateway/epp/bin/epp +``` + +`make dynamo-lib` builds `libdynamo_llm` from the repository root and copies: + +- `libdynamo_llm_capi.a` to `pkg/plugins/dynamo_kv_scorer/lib/` +- `llm_engine.h` to `pkg/plugins/dynamo_kv_scorer/include/` + +## Development Checks + +Run the Go checks from this directory: + +```bash +make fmt +make vet +make tidy +make test +``` + +`make test` runs with `CGO_ENABLED=1`, so it needs the same local C library artifacts as +host-native builds. + +## Cleaning + +Remove local Go build artifacts: + +```bash +make clean +``` + +This removes `bin/` and runs `go clean`. It does not clean the repository-level Cargo target +directory used by `make dynamo-lib`. diff --git a/deploy/inference-gateway/ext-proc/DEVEL.md b/deploy/inference-gateway/ext-proc/DEVEL.md new file mode 100644 index 000000000000..33724dae62cf --- /dev/null +++ b/deploy/inference-gateway/ext-proc/DEVEL.md @@ -0,0 +1,145 @@ + + +# Rust EPP Development + +This directory contains the native Rust Envoy `ext_proc` Endpoint Picker Plugin (EPP) for Gateway +API Inference Extension (GAIE). It builds a single Rust binary, `dynamo-ext-proc`, and does not use +the Go EPP or CGO bridge. + +Use this file for developer build, test, and image workflows. User-facing GAIE setup belongs in the +published Kubernetes Gateway API documentation. + +## Prerequisites + +- Rust and Cargo for host-native builds and tests. +- Docker with BuildKit/buildx for image builds. +- `kind` only when using `make image-kind`. + +The Makefile runs Cargo commands from the repository root so the crate participates in the full +Dynamo workspace. + +## Host-Native Build + +From this directory: + +```bash +cd deploy/inference-gateway/ext-proc +``` + +Build the release binary: + +```bash +make build +``` + +Build a debug binary: + +```bash +make build-debug +``` + +The release binary is written to: + +```text +target/release/dynamo-ext-proc +``` + +## Development Checks + +Run the Rust checks from this directory: + +```bash +make fmt +make check +make clippy +make test +``` + +These targets map to Cargo commands for the `dynamo-ext-proc` package. + +## Image Builds + +Build and load a local image: + +```bash +make image-load +``` + +Build and push an image: + +```bash +export DOCKER_SERVER=ghcr.io/nvidia/dynamo +export IMAGE_TAG=ghcr.io/nvidia/dynamo/dynamo-rust-epp: +make image-push +``` + +Build, load, and import into a kind cluster: + +```bash +export KIND_CLUSTER=kind +make image-kind +``` + +Build and push a multi-architecture image: + +```bash +export DOCKER_SERVER=ghcr.io/nvidia/dynamo +export IMAGE_TAG=ghcr.io/nvidia/dynamo/dynamo-rust-epp: +make image-multiarch-push +``` + +Useful image variables: + +| Variable | Default | Purpose | +|---|---|---| +| `DOCKER_SERVER` | `dynamo` | Registry or registry namespace used to form `IMAGE_REPO`. | +| `IMAGE_TAG` | `$(DOCKER_SERVER)/dynamo-rust-epp:$(git describe ...)` | Full image reference to build. | +| `DYNAMO_DIR` | Repository root auto-detected from this directory | Named Docker build context for the Dynamo workspace. | +| `PLATFORMS` | Host architecture | Platform for local image builds. | +| `MULTIARCH_PLATFORMS` | `linux/amd64,linux/arm64` | Platforms for multi-architecture builds. | +| `DOCKER_PROXY` | unset | Optional image prefix or mirror for base images. | +| `EXTRA_BUILD_ARGS` | unset | Extra arguments passed to `docker buildx build`. | +| `KIND_CLUSTER` | `kind` | kind cluster name for `make image-kind`. | + +Run `make info` to print the resolved build values. + +Common image targets: + +| Target | Purpose | +|---|---| +| `make image-build` | Build the image with the default buildx builder. | +| `make image-load` | Build the image and load it into the local Docker daemon. | +| `make image-push` | Build the image and push it to `IMAGE_TAG`. | +| `make image-kind` | Build, load, and import the image into `KIND_CLUSTER`. | +| `make image-multiarch-push` | Build and push a multi-architecture image. | +| `make image-local-build` | Build with a temporary local buildx builder. | +| `make image-local-load` | Build with a temporary local buildx builder and load locally. | +| `make image-local-push` | Build with a temporary local buildx builder and push. | + +## Runtime Notes for Developers + +The Rust EPP serves Envoy `ext_proc` gRPC on port `9002` and plaintext gRPC health on port `9003`. +It serves TLS on the `ext_proc` port by default. Set `DYN_SECURE_SERVING=false` only for local +debugging with a plaintext h2c gateway. + +The common local environment variables are: + +| Variable | Default | Purpose | +|---|---|---| +| `DYN_NAMESPACE_PREFIX` | unset | Preferred Dynamo discovery namespace prefix. | +| `DYN_NAMESPACE` | unset | Exact Dynamo discovery namespace fallback. If unset, the binary uses `vllm-agg`. | +| `DYN_COMPONENT_NAME` | `backend` | Dynamo component that exposes the `generate` endpoint. | +| `DYN_ENFORCE_DISAGG` | `false` | Fail when prefill routing is unavailable instead of falling back to aggregated routing. | +| `DYN_KUBE_DISCOVERY_MODE` | `pod` | Kubernetes discovery identity mode. The Rust EPP currently rejects `container`. | +| `RUST_LOG` | `info` | Tracing log filter. | + +## Cleaning + +Clean the Rust package build artifacts: + +```bash +make clean +``` From af3c1eef509e890b1ec31b2883bcaec6c0438037 Mon Sep 17 00:00:00 2001 From: Yan Ru Pei Date: Tue, 30 Jun 2026 08:17:19 -0700 Subject: [PATCH 003/320] feat(router): coordinate session affinity across replicas (#11079) Signed-off-by: PeaBrane --- Cargo.lock | 10 - .../configuration/groups/router_args.py | 8 +- docs/agents/session-ids.md | 19 +- docs/components/frontend/nvext.md | 12 +- .../components/router/router-configuration.md | 62 +- lib/bindings/kvbm/Cargo.lock | 10 - lib/bindings/python/Cargo.lock | 10 - .../src/kv_router/prefill_router/admission.rs | 4 +- lib/llm/src/kv_router/prefill_router/mod.rs | 9 +- lib/llm/src/kv_router/push_router.rs | 214 +++---- .../kv_router/push_router/request_guard.rs | 3 +- .../src/kv_router/push_router/selection.rs | 33 +- lib/llm/src/session_affinity/coordinator.rs | 478 ++++++++++----- lib/llm/src/session_affinity/mod.rs | 7 +- lib/llm/src/session_affinity/push_router.rs | 479 ++++++++------- lib/llm/src/session_affinity/tests.rs | 564 +++++++++++++----- lib/runtime/Cargo.toml | 3 +- lib/runtime/examples/Cargo.lock | 12 +- lib/runtime/src/discovery/kube.rs | 32 +- lib/runtime/src/discovery/kv_store.rs | 562 ++++++++++++++++- lib/runtime/src/discovery/mod.rs | 52 ++ .../pipeline/network/egress/push_router.rs | 43 ++ lib/runtime/src/storage/kv.rs | 4 + lib/runtime/src/storage/kv/file.rs | 96 ++- tests/router/common.py | 211 +++++++ tests/router/router_process.py | 8 +- tests/router/test_router_e2e_with_mockers.py | 39 ++ 27 files changed, 2231 insertions(+), 753 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ddba6950726e..7c2eb88d945c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3433,15 +3433,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -5776,7 +5767,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.11.1", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", diff --git a/components/src/dynamo/common/configuration/groups/router_args.py b/components/src/dynamo/common/configuration/groups/router_args.py index f4bb611b8769..1c9b334bcd94 100644 --- a/components/src/dynamo/common/configuration/groups/router_args.py +++ b/components/src/dynamo/common/configuration/groups/router_args.py @@ -244,9 +244,11 @@ def add_arguments(self, parser) -> None: env_var="DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", default=None, help=( - "Enable router-local session affinity with this idle TTL in seconds. " - "Affinity is disabled when this option is omitted. " - "This is independent of KV prediction TTL settings." + "Enable session affinity and set the process-local cache eviction TTL " + "in seconds. etcd and shared FileStore use immutable distributed claims " + "whose lifetime follows the creating frontend, not this TTL. Memory and " + "Kubernetes discovery remain process-local. Affinity is disabled when " + "this option is omitted." ), arg_type=int, dest="session_affinity_ttl_secs", diff --git a/docs/agents/session-ids.md b/docs/agents/session-ids.md index 6efb096f1ab3..c8aa3b4d36c3 100644 --- a/docs/agents/session-ids.md +++ b/docs/agents/session-ids.md @@ -7,7 +7,11 @@ subtitle: Identify agent sessions from supported coding agents and custom client A session ID is the stable identifier Dynamo uses for one agent reasoning/tool chain. A root agent, planner, researcher subagent, or OpenCode subtask can each have its own session. Every LLM request in that chain should carry the same `session_id`; child sessions can also carry a `parent_session_id` so traces and replay tools can rebuild the tree. Some academic papers also call this a `program_id`. -Session identity is passive metadata. Sending `X-Dynamo-Session-ID` does not enable sticky sessions or change request placement. Tracing records the identity when `DYN_REQUEST_TRACE` is enabled, and a session-aware routing policy can consume it only when that policy is configured separately. +Session identity is passive metadata unless session affinity is explicitly enabled. +Sending `X-Dynamo-Session-ID` alone does not change request placement. Tracing records +the identity when `DYN_REQUEST_TRACE` is enabled. When +`--router-session-affinity-ttl-secs` is configured, the router uses the ID for an +immutable endpoint- and phase-scoped worker binding. ## Session ID Inputs @@ -29,7 +33,18 @@ Dynamo also recognizes the current stable identity headers emitted by the follow | Codex | `session-id` | None | `session-id` becomes the `session_id`. | | OpenCode | `x-session-id` | `x-parent-session-id` | `x-session-id` becomes the `session_id`; `x-parent-session-id` becomes `parent_session_id` when present. | -`X-Dynamo-Session-Final` applies with either canonical or agent-native session identity. +`X-Dynamo-Session-Final` applies with either canonical or agent-native session +identity. With session affinity enabled, a final request routes normally and then +terminally closes its binding. Close invalidation across replicas is eventual. Do not +send more requests with that session ID after close. + +etcd and FileStore on a shared filesystem coordinate bindings across frontend +processes. MemoryStore and Kubernetes discovery retain process-local affinity only. +The affinity TTL controls local cache cleanup, not the distributed claim lifetime. +Claims follow the creating frontend's etcd lease or FileStore ownership. If a claim +expires or its bound worker disappears, create a new session with a new ID instead of +reusing or rebinding the old ID. See [Router session affinity](../components/router/router-configuration.md#session-affinity) +for the full contract. ### Custom Agent Harnesses diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index 0970e1ad9524..ad7b41457537 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -54,7 +54,7 @@ token IDs, pass integer IDs in the normal `stop` array, for example `"stop": [576]`. Strings such as `"token_id:576"` remain literal string stop sequences and are not parsed as token IDs. -### Header Overrides +### Header overrides Routing fields can also be set via HTTP headers, which take priority over `nvext` values: @@ -75,9 +75,13 @@ session headers described in [Session IDs](../../agents/session-ids.md); `nvext` does not accept session identity fields. When session affinity is enabled with `--router-session-affinity-ttl-secs`, the -router also uses `X-Dynamo-Session-ID` for router-local affinity. See -[Configuration and Tuning](../router/router-configuration.md#session-affinity) -for routing behavior and TTL settings. +router uses `X-Dynamo-Session-ID` for immutable endpoint- and phase-scoped affinity. +On etcd and shared FileStore, replicas coordinate through a distributed claim while +the request hot path uses a process-local cache. Existing local or shared bindings +override routing headers; the headers above are proposals only when no binding exists. +Memory and Kubernetes discovery do not provide cross-process affinity. See +[Configuration and Tuning](../router/router-configuration.md#session-affinity) for +claim lifetime, cache TTL, terminal close, and failure behavior. For trace sink configuration and JSONL schema details, see [Agent Tracing](../../agents/agent-tracing.md). diff --git a/docs/components/router/router-configuration.md b/docs/components/router/router-configuration.md index 09e8e13a9aa9..17be76ca6bb6 100644 --- a/docs/components/router/router-configuration.md +++ b/docs/components/router/router-configuration.md @@ -111,31 +111,43 @@ a value from `1` through `31536000` to enable it, then send `X-Dynamo-Session-ID` to keep related requests on one worker. Supplying the header without the TTL option provides session identity but does not enable router affinity. -The first successfully dispatched request binds the session ID to its selected -worker and, when available, data-parallel rank. Later requests exact-dispatch to -that target without transport fallback. Concurrent requests can share a binding. -Active requests prevent expiry. When a request lease ends after EOF, early drop, -error, or cancellation, the idle timer restarts. A missing bound worker or a -non-cancellation selection, setup, dispatch, or target-validation failure invalidates -the binding. - -The configured value is the idle timeout. It is independent of -`--router-ttl-secs` and `--router-predicted-ttl-secs`. Omit the session-affinity -option to keep affinity disabled. - -If the bound worker disappears, Dynamo invalidates the binding so a subsequent -selection can bind an available worker. Router restart clears all bindings. Bindings -are not shared between frontend replicas. - -Direct mode still requires the phase-appropriate explicit worker ID on every -affinity request. The stored binding validates that target but does not supply a -missing ID. In disaggregated serving, prefill and decode use separate phase-local -bindings. If no prefill router is active, only the decode or aggregated binding is -created. - -Session affinity does not create a backend session or send lifecycle RPCs. There is -no explicit unbind; idle expiry removes only router-local state. The same session -ID is available to tracing and other explicitly configured consumers. +The first affinity request creates one immutable binding from the session ID to a +worker and, when available, a data-parallel rank. The binding is scoped to the +existing endpoint and phase, so disaggregated prefill and decode routes remain +separate. Later requests exact-dispatch to that target without transport fallback. +An existing local or shared binding takes precedence over explicit routing headers; +those headers are proposals only while the claim is absent. Direct mode therefore +requires an explicit target for a new binding, but an existing binding supplies the +target for later requests. Query-only requests remain read-only and do not create or +close claims. + +With etcd or FileStore on a filesystem shared by all replicas, frontends coordinate +through an immutable distributed claim. The existing-session hot path reads only the +process-local cache. A cache miss reads shared storage first and attempts an atomic +insertion only when the claim is absent. Racing frontends all cache and dispatch to +the stored winner. Storage errors fail the request before scheduler bookkeeping or +dispatch. MemoryStore coordinates only callers sharing the same process and store. +Kubernetes discovery does not provide cross-process affinity and keeps process-local +behavior. + +For distributed backends, `--router-session-affinity-ttl-secs` controls only +process-local cache eviction. A cache miss after local eviction reloads the immutable +claim. The claim itself follows the creating frontend's existing etcd lease or +FileStore ownership lifetime; it is not a global idle-session timeout. Delete events +eventually invalidate other frontend caches. Watch lag, disconnect, or restart clears +the entire local affinity cache, and later requests reload claims on demand. + +`X-Dynamo-Session-Final: true` marks a terminal request. Dynamo routes that request +normally, then evicts the closing frontend's cache entry and idempotently deletes the +shared claim. Other replicas observe the delete eventually. Close must not race active +requests, and callers must not use that session ID again. The same no-reuse rule +applies after claim expiry. If the bound worker disappears while the claim exists, +exact dispatch fails; start a new session with a new session ID. + +Global idle-session TTL, rebinding, dead-worker replacement, compare-and-swap updates, +fencing, generations, broader `WorkerSet` affinity, and backend-tokenized path +expansion are outside this contract. The setting remains independent of +`--router-ttl-secs` and `--router-predicted-ttl-secs`; omit it to disable affinity. ### AIC Prefill Load Model diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index bf75e62b27d6..3b84de82f913 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -2252,15 +2252,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -4066,7 +4057,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.12.1", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index f1cbbb3c1801..710d5fc65fe8 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2994,15 +2994,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -5096,7 +5087,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.13.0", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", diff --git a/lib/llm/src/kv_router/prefill_router/admission.rs b/lib/llm/src/kv_router/prefill_router/admission.rs index b6ecf578bed9..d4e93ab1195a 100644 --- a/lib/llm/src/kv_router/prefill_router/admission.rs +++ b/lib/llm/src/kv_router/prefill_router/admission.rs @@ -20,7 +20,7 @@ use crate::{ llm_backend::{LLMEngineOutput, PreprocessedRequest}, timing::RequestTracker, }, - session_affinity::SessionAffinityPushRouter, + session_affinity::{AffinityTarget, SessionAffinityPushRouter}, }; pub(super) enum InnerPrefillRouter { @@ -35,7 +35,7 @@ impl InnerPrefillRouter { prepare: F, ) -> Result<(M, ManyOut>)> where - F: FnOnce(&mut PreprocessedRequest, u64, Option) -> Result, + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { match self { InnerPrefillRouter::KvRouter(router) => { diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index cc6c42d28cad..4d68b5ec689a 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -28,6 +28,7 @@ use crate::{ preprocessor::{BootstrapInfo, PrefillResult, TraceLink}, timing::{RequestPhase, RequestTracker}, }, + session_affinity::AffinityTarget, }; mod activation; @@ -224,8 +225,8 @@ impl .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; let prefill_result: Result<(PrefillOutcome, Option)> = async { let (prepared, prefill_stream) = router - .select_and_dispatch_prefill(prefill_context, |request, worker_id, dp_rank| { - self.prepare_prefill_dispatch(request, worker_id, dp_rank) + .select_and_dispatch_prefill(prefill_context, |request, target| { + self.prepare_prefill_dispatch(request, target) }) .await?; let topology_constraints = prepared.topology_constraints; @@ -327,9 +328,9 @@ impl PrefillRouter { fn prepare_prefill_dispatch( &self, request: &mut PreprocessedRequest, - worker_id: u64, - dp_rank: Option, + target: AffinityTarget, ) -> anyhow::Result { + let AffinityTarget { worker_id, dp_rank } = target; let endpoint_id = self.endpoint_id.get(); let topology_constraints = self.preflight_kv_transfer_constraints(endpoint_id, worker_id)?; diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 0c194ece4f10..aa35639a2fe6 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -5,13 +5,14 @@ use std::{sync::Arc, time::Duration}; use dynamo_kv_router::protocols::{TokensWithHashes, WorkerWithDpRank}; use dynamo_runtime::{ - error::{ErrorType, match_error_chain}, + discovery::ClaimPayloadFuture, metrics::frontend_perf::{STAGE_ROUTE, StageGuard}, pipeline::{ AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, PushRouter, ResponseStream, SingleIn, async_trait, }, protocols::annotated::Annotated, + traits::DistributedRuntimeProvider, }; use futures::stream::{self, StreamExt}; use tracing::Instrument; @@ -24,7 +25,7 @@ use crate::{ timing::{RequestPhase, RoutingData}, }, session_affinity::{ - AffinityAcquire, AffinityCoordinator, AffinityTarget, affinity_id, explicit_target, + AffinityCoordinator, AffinityTarget, ResolvedAffinity, affinity_id, session_final, }, }; @@ -52,7 +53,13 @@ impl KvPushRouter { session_affinity_ttl: Option, ) -> Result { let affinity = session_affinity_ttl - .map(AffinityCoordinator::new) + .map(|ttl| { + AffinityCoordinator::new_distributed( + ttl, + inner.client.endpoint.id().to_string(), + inner.client.endpoint.drt().discovery(), + ) + }) .transpose()?; // Eagerly register router request metrics (as zeros) so they are @@ -121,7 +128,7 @@ impl KvPushRouter { request: &SingleIn, phase: RequestPhase, is_query_only: bool, - ) -> Result<(WorkerSelection, Option), Error> { + ) -> Result<(WorkerSelection, Option), Error> { let Some(affinity) = self.affinity.as_ref() else { return Ok(( self.select_request(request, phase, is_query_only, None) @@ -136,9 +143,8 @@ impl KvPushRouter { None, )); }; - let explicit = explicit_target(request, phase)?; if is_query_only { - let target = affinity.query_target(&session_id, explicit)?; + let target = affinity.query_target(&session_id)?; let worker = target.and_then(affinity_worker); return Ok(( self.select_request(request, phase, true, worker).await?, @@ -148,32 +154,23 @@ impl KvPushRouter { let request_context = request.context(); let operation = affinity - .acquire_with_context(&session_id, explicit, request_context.as_ref()) + .acquire_with_context(&session_id, request_context.as_ref()) .await?; - let worker = operation.target().and_then(affinity_worker); - match self.select_request(request, phase, false, worker).await { - Ok(selection) => Ok((selection, Some(operation))), - Err(error) if match_error_chain(error.as_ref(), &[ErrorType::Cancelled], &[]) => { - Err(error) - } - Err(_) if operation.target().is_some() && explicit.is_none() => { - operation.invalidate(); - let retry = affinity - .acquire_with_context(&session_id, None, request_context.as_ref()) - .await?; - match self.select_request(request, phase, false, None).await { - Ok(selection) => Ok((selection, Some(retry))), - Err(retry_error) => { - retry.invalidate(); - Err(retry_error) - } - } - } - Err(error) => { - operation.invalidate(); - Err(error) - } - } + let resolved = operation + .resolve(|| -> ClaimPayloadFuture<'_> { + Box::pin(async { + let selection = self.select_request(request, phase, true, None).await?; + let target = AffinityTarget { + worker_id: selection.instance_id, + dp_rank: Some(selection.dp_rank), + }; + Ok(serde_json::to_value(target)?) + }) + }) + .await?; + let worker = affinity_worker(resolved.target()); + let selection = self.select_request(request, phase, false, worker).await?; + Ok((selection, Some(resolved))) } async fn track_selection( @@ -388,58 +385,38 @@ impl KvPushRouter { prepare: F, ) -> Result<(M, ManyOut>), Error> where - F: FnOnce(&mut PreprocessedRequest, u64, Option) -> Result, + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { let phase = RequestPhase::Prefill; let phase_label = phase.to_string(); let route_guard = StageGuard::new(STAGE_ROUTE, &phase_label); let is_query_only = request.get_annotation_value("query_instance_id").is_some(); - let (mut selection, mut operation) = self + let close_on_finish = !is_query_only && session_final(request.content()); + let (mut selection, operation) = self .select_with_affinity(&request, phase, is_query_only) .await?; - let mut guard = match self + let mut guard = self .track_selection(&request, &mut selection, is_query_only) - .await - { - Ok(guard) => guard, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } + .await?; + let target = AffinityTarget { + worker_id: selection.instance_id, + dp_rank: Some(selection.dp_rank), }; - let metadata = match prepare(&mut request, selection.instance_id, Some(selection.dp_rank)) { + let metadata = match prepare(&mut request, target) { Ok(metadata) => metadata, Err(error) => { guard.abort().await; - if let Some(operation) = operation.take() { - operation.invalidate(); - } return Err(error); } }; - let selected_target = AffinityTarget { - worker_id: selection.instance_id, - dp_rank: Some(selection.dp_rank), - }; drop(route_guard); - let stream = match self + let stream = self .dispatch_selection(request, selection, guard, true) - .await - { - Ok(stream) => stream, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } - }; + .await?; let Some(operation) = operation else { return Ok((metadata, stream)); }; - Ok((metadata, operation.into_stream(selected_target, stream)?)) + Ok((metadata, operation.into_stream(stream, close_on_finish))) } } @@ -478,7 +455,8 @@ impl AsyncEngine, ManyOut, ManyOut guard, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } - }; + let guard = self + .track_selection(&request, &mut selection, false) + .await?; drop(route_guard); - let selected_target = AffinityTarget { - worker_id: selection.instance_id, - dp_rank: Some(selection.dp_rank), - }; - let stream = match self + let stream = self .dispatch_selection(request, selection, guard, operation.is_some()) - .await - { - Ok(stream) => stream, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } - }; + .await?; match operation { - Some(operation) => operation.into_stream(selected_target, stream), + Some(operation) => Ok(operation.into_stream(stream, close_on_finish)), None => Ok(stream), } } @@ -653,7 +612,10 @@ mod tests { .unwrap(); let endpoint = component.endpoint("generate"); let client = endpoint.client().await.unwrap(); - let workers = HashMap::from([(7, ModelRuntimeConfig::default())]); + let workers = HashMap::from([ + (7, ModelRuntimeConfig::default()), + (8, ModelRuntimeConfig::default()), + ]); let (_tx, workers) = watch::channel(workers); let config = KvRouterConfig { skip_initial_worker_wait: true, @@ -701,17 +663,17 @@ mod tests { worker_id: 7, dp_rank: Some(0), }; - let AffinityAcquire::Initialize(initializer) = router + let resolved = router .affinity .as_ref() .unwrap() - .acquire(&session_id, None) + .acquire(&session_id) .await .unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(original_target).unwrap()); + .resolve(|| Box::pin(async move { Ok(serde_json::to_value(original_target)?) })) + .await + .unwrap(); + drop(resolved); let controller = Controller::new("cancelled-selection-request".to_string()); controller.stop(); @@ -734,23 +696,69 @@ mod tests { .affinity .as_ref() .unwrap() - .query_target(&session_id, None) + .query_target(&session_id) .unwrap(), Some(original_target) ); - let AffinityAcquire::Bound { target, lease } = router + let resolved = router .affinity .as_ref() .unwrap() - .acquire(&session_id, None) + .acquire(&session_id) .await .unwrap() - else { - panic!("cancellation must preserve the existing binding"); + .resolve(|| { + Box::pin(async move { + Ok(serde_json::to_value(AffinityTarget { + worker_id: 8, + dp_rank: Some(0), + })?) + }) + }) + .await + .unwrap(); + assert_eq!(resolved.target(), original_target); + + drop(router); + runtime.shutdown(); + } + + #[tokio::test] + async fn session_affinity_binding_overrides_conflicting_explicit_proposal() { + let (router, runtime) = router(Some(Duration::from_secs(10))).await; + let session_id = SessionAffinityId::new("conflicting-explicit-proposal"); + let bound_target = AffinityTarget { + worker_id: 7, + dp_rank: Some(0), }; - assert_eq!(target, original_target); - drop(lease); + let resolved = router + .affinity + .as_ref() + .unwrap() + .acquire(&session_id) + .await + .unwrap() + .resolve(|| Box::pin(async move { Ok(serde_json::to_value(bound_target)?) })) + .await + .unwrap(); + drop(resolved); + + let mut content = request(); + content.routing_mut().backend_instance_id = Some(8); + content.routing_mut().decode_worker_id = Some(8); + content.routing_mut().dp_rank = Some(0); + let mut request = Context::new(content); + request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id); + + let (selection, resolved) = router + .select_with_affinity(&request, RequestPhase::Aggregated, false) + .await + .unwrap(); + assert_eq!(selection.instance_id, 7); + assert_eq!(selection.dp_rank, 0); + assert_eq!(resolved.unwrap().target(), bound_target); + router.chooser.free(request.context().id()).await.unwrap(); drop(router); runtime.shutdown(); diff --git a/lib/llm/src/kv_router/push_router/request_guard.rs b/lib/llm/src/kv_router/push_router/request_guard.rs index 6f0df3cecf15..a944b22d0fc4 100644 --- a/lib/llm/src/kv_router/push_router/request_guard.rs +++ b/lib/llm/src/kv_router/push_router/request_guard.rs @@ -245,8 +245,7 @@ impl OutputBlockTracker { /// Coordinates scheduler cleanup, observability, and streamed load tracking. /// -/// Session-affinity lifetime is separate: `AffinityAcquire` and -/// `AffinityLease` own binding commit, release, and invalidation. +/// Session-affinity lifetime is separate: `ResolvedAffinity` owns the binding lease. pub(super) struct RequestGuard { cleanup: RequestCleanup, observability: RequestObservability, diff --git a/lib/llm/src/kv_router/push_router/selection.rs b/lib/llm/src/kv_router/push_router/selection.rs index 980c676383c0..5ac97db65182 100644 --- a/lib/llm/src/kv_router/push_router/selection.rs +++ b/lib/llm/src/kv_router/push_router/selection.rs @@ -140,9 +140,7 @@ impl KvPushRouter { let affinity_pin = options .affinity_worker .map(|worker| (worker.worker_id, Some(worker.dp_rank))); - let Some((pinned_worker_id, requested_dp_rank)) = - merge_affinity_pin(explicit_pin, affinity_pin) - else { + let Some((pinned_worker_id, requested_dp_rank)) = affinity_pin.or(explicit_pin) else { let _nvtx_kv = dynamo_nvtx_range!("route.kv_match"); let selection = self .select_best_match(BestMatchArgs { @@ -235,21 +233,6 @@ impl KvPushRouter { } } -fn merge_affinity_pin( - explicit: Option<(u64, Option)>, - affinity: Option<(u64, Option)>, -) -> Option<(u64, Option)> { - match (explicit, affinity) { - (Some((worker_id, None)), Some((affinity_worker_id, affinity_rank))) - if worker_id == affinity_worker_id => - { - Some((worker_id, affinity_rank)) - } - (Some(explicit), _) => Some(explicit), - (None, affinity) => affinity, - } -} - fn resolve_pinned_worker_rank( worker_id: WorkerId, requested_dp_rank: Option, @@ -295,7 +278,7 @@ mod tests { scheduling::{RoutingEligibility, WorkerEligibilityError}, }; - use super::{merge_affinity_pin, pinned_worker_hint, resolve_pinned_worker_rank}; + use super::{pinned_worker_hint, resolve_pinned_worker_rank}; use crate::{ local_model::runtime_config::ModelRuntimeConfig, protocols::common::{preprocessor::RoutingHints, timing::RequestPhase}, @@ -323,18 +306,6 @@ mod tests { assert!(error.contains("requires an explicit dp_rank")); } - #[test] - fn affinity_pin_supplies_rank_for_matching_explicit_worker() { - assert_eq!( - merge_affinity_pin(Some((7, None)), Some((7, Some(0)))), - Some((7, Some(0))) - ); - assert_eq!( - merge_affinity_pin(Some((7, Some(2))), Some((7, Some(3)))), - Some((7, Some(2))) - ); - } - #[test] fn pinned_worker_hint_prefill_uses_prefill_worker_before_backend() { let routing = RoutingHints { diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index 7702309b3bd1..f981fef307da 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -1,6 +1,21 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +//! Session affinity uses a process-local cache in front of immutable shared claims. +//! Cache hits exact-route without distributed I/O. On a cache miss, the discovery +//! backend reads the claim first and evaluates the query-only routing proposal only +//! when no claim exists. The payload returned by claim arbitration is authoritative: +//! distributed mode caches only `Created` or `Existing` payloads, and racing losers +//! discard their proposal, cache the winner, and dispatch to it. Explicit worker and +//! rank headers are proposals only while no binding exists. +//! +//! Shared claim deletion invalidates local caches eventually through `Delete` and +//! `Reset` events. The configured affinity TTL evicts only local cache entries; it +//! does not expire or replace a shared claim. Bindings are never rebound in v1. If a +//! bound worker disappears, exact dispatch fails without fallback and the caller must +//! use a new session ID. Explicit close requires no concurrent active requests, is +//! terminal, and the closed session ID must not be reused. + use std::{ pin::Pin, sync::{ @@ -13,12 +28,17 @@ use std::{ use dashmap::{DashMap, mapref::entry::Entry}; use dynamo_runtime::{ + discovery::{ClaimEvent, ClaimOutcome, ClaimPayloadFuture, Discovery}, engine::{AsyncEngineContext, AsyncEngineContextProvider}, error::{DynamoError, ErrorType}, pipeline::{Error, ManyOut, ResponseStream}, }; use futures::Stream; -use tokio::{sync::Notify, time::Instant}; +use serde::{Deserialize, Serialize}; +use tokio::{ + sync::{Notify, broadcast}, + time::Instant, +}; use tokio_util::sync::CancellationToken; use super::{ @@ -33,7 +53,7 @@ use crate::{ }, }; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] pub struct AffinityTarget { pub worker_id: u64, pub dp_rank: Option, @@ -54,6 +74,7 @@ enum AffinityEntry { struct AffinityCoordinatorInner { entries: DashMap, + claims: ClaimCoordination, ttl: Duration, max_entries: usize, max_session_id_bytes: usize, @@ -61,11 +82,71 @@ struct AffinityCoordinatorInner { next_revision: AtomicU64, cancel: CancellationToken, #[cfg(test)] + probe: AffinityCoordinatorProbe, +} + +#[derive(Clone)] +struct ClaimCoordination { + scope: String, + discovery: Option>, +} + +impl ClaimCoordination { + fn key(&self, session_id: &SessionAffinityId) -> String { + format!( + "{}/{}", + self.scope, + blake3::hash(session_id.as_str().as_bytes()).to_hex() + ) + } + + fn subscribe(&self) -> Option> { + self.discovery + .as_ref() + .and_then(|discovery| discovery.subscribe_claim_events()) + } + + async fn resolve( + &self, + key: &str, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result<(serde_json::Value, bool), Error> { + let Some(discovery) = self.discovery.as_ref() else { + return Ok((proposed_payload.as_mut().await?, true)); + }; + + match discovery.create_or_get_claim(key, proposed_payload).await? { + ClaimOutcome::Created(payload) => Ok((payload, true)), + ClaimOutcome::Existing(payload) => Ok((payload, false)), + ClaimOutcome::Unsupported => Ok((proposed_payload.as_mut().await?, true)), + } + } + + async fn close(&self, key: &str) -> anyhow::Result<()> { + let Some(discovery) = self.discovery.as_ref() else { + return Ok(()); + }; + discovery.close_claim(key).await?; + Ok(()) + } +} + +#[cfg(test)] +struct AffinityCoordinatorProbe { reaper_started: Arc, - #[cfg(test)] waiter_observed: Arc, } +#[cfg(test)] +impl AffinityCoordinatorProbe { + fn new() -> Self { + Self { + reaper_started: Arc::new(Notify::new()), + waiter_observed: Arc::new(Notify::new()), + } + } +} + impl Drop for AffinityCoordinatorInner { fn drop(&mut self) { self.cancel.cancel(); @@ -83,6 +164,22 @@ impl AffinityCoordinator { ttl, MAX_SESSION_AFFINITY_ENTRIES, MAX_SESSION_AFFINITY_ID_BYTES, + "local".to_string(), + None, + ) + } + + pub(crate) fn new_distributed( + ttl: Duration, + claim_scope: String, + discovery: Arc, + ) -> Result { + Self::new_with_limits( + ttl, + MAX_SESSION_AFFINITY_ENTRIES, + MAX_SESSION_AFFINITY_ID_BYTES, + claim_scope, + Some(discovery), ) } @@ -90,6 +187,8 @@ impl AffinityCoordinator { ttl: Duration, max_entries: usize, max_session_id_bytes: usize, + claim_scope: String, + discovery: Option>, ) -> Result { if !(Duration::from_secs(1)..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) .contains(&ttl) @@ -100,6 +199,10 @@ impl AffinityCoordinator { } let inner = Arc::new(AffinityCoordinatorInner { entries: DashMap::new(), + claims: ClaimCoordination { + scope: claim_scope, + discovery, + }, ttl, max_entries, max_session_id_bytes, @@ -107,20 +210,81 @@ impl AffinityCoordinator { next_revision: AtomicU64::new(1), cancel: CancellationToken::new(), #[cfg(test)] - reaper_started: Arc::new(Notify::new()), - #[cfg(test)] - waiter_observed: Arc::new(Notify::new()), + probe: AffinityCoordinatorProbe::new(), }); Self::spawn_reaper(&inner); + Self::spawn_claim_listener(&inner); Ok(Self { inner }) } + fn spawn_claim_listener(inner: &Arc) { + let Some(mut events) = inner.claims.subscribe() else { + return; + }; + let weak = Arc::downgrade(inner); + let cancel = inner.cancel.clone(); + + tokio::spawn(async move { + loop { + let event = tokio::select! { + _ = cancel.cancelled() => return, + event = events.recv() => event, + }; + let Some(inner) = weak.upgrade() else { + return; + }; + match event { + Ok(ClaimEvent::Delete(key)) => Self::evict_key(&inner, &key), + Ok(ClaimEvent::Reset) | Err(broadcast::error::RecvError::Lagged(_)) => { + Self::clear_entries(&inner); + } + Err(broadcast::error::RecvError::Closed) => { + Self::clear_entries(&inner); + return; + } + } + } + }); + } + + fn evict_key(inner: &AffinityCoordinatorInner, key: &str) { + let Some((_, entry)) = inner.entries.remove(key) else { + return; + }; + if let AffinityEntry::Initializing { notify, .. } = entry { + notify.notify_waiters(); + } + Self::decrement_entry_count(inner, 1); + tracing::debug!(claim_key = key, "evicted session affinity cache entry"); + } + + fn clear_entries(inner: &AffinityCoordinatorInner) { + let mut removed = 0; + inner.entries.retain(|_, entry| { + if let AffinityEntry::Initializing { notify, .. } = entry { + notify.notify_waiters(); + } + removed += 1; + false + }); + Self::decrement_entry_count(inner, removed); + tracing::debug!("cleared session affinity cache after claim watcher reset"); + } + + fn decrement_entry_count(inner: &AffinityCoordinatorInner, removed: usize) { + let _ = inner + .entry_count + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { + Some(count.saturating_sub(removed)) + }); + } + fn spawn_reaper(inner: &Arc) { let weak = Arc::downgrade(inner); let cancel = inner.cancel.clone(); let period = inner.ttl.min(Duration::from_secs(30)); #[cfg(test)] - let reaper_started = inner.reaper_started.clone(); + let reaper_started = inner.probe.reaper_started.clone(); tokio::spawn(async move { #[cfg(test)] reaper_started.notify_one(); @@ -146,54 +310,48 @@ impl AffinityCoordinator { removed += usize::from(!retain); retain }); - inner.entry_count.fetch_sub(removed, Ordering::Relaxed); + Self::decrement_entry_count(&inner, removed); } }); } #[cfg(test)] - pub async fn acquire( + pub(crate) async fn acquire( &self, session_id: &SessionAffinityId, - requested_target: Option, ) -> Result { - self.acquire_inner(session_id, requested_target, None).await + self.acquire_inner(session_id, None).await } - pub async fn acquire_with_context( + pub(crate) async fn acquire_with_context( &self, session_id: &SessionAffinityId, - requested_target: Option, request_context: &dyn AsyncEngineContext, ) -> Result { - self.acquire_inner(session_id, requested_target, Some(request_context)) - .await + self.acquire_inner(session_id, Some(request_context)).await } async fn acquire_inner( &self, session_id: &SessionAffinityId, - requested_target: Option, request_context: Option<&dyn AsyncEngineContext>, ) -> Result { self.validate_session_id(session_id)?; - let session_id = session_id.as_str().to_string(); + let claim_key = self.inner.claims.key(session_id); loop { let now = Instant::now(); - match self.inner.entries.entry(session_id.clone()) { + match self.inner.entries.entry(claim_key.clone()) { Entry::Vacant(entry) => { self.reserve_entry()?; - return Ok(AffinityAcquire::Initialize(entry.insert_initializing( - &self.inner, - session_id, - requested_target, - ))); + return Ok(AffinityAcquire::Initialize( + entry.insert_initializing(&self.inner, claim_key), + )); } Entry::Occupied(mut entry) => match entry.get_mut() { AffinityEntry::Initializing { notify, .. } => { #[cfg(test)] - self.inner.waiter_observed.notify_one(); + self.inner.probe.waiter_observed.notify_one(); let notified = notify.clone().notified_owned(); tokio::pin!(notified); notified.as_mut().enable(); @@ -228,10 +386,9 @@ impl AffinityCoordinator { drop(entry); return Ok(AffinityAcquire::Initialize(AffinityInitialization { coordinator: Arc::downgrade(&self.inner), - session_id, + claim_key, revision, notify, - requested_target, active: true, })); } @@ -241,11 +398,10 @@ impl AffinityCoordinator { active_leases, .. } => { - validate_bound_target(&session_id, *target, requested_target)?; *active_leases += 1; let lease = AffinityLease { coordinator: Arc::downgrade(&self.inner), - session_id, + claim_key, revision: *revision, active: true, }; @@ -262,10 +418,10 @@ impl AffinityCoordinator { pub fn query_target( &self, session_id: &SessionAffinityId, - requested_target: Option, ) -> Result, Error> { self.validate_session_id(session_id)?; - let Some(entry) = self.inner.entries.get(session_id.as_str()) else { + let claim_key = self.inner.claims.key(session_id); + let Some(entry) = self.inner.entries.get(&claim_key) else { return Ok(None); }; let AffinityEntry::Bound { @@ -280,7 +436,6 @@ impl AffinityCoordinator { if *active_leases == 0 && *idle_deadline <= Instant::now() { return Ok(None); } - validate_bound_target(session_id.as_str(), *target, requested_target)?; Ok(Some(*target)) } @@ -289,6 +444,11 @@ impl AffinityCoordinator { self.inner.entry_count.load(Ordering::Relaxed) } + #[cfg(test)] + pub(super) fn claim_key_for_test(&self, session_id: &SessionAffinityId) -> String { + self.inner.claims.key(session_id) + } + #[cfg(test)] pub(super) fn cancellation_token(&self) -> CancellationToken { self.inner.cancel.clone() @@ -296,17 +456,18 @@ impl AffinityCoordinator { #[cfg(test)] pub(super) async fn wait_for_reaper(&self) { - self.inner.reaper_started.notified().await; + self.inner.probe.reaper_started.notified().await; } #[cfg(test)] pub(super) async fn wait_for_initializing_waiter(&self) { - self.inner.waiter_observed.notified().await; + self.inner.probe.waiter_observed.notified().await; } #[cfg(test)] pub(super) fn expire_for_test(&self, session_id: &SessionAffinityId) { - let Some(mut entry) = self.inner.entries.get_mut(session_id.as_str()) else { + let claim_key = self.inner.claims.key(session_id); + let Some(mut entry) = self.inner.entries.get_mut(&claim_key) else { panic!("session affinity entry missing"); }; let AffinityEntry::Bound { @@ -323,7 +484,14 @@ impl AffinityCoordinator { #[cfg(test)] pub(super) fn with_test_limits(max_entries: usize, max_session_id_bytes: usize) -> Self { - Self::new_with_limits(Duration::from_secs(10), max_entries, max_session_id_bytes).unwrap() + Self::new_with_limits( + Duration::from_secs(10), + max_entries, + max_session_id_bytes, + "local".to_string(), + None, + ) + .unwrap() } fn validate_session_id(&self, session_id: &SessionAffinityId) -> Result<(), Error> { @@ -351,8 +519,7 @@ trait VacantEntryExt { fn insert_initializing( self, inner: &Arc, - session_id: String, - requested_target: Option, + claim_key: String, ) -> AffinityInitialization; } @@ -360,8 +527,7 @@ impl<'a> VacantEntryExt for dashmap::mapref::entry::VacantEntry<'a, String, Affi fn insert_initializing( self, inner: &Arc, - session_id: String, - requested_target: Option, + claim_key: String, ) -> AffinityInitialization { let revision = inner.next_revision.fetch_add(1, Ordering::Relaxed); let notify = Arc::new(Notify::new()); @@ -371,16 +537,15 @@ impl<'a> VacantEntryExt for dashmap::mapref::entry::VacantEntry<'a, String, Affi }); AffinityInitialization { coordinator: Arc::downgrade(inner), - session_id, + claim_key, revision, notify, - requested_target, active: true, } } } -pub enum AffinityAcquire { +pub(crate) enum AffinityAcquire { Initialize(AffinityInitialization), Bound { target: AffinityTarget, @@ -389,56 +554,57 @@ pub enum AffinityAcquire { } impl AffinityAcquire { - pub fn target(&self) -> Option { - match self { - Self::Initialize(_) => None, - Self::Bound { target, .. } => Some(*target), - } - } - - pub fn into_stream( - self, - selected_target: AffinityTarget, - stream: ManyOut, - ) -> Result, Error> { + pub(crate) async fn resolve<'a, F>(self, proposed_payload: F) -> Result + where + F: FnOnce() -> ClaimPayloadFuture<'a> + Send, + { match self { - Self::Initialize(initialization) => { - Ok(initialization.commit(selected_target)?.into_stream(stream)) - } - Self::Bound { target, mut lease } => { - if let Err(error) = validate_bound_target("session", target, Some(selected_target)) - { - lease.invalidate(); - return Err(error); - } - Ok(lease.into_stream(stream)) - } - } - } - - pub fn invalidate(self) { - if let Self::Bound { mut lease, .. } = self { - lease.invalidate(); + Self::Initialize(initialization) => initialization.resolve(proposed_payload()).await, + Self::Bound { target, lease } => Ok(ResolvedAffinity { + target, + lease, + created: false, + }), } } } -pub struct AffinityInitialization { +pub(crate) struct AffinityInitialization { coordinator: Weak, - session_id: String, + claim_key: String, revision: u64, notify: Arc, - requested_target: Option, active: bool, } impl AffinityInitialization { - pub fn commit(mut self, target: AffinityTarget) -> Result { - validate_bound_target(&self.session_id, target, self.requested_target)?; + async fn resolve( + self, + mut proposed_payload: ClaimPayloadFuture<'_>, + ) -> Result { + let Some(inner) = self.coordinator.upgrade() else { + return Err(anyhow::anyhow!("session affinity coordinator dropped")); + }; + + let (payload, created) = inner + .claims + .resolve(&self.claim_key, &mut proposed_payload) + .await?; + let target: AffinityTarget = serde_json::from_value(payload) + .map_err(|err| anyhow::anyhow!("invalid session affinity claim payload: {err}"))?; + let lease = self.commit(target)?; + Ok(ResolvedAffinity { + target, + lease, + created, + }) + } + + fn commit(mut self, target: AffinityTarget) -> Result { let Some(inner) = self.coordinator.upgrade() else { return Err(anyhow::anyhow!("session affinity coordinator dropped")); }; - let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { + let Some(mut entry) = inner.entries.get_mut(&self.claim_key) else { return Err(invalid_argument( "session affinity initialization was cancelled", )); @@ -460,7 +626,7 @@ impl AffinityInitialization { self.notify.notify_waiters(); Ok(AffinityLease { coordinator: Arc::downgrade(&inner), - session_id: self.session_id.clone(), + claim_key: self.claim_key.clone(), revision: self.revision, active: true, }) @@ -475,38 +641,27 @@ impl Drop for AffinityInitialization { let Some(inner) = self.coordinator.upgrade() else { return; }; - let removed = inner.entries.remove_if(&self.session_id, |_, entry| { + let removed = inner.entries.remove_if(&self.claim_key, |_, entry| { matches!( entry, AffinityEntry::Initializing { revision, .. } if *revision == self.revision ) }); if removed.is_some() { - inner.entry_count.fetch_sub(1, Ordering::Relaxed); + AffinityCoordinator::decrement_entry_count(&inner, 1); } self.notify.notify_waiters(); } } -pub struct AffinityLease { +pub(crate) struct AffinityLease { coordinator: Weak, - session_id: String, + claim_key: String, revision: u64, active: bool, } impl AffinityLease { - pub fn into_stream(self, stream: ManyOut) -> ManyOut { - let context = stream.context(); - ResponseStream::new( - Box::pin(AffinityTrackedStream { - stream, - lease: Some(self), - }), - context, - ) - } - fn release(&mut self) { if !self.active { return; @@ -515,7 +670,7 @@ impl AffinityLease { let Some(inner) = self.coordinator.upgrade() else { return; }; - let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { + let Some(mut entry) = inner.entries.get_mut(&self.claim_key) else { return; }; let AffinityEntry::Bound { @@ -533,25 +688,6 @@ impl AffinityLease { *active_leases -= 1; *idle_deadline = Instant::now() + inner.ttl; } - - fn invalidate(&mut self) { - if !self.active { - return; - } - self.active = false; - let Some(inner) = self.coordinator.upgrade() else { - return; - }; - let removed = inner.entries.remove_if(&self.session_id, |_, entry| { - matches!( - entry, - AffinityEntry::Bound { revision, .. } if *revision == self.revision - ) - }); - if removed.is_some() { - inner.entry_count.fetch_sub(1, Ordering::Relaxed); - } - } } impl Drop for AffinityLease { @@ -560,9 +696,85 @@ impl Drop for AffinityLease { } } +pub(crate) struct ResolvedAffinity { + target: AffinityTarget, + lease: AffinityLease, + created: bool, +} + +impl ResolvedAffinity { + pub(crate) fn target(&self) -> AffinityTarget { + self.target + } + + pub(crate) fn was_created(&self) -> bool { + self.created + } + + pub(crate) fn into_stream( + self, + stream: ManyOut, + close_on_finish: bool, + ) -> ManyOut { + let context = stream.context(); + let close = close_on_finish.then(|| CloseAction { + coordinator: self.lease.coordinator.clone(), + claims: self + .lease + .coordinator + .upgrade() + .map(|inner| inner.claims.clone()), + claim_key: self.lease.claim_key.clone(), + }); + ResponseStream::new( + Box::pin(AffinityTrackedStream { + stream, + lease: Some(self.lease), + close, + }), + context, + ) + } +} + +struct CloseAction { + coordinator: Weak, + claims: Option, + claim_key: String, +} + +impl CloseAction { + fn run(self) { + if let Some(inner) = self.coordinator.upgrade() { + AffinityCoordinator::evict_key(&inner, &self.claim_key); + } + let Some(claims) = self.claims else { + return; + }; + let claim_key = self.claim_key; + // TODO: Drive backend close to completion before returning stream EOF. This detached + // task keeps early stream drops best-effort and can be cancelled during runtime shutdown. + tokio::spawn(async move { + if let Err(error) = claims.close(&claim_key).await { + tracing::error!(%claim_key, %error, "failed to close session affinity claim"); + } + }); + } +} + struct AffinityTrackedStream { stream: ManyOut, lease: Option, + close: Option, +} + +impl AffinityTrackedStream { + fn finish(&mut self) { + drop(self.lease.take()); + if let Some(close) = self.close.take() { + close.run(); + } + } } impl Stream for AffinityTrackedStream { @@ -571,7 +783,7 @@ impl Stream for AffinityTrackedStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { match Pin::new(&mut self.stream).poll_next(cx) { Poll::Ready(None) => { - drop(self.lease.take()); + self.finish(); Poll::Ready(None) } Poll::Ready(Some(item)) => Poll::Ready(Some(item)), @@ -580,6 +792,12 @@ impl Stream for AffinityTrackedStream { } } +impl Drop for AffinityTrackedStream { + fn drop(&mut self) { + self.finish(); + } +} + pub fn affinity_id( request: &dynamo_runtime::pipeline::SingleIn, ) -> Result>, Error> { @@ -588,6 +806,13 @@ pub fn affinity_id( .map_err(|message| invalid_argument(format!("invalid session affinity context: {message}"))) } +pub fn session_final(request: &PreprocessedRequest) -> bool { + request + .agent_context + .as_ref() + .is_some_and(|context| context.session_final == Some(true)) +} + pub fn explicit_target( request: &PreprocessedRequest, phase: RequestPhase, @@ -617,31 +842,6 @@ pub fn explicit_target( Ok(worker_id.map(|worker_id| AffinityTarget { worker_id, dp_rank })) } -fn validate_bound_target( - session_id: &str, - bound: AffinityTarget, - requested: Option, -) -> Result<(), Error> { - let Some(requested) = requested else { - return Ok(()); - }; - if bound.worker_id != requested.worker_id { - return Err(invalid_argument(format!( - "session {session_id} is bound to worker {}, not {}", - bound.worker_id, requested.worker_id - ))); - } - match (bound.dp_rank, requested.dp_rank) { - (Some(bound), Some(requested)) if bound != requested => Err(invalid_argument(format!( - "session {session_id} is bound to DP rank {bound}, not {requested}" - ))), - (None, Some(requested)) => Err(invalid_argument(format!( - "session {session_id} has worker-only affinity and cannot add DP rank {requested}" - ))), - _ => Ok(()), - } -} - pub(crate) fn invalid_argument(message: impl Into) -> Error { DynamoError::builder() .error_type(ErrorType::InvalidArgument) diff --git a/lib/llm/src/session_affinity/mod.rs b/lib/llm/src/session_affinity/mod.rs index 08e6451d5533..1f841dab0a6b 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -4,11 +4,8 @@ mod coordinator; mod push_router; -pub(crate) use coordinator::affinity_id; -pub use coordinator::{ - AffinityAcquire, AffinityCoordinator, AffinityInitialization, AffinityLease, AffinityTarget, - explicit_target, -}; +pub use coordinator::{AffinityCoordinator, AffinityTarget, explicit_target}; +pub(crate) use coordinator::{ResolvedAffinity, affinity_id, session_final}; pub use push_router::SessionAffinityPushRouter; pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; diff --git a/lib/llm/src/session_affinity/push_router.rs b/lib/llm/src/session_affinity/push_router.rs index 48f4f6f0b8a5..8690c6df652c 100644 --- a/lib/llm/src/session_affinity/push_router.rs +++ b/lib/llm/src/session_affinity/push_router.rs @@ -3,15 +3,19 @@ use std::time::Duration; -use dynamo_runtime::pipeline::{ - AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Error, ManyOut, PushRouter, - SingleIn, async_trait as pipeline_async_trait, +use dynamo_runtime::{ + discovery::ClaimPayloadFuture, + pipeline::{ + AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Error, ManyOut, PushRouter, + SingleIn, async_trait as pipeline_async_trait, + }, + traits::DistributedRuntimeProvider, }; use super::{ - AffinityCoordinator, AffinityTarget, LlmResponse, + AffinityCoordinator, AffinityTarget, LlmResponse, ResolvedAffinity, coordinator::{affinity_id, invalid_argument}, - explicit_target, + explicit_target, session_final, }; use crate::{ preprocessor::PreprocessedRequest, @@ -30,9 +34,18 @@ impl SessionAffinityPushRouter { ttl: Option, direct: bool, ) -> Result { + let affinity = ttl + .map(|ttl| { + AffinityCoordinator::new_distributed( + ttl, + inner.client.endpoint.id().to_string(), + inner.client.endpoint.drt().discovery(), + ) + }) + .transpose()?; Ok(Self { inner, - affinity: ttl.map(AffinityCoordinator::new).transpose()?, + affinity, direct, }) } @@ -57,53 +70,104 @@ impl SessionAffinityPushRouter { tracker.record_worker(target.worker_id, target.dp_rank, worker_type); } - fn direct_target( - &self, - explicit: Option, - phase: RequestPhase, - ) -> Result, Error> { - if !self.direct { - return Ok(explicit); - } - explicit.map(Some).ok_or_else(|| { - invalid_argument(format!( - "worker ID required for {phase} request in Direct routing mode" - )) - }) - } - pub fn peek_next_worker(&self) -> Option { self.inner.peek_next_worker() } - async fn acquire_routable( + async fn resolve_affinity( &self, session_id: &crate::protocols::common::extensions::SessionAffinityId, - explicit: Option, + phase: RequestPhase, + request: &PreprocessedRequest, request_context: &dyn AsyncEngineContext, - ) -> Result { + ) -> Result<(ResolvedAffinity, bool), Error> { let affinity = self .affinity .as_ref() .expect("affinity acquisition requires an enabled coordinator"); let operation = affinity - .acquire_with_context(session_id, explicit, request_context) + .acquire_with_context(session_id, request_context) + .await?; + let resolved = operation + .resolve(|| -> ClaimPayloadFuture<'_> { + Box::pin(async move { + let target = explicit_target(request, phase)? + .or_else(|| { + self.inner + .peek_worker_for_request(request) + .map(|worker_id| AffinityTarget { + worker_id, + dp_rank: None, + }) + }) + .ok_or_else(|| { + if self.direct { + invalid_argument( + "worker ID required to create Direct session affinity", + ) + } else { + anyhow::anyhow!("no worker available for session affinity") + } + })?; + Ok(serde_json::to_value(target)?) + }) + }) .await?; - let Some(target) = operation.target() else { - return Ok(operation); + let proposal_was_explicit = if resolved.was_created() { + explicit_target(request, phase)?.is_some() + } else { + false }; - if self - .inner - .client - .instance_ids_avail() - .contains(&target.worker_id) - { - return Ok(operation); - } + Ok((resolved, proposal_was_explicit)) + } - operation.invalidate(); - affinity - .acquire_with_context(session_id, explicit, request_context) + /// Adapts the generic worker-only router while keeping a known rank attached + /// to its worker through preparation and exact dispatch. + async fn select_and_dispatch_exact_target( + &self, + request: SingleIn, + pinned_target: Option, + prepare: F, + ) -> Result<(M, ManyOut), Error> + where + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, + { + self.inner + .select_and_dispatch_exact( + request, + pinned_target.map(|target| target.worker_id), + move |request, worker_id| { + let target = pinned_target.unwrap_or(AffinityTarget { + worker_id, + dp_rank: None, + }); + debug_assert_eq!(target.worker_id, worker_id); + prepare(request, target) + }, + ) + .await + } + + async fn book_and_dispatch_exact_target( + &self, + request: SingleIn, + target: AffinityTarget, + advance_round_robin: bool, + prepare: F, + ) -> Result<(M, ManyOut), Error> + where + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, + { + self.inner + .book_and_dispatch_exact( + request, + target.worker_id, + advance_round_robin, + move |request, worker_id| { + debug_assert_eq!(target.worker_id, worker_id); + prepare(request, target) + }, + ) .await } @@ -113,7 +177,7 @@ impl SessionAffinityPushRouter { prepare: F, ) -> Result<(M, ManyOut), Error> where - F: FnOnce(&mut PreprocessedRequest, u64, Option) -> Result, + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { let session_id = if self.affinity.is_some() { affinity_id(&request)? @@ -121,88 +185,65 @@ impl SessionAffinityPushRouter { None }; if !self.direct && session_id.is_none() { - let pinned_worker = phase_worker_id(&request, RequestPhase::Prefill); + let explicit = explicit_target(&request, RequestPhase::Prefill)?; return self - .inner - .select_and_dispatch_exact(request, pinned_worker, move |request, worker_id| { - prepare(request, worker_id, None) - }) + .select_and_dispatch_exact_target(request, explicit, prepare) .await; } - let explicit = self.direct_target( - explicit_target(&request, RequestPhase::Prefill)?, - RequestPhase::Prefill, - )?; let Some(session_id) = session_id else { - let Some(pinned_worker) = explicit else { + let explicit = explicit_target(&request, RequestPhase::Prefill)?; + let Some(target) = explicit else { return Err(invalid_argument( - "Direct routing requires an explicit prefill target", + "worker ID required for prefill request in Direct routing mode", )); }; return self - .inner - .select_and_dispatch_exact( - request, - Some(pinned_worker.worker_id), - move |request, worker_id| prepare(request, worker_id, None), - ) + .select_and_dispatch_exact_target(request, Some(target), prepare) .await; }; let is_query_only = request.get_annotation_value("query_instance_id").is_some(); if is_query_only { - let selected = self + let bound = self .affinity .as_ref() .expect("affinity query requires an enabled coordinator") - .query_target(&session_id, explicit)? - .or(explicit); - let rank = selected.and_then(|target| target.dp_rank); + .query_target(&session_id)?; + let selected = match bound { + Some(target) => Some(target), + None => explicit_target(&request, RequestPhase::Prefill)?, + }; return self - .inner - .select_and_dispatch_exact( - request, - selected.map(|target| target.worker_id), - move |request, worker_id| { - let target = AffinityTarget { - worker_id, - dp_rank: rank, - }; - Self::record_target(request, target); - prepare(request, worker_id, rank) - }, - ) + .select_and_dispatch_exact_target(request, selected, move |request, target| { + Self::record_target(request, target); + prepare(request, target) + }) .await; } + let close_on_finish = session_final(request.content()); let request_context = request.context(); - let operation = self - .acquire_routable(&session_id, explicit, request_context.as_ref()) + let (resolved, proposal_was_explicit) = self + .resolve_affinity( + &session_id, + RequestPhase::Prefill, + request.content(), + request_context.as_ref(), + ) .await?; - let selected = operation.target().or(explicit); - let rank = selected.and_then(|target| target.dp_rank); - let dispatch = self - .inner - .select_and_dispatch_exact( + let target = resolved.target(); + let advance_round_robin = resolved.was_created() && !proposal_was_explicit; + let (metadata, stream) = self + .book_and_dispatch_exact_target( request, - selected.map(|target| target.worker_id), - move |request, worker_id| { - let target = AffinityTarget { - worker_id, - dp_rank: rank, - }; + target, + advance_round_robin, + move |request, target| { Self::record_target(request, target); - Ok((prepare(request, worker_id, rank)?, target)) + prepare(request, target) }, ) - .await; - let ((metadata, target), stream) = match dispatch { - Ok(result) => result, - Err(error) => { - operation.invalidate(); - return Err(error); - } - }; - Ok((metadata, operation.into_stream(target, stream)?)) + .await?; + Ok((metadata, resolved.into_stream(stream, close_on_finish))) } } @@ -223,11 +264,11 @@ impl AsyncEngine, ManyOut, Error> if !self.direct && session_id.is_none() { return self.inner.generate(request).await; } - let explicit = self.direct_target(explicit_target(&request, phase)?, phase)?; let Some(session_id) = session_id else { + let explicit = explicit_target(&request, phase)?; let Some(target) = explicit else { return Err(invalid_argument(format!( - "Direct routing requires an explicit {phase} target" + "worker ID required for {phase} request in Direct routing mode" ))); }; return self.inner.direct(request, target.worker_id).await; @@ -235,77 +276,54 @@ impl AsyncEngine, ManyOut, Error> let is_query_only = request.get_annotation_value("query_instance_id").is_some(); if is_query_only { - let target = self + let bound = self .affinity .as_ref() .expect("affinity query requires an enabled coordinator") - .query_target(&session_id, explicit)? - .or(explicit); - let rank = target.and_then(|target| target.dp_rank); + .query_target(&session_id)?; + let target = match bound { + Some(target) => Some(target), + None => explicit_target(&request, phase)?, + }; let (_, stream) = self - .inner - .select_and_dispatch_exact( - request, - target.map(|target| target.worker_id), - move |request, worker_id| { - if rank.is_some() { - request.routing_mut().dp_rank = rank; - } - Self::record_target( - request, - AffinityTarget { - worker_id, - dp_rank: rank, - }, - ); - Ok(()) - }, - ) + .select_and_dispatch_exact_target(request, target, move |request, target| { + if target.dp_rank.is_some() { + request.routing_mut().dp_rank = target.dp_rank; + } + Self::record_target(request, target); + Ok(()) + }) .await?; return Ok(stream); } + let close_on_finish = session_final(request.content()); let request_context = request.context(); - let operation = self - .acquire_routable(&session_id, explicit, request_context.as_ref()) + let (resolved, proposal_was_explicit) = self + .resolve_affinity( + &session_id, + phase, + request.content(), + request_context.as_ref(), + ) .await?; - let selected = operation.target().or(explicit); - let rank = selected.and_then(|target| target.dp_rank); - let dispatch = self - .inner - .select_and_dispatch_exact( + let target = resolved.target(); + let advance_round_robin = resolved.was_created() && !proposal_was_explicit; + let (_, stream) = self + .book_and_dispatch_exact_target( request, - selected.map(|target| target.worker_id), - move |request, worker_id| { - if rank.is_some() { - request.routing_mut().dp_rank = rank; + target, + advance_round_robin, + move |request, target| { + if target.dp_rank.is_some() { + request.routing_mut().dp_rank = target.dp_rank; } - let target = AffinityTarget { - worker_id, - dp_rank: rank, - }; Self::record_target(request, target); - Ok(target) + Ok(()) }, ) - .await; - let (target, stream) = match dispatch { - Ok(result) => result, - Err(error) => { - operation.invalidate(); - return Err(error); - } - }; - operation.into_stream(target, stream) - } -} - -fn phase_worker_id(request: &PreprocessedRequest, phase: RequestPhase) -> Option { - let routing = request.routing.as_ref()?; - match phase { - RequestPhase::Prefill => routing.prefill_worker_id.or(routing.backend_instance_id), - RequestPhase::Decode => routing.decode_worker_id.or(routing.backend_instance_id), - RequestPhase::Aggregated => routing.decode_worker_id.or(routing.backend_instance_id), + .await?; + Ok(resolved.into_stream(stream, close_on_finish)) } } @@ -322,7 +340,6 @@ mod tests { extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId}, preprocessor::RoutingHints, }; - use crate::session_affinity::AffinityAcquire; fn request(worker_id: Option, query_only: bool) -> PreprocessedRequest { PreprocessedRequest::builder() @@ -388,51 +405,34 @@ mod tests { } #[tokio::test] - async fn session_affinity_simple_modes_rollback_failed_initialization() { + async fn session_affinity_failed_dispatch_preserves_created_claim() { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) .await .unwrap(); - let namespace = distributed + let client = distributed .namespace("session_affinity_adapters".to_string()) + .unwrap() + .component("workers".to_string()) + .unwrap() + .endpoint("direct") + .client() + .await .unwrap(); - let component = namespace.component("workers".to_string()).unwrap(); - - for (index, mode) in [ - RouterMode::Random, - RouterMode::RoundRobin, - RouterMode::PowerOfTwoChoices, - RouterMode::LeastLoaded, - RouterMode::DeviceAwareWeighted, - RouterMode::Direct, - ] - .into_iter() - .enumerate() - { - let endpoint = component.endpoint(format!("mode-{index}")); - let client = endpoint.client().await.unwrap(); - let inner = PushRouter::from_client(client, mode).await.unwrap(); - let router = SessionAffinityPushRouter::new( - inner, - Some(Duration::from_secs(10)), - mode.is_direct_routing(), - ) + let inner = PushRouter::from_client(client, RouterMode::Direct) + .await .unwrap(); - let worker_id = mode.is_direct_routing().then_some(99); - - assert!( - router - .generate(affinity_request(worker_id, false)) - .await - .is_err() - ); - assert_eq!( - affinity(&router).entry_count(), - 0, - "failed {mode:?} dispatch must release initialization" - ); - } + let router = + SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), true).unwrap(); + + assert!( + router + .generate(affinity_request(Some(99), false)) + .await + .is_err() + ); + assert_eq!(affinity(&router).entry_count(), 1); runtime.shutdown(); } @@ -463,7 +463,7 @@ mod tests { assert_eq!(affinity(&router).entry_count(), 0); assert!( router - .select_and_dispatch_prefill(affinity_request(None, true), |_, _, _| Ok(())) + .select_and_dispatch_prefill(affinity_request(None, true), |_, _| Ok(())) .await .is_err() ); @@ -496,7 +496,7 @@ mod tests { .contains("worker ID required for aggregated request in Direct routing mode") ); let error = router - .select_and_dispatch_prefill(Context::new(request(None, false)), |_, _, _| Ok(())) + .select_and_dispatch_prefill(Context::new(request(None, false)), |_, _| Ok(())) .await .unwrap_err(); assert!( @@ -506,18 +506,55 @@ mod tests { ); assert_eq!(affinity(&router).entry_count(), 0); - let mut decode_only = request(None, false); - decode_only.routing_mut().decode_worker_id = Some(99); - assert_eq!( - phase_worker_id(&decode_only, RequestPhase::Aggregated), - Some(99) - ); + runtime.shutdown(); + } + + #[tokio::test] + async fn prefill_preparation_receives_explicit_rank_zero() { + let runtime = Runtime::from_current().unwrap(); + let distributed = + DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) + .await + .unwrap(); + let endpoint = distributed + .namespace("session_affinity_prefill_target".to_string()) + .unwrap() + .component("workers".to_string()) + .unwrap() + .endpoint("prefill".to_string()); + let client = endpoint.client().await.unwrap(); + endpoint.register_endpoint_instance().await.unwrap(); + let worker_id = client.wait_for_instances().await.unwrap()[0].id(); + let expected = AffinityTarget { + worker_id, + dp_rank: Some(0), + }; + + for (mode, direct) in [(RouterMode::Direct, true), (RouterMode::RoundRobin, false)] { + let inner = PushRouter::from_client(client.clone(), mode).await.unwrap(); + let router = SessionAffinityPushRouter::new(inner, None, direct).unwrap(); + let mut content = request(None, false); + content.routing_mut().prefill_worker_id = Some(worker_id); + content.routing_mut().prefill_dp_rank = Some(0); + let mut observed = None; + + let error = router + .select_and_dispatch_prefill(Context::new(content), |_, target| { + observed = Some(target); + Err::<(), _>(anyhow::anyhow!("stop before dispatch")) + }) + .await + .unwrap_err(); + + assert!(error.to_string().contains("stop before dispatch")); + assert_eq!(observed, Some(expected)); + } runtime.shutdown(); } #[tokio::test] - async fn session_affinity_unavailable_target_is_invalidated() { + async fn session_affinity_binding_wins_before_invalid_explicit_proposal() { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) @@ -537,29 +574,29 @@ mod tests { let router = SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), false).unwrap(); let session_id = SessionAffinityId::new("adapter-session"); - let AffinityAcquire::Initialize(initializer) = - affinity(&router).acquire(&session_id, None).await.unwrap() - else { - panic!("first request must initialize"); + let unavailable_target = AffinityTarget { + worker_id: 99, + dp_rank: None, }; - drop( - initializer - .commit(AffinityTarget { - worker_id: 99, - dp_rank: None, - }) - .unwrap(), - ); + let resolved = affinity(&router) + .acquire(&session_id) + .await + .unwrap() + .resolve(|| Box::pin(async move { Ok(serde_json::to_value(unavailable_target)?) })) + .await + .unwrap(); + drop(resolved); - assert!( - router - .generate(affinity_request(None, false)) - .await - .is_err() - ); + let mut request = affinity_request(None, false); + request.routing_mut().dp_rank = Some(0); + let error = router.generate(request).await.unwrap_err(); + assert!(!error.to_string().contains("DP rank requires")); assert_eq!( - affinity(&router).query_target(&session_id, None).unwrap(), - None + affinity(&router).query_target(&session_id).unwrap(), + Some(AffinityTarget { + worker_id: 99, + dp_rank: None, + }) ); runtime.shutdown(); diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index 044eea963356..5b9784f9e085 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -1,19 +1,32 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::{sync::Arc, time::Duration}; +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, + time::Duration, +}; +use async_trait::async_trait; use dynamo_runtime::{ + discovery::{ + ClaimCloseOutcome, ClaimEvent, ClaimOutcome, ClaimPayload, ClaimPayloadFuture, Discovery, + DiscoveryInstance, DiscoveryQuery, DiscoverySpec, DiscoveryStream, + }, engine::AsyncEngineContext, error::ErrorType, pipeline::{Context, ResponseStream, context::Controller}, protocols::maybe_error::MaybeError, }; use futures::{StreamExt, stream}; +use tokio::sync::broadcast; +use tokio_util::sync::CancellationToken; -use super::{ - AffinityAcquire, AffinityCoordinator, AffinityTarget, LlmResponse, affinity_id, explicit_target, -}; +use super::coordinator::AffinityAcquire; +use super::{AffinityCoordinator, AffinityTarget, LlmResponse, affinity_id, explicit_target}; use crate::{ preprocessor::PreprocessedRequest, protocols::common::{ @@ -37,6 +50,138 @@ fn coordinator() -> AffinityCoordinator { AffinityCoordinator::new(Duration::from_secs(10)).unwrap() } +struct ClaimTestDiscovery { + claims: Mutex>, + create_calls: AtomicUsize, + close_calls: AtomicUsize, + events: Mutex>>, +} + +impl ClaimTestDiscovery { + fn new(event_capacity: usize) -> Arc { + let (events, _) = broadcast::channel(event_capacity); + Arc::new(Self { + claims: Mutex::new(HashMap::new()), + create_calls: AtomicUsize::new(0), + close_calls: AtomicUsize::new(0), + events: Mutex::new(Some(events)), + }) + } + + fn emit(&self, event: ClaimEvent) { + if let Some(events) = self.events.lock().unwrap().as_ref() { + let _ = events.send(event); + } + } + + fn disconnect(&self) { + self.events.lock().unwrap().take(); + } +} + +#[async_trait] +impl Discovery for ClaimTestDiscovery { + fn instance_id(&self) -> u64 { + 1 + } + + async fn register_internal(&self, _spec: DiscoverySpec) -> anyhow::Result { + anyhow::bail!("not used by claim tests") + } + + async fn unregister(&self, _instance: DiscoveryInstance) -> anyhow::Result<()> { + Ok(()) + } + + async fn list(&self, _query: DiscoveryQuery) -> anyhow::Result> { + Ok(Vec::new()) + } + + async fn list_and_watch( + &self, + _query: DiscoveryQuery, + _cancel_token: Option, + ) -> anyhow::Result { + Ok(Box::pin(stream::pending())) + } + + async fn create_or_get_claim( + &self, + key: &str, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> anyhow::Result { + self.create_calls.fetch_add(1, Ordering::Relaxed); + if let Some(payload) = self.claims.lock().unwrap().get(key).cloned() { + return Ok(ClaimOutcome::Existing(payload)); + } + + let proposed = proposed_payload.as_mut().await?; + let mut claims = self.claims.lock().unwrap(); + if let Some(payload) = claims.get(key).cloned() { + return Ok(ClaimOutcome::Existing(payload)); + } + claims.insert(key.to_string(), proposed.clone()); + Ok(ClaimOutcome::Created(proposed)) + } + + async fn close_claim(&self, key: &str) -> anyhow::Result { + self.close_calls.fetch_add(1, Ordering::Relaxed); + if self.claims.lock().unwrap().remove(key).is_some() { + self.emit(ClaimEvent::Delete(key.to_string())); + } + Ok(ClaimCloseOutcome::Closed) + } + + fn subscribe_claim_events(&self) -> Option> { + self.events + .lock() + .unwrap() + .as_ref() + .map(broadcast::Sender::subscribe) + } +} + +fn distributed_coordinator(discovery: Arc) -> AffinityCoordinator { + AffinityCoordinator::new_distributed( + Duration::from_secs(10), + "ns/component/endpoint".to_string(), + discovery, + ) + .unwrap() +} + +fn target_payload(target: AffinityTarget) -> ClaimPayloadFuture<'static> { + Box::pin(async move { Ok(serde_json::to_value(target)?) }) +} + +async fn resolve_local( + coordinator: &AffinityCoordinator, + session_id: &SessionAffinityId, + selected: AffinityTarget, +) -> super::ResolvedAffinity { + coordinator + .acquire(session_id) + .await + .unwrap() + .resolve(|| target_payload(selected)) + .await + .unwrap() +} + +async fn wait_for_cached_target( + coordinator: &AffinityCoordinator, + session_id: &SessionAffinityId, + expected: Option, +) { + for _ in 0..100 { + if coordinator.query_target(session_id).unwrap() == expected { + return; + } + tokio::task::yield_now().await; + } + assert_eq!(coordinator.query_target(session_id).unwrap(), expected); +} + fn response_stream(items: usize) -> dynamo_runtime::pipeline::ManyOut { let items = (0..items).map(|_| Annotated::from_data(LLMEngineOutput::default())); ResponseStream::new( @@ -61,11 +206,11 @@ fn cancelled_response_stream() -> dynamo_runtime::pipeline::ManyOut async fn assert_binding_expires_after_refreshed_ttl(coordinator: &AffinityCoordinator) { tokio::time::advance(Duration::from_secs(9)).await; assert_eq!( - coordinator.query_target(&session_id(), None).unwrap(), + coordinator.query_target(&session_id()).unwrap(), Some(target(7, Some(0))) ); tokio::time::advance(Duration::from_secs(2)).await; - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); } fn request_with_routing(routing: RoutingHints) -> PreprocessedRequest { @@ -136,40 +281,208 @@ fn session_affinity_context_type_errors_are_preserved() { #[tokio::test(start_paused = true)] async fn session_affinity_initialization_is_atomic() { let coordinator = coordinator(); - let first = coordinator.acquire(&session_id(), None).await.unwrap(); - let AffinityAcquire::Initialize(first) = first else { - panic!("first request must initialize"); - }; + let first = coordinator.acquire(&session_id()).await.unwrap(); + assert!(matches!(&first, AffinityAcquire::Initialize(_))); let waiter_coordinator = coordinator.clone(); - let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id(), None).await }); + let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id()).await }); coordinator.wait_for_initializing_waiter().await; assert!(!waiter.is_finished()); - let first_lease = first.commit(target(7, Some(0))).unwrap(); - let second = waiter.await.unwrap().unwrap(); - let AffinityAcquire::Bound { - target: second_target, - lease: second_lease, - } = second - else { - panic!("waiter must acquire the committed binding"); - }; - assert_eq!(second_target, target(7, Some(0))); - drop(first_lease); - drop(second_lease); + let first = first + .resolve(|| target_payload(target(7, Some(0)))) + .await + .unwrap(); + let second = waiter + .await + .unwrap() + .unwrap() + .resolve(|| target_payload(target(8, Some(0)))) + .await + .unwrap(); + assert_eq!(second.target(), target(7, Some(0))); + drop(first); + drop(second); +} + +#[tokio::test] +async fn distributed_existing_winner_skips_proposal_and_cache_hits_skip_discovery() { + let discovery = ClaimTestDiscovery::new(16); + let first = distributed_coordinator(discovery.clone()); + let second = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + + let created = first + .acquire(&session_id) + .await + .unwrap() + .resolve(|| target_payload(target(7, Some(0)))) + .await + .unwrap(); + assert_eq!(created.target(), target(7, Some(0))); + assert!(created.was_created()); + drop(created); + + let proposal_polled = Arc::new(AtomicBool::new(false)); + let polled = proposal_polled.clone(); + let winner = second + .acquire(&session_id) + .await + .unwrap() + .resolve(|| { + Box::pin(async move { + polled.store(true, Ordering::Relaxed); + Ok(serde_json::to_value(target(8, Some(0)))?) + }) + }) + .await + .unwrap(); + assert_eq!(winner.target(), target(7, Some(0))); + assert!(!winner.was_created()); + assert!(!proposal_polled.load(Ordering::Relaxed)); + drop(winner); + assert_eq!(discovery.create_calls.load(Ordering::Relaxed), 2); + + let cached_proposal_constructed = Arc::new(AtomicBool::new(false)); + let constructed = cached_proposal_constructed.clone(); + let cached = second + .acquire(&session_id) + .await + .unwrap() + .resolve(|| { + constructed.store(true, Ordering::Relaxed); + target_payload(target(9, Some(0))) + }) + .await + .unwrap(); + assert_eq!(cached.target(), target(7, Some(0))); + assert!(!cached_proposal_constructed.load(Ordering::Relaxed)); + assert_eq!(discovery.create_calls.load(Ordering::Relaxed), 2); +} + +#[tokio::test] +async fn distributed_delete_evicts_one_entry_and_duplicate_is_harmless() { + let discovery = ClaimTestDiscovery::new(16); + let first = distributed_coordinator(discovery.clone()); + let second = distributed_coordinator(discovery.clone()); + let first_session = SessionAffinityId::new("first-session"); + let second_session = SessionAffinityId::new("second-session"); + + for (coordinator, session, worker_id) in + [(&first, &first_session, 7), (&second, &second_session, 8)] + { + drop( + coordinator + .acquire(session) + .await + .unwrap() + .resolve(|| target_payload(target(worker_id, Some(0)))) + .await + .unwrap(), + ); + } + + let first_key = first.claim_key_for_test(&first_session); + discovery.emit(ClaimEvent::Delete(first_key.clone())); + wait_for_cached_target(&first, &first_session, None).await; + assert_eq!( + second.query_target(&second_session).unwrap(), + Some(target(8, Some(0))) + ); + discovery.emit(ClaimEvent::Delete(first_key)); + wait_for_cached_target(&first, &first_session, None).await; +} + +#[tokio::test] +async fn distributed_subscriber_lag_clears_all_entries() { + let discovery = ClaimTestDiscovery::new(1); + let coordinator = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + drop( + coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(|| target_payload(target(8, Some(0)))) + .await + .unwrap(), + ); + + for index in 0..16 { + discovery.emit(ClaimEvent::Delete(format!("unrelated/{index}"))); + } + wait_for_cached_target(&coordinator, &session_id, None).await; +} + +#[tokio::test] +async fn distributed_reset_and_disconnect_clear_entries() { + let discovery = ClaimTestDiscovery::new(16); + let coordinator = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + + drop( + coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(|| target_payload(target(8, Some(0)))) + .await + .unwrap(), + ); + discovery.emit(ClaimEvent::Reset); + wait_for_cached_target(&coordinator, &session_id, None).await; + + let resolved = coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(|| target_payload(target(99, Some(0)))) + .await + .unwrap(); + assert_eq!(resolved.target(), target(8, Some(0))); + assert!(!resolved.was_created()); + drop(resolved); + discovery.disconnect(); + wait_for_cached_target(&coordinator, &session_id, None).await; +} + +#[tokio::test] +async fn terminal_close_evicts_synchronously() { + let discovery = ClaimTestDiscovery::new(16); + let coordinator = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + let resolved = coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(|| target_payload(target(7, Some(0)))) + .await + .unwrap(); + let key = coordinator.claim_key_for_test(&session_id); + let mut stream = resolved.into_stream(response_stream(0), true); + assert!(stream.next().await.is_none()); + assert_eq!(coordinator.query_target(&session_id).unwrap(), None); + + for _ in 0..100 { + if discovery.close_calls.load(Ordering::Relaxed) == 1 { + break; + } + tokio::task::yield_now().await; + } + assert_eq!(discovery.close_calls.load(Ordering::Relaxed), 1); + assert!(!discovery.claims.lock().unwrap().contains_key(&key)); } #[tokio::test(start_paused = true)] async fn session_affinity_initializer_cancellation_wakes_waiter() { let coordinator = coordinator(); - let first = coordinator.acquire(&session_id(), None).await.unwrap(); + let first = coordinator.acquire(&session_id()).await.unwrap(); let AffinityAcquire::Initialize(first) = first else { panic!("first request must initialize"); }; let waiter_coordinator = coordinator.clone(); - let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id(), None).await }); + let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id()).await }); coordinator.wait_for_initializing_waiter().await; drop(first); @@ -178,15 +491,43 @@ async fn session_affinity_initializer_cancellation_wakes_waiter() { drop(next); assert_eq!(coordinator.entry_count(), 0); assert!(matches!( - coordinator.acquire(&session_id(), None).await.unwrap(), + coordinator.acquire(&session_id()).await.unwrap(), AffinityAcquire::Initialize(_) )); } +#[tokio::test] +async fn distributed_reset_wakes_initializing_waiter_and_preserves_entry_count() { + let discovery = ClaimTestDiscovery::new(16); + let coordinator = distributed_coordinator(discovery.clone()); + let first = coordinator.acquire(&session_id()).await.unwrap(); + let AffinityAcquire::Initialize(first) = first else { + panic!("first request must initialize"); + }; + + let waiter_coordinator = coordinator.clone(); + let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id()).await }); + coordinator.wait_for_initializing_waiter().await; + discovery.emit(ClaimEvent::Reset); + + let next = tokio::time::timeout(Duration::from_secs(1), waiter) + .await + .expect("reset did not wake initializing waiter") + .unwrap() + .unwrap(); + assert!(matches!(&next, AffinityAcquire::Initialize(_))); + assert_eq!(coordinator.entry_count(), 1); + + drop(first); + assert_eq!(coordinator.entry_count(), 1); + drop(next); + assert_eq!(coordinator.entry_count(), 0); +} + #[tokio::test(start_paused = true)] async fn session_affinity_wait_stops_when_request_is_cancelled() { let coordinator = coordinator(); - let first = coordinator.acquire(&session_id(), None).await.unwrap(); + let first = coordinator.acquire(&session_id()).await.unwrap(); let AffinityAcquire::Initialize(first) = first else { panic!("first request must initialize"); }; @@ -196,7 +537,7 @@ async fn session_affinity_wait_stops_when_request_is_cancelled() { let waiter_coordinator = coordinator.clone(); let waiter = tokio::spawn(async move { waiter_coordinator - .acquire_with_context(&session_id(), None, waiter_context.as_ref()) + .acquire_with_context(&session_id(), waiter_context.as_ref()) .await }); coordinator.wait_for_initializing_waiter().await; @@ -214,70 +555,28 @@ async fn session_affinity_wait_stops_when_request_is_cancelled() { } #[tokio::test(start_paused = true)] -async fn session_affinity_validates_worker_and_rank_contract() { +async fn session_affinity_existing_binding_overrides_explicit_proposals() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = coordinator - .acquire(&session_id(), Some(target(7, None))) - .await - .unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, None)).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, None)).await); - assert!( - coordinator - .acquire(&session_id(), Some(target(8, None))) + for proposal in [target(8, None), target(7, Some(0)), target(7, None)] { + let resolved = coordinator + .acquire(&session_id()) .await - .is_err() - ); - assert!( - coordinator - .acquire(&session_id(), Some(target(7, Some(0)))) + .unwrap() + .resolve(|| target_payload(proposal)) .await - .is_err() - ); - assert!( - coordinator - .acquire(&session_id(), Some(target(7, None))) - .await - .is_ok() - ); -} - -#[tokio::test(start_paused = true)] -async fn session_affinity_failed_bound_operation_invalidates_binding() { - let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); - - let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - assert_eq!(operation.target(), Some(target(7, Some(0)))); - operation.invalidate(); - - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); - assert_eq!(coordinator.entry_count(), 0); - assert!(matches!( - coordinator.acquire(&session_id(), None).await.unwrap(), - AffinityAcquire::Initialize(_) - )); + .unwrap(); + assert_eq!(resolved.target(), target(7, None)); + } } #[tokio::test(start_paused = true)] async fn session_affinity_stream_drop_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(1)); + let mut stream = resolved.into_stream(response_stream(1), false); assert!(stream.next().await.is_some()); drop(stream); @@ -287,14 +586,9 @@ async fn session_affinity_stream_drop_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_empty_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(0)); + let mut stream = resolved.into_stream(response_stream(0), false); assert!(stream.next().await.is_none()); assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -303,23 +597,12 @@ async fn session_affinity_empty_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); tokio::time::advance(Duration::from_secs(9)).await; - let AffinityAcquire::Bound { - target: bound_target, - lease, - } = coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("continuation must acquire the existing binding"); - }; - assert_eq!(bound_target, target(7, Some(0))); - let mut stream = lease.into_stream(cancelled_response_stream()); + let resolved = resolve_local(&coordinator, &session_id(), target(8, Some(0))).await; + assert_eq!(resolved.target(), target(7, Some(0))); + let mut stream = resolved.into_stream(cancelled_response_stream(), false); assert!(stream.next().await.is_none()); assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -328,10 +611,8 @@ async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_committed_binding_survives_cancelled_stream_until_ttl() { let coordinator = coordinator(); - let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - let mut stream = operation - .into_stream(target(7, Some(0)), cancelled_response_stream()) - .unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; + let mut stream = resolved.into_stream(cancelled_response_stream(), false); tokio::time::advance(Duration::from_secs(9)).await; assert!(stream.next().await.is_none()); @@ -341,14 +622,9 @@ async fn session_affinity_committed_binding_survives_cancelled_stream_until_ttl( #[tokio::test(start_paused = true)] async fn session_affinity_error_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(error_response_stream()); + let mut stream = resolved.into_stream(error_response_stream(), false); assert!(stream.next().await.unwrap().is_err()); assert!(stream.next().await.is_none()); @@ -358,14 +634,9 @@ async fn session_affinity_error_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_stream_eof_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(1)); + let mut stream = resolved.into_stream(response_stream(1), false); while stream.next().await.is_some() {} assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -374,26 +645,17 @@ async fn session_affinity_stream_eof_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_bound_lease_drop_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); tokio::time::advance(Duration::from_secs(9)).await; - let AffinityAcquire::Bound { lease, .. } = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("continuation must acquire the binding"); - }; + let resolved = resolve_local(&coordinator, &session_id(), target(8, Some(0))).await; tokio::time::advance(Duration::from_secs(2)).await; tokio::task::yield_now().await; assert_eq!( - coordinator.query_target(&session_id(), None).unwrap(), + coordinator.query_target(&session_id()).unwrap(), Some(target(7, Some(0))) ); - drop(lease); + drop(resolved); assert_binding_expires_after_refreshed_ttl(&coordinator).await; } @@ -401,27 +663,22 @@ async fn session_affinity_bound_lease_drop_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_query_is_read_only() { let coordinator = coordinator(); - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); assert_eq!(coordinator.entry_count(), 0); - let initializing = coordinator.acquire(&session_id(), None).await.unwrap(); - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + let initializing = coordinator.acquire(&session_id()).await.unwrap(); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); assert_eq!(coordinator.entry_count(), 1); drop(initializing); assert_eq!(coordinator.entry_count(), 0); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); assert_eq!( - coordinator.query_target(&session_id(), None).unwrap(), + coordinator.query_target(&session_id()).unwrap(), Some(target(7, Some(0))) ); coordinator.expire_for_test(&session_id()); - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); assert_eq!(coordinator.entry_count(), 1); } @@ -429,12 +686,7 @@ async fn session_affinity_query_is_read_only() { async fn session_affinity_reaper_removes_idle_entries_and_stops_on_drop() { let coordinator = coordinator(); let cancellation = coordinator.cancellation_token(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); coordinator.wait_for_reaper().await; tokio::time::advance(Duration::from_secs(10)).await; @@ -467,7 +719,7 @@ fn session_affinity_rejects_invalid_ttl_before_starting_reaper() { async fn session_affinity_enforces_id_and_entry_limits() { let coordinator = AffinityCoordinator::with_test_limits(1, 8); let oversized = SessionAffinityId::new("123456789"); - let Err(error) = coordinator.acquire(&oversized, None).await else { + let Err(error) = coordinator.acquire(&oversized).await else { panic!("oversized session ID must fail"); }; assert!(dynamo_runtime::error::match_error_chain( @@ -478,9 +730,9 @@ async fn session_affinity_enforces_id_and_entry_limits() { assert_eq!(coordinator.entry_count(), 0); let first_id = SessionAffinityId::new("first"); - let first = coordinator.acquire(&first_id, None).await.unwrap(); + let first = coordinator.acquire(&first_id).await.unwrap(); let second_id = SessionAffinityId::new("second"); - let Err(error) = coordinator.acquire(&second_id, None).await else { + let Err(error) = coordinator.acquire(&second_id).await else { panic!("entry limit must reject a second session"); }; assert!(dynamo_runtime::error::match_error_chain( @@ -492,7 +744,7 @@ async fn session_affinity_enforces_id_and_entry_limits() { drop(first); assert_eq!(coordinator.entry_count(), 0); assert!(matches!( - coordinator.acquire(&second_id, None).await.unwrap(), + coordinator.acquire(&second_id).await.unwrap(), AffinityAcquire::Initialize(_) )); } diff --git a/lib/runtime/Cargo.toml b/lib/runtime/Cargo.toml index 7aa77c2b4151..54ca9be9e20a 100644 --- a/lib/runtime/Cargo.toml +++ b/lib/runtime/Cargo.toml @@ -78,7 +78,8 @@ bincode = { version = "1" } console-subscriber = { version = "0.4", optional = true } educe = { version = "0.6.0" } figment = { version = "0.10.19", features = ["env", "json", "toml", "test"] } -notify = { version = "6.1", default-features = false, features = ["macos_fsevent"] } +# FSEvents omits deletes for FileStore's revision-zero hard-link publication path. +notify = { version = "6.1", default-features = false, features = ["macos_kqueue"] } libc = { version = "0.2" } local-ip-address = { version = "0.6.3" } # `release_max_level_debug` compiles out `log::trace!` in release builds (a diff --git a/lib/runtime/examples/Cargo.lock b/lib/runtime/examples/Cargo.lock index 2bf02373f8fd..b9e12cda7066 100644 --- a/lib/runtime/examples/Cargo.lock +++ b/lib/runtime/examples/Cargo.lock @@ -1058,15 +1058,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -2062,7 +2053,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.11.1", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", @@ -2857,7 +2847,7 @@ dependencies = [ "pin-project-lite", "quinn", "rustls", - "rustls-native-certs 0.8.3", + "rustls-native-certs", "rustls-pki-types", "serde", "serde_json", diff --git a/lib/runtime/src/discovery/kube.rs b/lib/runtime/src/discovery/kube.rs index a07f0bee2137..611f2c2862c1 100644 --- a/lib/runtime/src/discovery/kube.rs +++ b/lib/runtime/src/discovery/kube.rs @@ -15,14 +15,16 @@ use utils::{KubeDiscoveryMode, PodInfo}; use crate::CancellationToken; use crate::discovery::{ - Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryMetadata, - DiscoveryQuery, DiscoverySpec, DiscoveryStream, MetadataSnapshot, + ClaimCloseOutcome, ClaimOutcome, ClaimPayloadFuture, Discovery, DiscoveryEvent, + DiscoveryInstance, DiscoveryInstanceId, DiscoveryMetadata, DiscoveryQuery, DiscoverySpec, + DiscoveryStream, MetadataSnapshot, }; use anyhow::Result; use async_trait::async_trait; use kube::{Api, Client as KubeClient, api::DeleteParams}; use std::collections::HashSet; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use tokio::sync::RwLock; /// Kubernetes-based discovery client @@ -33,6 +35,7 @@ pub struct KubeDiscoveryClient { metadata_watch: tokio::sync::watch::Receiver>, kube_client: KubeClient, pod_info: PodInfo, + claim_warning_emitted: Arc, } impl KubeDiscoveryClient { @@ -104,8 +107,19 @@ impl KubeDiscoveryClient { metadata_watch: watch_rx, kube_client, pod_info, + claim_warning_emitted: Arc::new(AtomicBool::new(false)), }) } + + fn warn_claims_unsupported(&self) { + if self.claim_warning_emitted.swap(true, Ordering::Relaxed) { + return; + } + + tracing::warn!( + "Kubernetes discovery does not coordinate session affinity across frontend processes; using process-local affinity" + ); + } } #[async_trait] @@ -496,4 +510,18 @@ impl Discovery for KubeDiscoveryClient { let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(event_rx); Ok(Box::pin(stream)) } + + async fn create_or_get_claim( + &self, + _key: &str, + _proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + self.warn_claims_unsupported(); + Ok(ClaimOutcome::Unsupported) + } + + async fn close_claim(&self, _key: &str) -> Result { + self.warn_claims_unsupported(); + Ok(ClaimCloseOutcome::Unsupported) + } } diff --git a/lib/runtime/src/discovery/kv_store.rs b/lib/runtime/src/discovery/kv_store.rs index c72e465f3283..528b20ff4c9f 100644 --- a/lib/runtime/src/discovery/kv_store.rs +++ b/lib/runtime/src/discovery/kv_store.rs @@ -3,27 +3,126 @@ use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; -use anyhow::Result; +use anyhow::{Context, Result}; use async_trait::async_trait; use futures::{Stream, StreamExt}; +use tokio::sync::{OnceCell, broadcast, oneshot}; use tokio_util::sync::CancellationToken; use super::{ - Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, - DiscoverySpec, DiscoveryStream, EndpointInstanceId, EventChannelInstanceId, - ModelCardInstanceId, + ClaimCloseOutcome, ClaimEvent, ClaimOutcome, ClaimPayload, ClaimPayloadFuture, Discovery, + DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, DiscoverySpec, + DiscoveryStream, EndpointInstanceId, EventChannelInstanceId, ModelCardInstanceId, }; use crate::storage::kv; const INSTANCES_BUCKET: &str = "v1/instances"; const MODELS_BUCKET: &str = "v1/mdc"; const EVENT_CHANNELS_BUCKET: &str = "v1/event_channels"; +const CLAIMS_BUCKET: &str = "v1/claims"; +const CLAIM_CREATE_ATTEMPTS: usize = 3; +const CLAIM_WATCH_RECONNECT_BACKOFF: Duration = Duration::from_millis(250); /// Discovery implementation backed by a kv::Store pub struct KVStoreDiscovery { store: Arc, cancel_token: CancellationToken, + claims: ClaimState, +} + +/// Process-local invalidation relay for the shared claims bucket. +/// +/// One backend watcher serves all affinity coordinators attached to this discovery +/// instance. `Put` events never populate caches. `Delete(key)` evicts one entry, while +/// watcher loss, restart, or subscriber lag produces `Reset` so coordinators clear all +/// entries rather than retain potentially stale bindings. +/// +/// TODO: `Bucket::watch` cannot yet surface etcd reconnect/compaction or FileStore +/// overflow errors, so those hidden backend failures cannot be converted into `Reset`. +struct ClaimState { + events: broadcast::Sender, + watcher_started: OnceCell<()>, + memory_warning_emitted: AtomicBool, + #[cfg(test)] + watcher_probe: ClaimWatcherProbe, +} + +impl ClaimState { + fn new() -> Self { + let (events, _) = broadcast::channel(1024); + Self { + events, + watcher_started: OnceCell::new(), + memory_warning_emitted: AtomicBool::new(false), + #[cfg(test)] + watcher_probe: ClaimWatcherProbe::new(), + } + } + + fn subscribe(&self) -> broadcast::Receiver { + self.events.subscribe() + } + + fn warn_if_memory(&self, is_memory: bool) { + if !is_memory || self.memory_warning_emitted.swap(true, Ordering::Relaxed) { + return; + } + + tracing::warn!( + "session affinity claims use MemoryStore and coordinate only within this process/store" + ); + } +} + +#[cfg(test)] +struct ClaimWatcherProbe { + start_count: Arc, + active_count: Arc, +} + +#[cfg(test)] +impl ClaimWatcherProbe { + fn new() -> Self { + Self { + start_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + active_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + + fn record_start(&self) -> Arc { + self.start_count + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.active_count.clone() + } + + fn starts(&self) -> usize { + self.start_count.load(std::sync::atomic::Ordering::Relaxed) + } + + fn active(&self) -> usize { + self.active_count.load(std::sync::atomic::Ordering::Relaxed) + } +} + +#[cfg(test)] +struct ClaimWatcherActiveGuard(Arc); + +#[cfg(test)] +impl ClaimWatcherActiveGuard { + fn new(active_count: Arc) -> Self { + active_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Self(active_count) + } +} + +#[cfg(test)] +impl Drop for ClaimWatcherActiveGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, std::sync::atomic::Ordering::Relaxed); + } } impl KVStoreDiscovery { @@ -31,7 +130,167 @@ impl KVStoreDiscovery { Self { store: Arc::new(store), cancel_token, + claims: ClaimState::new(), + } + } + + async fn ensure_claim_watcher(&self) -> Result<()> { + self.claims + .watcher_started + .get_or_try_init(|| async { + let (ready_tx, ready_rx) = oneshot::channel(); + let store = self.store.clone(); + let cancel_token = self.cancel_token.clone(); + let claim_events = self.claims.events.clone(); + #[cfg(test)] + let active_count = self.claims.watcher_probe.record_start(); + + tokio::spawn(async move { + #[cfg(test)] + let _active_guard = ClaimWatcherActiveGuard::new(active_count); + Self::run_claim_watcher(store, cancel_token, claim_events, ready_tx).await; + }); + + ready_rx + .await + .context("claim watcher stopped before startup completed")? + .map_err(anyhow::Error::msg) + }) + .await?; + Ok(()) + } + + async fn run_claim_watcher( + store: Arc, + cancel_token: CancellationToken, + claim_events: broadcast::Sender, + ready_tx: oneshot::Sender>, + ) { + let mut ready_tx = Some(ready_tx); + + loop { + if cancel_token.is_cancelled() { + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Err("claim watcher startup was cancelled".to_string())); + } + let _ = claim_events.send(ClaimEvent::Reset); + return; + } + + let bucket = match store.get_or_create_bucket(CLAIMS_BUCKET, None).await { + Ok(bucket) => bucket, + Err(err) => { + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Err(err.to_string())); + return; + } + tracing::error!(error = %err, "failed to reconnect session claim watcher"); + let _ = claim_events.send(ClaimEvent::Reset); + if Self::wait_for_claim_watcher_retry(&cancel_token).await { + return; + } + continue; + } + }; + + let mut stream = match bucket.watch().await { + Ok(stream) => stream, + Err(err) => { + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Err(err.to_string())); + return; + } + tracing::error!(error = %err, "failed to reconnect session claim watch stream"); + let _ = claim_events.send(ClaimEvent::Reset); + if Self::wait_for_claim_watcher_retry(&cancel_token).await { + return; + } + continue; + } + }; + + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Ok(())); + } else { + let _ = claim_events.send(ClaimEvent::Reset); + } + + loop { + let event = tokio::select! { + _ = cancel_token.cancelled() => { + let _ = claim_events.send(ClaimEvent::Reset); + return; + } + event = stream.next() => event, + }; + + let Some(event) = event else { + tracing::warn!( + "session claim watch stream ended; clearing local affinity caches" + ); + let _ = claim_events.send(ClaimEvent::Reset); + break; + }; + + if let kv::WatchEvent::Delete(key) = event { + let key = Self::strip_bucket_prefix(key.as_ref(), CLAIMS_BUCKET).to_string(); + let _ = claim_events.send(ClaimEvent::Delete(key)); + } + } + + if Self::wait_for_claim_watcher_retry(&cancel_token).await { + return; + } + } + } + + async fn wait_for_claim_watcher_retry(cancel_token: &CancellationToken) -> bool { + tokio::select! { + _ = cancel_token.cancelled() => true, + _ = tokio::time::sleep(CLAIM_WATCH_RECONNECT_BACKOFF) => false, + } + } + + fn warn_if_memory_claims(&self) { + self.claims.warn_if_memory(self.store.is_memory()); + } + + fn parse_claim(value: &[u8]) -> Result { + serde_json::from_slice(value).context("failed to deserialize session affinity claim") + } + + async fn create_or_get_in_bucket( + bucket: &dyn kv::Bucket, + key: &kv::Key, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + if let Some(payload) = bucket.get(key).await? { + return Ok(ClaimOutcome::Existing(Self::parse_claim(&payload)?)); + } + + let proposed_payload = proposed_payload.as_mut().await?; + let proposed_bytes = serde_json::to_vec(&proposed_payload)?; + + for attempt in 0..CLAIM_CREATE_ATTEMPTS { + match bucket.insert(key, proposed_bytes.clone().into(), 0).await? { + kv::StoreOutcome::Created(_) => { + return Ok(ClaimOutcome::Created(proposed_payload)); + } + kv::StoreOutcome::Exists(_) => { + if let Some(payload) = bucket.get(key).await? { + return Ok(ClaimOutcome::Existing(Self::parse_claim(&payload)?)); + } + + if attempt + 1 == CLAIM_CREATE_ATTEMPTS { + anyhow::bail!( + "session affinity claim disappeared after {CLAIM_CREATE_ATTEMPTS} competing insert attempts" + ); + } + } + } } + + unreachable!("claim creation loop always returns") } /// Build the key path for an endpoint (relative to bucket, not absolute) @@ -592,6 +851,41 @@ impl Discovery for KVStoreDiscovery { Ok(Box::pin(stream)) } + async fn create_or_get_claim( + &self, + key: &str, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + self.warn_if_memory_claims(); + self.ensure_claim_watcher().await?; + + let bucket = self.store.get_or_create_bucket(CLAIMS_BUCKET, None).await?; + let key = kv::Key::new(key.to_string()); + + Self::create_or_get_in_bucket(bucket.as_ref(), &key, proposed_payload).await + } + + async fn close_claim(&self, key: &str) -> Result { + self.warn_if_memory_claims(); + self.ensure_claim_watcher().await?; + + let Some(bucket) = self.store.get_bucket(CLAIMS_BUCKET).await? else { + return Ok(ClaimCloseOutcome::Closed); + }; + let key = kv::Key::new(key.to_string()); + + match bucket.delete(&key).await { + Ok(()) | Err(kv::StoreError::MissingBucket(_) | kv::StoreError::MissingKey(_)) => { + Ok(ClaimCloseOutcome::Closed) + } + Err(err) => Err(err.into()), + } + } + + fn subscribe_claim_events(&self) -> Option> { + Some(self.claims.subscribe()) + } + fn shutdown(&self) { self.store.shutdown(); } @@ -599,9 +893,269 @@ impl Discovery for KVStoreDiscovery { #[cfg(test)] mod tests { + use std::collections::HashMap; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use super::*; use crate::component::TransportType; + fn payload(worker_id: u64) -> ClaimPayload { + serde_json::json!({"worker_id": worker_id, "dp_rank": 0}) + } + + struct DisappearingBucket { + insert_calls: AtomicUsize, + create_on_call: Option, + } + + struct InsertBarrierBucket { + inner: Box, + barrier: tokio::sync::Barrier, + } + + #[async_trait] + impl kv::Bucket for InsertBarrierBucket { + async fn insert( + &self, + key: &kv::Key, + value: bytes::Bytes, + revision: u64, + ) -> std::result::Result { + self.barrier.wait().await; + self.inner.insert(key, value, revision).await + } + + async fn get( + &self, + key: &kv::Key, + ) -> std::result::Result, kv::StoreError> { + self.inner.get(key).await + } + + async fn delete(&self, key: &kv::Key) -> std::result::Result<(), kv::StoreError> { + self.inner.delete(key).await + } + + async fn watch( + &self, + ) -> std::result::Result< + Pin + Send + '_>>, + kv::StoreError, + > { + self.inner.watch().await + } + + async fn entries( + &self, + ) -> std::result::Result, kv::StoreError> { + self.inner.entries().await + } + } + + #[async_trait] + impl kv::Bucket for DisappearingBucket { + async fn insert( + &self, + _key: &kv::Key, + _value: bytes::Bytes, + _revision: u64, + ) -> std::result::Result { + let call = self.insert_calls.fetch_add(1, Ordering::Relaxed); + Ok(if self.create_on_call == Some(call) { + kv::StoreOutcome::Created(1) + } else { + kv::StoreOutcome::Exists(1) + }) + } + + async fn get( + &self, + _key: &kv::Key, + ) -> std::result::Result, kv::StoreError> { + Ok(None) + } + + async fn delete(&self, _key: &kv::Key) -> std::result::Result<(), kv::StoreError> { + Ok(()) + } + + async fn watch( + &self, + ) -> std::result::Result< + Pin + Send + '_>>, + kv::StoreError, + > { + Ok(Box::pin(futures::stream::pending())) + } + + async fn entries( + &self, + ) -> std::result::Result, kv::StoreError> { + Ok(HashMap::new()) + } + } + + #[tokio::test] + async fn existing_claim_does_not_poll_proposal() { + let client = KVStoreDiscovery::new(kv::Manager::memory(), CancellationToken::new()); + let mut first: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(7)) }); + assert_eq!( + client + .create_or_get_claim("scope/session", &mut first) + .await + .unwrap(), + ClaimOutcome::Created(payload(7)) + ); + + let polled = Arc::new(AtomicBool::new(false)); + let proposal_polled = polled.clone(); + let mut second: ClaimPayloadFuture<'_> = Box::pin(async move { + proposal_polled.store(true, Ordering::Relaxed); + Ok(payload(8)) + }); + assert_eq!( + client + .create_or_get_claim("scope/session", &mut second) + .await + .unwrap(), + ClaimOutcome::Existing(payload(7)) + ); + assert!(!polled.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn competing_claims_return_one_created_winner() { + let store = kv::Manager::memory(); + let bucket = Arc::new(InsertBarrierBucket { + inner: store + .get_or_create_bucket(CLAIMS_BUCKET, None) + .await + .unwrap(), + barrier: tokio::sync::Barrier::new(8), + }); + let key = Arc::new(kv::Key::new("scope/race".to_string())); + let mut tasks = Vec::new(); + for worker_id in 0..8 { + let bucket = bucket.clone(); + let key = key.clone(); + tasks.push(tokio::spawn(async move { + let mut proposal: ClaimPayloadFuture<'_> = + Box::pin(async move { Ok(payload(worker_id)) }); + KVStoreDiscovery::create_or_get_in_bucket( + bucket.as_ref(), + key.as_ref(), + &mut proposal, + ) + .await + .unwrap() + })); + } + + let outcomes = futures::future::join_all(tasks) + .await + .into_iter() + .map(Result::unwrap) + .collect::>(); + assert_eq!( + outcomes + .iter() + .filter(|outcome| matches!(outcome, ClaimOutcome::Created(_))) + .count(), + 1 + ); + let winner = match outcomes + .iter() + .find(|outcome| matches!(outcome, ClaimOutcome::Created(_))) + .unwrap() + { + ClaimOutcome::Created(payload) => payload, + _ => unreachable!(), + }; + assert!(outcomes.iter().all(|outcome| match outcome { + ClaimOutcome::Created(payload) | ClaimOutcome::Existing(payload) => payload == winner, + ClaimOutcome::Unsupported => false, + })); + } + + #[tokio::test] + async fn claim_disappearance_retries_and_is_bounded() { + let key = kv::Key::new("scope/disappearing".to_string()); + let recovering = DisappearingBucket { + insert_calls: AtomicUsize::new(0), + create_on_call: Some(1), + }; + let mut proposal: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(7)) }); + assert_eq!( + KVStoreDiscovery::create_or_get_in_bucket(&recovering, &key, &mut proposal) + .await + .unwrap(), + ClaimOutcome::Created(payload(7)) + ); + assert_eq!(recovering.insert_calls.load(Ordering::Relaxed), 2); + + let exhausting = DisappearingBucket { + insert_calls: AtomicUsize::new(0), + create_on_call: None, + }; + let mut proposal: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(8)) }); + let error = KVStoreDiscovery::create_or_get_in_bucket(&exhausting, &key, &mut proposal) + .await + .unwrap_err(); + assert!(error.to_string().contains("3 competing insert attempts")); + assert_eq!(exhausting.insert_calls.load(Ordering::Relaxed), 3); + } + + #[tokio::test] + async fn claim_watcher_ignores_put_emits_delete_and_close_is_idempotent() { + let cancel = CancellationToken::new(); + let client = KVStoreDiscovery::new(kv::Manager::memory(), cancel.clone()); + let mut events = client.subscribe_claim_events().unwrap(); + let mut proposal: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(7)) }); + client + .create_or_get_claim("scope/close", &mut proposal) + .await + .unwrap(); + + assert_eq!(client.claims.watcher_probe.starts(), 1); + assert_eq!(client.claims.watcher_probe.active(), 1); + + assert_eq!( + client.close_claim("scope/close").await.unwrap(), + ClaimCloseOutcome::Closed + ); + assert_eq!( + events.recv().await.unwrap(), + ClaimEvent::Delete("scope/close".to_string()) + ); + assert_eq!( + client.close_claim("scope/close").await.unwrap(), + ClaimCloseOutcome::Closed + ); + + cancel.cancel(); + } + + #[tokio::test] + async fn claim_watcher_stops_on_cancellation() { + let store = Arc::new(kv::Manager::memory()); + let cancel = CancellationToken::new(); + let (events, _) = broadcast::channel(16); + let (ready_tx, ready_rx) = oneshot::channel(); + let watcher = tokio::spawn(KVStoreDiscovery::run_claim_watcher( + store, + cancel.clone(), + events, + ready_tx, + )); + ready_rx.await.unwrap().unwrap(); + + cancel.cancel(); + tokio::time::timeout(Duration::from_secs(1), watcher) + .await + .expect("claim watcher did not stop after cancellation") + .unwrap(); + } + #[tokio::test] async fn test_kv_store_discovery_register_endpoint() { let store = kv::Manager::memory(); diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs index d4a03715f4a3..91ec2b76c911 100644 --- a/lib/runtime/src/discovery/mod.rs +++ b/lib/runtime/src/discovery/mod.rs @@ -5,7 +5,9 @@ use anyhow::{Context, Result}; use async_trait::async_trait; use futures::Stream; use serde::{Deserialize, Serialize}; +use std::future::Future; use std::pin::Pin; +use tokio::sync::broadcast; use tokio_util::sync::CancellationToken; mod metadata; @@ -23,6 +25,28 @@ pub mod utils; use crate::component::{DeviceType, TransportType}; pub use utils::watch_and_extract_field; +pub type ClaimPayload = serde_json::Value; +pub type ClaimPayloadFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Debug, Clone, PartialEq)] +pub enum ClaimOutcome { + Created(ClaimPayload), + Existing(ClaimPayload), + Unsupported, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClaimCloseOutcome { + Closed, + Unsupported, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ClaimEvent { + Delete(String), + Reset, +} + /// Transport kind for event plane - used for configuration and env var selection. /// /// This enum represents the *type* of transport without connection details. @@ -844,6 +868,34 @@ pub trait Discovery: Send + Sync { cancel_token: Option, ) -> Result; + /// Returns an existing immutable claim or atomically creates it from a deferred proposal. + /// + /// Implementations read before polling `proposed_payload`, atomically insert only when + /// absent, and return the winning stored payload after an insertion race. Payloads in + /// [`ClaimOutcome::Created`] and [`ClaimOutcome::Existing`] are authoritative. + /// [`ClaimOutcome::Unsupported`] leaves coordination process-local. Storage errors must + /// propagate to the caller before scheduler bookkeeping or dispatch. + async fn create_or_get_claim( + &self, + _key: &str, + _proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + Ok(ClaimOutcome::Unsupported) + } + + /// Idempotently closes an immutable claim. + /// + /// Close is terminal under the session-ID no-reuse contract; deleting an absent claim + /// succeeds. + async fn close_claim(&self, _key: &str) -> Result { + Ok(ClaimCloseOutcome::Unsupported) + } + + /// Subscribes to process-local claim invalidation events. + fn subscribe_claim_events(&self) -> Option> { + None + } + /// Clean up resources held by this discovery backend. /// For KV store backends, this deletes owned registrations immediately rather than /// waiting for TTL expiry. Default is a no-op for backends that don't need cleanup. diff --git a/lib/runtime/src/pipeline/network/egress/push_router.rs b/lib/runtime/src/pipeline/network/egress/push_router.rs index 217739cc5c84..42dddd5aa770 100644 --- a/lib/runtime/src/pipeline/network/egress/push_router.rs +++ b/lib/runtime/src/pipeline/network/egress/push_router.rs @@ -825,6 +825,32 @@ where Ok((metadata, stream)) } + /// Book a previously arbitrated worker and dispatch without reselection or fallback. + pub async fn book_and_dispatch_exact( + &self, + mut request: SingleIn, + instance_id: u64, + advance_round_robin: bool, + prepare: F, + ) -> anyhow::Result<(M, ManyOut)> + where + F: FnOnce(&mut T, u64) -> anyhow::Result, + { + if advance_round_robin && self.router_mode == RouterMode::RoundRobin { + self.round_robin_counter.fetch_add(1, Ordering::Relaxed); + } + let (instance_id, permit) = self + .select_exact_target(request.content(), Some(instance_id)) + .await?; + let metadata = prepare(&mut request, instance_id)?; + let stream = self.dispatch_exact(request, instance_id).await?; + let stream = match permit { + Some(permit) => permit.into_tracked_stream(stream), + None => stream, + }; + Ok((metadata, stream)) + } + /// Issue a request using device-aware weighted routing. /// /// Instances are partitioned by device type (CPU vs non-CPU), then the router @@ -1075,6 +1101,23 @@ where } } + /// Peek the worker this routing mode would choose for a request without booking it. + pub fn peek_worker_for_request(&self, request: &T) -> Option { + let instance_ids = self.client.routing_instances().free_ids().to_vec(); + if instance_ids.is_empty() { + return None; + } + + match self.router_mode { + RouterMode::DeviceAwareWeighted => { + let state = self.occupancy_state.as_deref()?; + let selection = self.device_aware_candidates(request, state, &instance_ids); + state.peek_min(&selection.candidates) + } + _ => self.peek_next_worker(), + } + } + async fn select_exact_target( &self, request: &T, diff --git a/lib/runtime/src/storage/kv.rs b/lib/runtime/src/storage/kv.rs index 544a056e27f1..207142842b50 100644 --- a/lib/runtime/src/storage/kv.rs +++ b/lib/runtime/src/storage/kv.rs @@ -273,6 +273,10 @@ impl Manager { Manager(Arc::new(s)) } + pub fn is_memory(&self) -> bool { + matches!(self.0.as_ref(), KeyValueStoreEnum::Memory(_)) + } + pub async fn get_or_create_bucket( &self, bucket_name: &str, diff --git a/lib/runtime/src/storage/kv/file.rs b/lib/runtime/src/storage/kv/file.rs index bbb138b4cb8a..8587b6c46388 100644 --- a/lib/runtime/src/storage/kv/file.rs +++ b/lib/runtime/src/storage/kv/file.rs @@ -195,8 +195,9 @@ pub struct Directory { impl Directory { fn new(root: PathBuf, p: PathBuf, ttl: Duration) -> Self { - // Canonicalize root to handle symlinks (e.g., /var -> /private/var on macOS) + // Keep watched paths and event paths in the same form across symlinked roots. let canonical_root = root.canonicalize().unwrap_or_else(|_| root.clone()); + let canonical_path = p.canonicalize().unwrap_or_else(|_| p.clone()); if ttl < MIN_KEEP_ALIVE { let h_ttl = humantime::format_duration(ttl); tracing::warn!(path = %p.display(), ttl = %h_ttl, "ttl is too short, increasing to {}", humantime::format_duration(MIN_KEEP_ALIVE)); @@ -204,7 +205,7 @@ impl Directory { let ttl = cmp::max(ttl, MIN_KEEP_ALIVE); Directory { root: canonical_root, - p, + p: canonical_path, ttl, owned_files: Arc::new(Mutex::new(HashSet::new())), } @@ -441,7 +442,6 @@ impl Bucket for Directory { continue; } }; - for item_path in event.paths { // Skip if the event is for the directory itself if item_path == dir { @@ -449,9 +449,7 @@ impl Bucket for Directory { continue; } - // Canonicalize paths to handle symlinks (e.g., /var -> /private/var on macOS) - // The unwrap_or_else path is for Remove case. - let canonical_item_path = item_path.canonicalize().unwrap_or_else(|_| item_path.clone()); + let canonical_item_path = canonicalize_event_path(&item_path); let key = match canonical_item_path.strip_prefix(&root) { Ok(stripped) => Key::from_url_safe(&stripped.display().to_string()), @@ -490,7 +488,7 @@ impl Bucket for Directory { let item = KeyValue::new(key, data); yield WatchEvent::Put(item); } - EventKind::Remove(event::RemoveKind::File) => { + EventKind::Remove(_) => { yield WatchEvent::Delete(key); } _ => { @@ -589,6 +587,19 @@ fn write_temp_file_at(temp_path: &Path, value: &[u8]) -> Result PathBuf { + if let Ok(canonical_path) = path.canonicalize() { + return canonical_path; + } + let (Some(parent), Some(file_name)) = (path.parent(), path.file_name()) else { + return path.to_path_buf(); + }; + let Ok(canonical_parent) = parent.canonicalize() else { + return path.to_path_buf(); + }; + canonical_parent.join(file_name) +} + // For anyhow preserve the context fn a_to_fs_err(err: anyhow::Error) -> StoreError { StoreError::FilesystemError(format!("{err:#}")) @@ -602,11 +613,82 @@ fn to_fs_err(err: E) -> StoreError { mod tests { use std::collections::HashSet; use std::fs; + use std::os::unix::fs::symlink; + use std::time::Duration; + use futures::StreamExt; use tokio_util::sync::CancellationToken; use crate::storage::kv::{Bucket as _, FileStore, Key, Store as _, StoreOutcome}; + #[test] + fn deleted_event_path_canonicalizes_existing_parent() { + let t = tempfile::tempdir().unwrap(); + let canonical_root = t.path().join("canonical"); + let bucket = canonical_root.join("v1/claims"); + fs::create_dir_all(&bucket).unwrap(); + let linked_root = t.path().join("linked"); + symlink(&canonical_root, &linked_root).unwrap(); + + assert_eq!( + super::canonicalize_event_path(&linked_root.join("v1/claims/deleted")), + canonical_root + .canonicalize() + .unwrap() + .join("v1/claims/deleted") + ); + } + + #[tokio::test] + async fn external_delete_is_observed_under_noncanonical_root() { + let t = tempfile::tempdir().unwrap(); + let canonical_root = t.path().join("canonical"); + fs::create_dir_all(&canonical_root).unwrap(); + let linked_root = t.path().join("linked"); + symlink(&canonical_root, &linked_root).unwrap(); + let watcher_cancel = CancellationToken::new(); + let creator_cancel = CancellationToken::new(); + let watcher_store = FileStore::new(watcher_cancel.clone(), &linked_root); + let creator_store = FileStore::new(creator_cancel.clone(), &canonical_root); + let watcher_bucket = watcher_store + .get_or_create_bucket("v1/claims", None) + .await + .unwrap(); + let creator_bucket = creator_store + .get_or_create_bucket("v1/claims", None) + .await + .unwrap(); + let mut events = watcher_bucket.watch().await.unwrap(); + let key = Key::new("scope/session".to_string()); + + creator_bucket + .insert(&key, "value".into(), 0) + .await + .unwrap(); + loop { + let event = tokio::time::timeout(Duration::from_secs(2), events.next()) + .await + .expect("FileStore watcher did not observe claim creation") + .expect("FileStore watcher ended after claim creation"); + if matches!(event, super::WatchEvent::Put(ref item) if item.key_str() == "v1/claims/scope/session") + { + break; + } + } + + creator_bucket.delete(&key).await.unwrap(); + let event = tokio::time::timeout(Duration::from_secs(2), events.next()) + .await + .expect("FileStore watcher did not observe claim deletion") + .expect("FileStore watcher ended after claim deletion"); + assert!( + matches!(event, super::WatchEvent::Delete(ref deleted) if deleted == &Key::new("v1/claims/scope/session".to_string())) + ); + + watcher_cancel.cancel(); + creator_cancel.cancel(); + } + #[tokio::test] async fn test_entries_full_path() { let t = tempfile::tempdir().unwrap(); diff --git a/tests/router/common.py b/tests/router/common.py index 61c42aaba46f..2222295b676c 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -9,6 +9,7 @@ import random import threading import time +import uuid from typing import TYPE_CHECKING, Any, Callable, Optional import aiohttp @@ -521,6 +522,216 @@ async def verify_consumer_lifecycle(): kv_router.__exit__(None, None, None) +def _test_distributed_session_affinity( + engine_workers, + block_size: int, + request, + router_ports: list[int], + test_payload: dict[str, Any], + store_backend: str = "etcd", +): + """Verify shared affinity claims override conflicting KV-prefix placement.""" + with ( + FrontendRouterProcess( + request, + block_size, + router_ports[0], + engine_workers.namespace, + store_backend, + router_mode="kv", + min_initial_workers=engine_workers.num_workers, + event_plane="nats", + session_affinity_ttl_secs=300, + ) as first_router, + FrontendRouterProcess( + request, + block_size, + router_ports[1], + engine_workers.namespace, + store_backend, + router_mode="kv", + min_initial_workers=engine_workers.num_workers, + event_plane="nats", + session_affinity_ttl_secs=300, + ) as second_router, + ): + urls = [f"http://localhost:{port}/v1/chat/completions" for port in router_ports] + + async def run_test() -> None: + runtime = get_runtime(store_backend, "nats") + endpoint = runtime.endpoint( + f"{engine_workers.namespace}.{engine_workers.component_name}.generate" + ) + worker_ids = sorted( + await poll_for_worker_instances(endpoint, engine_workers.num_workers) + ) + assert len(worker_ids) >= 2 + worker_a, worker_b = worker_ids[:2] + + for port in router_ports: + await wait_for_frontend_ready( + frontend_url=f"http://localhost:{port}", + expected_num_workers=engine_workers.num_workers, + timeout=120, + engine_workers=engine_workers, + store_backend=store_backend, + request_plane="nats", + ) + + suffix = uuid.uuid4().hex + prefix_a = " ".join([f"affinity-alpha-{suffix}"] * (block_size * 2)) + prefix_b = " ".join([f"affinity-beta-{suffix}"] * (block_size * 2)) + session_a = f"distributed-affinity-a-{uuid.uuid4()}" + session_b = f"distributed-affinity-b-{uuid.uuid4()}" + + def payload(content: str, *, query_only: bool = False) -> dict[str, Any]: + annotations = ["query_instance_id:"] if query_only else [] + return { + **test_payload, + "messages": [{"role": "user", "content": content}], + "stream": True, + "max_tokens": 1, + "nvext": { + "annotations": annotations, + "extra_fields": ["worker_id"], + }, + } + + async def send( + client: aiohttp.ClientSession, + url: str, + request_payload: dict[str, Any], + headers: dict[str, str] | None = None, + ) -> tuple[int, int]: + async with client.post( + url, json=request_payload, headers=headers + ) as response: + body = await response.text() + assert response.status == 200, body + + worker_info = None + for line in body.splitlines(): + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + continue + candidate = json.loads(data).get("nvext", {}).get("worker_id") + if candidate: + worker_info = candidate + + assert worker_info is not None, body + return ( + worker_info["decode_worker_id"], + worker_info["decode_dp_rank"], + ) + + async def wait_for_prefix_target( + client: aiohttp.ClientSession, + url: str, + content: str, + expected: tuple[int, int], + ) -> None: + for _ in range(50): + if ( + await send(client, url, payload(content, query_only=True)) + == expected + ): + return + await asyncio.sleep(0.1) + raise AssertionError( + f"KV events did not make prefix target {expected} visible" + ) + + session_a_headers = {"x-dynamo-session-id": session_a} + session_b_headers = {"x-dynamo-session-id": session_b} + proposal_a = { + **session_a_headers, + "x-dynamo-worker-instance-id": str(worker_a), + "x-dynamo-dp-rank": "0", + } + proposal_b = { + **session_b_headers, + "x-dynamo-worker-instance-id": str(worker_b), + "x-dynamo-dp-rank": "0", + } + + async with aiohttp.ClientSession() as client: + assert await send(client, urls[0], payload(prefix_a), proposal_a) == ( + worker_a, + 0, + ) + assert await send(client, urls[1], payload(prefix_b), proposal_b) == ( + worker_b, + 0, + ) + + await wait_for_prefix_target(client, urls[0], prefix_a, (worker_a, 0)) + await wait_for_prefix_target(client, urls[1], prefix_b, (worker_b, 0)) + + assert await send( + client, urls[0], payload(prefix_a), session_b_headers + ) == (worker_b, 0) + assert await send( + client, urls[1], payload(prefix_b), session_a_headers + ) == (worker_a, 0) + + first_evictions = first_router.read_logs().count( + "evicted session affinity cache entry" + ) + assert await send( + client, + urls[1], + payload(prefix_b), + { + **session_a_headers, + "x-dynamo-session-final": "true", + }, + ) == (worker_a, 0) + + for _ in range(50): + if ( + first_router.read_logs().count( + "evicted session affinity cache entry" + ) + > first_evictions + ): + break + await asyncio.sleep(0.1) + else: + raise AssertionError( + "first frontend did not observe session A claim deletion" + ) + + second_evictions = second_router.read_logs().count( + "evicted session affinity cache entry" + ) + assert await send( + client, + urls[0], + payload(prefix_a), + { + **session_b_headers, + "x-dynamo-session-final": "true", + }, + ) == (worker_b, 0) + + for _ in range(50): + if ( + second_router.read_logs().count( + "evicted session affinity cache entry" + ) + > second_evictions + ): + return + await asyncio.sleep(0.1) + raise AssertionError( + "second frontend did not observe session B claim deletion" + ) + + asyncio.run(run_test()) + + def _test_remote_indexer_decisions( engine_workers, model_name: str, diff --git a/tests/router/router_process.py b/tests/router/router_process.py index 330453eb5081..988fbbf19886 100644 --- a/tests/router/router_process.py +++ b/tests/router/router_process.py @@ -58,6 +58,7 @@ def __init__( serve_indexer: bool = False, use_remote_indexer: bool = False, event_plane: str | None = None, + session_affinity_ttl_secs: int | None = None, ): command = [ sys.executable, @@ -102,6 +103,11 @@ def __init__( if use_remote_indexer: command.append("--use-remote-indexer") + if session_affinity_ttl_secs is not None: + command.extend( + ["--router-session-affinity-ttl-secs", str(session_affinity_ttl_secs)] + ) + if router_aic_config is not None: command.extend( [ @@ -146,7 +152,7 @@ def __init__( ], log_dir=request.node.name, terminate_all_matching_process_names=False, - display_name=f"dynamo-frontend-{router_mode}", + display_name=f"dynamo-frontend-{router_mode}-{frontend_port}", ) self.port = frontend_port self.router_mode = router_mode diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index 273b58788bd4..c848bade8a5f 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -23,6 +23,7 @@ _test_disagg_direct_mode, _test_disagg_router_overload_529, _test_disagg_topology_required_prefill_pin_match_and_mismatch, + _test_distributed_session_affinity, _test_python_router_bindings, _test_remote_indexer_decisions, _test_router_decisions_disagg_round_robin_prefill_dp_rank, @@ -490,6 +491,44 @@ def test_mocker_two_kv_router( ) +@pytest.mark.parametrize("store_backend", ["etcd", "file"]) +@pytest.mark.timeout(180) +def test_mocker_distributed_session_affinity( + request, + runtime_services_dynamic_ports, + predownload_tokenizers, + file_storage_backend, + store_backend, + monkeypatch, +): + """Shared claims override conflicting KV-prefix routing on another frontend.""" + current_log = os.environ.get("DYN_LOG", "info") + monkeypatch.setenv( + "DYN_LOG", + f"{current_log},dynamo_llm::session_affinity::coordinator=debug", + ) + mocker_args = { + "speedup_ratio": SPEEDUP_RATIO, + "block_size": BLOCK_SIZE, + "durable_kv_events": False, + } + + with MockerProcess( + request, + mocker_args=mocker_args, + num_mockers=NUM_MOCKERS, + store_backend=store_backend, + ) as mockers: + _test_distributed_session_affinity( + engine_workers=mockers, + block_size=BLOCK_SIZE, + request=request, + router_ports=allocate_frontend_ports(request, 2), + test_payload=TEST_PAYLOAD, + store_backend=store_backend, + ) + + @pytest.mark.parametrize( "durable_kv_events", [False], ids=["nondurable"], indirect=True ) # Use NATS Core (local indexer) From 97b8acfb817df24f1ae9cdaa8b697ee280e54e30 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Tue, 30 Jun 2026 08:50:30 -0700 Subject: [PATCH 004/320] perf(runtime): use parking_lot for TCP stream registry (#11065) Signed-off-by: jthomson04 --- .../src/pipeline/network/tcp/server.rs | 197 ++++++++++++------ 1 file changed, 133 insertions(+), 64 deletions(-) diff --git a/lib/runtime/src/pipeline/network/tcp/server.rs b/lib/runtime/src/pipeline/network/tcp/server.rs index fbafc8a228a0..45d9cfa4f14f 100644 --- a/lib/runtime/src/pipeline/network/tcp/server.rs +++ b/lib/runtime/src/pipeline/network/tcp/server.rs @@ -9,7 +9,6 @@ use std::{ sync::Arc, time::Duration, }; -use tokio::sync::Mutex; use tokio::time::Instant; /// Tombstone lifetime. Bridges the `register()` → `associate_instance()` @@ -22,6 +21,7 @@ use bytes::Bytes; use derive_builder::Builder; use futures::{SinkExt, StreamExt}; use local_ip_address::{Error, list_afinet_netifas, local_ip, local_ipv6}; +use parking_lot::Mutex; use serde::{Deserialize, Serialize}; use tokio::{ @@ -259,7 +259,7 @@ impl TcpStreamServer { send_subject: Option<&str>, id: &EndpointInstanceId, ) -> bool { - let mut state = self.state.lock().await; + let mut state = self.state.lock(); let now = Instant::now(); prune_tombstones(&mut state.removed_instances, now); if state.removed_instances.contains_key(id) { @@ -296,7 +296,7 @@ impl TcpStreamServer { /// Cancel one pending response-stream registration. Drops the /// `oneshot::Sender` so the waiting receiver resolves with `RecvError`. pub async fn cancel_recv_stream(&self, subject: &str) { - let mut state = self.state.lock().await; + let mut state = self.state.lock(); state.rx_subjects.remove(subject); if let Some(key) = state.subject_instance.remove(subject) && let Some(subjects) = state.instance_subjects.get_mut(&key) @@ -314,7 +314,7 @@ impl TcpStreamServer { /// `(StreamType::Request, _)` tag from `instance_subjects` so the per- /// instance bookkeeping stays consistent. pub async fn cancel_send_stream(&self, subject: &str) { - let mut state = self.state.lock().await; + let mut state = self.state.lock(); state.tx_subjects.remove(subject); if let Some(key) = state.subject_instance.remove(subject) && let Some(subjects) = state.instance_subjects.get_mut(&key) @@ -331,7 +331,7 @@ impl TcpStreamServer { /// `associate_instance` — and tombstone the id so any racing associate /// for the same id cancels too. Returns the number of streams cancelled. pub async fn cancel_instance_streams(&self, id: &EndpointInstanceId) -> usize { - let mut state = self.state.lock().await; + let mut state = self.state.lock(); let now = Instant::now(); prune_tombstones(&mut state.removed_instances, now); state.removed_instances.insert(id.clone(), now); @@ -357,25 +357,63 @@ impl TcpStreamServer { /// Drop the tombstone for an instance that has reappeared in discovery, /// so future subjects for that identity are tracked normally. pub async fn clear_instance_tombstone(&self, id: &EndpointInstanceId) { - let mut state = self.state.lock().await; + let mut state = self.state.lock(); state.removed_instances.remove(id); } - #[allow(clippy::await_holding_lock)] async fn start(local_ip: String, local_port: u16, state: Arc>) -> Result { let addr = format!("{}:{}", local_ip, local_port); let state_clone = state.clone(); - let mut guard = state.lock().await; - if guard.handle.is_some() { - panic!("TcpStreamServer already started"); - } let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::>(); - let handle = tokio::spawn(tcp_listener(addr, state_clone, ready_tx)); - guard.handle = Some(handle); - drop(guard); + { + let mut guard = state.lock(); + if guard.handle.is_some() { + panic!("TcpStreamServer already started"); + } + guard.handle = Some(tokio::spawn(tcp_listener(addr, state_clone, ready_tx))); + } let local_port = ready_rx.await??; Ok(local_port) } + + fn insert_request_stream(&self, subject: String, connection: RequestedSendConnection) { + self.state.lock().tx_subjects.insert(subject, connection); + } + + fn insert_response_stream(&self, subject: String, connection: RequestedRecvConnection) { + self.state.lock().rx_subjects.insert(subject, connection); + } + + fn take_request_stream(state: &Mutex, subject: &str) -> Option { + let mut state = state.lock(); + let connection = state.tx_subjects.remove(subject); + if let Some(key) = state.subject_instance.remove(subject) + && let Some(subjects) = state.instance_subjects.get_mut(&key) + { + subjects.remove(&(StreamType::Request, subject.to_string())); + if subjects.is_empty() { + state.instance_subjects.remove(&key); + } + } + connection + } + + fn take_response_stream( + state: &Mutex, + subject: &str, + ) -> Option { + let mut state = state.lock(); + let connection = state.rx_subjects.remove(subject); + if let Some(key) = state.subject_instance.remove(subject) + && let Some(subjects) = state.instance_subjects.get_mut(&key) + { + subjects.remove(&(StreamType::Response, subject.to_string())); + if subjects.is_empty() { + state.instance_subjects.remove(&key); + } + } + connection + } } // todo - possible rename ResponseService to ResponseServer @@ -409,6 +447,7 @@ impl ResponseService for TcpStreamServer { let send_stream = if options.enable_request_stream { let sender_subject = uuid::Uuid::new_v4().to_string(); + let registry_subject = sender_subject.clone(); let (pending_sender_tx, pending_sender_rx) = oneshot::channel(); @@ -418,11 +457,6 @@ impl ResponseService for TcpStreamServer { send_buffer_count: options.send_buffer_count, }; - let mut state = self.state.lock().await; - state - .tx_subjects - .insert(sender_subject.clone(), connection_info); - let cleanup_subject = sender_subject.clone(); let cleanup_state = self.state.clone(); let registered_stream = RegisteredStream::new( @@ -438,7 +472,7 @@ impl ResponseService for TcpStreamServer { .with_cleanup(move || { // Drop is sync; fire-and-forget the lock acquisition. tokio::spawn(async move { - let mut state = cleanup_state.lock().await; + let mut state = cleanup_state.lock(); state.tx_subjects.remove(&cleanup_subject); if let Some(key) = state.subject_instance.remove(&cleanup_subject) && let Some(subjects) = state.instance_subjects.get_mut(&key) @@ -451,6 +485,8 @@ impl ResponseService for TcpStreamServer { }); }); + self.insert_request_stream(registry_subject, connection_info); + Some(registered_stream) } else { None @@ -459,6 +495,7 @@ impl ResponseService for TcpStreamServer { let recv_stream = if options.enable_response_stream { let (pending_recver_tx, pending_recver_rx) = oneshot::channel(); let receiver_subject = uuid::Uuid::new_v4().to_string(); + let registry_subject = receiver_subject.clone(); let connection_info = RequestedRecvConnection { context: options.context.clone(), @@ -466,11 +503,6 @@ impl ResponseService for TcpStreamServer { send_buffer_count: options.send_buffer_count, }; - let mut state = self.state.lock().await; - state - .rx_subjects - .insert(receiver_subject.clone(), connection_info); - let cleanup_subject = receiver_subject.clone(); let cleanup_state = self.state.clone(); let registered_stream = RegisteredStream::new( @@ -486,7 +518,7 @@ impl ResponseService for TcpStreamServer { .with_cleanup(move || { // Drop is sync; fire-and-forget the lock acquisition. tokio::spawn(async move { - let mut state = cleanup_state.lock().await; + let mut state = cleanup_state.lock(); state.rx_subjects.remove(&cleanup_subject); if let Some(key) = state.subject_instance.remove(&cleanup_subject) && let Some(subjects) = state.instance_subjects.get_mut(&key) @@ -499,6 +531,8 @@ impl ResponseService for TcpStreamServer { }); }); + self.insert_response_stream(registry_subject, connection_info); + Some(registered_stream) } else { None @@ -653,22 +687,12 @@ async fn tcp_listener( // Request stream is unidirectional; we don't read from the downstream. drop(reader); - let request_stream = { - let mut guard = state.lock().await; - let conn = guard.tx_subjects.remove(&subject).ok_or(error!( + let request_stream = TcpStreamServer::take_request_stream(&state, &subject).ok_or_else(|| { + error!( "Subject not found: {}; downstream subscriber specified a subject unknown to the upstream publisher", subject - ))?; - if let Some(key) = guard.subject_instance.remove(&subject) - && let Some(subjects) = guard.instance_subjects.get_mut(&key) - { - subjects.remove(&(StreamType::Request, subject.clone())); - if subjects.is_empty() { - guard.instance_subjects.remove(&key); - } - } - conn - }; + ) + })?; let RequestedSendConnection { context, @@ -771,22 +795,9 @@ async fn tcp_listener( mut reader: FramedRead, TwoPartCodec>, writer: FramedWrite, TwoPartCodec>, ) -> Result<()> { - let response_stream = { - let mut guard = state.lock().await; - let conn = guard - .rx_subjects - .remove(&subject) - .ok_or(error!("Subject not found: {}; upstream publisher specified a subject unknown to the downsteam subscriber", subject))?; - if let Some(key) = guard.subject_instance.remove(&subject) - && let Some(subjects) = guard.instance_subjects.get_mut(&key) - { - subjects.remove(&(StreamType::Response, subject.clone())); - if subjects.is_empty() { - guard.instance_subjects.remove(&key); - } - } - conn - }; + let response_stream = TcpStreamServer::take_response_stream(&state, &subject).ok_or_else(|| { + error!("Subject not found: {}; upstream publisher specified a subject unknown to the downsteam subscriber", subject) + })?; // unwrap response_stream let RequestedRecvConnection { @@ -1036,6 +1047,7 @@ mod tests { use crate::engine::AsyncEngineContextProvider; use crate::pipeline::Context; use crate::pipeline::network::DEFAULT_SEND_BUFFER_COUNT; + use crate::pipeline::network::tcp::client::TcpClient; use tokio::io::{AsyncWriteExt, ReadHalf, WriteHalf}; use tokio::net::TcpStream; @@ -1138,7 +1150,7 @@ mod tests { let _pending = server.register(options).await; - let state = server.state.lock().await; + let state = server.state.lock(); assert_eq!(state.tx_subjects.len(), 1, "one request stream registered"); assert_eq!(state.rx_subjects.len(), 1, "one response stream registered"); assert!( @@ -1439,7 +1451,7 @@ mod tests { // Verify it's in rx_subjects { - let state = server.state.lock().await; + let state = server.state.lock(); assert!(state.rx_subjects.contains_key(&subject)); } @@ -1451,7 +1463,7 @@ mod tests { // Verify it's been removed from rx_subjects { - let state = server.state.lock().await; + let state = server.state.lock(); assert!( !state.rx_subjects.contains_key(&subject), "RAII cleanup should have removed the rx_subjects entry" @@ -1486,7 +1498,7 @@ mod tests { // The entry should still be in rx_subjects (cleanup was disarmed) { - let state = server.state.lock().await; + let state = server.state.lock(); assert!( state.rx_subjects.contains_key(&subject), "into_parts() should disarm the RAII cleanup" @@ -1643,7 +1655,7 @@ mod tests { // Tombstone the identity. server.cancel_instance_streams(&id).await; { - let state = server.state.lock().await; + let state = server.state.lock(); assert!(state.removed_instances.contains_key(&id)); } @@ -1661,7 +1673,7 @@ mod tests { // The expired tombstone must have been pruned (lazy pruning fires on // every associate_instance/cancel_instance_streams call). { - let state = server.state.lock().await; + let state = server.state.lock(); assert!( !state.removed_instances.contains_key(&id), "expired tombstone should be pruned, not retained" @@ -1702,7 +1714,7 @@ mod tests { tokio::time::advance(TOMBSTONE_TTL + Duration::from_secs(1)).await; server.cancel_instance_streams(&id_new).await; - let state = server.state.lock().await; + let state = server.state.lock(); assert!( !state.removed_instances.contains_key(&id_old), "old tombstone should be pruned by the next cancel_instance_streams call" @@ -1728,7 +1740,7 @@ mod tests { server.cancel_instance_streams(&id_a).await; server.clear_instance_tombstone(&id_b).await; - let state = server.state.lock().await; + let state = server.state.lock(); assert!( state.removed_instances.contains_key(&id_a), "clearing a different identity must not remove id_a's tombstone" @@ -1957,6 +1969,63 @@ mod tests { } } + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_concurrent_response_registration_and_call_home() { + const STREAMS: usize = 128; + + let result = time::timeout(Duration::from_secs(20), async { + let server = test_server().await; + let mut pending_streams = Vec::with_capacity(STREAMS); + let mut client_tasks = Vec::with_capacity(STREAMS); + + for idx in 0..STREAMS { + let context = Context::new(()); + let options = StreamOptions::builder() + .context(context.context()) + .enable_request_stream(false) + .enable_response_stream(true) + .build() + .unwrap(); + + let pending = server.register(options).await; + let registered_stream = pending.recv_stream.unwrap(); + let (connection_info, stream_provider) = registered_stream.into_parts(); + let client_context = + Context::with_id_and_metadata((), context.id().to_string(), Default::default()); + let payload = Bytes::from(format!("payload-{idx}")); + + pending_streams.push((idx, payload.clone(), stream_provider)); + client_tasks.push(tokio::spawn(async move { + let mut sender = TcpClient::create_response_stream( + client_context.context(), + connection_info, + None, + ) + .await + .unwrap(); + sender.send_prologue(None).await.unwrap(); + sender.send(payload).await.unwrap(); + })); + } + + for task in client_tasks { + task.await.unwrap(); + } + + for (idx, expected, stream_provider) in pending_streams { + let mut stream = stream_provider.await.unwrap().unwrap(); + let actual = stream.rx.recv().await.unwrap(); + assert_eq!(actual, expected, "payload mismatch for stream {idx}"); + } + }) + .await; + + assert!( + result.is_ok(), + "concurrent response registration and call-home timed out" + ); + } + // ==================== request_stream_send_handler integration tests ==================== // // These exercise the closing-message contract of `request_stream_send_handler` From 3c061903cb13369f89359c5e68deed4a3a2ebcd9 Mon Sep 17 00:00:00 2001 From: Xiao Yang Date: Tue, 30 Jun 2026 23:52:59 +0800 Subject: [PATCH 005/320] fix(operator): continue startup when docker secret index refresh fails (#10789) Signed-off-by: yang.xiao Signed-off-by: Thomas Montfort <61255722+tmonty12@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: Thomas Montfort <61255722+tmonty12@users.noreply.github.com> Co-authored-by: Dr. Stefan Schimanski --- deploy/operator/cmd/main.go | 10 ++-- deploy/operator/internal/secrets/docker.go | 10 ++-- .../operator/internal/secrets/docker_test.go | 53 +++++++++++++++++++ 3 files changed, 64 insertions(+), 9 deletions(-) diff --git a/deploy/operator/cmd/main.go b/deploy/operator/cmd/main.go index 5df7f7bb0b23..506893e70bc5 100644 --- a/deploy/operator/cmd/main.go +++ b/deploy/operator/cmd/main.go @@ -572,10 +572,10 @@ func main() { os.Exit(1) } if err := dockerSecretRetriever.RefreshIndex(mainCtx); err != nil { - setupLog.Error(err, "initial docker secrets index refresh failed") - os.Exit(1) + setupLog.Error(err, "initial docker secrets index refresh completed with errors; continuing startup") + } else { + setupLog.Info("initial docker secrets index refreshed") } - setupLog.Info("initial docker secrets index refreshed") // launch a goroutine to refresh the docker secret indexer in any case every minute go func() { ticker := time.NewTicker(60 * time.Second) @@ -585,11 +585,9 @@ func main() { case <-mainCtx.Done(): return case <-ticker.C: - setupLog.Info("refreshing docker secrets index...") if err := dockerSecretRetriever.RefreshIndex(mainCtx); err != nil { - setupLog.Error(err, "unable to refresh docker secrets index") + setupLog.Error(err, "failed to refresh docker secrets index") } - setupLog.Info("docker secrets index refreshed") } } }() diff --git a/deploy/operator/internal/secrets/docker.go b/deploy/operator/internal/secrets/docker.go index 2e8846618d3e..35cdeda3d720 100644 --- a/deploy/operator/internal/secrets/docker.go +++ b/deploy/operator/internal/secrets/docker.go @@ -3,6 +3,7 @@ package secrets import ( "context" "encoding/json" + "errors" "fmt" "slices" "sync" @@ -51,6 +52,7 @@ func (i *DockerSecretIndexer) RefreshIndex(ctx context.Context) error { return 0 }) tmpSecrets := make(map[string]map[string][]string) + var refreshErrors []error for _, secret := range secrets.Items { if secret.Type == corev1.SecretTypeDockerConfigJson { // unmarshal the secret data @@ -58,7 +60,8 @@ func (i *DockerSecretIndexer) RefreshIndex(ctx context.Context) error { Auths map[string]any `json:"auths"` }{} if err := json.Unmarshal(secret.Data[corev1.DockerConfigJsonKey], dockerConfig); err != nil { - return fmt.Errorf("unable to unmarshal docker config json for secret %s: %w", secret.Name, err) + refreshErrors = append(refreshErrors, fmt.Errorf("unable to unmarshal docker config json for secret %s/%s: %w", secret.Namespace, secret.Name, err)) + continue } namespace := secret.Namespace if _, ok := tmpSecrets[namespace]; !ok { @@ -68,7 +71,8 @@ func (i *DockerSecretIndexer) RefreshIndex(ctx context.Context) error { // retrieve the registry host registry, err := common.GetHost(auth) if err != nil { - return fmt.Errorf("unable to get host for registry %s for secret %s: %w", auth, secret.Name, err) + refreshErrors = append(refreshErrors, fmt.Errorf("unable to get host for registry %q for secret %s/%s: %w", auth, secret.Namespace, secret.Name, err)) + continue } tmpSecrets[namespace][registry] = append(tmpSecrets[namespace][registry], secret.Name) } @@ -83,7 +87,7 @@ func (i *DockerSecretIndexer) RefreshIndex(ctx context.Context) error { i.mu.Lock() defer i.mu.Unlock() i.secrets = tmpSecrets - return nil + return errors.Join(refreshErrors...) } func (i *DockerSecretIndexer) listOptions() []client.ListOption { diff --git a/deploy/operator/internal/secrets/docker_test.go b/deploy/operator/internal/secrets/docker_test.go index 9c07a2f04cd1..1c4d6037ca36 100644 --- a/deploy/operator/internal/secrets/docker_test.go +++ b/deploy/operator/internal/secrets/docker_test.go @@ -146,6 +146,59 @@ func TestDockerSecretIndexer_DeterministicSecrets(t *testing.T) { } } +func TestDockerSecretIndexer_RefreshIndexSkipsMalformedSecrets(t *testing.T) { + mockSecrets := []corev1.Secret{ + { + ObjectMeta: metav1.ObjectMeta{ + Name: "good-secret", + Namespace: "default", + }, + Type: corev1.SecretTypeDockerConfigJson, + Data: map[string][]byte{ + ".dockerconfigjson": []byte(`{"auths":{"docker.io":{}}}`), + }, + }, + { + ObjectMeta: metav1.ObjectMeta{ + Name: "bad-secret", + Namespace: "default", + }, + Type: corev1.SecretTypeDockerConfigJson, + Data: map[string][]byte{ + ".dockerconfigjson": []byte(`not-json`), + }, + }, + { + ObjectMeta: metav1.ObjectMeta{ + Name: "invalid-registry-secret", + Namespace: "default", + }, + Type: corev1.SecretTypeDockerConfigJson, + Data: map[string][]byte{ + ".dockerconfigjson": []byte(`{"auths":{"pip install torch --index-url https:":{}}}`), + }, + }, + } + + fakeClient := fake.NewClientBuilder(). + WithScheme(scheme.Scheme). + WithObjects(&mockSecrets[0], &mockSecrets[1], &mockSecrets[2]). + Build() + + i := NewDockerSecretIndexer(fakeClient, "") + if err := i.RefreshIndex(t.Context()); err == nil { + t.Fatal("DockerSecretIndexer.RefreshIndex() error = nil, want errors for malformed secrets") + } + + secrets, err := i.GetSecrets("default", "docker.io") + if err != nil { + t.Fatalf("DockerSecretIndexer.GetSecrets() error = %v", err) + } + if got, want := secrets, []string{"good-secret"}; !slices.Equal(got, want) { + t.Fatalf("DockerSecretIndexer.GetSecrets() = %v, want %v", got, want) + } +} + func TestDockerSecretIndexer_RefreshIndexRespectsNamespaceScope(t *testing.T) { mockSecrets := []corev1.Secret{ { From 7c8b4b1928841db0dbe0ac53bf19e8afc4b86429 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Tue, 30 Jun 2026 10:19:38 -0700 Subject: [PATCH 006/320] perf(frontend): enable tokenizer cache by default (#11078) Signed-off-by: jthomson04 --- lib/llm/src/model_card.rs | 52 ++++++++++++++++----- lib/runtime/src/metrics/frontend_perf.rs | 4 +- lib/runtime/src/metrics/prometheus_names.rs | 4 +- 3 files changed, 44 insertions(+), 16 deletions(-) diff --git a/lib/llm/src/model_card.rs b/lib/llm/src/model_card.rs index a69e6a9ce71a..6ec04b7a2966 100644 --- a/lib/llm/src/model_card.rs +++ b/lib/llm/src/model_card.rs @@ -30,6 +30,18 @@ use tokenizers::Tokenizer as HfTokenizer; use crate::preprocessor::media::{MediaDecoder, MediaFetcher}; use crate::protocols::TokenIdType; +const DEFAULT_TOKENIZER_CACHE_BYTES: usize = 64 * 1024 * 1024; + +fn tokenizer_cache_enabled(value: Option<&str>) -> bool { + !matches!(value, Some("0")) +} + +fn tokenizer_cache_bytes(value: Option<&str>) -> usize { + value + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_TOKENIZER_CACHE_BYTES) +} + /// Identify model deployment cards in the key-value store pub const ROOT_PATH: &str = "v1/mdc"; @@ -1019,10 +1031,9 @@ impl ModelDeploymentCard { /// Tokenizer backend controls: /// - `runtime_config.tokenizer_backend=fastokens` — use `fastokens` as the encoding backend /// - `DYN_TOKENIZER=fastokens` — fallback backend for callers without explicit runtime config - /// - `DYN_TOKENIZER_CACHE=1` — wrap the tokenizer in an L1 prefix cache that records - /// tokenizations at special-token boundaries (massive speed-up for shared chat - /// prefixes; default off, zero cost when unset) - /// - `DYN_TOKENIZER_CACHE_BYTES=` — L1 cache byte budget (default 50 MB) + /// - `DYN_TOKENIZER_CACHE=0` — disable the L1 prefix cache that records tokenizations + /// at special-token boundaries (enabled by default; any other value keeps it enabled) + /// - `DYN_TOKENIZER_CACHE_BYTES=` — L1 cache byte budget (default 64 MiB) /// - `DYN_TOKENIZER_CACHE_EXTEND=0` — disable partial-hit extension. By default /// (when the cache is enabled) a partial hit also caches the new suffix so each /// turn of a growing multi-turn conversation hits deeper than the last, keeping @@ -1034,14 +1045,10 @@ impl ModelDeploymentCard { .effective_tokenizer_backend() .is_fastokens(); - let cache_enabled = matches!( - std::env::var("DYN_TOKENIZER_CACHE").ok().as_deref(), - Some("1") - ); - let cache_bytes = std::env::var("DYN_TOKENIZER_CACHE_BYTES") - .ok() - .and_then(|s| s.parse::().ok()) - .unwrap_or(50 * 1024 * 1024); + let cache_enabled = + tokenizer_cache_enabled(std::env::var("DYN_TOKENIZER_CACHE").ok().as_deref()); + let cache_bytes = + tokenizer_cache_bytes(std::env::var("DYN_TOKENIZER_CACHE_BYTES").ok().as_deref()); // Partial-hit extension is on by default; disable with DYN_TOKENIZER_CACHE_EXTEND=0. let cache_extend = !matches!( std::env::var("DYN_TOKENIZER_CACHE_EXTEND").ok().as_deref(), @@ -2076,6 +2083,27 @@ mod tests { use std::collections::HashSet; use std::path::{Path, PathBuf}; + #[test] + fn tokenizer_cache_is_enabled_by_default_and_disabled_only_by_zero() { + assert!(super::tokenizer_cache_enabled(None)); + assert!(super::tokenizer_cache_enabled(Some("1"))); + assert!(!super::tokenizer_cache_enabled(Some("0"))); + assert!(super::tokenizer_cache_enabled(Some("true"))); + } + + #[test] + fn tokenizer_cache_bytes_defaults_to_64_mib_and_accepts_valid_overrides() { + assert_eq!( + super::tokenizer_cache_bytes(None), + super::DEFAULT_TOKENIZER_CACHE_BYTES + ); + assert_eq!(super::tokenizer_cache_bytes(Some("1024")), 1024); + assert_eq!( + super::tokenizer_cache_bytes(Some("invalid")), + super::DEFAULT_TOKENIZER_CACHE_BYTES + ); + } + #[test] pub fn test_config_json_llama3() -> anyhow::Result<()> { let config_file = Path::new(env!("CARGO_MANIFEST_DIR")) diff --git a/lib/runtime/src/metrics/frontend_perf.rs b/lib/runtime/src/metrics/frontend_perf.rs index f5a19d370886..2836e2c90895 100644 --- a/lib/runtime/src/metrics/frontend_perf.rs +++ b/lib/runtime/src/metrics/frontend_perf.rs @@ -120,7 +120,7 @@ pub static DETOKENIZE_TOKEN_COUNT: Lazy = Lazy::new(|| { .expect("detokenize_token_count counter") }); -/// Cumulative L1 tokenizer cache hits. Only nonzero when `DYN_TOKENIZER_CACHE=1`. +/// Cumulative L1 tokenizer cache hits. The cache is enabled unless `DYN_TOKENIZER_CACHE=0`. pub static TOKENIZER_CACHE_HITS_TOTAL: Lazy = Lazy::new(|| { Counter::with_opts(Opts::new( frontend_metric_name(frontend_perf::TOKENIZER_CACHE_HITS_TOTAL), @@ -129,7 +129,7 @@ pub static TOKENIZER_CACHE_HITS_TOTAL: Lazy = Lazy::new(|| { .expect("tokenizer_cache_hits_total counter") }); -/// Cumulative L1 tokenizer cache misses. Only nonzero when `DYN_TOKENIZER_CACHE=1`. +/// Cumulative L1 tokenizer cache misses. The cache is enabled unless `DYN_TOKENIZER_CACHE=0`. pub static TOKENIZER_CACHE_MISSES_TOTAL: Lazy = Lazy::new(|| { Counter::with_opts(Opts::new( frontend_metric_name(frontend_perf::TOKENIZER_CACHE_MISSES_TOTAL), diff --git a/lib/runtime/src/metrics/prometheus_names.rs b/lib/runtime/src/metrics/prometheus_names.rs index 5fd65dba1d47..8e9e5ebd35f3 100644 --- a/lib/runtime/src/metrics/prometheus_names.rs +++ b/lib/runtime/src/metrics/prometheus_names.rs @@ -642,9 +642,9 @@ pub mod frontend_perf { pub const TOKENIZE_SECONDS: &str = "tokenize_seconds"; /// Template application time in preprocessor pub const TEMPLATE_SECONDS: &str = "template_seconds"; - /// L1 tokenizer cache hits (cumulative); only incremented when DYN_TOKENIZER_CACHE is enabled + /// L1 tokenizer cache hits (cumulative); enabled unless DYN_TOKENIZER_CACHE=0 pub const TOKENIZER_CACHE_HITS_TOTAL: &str = "tokenizer_cache_hits_total"; - /// L1 tokenizer cache misses (cumulative); only incremented when DYN_TOKENIZER_CACHE is enabled + /// L1 tokenizer cache misses (cumulative); enabled unless DYN_TOKENIZER_CACHE=0 pub const TOKENIZER_CACHE_MISSES_TOTAL: &str = "tokenizer_cache_misses_total"; /// Cumulative detokenization time (microseconds); pair with DETOKENIZE_TOKEN_COUNT pub const DETOKENIZE_TOTAL_US: &str = "detokenize_total_us"; From 0b973594e2cac8dfcf826401333a650349abea0c Mon Sep 17 00:00:00 2001 From: Tzu-Ling Kan Date: Tue, 30 Jun 2026 12:48:39 -0500 Subject: [PATCH 007/320] fix(vllm): pool embedding worker output via PoolingParams(task="embed") (#10248) Signed-off-by: Tzu-Ling Co-authored-by: Claude Opus 4.8 (1M context) --- components/src/dynamo/vllm/handlers.py | 66 +++++++++---- .../vllm/tests/test_vllm_worker_handler.py | 93 +++++++++++++++++++ examples/backends/vllm/launch/agg_embed.sh | 10 ++ tests/serve/test_vllm.py | 8 +- 4 files changed, 157 insertions(+), 20 deletions(-) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 049689310a93..c26178cb54c8 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -30,6 +30,7 @@ ) import torch +from vllm import PoolingParams from vllm.config import ModelConfig, VllmConfig from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt from vllm.lora.request import LoRARequest @@ -3874,18 +3875,15 @@ async def generate( The Rust frontend forwards the request dict directly. Expected keys: ``model: str``, ``input: str | list[str] | list[int] | list[list[int]]``. - Optional ``dimensions`` (Matryoshka truncation; first N floats of each - embedding). Optional ``encoding_format`` (``"float"`` -- default -- - or ``"base64"``); when ``"base64"`` is requested, each per-input - vector is serialized as a base64-encoded string of little-endian - ``f32`` bytes per the OpenAI spec, applied after any - ``dimensions`` truncation so the byte count matches the requested + Optional ``dimensions`` (Matryoshka dimensionality reduction): + forwarded to vLLM's pooler, which truncates to N dims and + re-normalizes; vLLM requires the model to declare Matryoshka support. + Optional ``encoding_format`` (``"float"`` -- default -- or + ``"base64"``); when ``"base64"`` is requested, each per-input vector is + serialized as a base64-encoded string of little-endian ``f32`` bytes + per the OpenAI spec, so the byte count matches the (possibly reduced) dimensionality. """ - # Lazy import to avoid pulling PoolingParams into handlers.py at module - # load time for non-embedding workers. - from vllm import PoolingParams - model_name = request.get("model") or self.config.served_model_name or "" input_field = request.get("input") if input_field is None: @@ -3916,7 +3914,29 @@ async def generate( "expected 'float' or 'base64'" ) - pooling_params = PoolingParams() + # Request the pooled sentence embedding. With no task, vLLM's + # encode() resolves to per-token output (the full ``n_tokens x + # hidden`` hidden-state matrix), so the OpenAI ``/v1/embeddings`` + # response ends up with the wrong shape (dim scales with input + # length) instead of one vector per input. ``task="embed"`` selects + # the pooled embedding and runs the model's configured pooler + # (normalization included for models like Qwen3-Embedding), matching + # vLLM's own embedding server. ``use_activation`` is intentionally + # left at the pooler default so per-model behaviour isn't overridden. + # + # ``dimensions`` (OpenAI Matryoshka truncation) is forwarded to vLLM + # rather than applied here: vLLM's pooler truncates to ``dimensions`` + # and then re-normalizes (the correct MRL behaviour) and validates + # that the model actually supports Matryoshka -- raising rather than + # silently returning a degraded, un-normalized vector for models that + # don't. This matches bare ``vllm serve``. Models whose HF config + # doesn't declare Matryoshka support (e.g. Qwen3-Embedding) must be + # launched with ``--hf-overrides '{"is_matryoshka": true}'`` for + # ``dimensions`` requests to be accepted. + pooling_kwargs: dict[str, Any] = {"task": "embed"} + if dimensions is not None: + pooling_kwargs["dimensions"] = dimensions + pooling_params = PoolingParams(**pooling_kwargs) # Use the per-request context id (same as the chat/completion paths # in this file) so concurrent embeddings never collide inside # ``AsyncLLM``. ``context.trace_id`` is a distributed-trace id @@ -3969,14 +3989,24 @@ async def _encode_one(idx: int, prompt: Any): embedding_objects: list[Dict[str, Any]] = [] prompt_tokens = 0 for idx, final_output in enumerate(outputs): + # vLLM has already applied any ``dimensions`` Matryoshka reduction + # (truncate + re-normalize) inside the pooler, so this is the + # final per-input vector -- no post-hoc truncation here. embedding = _pooling_output_to_list(final_output.outputs.data) - if dimensions is not None: - if dimensions > len(embedding): - raise ValueError( - f"dimensions={dimensions} exceeds model embedding " - f"dimension {len(embedding)}" - ) - embedding = embedding[:dimensions] + + # vLLM rejects an unsupported ``dimensions`` for models that + # declare a ``matryoshka_dimensions`` list, but a model enabled + # via ``--hf-overrides '{"is_matryoshka": true}'`` (no explicit + # list) is only validated for ``dimensions >= 1`` -- the pooler + # then silently clamps an oversized request to the model's native + # size (``embeddings[..., :dimensions]``). Surface the same clear + # error the old post-hoc path raised instead of returning a + # shorter-than-requested vector. + if dimensions is not None and len(embedding) < dimensions: + raise ValueError( + f"dimensions={dimensions} exceeds model embedding " + f"dimension {len(embedding)}" + ) # Always emit base64 over the worker->frontend wire format. The # Rust frontend decodes back to float when the client's diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py index 243f8b1c5b04..6eeffe38f89e 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py @@ -1419,6 +1419,99 @@ async def fake_encode(prompt, pooling_params, request_id): # cancel-and-await pass must not have touched the engine. assert aborted == [] + @pytest.mark.asyncio + @pytest.mark.timeout(5) + async def test_dimensions_forwarded_to_pooling_params(self): + """``dimensions`` is forwarded to vLLM via ``PoolingParams`` rather + than applied as post-hoc truncation in the handler. + + vLLM's pooler then does the Matryoshka reduction (truncate + + re-normalize) and validates that the model supports it. The handler + must NOT slice the returned vector itself, so it emits exactly what + the engine produced. + """ + handler = self._make_embedding_handler() + context = self._make_context() + captured: dict = {} + # vLLM's pooler has already reduced to the requested ``dimensions``, so + # the stub returns a 128-dim vector (not 3) -- otherwise the handler's + # oversized-dimensions guard would (correctly) reject it. + vec = [i * 0.01 for i in range(128)] + + async def fake_encode(prompt, pooling_params, request_id): + captured["pooling_params"] = pooling_params + output = MagicMock() + output.outputs.data = torch.tensor(vec) + output.prompt_token_ids = [1, 2, 3] + yield output + + handler.engine_client.encode = fake_encode + + request = {"input": ["hello"], "model": "test-model", "dimensions": 128} + responses = [r async for r in handler.generate(request, context)] + + pp = captured["pooling_params"] + assert pp.task == "embed" + assert pp.dimensions == 128 + # No post-hoc truncation: the handler returns exactly the vector vLLM + # produced (the 128-float stub here), trusting the pooler to have + # already applied the dimensionality reduction. + expected_b64 = mod._encode_floats_to_base64(vec) + assert responses[0]["data"][0]["embedding"] == expected_b64 + + @pytest.mark.asyncio + @pytest.mark.timeout(5) + async def test_no_dimensions_omits_pooling_dimensions(self): + """Without ``dimensions`` the handler requests ``task="embed"`` only, + leaving ``PoolingParams.dimensions`` unset so vLLM returns the model's + native embedding size. + """ + handler = self._make_embedding_handler() + context = self._make_context() + captured: dict = {} + + async def fake_encode(prompt, pooling_params, request_id): + captured["pooling_params"] = pooling_params + output = MagicMock() + output.outputs.data = torch.tensor([0.1, 0.2, 0.3]) + output.prompt_token_ids = [1, 2, 3] + yield output + + handler.engine_client.encode = fake_encode + + request = {"input": ["hello"], "model": "test-model"} + _ = [r async for r in handler.generate(request, context)] + + pp = captured["pooling_params"] + assert pp.task == "embed" + assert pp.dimensions is None + + @pytest.mark.asyncio + @pytest.mark.timeout(5) + async def test_oversized_dimensions_raises(self): + """When vLLM silently clamps an oversized ``dimensions`` request (a + model enabled via ``--hf-overrides '{"is_matryoshka": true}'`` with no + ``matryoshka_dimensions`` list), the handler raises a clear error + instead of returning a shorter-than-requested vector. + """ + handler = self._make_embedding_handler() + context = self._make_context() + + async def fake_encode(prompt, pooling_params, request_id): + output = MagicMock() + # vLLM clamped to the model's native size (3 dims here) even though + # 2048 was requested. + output.outputs.data = torch.tensor([0.1, 0.2, 0.3]) + output.prompt_token_ids = [1, 2, 3] + yield output + + handler.engine_client.encode = fake_encode + + request = {"input": ["hello"], "model": "test-model", "dimensions": 2048} + with pytest.raises(ValueError, match="exceeds model embedding dimension"): + async for _ in handler.generate(request, context): + pass + class TestPadMmHashesTo64: """The frontend forwards canonical 16-char hex mm_hashes; vLLM must pad diff --git a/examples/backends/vllm/launch/agg_embed.sh b/examples/backends/vllm/launch/agg_embed.sh index 5224a2ec9ae3..fee4be87ee4a 100755 --- a/examples/backends/vllm/launch/agg_embed.sh +++ b/examples/backends/vllm/launch/agg_embed.sh @@ -69,6 +69,15 @@ python3 -m dynamo.frontend & # unusually long embedding inputs. MAX_MODEL_LEN="${MAX_MODEL_LEN:-2048}" +# Qwen3-Embedding supports Matryoshka (flexible output dims 32-1024) but its +# HF config does not declare it, so vLLM rejects OpenAI `dimensions` requests +# unless told the model is Matryoshka. Inject the flag only for the default +# model; other models must declare their own support (or omit `dimensions`). +HF_OVERRIDES_ARGS=() +if [[ "$MODEL" == "Qwen/Qwen3-Embedding-0.6B" ]]; then + HF_OVERRIDES_ARGS=(--hf-overrides '{"is_matryoshka": true}') +fi + # run worker # --runner pooling: required for embedding models. # --pooler-config: MEAN pool, no activation — the Qwen3-Embedding default. @@ -84,6 +93,7 @@ DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ --max-model-len "$MAX_MODEL_LEN" \ --no-enable-prefix-caching \ --trust-remote-code \ + "${HF_OVERRIDES_ARGS[@]}" \ $GPU_MEM_ARGS \ "${EXTRA_ARGS[@]}" & diff --git a/tests/serve/test_vllm.py b/tests/serve/test_vllm.py index 64b4fcd7346c..673ca8af2e02 100644 --- a/tests/serve/test_vllm.py +++ b/tests/serve/test_vllm.py @@ -688,8 +688,12 @@ class VLLMConfig(EngineConfig): repeat_count=1, expected_response=["Generated 3 embeddings with dimension"], ), - # `dimensions` truncation (Matryoshka). Qwen3-Embedding-0.6B has a - # hidden dim of 1024, so the truncated vector should be exactly 128. + # `dimensions` reduction (Matryoshka). Qwen3-Embedding-0.6B has a + # hidden dim of 1024, so the reduced vector should be exactly 128. + # The worker forwards `dimensions` to vLLM's pooler (truncate + + # re-normalize); `agg_embed.sh` launches this model with + # `--hf-overrides '{"is_matryoshka": true}'` so vLLM accepts the + # request (Qwen3-Embedding's config doesn't declare Matryoshka). # Built inline because the `embedding_payload()` helper doesn't # expose an `extra_body` kwarg yet. EmbeddingPayload( From 938e4467fad2b592400a449e11175a94e01a6829 Mon Sep 17 00:00:00 2001 From: JinYan Su Date: Wed, 1 Jul 2026 02:22:50 +0800 Subject: [PATCH 008/320] feat(kv-hashing): expose compute_salt_hash for producers without a Request (#10652) Signed-off-by: xiaguan <751080330@qq.com> Co-authored-by: Ryan McCormick --- lib/kv-hashing/src/lib.rs | 1 + lib/kv-hashing/src/salt.rs | 8 ++++++- lib/kv-hashing/tests/request_hashing.rs | 31 ++++++++++++++++++++++++- 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/lib/kv-hashing/src/lib.rs b/lib/kv-hashing/src/lib.rs index f9ab655a8517..674b3cf4b0f8 100644 --- a/lib/kv-hashing/src/lib.rs +++ b/lib/kv-hashing/src/lib.rs @@ -25,6 +25,7 @@ mod salt; pub use block::UniversalBlock; pub use error::KvHashingError; pub use request::{Request, RequestBuilder, RequestMmObjectInfo}; +pub use salt::compute_salt_hash; // Re-export the underlying primitives so consumers can depend solely on this crate. pub use dynamo_tokens::{ diff --git a/lib/kv-hashing/src/salt.rs b/lib/kv-hashing/src/salt.rs index 8e9ff141af4b..f409e1995b14 100644 --- a/lib/kv-hashing/src/salt.rs +++ b/lib/kv-hashing/src/salt.rs @@ -39,7 +39,13 @@ use crate::error::KvHashingError; /// `lib/kv-router/src/protocols.rs:79` (`options.lora_name.filter(|n| !n.is_empty())`) /// so a client that sends `lora_name = ""` shares the cache with a client that sends /// `lora_name = None`. -pub(crate) fn compute_salt_hash( +/// +/// Prefer [`crate::Request::salt_hash`] when you already hold a [`crate::Request`]. +/// This free function is the seam for producers that drive a +/// `dynamo_tokens::TokenBlockSequence` directly (e.g. incremental block formation +/// during decode) and need the salt without constructing a throwaway `Request` +/// around tokens they own elsewhere. +pub fn compute_salt_hash( salt: Option<&str>, lora_name: Option<&str>, ) -> Result { diff --git a/lib/kv-hashing/tests/request_hashing.rs b/lib/kv-hashing/tests/request_hashing.rs index b44427d69088..978c3997d2b6 100644 --- a/lib/kv-hashing/tests/request_hashing.rs +++ b/lib/kv-hashing/tests/request_hashing.rs @@ -5,7 +5,7 @@ use dynamo_kv_hashing::{ KvHashingError, Request, RequestMmObjectInfo, SaltHash, Token, TokenBlockMmInfo, - compute_block_hash, + compute_block_hash, compute_salt_hash, }; use dynamo_tokens::{TokenBlockSequence, Tokens}; @@ -444,3 +444,32 @@ fn consuming_sequence_hashes_match_borrowed_path() { assert_eq!(consuming, borrowed); } + +// ----------------------------------------------------------------------------- +// compute_salt_hash is the public seam for producers without a Request +// ----------------------------------------------------------------------------- +#[test] +fn compute_salt_hash_matches_request_driven_hashing() { + let tokens: Vec = (1..=12).collect(); + let no_mm: &[TokenBlockMmInfo] = &[]; + + for lora in [None, Some("lora-a")] { + let request = req(tokens.clone(), lora, None, vec![]); + let via_request = request.positional_lineage_hashes(BS).unwrap(); + + // A producer driving TokenBlockSequence directly (no Request) must land + // on the same chain when seeded with compute_salt_hash. + let salt = compute_salt_hash(None, lora).unwrap(); + assert_eq!(salt, request.salt_hash().unwrap()); + + let seq = + TokenBlockSequence::new_with_mm(Tokens::from(tokens.clone()), no_mm, BS, Some(salt)) + .unwrap(); + let via_sequence: Vec<_> = seq + .blocks() + .iter() + .map(|b| b.positional_lineage_hash()) + .collect(); + assert_eq!(via_sequence, via_request, "lora={lora:?}"); + } +} From dac7dfb2ec7f502b76501fe2dd99f51e841806a6 Mon Sep 17 00:00:00 2001 From: Karen Chung Date: Tue, 30 Jun 2026 12:25:59 -0700 Subject: [PATCH 009/320] chore(deps): bump to vLLM 0.24.0 (#11076) Signed-off-by: Karen Chung Co-authored-by: Dmitry Tokarev --- Cargo.lock | 10 +- components/src/dynamo/frontend/prepost.py | 65 ++++++-- .../tests/test_vllm_processor_unit.py | 4 +- .../src/dynamo/vllm/instrumented_scheduler.py | 12 +- .../vllm/tests/test_vllm_kv_events_api.py | 140 +++++------------- container/context.yaml | 4 +- container/deps/vllm/protected_packages.txt | 2 +- docs/reference/support-matrix.md | 2 +- lib/vllm-rs-backend/Cargo.toml | 8 +- lib/vllm-rs-backend/src/backend.rs | 24 ++- lib/vllm-rs-backend/src/control.rs | 13 +- lib/vllm-rs-backend/src/convert.rs | 6 + pyproject.toml | 2 +- tests/frontend/test_prepost.py | 90 +++++++++-- tests/report_pytest_markers.py | 4 +- 15 files changed, 229 insertions(+), 157 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7c2eb88d945c..2b42f8a7c7fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10432,7 +10432,7 @@ checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" [[package]] name = "vllm-engine-core-client" version = "0.1.0" -source = "git+https://github.com/vllm-project/vllm.git?tag=v0.22.0#0b3ba88f165976e77ca5e6a7a3f5bba4562b80af" +source = "git+https://github.com/vllm-project/vllm.git?tag=v0.24.0#ee0da84ab9e04ac7610e28580af62c365e898389" dependencies = [ "arc-swap", "bytemuck", @@ -10466,11 +10466,12 @@ dependencies = [ [[package]] name = "vllm-llm" version = "0.1.0" -source = "git+https://github.com/vllm-project/vllm.git?tag=v0.22.0#0b3ba88f165976e77ca5e6a7a3f5bba4562b80af" +source = "git+https://github.com/vllm-project/vllm.git?tag=v0.24.0#ee0da84ab9e04ac7610e28580af62c365e898389" dependencies = [ "easy-ext", "enum-as-inner", "futures", + "parking_lot", "serde", "serde_json", "thiserror 2.0.18", @@ -10485,7 +10486,7 @@ dependencies = [ [[package]] name = "vllm-managed-engine" version = "0.1.0" -source = "git+https://github.com/vllm-project/vllm.git?tag=v0.22.0#0b3ba88f165976e77ca5e6a7a3f5bba4562b80af" +source = "git+https://github.com/vllm-project/vllm.git?tag=v0.24.0#ee0da84ab9e04ac7610e28580af62c365e898389" dependencies = [ "anyhow", "clap", @@ -10497,8 +10498,9 @@ dependencies = [ [[package]] name = "vllm-metrics" version = "0.1.0" -source = "git+https://github.com/vllm-project/vllm.git?tag=v0.22.0#0b3ba88f165976e77ca5e6a7a3f5bba4562b80af" +source = "git+https://github.com/vllm-project/vllm.git?tag=v0.24.0#ee0da84ab9e04ac7610e28580af62c365e898389" dependencies = [ + "itertools 0.14.0", "prometheus-client", ] diff --git a/components/src/dynamo/frontend/prepost.py b/components/src/dynamo/frontend/prepost.py index 7fa52f345ac3..003813e99ce6 100644 --- a/components/src/dynamo/frontend/prepost.py +++ b/components/src/dynamo/frontend/prepost.py @@ -4,7 +4,8 @@ from __future__ import annotations import os -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Sequence +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from typing import Any, Protocol @@ -20,7 +21,7 @@ from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tool_parsers import ToolParser -from vllm.utils.async_utils import AsyncMicrobatchTokenizer +from vllm.utils.async_utils import make_async class _Renderer(Protocol): @@ -41,15 +42,17 @@ class PreprocessResult: prompt_token_ids: list[int] -_ASYNC_TOKENIZER_POOL: dict[int, AsyncMicrobatchTokenizer] = {} +_ASYNC_TOKENIZER_POOL: dict[int, Callable[..., Awaitable[Any]]] = {} SKIP_REQUEST_VALIDATION = os.getenv("DYN_VLLM_SKIP_REQUEST_VALIDATION", "1") == "1" -def _get_async_tokenizer(tokenizer: TokenizerLike) -> AsyncMicrobatchTokenizer: +def _get_async_tokenizer(tokenizer: TokenizerLike) -> Callable[..., Awaitable[Any]]: key = id(tokenizer) async_tokenizer = _ASYNC_TOKENIZER_POOL.get(key) if async_tokenizer is None: - async_tokenizer = AsyncMicrobatchTokenizer(tokenizer) + async_tokenizer = make_async( + tokenizer, executor=ThreadPoolExecutor(max_workers=1) + ) _ASYNC_TOKENIZER_POOL[key] = async_tokenizer return async_tokenizer @@ -320,6 +323,44 @@ def _should_parse_tools(self) -> bool: and self.request_for_sampling.tool_choice != "none" ) + def _tool_parser_terminal_markers(self, names: tuple[str, ...]) -> tuple[str, ...]: + parser_engine = getattr(self.tool_parser, "_parser_engine", None) + parser_engine_config = getattr(parser_engine, "parser_engine_config", None) + terminals = getattr(parser_engine_config, "terminals", None) + if not isinstance(terminals, dict): + return () + + markers: list[str] = [] + for name in names: + marker = terminals.get(name) + if isinstance(marker, str) and marker: + markers.append(marker) + return tuple(markers) + + def _tool_start_markers(self) -> tuple[str, ...]: + markers = [ + getattr(self.tool_parser, "tool_call_start_token", None), + # MistralToolParser names its [TOOL_CALLS] marker bot_token. + getattr(self.tool_parser, "bot_token", None), + *self._tool_parser_terminal_markers(("TOOL_START", "FUNC_PREFIX")), + ] + return tuple( + dict.fromkeys( + marker for marker in markers if isinstance(marker, str) and marker + ) + ) + + def _tool_end_markers(self) -> tuple[str, ...]: + markers = [ + getattr(self.tool_parser, "tool_call_end_token", None), + *self._tool_parser_terminal_markers(("TOOL_END", "FUNC_END")), + ] + return tuple( + dict.fromkeys( + marker for marker in markers if isinstance(marker, str) and marker + ) + ) + @staticmethod def _compose_delta_message( reasoning: str | None, content: str | None @@ -507,9 +548,11 @@ def process_output(self, output: Any) -> dict[str, Any] | None: # ------------------------------------------------------------------ if self._tool_text_buffer is not None: self._tool_text_buffer += delta_text - tool_call_end = getattr(self.tool_parser, "tool_call_end_token", None) buffer_complete = ( - tool_call_end and tool_call_end in self._tool_text_buffer + any( + marker in self._tool_text_buffer + for marker in self._tool_end_markers() + ) ) or output.finish_reason if buffer_complete: buffered_text = self._tool_text_buffer @@ -548,10 +591,10 @@ def process_output(self, output: Any) -> dict[str, Any] | None: current_text = "" current_token_ids = [] - tool_call_start = getattr( - self.tool_parser, "tool_call_start_token", None - ) - if post_content and tool_call_start and tool_call_start in post_content: + tool_start_markers = self._tool_start_markers() + if post_content and any( + marker in post_content for marker in tool_start_markers + ): # Tool call markup present — buffer for non-streaming # extraction (streaming parser can't handle the combined # reasoning-end + tool-start in a single chunk). diff --git a/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py b/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py index df32d6966854..84db2b94e85f 100644 --- a/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py @@ -13,7 +13,7 @@ import pytest from _routed_engine_fakes import FakeRoutedEngine as _FakeRoutedEngine from transformers import AutoTokenizer -from vllm.tool_parsers.qwen3coder_tool_parser import Qwen3CoderToolParser +from vllm.tool_parsers.qwen3_engine_tool_parser import Qwen3EngineToolParser from dynamo.frontend.prepost import _prepare_request @@ -409,7 +409,7 @@ def test_qwen3_coder_coerces_object_typed_arg(self, tokenizer): request_for_sampling, parser, _, _, _ = _prepare_request( OBJECT_TYPED_TOOL_REQUEST, tokenizer=tokenizer, - tool_parser_class=Qwen3CoderToolParser, + tool_parser_class=Qwen3EngineToolParser, ) assert parser is not None, "Expected _prepare_request to construct the parser" diff --git a/components/src/dynamo/vllm/instrumented_scheduler.py b/components/src/dynamo/vllm/instrumented_scheduler.py index 5f9f3bf59a6e..0af4bbf9715c 100644 --- a/components/src/dynamo/vllm/instrumented_scheduler.py +++ b/components/src/dynamo/vllm/instrumented_scheduler.py @@ -324,7 +324,7 @@ def has_requests(self) -> bool: return True return super().has_requests() - def schedule(self) -> SchedulerOutput: + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if self._bench_active and self._bench_phase != _BenchPhase.IDLE: try: output = self._bench_step() @@ -333,7 +333,7 @@ def schedule(self) -> SchedulerOutput: self._bench_cleanup_requests() self._bench_active = False self._bench_phase = _BenchPhase.IDLE - return self._schedule_and_record_time() + return self._schedule_and_record_time(throttle_prefills) if output is not None: self.kv_cache_manager.new_step_starts() self._update_after_schedule(output) @@ -376,10 +376,12 @@ def schedule(self) -> SchedulerOutput: self._update_after_schedule(empty) return empty - return self._schedule_and_record_time() + return self._schedule_and_record_time(throttle_prefills) - def _schedule_and_record_time(self) -> SchedulerOutput: - output = super().schedule() + def _schedule_and_record_time( + self, throttle_prefills: bool = False + ) -> SchedulerOutput: + output = super().schedule(throttle_prefills) if output.total_num_scheduled_tokens > 0: self._schedule_times.append(time.monotonic()) return output diff --git a/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py b/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py index e3f82eb168cc..87f4564bd4d0 100755 --- a/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py +++ b/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py @@ -11,12 +11,14 @@ This test is the early warning for vLLM KV-event wire-format changes. -In the normal case, if this fails, update `lib/kv-router/src/zmq_wire.rs` to -match the new upstream vLLM event shape, then update this test. +In the normal case, if this fails, first check whether +`lib/kv-router/src/zmq_wire.rs` already accepts the new upstream vLLM event +shape. If not, update that compatibility layer before updating this test. That file is Dynamo's compatibility layer for vLLM KV events: -- it decodes vLLM's msgpack `array_like=True` wire format -- it handles field order changes in `BlockStored` / `BlockRemoved` / `EventBatch` +- it decodes vLLM's msgpack tagged-map wire format and legacy + `array_like=True` payloads +- it handles field changes in `BlockStored` / `BlockRemoved` / `EventBatch` - it translates upstream `extra_keys` into Dynamo's internal `block_mm_infos` Only touch consolidator files if we explicitly need the consolidator publisher @@ -104,7 +106,7 @@ def test_block_stored_fields(self): f"Actual: {actual_fields}\n" f"Required follow-up:\n" f" - Update lib/kv-router/src/zmq_wire.rs to match the new BlockStored wire format.\n" - f" - Update this test's expected_fields and msgpack position checks.\n" + f" - Update this test's expected_fields and msgpack shape checks.\n" f" - If needed, add or update a regression test in lib/llm/src/kv_router/publisher.rs." ) @@ -151,24 +153,19 @@ def test_event_batch_fields(self): f" - Update this test's expected_fields." ) - def test_kv_cache_event_uses_array_like(self): - """Verify KVCacheEvent uses array_like=True serialization. - - Our Rust deserializers expect msgpack arrays, not objects. - If this changes, deserialization will break. - """ - # msgspec structs with array_like=True have this attribute + def test_kv_cache_event_uses_tagged_map(self): + """Verify KVCacheEvent uses tagged-map serialization.""" struct_config = getattr(KVCacheEvent, "__struct_config__", None) assert struct_config is not None, "KVCacheEvent is not a msgspec Struct" - assert struct_config.array_like is True, ( - "KVCacheEvent no longer uses array_like=True! " - "This will break Rust deserialization." + assert struct_config.array_like is False, ( + "KVCacheEvent changed away from tagged-map serialization. " + "Check lib/kv-router/src/zmq_wire/deserialize.rs compatibility." ) def test_kv_cache_event_uses_tag(self): """Verify KVCacheEvent uses tag=True for variant identification. - The tag (e.g., 'BlockStored') is the first element in the msgpack array. + The tag is encoded in the msgpack map's 'type' field. """ struct_config = getattr(KVCacheEvent, "__struct_config__", None) assert struct_config is not None, "KVCacheEvent is not a msgspec Struct" @@ -180,7 +177,7 @@ def test_kv_cache_event_uses_tag(self): ) def test_block_stored_serialization_format(self): - """Verify BlockStored serializes to expected msgpack array format. + """Verify BlockStored serializes to expected msgpack map format. This is the ultimate test - if the serialized format changes, Rust deserialization will fail. @@ -208,48 +205,22 @@ def test_block_stored_serialization_format(self): encoded = msgspec.msgpack.encode(event) decoded = msgspec.msgpack.decode(encoded) - # Should be an array with tag as first element - assert isinstance(decoded, list), f"Expected list, got {type(decoded)}" - assert ( - decoded[0] == "BlockStored" - ), f"Expected tag 'BlockStored', got {decoded[0]}" - - expected_len = ( - 9 - + int(_has_group_idx(BlockStored)) - + int(_has_kv_cache_spec_kind(BlockStored)) - + int(_has_kv_cache_spec_sliding_window(BlockStored)) - ) - assert len(decoded) == expected_len, ( - f"Expected {expected_len} elements, got {len(decoded)}.\n" - f"Decoded: {decoded}\n" - f"If field count changed, update Rust deserializers." - ) - - # Verify field positions - assert decoded[1] == [123, 456], f"block_hashes at wrong position: {decoded[1]}" - assert decoded[2] == 789, f"parent_block_hash at wrong position: {decoded[2]}" - assert decoded[3] == [1, 2, 3, 4], f"token_ids at wrong position: {decoded[3]}" - assert decoded[4] == 16, f"block_size at wrong position: {decoded[4]}" - assert decoded[5] is None, f"lora_id at wrong position: {decoded[5]}" - assert decoded[6] == "GPU", f"medium at wrong position: {decoded[6]}" - assert decoded[7] is None, f"lora_name at wrong position: {decoded[7]}" - assert decoded[8] is None, f"extra_keys at wrong position: {decoded[8]}" - next_idx = 9 + assert isinstance(decoded, dict), f"Expected dict, got {type(decoded)}" + assert decoded["type"] == "BlockStored" + assert decoded["block_hashes"] == [123, 456] + assert decoded["parent_block_hash"] == 789 + assert decoded["token_ids"] == [1, 2, 3, 4] + assert decoded["block_size"] == 16 + assert decoded["lora_id"] is None + assert decoded["medium"] == "GPU" + assert decoded["lora_name"] is None + assert decoded.get("extra_keys") is None if _has_group_idx(BlockStored): - assert ( - decoded[next_idx] == 0 - ), f"group_idx at wrong position: {decoded[next_idx]}" - next_idx += 1 + assert decoded["group_idx"] == 0 if _has_kv_cache_spec_kind(BlockStored): - assert ( - decoded[next_idx] == "full_attention" - ), f"kv_cache_spec_kind at wrong position: {decoded[next_idx]}" - next_idx += 1 + assert decoded["kv_cache_spec_kind"] == "full_attention" if _has_kv_cache_spec_sliding_window(BlockStored): - assert ( - decoded[next_idx] == 128 - ), f"kv_cache_spec_sliding_window at wrong position: {decoded[next_idx]}" + assert decoded["kv_cache_spec_sliding_window"] == 128 def test_block_stored_tuple_extra_keys_serialization_format(self): """Verify multimodal tuple extra_keys keep the vLLM 0.19 wire shape.""" @@ -276,31 +247,20 @@ def test_block_stored_tuple_extra_keys_serialization_format(self): decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event)) - assert decoded[0] == "BlockStored" - assert decoded[8] == [[[mm_hash, 7]]], ( + assert decoded["type"] == "BlockStored" + assert decoded["extra_keys"] == [[[mm_hash, 7]]], ( "vLLM multimodal extra_keys no longer serialize as nested tuple/list " - f"payloads. Decoded: {decoded[8]!r}" + f"payloads. Decoded: {decoded['extra_keys']!r}" ) if _has_group_idx(BlockStored): - assert decoded[9] == 0, f"group_idx at wrong position: {decoded[9]}" + assert decoded["group_idx"] == 0 if _has_kv_cache_spec_kind(BlockStored): - kind_idx = 10 if _has_group_idx(BlockStored) else 9 - assert ( - decoded[kind_idx] == "full_attention" - ), f"kv_cache_spec_kind at wrong position: {decoded[kind_idx]}" + assert decoded["kv_cache_spec_kind"] == "full_attention" if _has_kv_cache_spec_sliding_window(BlockStored): - window_idx = 9 - if _has_group_idx(BlockStored): - window_idx += 1 - if _has_kv_cache_spec_kind(BlockStored): - window_idx += 1 - assert decoded[window_idx] == 128, ( - "kv_cache_spec_sliding_window at wrong position: " - f"{decoded[window_idx]}" - ) + assert decoded["kv_cache_spec_sliding_window"] == 128 def test_block_removed_serialization_format(self): - """Verify BlockRemoved serializes to expected msgpack array format.""" + """Verify BlockRemoved serializes to expected msgpack map format.""" import msgspec event_kwargs = { @@ -317,32 +277,14 @@ def test_block_removed_serialization_format(self): decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event)) - assert decoded[0] == "BlockRemoved" - expected_len = ( - 3 - + int(_has_group_idx(BlockRemoved)) - + int(_has_kv_cache_spec_kind(BlockRemoved)) - + int(_has_kv_cache_spec_sliding_window(BlockRemoved)) - ) - assert len(decoded) == expected_len, ( - f"Expected {expected_len} elements, got {len(decoded)}.\n" - f"Decoded: {decoded}\n" - f"If field count changed, update Rust deserializers." - ) - assert decoded[1] == [123, 456], f"block_hashes at wrong position: {decoded[1]}" - assert decoded[2] == "GPU", f"medium at wrong position: {decoded[2]}" - next_idx = 3 + assert decoded["type"] == "BlockRemoved" + assert decoded["block_hashes"] == [123, 456] + assert decoded["medium"] == "GPU" if _has_group_idx(BlockRemoved): - assert ( - decoded[next_idx] == 0 - ), f"group_idx at wrong position: {decoded[next_idx]}" - next_idx += 1 + assert decoded["group_idx"] == 0 if _has_kv_cache_spec_kind(BlockRemoved): - assert ( - decoded[next_idx] == "full_attention" - ), f"kv_cache_spec_kind at wrong position: {decoded[next_idx]}" - next_idx += 1 + assert decoded["kv_cache_spec_kind"] == "full_attention" if _has_kv_cache_spec_sliding_window(BlockRemoved): assert ( - decoded[next_idx] == 128 - ), f"kv_cache_spec_sliding_window at wrong position: {decoded[next_idx]}" + decoded["kv_cache_spec_sliding_window"] == 128 + ), "kv_cache_spec_sliding_window has wrong value" diff --git a/container/context.yaml b/container/context.yaml index 974d57d07662..2f2f40b24cd9 100644 --- a/container/context.yaml +++ b/container/context.yaml @@ -55,7 +55,7 @@ vllm: base_image: nvcr.io/nvidia/cuda-dl-base runtime_image: vllm/vllm-openai base_image_tag: 25.11-cuda13.0-devel-ubuntu24.04 - runtime_image_tag: v0.23.0-ubuntu2404 + runtime_image_tag: v0.24.0-ubuntu2404 # Baseline is vllm-openai's TRUE base, nvidia/cuda:13.0.2-base-ubuntu24.04 # (Docker Hub), NOT a self-baseline of the third-party vllm/vllm-openai image. # vllm-openai is built FROM nvidia/cuda, so this attributes everything vLLM @@ -74,7 +74,7 @@ vllm: base_image: ubuntu runtime_image: vllm/vllm-openai-cpu base_image_tag: 22.04 - runtime_image_tag: v0.23.0 + runtime_image_tag: v0.24.0 # baseline_sbom: not yet captured for cpu — runtime build runs without subtraction flashinf_ref: v0.6.8.post1 vllm_omni_ref: "v0.21.0rc1" diff --git a/container/deps/vllm/protected_packages.txt b/container/deps/vllm/protected_packages.txt index be3421a328f8..6c5d2590a6d4 100644 --- a/container/deps/vllm/protected_packages.txt +++ b/container/deps/vllm/protected_packages.txt @@ -11,7 +11,7 @@ vllm transformers tokenizers # vLLM-Omni (v0.21.0rc1) may require a newer safetensors than the upstream -# vLLM 0.23.0 image ships, so safetensors is intentionally left unfrozen here. +# vLLM 0.24.0 image ships, so safetensors is intentionally left unfrozen here. # safetensors msgspec flashinfer-python diff --git a/docs/reference/support-matrix.md b/docs/reference/support-matrix.md index 1e24937c3312..e7b53de7d48b 100644 --- a/docs/reference/support-matrix.md +++ b/docs/reference/support-matrix.md @@ -31,7 +31,7 @@ The following table shows the backend framework versions included with each Dyna | **Dynamo** | **SGLang** | **TensorRT-LLM** | **vLLM** | **NIXL** | | :--- | :--- | :--- | :--- | :--- | -| **main (ToT)** | `0.5.11` | `1.3.0rc19` | `0.23.0` | `1.0.1` (TRT-LLM); `1.1.0` (vLLM); `1.0.1` (SGLang) | +| **main (ToT)** | `0.5.11` | `1.3.0rc19` | `0.24.0` | `1.0.1` (TRT-LLM); `1.1.0` (vLLM); `1.0.1` (SGLang) | | **v1.3.0-dev.1** *(experimental)* | `0.5.12.post1` | `1.3.0rc17` | `0.22.0` | `0.10.1` (TRT-LLM); `1.1.0` (vLLM); `1.0.1` (SGLang) | | **v1.2.1** | `0.5.11` | `1.3.0rc14` | `0.20.1` | `0.10.1` (TRT-LLM, vLLM); `1.0.1` (SGLang) | | **v1.2.0** | `0.5.11` | `1.3.0rc14` | `0.20.1` | `0.10.1` (TRT-LLM, vLLM); `1.0.1` (SGLang) | diff --git a/lib/vllm-rs-backend/Cargo.toml b/lib/vllm-rs-backend/Cargo.toml index b1ed83f3fa98..e805e60fdcb1 100644 --- a/lib/vllm-rs-backend/Cargo.toml +++ b/lib/vllm-rs-backend/Cargo.toml @@ -48,10 +48,10 @@ serde_json = { workspace = true } tokio = { workspace = true } tokio-util = { workspace = true } tracing = { workspace = true } -vllm-engine-core-client = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.22.0", optional = true } -vllm-llm = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.22.0", optional = true } -vllm-managed-engine = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.22.0", optional = true } -vllm-metrics = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.22.0", optional = true } +vllm-engine-core-client = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.24.0", optional = true } +vllm-llm = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.24.0", optional = true } +vllm-managed-engine = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.24.0", optional = true } +vllm-metrics = { git = "https://github.com/vllm-project/vllm.git", tag = "v0.24.0", optional = true } [dev-dependencies] dynamo-backend-common = { workspace = true, features = ["testing"] } diff --git a/lib/vllm-rs-backend/src/backend.rs b/lib/vllm-rs-backend/src/backend.rs index 81473cc75aa6..f99f0dc9502f 100644 --- a/lib/vllm-rs-backend/src/backend.rs +++ b/lib/vllm-rs-backend/src/backend.rs @@ -208,10 +208,22 @@ impl LLMEngine for VllmBackend { .map_err(|e| cannot_connect(format!("failed to resolve handshake port: {e:#}")))?; let managed_config = { - let mut config = - self.managed_engine - .clone() - .into_config(self.model.clone(), None, handshake_port); + // vLLM 0.24 into_config takes 7 positional args; name them so a future + // value/signature change can't silently land in the wrong slot. + let max_model_len: Option = None; // let vLLM auto-fit from KV profiling + let max_logprobs: Option = None; + let language_model_only = false; + let disable_log_stats = false; // keep stats on (mirrors with_log_stats(true) below) + let shutdown_timeout: u64 = 0; // NOTE: 0 = abort in-flight requests on shutdown + let mut config = self.managed_engine.clone().into_config( + self.model.clone(), + max_model_len, + max_logprobs, + language_model_only, + disable_log_stats, + shutdown_timeout, + handshake_port, + ); self.extra.append_python_args(&mut config.python_args); config }; @@ -260,9 +272,7 @@ impl LLMEngine for VllmBackend { } }; - let context_length = client - .max_model_len() - .ok_or_else(|| backend_unknown("vLLM engine-core did not report max_model_len"))?; + let context_length = client.max_model_len(); let total_kv_blocks = match client.total_num_gpu_blocks() { 0 => None, blocks => Some(blocks), diff --git a/lib/vllm-rs-backend/src/control.rs b/lib/vllm-rs-backend/src/control.rs index 601c89b91ccf..f7f032ba61ff 100644 --- a/lib/vllm-rs-backend/src/control.rs +++ b/lib/vllm-rs-backend/src/control.rs @@ -6,6 +6,7 @@ use dynamo_backend_common::DynamoError; use serde::Deserialize; use vllm_engine_core_client::EngineCoreClient; +use vllm_engine_core_client::protocol::utility::PauseMode; const SUPPORTED_CONTROLS: [&str; 3] = ["sleep", "wake_up", "reset_prefix_cache"]; @@ -58,9 +59,17 @@ async fn sleep( } }; let level = body.level.unwrap_or(1); - let mode = body.mode.unwrap_or_else(|| "abort".to_string()); + let mode = match body.mode { + Some(mode) => match mode.parse::() { + Ok(mode) => mode, + Err(error) => { + return Ok(error_response(format!("invalid sleep mode: {error}"))); + } + }, + None => PauseMode::Abort, + }; - match client.sleep(level, &mode).await { + match client.sleep(level, mode).await { Ok(()) => Ok(serde_json::json!({ "status": "ok", "message": format!("Engine slept (level={level})"), diff --git a/lib/vllm-rs-backend/src/convert.rs b/lib/vllm-rs-backend/src/convert.rs index a12f60f2067e..512f17060dad 100644 --- a/lib/vllm-rs-backend/src/convert.rs +++ b/lib/vllm-rs-backend/src/convert.rs @@ -88,6 +88,7 @@ pub(crate) fn lower_request( .map(structured_outputs_from_guided_decoding), logprob_token_ids: None, skip_reading_prefix_cache: None, + thinking_token_budget: request.stop_conditions.max_thinking_tokens.map(u64::from), extra_args: extra_args_as_object(request.extra_args)?, }; apply_disaggregation_mode( @@ -382,6 +383,7 @@ mod tests { let mut request = sample_request(); request.stop_conditions.max_tokens = Some(7); request.stop_conditions.min_tokens = Some(2); + request.stop_conditions.max_thinking_tokens = Some(1024); request.stop_conditions.stop = Some(vec!["".to_string()]); request.stop_conditions.stop_token_ids_hidden = Some(vec![99]); request.eos_token_ids = vec![2, 3]; @@ -433,6 +435,7 @@ mod tests { let sampling = generate.sampling_params; assert_eq!(sampling.max_tokens, 7); assert_eq!(sampling.min_tokens, 2); + assert_eq!(sampling.thinking_token_budget, Some(1024)); assert_eq!(sampling.top_k, 0); assert_eq!(sampling.logprobs, Some(5)); assert_eq!(sampling.prompt_logprobs, Some(2)); @@ -680,6 +683,7 @@ mod tests { }], }), finish_reason: Some(VllmFinishReason::Stop(Some(VllmStopReason::TokenId(42)))), + cached_token_count: 0, kv_transfer_params: Some(json!({"connector": "kv"})), }; @@ -712,6 +716,7 @@ mod tests { positions: vec![PositionLogprobs { entries: vec![] }], }), finish_reason: None, + cached_token_count: 0, kv_transfer_params: None, }; @@ -787,6 +792,7 @@ mod tests { token_ids: vec![1, 2], logprobs: None, finish_reason: Some(reason), + cached_token_count: 0, kv_transfer_params: None, } } diff --git a/pyproject.toml b/pyproject.toml index afcb236588d2..49050df7a111 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,7 @@ trtllm =[ vllm = [ "uvloop", "nixl[cu13]==1.1.0", - "vllm[flashinfer,runai,otel]==0.23.0", + "vllm[flashinfer,runai,otel]==0.24.0", # vllm-omni is not part of ai-dynamo[vllm]: container builds inherit the # framework stack from vllm/vllm-openai, and pip/uv dependency resolution # for omni can override the vLLM torch stack. diff --git a/tests/frontend/test_prepost.py b/tests/frontend/test_prepost.py index 5a54cf990cbf..4f34c34a718b 100644 --- a/tests/frontend/test_prepost.py +++ b/tests/frontend/test_prepost.py @@ -20,10 +20,10 @@ ) from vllm.entrypoints.openai.engine.protocol import FunctionDefinition from vllm.outputs import CompletionOutput - from vllm.reasoning.qwen3_reasoning_parser import Qwen3ReasoningParser + from vllm.reasoning.qwen3_engine_reasoning_parser import Qwen3ParserReasoningAdapter from vllm.sampling_params import SamplingParams from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser - from vllm.tool_parsers.qwen3coder_tool_parser import Qwen3CoderToolParser + from vllm.tool_parsers.qwen3_engine_tool_parser import Qwen3EngineToolParser from dynamo.frontend.prepost import StreamingPostProcessor else: @@ -1378,7 +1378,7 @@ def processor(tokenizer, request_for_sampling, sampling_params): sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, tool_parser=tool_parser, - reasoning_parser_class=Qwen3ReasoningParser, + reasoning_parser_class=Qwen3ParserReasoningAdapter, chat_template_kwargs={"reasoning_effort": None}, ) @@ -1478,7 +1478,7 @@ def test_qwen3_coder_non_streaming_uses_batch_tool_parse( request_for_sampling=qwen3_coder_request_for_sampling, sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, - tool_parser=Qwen3CoderToolParser( + tool_parser=Qwen3EngineToolParser( tokenizer, qwen3_coder_request_for_sampling.tools ), reasoning_parser_class=None, @@ -1489,15 +1489,20 @@ def test_qwen3_coder_non_streaming_uses_batch_tool_parse( streaming_content = "".join( r.get("delta", {}).get("content", "") for r in streaming_results ) - assert "" in streaming_content - assert _collect_tool_calls(streaming_results) == [] + assert "" not in streaming_content + streaming_tool_calls = _collect_tool_calls(streaming_results) + assert len(streaming_tool_calls) == 1 + assert streaming_tool_calls[0]["function"]["name"] == "get_weather" + assert json.loads(streaming_tool_calls[0]["function"]["arguments"]) == { + "location": "NYC" + } non_streaming_proc = StreamingPostProcessor( tokenizer=tokenizer, request_for_sampling=qwen3_coder_request_for_sampling, sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, - tool_parser=Qwen3CoderToolParser( + tool_parser=Qwen3EngineToolParser( tokenizer, qwen3_coder_request_for_sampling.tools ), reasoning_parser_class=None, @@ -1547,7 +1552,7 @@ def test_qwen3_coder_non_streaming_preserves_content_before_tool_call( request_for_sampling=qwen3_coder_request_for_sampling, sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, - tool_parser=Qwen3CoderToolParser( + tool_parser=Qwen3EngineToolParser( tokenizer, qwen3_coder_request_for_sampling.tools ), reasoning_parser_class=None, @@ -1557,7 +1562,7 @@ def test_qwen3_coder_non_streaming_preserves_content_before_tool_call( results = _collect_results(proc, outputs) assert len(results) == 1 - assert results[0]["delta"]["content"] == "I can check that.\n" + assert results[0]["delta"]["content"] == "I can check that." tool_calls = _collect_tool_calls(results) assert len(tool_calls) == 1 @@ -1593,10 +1598,10 @@ def test_qwen3_coder_non_streaming_batches_reasoning_before_tool_parse( request_for_sampling=qwen3_coder_request_for_sampling, sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, - tool_parser=Qwen3CoderToolParser( + tool_parser=Qwen3EngineToolParser( tokenizer, qwen3_coder_request_for_sampling.tools ), - reasoning_parser_class=Qwen3ReasoningParser, + reasoning_parser_class=Qwen3ParserReasoningAdapter, chat_template_kwargs={"reasoning_effort": None}, stream_response=False, ) @@ -1616,6 +1621,59 @@ def test_qwen3_coder_non_streaming_batches_reasoning_before_tool_parse( assert results[0]["finish_reason"] == "tool_calls" +@pytest.mark.vllm +def test_qwen3_streaming_buffers_function_marker_after_reasoning_end( + tokenizer, qwen3_coder_request_for_sampling, sampling_params +): + outputs = [ + CompletionOutput( + index=0, + text=( + "Need the weather.\n" + "\n" + "\n" + "NYC\n" + "\n" + ), + token_ids=[151667, 151668], + cumulative_logprob=None, + logprobs=None, + ), + CompletionOutput( + index=0, + text="", + token_ids=[1002], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ), + ] + + proc = StreamingPostProcessor( + tokenizer=tokenizer, + request_for_sampling=qwen3_coder_request_for_sampling, + sampling_params=sampling_params, + prompt_token_ids=PROMPT_TOKEN_IDS, + tool_parser=Qwen3EngineToolParser( + tokenizer, qwen3_coder_request_for_sampling.tools + ), + reasoning_parser_class=Qwen3ParserReasoningAdapter, + chat_template_kwargs={"reasoning_effort": None}, + stream_response=True, + ) + + results = _collect_results(proc, outputs) + assert _collect_reasoning(results) == "Need the weather." + all_content = "".join(r.get("delta", {}).get("content", "") for r in results) + assert "" not in all_content + assert "" not in all_content + + tool_calls = _collect_tool_calls(results) + assert len(tool_calls) == 1 + assert tool_calls[0]["function"]["name"] == "get_weather" + assert json.loads(tool_calls[0]["function"]["arguments"]) == {"location": "NYC"} + + @pytest.mark.vllm def test_stream_interval_1(processor): """stream_interval=1: one token per chunk. Baseline that works.""" @@ -1676,7 +1734,7 @@ def test_stream_interval_20(tokenizer, request_for_sampling, sampling_params): sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, tool_parser=tool_parser, - reasoning_parser_class=Qwen3ReasoningParser, + reasoning_parser_class=Qwen3ParserReasoningAdapter, chat_template_kwargs={"reasoning_effort": None}, ) @@ -1729,7 +1787,7 @@ def test_stream_interval_20_reasoning_and_tool_finish_same_chunk( sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, tool_parser=tool_parser, - reasoning_parser_class=Qwen3ReasoningParser, + reasoning_parser_class=Qwen3ParserReasoningAdapter, chat_template_kwargs={"reasoning_effort": None}, ) @@ -1782,7 +1840,7 @@ def test_stream_terminal_single_chunk(tokenizer, request_for_sampling, sampling_ sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, tool_parser=tool_parser, - reasoning_parser_class=Qwen3ReasoningParser, + reasoning_parser_class=Qwen3ParserReasoningAdapter, chat_template_kwargs={"reasoning_effort": None}, ) @@ -1847,7 +1905,7 @@ def test_no_tool_call(tokenizer, request_for_sampling, sampling_params): sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, tool_parser=tool_parser, - reasoning_parser_class=Qwen3ReasoningParser, + reasoning_parser_class=Qwen3ParserReasoningAdapter, chat_template_kwargs={"reasoning_effort": None}, ) @@ -1949,7 +2007,7 @@ def test_streaming_parallel_tool_calls_no_think( sampling_params=sampling_params, prompt_token_ids=PROMPT_TOKEN_IDS, tool_parser=tool_parser, - reasoning_parser_class=Qwen3ReasoningParser, + reasoning_parser_class=Qwen3ParserReasoningAdapter, chat_template_kwargs={"enable_thinking": False}, ) diff --git a/tests/report_pytest_markers.py b/tests/report_pytest_markers.py index 4d47cde47af7..335667e4fefb 100755 --- a/tests/report_pytest_markers.py +++ b/tests/report_pytest_markers.py @@ -232,7 +232,7 @@ "vllm.outputs", "vllm.reasoning", "vllm.reasoning.mistral_reasoning_parser", - "vllm.reasoning.qwen3_reasoning_parser", + "vllm.reasoning.qwen3_engine_reasoning_parser", "vllm.renderers", "vllm.renderers.embed_utils", "vllm.sampling_params", @@ -241,7 +241,7 @@ "vllm.tool_parsers", "vllm.tool_parsers.hermes_tool_parser", "vllm.tool_parsers.mistral_tool_parser", - "vllm.tool_parsers.qwen3coder_tool_parser", + "vllm.tool_parsers.qwen3_engine_tool_parser", "vllm.utils", "vllm.utils.async_utils", "vllm.utils.hashing", From 62f0c8cc40954007732900812ca635c4e1cfc9ff Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 30 Jun 2026 13:07:12 -0700 Subject: [PATCH 010/320] feat(backend): add sample multimodal smoke coverage (#11001) Signed-off-by: Connor Carpenter --- .../src/dynamo/common/backend/README.md | 17 +- .../dynamo/common/backend/sample_engine.py | 112 ++++- .../backend/tests/test_sample_engine.py | 199 ++++++++- .../backends/sample/launch/multimodal_agg.sh | 59 +++ .../sample/launch/multimodal_disagg.sh | 83 ++++ .../sample/launch/multimodal_smoke_client.py | 170 ++++++++ lib/backend-common/CLAUDE.md | 9 +- lib/backend-common/README.md | 17 + lib/backend-common/src/testing.rs | 397 +++++++++++++++++- tests/runtime/test_sample_multimodal_smoke.py | 69 +++ 10 files changed, 1103 insertions(+), 29 deletions(-) create mode 100755 examples/backends/sample/launch/multimodal_agg.sh create mode 100755 examples/backends/sample/launch/multimodal_disagg.sh create mode 100644 examples/backends/sample/launch/multimodal_smoke_client.py create mode 100644 tests/runtime/test_sample_multimodal_smoke.py diff --git a/components/src/dynamo/common/backend/README.md b/components/src/dynamo/common/backend/README.md index eb5371e6ad13..0d4918ddeb4a 100644 --- a/components/src/dynamo/common/backend/README.md +++ b/components/src/dynamo/common/backend/README.md @@ -1,9 +1,10 @@ # Dynamo Python Backend -**Supported today:** aggregated and disaggregated (prefill/decode) -inference, metrics + Prometheus bridging, KV event publishing, -KV-aware (DP-rank) routing, health-check canaries, OpenTelemetry -tracing, and request-side guided decoding / structural tag. +**Supported today:** aggregated and disaggregated (prefill/decode/encode) +inference, the shared multimodal request and encoder-handoff contract, +metrics + Prometheus bridging, KV event publishing, KV-aware (DP-rank) +routing, health-check canaries, OpenTelemetry tracing, and request-side +guided decoding / structural tag. > **Work in progress.** Multimodal, diffusion (image/video/DLLM), > LoRA (SGLang / TRT-LLM — vLLM is supported), @@ -142,6 +143,12 @@ def main(): ``` See `sample_engine.py` for a complete, runnable reference implementation. +The sample engine includes synthetic multimodal handling for aggregated and +Encode/Prefill/Decode deployments. CPU-only direct worker-handoff smokes live in +`examples/backends/sample/launch/multimodal_agg.sh` and +`examples/backends/sample/launch/multimodal_disagg.sh`. These smokes exercise +distinct worker processes and TCP request transport; they intentionally bypass +the frontend and do not claim frontend routing coverage. ## Request / Response Types @@ -563,7 +570,7 @@ Request handling: | Feature | Description | |---------|-------------| | Text-in-text-out mode | OpenAI-compatible chat/completion with engine-side tokenization. Unified hardcodes `ModelInput.Tokens`. | -| Multimodal | Images / video / embeddings, NIXL embedding transfer, encode workers. `worker.py:_to_rust_disaggregation_mode` rejects the `ENCODE` role. | +| Multimodal | The shared request and `encoder_result` contract, Encode role, and discovery wiring are available. Frontend Encode-to-Prefill/Aggregated request routing and backend-specific encoder implementations remain separate work. | | Diffusion | Image (FLUX), video (Wan2.1), LLM diffusion (DLLM) workers; no diffusion engine, MediaOutput, or media scheduling on the unified path. | | LoRA adapters (SGLang / TRT-LLM) | Dynamic load / unload / list, ModelDeploymentCard publishing, per-adapter serialization locks, per-request adapter threading. **vLLM is supported on the unified path** — see [What works today](#what-works-today); SGLang and TRT-LLM advertise no LoRA updates yet. | | Snapshot / checkpoint | CRIU-based engine state save/restore + identity reload. | diff --git a/components/src/dynamo/common/backend/sample_engine.py b/components/src/dynamo/common/backend/sample_engine.py index db1a9a18b4db..af2bb3b1b673 100644 --- a/components/src/dynamo/common/backend/sample_engine.py +++ b/components/src/dynamo/common/backend/sample_engine.py @@ -28,6 +28,11 @@ ) from .health_check import build_health_check_payload, is_probe from .logprobs import parse_logprob_options +from .multimodal import ( + encoder_terminal_chunk, + extract_multimodal_kwargs, + require_encoder_result, +) from .publisher import ComponentSnapshot, KvEventSource, PushSource from .worker import WorkerConfig @@ -86,14 +91,15 @@ class SampleLLMEngine(LLMEngine): for engine leads implementing real backends. Disaggregation: - ``--disaggregation-mode {agg,prefill,decode}`` selects the role. + ``--disaggregation-mode {agg,prefill,decode,encode}`` selects the role. AGGREGATED is the default and produces ``max_tokens`` rotating tokens. PREFILL caps generation at one token and stamps a synthetic ``disaggregated_params`` payload on the terminal so the frontend's PrefillRouter has something to forward. DECODE requires the request to carry ``prefill_result`` (otherwise the frontend forgot to route through the prefill peer); on success - it generates normally. + it generates normally. ENCODE emits one terminal chunk containing + an object-shaped synthetic ``encoder_result`` and no generated tokens. """ def __init__( @@ -102,11 +108,13 @@ def __init__( max_tokens: int = 16, delay: float = 0.01, disaggregation_mode: DisaggregationMode = DisaggregationMode.AGGREGATED, + route_to_encoder: bool = False, ): self.model_name = model_name self.max_tokens = max_tokens self.delay = delay self.disaggregation_mode = disaggregation_mode + self.route_to_encoder = route_to_encoder self._kv_used_blocks = 0 self._publish_queue: queue.SimpleQueue[tuple[str, dict]] = queue.SimpleQueue() self._publish_stop = threading.Event() @@ -135,11 +143,19 @@ async def from_args( parser.add_argument("--event-plane", default=None) parser.add_argument( "--disaggregation-mode", - choices=[ - m.value for m in DisaggregationMode if m != DisaggregationMode.ENCODE - ], + choices=[m.value for m in DisaggregationMode], default=DisaggregationMode.AGGREGATED.value, - help="Disaggregation role: 'agg' (default), 'prefill', or 'decode'.", + help="Disaggregation role: 'agg' (default), 'prefill', 'decode', or 'encode'.", + ) + parser.add_argument( + "--route-to-encoder", + action="store_true", + help="Require an upstream Encode worker (valid for agg/prefill).", + ) + parser.add_argument( + "--disable-kv-routing", + action="store_true", + help="Disable KV event and load publishers (useful for isolated smokes).", ) args = parser.parse_args(argv) @@ -149,6 +165,7 @@ async def from_args( max_tokens=args.max_tokens, delay=args.delay, disaggregation_mode=mode, + route_to_encoder=args.route_to_encoder, ) worker_config = WorkerConfig( namespace=args.namespace, @@ -161,6 +178,8 @@ async def from_args( request_plane=args.request_plane, event_plane=args.event_plane, disaggregation_mode=mode, + route_to_encoder=args.route_to_encoder, + enable_kv_routing=not args.disable_kv_routing, ) return engine, worker_config @@ -179,9 +198,13 @@ async def start(self, worker_id: int) -> EngineConfig: ) async def kv_event_sources(self) -> list[KvEventSource]: + if self.disaggregation_mode == DisaggregationMode.ENCODE: + return [] return [PushSource(on_ready=self._start_publisher_thread, dp_rank=0)] def component_metrics_dp_ranks(self) -> list[int]: + if self.disaggregation_mode == DisaggregationMode.ENCODE: + return [] return [0] def attach_snapshot_publisher(self, publisher) -> None: @@ -258,9 +281,64 @@ def _release_synthetic_blocks(self, hashes: list[int]) -> None: self._publish_queue.put(("removed", {"block_hashes": hashes})) self._kv_used_blocks = max(0, self._kv_used_blocks - len(hashes)) + async def _encode_multimodal(self, request: GenerateRequest) -> dict[str, Any]: + """Return a deterministic-shape synthetic encoder handoff payload.""" + await asyncio.sleep(self.delay) + multimodal_kwargs = extract_multimodal_kwargs(request) or {} + return { + "handle": f"sample-encoder:{uuid.uuid4().hex}", + "multimodal_kwargs": multimodal_kwargs, + } + + def _validate_encoder_result(self, request: GenerateRequest) -> dict[str, Any]: + encoder_result = require_encoder_result(request, self.disaggregation_mode) + handle = encoder_result.get("handle") + if not isinstance(handle, str) or not handle.startswith("sample-encoder:"): + raise ValueError( + "encoder_result.handle must be a string starting with 'sample-encoder:'" + ) + return encoder_result + async def generate( self, request: GenerateRequest, context: Context ) -> AsyncGenerator[GenerateChunk, None]: + if self.disaggregation_mode == DisaggregationMode.ENCODE: + prompt_len = len(request.get("token_ids", [])) + if context.is_stopped(): + yield { + "token_ids": [], + "index": 0, + "finish_reason": "cancelled", + "completion_usage": { + "prompt_tokens": prompt_len, + "completion_tokens": 0, + "total_tokens": prompt_len, + }, + } + return + encoder_result = await self._encode_multimodal(request) + if context.is_stopped(): + yield { + "token_ids": [], + "index": 0, + "finish_reason": "cancelled", + "completion_usage": { + "prompt_tokens": prompt_len, + "completion_tokens": 0, + "total_tokens": prompt_len, + }, + } + return + yield encoder_terminal_chunk( + encoder_result, + completion_usage={ + "prompt_tokens": prompt_len, + "completion_tokens": 0, + "total_tokens": prompt_len, + }, + ) + return + # Canary probes bypass cross-worker coordination — run as aggregated. if self.disaggregation_mode == DisaggregationMode.DECODE and not is_probe( request @@ -269,6 +347,21 @@ async def generate( if self.disaggregation_mode == DisaggregationMode.PREFILL: enforce_prefill_max_tokens(request) + multimodal_kwargs = extract_multimodal_kwargs(request) + sample_multimodal_result: Optional[dict[str, Any]] = None + if multimodal_kwargs is not None: + if self.disaggregation_mode == DisaggregationMode.DECODE: + raise ValueError( + "decode worker should not receive raw multimodal payloads; " + "encoder inputs are consumed by the upstream prefill worker" + ) + if self.route_to_encoder: + sample_multimodal_result = self._validate_encoder_result(request) + else: + # Aggregated workers process multimodal inputs locally when + # encoder routing is disabled. + sample_multimodal_result = await self._encode_multimodal(request) + token_ids = request.get("token_ids", []) prompt_len = len(token_ids) stop_conditions = request.get("stop_conditions", {}) @@ -284,6 +377,13 @@ async def generate( async for chunk in self._generate_tokens( prompt_len, max_new, context, logprobs_k ): + if ( + chunk.get("finish_reason") + and sample_multimodal_result is not None + ): + chunk["engine_data"] = { + "sample_multimodal": sample_multimodal_result + } yield chunk finally: self._release_synthetic_blocks(block_hashes) diff --git a/components/src/dynamo/common/backend/tests/test_sample_engine.py b/components/src/dynamo/common/backend/tests/test_sample_engine.py index b3d7ccbf0e52..0706b550ac3e 100644 --- a/components/src/dynamo/common/backend/tests/test_sample_engine.py +++ b/components/src/dynamo/common/backend/tests/test_sample_engine.py @@ -5,7 +5,7 @@ from __future__ import annotations -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -106,6 +106,22 @@ async def test_decode_mode_runs_to_completion_when_prefill_result_provided(): assert "disaggregated_params" not in terminal +async def test_decode_mode_rejects_raw_multimodal_payload(): + engine = SampleLLMEngine( + max_tokens=1, + delay=0.0, + disaggregation_mode=DisaggregationMode.DECODE, + ) + request = { + "token_ids": [1, 2, 3], + "multi_modal_data": {"image": [{"url": "data:image/png;base64,AA=="}]}, + "prefill_result": {"disaggregated_params": {"sample_handle": "from-test"}}, + } + + with pytest.raises(ValueError, match="decode worker should not receive raw"): + await _collect(engine, request) + + async def test_from_args_propagates_mode_to_worker_config(): engine, worker_config = await SampleLLMEngine.from_args( ["--disaggregation-mode", "prefill"] @@ -114,6 +130,187 @@ async def test_from_args_propagates_mode_to_worker_config(): assert worker_config.disaggregation_mode is DisaggregationMode.PREFILL +async def test_aggregated_mode_processes_multimodal_kwargs_locally(monkeypatch): + engine = SampleLLMEngine(max_tokens=1, delay=0.0) + encode = AsyncMock(wraps=engine._encode_multimodal) + monkeypatch.setattr(engine, "_encode_multimodal", encode) + request = { + "token_ids": [1, 2], + "multi_modal_data": {"image": [{"url": "data:image/png;base64,AA=="}]}, + "mm_processor_kwargs": {"min_pixels": 64}, + } + + chunks = await _collect(engine, request) + + assert chunks[-1]["finish_reason"] == "length" + assert chunks[-1]["engine_data"]["sample_multimodal"]["multimodal_kwargs"] == { + "multi_modal_data": request["multi_modal_data"], + "mm_processor_kwargs": request["mm_processor_kwargs"], + } + encode.assert_awaited_once_with(request) + + +async def test_encode_mode_emits_single_terminal_with_encoder_result(): + engine = SampleLLMEngine( + delay=0.0, + disaggregation_mode=DisaggregationMode.ENCODE, + ) + request = { + "token_ids": [1, 2, 3], + "multi_modal_data": {"image": [{"url": "data:image/png;base64,AA=="}]}, + } + + chunks = await _collect(engine, request) + + assert len(chunks) == 1 + terminal = chunks[0] + assert terminal["token_ids"] == [] + assert terminal["finish_reason"] == "stop" + assert terminal["completion_usage"] == { + "prompt_tokens": 3, + "completion_tokens": 0, + "total_tokens": 3, + } + encoder_result = terminal["encoder_result"] + assert encoder_result["handle"].startswith("sample-encoder:") + assert ( + encoder_result["multimodal_kwargs"]["multi_modal_data"] + == request["multi_modal_data"] + ) + + +@pytest.mark.parametrize("stop_checks", [[True], [False, True]]) +async def test_encode_mode_observes_cancellation(stop_checks): + engine = SampleLLMEngine( + delay=0.0, + disaggregation_mode=DisaggregationMode.ENCODE, + ) + context = _ctx() + context.is_stopped.side_effect = stop_checks + + chunks = [ + chunk + async for chunk in engine.generate( + {"token_ids": [1, 2, 3], "multi_modal_data": {"image": []}}, context + ) + ] + + assert chunks == [ + { + "token_ids": [], + "index": 0, + "finish_reason": "cancelled", + "completion_usage": { + "prompt_tokens": 3, + "completion_tokens": 0, + "total_tokens": 3, + }, + } + ] + + +async def test_encoder_routed_worker_requires_encoder_result(): + engine = SampleLLMEngine( + max_tokens=1, + delay=0.0, + route_to_encoder=True, + ) + request = { + "token_ids": [1], + "multi_modal_data": {"image": [{"url": "data:image/png;base64,AA=="}]}, + } + + with pytest.raises(ValueError, match="no encoder_result"): + await _collect(engine, request) + + +async def test_encoder_routed_worker_rejects_malformed_encoder_result(): + engine = SampleLLMEngine( + max_tokens=1, + delay=0.0, + route_to_encoder=True, + ) + request = { + "token_ids": [1], + "multi_modal_data": {"image": [{"url": "data:image/png;base64,AA=="}]}, + "encoder_result": {"handle": "not-a-sample-handle"}, + } + + with pytest.raises(ValueError, match=r"encoder_result\.handle"): + await _collect(engine, request) + + +async def test_multimodal_epd_handoff_contract(): + """Exercise Encode -> Prefill -> Decode using separate role instances.""" + encode = SampleLLMEngine( + delay=0.0, + disaggregation_mode=DisaggregationMode.ENCODE, + ) + prefill = SampleLLMEngine( + delay=0.0, + disaggregation_mode=DisaggregationMode.PREFILL, + route_to_encoder=True, + ) + decode = SampleLLMEngine( + max_tokens=2, + delay=0.0, + disaggregation_mode=DisaggregationMode.DECODE, + ) + multimodal_request = { + "token_ids": [1, 2, 3], + "multi_modal_data": {"image": [{"url": "data:image/png;base64,AA=="}]}, + } + + [encode_terminal] = await _collect(encode, multimodal_request) + [prefill_terminal] = await _collect( + prefill, + { + **multimodal_request, + "encoder_result": encode_terminal["encoder_result"], + "stop_conditions": {"max_tokens": 8}, + }, + ) + decode_chunks = await _collect( + decode, + { + "token_ids": multimodal_request["token_ids"], + "prefill_result": { + "disaggregated_params": prefill_terminal["disaggregated_params"] + }, + }, + ) + + assert prefill_terminal["finish_reason"] == "length" + assert ( + prefill_terminal["engine_data"]["sample_multimodal"] + == encode_terminal["encoder_result"] + ) + assert len(decode_chunks) == 2 + assert decode_chunks[-1]["finish_reason"] == "length" + + +async def test_from_args_propagates_encode_routing(): + engine, worker_config = await SampleLLMEngine.from_args( + ["--disaggregation-mode", "prefill", "--route-to-encoder"] + ) + + assert engine.route_to_encoder is True + assert worker_config.route_to_encoder is True + + +async def test_from_args_can_disable_kv_routing(): + _, worker_config = await SampleLLMEngine.from_args(["--disable-kv-routing"]) + + assert worker_config.enable_kv_routing is False + + +async def test_encode_mode_opts_out_of_kv_publishers(): + engine = SampleLLMEngine(disaggregation_mode=DisaggregationMode.ENCODE) + + assert await engine.kv_event_sources() == [] + assert engine.component_metrics_dp_ranks() == [] + + async def test_source_descriptors_have_expected_shape(): engine = SampleLLMEngine(max_tokens=4, delay=0.0) [kv_src] = await engine.kv_event_sources() diff --git a/examples/backends/sample/launch/multimodal_agg.sh b/examples/backends/sample/launch/multimodal_agg.sh new file mode 100755 index 000000000000..83a4a1cdedc0 --- /dev/null +++ b/examples/backends/sample/launch/multimodal_agg.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# CPU-only aggregated multimodal smoke for the sample backend. + +set -e +trap 'echo Cleaning up...; kill 0' EXIT + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +REPO_ROOT="$(readlink -f "$SCRIPT_DIR/../../../..")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" + +MODEL_NAME="${MODEL_NAME:-$REPO_ROOT/lib/llm/tests/data/sample-models/TinyLlama_v1.1}" +NAMESPACE="${NAMESPACE:-dynamo}" +COMPONENT="${COMPONENT:-sample-multimodal-agg}" + +EXTRA_ARGS=() +while [[ $# -gt 0 ]]; do + case $1 in + --model-name) + MODEL_NAME="$2" + shift 2 + ;; + --namespace) + NAMESPACE="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [--model-name NAME] [--namespace NAMESPACE] [WORKER OPTIONS]" + exit 0 + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +# This direct worker smoke intentionally has no frontend. print_launch_banner +# always advertises a frontend URL, so use a scoped message instead. +echo "Running direct aggregated-worker multimodal handoff smoke with $MODEL_NAME" + +DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ +python3 -m dynamo.common.backend.sample_main \ + --model-name "$MODEL_NAME" \ + --namespace "$NAMESPACE" \ + --component "$COMPONENT" \ + --disable-kv-routing \ + "${EXTRA_ARGS[@]}" & + +python3 "$SCRIPT_DIR/multimodal_smoke_client.py" \ + --mode aggregated \ + --model-name "$MODEL_NAME" \ + --namespace "$NAMESPACE" \ + --aggregated-component "$COMPONENT" & + +wait_any_exit diff --git a/examples/backends/sample/launch/multimodal_disagg.sh b/examples/backends/sample/launch/multimodal_disagg.sh new file mode 100755 index 000000000000..fd5d3f70556e --- /dev/null +++ b/examples/backends/sample/launch/multimodal_disagg.sh @@ -0,0 +1,83 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# CPU-only Encode/Prefill/Decode multimodal handoff smoke for the sample backend. + +set -e +trap 'echo Cleaning up...; kill 0' EXIT + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +REPO_ROOT="$(readlink -f "$SCRIPT_DIR/../../../..")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" + +MODEL_NAME="${MODEL_NAME:-$REPO_ROOT/lib/llm/tests/data/sample-models/TinyLlama_v1.1}" +NAMESPACE="${NAMESPACE:-dynamo}" +ENCODE_COMPONENT="${ENCODE_COMPONENT:-sample-multimodal-encode}" +PREFILL_COMPONENT="${PREFILL_COMPONENT:-sample-multimodal-prefill}" +DECODE_COMPONENT="${DECODE_COMPONENT:-sample-multimodal-decode}" + +EXTRA_ARGS=() +while [[ $# -gt 0 ]]; do + case $1 in + --model-name) + MODEL_NAME="$2" + shift 2 + ;; + --namespace) + NAMESPACE="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [--model-name NAME] [--namespace NAMESPACE] [WORKER OPTIONS]" + exit 0 + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +# This direct worker smoke intentionally has no frontend. print_launch_banner +# always advertises a frontend URL, so use a scoped message instead. +echo "Running direct Encode -> Prefill -> Decode multimodal handoff smoke with $MODEL_NAME" + +DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT1:-8081} \ +python3 -m dynamo.common.backend.sample_main \ + --model-name "$MODEL_NAME" \ + --namespace "$NAMESPACE" \ + --component "$ENCODE_COMPONENT" \ + --disaggregation-mode encode \ + --disable-kv-routing \ + "${EXTRA_ARGS[@]}" & + +DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT2:-8082} \ +python3 -m dynamo.common.backend.sample_main \ + --model-name "$MODEL_NAME" \ + --namespace "$NAMESPACE" \ + --component "$PREFILL_COMPONENT" \ + --disaggregation-mode prefill \ + --route-to-encoder \ + --disable-kv-routing \ + "${EXTRA_ARGS[@]}" & + +DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT3:-8083} \ +python3 -m dynamo.common.backend.sample_main \ + --model-name "$MODEL_NAME" \ + --namespace "$NAMESPACE" \ + --component "$DECODE_COMPONENT" \ + --disaggregation-mode decode \ + --disable-kv-routing \ + "${EXTRA_ARGS[@]}" & + +python3 "$SCRIPT_DIR/multimodal_smoke_client.py" \ + --mode epd \ + --model-name "$MODEL_NAME" \ + --namespace "$NAMESPACE" \ + --encode-component "$ENCODE_COMPONENT" \ + --prefill-component "$PREFILL_COMPONENT" \ + --decode-component "$DECODE_COMPONENT" & + +wait_any_exit diff --git a/examples/backends/sample/launch/multimodal_smoke_client.py b/examples/backends/sample/launch/multimodal_smoke_client.py new file mode 100644 index 000000000000..1f62df1b3c3a --- /dev/null +++ b/examples/backends/sample/launch/multimodal_smoke_client.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct CPU smoke client for the sample multimodal worker roles.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +from typing import Any + +from dynamo.runtime import DistributedRuntime + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--mode", choices=("aggregated", "epd"), required=True) + parser.add_argument("--namespace", default="dynamo") + parser.add_argument("--model-name", default="sample-model") + parser.add_argument("--aggregated-component", default="sample-multimodal-agg") + parser.add_argument("--encode-component", default="sample-multimodal-encode") + parser.add_argument("--prefill-component", default="sample-multimodal-prefill") + parser.add_argument("--decode-component", default="sample-multimodal-decode") + parser.add_argument("--timeout", type=float, default=60.0) + return parser.parse_args() + + +async def connect(runtime: DistributedRuntime, endpoint_name: str, timeout: float): + endpoint = runtime.endpoint(endpoint_name) + client = await endpoint.client() + await asyncio.wait_for(client.wait_for_instances(), timeout=timeout) + return client + + +async def collect( + client, request: dict[str, Any], timeout: float +) -> list[dict[str, Any]]: + chunks: list[dict[str, Any]] = [] + async with asyncio.timeout(timeout): + stream = await client.generate(request) + async for response in stream: + if response.is_error(): + comments = response.comments() or [] + raise RuntimeError("worker returned an error: " + "; ".join(comments)) + data = response.data() + if isinstance(data, str): + data = json.loads(data) + if data is not None: + chunks.append(data) + return chunks + + +def multimodal_request(model_name: str) -> dict[str, Any]: + return { + "model": model_name, + "token_ids": [1, 2, 3], + "sampling_options": {}, + "stop_conditions": {"max_tokens": 2}, + "output_options": {}, + "multi_modal_data": {"image_url": [{"Url": "data:image/png;base64,AA=="}]}, + "mm_processor_kwargs": {"min_pixels": 64}, + } + + +def _require_equal(actual: Any, expected: Any, field: str) -> None: + if actual != expected: + raise RuntimeError(f"{field} mismatch: expected {expected!r}, got {actual!r}") + + +async def run_aggregated(runtime: DistributedRuntime, args: argparse.Namespace) -> None: + client = await connect( + runtime, + f"{args.namespace}.{args.aggregated_component}.generate", + args.timeout, + ) + request = multimodal_request(args.model_name) + chunks = await collect(client, request, args.timeout) + terminal = chunks[-1] + observed = terminal["engine_data"]["sample_multimodal"]["multimodal_kwargs"] + _require_equal( + observed["multi_modal_data"], + request["multi_modal_data"], + "aggregated multi_modal_data", + ) + _require_equal( + observed["mm_processor_kwargs"], + request["mm_processor_kwargs"], + "aggregated mm_processor_kwargs", + ) + + +async def run_epd(runtime: DistributedRuntime, args: argparse.Namespace) -> None: + encode, prefill, decode = await asyncio.gather( + connect( + runtime, + f"{args.namespace}.{args.encode_component}.generate", + args.timeout, + ), + connect( + runtime, + f"{args.namespace}.{args.prefill_component}.generate", + args.timeout, + ), + connect( + runtime, + f"{args.namespace}.{args.decode_component}.generate", + args.timeout, + ), + ) + request = multimodal_request(args.model_name) + + [encode_terminal] = await collect(encode, request, args.timeout) + _require_equal(encode_terminal["token_ids"], [], "encode token_ids") + _require_equal(encode_terminal["finish_reason"], "stop", "encode finish_reason") + + [prefill_terminal] = await collect( + prefill, + {**request, "encoder_result": encode_terminal["encoder_result"]}, + args.timeout, + ) + _require_equal( + prefill_terminal["engine_data"]["sample_multimodal"], + encode_terminal["encoder_result"], + "prefill encoder_result handoff", + ) + _require_equal( + encode_terminal["encoder_result"]["multimodal_kwargs"], + { + "multi_modal_data": request["multi_modal_data"], + "mm_processor_kwargs": request["mm_processor_kwargs"], + }, + "encode multimodal_kwargs", + ) + + decode_request = { + "model": args.model_name, + "token_ids": request["token_ids"], + "sampling_options": {}, + "stop_conditions": {"max_tokens": 2}, + "output_options": {}, + "prefill_result": { + "disaggregated_params": prefill_terminal["disaggregated_params"] + }, + } + decode_chunks = await collect(decode, decode_request, args.timeout) + _require_equal(decode_chunks[-1]["finish_reason"], "length", "decode finish_reason") + + +async def main() -> None: + args = parse_args() + # Launch scripts assign DYN_SYSTEM_PORT to workers. The direct client is a + # separate runtime and must allocate its own system-server port. + os.environ.pop("DYN_SYSTEM_PORT", None) + runtime = DistributedRuntime(asyncio.get_running_loop(), "etcd", "tcp") + try: + # Bound the whole smoke, rather than granting each sequential EPD stage + # a fresh timeout that can exceed the outer test harness deadline. + async with asyncio.timeout(args.timeout): + if args.mode == "aggregated": + await run_aggregated(runtime, args) + else: + await run_epd(runtime, args) + finally: + runtime.shutdown() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/lib/backend-common/CLAUDE.md b/lib/backend-common/CLAUDE.md index f246e9a8afed..c0dd900a3b63 100644 --- a/lib/backend-common/CLAUDE.md +++ b/lib/backend-common/CLAUDE.md @@ -510,6 +510,13 @@ async fn my_engine_satisfies_contract() { constructs one engine for the main lifecycle test and a second pristine engine for the "cleanup before start" check. +Encode-role engines call `run_encode_conformance`. It sends a multimodal +request and checks the narrower handoff response contract: one terminal `Stop` +chunk, no generated tokens, and an object-shaped `encoder_result`. Terminal +usage remains optional; when provided, it must report zero completion tokens +consistently. Encode engines must also satisfy the common KV source, metrics, +concurrency, cancellation, and cleanup lifecycle guarantees. + The kit asserts: - `start()` returns a non-empty `EngineConfig.model`. @@ -550,7 +557,7 @@ Also available: `testing::mock_context()` and | `adapter.rs` | `EngineAdapter` — bridges `LLMEngine` to `AsyncEngine` (token telemetry, disagg first-token, debug validator). `RawEngineAdapter` — bridges `RawEngine` to `AsyncEngine` (JSON passthrough, cancellation monitor; no token telemetry/disagg). `JsonProbeAdapter` — JSON health-check wrapper for the LLM path (the raw path is already JSON-shaped). | | `run.rs` | `pub fn run(engine, config)` (LLM) and `pub fn run_raw(engine, config)` (raw media) — entry points used by per-backend `main.rs`. Non-generic. | | `args.rs` | `CommonArgs` — shared CLI flags (`--namespace`, `--component`, `--disaggregation-mode`, etc.) that every engine's `Args` flattens in. | -| `disagg.rs` | `DisaggregationMode` enum (`Aggregated` / `Prefill` / `Decode`) with `clap::ValueEnum` derive. | +| `disagg.rs` | `DisaggregationMode` enum (`Aggregated` / `Prefill` / `Decode` / `Encode`) with `clap::ValueEnum` derive. | | `error.rs` | Re-exports `DynamoError`, `ErrorType`, `BackendError` from `dynamo-runtime`. No custom error types. | | `validate.rs` | Debug-build stream validator. Compiled out in release. | | `testing.rs` | Conformance test kit. Gated behind the `testing` feature. | diff --git a/lib/backend-common/README.md b/lib/backend-common/README.md index fb440da7f26b..55bb68443da7 100644 --- a/lib/backend-common/README.md +++ b/lib/backend-common/README.md @@ -279,6 +279,23 @@ async fn my_engine_passes_conformance() { } ``` +Encode-role engines use the narrower handoff contract: + +```rust +#[tokio::test] +async fn my_encoder_passes_conformance() { + dynamo_backend_common::testing::run_encode_conformance(MyEncoder::new_for_test) + .await + .expect("encode conformance"); +} +``` + +`run_encode_conformance` sends a multimodal request and requires one terminal +`FinishReason::Stop` chunk, empty `token_ids`, and an object-shaped +`encoder_result`. When terminal usage is provided, it must consistently report +zero completion tokens. The suite also applies the same KV source, metrics, +concurrency, cancellation, and cleanup checks as token-engine conformance. + The kit asserts: | Check | Failure mode | diff --git a/lib/backend-common/src/testing.rs b/lib/backend-common/src/testing.rs index 84778cf95a4b..acc28cb021f5 100644 --- a/lib/backend-common/src/testing.rs +++ b/lib/backend-common/src/testing.rs @@ -4,8 +4,8 @@ //! Conformance test kit for [`LLMEngine`] and [`RawEngine`] implementations. //! //! Engines wire themselves into the test suite with one call — -//! [`run_conformance`] for token engines, [`run_raw_conformance`] for raw -//! media engines: +//! [`run_conformance`] for token engines, [`run_encode_conformance`] for +//! Encode-role engines, and [`run_raw_conformance`] for raw media engines: //! //! ```ignore //! #[tokio::test] @@ -27,7 +27,7 @@ use std::sync::Arc; use std::time::Duration; -use dynamo_llm::protocols::common::preprocessor::PreprocessedRequest; +use dynamo_llm::protocols::common::preprocessor::{MultimodalData, PreprocessedRequest}; use dynamo_llm::protocols::common::{FinishReason, OutputOptions, SamplingOptions, StopConditions}; use dynamo_runtime::engine::AsyncEngineContext; use dynamo_runtime::pipeline::{AsyncEngineContextProvider, Context}; @@ -59,6 +59,7 @@ pub fn cancelling_context(after: Duration) -> Arc { /// Which conformance check failed, and why. #[derive(Debug)] +#[non_exhaustive] pub enum ConformanceFailure { StartFailed(String), EmptyModelInConfig, @@ -91,6 +92,25 @@ pub enum ConformanceFailure { chunked: usize, reported: u32, }, + /// Encode conformance requires exactly one terminal handoff chunk. + EncodeChunkCount { + count: usize, + }, + /// Encode terminals use the standard `Stop` finish reason. + EncodeTerminalExpected, + /// Encode workers produce a handoff payload, not generated tokens. + EncodeTokensNotEmpty { + count: usize, + }, + /// Encode terminal usage must describe a zero-token handoff response. + EncodeUsageMismatch { + expected_prompt: u32, + prompt: u32, + completion: u32, + total: u32, + }, + /// The encoder handoff contract is object-shaped at every boundary. + EncoderResultExpectedObject, } impl std::fmt::Display for ConformanceFailure { @@ -148,6 +168,34 @@ impl std::fmt::Display for ConformanceFailure { completion_usage.completion_tokens = {reported} on the terminal \ (engine bookkeeping diverges from streamed output)" ), + EncodeChunkCount { count } => write!( + f, + "encode generate() must yield exactly one terminal chunk; got {count}" + ), + EncodeTerminalExpected => write!( + f, + "encode generate() must yield a terminal chunk with \ + finish_reason = FinishReason::Stop" + ), + EncodeTokensNotEmpty { count } => write!( + f, + "encode terminal chunk must have empty token_ids; got {count} tokens" + ), + EncodeUsageMismatch { + expected_prompt, + prompt, + completion, + total, + } => write!( + f, + "encode terminal usage must report prompt={expected_prompt}, completion=0, \ + total={expected_prompt}; got prompt={prompt}, completion={completion}, \ + total={total}" + ), + EncoderResultExpectedObject => write!( + f, + "encode terminal chunk must carry an object-shaped encoder_result" + ), } } } @@ -190,10 +238,16 @@ where // 5. Interleaved generate() calls both complete — catches shared-state bugs. // Uses tokio::join! under the test runtime (single-threaded by default), // so this is interleaving rather than true parallelism. - check_concurrent_generates(&engine, &config.model).await?; + check_concurrent_generates(&engine, &config.model, LlmConformanceMode::Token).await?; // 6. Cancellation is observed within a bounded deadline. - check_cancellation(&engine, &config.model, DEFAULT_CANCEL_DEADLINE).await?; + check_cancellation( + &engine, + &config.model, + LlmConformanceMode::Token, + DEFAULT_CANCEL_DEADLINE, + ) + .await?; // 7. cleanup() succeeds and is idempotent. engine @@ -217,6 +271,114 @@ where Ok(()) } +/// Run the Encode-role conformance suite against an [`LLMEngine`]. +/// +/// Encode engines have a narrower response shape than token generators, but +/// share the same routing, metrics, concurrency, cancellation, and cleanup +/// lifecycle contracts. +pub async fn run_encode_conformance(mut factory: F) -> Result<(), ConformanceFailure> +where + E: LLMEngine, + F: FnMut() -> E, +{ + let engine = factory(); + let config = engine + .start(0) + .await + .map_err(|e| StartFailed(e.to_string()))?; + if config.model.is_empty() { + return Err(EmptyModelInConfig); + } + + check_kv_event_sources(&engine).await?; + check_setup_metrics(&engine).await?; + check_encode_generate(&engine, &config.model).await?; + check_concurrent_generates(&engine, &config.model, LlmConformanceMode::Encode).await?; + check_cancellation( + &engine, + &config.model, + LlmConformanceMode::Encode, + DEFAULT_CANCEL_DEADLINE, + ) + .await?; + + engine + .cleanup() + .await + .map_err(|e| CleanupFailed(e.to_string()))?; + engine + .cleanup() + .await + .map_err(|e| SecondCleanupFailed(e.to_string()))?; + + let fresh = factory(); + fresh + .cleanup() + .await + .map_err(|e| CleanupWithoutStartFailed(e.to_string()))?; + + Ok(()) +} + +async fn check_encode_generate( + engine: &E, + model: &str, +) -> Result<(), ConformanceFailure> { + let request = encode_request(model); + let expected_prompt = request.token_ids.len() as u32; + let stream = engine + .generate(request, GenerateContext::new(mock_context(), None)) + .await + .map_err(|e| GenerateFailed(e.to_string()))?; + let items: Vec<_> = stream.collect().await; + validate_encode_items(items, expected_prompt) +} + +fn validate_encode_items( + items: Vec>, + expected_prompt: u32, +) -> Result<(), ConformanceFailure> { + if items.len() != 1 { + return Err(EncodeChunkCount { count: items.len() }); + } + let chunk = items + .into_iter() + .next() + .expect("length checked") + .map_err(|e| StreamYieldedError(e.to_string()))?; + + if !matches!(chunk.finish_reason, Some(FinishReason::Stop)) { + return Err(EncodeTerminalExpected); + } + if !chunk.token_ids.is_empty() { + return Err(EncodeTokensNotEmpty { + count: chunk.token_ids.len(), + }); + } + if !chunk.encoder_result.as_ref().is_some_and(|v| v.is_object()) { + return Err(EncoderResultExpectedObject); + } + if let Some(usage) = chunk.completion_usage.as_ref() + && (usage.prompt_tokens != expected_prompt + || usage.completion_tokens != 0 + || usage.total_tokens != expected_prompt) + { + return Err(EncodeUsageMismatch { + expected_prompt, + prompt: usage.prompt_tokens, + completion: usage.completion_tokens, + total: usage.total_tokens, + }); + } + Ok(()) +} + +#[derive(Clone, Copy)] +enum LlmConformanceMode { + Token, + Encode, +} + fn request(model: &str) -> PreprocessedRequest { // Keep conformance smokes bounded for real LLM engines. The separate // cancellation check still requests enough tokens to catch ignored cancels. @@ -237,6 +399,28 @@ fn request_with_max_tokens(model: &str, max_tokens: Option) -> Preprocessed .expect("build request") } +fn encode_request(model: &str) -> PreprocessedRequest { + let multi_modal_data = std::collections::HashMap::from([( + "image".to_string(), + vec![MultimodalData::RawUrl( + "data:image/png;base64,AA==".to_string(), + )], + )]); + PreprocessedRequest::builder() + .model(model.to_string()) + .token_ids(vec![1, 2, 3]) + .multi_modal_data(Some(multi_modal_data)) + .mm_processor_kwargs(Some(serde_json::json!({ "min_pixels": 64 }))) + .stop_conditions(StopConditions { + max_tokens: Some(8), + ..Default::default() + }) + .sampling_options(SamplingOptions::default()) + .output_options(OutputOptions::default()) + .build() + .expect("build encode request") +} + async fn check_single_generate( engine: &E, model: &str, @@ -291,6 +475,7 @@ async fn check_single_generate( async fn check_concurrent_generates( engine: &E, model: &str, + mode: LlmConformanceMode, ) -> Result<(), ConformanceFailure> { // 8 in-flight streams — enough to catch state-tramping under interleaved // polls. Under a single-threaded test runtime this is interleaving rather @@ -298,15 +483,27 @@ async fn check_concurrent_generates( const CONCURRENT: usize = 8; let futs = (0..CONCURRENT).map(|_| async { let ctx = mock_context(); + let request = match mode { + LlmConformanceMode::Token => request(model), + LlmConformanceMode::Encode => encode_request(model), + }; + let expected_prompt = request.token_ids.len() as u32; let stream = engine - .generate(request(model), GenerateContext::new(ctx, None)) + .generate(request, GenerateContext::new(ctx, None)) .await .map_err(|e| ConcurrentGenerateFailed(e.to_string()))?; - let n = stream.count().await; - if n == 0 { - Err(ConcurrentGenerateFailed("stream was empty".to_string())) - } else { - Ok(()) + match mode { + LlmConformanceMode::Token => { + let n = stream.count().await; + if n == 0 { + Err(ConcurrentGenerateFailed("stream was empty".to_string())) + } else { + Ok(()) + } + } + LlmConformanceMode::Encode => { + validate_encode_items(stream.collect().await, expected_prompt) + } } }); for result in futures::future::join_all(futs).await { @@ -367,6 +564,7 @@ async fn check_setup_metrics(engine: &E) -> Result<(), Conformance async fn check_cancellation( engine: &E, model: &str, + mode: LlmConformanceMode, deadline: Duration, ) -> Result<(), ConformanceFailure> { // Request enough tokens that an engine which ignores cancellation @@ -374,11 +572,12 @@ async fn check_cancellation( const LONG_MAX_TOKENS: u32 = 10_000; let ctx = mock_context(); + let request = match mode { + LlmConformanceMode::Token => request_with_max_tokens(model, Some(LONG_MAX_TOKENS)), + LlmConformanceMode::Encode => encode_request(model), + }; let stream = engine - .generate( - request_with_max_tokens(model, Some(LONG_MAX_TOKENS)), - GenerateContext::new(ctx.clone(), None), - ) + .generate(request, GenerateContext::new(ctx.clone(), None)) .await .map_err(|e| GenerateFailed(e.to_string()))?; @@ -556,7 +755,9 @@ async fn check_cancellation_raw( #[cfg(test)] mod tests { use super::*; - use crate::engine::{EngineConfig, PreprocessedRequest}; + use crate::engine::{ + EngineConfig, LLMEngineOutput, LLMEngineOutputExt, PreprocessedRequest, usage, + }; use crate::error::DynamoError; use async_trait::async_trait; use futures::stream::BoxStream; @@ -618,6 +819,170 @@ mod tests { assert!(check_setup_metrics(&engine).await.is_ok()); } + #[derive(Clone, Copy)] + enum EncodeMockResponse { + Valid, + MissingEncoderResult, + WrongCount, + WrongFinish, + Tokens, + BadUsage, + Empty, + } + + struct EncodeConformanceMock { + response: EncodeMockResponse, + honor_cancel: bool, + } + + fn encode_mock(response: EncodeMockResponse) -> EncodeConformanceMock { + EncodeConformanceMock { + response, + honor_cancel: true, + } + } + + fn valid_encode_chunk() -> LLMEngineOutput { + LLMEngineOutput::encode_terminal(serde_json::Map::from_iter([( + "handle".to_string(), + serde_json::json!("sample-encoder:test"), + )])) + } + + #[async_trait] + impl LLMEngine for EncodeConformanceMock { + async fn start(&self, _worker_id: u64) -> Result { + Ok(EngineConfig { + model: "encode-mock".to_string(), + ..EngineConfig::default() + }) + } + + async fn generate( + &self, + request: PreprocessedRequest, + ctx: GenerateContext, + ) -> Result>, DynamoError> { + assert!( + request + .multi_modal_data + .as_ref() + .is_some_and(|data| !data.is_empty()), + "encode conformance request must contain multi_modal_data" + ); + assert_eq!( + request.mm_processor_kwargs, + Some(serde_json::json!({ "min_pixels": 64 })), + "encode conformance request must contain mm_processor_kwargs" + ); + let chunks = match self.response { + EncodeMockResponse::Valid => vec![Ok(valid_encode_chunk())], + EncodeMockResponse::MissingEncoderResult => vec![Ok(LLMEngineOutput::stop())], + EncodeMockResponse::WrongCount => { + vec![Ok(valid_encode_chunk()), Ok(valid_encode_chunk())] + } + EncodeMockResponse::WrongFinish => { + let mut chunk = valid_encode_chunk(); + chunk.finish_reason = Some(FinishReason::Length); + vec![Ok(chunk)] + } + EncodeMockResponse::Tokens => { + let mut chunk = valid_encode_chunk(); + chunk.token_ids = vec![1]; + vec![Ok(chunk)] + } + EncodeMockResponse::BadUsage => { + vec![Ok(valid_encode_chunk().with_usage(usage(3, 1)))] + } + EncodeMockResponse::Empty => vec![], + }; + let honor_cancel = self.honor_cancel; + let ctx = ctx.inner_arc(); + Ok(Box::pin(async_stream::stream! { + if honor_cancel && ctx.is_stopped() { + yield Ok(LLMEngineOutput::cancelled()); + return; + } + for chunk in chunks { + yield chunk; + } + })) + } + + async fn cleanup(&self) -> Result<(), DynamoError> { + Ok(()) + } + } + + #[tokio::test] + async fn encode_mock_without_usage_satisfies_conformance() { + run_encode_conformance(|| encode_mock(EncodeMockResponse::Valid)) + .await + .expect("encode conformance"); + } + + #[tokio::test] + async fn encode_conformance_rejects_missing_encoder_result() { + let result = + run_encode_conformance(|| encode_mock(EncodeMockResponse::MissingEncoderResult)).await; + assert!( + matches!(result, Err(EncoderResultExpectedObject)), + "expected EncoderResultExpectedObject, got {result:?}" + ); + } + + #[tokio::test] + async fn encode_conformance_rejects_wrong_chunk_count() { + let engine = encode_mock(EncodeMockResponse::WrongCount); + let result = check_encode_generate(&engine, "encode-mock").await; + assert!(matches!(result, Err(EncodeChunkCount { count: 2 }))); + } + + #[tokio::test] + async fn encode_conformance_rejects_wrong_finish_reason() { + let engine = encode_mock(EncodeMockResponse::WrongFinish); + let result = check_encode_generate(&engine, "encode-mock").await; + assert!(matches!(result, Err(EncodeTerminalExpected))); + } + + #[tokio::test] + async fn encode_conformance_rejects_generated_tokens() { + let engine = encode_mock(EncodeMockResponse::Tokens); + let result = check_encode_generate(&engine, "encode-mock").await; + assert!(matches!(result, Err(EncodeTokensNotEmpty { count: 1 }))); + } + + #[tokio::test] + async fn encode_conformance_rejects_inconsistent_usage() { + let engine = encode_mock(EncodeMockResponse::BadUsage); + let result = check_encode_generate(&engine, "encode-mock").await; + assert!(matches!(result, Err(EncodeUsageMismatch { .. }))); + } + + #[tokio::test] + async fn encode_conformance_rejects_ignored_cancellation() { + let engine = EncodeConformanceMock { + response: EncodeMockResponse::Valid, + honor_cancel: false, + }; + let result = check_cancellation( + &engine, + "encode-mock", + LlmConformanceMode::Encode, + Duration::from_millis(150), + ) + .await; + assert!(matches!(result, Err(CancellationIgnored))); + } + + #[tokio::test] + async fn encode_conformance_validates_concurrent_streams() { + let engine = encode_mock(EncodeMockResponse::Empty); + let result = + check_concurrent_generates(&engine, "encode-mock", LlmConformanceMode::Encode).await; + assert!(matches!(result, Err(EncodeChunkCount { count: 0 }))); + } + /// Minimal `RawEngine` for exercising the raw conformance kit itself. /// With `honor_cancel = false`, `generate` ignores `is_stopped()` so the /// cancellation check can be shown to have teeth. diff --git a/tests/runtime/test_sample_multimodal_smoke.py b/tests/runtime/test_sample_multimodal_smoke.py new file mode 100644 index 000000000000..5fd4813a36ef --- /dev/null +++ b/tests/runtime/test_sample_multimodal_smoke.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only process smokes for sample multimodal worker handoffs.""" + +from __future__ import annotations + +import os +import signal +import subprocess +import uuid +from pathlib import Path + +import pytest + +from tests.utils.port_utils import allocate_ports, deallocate_ports + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unified, + pytest.mark.timeout(270), +] + +REPO_ROOT = Path(__file__).parents[2] +LAUNCH_DIR = REPO_ROOT / "examples" / "backends" / "sample" / "launch" + + +@pytest.mark.parametrize( + "script_name", + ["multimodal_agg.sh", "multimodal_disagg.sh"], +) +def test_sample_multimodal_smoke(script_name, request, runtime_services_dynamic_ports): + del runtime_services_dynamic_ports + ports = allocate_ports(3, 18000) + request.addfinalizer(lambda: deallocate_ports(ports)) + + env = os.environ.copy() + env.update( + { + "DYN_SYSTEM_PORT": str(ports[0]), + "DYN_SYSTEM_PORT1": str(ports[0]), + "DYN_SYSTEM_PORT2": str(ports[1]), + "DYN_SYSTEM_PORT3": str(ports[2]), + "NAMESPACE": f"sample-mm-{uuid.uuid4().hex}", + } + ) + process = subprocess.Popen( + ["bash", str(LAUNCH_DIR / script_name)], + cwd=REPO_ROOT, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + try: + output, _ = process.communicate(timeout=90) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGTERM) + try: + output, _ = process.communicate(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + output, _ = process.communicate(timeout=5) + pytest.fail(f"{script_name} timed out\n{output}") + + assert process.returncode == 0, f"{script_name} failed\n{output}" From 95530f12e0b0a510102aeea517ba362b986523e6 Mon Sep 17 00:00:00 2001 From: Julien Mancuso <161955438+julienmancuso@users.noreply.github.com> Date: Tue, 30 Jun 2026 14:50:20 -0600 Subject: [PATCH 011/320] chore(operator): pin PR deploy tests to commit image (#11097) --- .../actions/setup-dynamo-operator/action.yml | 119 ++++++++++++++++-- .github/workflows/pr.yaml | 20 +-- 2 files changed, 122 insertions(+), 17 deletions(-) diff --git a/.github/actions/setup-dynamo-operator/action.yml b/.github/actions/setup-dynamo-operator/action.yml index b9f1bcadb926..2af8e2bd2789 100644 --- a/.github/actions/setup-dynamo-operator/action.yml +++ b/.github/actions/setup-dynamo-operator/action.yml @@ -19,9 +19,9 @@ inputs: required: false default: '' operator_tag: - description: 'Operator image tag (default: main-operator)' + description: 'Operator image tag. Required when registry is set.' required: false - default: 'main-operator' + default: '' hf_token: description: 'HuggingFace token for model access' required: false @@ -340,16 +340,16 @@ runs: IMAGE_ARGS=() if [ -n "${REGISTRY}" ]; then + if [ -z "${OPERATOR_TAG}" ]; then + echo "::error::operator_tag is required when registry is set" + exit 1 + fi OPERATOR_REPO="${REGISTRY}/ai-dynamo/dynamo" echo "Using operator image: ${OPERATOR_REPO}:${OPERATOR_TAG}" - # OPERATOR_TAG is a mutable, floating tag (e.g. main-operator) that is - # re-pushed on every main build. The chart default - # imagePullPolicy=IfNotPresent means a CI node that cached an older copy - # never re-pulls it, so a PR that doesn't rebuild the operator can run a - # stale crd-apply/manager binary (e.g. missing a newly-added flag) and - # land in CrashLoopBackOff. Force Always so the floating tag is refreshed - # on every pod start; the chart applies this value to both the manager - # and the crd-apply init container. + # Automated CI passes a commit-scoped tag so the chart, CRDs, + # and manager binary stay in sync. Manual workflows may still pass a + # mutable tag, so force a registry lookup on every pod start. The chart + # applies this policy to both the manager and crd-apply init container. IMAGE_ARGS+=( --set dynamo-operator.controllerManager.manager.image.repository="${OPERATOR_REPO}" --set dynamo-operator.controllerManager.manager.image.tag="${OPERATOR_TAG}" @@ -440,6 +440,105 @@ runs: crd/dynamographdeployments.nvidia.com \ --timeout=120s + - name: Verify operator image and admission webhooks + shell: bash + env: + NAMESPACE: ${{ steps.resolve-names.outputs.namespace }} + run: | + set -euo pipefail + VKUBECONFIG=${{ github.workspace }}/.kubeconfig-vcluster + + OPERATOR_POD=$(kubectl get pods -n "${NAMESPACE}" -o name \ + | grep 'dynamo-platform-dynamo-operator-controller-manager' \ + | head -1 || true) + if [ -z "${OPERATOR_POD}" ]; then + echo "::error::Ready operator pod disappeared before verification" + exit 1 + fi + + echo "::group::Resolved operator image identity" + kubectl get "${OPERATOR_POD}" -n "${NAMESPACE}" \ + -o jsonpath='{range .status.initContainerStatuses[*]}initContainer={.name} image={.image} imageID={.imageID}{"\n"}{end}{range .status.containerStatuses[*]}container={.name} image={.image} imageID={.imageID}{"\n"}{end}' + echo "::endgroup::" + + echo "::group::DynamoGraphDeployment admission smoke tests" + smoke_v1alpha1() { + kubectl --kubeconfig="${VKUBECONFIG}" --request-timeout=15s \ + create --namespace=default --dry-run=server \ + -o json \ + -f - <<'EOF' + apiVersion: nvidia.com/v1alpha1 + kind: DynamoGraphDeployment + metadata: + name: ci-webhook-smoke-v1alpha1 + spec: + services: + Smoke: + componentType: worker + replicas: 0 + EOF + } + + smoke_v1beta1() { + kubectl --kubeconfig="${VKUBECONFIG}" --request-timeout=15s \ + create --namespace=default --dry-run=server \ + -o json \ + -f - <<'EOF' + apiVersion: nvidia.com/v1beta1 + kind: DynamoGraphDeployment + metadata: + name: ci-webhook-smoke-v1beta1 + spec: + components: + - name: Smoke + type: worker + replicas: 0 + EOF + } + + run_admission_smoke_test() { + local version="$1" + local output + local origin_version + local deadline=$((SECONDS + 120)) + local attempt=1 + shift + + while true; do + if output=$("$@" 2>&1); then + origin_version=$(sed -n 's/.*"nvidia\.com\/dynamo-operator-origin-version":[[:space:]]*"\([^"]*\)".*/\1/p' <<<"${output}" | tail -1) + if [ -z "${origin_version}" ]; then + echo "::error::${version} request succeeded without the operator origin-version mutation" + printf '%s\n' "${output}" >&2 + return 1 + fi + echo "${version} admission succeeded (operator origin version: ${origin_version})" + return 0 + fi + + if ! grep -Eiq \ + 'no endpoints available|connection refused|connection reset by peer|TLS handshake timeout|service .* not found|unexpected EOF|^EOF$' \ + <<<"${output}"; then + printf '%s\n' "${output}" >&2 + return 1 + fi + + if [ "${SECONDS}" -ge "${deadline}" ]; then + echo "::error::${version} admission smoke test did not become reachable within 120s" + printf '%s\n' "${output}" >&2 + return 1 + fi + + echo "::warning::${version} admission endpoint not ready (attempt ${attempt}); retrying in 2s" + attempt=$((attempt + 1)) + sleep 2 + done + } + + run_admission_smoke_test v1alpha1 smoke_v1alpha1 + run_admission_smoke_test v1beta1 smoke_v1beta1 + echo "::endgroup::" + - name: Debug deployment failure if: failure() shell: bash diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 41c0ef27bdf9..0c8db5c8cb86 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -131,7 +131,13 @@ jobs: operator: needs: changed-files - if: needs.changed-files.outputs.operator == 'true' + if: | + needs.changed-files.outputs.operator == 'true' || + needs.changed-files.outputs.vllm == 'true' || + needs.changed-files.outputs.sglang == 'true' || + needs.changed-files.outputs.trtllm == 'true' || + needs.changed-files.outputs.deploy == 'true' || + needs.changed-files.outputs.snapshot == 'true' name: Operator runs-on: prod-default-v2 outputs: @@ -859,7 +865,7 @@ jobs: needs.changed-files.outputs.sglang == 'true' || needs.changed-files.outputs.trtllm == 'true' || needs.changed-files.outputs.deploy == 'true') && - (needs.operator.result == 'success' || needs.operator.result == 'skipped') + needs.operator.result == 'success' needs: [changed-files, operator] runs-on: prod-deploy-tester-v1 permissions: @@ -877,7 +883,7 @@ jobs: uses: ./.github/actions/setup-dynamo-operator with: registry: ${{ secrets.AZURE_ACR_HOSTNAME }} - operator_tag: ${{ needs.operator.result == 'success' && needs.operator.outputs.operator_default_tag || 'main-operator' }} + operator_tag: ${{ needs.operator.outputs.operator_default_tag }} hf_token: ${{ secrets.HF_TOKEN }} dockerhub_username: ${{ secrets.DOCKERHUB_LOGIN_USER }} dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} @@ -937,7 +943,7 @@ jobs: (needs.changed-files.outputs.operator == 'true' || needs.changed-files.outputs.snapshot == 'true' || needs.changed-files.outputs.deploy == 'true') && - (needs.operator.result == 'success' || needs.operator.result == 'skipped') + needs.operator.result == 'success' needs: [changed-files, operator] runs-on: prod-deploy-tester-v1 permissions: @@ -955,7 +961,7 @@ jobs: vcluster_name: ci-${{ github.run_id }}-checkpoint-vllm vcluster_namespace: gh-id-${{ github.run_id }}-checkpoint-vllm-dt registry: ${{ secrets.AZURE_ACR_HOSTNAME }} - operator_tag: ${{ needs.operator.result == 'success' && needs.operator.outputs.operator_default_tag || 'main-operator' }} + operator_tag: ${{ needs.operator.outputs.operator_default_tag }} hf_token: ${{ secrets.HF_TOKEN }} dockerhub_username: ${{ secrets.DOCKERHUB_LOGIN_USER }} dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} @@ -969,7 +975,7 @@ jobs: (needs.changed-files.outputs.operator == 'true' || needs.changed-files.outputs.snapshot == 'true' || needs.changed-files.outputs.deploy == 'true') && - (needs.operator.result == 'success' || needs.operator.result == 'skipped') + needs.operator.result == 'success' needs: [changed-files, operator] runs-on: prod-deploy-tester-v1 permissions: @@ -987,7 +993,7 @@ jobs: vcluster_name: ci-${{ github.run_id }}-checkpoint-sglang vcluster_namespace: gh-id-${{ github.run_id }}-checkpoint-sglang-dt registry: ${{ secrets.AZURE_ACR_HOSTNAME }} - operator_tag: ${{ needs.operator.result == 'success' && needs.operator.outputs.operator_default_tag || 'main-operator' }} + operator_tag: ${{ needs.operator.outputs.operator_default_tag }} hf_token: ${{ secrets.HF_TOKEN }} dockerhub_username: ${{ secrets.DOCKERHUB_LOGIN_USER }} dockerhub_password: ${{ secrets.DOCKERHUB_ACCESS_TOKEN }} From 3f7716378dc4e25668fa404df710df90591cd1b4 Mon Sep 17 00:00:00 2001 From: juju <45225513+jooe0824@users.noreply.github.com> Date: Wed, 1 Jul 2026 05:50:44 +0900 Subject: [PATCH 012/320] docs(benchmarks): fix misleading isl/osl preset docs in sin_load_generator (#10791) Signed-off-by: jooe0824 Co-authored-by: jooe0824 Co-authored-by: Ryan McCormick --- benchmarks/sin_load_generator/README.md | 10 +++++----- benchmarks/sin_load_generator/sin_synth.py | 20 ++++++++++++++++---- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/benchmarks/sin_load_generator/README.md b/benchmarks/sin_load_generator/README.md index 3a16eef3799f..d3b917e2c8e0 100644 --- a/benchmarks/sin_load_generator/README.md +++ b/benchmarks/sin_load_generator/README.md @@ -54,7 +54,7 @@ Note the phase shift of `-π/2` is to make the request rate start from the minim ### Input/Output Sequence Length Parameters The script will generate load with requests sampled from two preset ISL/OSL combinations. -The ISL/OSL ratio defines how much of requests follow the first preset ISL/OSL pattern. ISl/OSL 0 means all requests follow the first preset ISL/OSL pattern, while ISL/OSL 1 means all requests follow the second preset ISL/OSL pattern. +The ISL/OSL ratio defines how much of requests follow the first preset ISL/OSL pattern. An ISL/OSL ratio of 1 means all requests follow the first preset ISL/OSL pattern, while a ratio of 0 means all requests follow the second preset ISL/OSL pattern. The ISL/OSL ratio follows a sinusoidal pattern: ``` @@ -64,16 +64,16 @@ isl-osl-ratio(t) = (min + max) / 2 + (max - min) / 2 * sin(2 * π / period * t - Similarly, the phase shift of `-π/2` is to make the ISL/OSL ratio start from the minimum at `t = 0`. - `--isl1 INT` (default: 100) - - Minimum input sequence length + - Input sequence length of the first preset ISL/OSL pair - `--osl1 INT` (default: 2000) - - Minimum output sequence length + - Output sequence length of the first preset ISL/OSL pair - `--isl2 INT` (default: 5000) - - Maximum input sequence length + - Input sequence length of the second preset ISL/OSL pair - `--osl2 INT` (default: 100) - - Maximum output sequence length + - Output sequence length of the second preset ISL/OSL pair - `--isl-osl-ratio-min FLOAT` (default: 0.2) - Minimum ratio of input sequence length to output sequence length diff --git a/benchmarks/sin_load_generator/sin_synth.py b/benchmarks/sin_load_generator/sin_synth.py index 76212d11b55c..dbac3f386172 100644 --- a/benchmarks/sin_load_generator/sin_synth.py +++ b/benchmarks/sin_load_generator/sin_synth.py @@ -135,16 +135,28 @@ def get_isl_osl(t): # isl-osl-ratio(t) = (min + max) / 2 + (max - min) / 2 * sin(2 * pi / period * t - pi / 2) # Then, we sample [isl1/osl1, isl2/osl2] from the distribution [isl-osl-ratio(t), 1 - isl-osl-ratio(t)] parser.add_argument( - "--isl1", type=int, default=100, help="Minimum input sequence length" + "--isl1", + type=int, + default=100, + help="Input sequence length of the first preset ISL/OSL pair", ) parser.add_argument( - "--osl1", type=int, default=2000, help="Minimum output sequence length" + "--osl1", + type=int, + default=2000, + help="Output sequence length of the first preset ISL/OSL pair", ) parser.add_argument( - "--isl2", type=int, default=5000, help="Maximum input sequence length" + "--isl2", + type=int, + default=5000, + help="Input sequence length of the second preset ISL/OSL pair", ) parser.add_argument( - "--osl2", type=int, default=100, help="Maximum output sequence length" + "--osl2", + type=int, + default=100, + help="Output sequence length of the second preset ISL/OSL pair", ) parser.add_argument( "--isl-osl-ratio-min", From fa6894c9a8a2330a85151ddc8f64df7e9326a57c Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Tue, 30 Jun 2026 13:56:28 -0700 Subject: [PATCH 013/320] perf(mocker): reduce KV block bookkeeping overhead (#11095) Signed-off-by: jthomson04 --- lib/kvbm-logical/README.md | 8 +- lib/kvbm-logical/src/manager/mod.rs | 61 +- lib/kvbm-logical/src/manager/tests.rs | 311 ++++++++++ lib/kvbm-logical/src/metrics/collector.rs | 8 +- lib/kvbm-logical/src/metrics/stats.rs | 6 +- lib/kvbm-logical/src/pools/store.rs | 90 ++- lib/kvbm-logical/src/registry/mod.rs | 112 +++- lib/kvbm-logical/src/registry/tests.rs | 130 +++- lib/mocker/src/common/sequence.rs | 86 +-- lib/mocker/src/kv_manager/kvbm_backend.rs | 687 +++++++++++++++++++--- 10 files changed, 1320 insertions(+), 179 deletions(-) diff --git a/lib/kvbm-logical/README.md b/lib/kvbm-logical/README.md index 0a30513b5361..0b6147773209 100644 --- a/lib/kvbm-logical/README.md +++ b/lib/kvbm-logical/README.md @@ -102,10 +102,10 @@ All metrics carry a `pool` label identifying the storage tier. | `kvbm_duplicate_blocks_total` | Total duplicate blocks created (Allow policy) | | `kvbm_registration_dedup_total` | Total block registrations deduplicated (Reject policy) | | `kvbm_stagings_total` | Total MutableBlock → CompleteBlock transitions | -| `kvbm_match_hashes_requested_total` | Total hashes requested in match_blocks calls | -| `kvbm_match_blocks_returned_total` | Total blocks returned from match_blocks calls | -| `kvbm_scan_hashes_requested_total` | Total hashes requested in scan_matches calls | -| `kvbm_scan_blocks_returned_total` | Total blocks returned from scan_matches calls | +| `kvbm_match_hashes_requested_total` | Total input hash occurrences requested in `match_blocks` and `match_blocks_scattered` calls | +| `kvbm_match_blocks_returned_total` | Total block hit occurrences returned from `match_blocks` and `match_blocks_scattered` calls | +| `kvbm_scan_hashes_requested_total` | Total input hash occurrences requested in `scan_matches` calls | +| `kvbm_scan_blocks_returned_total` | Total distinct matching hashes returned from `scan_matches` calls | ### Gauges diff --git a/lib/kvbm-logical/src/manager/mod.rs b/lib/kvbm-logical/src/manager/mod.rs index 5f7150fb78a7..c2b8b20f036b 100644 --- a/lib/kvbm-logical/src/manager/mod.rs +++ b/lib/kvbm-logical/src/manager/mod.rs @@ -99,10 +99,9 @@ impl BlockManager { .block_registry .register_sequence_hashes(blocks.iter().map(CompleteBlock::sequence_hash)); let batch_size = blocks.len(); - let registered = self.store.register_completed_blocks( - blocks.into_iter().zip(handles).collect(), - self.duplication_policy, - ); + let registered = + self.store + .register_completed_blocks(blocks, handles, self.duplication_policy); // The offline settlement bridge observes this counter as a // publication watermark, so publish only after every store transition // and presence marker in the batch is complete. @@ -168,8 +167,60 @@ impl BlockManager { matched } + /// Scattered batch match: resolves every input hash against the active or + /// inactive pool without stopping at a miss. + /// + /// The returned vector is aligned with `seq_hash`: each hit is `Some`, + /// each miss is `None`, and input order and duplicates are preserved. The + /// complete batch is resolved under one store-mutex acquisition. Frequency + /// tracking is applied after releasing that lock, exactly once per hit + /// (including repeated hashes and inactive resurrections). + /// + /// This operation contributes to the existing match metrics. Requested + /// and returned values are counted as occurrences, so repeated input + /// hashes and their repeated hits are counted repeatedly. + pub fn match_blocks_scattered( + &self, + seq_hash: &[SequenceHash], + ) -> Vec>> { + self.metrics + .inc_match_hashes_requested(seq_hash.len() as u64); + + if seq_hash.is_empty() { + self.metrics.inc_match_blocks_returned(0); + return Vec::new(); + } + + // ONE store-lock acquisition for all active+inactive probes, including + // misses and repeated hashes. + let inners = self.store.match_scattered_locked_batch(seq_hash); + + // Keep TinyLFU work outside the store critical section. A duplicate + // input is a duplicate access, so each returned occurrence is touched. + if self.block_registry.has_frequency_tracking() { + for inner in inners.iter().flatten() { + self.block_registry.touch(inner.sequence_hash()); + } + } + + let hit_count = inners.iter().filter(|inner| inner.is_some()).count(); + let matched = inners + .into_iter() + .map(|inner| inner.map(ImmutableBlock::from_inner)) + .collect(); + + self.metrics.inc_match_blocks_returned(hit_count as u64); + tracing::debug!( + num_hashes = seq_hash.len(), + total_matched = hit_count, + "match_blocks_scattered result" + ); + matched + } + /// Scatter-gather scan: finds all blocks matching any hash, without - /// stopping on misses. + /// stopping on misses. Requested hashes are counted as input occurrences, + /// while returned blocks are counted as distinct hashes in the result map. pub fn scan_matches( &self, seq_hashes: &[SequenceHash], diff --git a/lib/kvbm-logical/src/manager/tests.rs b/lib/kvbm-logical/src/manager/tests.rs index 6f60b2bad205..7c6f3432ea0e 100644 --- a/lib/kvbm-logical/src/manager/tests.rs +++ b/lib/kvbm-logical/src/manager/tests.rs @@ -896,6 +896,129 @@ mod registration_tests { assert_eq!(snap.stagings, 3); } + #[test] + fn test_batch_registration_mixed_rejects_preserves_order_presence_and_guards() { + let registry = BlockRegistry::new(); + let manager = BlockManager::::builder() + .block_count(5) + .block_size(4) + .registry(registry) + .duplication_policy(BlockDuplicationPolicy::Reject) + .build() + .expect("Should build manager"); + + let token_a = create_test_token_block_from_iota(100); + let token_b = create_test_token_block_from_iota(200); + let token_c = create_test_token_block_from_iota(300); + let hash_a = token_a.kvbm_sequence_hash(); + let hash_b = token_b.kvbm_sequence_hash(); + let hash_c = token_c.kvbm_sequence_hash(); + + let primary_a = manager + .allocate_blocks(1) + .expect("allocate primary") + .pop() + .unwrap() + .complete(&token_a) + .unwrap(); + let primary_a_id = primary_a.block_id(); + let primary_a = manager.register_blocks(vec![primary_a]).pop().unwrap(); + + let mut candidates = manager + .allocate_blocks(4) + .expect("allocate batch candidates") + .into_iter(); + let candidate_b = candidates.next().unwrap(); + let duplicate_a = candidates.next().unwrap(); + let duplicate_b = candidates.next().unwrap(); + let candidate_c = candidates.next().unwrap(); + let candidate_b_id = candidate_b.block_id(); + let duplicate_a_id = duplicate_a.block_id(); + let duplicate_b_id = duplicate_b.block_id(); + let candidate_c_id = candidate_c.block_id(); + + let registered = manager.register_blocks(vec![ + candidate_b.complete(&token_b).unwrap(), + duplicate_a.complete(&token_a).unwrap(), + duplicate_b.complete(&token_b).unwrap(), + candidate_c.complete(&token_c).unwrap(), + ]); + + assert_eq!( + registered + .iter() + .map(|block| (block.sequence_hash(), block.block_id())) + .collect::>(), + vec![ + (hash_b, candidate_b_id), + (hash_a, primary_a_id), + (hash_b, candidate_b_id), + (hash_c, candidate_c_id), + ] + ); + assert_eq!( + manager + .block_registry() + .check_presence::(&[hash_a, hash_b, hash_c]), + vec![(hash_a, true), (hash_b, true), (hash_c, true)] + ); + assert_eq!(manager.metrics().snapshot().registration_dedup, 2); + + // Both rejected candidates must have dropped their re-armed guards + // after the store critical section and returned to the reset pool. + assert_eq!(manager.available_blocks(), 2); + let mut returned_ids = manager + .allocate_blocks(2) + .expect("rejected candidates returned to reset") + .into_iter() + .map(|block| block.block_id()) + .collect::>(); + returned_ids.sort_unstable(); + let mut rejected_ids = vec![duplicate_a_id, duplicate_b_id]; + rejected_ids.sort_unstable(); + assert_eq!(returned_ids, rejected_ids); + + drop(primary_a); + drop(registered); + } + + #[test] + fn test_batch_registration_allow_marks_each_same_hash_slot_present() { + let registry = BlockRegistry::new(); + let manager = BlockManager::::builder() + .block_count(2) + .block_size(4) + .registry(registry) + .duplication_policy(BlockDuplicationPolicy::Allow) + .build() + .expect("Should build manager"); + + let token = create_test_token_block_from_iota(400); + let seq_hash = token.kvbm_sequence_hash(); + let completed = manager + .allocate_blocks(2) + .expect("allocate same-hash batch") + .into_iter() + .map(|block| block.complete(&token).unwrap()) + .collect(); + + let mut registered = manager.register_blocks(completed); + assert_eq!(registered.len(), 2); + assert_ne!(registered[0].block_id(), registered[1].block_id()); + assert_eq!(manager.metrics().snapshot().duplicate_blocks, 1); + + // The duplicate releases one presence reference. Presence must remain + // set for the primary, proving both same-batch outcomes were marked. + let duplicate = registered.pop().unwrap(); + drop(duplicate); + assert_eq!( + manager + .block_registry() + .check_presence::(&[seq_hash]), + vec![(seq_hash, true)] + ); + } + #[rstest] #[case(BlockDuplicationPolicy::Allow, 200, "allow", false)] #[case(BlockDuplicationPolicy::Reject, 300, "reject", true)] @@ -1376,6 +1499,149 @@ mod single_lock_match_tests { } } +// ============================================================================ +// ALIGNED SCATTERED MATCH TESTS +// ============================================================================ + +mod scattered_match_tests { + use super::*; + + fn create_backend_manager( + block_count: usize, + backend_builder: fn( + BlockManagerConfigBuilder, + ) -> BlockManagerConfigBuilder, + ) -> BlockManager { + let registry = BlockRegistry::builder() + .frequency_tracker(FrequencyTrackingCapacity::default().create_tracker()) + .build(); + backend_builder( + BlockManager::::builder() + .block_count(block_count) + .block_size(4) + .registry(registry), + ) + .build() + .expect("build manager") + } + + fn register_one( + manager: &BlockManager, + base: u32, + ) -> (SequenceHash, ImmutableBlock) { + let token_block = create_token_block(&[base, base + 1, base + 2, base + 3]); + let seq_hash = token_block.kvbm_sequence_hash(); + let mutable = manager + .allocate_blocks(1) + .expect("allocate") + .into_iter() + .next() + .unwrap(); + let complete = mutable.complete(&token_block).expect("complete"); + let immutable = manager + .register_blocks(vec![complete]) + .into_iter() + .next() + .unwrap(); + (seq_hash, immutable) + } + + /// Active and inactive hits remain aligned around a miss; duplicate + /// inactive hashes are returned at both positions and later hits are not + /// truncated. Run against every supported inactive backend. + #[rstest] + #[case("lru", |b: BlockManagerConfigBuilder| b.with_lru_backend())] + #[case("multi_lru", |b: BlockManagerConfigBuilder| b.with_multi_lru_backend())] + #[case("lineage", |b: BlockManagerConfigBuilder| b.with_lineage_backend())] + fn scattered_preserves_alignment_misses_duplicates_and_later_hits( + #[case] _backend_name: &str, + #[case] backend_builder: fn( + BlockManagerConfigBuilder, + ) -> BlockManagerConfigBuilder, + ) { + let manager = create_backend_manager(4, backend_builder); + let (active_hash, _active) = register_one(&manager, 20_000); + let (inactive_hash, inactive) = register_one(&manager, 20_010); + let (later_hash, _later) = register_one(&manager, 20_020); + drop(inactive); + + let miss = create_token_block(&[90_000, 90_001, 90_002, 90_003]).kvbm_sequence_hash(); + let input = [active_hash, miss, inactive_hash, inactive_hash, later_hash]; + let before = manager.metrics().snapshot(); + + let matched = manager.match_blocks_scattered(&input); + + assert_eq!(matched.len(), input.len()); + assert_eq!(matched[0].as_ref().unwrap().sequence_hash(), active_hash); + assert!(matched[1].is_none(), "miss must retain its aligned slot"); + assert_eq!(matched[2].as_ref().unwrap().sequence_hash(), inactive_hash); + assert_eq!(matched[3].as_ref().unwrap().sequence_hash(), inactive_hash); + assert_eq!( + matched[2].as_ref().unwrap().block_id(), + matched[3].as_ref().unwrap().block_id(), + "duplicate hashes must resolve to the same physical primary" + ); + assert_eq!(matched[4].as_ref().unwrap().sequence_hash(), later_hash); + + let after = manager.metrics().snapshot(); + assert_eq!( + after.match_hashes_requested - before.match_hashes_requested, + input.len() as u64 + ); + assert_eq!( + after.match_blocks_returned - before.match_blocks_returned, + 4, + "match return metric counts hit occurrences, including duplicates" + ); + assert_eq!( + after.scan_hashes_requested - before.scan_hashes_requested, + 0 + ); + assert_eq!(after.scan_blocks_returned - before.scan_blocks_returned, 0); + } + + #[test] + fn scattered_empty_input_returns_empty() { + let manager = create_test_manager(1); + let before = manager.metrics().snapshot(); + assert!(manager.match_blocks_scattered(&[]).is_empty()); + let after = manager.metrics().snapshot(); + assert_eq!( + after.match_hashes_requested - before.match_hashes_requested, + 0 + ); + assert_eq!( + after.match_blocks_returned - before.match_blocks_returned, + 0 + ); + assert_eq!( + after.scan_hashes_requested - before.scan_hashes_requested, + 0 + ); + assert_eq!(after.scan_blocks_returned - before.scan_blocks_returned, 0); + } + + #[test] + fn scattered_touches_once_per_hit_occurrence() { + let (manager, metered) = crate::testing::create_test_manager_metered::(3); + let (active_hash, _active) = register_one(&manager, 30_000); + let (inactive_hash, inactive) = register_one(&manager, 30_010); + drop(inactive); + let miss = create_token_block(&[91_000, 91_001, 91_002, 91_003]).kvbm_sequence_hash(); + + metered.reset(); + let matched = + manager.match_blocks_scattered(&[active_hash, miss, inactive_hash, inactive_hash]); + + assert_eq!(matched.iter().filter(|block| block.is_some()).count(), 3); + assert_eq!( + metered.touches(), + 3, + "each aligned hit occurrence must touch exactly once; misses never touch" + ); + } +} + // ============================================================================ // IMMUTABLE BLOCK AND WEAK BLOCK TESTS // ============================================================================ @@ -2287,6 +2553,51 @@ mod capacity_lifecycle_tests { mod scan_matches_tests { use super::*; + #[test] + fn scan_metrics_count_input_occurrences_and_distinct_results() { + let manager = create_test_manager(2); + let token_block = create_iota_token_block(12_000, 4); + let hash = token_block.kvbm_sequence_hash(); + let mutable = manager + .allocate_blocks(1) + .expect("allocate") + .into_iter() + .next() + .unwrap(); + let complete = mutable.complete(&token_block).expect("complete"); + let _active = manager + .register_blocks(vec![complete]) + .into_iter() + .next() + .unwrap(); + let missing_hash = create_iota_token_block(99_000, 4).kvbm_sequence_hash(); + let before = manager.metrics().snapshot(); + + let found = manager.scan_matches(&[hash, hash, missing_hash], true); + + assert_eq!(found.len(), 1); + assert!(found.contains_key(&hash)); + let after = manager.metrics().snapshot(); + assert_eq!( + after.scan_hashes_requested - before.scan_hashes_requested, + 3, + "scan request metrics count duplicate input occurrences" + ); + assert_eq!( + after.scan_blocks_returned - before.scan_blocks_returned, + 1, + "scan return metrics count distinct result-map entries" + ); + assert_eq!( + after.match_hashes_requested - before.match_hashes_requested, + 0 + ); + assert_eq!( + after.match_blocks_returned - before.match_blocks_returned, + 0 + ); + } + #[test] fn test_scan_matches_with_pool_size_gauges() { let manager = create_test_manager(10); diff --git a/lib/kvbm-logical/src/metrics/collector.rs b/lib/kvbm-logical/src/metrics/collector.rs index 885983e1bda5..92af27bad83c 100644 --- a/lib/kvbm-logical/src/metrics/collector.rs +++ b/lib/kvbm-logical/src/metrics/collector.rs @@ -45,19 +45,19 @@ const COUNTER_DEFS: &[(&str, &str)] = &[ ), ( "kvbm_match_hashes_requested_total", - "Total hashes requested in match_blocks calls", + "Total input hash occurrences requested in match_blocks and match_blocks_scattered calls", ), ( "kvbm_match_blocks_returned_total", - "Total blocks returned from match_blocks calls", + "Total block hit occurrences returned from match_blocks and match_blocks_scattered calls", ), ( "kvbm_scan_hashes_requested_total", - "Total hashes requested in scan_matches calls", + "Total input hash occurrences requested in scan_matches calls", ), ( "kvbm_scan_blocks_returned_total", - "Total blocks returned from scan_matches calls", + "Total distinct matching hashes returned from scan_matches calls", ), ( "kvbm_eager_primary_to_inactive_total", diff --git a/lib/kvbm-logical/src/metrics/stats.rs b/lib/kvbm-logical/src/metrics/stats.rs index 1dcef201f7f9..9279ca2502a0 100644 --- a/lib/kvbm-logical/src/metrics/stats.rs +++ b/lib/kvbm-logical/src/metrics/stats.rs @@ -38,9 +38,11 @@ pub struct StatsSnapshot { pub allocation_rate: f64, /// Evictions per second. pub eviction_rate: f64, - /// Ratio of blocks returned to hashes requested in match_blocks. + /// Ratio of block hit occurrences returned to input hash occurrences + /// requested in match_blocks and match_blocks_scattered. pub match_hit_rate: f64, - /// Ratio of blocks returned to hashes requested in scan_matches. + /// Ratio of distinct matching hashes returned to input hash occurrences + /// requested in scan_matches. pub scan_hit_rate: f64, /// Rate of change of allocation_rate (d(alloc_rate)/dt). pub allocation_gradient: f64, diff --git a/lib/kvbm-logical/src/pools/store.rs b/lib/kvbm-logical/src/pools/store.rs index 19596834d436..f6b49eb3dc2a 100644 --- a/lib/kvbm-logical/src/pools/store.rs +++ b/lib/kvbm-logical/src/pools/store.rs @@ -612,6 +612,31 @@ impl BlockStore { out } + /// Batched active-or-inactive scattered lookup under **one** store-mutex + /// acquisition. Unlike [`match_prefix_locked_batch`](Self::match_prefix_locked_batch), + /// this preserves one output position per input hash and continues after + /// misses. + /// + /// Each position uses [`acquire_for_hash_locked`](Self::acquire_for_hash_locked), + /// so active lookup, eager `Primary -> Inactive` recovery, and inactive + /// resurrection remain atomic with respect to every other position in the + /// batch. Repeated hashes are resolved independently: an inactive hit in + /// the first occurrence is resurrected, and subsequent occurrences clone + /// that now-active primary. + /// + /// Frequency tracking is deliberately disabled while the store lock is + /// held. The caller applies one touch per hit after this method returns. + pub(crate) fn match_scattered_locked_batch( + self: &Arc, + hashes: &[SequenceHash], + ) -> Vec>>> { + let mut inner = self.inner.lock(); + hashes + .iter() + .map(|&hash| self.acquire_for_hash_locked(&mut inner, hash, /*touch*/ false)) + .collect() + } + /// Atomic registration of a [`CompleteBlock`]: lookup-then-transition /// under one store-mutex acquisition. Closes the register-vs-register /// race for the same sequence hash. @@ -632,15 +657,17 @@ impl BlockStore { block.disarm(); let mut inner = self.inner.lock(); - let (result, presence_added) = - self.register_completed_block_locked(&mut inner, &mut block, &handle, policy); + let result = self.register_completed_block_locked(&mut inner, &mut block, &handle, policy); drop(inner); // mark_present takes the attachments lock; lock-order // (attachments → store) is satisfied because the store lock has // already been released. Skip on Reject — no new presence-bearing - // slot was created. + // slot was created. The locked helper guarantees that fresh/Allow + // outcomes use the candidate slot, while Reject returns the distinct + // existing primary (enforced by its same-block collision assertion). + let presence_added = result.block_id() == block.block_id(); if presence_added { handle.mark_present::(); } @@ -655,37 +682,56 @@ impl BlockStore { /// Register a batch while acquiring the store mutex only once. pub(crate) fn register_completed_blocks( self: &Arc, - blocks: Vec<(CompleteBlock, BlockRegistrationHandle)>, + mut blocks: Vec>, + handles: Vec, policy: BlockDuplicationPolicy, ) -> Vec>> { - let mut registrations = Vec::with_capacity(blocks.len()); - let mut inner = self.inner.lock(); - for (mut block, handle) in blocks { - block.disarm(); - let (result, presence_added) = - self.register_completed_block_locked(&mut inner, &mut block, &handle, policy); - registrations.push((block, handle, presence_added, result)); - } - drop(inner); + assert_eq!( + blocks.len(), + handles.len(), + "each completed block must have a registration handle" + ); + + let outcomes = { + let mut outcomes = Vec::with_capacity(blocks.len()); + let mut inner = self.inner.lock(); + for (block, handle) in blocks.iter_mut().zip(&handles) { + block.disarm(); + outcomes + .push(self.register_completed_block_locked(&mut inner, block, handle, policy)); + } + outcomes + }; - let mut results = Vec::with_capacity(registrations.len()); - for (block, handle, presence_added, result) in registrations { + // Presence attachments and re-armed Reject-guard drops both acquire + // locks outside the store. Keep them after the one batch critical + // section to preserve attachments -> store lock ordering. + for ((block, handle), outcome) in blocks.iter().zip(&handles).zip(&outcomes) { + // Fresh and Allow outcomes retain the candidate block ID; Reject + // returns the different existing primary ID and re-arms `block`. + let presence_added = outcome.block_id() == block.block_id(); if presence_added { handle.mark_present::(); } - drop(block); - results.push(result); } - results + drop(blocks); + outcomes } + /// Register a completed candidate while the store mutex is held. + /// + /// The returned block ID equals the candidate block ID exactly when this + /// call creates a presence-bearing slot (fresh primary or allowed + /// duplicate). A rejected duplicate re-arms the candidate guard and + /// returns the existing primary, whose distinct ID is enforced by the + /// same-block collision assertion below. fn register_completed_block_locked( self: &Arc, inner: &mut BlockStoreInner, block: &mut CompleteBlock, handle: &BlockRegistrationHandle, policy: BlockDuplicationPolicy, - ) -> (Arc>, bool) { + ) -> Arc> { let block_id = block.block_id(); let seq_hash = block.sequence_hash(); debug_assert_eq!(seq_hash, handle.seq_hash()); @@ -716,12 +762,12 @@ impl BlockStore { inner: Arc::downgrade(&inner_arc), }; self.metrics.inc_duplicate_blocks(); - (inner_arc, true) + inner_arc } BlockDuplicationPolicy::Reject => { self.metrics.inc_registration_dedup(); block.rearm(); - (existing_primary, false) + existing_primary } }; } @@ -738,7 +784,7 @@ impl BlockStore { inner: Arc::downgrade(&inner_arc), }; inner.active_by_hash.insert(seq_hash, block_id); - (inner_arc, true) + inner_arc } /// Internal helper: under the store lock, transition a Primary slot to diff --git a/lib/kvbm-logical/src/registry/mod.rs b/lib/kvbm-logical/src/registry/mod.rs index 862131517b85..c84b15f0fc6d 100644 --- a/lib/kvbm-logical/src/registry/mod.rs +++ b/lib/kvbm-logical/src/registry/mod.rs @@ -43,7 +43,6 @@ use crate::{events::EventsManager, tinylfu::FrequencyTracker}; use crate::blocks::SequenceHash; -use std::collections::BTreeMap; use std::sync::{Arc, Weak}; use handle::BlockRegistrationHandleInner; @@ -261,38 +260,16 @@ impl BlockRegistry { return Vec::new(); } - let mut by_position = BTreeMap::>::new(); - for (index, seq_hash) in seq_hashes.iter().copied().enumerate() { - by_position - .entry(seq_hash.position()) - .or_default() - .push((index, seq_hash)); - } - - let mut registered = vec![None; seq_hashes.len()]; - let mut newly_created = vec![false; seq_hashes.len()]; - for hashes in by_position.into_values() { - let map = self.prt.prefix(&hashes[0].1); - for (index, seq_hash) in hashes { - let mut weak = map.entry(seq_hash).or_default(); - let (inner, is_new) = match weak.upgrade() { - Some(inner) => (inner, false), - None => { - let inner = self.create_registration(seq_hash); - *weak = Arc::downgrade(&inner); - (inner, true) - } - }; - registered[index] = Some(BlockRegistrationHandle::from_inner(inner)); - newly_created[index] = is_new; - } - } + let positions_are_monotonic = seq_hashes.is_sorted_by_key(|seq_hash| seq_hash.position()); + let registered = if positions_are_monotonic { + self.register_monotonic_sequence_hashes(&seq_hashes) + } else { + self.register_grouped_sequence_hashes(&seq_hashes) + }; registered .into_iter() - .zip(newly_created) .map(|(handle, is_new)| { - let handle = handle.expect("every batched sequence hash must be registered"); if is_new { if let Some(event_manager) = &self.event_manager && let Err(e) = event_manager.on_block_registered(&handle) @@ -306,6 +283,83 @@ impl BlockRegistry { .collect() } + /// Fast path for the normal sequence-registration shape: non-decreasing + /// block positions. Equal positions are adjacent, so each position needs + /// only one radix-prefix guard and results can be appended in input order. + fn register_monotonic_sequence_hashes( + &self, + seq_hashes: &[SequenceHash], + ) -> Vec<(BlockRegistrationHandle, bool)> { + let mut registered = Vec::with_capacity(seq_hashes.len()); + self.register_ordered_sequence_hashes( + seq_hashes.iter().copied().map(|seq_hash| ((), seq_hash)), + |(), handle, is_new| registered.push((handle, is_new)), + ); + registered + } + + /// Fallback for callers that supply positions out of order. A flat index + /// vector replaces the previous `BTreeMap` of per-position vectors. It is + /// sorted by `(position, original_index)` so registration stays grouped by + /// radix prefix without changing which duplicate occurrence is considered + /// new. Results are restored to input order before observers run. + fn register_grouped_sequence_hashes( + &self, + seq_hashes: &[SequenceHash], + ) -> Vec<(BlockRegistrationHandle, bool)> { + let mut ordered: Vec<_> = seq_hashes.iter().copied().enumerate().collect(); + ordered.sort_unstable_by_key(|(index, seq_hash)| (seq_hash.position(), *index)); + + let mut registered = Vec::with_capacity(seq_hashes.len()); + self.register_ordered_sequence_hashes(ordered.iter().copied(), |index, handle, is_new| { + registered.push((index, handle, is_new)) + }); + + registered.sort_unstable_by_key(|(index, _, _)| *index); + registered + .into_iter() + .map(|(_, handle, is_new)| (handle, is_new)) + .collect() + } + + /// Register position-ordered entries, holding exactly one outer radix + /// guard for each adjacent same-position group. `record` only stages the + /// result; user-visible callbacks and frequency touches run after this + /// method returns and releases every guard. + fn register_ordered_sequence_hashes( + &self, + entries: impl IntoIterator, + mut record: impl FnMut(K, BlockRegistrationHandle, bool), + ) { + let mut entries = entries.into_iter().peekable(); + while let Some(first) = entries.next() { + let position = first.1.position(); + let map = self.prt.prefix(&first.1); + let mut current = first; + + loop { + let (key, seq_hash) = current; + debug_assert_eq!(seq_hash.position(), position); + let mut weak = map.entry(seq_hash).or_default(); + let (inner, is_new) = match weak.upgrade() { + Some(inner) => (inner, false), + None => { + let inner = self.create_registration(seq_hash); + *weak = Arc::downgrade(&inner); + (inner, true) + } + }; + record(key, BlockRegistrationHandle::from_inner(inner), is_new); + + let Some(next) = entries.next_if(|(_, seq_hash)| seq_hash.position() == position) + else { + break; + }; + current = next; + } + } + } + /// Internal method for transferring block registration without triggering frequency tracking. /// Used when copying blocks between pools where we don't want to count the transfer as a new access. #[allow(dead_code)] diff --git a/lib/kvbm-logical/src/registry/tests.rs b/lib/kvbm-logical/src/registry/tests.rs index c0ded24e5411..3b74f84a72d5 100644 --- a/lib/kvbm-logical/src/registry/tests.rs +++ b/lib/kvbm-logical/src/registry/tests.rs @@ -1,7 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use crate::{KvbmSequenceHashProvider, tinylfu::TinyLFUTracker}; +use crate::{ + KvbmSequenceHashProvider, + events::EventEmissionPolicy, + tinylfu::{FrequencyTracker, TinyLFUTracker}, +}; use super::attachments::AttachmentError; use super::*; @@ -9,6 +13,7 @@ use super::*; use crate::testing::{self, MetadataA, MetadataB, MetadataC, TestMeta}; use crate::{BlockManager, blocks::BlockDuplicationPolicy}; +use parking_lot::Mutex; use std::any::TypeId; use std::sync::Arc; @@ -19,6 +24,43 @@ fn create_test_token_block(tokens: &[u32]) -> dynamo_tokens::TokenBlock { testing::create_test_token_block(tokens, tokens.len() as u32) } +fn sequence_hash_at(position: u64, current: u64) -> SequenceHash { + SequenceHash::new(current, position.checked_sub(1), position) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RegistrationHook { + Event(u128), + Touch(u128), +} + +struct RecordingEventPolicy { + hooks: Arc>>, +} + +impl EventEmissionPolicy for RecordingEventPolicy { + fn should_emit(&self, seq_hash: SequenceHash) -> bool { + self.hooks + .lock() + .push(RegistrationHook::Event(seq_hash.as_u128())); + true + } +} + +struct RecordingFrequencyTracker { + hooks: Arc>>, +} + +impl FrequencyTracker for RecordingFrequencyTracker { + fn touch(&self, key: u128) { + self.hooks.lock().push(RegistrationHook::Touch(key)); + } + + fn count(&self, _key: u128) -> u32 { + 0 + } +} + /// Helper to construct a manager seeded with a registry that the test owns. fn manager_with_registry( registry: BlockRegistry, @@ -67,6 +109,90 @@ fn test_batch_registration_preserves_order_and_reuses_duplicate_handle() { ); } +#[test] +fn test_batch_registration_monotonic_positions_groups_adjacent_equals() { + let registry = BlockRegistry::new(); + let hashes = [ + sequence_hash_at(0, 10), + sequence_hash_at(1, 20), + sequence_hash_at(1, 21), + sequence_hash_at(2, 30), + ]; + + let handles = registry.register_sequence_hashes(hashes); + + assert_eq!( + handles + .iter() + .map(BlockRegistrationHandle::seq_hash) + .collect::>(), + hashes + ); + assert_eq!(registry.registered_count(), hashes.len()); +} + +#[test] +fn test_batch_registration_out_of_order_reuses_existing_and_duplicate_handles() { + let registry = BlockRegistry::new(); + let first = sequence_hash_at(0, 10); + let second = sequence_hash_at(1, 20); + let third = sequence_hash_at(2, 30); + let existing = registry.register_sequence_hash(second); + + let hashes = [third, first, second, third]; + let handles = registry.register_sequence_hashes(hashes); + + assert_eq!( + handles + .iter() + .map(BlockRegistrationHandle::seq_hash) + .collect::>(), + hashes + ); + assert!(Arc::ptr_eq(&handles[2].inner, &existing.inner)); + assert!(Arc::ptr_eq(&handles[0].inner, &handles[3].inner)); + assert_eq!(registry.registered_count(), 3); +} + +#[test] +fn test_batch_registration_runs_observers_in_original_input_order() { + let hooks = Arc::new(Mutex::new(Vec::new())); + let event_manager = Arc::new( + EventsManager::builder() + .policy(Arc::new(RecordingEventPolicy { + hooks: hooks.clone(), + })) + .build(), + ); + let tracker = Arc::new(RecordingFrequencyTracker { + hooks: hooks.clone(), + }); + let registry = BlockRegistry::builder() + .event_manager(event_manager) + .frequency_tracker(tracker) + .build(); + + let existing_hash = sequence_hash_at(0, 10); + let later_hash = sequence_hash_at(2, 30); + let middle_hash = sequence_hash_at(1, 20); + let _existing = registry.register_sequence_hash(existing_hash); + hooks.lock().clear(); + + let handles = + registry.register_sequence_hashes([later_hash, existing_hash, middle_hash, later_hash]); + + assert_eq!( + *hooks.lock(), + vec![ + RegistrationHook::Event(later_hash.as_u128()), + RegistrationHook::Touch(later_hash.as_u128()), + RegistrationHook::Event(middle_hash.as_u128()), + RegistrationHook::Touch(middle_hash.as_u128()), + ] + ); + assert!(Arc::ptr_eq(&handles[0].inner, &handles[3].inner)); +} + #[test] fn test_type_tracking_enforcement() { let registry = BlockRegistry::new(); @@ -510,8 +636,6 @@ fn test_touch_no_callbacks_is_noop() { #[test] fn test_touch_callback_receives_correct_hash() { - use parking_lot::Mutex; - let registry = BlockRegistry::new(); let seq_hash = create_test_token_block(&[13, 14, 15, 16]).kvbm_sequence_hash(); let handle = registry.register_sequence_hash(seq_hash); diff --git a/lib/mocker/src/common/sequence.rs b/lib/mocker/src/common/sequence.rs index d371aacc8e5c..104470cb10c8 100644 --- a/lib/mocker/src/common/sequence.rs +++ b/lib/mocker/src/common/sequence.rs @@ -380,18 +380,16 @@ impl ActiveSequence { let active_blocks = active_tokens .div_ceil(self.block_size) .min(self.unique_blocks.len()); - self.unique_blocks[..active_blocks] + if active_blocks == 0 { + return Vec::new(); + } + + let blocks = self.unique_blocks[..active_blocks] .iter() .rev() - .map(|block| match block { - UniqueBlock::PartialBlock(uuid) => { - MoveBlock::Deref(vec![UniqueBlock::PartialBlock(*uuid)]) - } - UniqueBlock::FullBlock(hash) => { - MoveBlock::Deref(vec![UniqueBlock::FullBlock(*hash)]) - } - }) - .collect() + .cloned() + .collect(); + vec![MoveBlock::Deref(blocks)] } /// Free the currently active allocation footprint. @@ -485,23 +483,12 @@ mod tests { } } - fn assert_deref_partial(signal: &MoveBlock) { + fn assert_deref_blocks(signal: &MoveBlock, expected: &[UniqueBlock]) { match signal { MoveBlock::Deref(blocks) => { - assert_eq!(blocks.len(), 1); - assert!(matches!(blocks[0], UniqueBlock::PartialBlock(_))); + assert_eq!(blocks, expected); } - _ => panic!("Expected MoveBlock::Deref for partial block"), - } - } - - fn assert_deref_full(signal: &MoveBlock) { - match signal { - MoveBlock::Deref(blocks) => { - assert_eq!(blocks.len(), 1); - assert!(matches!(blocks[0], UniqueBlock::FullBlock(_))); - } - _ => panic!("Expected MoveBlock::Deref for full block"), + _ => panic!("Expected MoveBlock::Deref"), } } @@ -614,11 +601,42 @@ mod tests { let free_signals = seq.reset_with_signal(); - assert!(!free_signals.is_empty()); + assert_eq!(free_signals.len(), 1); + let expected = seq + .unique_blocks() + .iter() + .rev() + .cloned() + .collect::>(); + assert_deref_blocks(&free_signals[0], &expected); assert_eq!(seq.num_allocated_tokens(), 0); assert_eq!(seq.generated_tokens(), 2); } + #[test] + fn test_free_signal_is_empty_without_an_active_allocation() { + let seq = ActiveSequence::new((0..10).collect(), 4, Some(4), true, false); + + assert!(seq.free_signal().is_empty()); + } + + #[test] + fn test_free_signal_batches_allocated_blocks_in_reverse_order() { + let mut seq = ActiveSequence::new((0..10).collect(), 4, Some(4), true, false); + seq.commit_allocation(seq.len()); + + let expected = seq + .unique_blocks() + .iter() + .rev() + .cloned() + .collect::>(); + let signals = seq.free_signal(); + + assert_eq!(signals.len(), 1); + assert_deref_blocks(&signals[0], &expected); + } + #[test] fn test_active_sequence_generate_signals() { // Create a sequence with block size 16, max_output_tokens 4, initialized with tokens [0..14) @@ -647,15 +665,17 @@ mod tests { let signals_third = seq.generate(); assert_eq!(signals_third.len(), 0); - // Generate last token - we reach max_output_tokens, should trigger Deref signals + // Generate last token - we reach max_output_tokens, so all blocks should + // be dereferenced in one reverse-ordered batch. + let expected = seq + .unique_blocks() + .iter() + .rev() + .cloned() + .collect::>(); let signals_last = seq.generate(); - assert_eq!(signals_last.len(), 2); - - // First signal should be Deref for the partial block - assert_deref_partial(&signals_last[0]); - - // Second signal should be Deref for the full block - assert_deref_full(&signals_last[1]); + assert_eq!(signals_last.len(), 1); + assert_deref_blocks(&signals_last[0], &expected); } #[test] diff --git a/lib/mocker/src/kv_manager/kvbm_backend.rs b/lib/mocker/src/kv_manager/kvbm_backend.rs index 97ce183f83ae..da2b11152d7a 100644 --- a/lib/mocker/src/kv_manager/kvbm_backend.rs +++ b/lib/mocker/src/kv_manager/kvbm_backend.rs @@ -13,10 +13,11 @@ //! transaction, then commit the whole request atomically. Capacity exhaustion //! leaves ownership and sequence state unchanged so the scheduler can decide //! whether to preempt a running request. -//! - **Deref**: release one request-owned handle. For `PartialBlock` this drops -//! the unique `MutableBlock` and returns it to the reset pool. For -//! `FullBlock` this pops one `ImmutableBlock` clone; when the vec empties, -//! the block transitions to kvbm-logical's inactive pool (RAII return). +//! - **Deref**: release one logical request owner. For `PartialBlock` this +//! drops the unique `MutableBlock` and returns it to the reset pool. For +//! `FullBlock` this decrements an explicit logical refcount; the final +//! release drops the canonical `ImmutableBlock` and transitions the block to +//! kvbm-logical's inactive pool (RAII return). //! - **Promote**: PartialBlock (`MutableBlock`) → FullBlock (`ImmutableBlock`). //! Collapses onto an existing registered handle if the PLH / SequenceHash is //! already present; otherwise stages + registers a new block. @@ -29,6 +30,7 @@ //! - `Lru` — simple recency-based LRU. //! - `MultiLru` — 4-tier frequency-aware LRU (requires TinyLFU tracker). +use std::collections::hash_map::Entry; use std::sync::Arc; #[cfg(feature = "kvbm-offload")] use std::sync::Mutex; @@ -98,9 +100,9 @@ enum SwapInSlotReservation { /// Classification for each block processed inside `Use`. /// /// - `ActiveHit`: block is already pinned in `active_full` / `active_partial`; -/// we just bump our local refcount (handle clone). +/// commit bumps its explicit logical refcount without cloning a handle. /// - `InactiveHit`: block was in kvbm-logical's inactive pool and was -/// reactivated via `match_blocks(plh)`. +/// reactivated by the aligned scattered batch lookup. /// - `NewStore`: block was freshly allocated, staged, and registered. /// /// The router radix tree already knows about `ActiveHit` and `InactiveHit` @@ -135,10 +137,23 @@ pub struct OffloadDependency { } enum PreparedUseBlock { - ExistingFull { + /// Already represented in `active_full`; no temporary RAII clone is + /// needed while the transaction reserves its fresh suffix. + ExistingActiveFull { + seq_hash: SequenceHash, + }, + /// Resurrected from KVBM's inactive pool (or otherwise matched outside the + /// mocker's active map). The handle pins it until commit or rollback. + ExistingMatchedFull { seq_hash: SequenceHash, handle: ImmutableBlock, }, + /// Not present in `active_full`; resolved by the single aligned scattered + /// lookup after the initial classification pass. + PendingNonLocalFull { + seq_hash: SequenceHash, + full_idx: usize, + }, ExistingPartial, FreshFull { seq_hash: SequenceHash, @@ -161,6 +176,7 @@ struct UseSignalRef<'a> { struct UseTransaction<'a> { signal: UseSignalRef<'a>, prepared: Vec, + fresh_full_blocks: usize, evicted_plhs: Vec, } @@ -225,6 +241,34 @@ struct FullBlockMetadata { token_ids: Option>, } +/// One physical full-block pin plus the number of logical request owners. +/// +/// `ImmutableBlock` clones are physical-lifetime guards, not request block +/// tables. Keeping a clone per logical owner needlessly makes KVBM's handle +/// count, allocator traffic, and Arc traffic scale with prefix sharing. The +/// canonical handle pins the physical block while `logical_refs` tracks the +/// ownership semantics the mocker needs for Deref. +struct ActiveFullBlock { + handle: ImmutableBlock, + logical_refs: usize, +} + +impl ActiveFullBlock { + fn new(handle: ImmutableBlock) -> Self { + Self { + handle, + logical_refs: 1, + } + } + + fn retain(&mut self) { + self.logical_refs = self + .logical_refs + .checked_add(1) + .expect("active full-block logical reference count overflowed"); + } +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum FullBlockCommit { Reused, @@ -244,11 +288,11 @@ pub struct KvManager { /// Dropped blocks return to kvbm-logical's reset pool. active_partial: FxHashMap>, - /// FullBlocks held as `ImmutableBlock`, keyed by `SequenceHash`. The vec - /// length is the mocker's reference count — each `Use` pushes a clone, - /// each `Deref` pops one. When the vec empties, the block transitions to - /// kvbm-logical's inactive pool (RAII return on drop of the last clone). - active_full: FxHashMap>>, + /// FullBlocks held as one canonical `ImmutableBlock` per physical block, + /// keyed by `SequenceHash`, plus the number of logical request owners. + /// The final logical `Deref` drops the canonical handle and transitions the + /// block to kvbm-logical's inactive pool. + active_full: FxHashMap, /// Shadow registry for every block registered in kvbm-logical. The logical /// registry is keyed by `PositionalLineageHash`, while the router's radix @@ -352,6 +396,55 @@ impl KvManager { } } + /// Install a newly acquired physical handle or merge it into an entry that + /// became active earlier in the same serial commit. + fn insert_or_retain_active_full(&mut self, seq_hash: SequenceHash, handle: ImmutableBlock) { + match self.active_full.entry(seq_hash) { + Entry::Vacant(entry) => { + entry.insert(ActiveFullBlock::new(handle)); + } + Entry::Occupied(mut entry) => { + assert_eq!( + entry.get().handle.block_id(), + handle.block_id(), + "active full-block hash resolved to a different physical block" + ); + entry.get_mut().retain(); + // `handle` is a redundant physical pin. Dropping it leaves the + // canonical entry alive while logical ownership is tracked by + // `logical_refs`. + drop(handle); + } + } + } + + /// Add one logical owner for a block known to be present in the active map. + /// This is called only after every fallible reservation for the surrounding + /// Use transaction has succeeded. + fn retain_active_full(&mut self, seq_hash: SequenceHash) { + self.active_full + .get_mut(&seq_hash) + .unwrap_or_else(|| panic!("active full block {seq_hash:?} disappeared before commit")) + .retain(); + } + + /// Release one logical owner. Removing the final entry drops the sole + /// physical handle and lets KVBM transition the block to inactive. + fn release_active_full(&mut self, seq_hash: SequenceHash) { + let Entry::Occupied(mut entry) = self.active_full.entry(seq_hash) else { + panic!("Deref: full block not in active pool"); + }; + assert!( + entry.get().logical_refs > 0, + "active full block must retain at least one logical owner" + ); + if entry.get().logical_refs == 1 { + entry.remove(); + } else { + entry.get_mut().logical_refs -= 1; + } + } + /// Wrap `engine` in `Arc>`, install it onto this /// `KvManager`, and return a clone of the Arc to the caller. /// Called once after construction by the scheduler's init helper; @@ -1002,7 +1095,7 @@ impl KvManager { plh: PositionalLineageHash, ) -> Option> { if let Some(active) = self.active_full.get(&seq_hash) { - return Some(active[0].clone()); + return Some(active.handle.clone()); } self.block_manager.match_blocks(&[plh]).into_iter().next() } @@ -1022,10 +1115,7 @@ impl KvManager { if let Some(canonical) = self.acquire_existing_full(seq_hash, plh) { drop(candidate); - self.active_full - .entry(seq_hash) - .or_default() - .push(canonical); + self.insert_or_retain_active_full(seq_hash, canonical); return FullBlockCommit::Reused; } @@ -1035,10 +1125,7 @@ impl KvManager { .expect("full block stage failed"); let canonical = self.block_manager.register_block(complete); let canonical_block_id = canonical.block_id(); - self.active_full - .entry(seq_hash) - .or_default() - .push(canonical); + self.insert_or_retain_active_full(seq_hash, canonical); if canonical_block_id != candidate_block_id { return FullBlockCommit::Reused; @@ -1202,7 +1289,7 @@ impl KvManager { let (_, handle) = cached_prefix .next() .expect("reserved prefix handle must exist"); - self.active_full.entry(seq_hash).or_default().push(handle); + self.insert_or_retain_active_full(seq_hash, handle); metadata_parent_hash = Some(seq_hash); continue; } @@ -1460,35 +1547,39 @@ impl KvManager { expected_full_blocks, ); + // Classify locally active blocks once, and preserve the existing + // per-block scattered reuse semantics while collapsing all non-local + // lookups into one store-lock acquisition. `match_blocks` cannot be + // used here because it stops at the first miss, whereas the existing + // singleton loop can still reuse a later registered block. + // + // Start empty rather than reserving for every full block: an all-active + // request never needs storage for non-local lookup inputs. let mut prepared = Vec::with_capacity(blocks.len()); - let mut evicted_plhs = Vec::new(); + let mut nonlocal_plhs = Vec::new(); let mut fresh_blocks = 0usize; + let mut fresh_full_blocks = 0usize; let mut full_idx = 0usize; for block in blocks { match block { UniqueBlock::FullBlock(seq_hash) => { - let plh = plhs[full_idx]; - let entry = if let Some(active) = self.active_full.get(seq_hash) { - PreparedUseBlock::ExistingFull { - seq_hash: *seq_hash, - handle: active[0].clone(), - } - } else if let Some(handle) = - self.block_manager.match_blocks(&[plh]).into_iter().next() - { - PreparedUseBlock::ExistingFull { + if self.active_full.contains_key(seq_hash) { + prepared.push(PreparedUseBlock::ExistingActiveFull { seq_hash: *seq_hash, - handle, - } + }); } else { - fresh_blocks += 1; - PreparedUseBlock::FreshFull { + prepared.push(PreparedUseBlock::PendingNonLocalFull { seq_hash: *seq_hash, full_idx, - mutable: None, + }); + // Allocate once, but only when the request actually + // contains a non-local full block. Every remaining + // full block is the largest possible suffix here. + if nonlocal_plhs.is_empty() { + nonlocal_plhs.reserve_exact(expected_full_blocks - full_idx); } - }; - prepared.push(entry); + nonlocal_plhs.push(plhs[full_idx]); + } full_idx += 1; } UniqueBlock::PartialBlock(uuid) => { @@ -1505,6 +1596,40 @@ impl KvManager { } } + if !nonlocal_plhs.is_empty() { + let mut nonlocal_matches = self + .block_manager + .match_blocks_scattered(&nonlocal_plhs) + .into_iter(); + for entry in &mut prepared { + let PreparedUseBlock::PendingNonLocalFull { seq_hash, full_idx } = entry else { + continue; + }; + let seq_hash = *seq_hash; + let full_idx = *full_idx; + *entry = if let Some(handle) = nonlocal_matches + .next() + .expect("scattered match result must align with non-local full blocks") + { + PreparedUseBlock::ExistingMatchedFull { seq_hash, handle } + } else { + fresh_blocks += 1; + fresh_full_blocks += 1; + PreparedUseBlock::FreshFull { + seq_hash, + full_idx, + mutable: None, + } + }; + } + assert!( + nonlocal_matches.next().is_none(), + "scattered match returned more entries than non-local full blocks" + ); + } + + let mut evicted_plhs = Vec::new(); + if let Some(reservation) = reservation.as_mut() { if reservation.len() < fresh_blocks { return G1Acquire::CapacityExhausted; @@ -1519,7 +1644,12 @@ impl KvManager { .expect("prechecked decode reservation must contain a slot"), ); } - PreparedUseBlock::ExistingFull { .. } | PreparedUseBlock::ExistingPartial => {} + PreparedUseBlock::ExistingActiveFull { .. } + | PreparedUseBlock::ExistingMatchedFull { .. } + | PreparedUseBlock::ExistingPartial => {} + PreparedUseBlock::PendingNonLocalFull { .. } => { + unreachable!("non-local full block must be resolved before reservation") + } } } } else { @@ -1557,7 +1687,12 @@ impl KvManager { .expect("atomic Use reservation returned too few slots"), ); } - PreparedUseBlock::ExistingFull { .. } | PreparedUseBlock::ExistingPartial => {} + PreparedUseBlock::ExistingActiveFull { .. } + | PreparedUseBlock::ExistingMatchedFull { .. } + | PreparedUseBlock::ExistingPartial => {} + PreparedUseBlock::PendingNonLocalFull { .. } => { + unreachable!("non-local full block must be resolved before reservation") + } } } assert!( @@ -1574,6 +1709,7 @@ impl KvManager { parent, }, prepared, + fresh_full_blocks, evicted_plhs, }) } @@ -1581,9 +1717,45 @@ impl KvManager { fn commit_use(&mut self, transaction: UseTransaction<'_>) { let UseTransaction { signal, - prepared, + mut prepared, + fresh_full_blocks, evicted_plhs, } = transaction; + + // Complete every fresh full block first, then register the whole set + // under one BlockStore lock. Registration results preserve input order, + // so the second pass can consume them alongside the fresh prepared + // entries while preserving router-event segmentation and metadata. + let mut completed_blocks = Vec::with_capacity(fresh_full_blocks); + let mut candidate_block_ids = Vec::with_capacity(fresh_full_blocks); + for entry in &mut prepared { + if let PreparedUseBlock::FreshFull { + full_idx, mutable, .. + } = entry + { + let mutable = mutable + .take() + .expect("committing Use must own every fresh full slot"); + candidate_block_ids.push(mutable.block_id()); + let complete = mutable + .stage(signal.plhs[*full_idx], self.block_size) + .expect("Use full block stage failed"); + completed_blocks.push(complete); + } + } + let registered_blocks = self.block_manager.register_blocks(completed_blocks); + assert_eq!( + candidate_block_ids.len(), + fresh_full_blocks, + "prepared fresh full count must match staged candidate IDs" + ); + assert_eq!( + candidate_block_ids.len(), + registered_blocks.len(), + "fresh candidate IDs must align with batch registration results" + ); + let mut fresh_registrations = candidate_block_ids.into_iter().zip(registered_blocks); + let mut metadata_parent_hash = match signal.parent { None => None, Some(UniqueBlock::FullBlock(seq_hash)) => Some(*seq_hash), @@ -1596,7 +1768,24 @@ impl KvManager { for entry in prepared { match entry { - PreparedUseBlock::ExistingFull { seq_hash, handle } => { + PreparedUseBlock::ExistingActiveFull { seq_hash } => { + if !blocks_stored.is_empty() { + let hashes = std::mem::take(&mut blocks_stored); + let local_hashes = std::mem::take(&mut stored_local_hashes); + let token_ids = stored_token_ids.as_mut().map(std::mem::take); + self.publish_kv_event( + hashes, + &local_hashes, + first_store_parent, + true, + token_ids, + ); + } + self.retain_active_full(seq_hash); + metadata_parent_hash = Some(seq_hash); + first_store_parent = metadata_parent_hash; + } + PreparedUseBlock::ExistingMatchedFull { seq_hash, handle } => { if !blocks_stored.is_empty() { let hashes = std::mem::take(&mut blocks_stored); let local_hashes = std::mem::take(&mut stored_local_hashes); @@ -1609,10 +1798,13 @@ impl KvManager { token_ids, ); } - self.active_full.entry(seq_hash).or_default().push(handle); + self.insert_or_retain_active_full(seq_hash, handle); metadata_parent_hash = Some(seq_hash); first_store_parent = metadata_parent_hash; } + PreparedUseBlock::PendingNonLocalFull { .. } => { + unreachable!("non-local full block must be resolved before commit") + } PreparedUseBlock::ExistingPartial => {} PreparedUseBlock::FreshFull { seq_hash, @@ -1623,21 +1815,38 @@ impl KvManager { first_store_parent = metadata_parent_hash; } let plh = signal.plhs[full_idx]; - let mutable = mutable.expect("committing Use must own every fresh full slot"); - let candidate_block_id = mutable.block_id(); - let complete = mutable - .stage(plh, self.block_size) - .expect("Use full block stage failed"); - let immutable = self.block_manager.register_block(complete); - assert_eq!( - immutable.block_id(), - candidate_block_id, - "prepared fresh Use block unexpectedly resolved to an existing registration" + assert!( + mutable.is_none(), + "fresh full slot must be consumed by batch staging" ); - self.active_full - .entry(seq_hash) - .or_default() - .push(immutable); + let (candidate_block_id, immutable) = fresh_registrations + .next() + .expect("fresh full block must have a registration result"); + if immutable.block_id() != candidate_block_id { + // Reject deduplication can resolve two fresh entries in + // this same batch to one canonical block. Finish the + // preceding Stored group, retain the returned handle as + // another logical owner, and advance the lineage cursor + // without replacing canonical shadow metadata or + // publishing a duplicate Stored event. + if !blocks_stored.is_empty() { + let hashes = std::mem::take(&mut blocks_stored); + let local_hashes = std::mem::take(&mut stored_local_hashes); + let token_ids = stored_token_ids.as_mut().map(std::mem::take); + self.publish_kv_event( + hashes, + &local_hashes, + first_store_parent, + true, + token_ids, + ); + } + self.insert_or_retain_active_full(seq_hash, immutable); + metadata_parent_hash = Some(seq_hash); + first_store_parent = metadata_parent_hash; + continue; + } + self.insert_or_retain_active_full(seq_hash, immutable); let local_hash = signal.local_hashes.get(full_idx).copied(); let registry_token_ids = signal @@ -1678,6 +1887,10 @@ impl KvManager { } } } + assert!( + fresh_registrations.next().is_none(), + "unused fresh full registration result" + ); if !blocks_stored.is_empty() { self.publish_kv_event( @@ -1801,14 +2014,7 @@ impl KvManager { .expect("Deref: partial block not in active pool"); } UniqueBlock::FullBlock(seq_hash) => { - let vec = self - .active_full - .get_mut(seq_hash) - .expect("Deref: full block not in active pool"); - vec.pop(); - if vec.is_empty() { - self.active_full.remove(seq_hash); - } + self.release_active_full(*seq_hash); } } } @@ -1869,11 +2075,17 @@ impl KvManager { self.block_manager.total_blocks() - self.block_manager.available_blocks() } - /// Total number of held RAII handles (refcount-style): one per held - /// `MutableBlock` plus one per cloned `ImmutableBlock` in `active_full`. - /// Shared-prefix reuse inflates this above the distinct-block count. + /// Total number of logical block owners: one per held `MutableBlock` plus + /// the explicit logical reference count of every full block. This remains + /// a request-ownership metric even though KVBM's `inflight_immutable` + /// metric now counts only the canonical physical handles retained here. pub fn num_active_block_refs(&self) -> usize { - self.active_partial.len() + self.active_full.values().map(|v| v.len()).sum::() + self.active_partial.len() + + self + .active_full + .values() + .map(|active| active.logical_refs) + .sum::() } #[cfg(test)] @@ -1885,8 +2097,7 @@ impl KvManager { UniqueBlock::FullBlock(hash) => self .active_full .get(hash) - .and_then(|handles| handles.last()) - .map(ImmutableBlock::block_id), + .map(|active| active.handle.block_id()), UniqueBlock::PartialBlock(uuid) => { self.active_partial.get(uuid).map(MutableBlock::block_id) } @@ -1976,7 +2187,7 @@ mod tests { use std::sync::Mutex; use super::*; - use crate::common::protocols::KvCacheEventSink; + use crate::common::protocols::{KvCacheEventSink, RawKvEvent, RawKvEventSink}; /// Capturing event sink for router-publication assertions. #[derive(Default)] @@ -1990,6 +2201,18 @@ mod tests { } } + #[derive(Default)] + struct CapturingRawSink { + events: Mutex>, + } + + impl RawKvEventSink for CapturingRawSink { + fn publish(&self, event: RawKvEvent) -> anyhow::Result<()> { + self.events.lock().unwrap().push(event); + Ok(()) + } + } + fn make_mgr(capacity: usize, block_size: usize) -> KvManager { KvManager::new_with_event_sink(capacity, block_size, KvEventPublishers::default(), 0) } @@ -2014,6 +2237,21 @@ mod tests { ) } + fn make_mgr_capturing_with_raw( + capacity: usize, + block_size: usize, + ) -> (KvManager, Arc, Arc) { + let sink = Arc::new(CapturingSink::default()); + let raw_sink = Arc::new(CapturingRawSink::default()); + let publishers = + KvEventPublishers::new(Some(sink.clone() as _), Some(raw_sink.clone() as _)); + ( + KvManager::new_with_event_sink(capacity, block_size, publishers, 0), + sink, + raw_sink, + ) + } + fn make_mgr_capturing_with_backend( capacity: usize, block_size: usize, @@ -2164,9 +2402,65 @@ mod tests { use_full(&mut mgr, 1, plh(100)); use_full(&mut mgr, 1, plh(100)); // Same seq_hash used twice: only one distinct physical block is - // resident, but the mocker holds two RAII handles. + // resident and pinned by one canonical RAII handle, while the mocker + // tracks two logical request owners. assert_eq!(mgr.num_active_blocks(), 1); assert_eq!(mgr.num_active_block_refs(), 2); + assert_eq!(mgr.block_manager.metrics().snapshot().inflight_immutable, 1); + + deref_full(&mut mgr, 1); + assert_eq!(mgr.num_active_blocks(), 1); + assert_eq!(mgr.num_active_block_refs(), 1); + assert_eq!(mgr.block_manager.metrics().snapshot().inflight_immutable, 1); + + deref_full(&mut mgr, 1); + assert_eq!(mgr.num_active_blocks(), 0); + assert_eq!(mgr.num_active_block_refs(), 0); + assert_eq!(mgr.block_manager.metrics().snapshot().inflight_immutable, 0); + } + + #[test] + fn all_active_multi_block_use_only_retains_logical_owners() { + let (mut mgr, sink) = make_mgr_capturing(4, 4); + let blocks = vec![UniqueBlock::FullBlock(10), UniqueBlock::FullBlock(20)]; + let plhs = vec![plh(100), plh(200)]; + + assert_eq!( + expect_ready(mgr.process(&MoveBlock::Use( + blocks.clone(), + vec![101, 201], + plhs.clone(), + None, + None, + ))), + 2 + ); + let available_before = mgr.block_manager.available_blocks(); + let first_block_id = mgr.active_full[&10].handle.block_id(); + let second_block_id = mgr.active_full[&20].handle.block_id(); + sink.events.lock().unwrap().clear(); + + assert_eq!( + expect_ready(mgr.process(&MoveBlock::Use(blocks, vec![101, 201], plhs, None, None,))), + 2 + ); + + assert_eq!(mgr.block_manager.available_blocks(), available_before); + assert_eq!(mgr.num_active_blocks(), 2); + assert_eq!(mgr.num_active_block_refs(), 4); + assert_eq!(mgr.active_full[&10].handle.block_id(), first_block_id); + assert_eq!(mgr.active_full[&20].handle.block_id(), second_block_id); + assert!( + sink.events.lock().unwrap().is_empty(), + "retaining active blocks must not publish another Stored event" + ); + + for _ in 0..2 { + deref_full(&mut mgr, 10); + deref_full(&mut mgr, 20); + } + assert_eq!(mgr.num_active_blocks(), 0); + assert_eq!(mgr.num_active_block_refs(), 0); } #[test] @@ -2189,6 +2483,27 @@ mod tests { assert_eq!(mgr.num_active_block_refs(), refs_before); } + #[test] + fn failed_mixed_use_does_not_retain_existing_active_blocks() { + let mut mgr = make_mgr(1, 16); + use_full(&mut mgr, 1, plh(100)); + + assert!(matches!( + mgr.process(&MoveBlock::Use( + vec![UniqueBlock::FullBlock(1), UniqueBlock::FullBlock(2)], + vec![], + vec![plh(100), plh(200)], + None, + None, + )), + G1Acquire::CapacityExhausted + )); + assert_eq!(mgr.num_active_blocks(), 1); + assert_eq!(mgr.num_active_block_refs(), 1); + assert_eq!(mgr.active_full[&1].logical_refs, 1); + assert!(!mgr.active_full.contains_key(&2)); + } + #[test] fn test_deref_returns_to_inactive() { let mut mgr = make_mgr(4, 16); @@ -2207,6 +2522,221 @@ mod tests { assert_eq!(use_full(&mut mgr, 2, p), 1); } + #[test] + fn scattered_use_reuses_later_hit_after_a_miss() { + let (mut mgr, sink) = make_mgr_capturing(10, 16); + let first_plh = plh(100); + let missing_plh = plh(200); + let later_plh = plh(300); + + use_full(&mut mgr, 10, first_plh); + use_full(&mut mgr, 30, later_plh); + let later_block_id = mgr.active_full[&30].handle.block_id(); + deref_full(&mut mgr, 10); + deref_full(&mut mgr, 30); + sink.events.lock().unwrap().clear(); + + assert_eq!( + expect_ready(mgr.process(&MoveBlock::Use( + vec![ + UniqueBlock::FullBlock(10), + UniqueBlock::FullBlock(20), + UniqueBlock::FullBlock(30), + ], + vec![], + vec![first_plh, missing_plh, later_plh], + None, + None, + ))), + 3 + ); + + assert_eq!(mgr.num_active_blocks(), 3); + assert_eq!(mgr.num_active_block_refs(), 3); + assert_eq!( + mgr.active_full[&30].handle.block_id(), + later_block_id, + "the registered block after a miss must still be reused" + ); + + let events = sink.events.lock().unwrap(); + assert_eq!(events.len(), 1, "only the missing middle block is stored"); + let KvCacheEventData::Stored(stored) = &events[0].data else { + panic!("expected Stored event, got {:?}", events[0].data); + }; + assert_eq!(stored.parent_hash.map(|hash| hash.0), Some(10)); + assert_eq!(stored.blocks.len(), 1); + assert_eq!(stored.blocks[0].block_hash.0, 20); + } + + #[test] + fn mixed_use_keeps_fresh_registration_and_event_order() { + let (mut mgr, sink) = make_mgr_capturing(8, 4); + let reused_plh = plh(200); + + use_full(&mut mgr, 20, reused_plh); + sink.events.lock().unwrap().clear(); + + let seq_hashes = [10, 11, 20, 30, 31]; + let plhs = [plh(100), plh(110), reused_plh, plh(300), plh(310)]; + let local_hashes = vec![1010, 1011, 1020, 1030, 1031]; + let token_ids = vec![ + vec![10, 10, 10, 10], + vec![11, 11, 11, 11], + vec![20, 20, 20, 20], + vec![30, 30, 30, 30], + vec![31, 31, 31, 31], + ]; + let blocks = seq_hashes.into_iter().map(UniqueBlock::FullBlock).collect(); + + assert_eq!( + expect_ready(mgr.process(&MoveBlock::Use( + blocks, + local_hashes.clone(), + plhs.to_vec(), + Some(token_ids.clone()), + Some(UniqueBlock::FullBlock(5)), + ))), + seq_hashes.len() + ); + + for (idx, (seq_hash, plh)) in [(10, plhs[0]), (11, plhs[1]), (30, plhs[3]), (31, plhs[4])] + .into_iter() + .enumerate() + { + let signal_idx = [0, 1, 3, 4][idx]; + let info = mgr + .registered_blocks + .get(&plh) + .expect("fresh block must retain registration metadata"); + assert_eq!(info.seq_hash, seq_hash); + assert_eq!(info.block_id, mgr.active_full[&seq_hash].handle.block_id()); + assert_eq!(info.local_hash, Some(local_hashes[signal_idx])); + assert_eq!(info.token_ids.as_ref(), Some(&token_ids[signal_idx])); + } + assert_eq!(mgr.registered_blocks[&plhs[0]].parent_hash, Some(5)); + assert_eq!(mgr.registered_blocks[&plhs[1]].parent_hash, Some(10)); + assert_eq!(mgr.registered_blocks[&plhs[3]].parent_hash, Some(20)); + assert_eq!(mgr.registered_blocks[&plhs[4]].parent_hash, Some(30)); + + let events = sink.events.lock().unwrap(); + assert_eq!(events.len(), 2, "the reused middle block splits stores"); + for (event, expected_hashes, expected_local_hashes, expected_parent) in [ + (&events[0], &[10, 11][..], &[1010, 1011][..], Some(5)), + (&events[1], &[30, 31][..], &[1030, 1031][..], Some(20)), + ] { + let KvCacheEventData::Stored(stored) = &event.data else { + panic!("expected Stored event, got {:?}", event.data); + }; + assert_eq!(stored.parent_hash.map(|hash| hash.0), expected_parent); + assert_eq!( + stored + .blocks + .iter() + .map(|block| block.block_hash.0) + .collect::>(), + expected_hashes + ); + assert_eq!( + stored + .blocks + .iter() + .map(|block| block.tokens_hash.0) + .collect::>(), + expected_local_hashes + ); + } + } + + #[test] + fn duplicate_fresh_registration_reuses_canonical_without_duplicate_event() { + const CAPACITY: usize = 6; + let (mut mgr, sink, raw_sink) = make_mgr_capturing_with_raw(CAPACITY, 4); + let a_plh = plh(100); + let b_plh = plh(200); + let local_hashes = vec![101, 102, 201]; + let token_ids = vec![vec![1; 4], vec![2; 4], vec![3; 4]]; + let dedup_before = mgr.block_manager.metrics().snapshot().registration_dedup; + + assert_eq!( + expect_ready(mgr.process(&MoveBlock::Use( + vec![ + UniqueBlock::FullBlock(10), + UniqueBlock::FullBlock(10), + UniqueBlock::FullBlock(20), + ], + local_hashes.clone(), + vec![a_plh, a_plh, b_plh], + Some(token_ids.clone()), + Some(UniqueBlock::FullBlock(5)), + ))), + 3 + ); + + assert_eq!(mgr.num_active_blocks(), 2); + assert_eq!(mgr.num_active_block_refs(), 3); + assert_eq!(mgr.block_manager.available_blocks(), CAPACITY - 2); + assert_eq!( + mgr.block_manager.metrics().snapshot().registration_dedup, + dedup_before + 1 + ); + assert_eq!(mgr.registered_blocks.len(), 2); + + let a_info = &mgr.registered_blocks[&a_plh]; + assert_eq!(a_info.seq_hash, 10); + assert_eq!(a_info.block_id, mgr.active_full[&10].handle.block_id()); + assert_eq!(a_info.parent_hash, Some(5)); + assert_eq!(a_info.local_hash, Some(local_hashes[0])); + assert_eq!(a_info.token_ids.as_ref(), Some(&token_ids[0])); + + let b_info = &mgr.registered_blocks[&b_plh]; + assert_eq!(b_info.seq_hash, 20); + assert_eq!(b_info.block_id, mgr.active_full[&20].handle.block_id()); + assert_eq!(b_info.parent_hash, Some(10)); + assert_eq!(b_info.local_hash, Some(local_hashes[2])); + assert_eq!(b_info.token_ids.as_ref(), Some(&token_ids[2])); + + let events = sink.events.lock().unwrap(); + assert_eq!(events.len(), 2, "the duplicate must split Stored groups"); + for (event, expected_hash, expected_local_hash, expected_parent) in [ + (&events[0], 10, local_hashes[0], Some(5)), + (&events[1], 20, local_hashes[2], Some(10)), + ] { + let KvCacheEventData::Stored(stored) = &event.data else { + panic!("expected Stored event, got {:?}", event.data); + }; + assert_eq!(stored.parent_hash.map(|hash| hash.0), expected_parent); + assert_eq!(stored.blocks.len(), 1); + assert_eq!(stored.blocks[0].block_hash.0, expected_hash); + assert_eq!(stored.blocks[0].tokens_hash.0, expected_local_hash); + } + drop(events); + + let raw_events = raw_sink.events.lock().unwrap(); + assert_eq!( + raw_events.len(), + 2, + "the duplicate must split raw Stored groups" + ); + assert_eq!( + raw_events[0].block_token_ids.as_deref(), + Some(std::slice::from_ref(&token_ids[0])) + ); + assert_eq!( + raw_events[1].block_token_ids.as_deref(), + Some(std::slice::from_ref(&token_ids[2])) + ); + drop(raw_events); + + deref_full(&mut mgr, 10); + deref_full(&mut mgr, 10); + deref_full(&mut mgr, 20); + assert_eq!(mgr.num_active_blocks(), 0); + assert_eq!(mgr.num_active_block_refs(), 0); + assert!(mgr.active_full.is_empty()); + assert_eq!(mgr.block_manager.available_blocks(), CAPACITY); + } + #[test] fn failed_decode_reservation_preserves_inactive_cache() { let (mut mgr, sink) = make_mgr_capturing(2, 16); @@ -2350,7 +2880,10 @@ mod tests { mgr.process(&MoveBlock::Deref(blocks)); } fn refcount(mgr: &KvManager, id: u64) -> usize { - mgr.active_full.get(&id).map(|v| v.len()).unwrap_or(0) + mgr.active_full + .get(&id) + .map(|active| active.logical_refs) + .unwrap_or(0) } fn assert_active(mgr: &KvManager, expected: &[(u64, usize)]) { let distinct = expected.len(); @@ -3303,7 +3836,7 @@ mod tests { assert!(target_hashes.iter().all(|target| { mgr.active_full .get(target) - .is_some_and(|refs| refs.len() == 1) + .is_some_and(|active| active.logical_refs == 1) })); let committed = sink.take(); From 8c9b744e69e91f4b01ede3a89398593a76514547 Mon Sep 17 00:00:00 2001 From: Michael Feil <63565275+michaelfeil@users.noreply.github.com> Date: Tue, 30 Jun 2026 14:29:22 -0700 Subject: [PATCH 014/320] fix(runtime): retry etcd startup connection (#10799) Signed-off-by: Michael Feil <63565275+michaelfeil@users.noreply.github.com> --- lib/runtime/src/transports/etcd.rs | 110 ++++++++++++++++------- lib/runtime/src/transports/etcd/lease.rs | 15 ++++ 2 files changed, 95 insertions(+), 30 deletions(-) diff --git a/lib/runtime/src/transports/etcd.rs b/lib/runtime/src/transports/etcd.rs index 6db385c143c5..65764d7f40b8 100644 --- a/lib/runtime/src/transports/etcd.rs +++ b/lib/runtime/src/transports/etcd.rs @@ -19,7 +19,7 @@ use etcd_client::{ WatchStream, Watcher, }; pub use etcd_client::{ConnectOptions, KeyValue, LeaseClient}; -use tokio::time::{Duration, interval}; +use tokio::time::{Duration, Instant, interval}; use tokio_util::sync::CancellationToken; mod connector; @@ -33,6 +33,10 @@ pub use lock::*; use super::utils::build_in_runtime; use crate::config::environment_names::etcd as env_etcd; +const STARTUP_CONNECT_TIMEOUT: Duration = Duration::from_secs(120); +const STARTUP_CONNECT_INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const STARTUP_CONNECT_MAX_BACKOFF: Duration = Duration::from_secs(30); + /// ETCD Client #[derive(Clone)] pub struct Client { @@ -68,35 +72,7 @@ impl Client { let token = runtime.primary_token(); let ((connector, lease_id), rt) = build_in_runtime( - async move { - let etcd_urls = config.etcd_url.clone(); - let connect_options = config.etcd_connect_options.clone(); - - // Create the connector - let connector = Connector::new(etcd_urls, connect_options) - .await - .with_context(|| { - format!( - "Unable to connect to etcd server at {}. Check etcd server status", - config.etcd_url.join(", ") - ) - })?; - - let lease_id = if config.attach_lease { - create_lease(connector.clone(), config.lease_ttl, token) - .await - .with_context(|| { - format!( - "Unable to create lease. Check etcd server status at {}", - config.etcd_url.join(", ") - ) - })? - } else { - 0 - }; - - Ok((connector, lease_id)) - }, + async move { Self::connect_with_startup_retry(&config, token).await }, 1, ) .await?; @@ -109,6 +85,80 @@ impl Client { }) } + /// Connect to etcd during startup, retrying with exponential backoff for up to 2 minutes. + async fn connect_with_startup_retry( + config: &ClientOptions, + token: CancellationToken, + ) -> Result<(Arc, u64)> { + let deadline = Instant::now() + STARTUP_CONNECT_TIMEOUT; + let mut backoff = STARTUP_CONNECT_INITIAL_BACKOFF; + + loop { + if token.is_cancelled() { + anyhow::bail!("etcd startup connection cancelled"); + } + + let attempt = Self::connect_startup_attempt(config, &token).await; + + match attempt { + Ok(connection) => return Ok(connection), + Err(err) => { + let now = Instant::now(); + if now >= deadline { + return Err(err); + } + + let sleep_duration = backoff.min(deadline.saturating_duration_since(now)); + + tracing::warn!( + error = %err, + retry_in = ?sleep_duration, + remaining = ?deadline.saturating_duration_since(now), + "etcd not reachable yet; retrying startup connection" + ); + + tokio::select! { + biased; + + _ = token.cancelled() => { + anyhow::bail!("etcd startup connection cancelled"); + } + + _ = tokio::time::sleep(sleep_duration) => {} + } + backoff = backoff.saturating_mul(2).min(STARTUP_CONNECT_MAX_BACKOFF); + } + } + } + } + + async fn connect_startup_attempt( + config: &ClientOptions, + token: &CancellationToken, + ) -> Result<(Arc, u64)> { + let connector = + Connector::new(config.etcd_url.clone(), config.etcd_connect_options.clone()).await?; + + let lease_id = if config.attach_lease { + create_lease(connector.clone(), config.lease_ttl, token.clone()) + .await + .with_context(|| { + format!( + "Unable to create lease. Check etcd server status at {}", + config.etcd_url.join(", ") + ) + })? + } else { + 0 + }; + + if token.is_cancelled() { + anyhow::bail!("etcd startup connection cancelled"); + } + + Ok((connector, lease_id)) + } + /// Get a clone of the underlying [`etcd_client::Client`] instance. /// This returns a clone since the client is behind an RwLock. fn etcd_client(&self) -> etcd_client::Client { diff --git a/lib/runtime/src/transports/etcd/lease.rs b/lib/runtime/src/transports/etcd/lease.rs index f92049bb34cb..a34127c9e9a2 100644 --- a/lib/runtime/src/transports/etcd/lease.rs +++ b/lib/runtime/src/transports/etcd/lease.rs @@ -17,10 +17,25 @@ pub async fn create_lease( ttl: u64, token: CancellationToken, ) -> anyhow::Result { + if token.is_cancelled() { + anyhow::bail!("lease creation cancelled"); + } + let mut lease_client = connector.get_client().lease_client(); let lease = lease_client.grant(ttl as i64, None).await?; let id = lease.id() as u64; + if token.is_cancelled() { + if let Err(e) = lease_client.revoke(id as i64).await { + tracing::warn!( + lease_id = id, + error = %e, + "Failed to revoke lease after cancellation during creation" + ); + } + anyhow::bail!("lease creation cancelled"); + } + let ttl = lease.ttl() as u64; let child = token.child_token(); From f285e9185588eaf4a5a656863ec93ea35d703b26 Mon Sep 17 00:00:00 2001 From: Xianlu Bird Date: Wed, 1 Jul 2026 05:31:25 +0800 Subject: [PATCH 015/320] fix(kvbm-physical): Linux-only DiskStorage handling in builder (#7518) Signed-off-by: xianlubird Co-authored-by: Ryan McCormick --- lib/memory/src/disk.rs | 59 ++++++++++++++++++++++++++++-------------- lib/memory/src/lib.rs | 2 -- 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/lib/memory/src/disk.rs b/lib/memory/src/disk.rs index 774cdf77fb79..9229b0a38d73 100644 --- a/lib/memory/src/disk.rs +++ b/lib/memory/src/disk.rs @@ -8,7 +8,10 @@ use std::any::Any; use std::path::{Path, PathBuf}; use core::ffi::c_char; +#[cfg(target_os = "linux")] use nix::fcntl::{FallocateFlags, fallocate}; +#[cfg(not(target_os = "linux"))] +use nix::unistd::ftruncate; use nix::unistd::unlink; use std::ffi::CString; use std::os::fd::BorrowedFd; @@ -16,6 +19,11 @@ use std::os::fd::BorrowedFd; const DISK_CACHE_KEY: &str = "DYN_KVBM_DISK_CACHE_DIR"; const DEFAULT_DISK_CACHE_DIR: &str = "/tmp/"; +#[cfg(target_os = "linux")] +const DISK_OPEN_DIRECT_FLAG: i32 = nix::libc::O_DIRECT; +#[cfg(not(target_os = "linux"))] +const DISK_OPEN_DIRECT_FLAG: i32 = 0; + /// Disk-backed storage using memory-mapped files with O_DIRECT support. #[derive(Debug)] pub struct DiskStorage { @@ -82,12 +90,7 @@ impl DiskStorage { let template = CString::new(path_str).unwrap(); let mut template_bytes = template.into_bytes_with_nul(); - let fd = unsafe { - nix::libc::mkostemp( - template_bytes.as_mut_ptr() as *mut c_char, - nix::libc::O_RDWR | nix::libc::O_DIRECT, - ) - }; + let fd = unsafe { create_temp_file(template_bytes.as_mut_ptr() as *mut c_char) }; if fd == -1 { return Err(StorageError::AllocationFailed(format!( @@ -111,7 +114,7 @@ impl DiskStorage { let fd = unsafe { nix::libc::open( path_cstr.as_ptr(), - nix::libc::O_CREAT | nix::libc::O_RDWR | nix::libc::O_DIRECT, + nix::libc::O_CREAT | nix::libc::O_RDWR | DISK_OPEN_DIRECT_FLAG, 0o644, ) }; @@ -126,18 +129,7 @@ impl DiskStorage { (fd, file_path) }; - // We need to use fallocate to actually allocate the storage and create the blocks on disk. - unsafe { - fallocate( - BorrowedFd::borrow_raw(raw_fd), - FallocateFlags::empty(), - 0, - len as i64, - ) - .map_err(|e| { - StorageError::AllocationFailed(format!("Failed to allocate temp file: {}", e)) - })? - }; + allocate_file(raw_fd, len)?; Ok(Self { fd: raw_fd as u64, @@ -178,6 +170,35 @@ impl DiskStorage { } } +fn allocate_file(raw_fd: i32, len: usize) -> Result<()> { + #[cfg(target_os = "linux")] + unsafe { + fallocate( + BorrowedFd::borrow_raw(raw_fd), + FallocateFlags::empty(), + 0, + len as i64, + ) + .map_err(|e| StorageError::AllocationFailed(format!("Failed to allocate temp file: {}", e))) + } + + #[cfg(not(target_os = "linux"))] + unsafe { + ftruncate(BorrowedFd::borrow_raw(raw_fd), len as i64) + .map_err(|e| StorageError::AllocationFailed(format!("Failed to size temp file: {}", e))) + } +} + +#[cfg(target_os = "linux")] +unsafe fn create_temp_file(template: *mut c_char) -> i32 { + unsafe { nix::libc::mkostemp(template, nix::libc::O_RDWR | DISK_OPEN_DIRECT_FLAG) } +} + +#[cfg(not(target_os = "linux"))] +unsafe fn create_temp_file(template: *mut c_char) -> i32 { + unsafe { nix::libc::mkstemp(template) } +} + impl Drop for DiskStorage { fn drop(&mut self) { let _ = self.unlink(); diff --git a/lib/memory/src/lib.rs b/lib/memory/src/lib.rs index f8749b3f1356..3067ae24cf29 100644 --- a/lib/memory/src/lib.rs +++ b/lib/memory/src/lib.rs @@ -27,7 +27,6 @@ pub mod pool; pub mod prelude; mod device; -#[cfg(target_os = "linux")] mod disk; mod external; mod pinned; @@ -39,7 +38,6 @@ mod tests; pub use arena::{ArenaAllocator, ArenaBuffer, ArenaError}; pub use device::DeviceStorage; -#[cfg(target_os = "linux")] pub use disk::DiskStorage; pub use external::ExternalDeviceMemory; #[cfg(target_os = "linux")] From ec348af0626357a053177d550840add890dd9f83 Mon Sep 17 00:00:00 2001 From: weizhoublue <45163302+weizhoublue@users.noreply.github.com> Date: Wed, 1 Jul 2026 06:13:11 +0800 Subject: [PATCH 016/320] fix: Preventing Silent Data Poisoning and Vector Integrity Loss from Malicious Boolean Inputs (#9873) Signed-off-by: weizhou.lan Signed-off-by: weizhoublue --- components/src/dynamo/vllm/handlers.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index c26178cb54c8..478da740be49 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -3900,7 +3900,9 @@ async def generate( prompts: list[Any] = _classify_embedding_input(input_field) dimensions = request.get("dimensions") - if dimensions is not None and not isinstance(dimensions, int): + if dimensions is not None and ( + not isinstance(dimensions, int) or isinstance(dimensions, bool) + ): raise TypeError( f"Invalid 'dimensions' type {type(dimensions).__name__}; expected int" ) From ccc835a80cf8e5ccc022b30db7df1f9828930333 Mon Sep 17 00:00:00 2001 From: GuanLuo <41310872+GuanLuo@users.noreply.github.com> Date: Tue, 30 Jun 2026 15:17:50 -0700 Subject: [PATCH 017/320] feat(vllm/omni): realtime (bidirectional) vLLM-Omni worker (#10166) Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) --- components/src/dynamo/vllm/omni/__init__.py | 3 +- components/src/dynamo/vllm/omni/args.py | 18 + components/src/dynamo/vllm/omni/main.py | 4 + .../src/dynamo/vllm/omni/realtime_handler.py | 627 ++++++++++++++++++ .../src/dynamo/vllm/omni/realtime_utils.py | 146 ++++ .../vllm/tests/omni/test_omni_base_handler.py | 1 + .../tests/omni/test_omni_realtime_handler.py | 424 ++++++++++++ .../vllm/tests/test_vllm_api_contract.py | 12 +- .../vllm/tests/test_vllm_renderer_api.py | 6 +- container/context.yaml | 2 +- container/deps/vllm/install_vllm_omni.sh | 41 ++ .../backends/vllm/launch/agg_omni_realtime.sh | 62 ++ .../vllm/launch/realtime_omni_client.py | 285 ++++++++ tests/frontend/realtime_omni_mock_worker.py | 95 +++ tests/frontend/test_realtime_omni_bridge.py | 211 ++++++ 15 files changed, 1932 insertions(+), 5 deletions(-) create mode 100644 components/src/dynamo/vllm/omni/realtime_handler.py create mode 100644 components/src/dynamo/vllm/omni/realtime_utils.py create mode 100644 components/src/dynamo/vllm/tests/omni/test_omni_realtime_handler.py create mode 100755 examples/backends/vllm/launch/agg_omni_realtime.sh create mode 100644 examples/backends/vllm/launch/realtime_omni_client.py create mode 100644 tests/frontend/realtime_omni_mock_worker.py create mode 100644 tests/frontend/test_realtime_omni_bridge.py diff --git a/components/src/dynamo/vllm/omni/__init__.py b/components/src/dynamo/vllm/omni/__init__.py index dedcd71e7306..0fd55184ce37 100644 --- a/components/src/dynamo/vllm/omni/__init__.py +++ b/components/src/dynamo/vllm/omni/__init__.py @@ -5,5 +5,6 @@ from .base_handler import BaseOmniHandler from .omni_handler import OmniHandler +from .realtime_handler import RealtimeOmniHandler -__all__ = ["BaseOmniHandler", "OmniHandler"] +__all__ = ["BaseOmniHandler", "OmniHandler", "RealtimeOmniHandler"] diff --git a/components/src/dynamo/vllm/omni/args.py b/components/src/dynamo/vllm/omni/args.py index fd64cd5ec7f6..024f1c78d292 100644 --- a/components/src/dynamo/vllm/omni/args.py +++ b/components/src/dynamo/vllm/omni/args.py @@ -318,6 +318,17 @@ def add_arguments(self, parser) -> None: "Requires --stage-configs-path. Mutually exclusive with --stage-id." ), ) + add_negatable_bool_argument( + g, + flag_name="--realtime", + env_var="DYN_OMNI_REALTIME", + default=False, + help=( + "Serve a ModelType.Realtime bidirectional endpoint (OpenAI " + "Realtime API) backed by vLLM-Omni streaming generation, instead " + "of the unary multimodal endpoint." + ), + ) class OmniConfig(DynamoRuntimeConfig): @@ -350,6 +361,9 @@ class OmniConfig(DynamoRuntimeConfig): stage_id: Optional[int] = None omni_router: bool = False + # Realtime (bidirectional) serving mode + realtime: bool = False + @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "OmniConfig": config = super().from_cli_args(args) @@ -390,6 +404,10 @@ def validate(self) -> None: raise ValueError("--stage-id must be >= 0") if self.stage_id is not None and self.omni_router: raise ValueError("--stage-id and --omni-router are mutually exclusive") + if self.realtime and (self.stage_id is not None or self.omni_router): + raise ValueError( + "--realtime cannot be combined with --stage-id or --omni-router" + ) def parse_omni_args() -> OmniConfig: diff --git a/components/src/dynamo/vllm/omni/main.py b/components/src/dynamo/vllm/omni/main.py index 0fa91e259ce6..0b03ee989d85 100644 --- a/components/src/dynamo/vllm/omni/main.py +++ b/components/src/dynamo/vllm/omni/main.py @@ -20,6 +20,7 @@ from dynamo.runtime.logging import configure_dynamo_logging from dynamo.vllm.health_check import VllmOmniHealthCheckPayload from dynamo.vllm.main import setup_metrics_collection +from dynamo.vllm.omni.realtime_utils import init_omni_realtime from dynamo.vllm.omni.stage_router import init_omni_stage_router from dynamo.vllm.omni.stage_worker import init_omni_stage @@ -141,6 +142,9 @@ async def worker(): elif config.omni_router: await init_omni_stage_router(runtime, config, shutdown_endpoints) logger.debug("init_omni_stage_router completed") + elif config.realtime: + await init_omni_realtime(runtime, config, shutdown_endpoints, shutdown_event) + logger.debug("init_omni_realtime completed, exiting...") else: await init_omni(runtime, config, shutdown_event) logger.debug("Omni worker completed, exiting...") diff --git a/components/src/dynamo/vllm/omni/realtime_handler.py b/components/src/dynamo/vllm/omni/realtime_handler.py new file mode 100644 index 000000000000..468156038eca --- /dev/null +++ b/components/src/dynamo/vllm/omni/realtime_handler.py @@ -0,0 +1,627 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Realtime (bidirectional) handler backed by vLLM-Omni's streaming engine. + +This handler expects ``request_stream`` to yield ``RealtimeServerEvent`` frames. + +Turn model: + + * ``session.update`` -> ``session.updated`` echoing the session; + also captures the requested ``output_modalities`` for later turns. + * ``input_audio_buffer.append`` -> base64 PCM16 chunk decoded to a float32 + waveform and queued for the turn's audio stream. The first ``append`` (or + ``commit``) opens the turn and the engine begins draining audio. + * ``input_audio_buffer.commit`` -> a final ``commit`` closes the audio stream + so the engine drains and produces the response. + +Turns run concurrently -- a commit's turn drives the engine as soon as it opens +-- but their responses are forwarded to the client in turn order: each turn +buffers its server events, and a later turn's buffer is held until the previous +turn's response completes. Responses never interleave and each is identified by +its ``response_id``. + +Each turn emits ``response.created`` -> ``response.output_audio.delta``* (+ +optional ``response.output_audio_transcript.delta`` for the thinker text) -> +``response.output_audio.done`` -> ``response.done``. These are the OpenAI-spec +event names the frontend's typed reader requires, which differ from +vLLM-Omni's own ``response.audio.delta`` / ``transcription.delta`` names; the +PCM16 and cumulative-vs-delta waveform handling below is ported from +vLLM-Omni's ``realtime_connection.py`` and only the event tags change. + +Limitations (MVP): each ``input_audio_buffer.commit`` is transcribed and +answered independently -- a turn's generation is seeded only by its own audio. +Prior turns' transcripts/responses are not fed into later turns, and +``conversation.item.*`` / ``response.create`` are accepted-and-ignored. This is +a single-utterance transcribe-and-respond bridge, not a stateful multi-turn +dialogue. +""" + +from __future__ import annotations + +import asyncio +import base64 +import logging +import uuid +from typing import Any, AsyncGenerator, Callable, Optional, Sequence + +import numpy as np + +from dynamo._core import Context + +logger = logging.getLogger(__name__) + +# ``streaming_input_factory(audio_stream, input_stream) -> AsyncGenerator`` — +# mirrors ``OpenAIServingRealtime.transcribe_realtime``: it consumes float32 +# audio chunks and an ``asyncio.Queue`` of context token ids, yielding engine +# ``StreamingInput`` prompts. +# The factory and engine are injected so the worker passes the real serving/AsyncOmni +# while tests pass lightweight fakes. +StreamingInputFactory = Callable[ + [AsyncGenerator[np.ndarray, None], "asyncio.Queue[list[int]]"], + AsyncGenerator[Any, None], +] + + +def event_id() -> str: + return f"event_{uuid.uuid4().hex}" + + +def session_updated_event(session: Any) -> dict: + """Echo a client ``session.update`` back as the spec ``session.updated``.""" + return { + "type": "session.updated", + "event_id": event_id(), + "session": session, + } + + +class Turn: + """One request->response cycle: a committed span of input audio and the + single OpenAI-spec ``response`` it produces. + + Owns its engine drive and output extraction (``drive_engine``: it feeds its + buffered audio into the engine and yields ``(transcript, audio_chunks)`` per + step); ``RealtimeOmniHandler`` orchestrates turns and translates those into + OpenAI-spec server events (see ``RealtimeOmniHandler.run_turn``). + + Fields: ``response_id`` / ``item_id`` tag every event of this turn; + ``audio_queue`` carries the float32 input (``None`` = end of input); + ``audio_ref`` tracks the last emitted waveform so cumulative engine outputs + are de-duplicated into true deltas; ``output_modalities`` is snapshotted from + the latest ``session.update``; ``task`` runs the turn's processing; + ``events`` buffers its server events for in-order forwarding (``None`` = end + of the turn's response). + """ + + def __init__( + self, + *, + engine_client: Any, + streaming_input_factory: StreamingInputFactory, + default_sampling_params_list: Optional[Sequence[Any]] = None, + output_modalities: list[str] | None = None, + ) -> None: + self.response_id = f"resp_{uuid.uuid4().hex}" + self.item_id = f"item_{uuid.uuid4().hex}" + # Unbounded on purpose: filled by non-blocking put_nowait so the inbound + # demux never stalls control events (commit/clear/session.update) behind + # audio backpressure; paced by the client's input rate and drained by the + # engine. + self.audio_queue: asyncio.Queue[Optional[np.ndarray]] = asyncio.Queue() + self.audio_ref: np.ndarray | None = None + self.task: asyncio.Task | None = None + # This turn's server events, forwarded to the client in turn order by + # ``RealtimeOmniHandler.generate`` (``None`` = end of response). Bounded + # so a turn waiting for its forwarding slot exerts backpressure on its + # own engine drive rather than buffering unbounded. + self.events: asyncio.Queue[Optional[dict]] = asyncio.Queue(maxsize=256) + self.output_modalities = output_modalities + self._engine_client = engine_client + self._streaming_input_factory = streaming_input_factory + self._default_sampling_params_list = default_sampling_params_list + + async def drive_engine( + self, + ) -> AsyncGenerator[tuple[str | None, list[np.ndarray] | None], None]: + """Feed this turn's buffered audio into the engine and yield, per engine + step, the extracted ``(transcript, audio_chunks)`` -- either element is + ``None`` when that step produced none. + + The float32 ``audio_stream`` (ending on the ``None`` sentinel) and an + ``input_stream`` token queue are handed to the streaming-input factory + (``transcribe_realtime``), whose ``StreamingInput`` generator drives + ``AsyncOmni.generate``. + """ + + # build input stream generator + async def audio_stream() -> AsyncGenerator[np.ndarray, None]: + while True: + waveform = await self.audio_queue.get() + if waveform is None: + return + yield waveform + + input_stream: asyncio.Queue[list[int]] = asyncio.Queue() + streaming_input_gen = self._streaming_input_factory( + audio_stream(), input_stream + ) + + generate_kwargs: dict[str, Any] = { + "prompt": streaming_input_gen, + "request_id": self.response_id, + } + if self._default_sampling_params_list is not None: + generate_kwargs["sampling_params_list"] = list( + self._default_sampling_params_list + ) + # Client-requested output modalities (session.update) select the final + # pipeline stage: include "audio" to drive the talker. Omitted -> + # AsyncOmni.generate uses the engine's launch-time default. + if self.output_modalities is not None: + generate_kwargs["output_modalities"] = self.output_modalities + + # drive the engine + async for output in self._engine_client.generate(**generate_kwargs): + token_ids = self.thinker_token_ids(output) + if token_ids: + input_stream.put_nowait(token_ids) + transcript = self.extract_transcript(output) or None + audio_chunks = self.extract_audio_chunks(output) or None + yield transcript, audio_chunks + + # Generation done: release the cumulative reference waveform. It only + # de-dups deltas within this drive; nothing client-bound survives on it + # (the deltas are already encoded into the queued events), and for long + # voice sessions a ~10 MB/turn float32 buffer pinned per turn adds up. + self.audio_ref = None + + @staticmethod + def thinker_token_ids(output: Any) -> list[int]: + """Stage-0 (thinker) per-step token ids to feed back to the talker.""" + if getattr(output, "stage_id", None) != 0: + return [] + outputs = getattr(output, "outputs", None) + if not outputs: + return [] + token_ids = getattr(outputs[0], "token_ids", None) + return list(token_ids) if token_ids else [] + + @staticmethod + def extract_transcript(output: Any) -> str: + """Pull incremental thinker text from a stage-0 LLM output, if any.""" + if getattr(output, "stage_id", None) != 0: + return "" + outputs = getattr(output, "outputs", None) + if not outputs: + return "" + return getattr(outputs[0], "text", "") or "" + + def extract_audio_chunks(self, output: Any) -> list[np.ndarray]: + """Extract per-step audio deltas from an engine output. + + Audio lives in ``output.multimodal_output['audio'|'model_outputs']`` as a + float32 waveform (or list of them). Some engine paths emit a growing + cumulative waveform; ``waveform_to_deltas`` reconciles both shapes + against ``self.audio_ref`` so the client never hears duplicates. + """ + mm = getattr(output, "multimodal_output", None) + if mm is None: + return [] + + raw_audio = mm.tensors["audio"] if "audio" in mm.tensors.keys() else None + if raw_audio is None: + return [] + + if isinstance(raw_audio, (list, tuple)): + if not raw_audio: + return [] + arr = tensor_to_numpy(raw_audio[-1]) + else: + arr = tensor_to_numpy(raw_audio) + + if arr is None or arr.size == 0: + return [] + return self.waveform_to_deltas(arr) + + def waveform_to_deltas(self, arr: np.ndarray) -> list[np.ndarray]: + """Convert one streaming PCM f32 chunk into incremental piece(s).""" + if arr.size == 0: + return [] + ref = self.audio_ref + if ref is None: + self.audio_ref = arr.copy() + return [arr] + if numpy_audio_prefix_match(ref, arr): + delta = arr[ref.shape[0] :] + self.audio_ref = arr.copy() + return [delta] if delta.size > 0 else [] + # True per-step delta (not a prefix extension of what we have seen). The + # growing concat makes this O(n^2) over a response, but per-response audio + # is bounded (seconds); kept to mirror vLLM-Omni's realtime_connection. + self.audio_ref = np.concatenate([ref, arr]) + return [arr] + + +class RealtimeOmniHandler: + """Bridge OpenAI Realtime client events to vLLM-Omni streaming generation. + + Owns the per-connection orchestration (event demux, concurrent turns whose + responses are forwarded in turn order) and the conversion between the + realtime API and model output: it drives each ``Turn``'s engine generation + and translates the engine's stage outputs into OpenAI-spec server events. + """ + + def __init__( + self, + *, + engine_client: Any, + model_name: str, + streaming_input_factory: StreamingInputFactory, + default_sampling_params_list: Optional[Sequence[Any]] = None, + emit_transcript: bool = True, + max_concurrent_turns: int = 8, + ) -> None: + self.engine_client = engine_client + self.model_name = model_name + self._streaming_input_factory = streaming_input_factory + self._default_sampling_params_list = default_sampling_params_list + self._emit_transcript = emit_transcript + # Upper bound on in-flight turns per connection: the pump blocks before + # opening a new turn once this many are running, so a pipelining or + # abusive client cannot spawn unbounded concurrent engine generations. + self._max_concurrent_turns = max_concurrent_turns + + def new_turn(self, output_modalities: list[str] | None) -> Turn: + return Turn( + engine_client=self.engine_client, + streaming_input_factory=self._streaming_input_factory, + default_sampling_params_list=self._default_sampling_params_list, + output_modalities=output_modalities, + ) + + async def run_turn(self, turn: Turn, context: Context) -> None: + """Drive one turn's engine generation and buffer its server events. + + Events are appended to ``turn.events`` rather than sent to the client + directly; ``generate`` forwards each turn's buffer in turn order, so + turns may run concurrently while their responses never interleave. A + ``None`` sentinel is always appended last to mark the turn's response + complete and let the forwarder advance to the next turn. + """ + events = turn.events + try: + await events.put(self.response_created_event(turn)) + + sent_audio = False + async for transcript, audio_chunks in turn.drive_engine(): + if context.is_stopped(): + break + + if transcript and self._emit_transcript: + await events.put(self.transcript_delta_event(turn, transcript)) + + for chunk in audio_chunks or (): + sent_audio = True + await events.put(self.audio_delta_event(turn, chunk)) + + if context.is_stopped(): + # Connection torn down mid-turn; don't claim a completed response. + return + + if sent_audio: + await events.put(self.audio_done_event(turn)) + await events.put(self.response_done_event(turn)) + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 - surface engine errors on the wire + logger.exception("realtime omni turn failed: %s", exc) + # The top-level ``error`` event carries the human-readable message but + # no ``response_id``, so also close the dangling in-progress response + # with a terminal ``response.done(status=failed)`` -- that event + # carries the id, so the client can correlate and the response reaches + # a terminal state instead of hanging. + await events.put(self.error_event(exc)) + await events.put( + self.response_done_event( + turn, + status="failed", + status_details={ + "type": "failed", + "error": { + "code": "omni_generation_error", + "type": "server_error", + }, + }, + ) + ) + finally: + await events.put(None) + + # -- response lifecycle events -------------------------------------------- + + def response_created_event(self, turn: Turn) -> dict: + return { + "type": "response.created", + "event_id": event_id(), + "response": { + "id": turn.response_id, + "max_output_tokens": "inf", + "object": "realtime.response", + "output": [], + "output_modalities": ["audio"], + "status": "in_progress", + }, + } + + def response_done_event( + self, turn: Turn, status: str = "completed", status_details: dict | None = None + ) -> dict: + response: dict[str, Any] = { + "id": turn.response_id, + "max_output_tokens": "inf", + "object": "realtime.response", + "output": [], + "output_modalities": ["audio"], + "status": status, + } + if status_details is not None: + response["status_details"] = status_details + return { + "type": "response.done", + "event_id": event_id(), + "response": response, + } + + def error_event(self, exc: Exception) -> dict: + return { + "type": "error", + "event_id": event_id(), + "error": { + "type": "server_error", + "code": "omni_generation_error", + "message": str(exc), + }, + } + + # -- output translation (ported from vllm-omni realtime_connection.py) ---- + + def audio_delta_event(self, turn: Turn, chunk: np.ndarray) -> dict: + return { + "type": "response.output_audio.delta", + "event_id": event_id(), + "response_id": turn.response_id, + "item_id": turn.item_id, + "output_index": 0, + "content_index": 0, + "delta": pcm16_b64(chunk), + } + + def audio_done_event(self, turn: Turn) -> dict: + return { + "type": "response.output_audio.done", + "event_id": event_id(), + "response_id": turn.response_id, + "item_id": turn.item_id, + "output_index": 0, + "content_index": 0, + } + + def transcript_delta_event(self, turn: Turn, delta: str) -> dict: + return { + "type": "response.output_audio_transcript.delta", + "event_id": event_id(), + "response_id": turn.response_id, + "item_id": turn.item_id, + "output_index": 0, + "content_index": 0, + "delta": delta, + } + + async def generate( + self, request_stream: AsyncGenerator[Any, None], context: Context + ) -> AsyncGenerator[dict, None]: + """Serve one realtime connection. + + Each turn is spawned as a task that drives its engine and buffers its + server events on the turn's own queue. Standalone events (e.g. + ``session.updated``) and turns are placed on ``out_stream`` in arrival + order; this coroutine forwards them in that order, draining a turn's + buffered events in full before the next turn -- so turns run concurrently + while their responses never interleave. + """ + # Ordered hand-off: each item is either a standalone server event (dict) + # to forward as-is, or a ``Turn`` whose buffered events are drained in + # order once reached. Low-volume (one item per turn / session.update) so + # unbounded; per-turn audio backpressure lives on each turn's ``events``. + out_stream: asyncio.Queue[Any] = asyncio.Queue() + active_turn: Turn | None = None + turns: list[Turn] = [] + # Latest output modalities requested by the client via session.update; + # snapshotted into each turn so the engine emits text/audio accordingly. + session_output_modalities: list[str] | None = None + # Cap concurrent in-flight turns: a slot is acquired before a turn opens + # and released when its task finishes, so the pump back-pressures rather + # than spawning unbounded engine generations on one connection. + turn_slots = asyncio.Semaphore(self._max_concurrent_turns) + + async def ensure_turn() -> Turn: + nonlocal active_turn + if active_turn is None: + await turn_slots.acquire() + active_turn = self.new_turn(session_output_modalities) + turns.append(active_turn) + out_stream.put_nowait(active_turn) + active_turn.task = asyncio.create_task( + self.run_turn(active_turn, context) + ) + active_turn.task.add_done_callback(lambda _: turn_slots.release()) + return active_turn + + async def pump() -> None: + nonlocal active_turn, session_output_modalities + try: + async for client_event in request_stream: + if context.is_stopped(): + break + etype = ( + client_event.get("type") + if isinstance(client_event, dict) + else None + ) + + if etype == "session.update": + session = client_event.get("session") + modalities = parse_output_modalities(session) + if modalities is not None: + session_output_modalities = modalities + out_stream.put_nowait(session_updated_event(session)) + elif etype == "input_audio_buffer.append": + turn = await ensure_turn() + waveform = decode_pcm16(client_event.get("audio", "")) + if waveform is not None: + turn.audio_queue.put_nowait(waveform) + elif etype == "input_audio_buffer.commit": + turn = await ensure_turn() + # `final` absent defaults to True: a bare commit means + # "buffer complete, generate". A non-final commit opens + # the turn (engine starts) but keeps the input open. + if client_event.get("final", True): + turn.audio_queue.put_nowait(None) + active_turn = None + elif etype == "input_audio_buffer.clear": + # Per the spec, clear only discards the *input* buffer (not + # yet committed); it does not cancel an in-flight response + # (that's response.cancel). After a final commit active_turn + # is None, so a clear correctly no-ops. + if active_turn is not None: + drain_queue(active_turn.audio_queue) + else: + # The frontend forwards every client event; ones we don't + # drive (conversation.item.*, response.*, etc.) are logged + # and ignored so a well-behaved session is not torn down. + logger.debug("realtime omni: ignoring client event %s", etype) + + # Input stream ended: close any still-open turn so the engine + # sees end-of-input rather than hanging. + if active_turn is not None: + active_turn.audio_queue.put_nowait(None) + active_turn = None + finally: + # No more turns/events will be queued; lets the forwarder stop + # once it has drained every turn already on out_stream. + out_stream.put_nowait(None) + + pump_task = asyncio.create_task(pump()) + try: + while True: + item = await out_stream.get() + if item is None: + break + if isinstance(item, Turn): + # Forward this turn's response in full before the next turn; + # run_turn always closes the buffer with a None sentinel. + while True: + event = await item.events.get() + if event is None: + break + yield event + # Fully forwarded (its task is finishing): drop our reference + # so the turn's buffers are released instead of pinned in + # `turns` for the whole connection. + item.audio_ref = None + turns.remove(item) + else: + yield item + finally: + pump_task.cancel() + for turn in turns: + if turn.task is not None: + turn.task.cancel() + # Unblock any turn task parked on a full events queue so its + # cancellation can propagate, then drive every task to completion -- + # this closes each engine generate() async-gen instead of leaking it. + for turn in turns: + drain_queue(turn.events) + await asyncio.gather( + pump_task, + *(turn.task for turn in turns if turn.task is not None), + return_exceptions=True, + ) + + +def parse_output_modalities(session: Any) -> list[str] | None: + """Extract requested output modalities from a session.update `session` block. + + Reads OpenAI Realtime ``output_modalities`` (falling back to the older + ``modalities``); returns a list of strings, or None when unset/malformed so + the engine's launch-time default applies. + """ + if not isinstance(session, dict): + return None + modalities = session.get("output_modalities") + if modalities is None: + modalities = session.get("modalities") + if isinstance(modalities, list) and all(isinstance(m, str) for m in modalities): + return modalities + return None + + +def decode_pcm16(audio_b64: str) -> np.ndarray | None: + """Decode a base64 PCM16 chunk to a float32 waveform in [-1, 1]. + + Mirrors vLLM's realtime connection decode (int16 / 32768). Empty / blank + payloads yield ``None`` so they are not queued as audio. + """ + if not audio_b64: + return None + raw = base64.b64decode(audio_b64) + if len(raw) % 2: + # PCM16 is 2-byte aligned; np.frombuffer would raise. Drop the malformed + # chunk rather than let one bad frame tear down the whole session. + logger.warning( + "realtime omni: dropping odd-length (%d-byte) audio chunk", len(raw) + ) + return None + waveform = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + return waveform if waveform.size else None + + +def drain_queue(queue: asyncio.Queue) -> None: + while not queue.empty(): + try: + queue.get_nowait() + except asyncio.QueueEmpty: + break + + +def tensor_to_numpy(value: Any) -> np.ndarray | None: + if value is None: + return None + if isinstance(value, np.ndarray): + arr = value + elif hasattr(value, "detach"): + arr = value.detach().float().cpu().numpy() + else: + try: + arr = np.asarray(value) + except Exception: # noqa: BLE001 - non-array engine payloads are skipped + return None + if arr.ndim > 1: + arr = arr.reshape(-1) + return arr.astype(np.float32, copy=False) + + +def numpy_audio_prefix_match(prev: np.ndarray, curr: np.ndarray) -> bool: + n = prev.shape[0] + if n == 0: + return True + if curr.shape[0] < n: + return False + return bool(np.allclose(curr[:n], prev, rtol=1e-3, atol=2e-4)) + + +def pcm16_b64(audio_f32: np.ndarray) -> str: + clipped = np.clip(audio_f32, -1.0, 1.0) + pcm16 = (clipped * 32767.0).astype(np.int16) + return base64.b64encode(pcm16.tobytes()).decode("utf-8") diff --git a/components/src/dynamo/vllm/omni/realtime_utils.py b/components/src/dynamo/vllm/omni/realtime_utils.py new file mode 100644 index 000000000000..427aebdaf14a --- /dev/null +++ b/components/src/dynamo/vllm/omni/realtime_utils.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Realtime (bidirectional) Omni worker initialization. + +Serves a ``ModelType.Realtime`` model backed by vLLM-Omni's streaming engine +via ``serve_bidirectional_endpoint``. The frontend discovers it and installs a +typed realtime PushRouter; see ``realtime_handler.RealtimeOmniHandler`` for the +event translation. +""" + +import asyncio +import logging + +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime + +from dynamo import prometheus_names +from dynamo.llm import ModelInput, ModelType, WorkerType, register_model +from dynamo.runtime import DistributedRuntime +from dynamo.vllm.main import setup_metrics_collection +from dynamo.vllm.omni.base_handler import BaseOmniHandler +from dynamo.vllm.omni.realtime_handler import RealtimeOmniHandler + +from .args import OmniConfig + +logger = logging.getLogger(__name__) + + +async def init_omni_realtime( + runtime: DistributedRuntime, + config: OmniConfig, + shutdown_endpoints: list, + shutdown_event: asyncio.Event, +) -> None: + """Initialize and serve the realtime bidirectional Omni worker.""" + generate_endpoint = runtime.endpoint( + f"{config.namespace}.{config.component}.{config.endpoint}" + ) + shutdown_endpoints[:] = [generate_endpoint] + + # BaseOmniHandler builds the AsyncOmni engine from the same kwargs the unary + # Omni worker uses; we only need its engine_client for the realtime bridge. + base = BaseOmniHandler( + runtime=runtime, + config=config, + default_sampling_params={}, + shutdown_event=shutdown_event, + ) + + sampling_params_list = streaming_sampling_params(base.engine_client) + streaming_input_factory = build_streaming_input_factory(config, base.engine_client) + + handler = RealtimeOmniHandler( + engine_client=base.engine_client, + model_name=config.served_model_name or config.model, + streaming_input_factory=streaming_input_factory, + default_sampling_params_list=sampling_params_list, + ) + + logger.info("Realtime Omni worker initialized for model: %s", config.model) + + setup_metrics_collection(config, generate_endpoint, logger) + + if config.engine_args.data_parallel_rank: + logger.info( + "Non-leader DP rank %d; skipping endpoint registration", + config.engine_args.data_parallel_rank, + ) + await shutdown_event.wait() + return + + model_label = config.served_model_name or config.model + try: + await register_model( + ModelInput.Text, + ModelType.Realtime, + generate_endpoint, + config.model, + config.served_model_name, + kv_cache_block_size=config.engine_args.block_size, + # The realtime worker serves the full multi-stage pipeline behind one + # endpoint, so it registers as Aggregated like the unary Omni worker. + worker_type=WorkerType.Aggregated, + needs=[], + ) + + logger.info("Starting to serve realtime Omni worker endpoint...") + + # No health_check_payload: serve_bidirectional_endpoint does not yet + # support canary probes (the bidirectional engine is stateful and needs + # a session.update-shaped payload); see the Rust binding's doc comment. + await generate_endpoint.serve_bidirectional_endpoint( + handler.generate, + graceful_shutdown=True, + metrics_labels=[ + (prometheus_names.labels.MODEL, model_label), + (prometheus_names.labels.MODEL_NAME, model_label), + ], + ) + except Exception as e: + logger.error("Realtime Omni worker failed: %s", e) + raise + finally: + logger.debug("Cleaning up realtime Omni worker") + base.cleanup() + + +def build_streaming_input_factory(config: OmniConfig, engine_client): + """Build the audio -> StreamingInput factory from vLLM utils.""" + model_name = config.served_model_name or config.model + base_model_paths = [BaseModelPath(name=model_name, model_path=config.model)] + serving_models = OpenAIServingModels( + engine_client=engine_client, + base_model_paths=base_model_paths, + lora_modules=None, + ) + serving_realtime = OpenAIServingRealtime( + engine_client=engine_client, + models=serving_models, + request_logger=None, + ) + return serving_realtime.transcribe_realtime + + +def streaming_sampling_params(engine_client) -> list | None: + """Default per-stage sampling params coerced for streaming generation. + + vLLM-Omni requires streaming requests to emit incremental (delta) outputs; + ``coerce_param_message_types`` flips the engine defaults accordingly. Falls + back to ``None`` (engine defaults) if anything is unavailable. + """ + try: + from vllm_omni.entrypoints.utils import coerce_param_message_types + + defaults = list(engine_client.default_sampling_params_list or []) + if not defaults: + return None + return coerce_param_message_types(defaults, is_streaming=True) + except Exception as e: # noqa: BLE001 - fall back to engine defaults + logger.warning( + "Could not coerce streaming sampling params; using engine defaults: %s", + e, + ) + return None diff --git a/components/src/dynamo/vllm/tests/omni/test_omni_base_handler.py b/components/src/dynamo/vllm/tests/omni/test_omni_base_handler.py index f27dddda4c27..36a5f0314d57 100644 --- a/components/src/dynamo/vllm/tests/omni/test_omni_base_handler.py +++ b/components/src/dynamo/vllm/tests/omni/test_omni_base_handler.py @@ -30,6 +30,7 @@ "sequence_parallel_size", "enable_expert_parallel", "ulysses_mode", + "mask_sp_padding", } diff --git a/components/src/dynamo/vllm/tests/omni/test_omni_realtime_handler.py b/components/src/dynamo/vllm/tests/omni/test_omni_realtime_handler.py new file mode 100644 index 000000000000..7f14dd890830 --- /dev/null +++ b/components/src/dynamo/vllm/tests/omni/test_omni_realtime_handler.py @@ -0,0 +1,424 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the realtime Omni handler's event translation. + +These exercise RealtimeOmniHandler in isolation (no frontend, no vLLM): a fake +engine yields OmniRequestOutput-shaped frames and we assert the OpenAI-spec +server-event sequence the handler emits, including PCM16 round-tripping. +""" + +from __future__ import annotations + +import asyncio +import base64 +from types import SimpleNamespace + +import numpy as np +import pytest + +try: + # Importing the omni package pulls omni_handler -> vllm_omni; the handler + # logic itself is vllm-free, but the package import is not. + # The handler reads audio off a vLLM-Omni MultimodalPayload (``mm.tensors``), + # so the fake engine outputs below must use the real type, not a plain dict. + from vllm_omni.engine.mm_outputs import MultimodalPayload + + from dynamo.vllm.omni.realtime_handler import RealtimeOmniHandler +except (ImportError, ModuleNotFoundError): + pytest.skip("vLLM omni dependencies not available", allow_module_level=True) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.vllm, + pytest.mark.multimodal, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + +MODEL_NAME = "omni-realtime-unit" + + +class _FakeContext: + """Minimal Context stand-in; the handler only calls is_stopped().""" + + def __init__(self, stopped: bool = False) -> None: + self._stopped = stopped + + def is_stopped(self) -> bool: + return self._stopped + + +def _audio_output(samples: np.ndarray, sample_rate: int = 16000): + return SimpleNamespace( + stage_id=1, + outputs=[], + multimodal_output=MultimodalPayload( + tensors={"audio": samples}, metadata={"sr": sample_rate} + ), + ) + + +def _text_output(text: str): + return SimpleNamespace( + stage_id=0, + outputs=[SimpleNamespace(text=text, token_ids=[1, 2, 3])], + prompt_token_ids=[0], + multimodal_output=MultimodalPayload(), + ) + + +class _FakeEngine: + """Echoes appended audio back as two output chunks, preceded by a text delta. + + Consuming the streaming input generator proves the audio-in plumbing; the + canned text and split audio exercise transcript + multi-delta translation. + """ + + def __init__(self, text: str = "hello") -> None: + self.text = text + self.seen_chunks: list = [] + self.seen_output_modalities: list = [] + + async def generate( + self, *, prompt, request_id, sampling_params_list=None, output_modalities=None + ): + self.seen_output_modalities.append(output_modalities) + async for chunk in prompt: + self.seen_chunks.append(chunk) + full = ( + np.concatenate(self.seen_chunks) + if self.seen_chunks + else np.array([], np.float32) + ) + yield _text_output(self.text) + half = len(full) // 2 + yield _audio_output(full[:half]) + yield _audio_output(full[half:]) + + +async def _passthrough_factory(audio_stream, input_stream): + """Stand-in for transcribe_realtime: yield raw float32 chunks unchanged.""" + async for waveform in audio_stream: + yield waveform + + +def _make_handler(engine, **kwargs): + return RealtimeOmniHandler( + engine_client=engine, + model_name=MODEL_NAME, + streaming_input_factory=_passthrough_factory, + **kwargs, + ) + + +async def _drive(handler, events, context): + async def request_stream(): + for ev in events: + yield ev + + return [event async for event in handler.generate(request_stream(), context)] + + +def test_full_turn_event_sequence(): + # Input PCM16 chunk: a short ramp, base64-encoded like the wire format. + pcm16 = np.linspace(-8000, 8000, 64, dtype=np.int16).tobytes() + audio_b64 = base64.b64encode(pcm16).decode("utf-8") + + engine = _FakeEngine(text="hi there") + handler = _make_handler(engine) + + events = [ + { + "type": "session.update", + "session": {"type": "realtime", "model": MODEL_NAME}, + }, + {"type": "input_audio_buffer.append", "audio": audio_b64}, + {"type": "input_audio_buffer.commit"}, + ] + + out = asyncio.run(_drive(handler, events, _FakeContext())) + types = [e["type"] for e in out] + + assert types[0] == "session.updated" + assert out[0]["session"]["model"] == MODEL_NAME + assert "response.created" in types + assert "response.output_audio.delta" in types + assert "response.output_audio.done" in types + assert types[-1] == "response.done" + + # created precedes audio precedes done precedes response.done. + assert types.index("response.created") < types.index("response.output_audio.delta") + assert types.index("response.output_audio.delta") < types.index( + "response.output_audio.done" + ) + assert types.index("response.output_audio.done") < types.index("response.done") + + # response ids are consistent across the turn's frames. + created = next(e for e in out if e["type"] == "response.created") + response_id = created["response"]["id"] + for e in out: + if e["type"] in ("response.output_audio.delta", "response.output_audio.done"): + assert e["response_id"] == response_id + if e["type"] == "response.done": + assert e["response"]["id"] == response_id + assert e["response"]["status"] == "completed" + + # Transcript delta carries the thinker text. + transcripts = [ + e["delta"] for e in out if e["type"] == "response.output_audio_transcript.delta" + ] + assert "".join(transcripts) == "hi there" + + # Concatenated audio deltas decode back to the input PCM16 (echo round-trip). + deltas = b"".join( + base64.b64decode(e["delta"]) + for e in out + if e["type"] == "response.output_audio.delta" + ) + in_f32 = np.frombuffer(pcm16, dtype=np.int16).astype(np.float32) / 32768.0 + out_f32 = np.frombuffer(deltas, dtype=np.int16).astype(np.float32) / 32767.0 + assert out_f32.shape == in_f32.shape + assert np.allclose(out_f32, in_f32, atol=2e-4) + + +def test_unknown_client_events_are_ignored(): + engine = _FakeEngine() + handler = _make_handler(engine, emit_transcript=False) + events = [ + {"type": "session.update", "session": {"model": MODEL_NAME}}, + {"type": "conversation.item.create", "item": {}}, + {"type": "response.cancel"}, + ] + out = asyncio.run(_drive(handler, events, _FakeContext())) + # Only the session.updated echo; no turn started, no error frame. + assert [e["type"] for e in out] == ["session.updated"] + + +def test_stopped_context_emits_no_turn(): + engine = _FakeEngine() + handler = _make_handler(engine) + events = [ + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(b"\x00\x00").decode(), + }, + {"type": "input_audio_buffer.commit"}, + ] + out = asyncio.run(_drive(handler, events, _FakeContext(stopped=True))) + assert out == [] + + +@pytest.mark.parametrize( + "session, expected", + [ + ({"model": MODEL_NAME, "output_modalities": ["audio"]}, ["audio"]), + ({"model": MODEL_NAME}, None), # unset -> engine's launch default (None) + ], +) +def test_session_output_modalities_forwarded_to_engine(session, expected): + audio_b64 = base64.b64encode( + np.linspace(-8000, 8000, 16, dtype=np.int16).tobytes() + ).decode() + engine = _FakeEngine() + handler = _make_handler(engine) + events = [ + {"type": "session.update", "session": session}, + {"type": "input_audio_buffer.append", "audio": audio_b64}, + {"type": "input_audio_buffer.commit"}, + ] + asyncio.run(_drive(handler, events, _FakeContext())) + assert engine.seen_output_modalities == [expected] + + +def test_turn_responses_forwarded_in_order(): + # Two commits => two turns. Turns run concurrently, but their responses are + # forwarded in turn order: distinct response ids, and turn 2's events never + # interleave with turn 1's (turn 1 fully completes before turn 2 emits). + audio_b64 = base64.b64encode( + np.linspace(-8000, 8000, 16, dtype=np.int16).tobytes() + ).decode() + handler = _make_handler(_FakeEngine()) + events = [ + {"type": "session.update", "session": {"model": MODEL_NAME}}, + {"type": "input_audio_buffer.append", "audio": audio_b64}, + {"type": "input_audio_buffer.commit"}, + {"type": "input_audio_buffer.append", "audio": audio_b64}, + {"type": "input_audio_buffer.commit"}, + ] + out = asyncio.run(_drive(handler, events, _FakeContext())) + + created = [e["response"]["id"] for e in out if e["type"] == "response.created"] + done = [e["response"]["id"] for e in out if e["type"] == "response.done"] + assert len(created) == 2 and len(set(created)) == 2, created # two distinct turns + assert created == done, (created, done) # same ids, completed in start order + + def rid(e: dict) -> str | None: + return e.get("response_id") or e.get("response", {}).get("id") + + # Turn 1 fully completes (its response.done) before turn 2 starts, and no + # turn-2-tagged event appears before that point. + first_done = next( + i + for i, e in enumerate(out) + if e["type"] == "response.done" and e["response"]["id"] == created[0] + ) + second_created = next( + i + for i, e in enumerate(out) + if e["type"] == "response.created" and e["response"]["id"] == created[1] + ) + assert first_done < second_created + assert all(rid(e) != created[1] for e in out[: first_done + 1]) + + +def test_turns_run_concurrently(): + # Both turns' engine drives must be in flight at once: a barrier engine + # blocks each generate() until two have started. This only completes if + # turns are not serialized -- a one-at-a-time model would deadlock the + # barrier and time out. Responses are still forwarded in turn order. + class _BarrierEngine: + def __init__(self, n: int) -> None: + self.n = n + self.started = 0 + self.gate = asyncio.Event() + + async def generate( + self, + *, + prompt, + request_id, + sampling_params_list=None, + output_modalities=None, + ): + async for _ in prompt: # drain this turn's audio + pass + self.started += 1 + if self.started >= self.n: + self.gate.set() + await asyncio.wait_for(self.gate.wait(), timeout=5) + yield _text_output("ok") + yield _audio_output(np.zeros(4, dtype=np.float32)) + + audio_b64 = base64.b64encode( + np.linspace(-8000, 8000, 16, dtype=np.int16).tobytes() + ).decode() + engine = _BarrierEngine(2) + handler = _make_handler(engine) + events = [ + {"type": "session.update", "session": {"model": MODEL_NAME}}, + {"type": "input_audio_buffer.append", "audio": audio_b64}, + {"type": "input_audio_buffer.commit"}, + {"type": "input_audio_buffer.append", "audio": audio_b64}, + {"type": "input_audio_buffer.commit"}, + ] + out = asyncio.run(_drive(handler, events, _FakeContext())) + + # Both drives passed the barrier => they were concurrently in flight. + assert engine.started == 2 + assert not any(e["type"] == "error" for e in out) + created = [e["response"]["id"] for e in out if e["type"] == "response.created"] + done = [e["response"]["id"] for e in out if e["type"] == "response.done"] + assert len(set(created)) == 2 # two distinct turns, both completed + assert created == done # forwarded in turn order, non-interleaved + + +def test_engine_failure_emits_error_and_failed_response_done(): + # When generation throws, the turn surfaces a top-level `error` event (human + # readable, but no response_id) AND a terminal response.done(status=failed) + # so the in-progress response closes and is correlatable by response id. + class _FailingEngine: + async def generate( + self, + *, + prompt, + request_id, + sampling_params_list=None, + output_modalities=None, + ): + async for _ in prompt: # drain this turn's audio + pass + yield _text_output("partial") + raise RuntimeError("boom") + + audio_b64 = base64.b64encode( + np.linspace(-8000, 8000, 16, dtype=np.int16).tobytes() + ).decode() + handler = _make_handler(_FailingEngine()) + events = [ + {"type": "session.update", "session": {"model": MODEL_NAME}}, + {"type": "input_audio_buffer.append", "audio": audio_b64}, + {"type": "input_audio_buffer.commit"}, + ] + out = asyncio.run(_drive(handler, events, _FakeContext())) + types = [e["type"] for e in out] + + created = next(e for e in out if e["type"] == "response.created") + response_id = created["response"]["id"] + + error = next((e for e in out if e["type"] == "error"), None) + assert error is not None + assert error["error"]["code"] == "omni_generation_error" + + done = next((e for e in out if e["type"] == "response.done"), None) + assert done is not None + assert done["response"]["id"] == response_id # correlatable, unlike `error` + assert done["response"]["status"] == "failed" + assert done["response"]["status_details"]["type"] == "failed" + assert ( + done["response"]["status_details"]["error"]["code"] == "omni_generation_error" + ) + assert types.index("error") < types.index("response.done") + + +def test_concurrent_turns_capped(): + # With max_concurrent_turns=2, four commits must never put more than two + # engine drives in flight at once: the pump blocks opening turn 3 until an + # earlier turn finishes (and frees its slot). A barrier that releases once + # `cap` drives have started lets the two in-flight turns unblock each other, + # so the test still completes -- peak concurrency pinned at the cap. + cap = 2 + + class _PeakEngine: + def __init__(self) -> None: + self.inflight = 0 + self.peak = 0 + self.started = 0 + self.gate = asyncio.Event() + + async def generate( + self, + *, + prompt, + request_id, + sampling_params_list=None, + output_modalities=None, + ): + async for _ in prompt: # drain this turn's audio + pass + self.inflight += 1 + self.peak = max(self.peak, self.inflight) + self.started += 1 + if self.started >= cap: + self.gate.set() + await asyncio.wait_for(self.gate.wait(), timeout=5) + yield _text_output("ok") + yield _audio_output(np.zeros(4, dtype=np.float32)) + self.inflight -= 1 + + audio_b64 = base64.b64encode( + np.linspace(-8000, 8000, 16, dtype=np.int16).tobytes() + ).decode() + engine = _PeakEngine() + handler = _make_handler(engine, max_concurrent_turns=cap) + events = [{"type": "session.update", "session": {"model": MODEL_NAME}}] + for _ in range(4): # four commits => four turns + events.append({"type": "input_audio_buffer.append", "audio": audio_b64}) + events.append({"type": "input_audio_buffer.commit"}) + + out = asyncio.run(_drive(handler, events, _FakeContext())) + + assert engine.started == 4 # all four turns ran + assert engine.peak == cap # never more than the cap in flight at once + done = [e for e in out if e["type"] == "response.done"] + assert len(done) == 4 and all(e["response"]["status"] == "completed" for e in done) diff --git a/components/src/dynamo/vllm/tests/test_vllm_api_contract.py b/components/src/dynamo/vllm/tests/test_vllm_api_contract.py index eb38ddbd24c2..3e4d142b4849 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_api_contract.py +++ b/components/src/dynamo/vllm/tests/test_vllm_api_contract.py @@ -129,7 +129,17 @@ def test_request_exposes_all_token_ids(): private attribute so a rename is caught here, not at runtime.""" from vllm.v1.request import Request - assert "_all_token_ids" in inspect.getsource(Request), ( + src = inspect.getsource(Request) + # [gluo NOTE] the test suit will attempt to import vllm-omni at conftest for test + # selection. However, omni will monkeypatch Request so naive source inspection will + # fail (only see OmniRequest's source) walk MRO so a subclass (OmniRequest) still gets + # the base that defines it + src = "".join( + inspect.getsource(c) for c in Request.__mro__ if c.__module__.startswith("vllm") + ) + assert "_all_token_ids" in src + + assert "_all_token_ids" in src, ( "vllm.v1.request.Request no longer exposes `_all_token_ids` — " "InstrumentedScheduler relies on it for NewRequestData.prefill_token_ids." ) diff --git a/components/src/dynamo/vllm/tests/test_vllm_renderer_api.py b/components/src/dynamo/vllm/tests/test_vllm_renderer_api.py index f40b48617961..6d47b97dc056 100755 --- a/components/src/dynamo/vllm/tests/test_vllm_renderer_api.py +++ b/components/src/dynamo/vllm/tests/test_vllm_renderer_api.py @@ -431,9 +431,11 @@ def test_engine_core_struct_contract(self): "routed_experts", "num_nans_in_logits", ) - # vllm-omni extends EngineCoreOutput with streaming segment metadata - # (only installed on amd64, not arm64). + # vllm-omni extends EngineCoreOutput with a multimodal output channel + # and streaming segment metadata (only installed on amd64, not arm64). + # Declaration order in OmniEngineCoreOutput determines wire position. omni_output_extra_fields = ( + "multimodal_output", "is_segment_finished", "new_prompt_len_snapshot", ) diff --git a/container/context.yaml b/container/context.yaml index 2f2f40b24cd9..5fab16fceeb3 100644 --- a/container/context.yaml +++ b/container/context.yaml @@ -77,7 +77,7 @@ vllm: runtime_image_tag: v0.24.0 # baseline_sbom: not yet captured for cpu — runtime build runs without subtraction flashinf_ref: v0.6.8.post1 - vllm_omni_ref: "v0.21.0rc1" + vllm_omni_ref: "v0.23.0rc1" nixl_ref: v1.1.0 max_jobs: "10" enable_media_ffmpeg: "false" diff --git a/container/deps/vllm/install_vllm_omni.sh b/container/deps/vllm/install_vllm_omni.sh index 929e9a4821cd..527a25fe7c25 100755 --- a/container/deps/vllm/install_vllm_omni.sh +++ b/container/deps/vllm/install_vllm_omni.sh @@ -48,3 +48,44 @@ else --constraints "${PROTECTED_CONSTRAINTS}" \ "vllm-omni==${VLLM_OMNI_VERSION}" fi + +# Cherry-pick vllm-project/vllm-omni#4568 onto the released wheel. +# +# vLLM-Omni globally monkeypatches vllm.v1.request.Request with its OmniRequest +# subclass at import time, and the test suite imports vllm_omni for collection, +# so this applies to every vLLM worker in the image -- not just omni modes. In +# the released v0.23.0rc1, OmniRequest.__init__ still declares `*args` after its +# named parameters, so vLLM 0.23's positional Request(...) construction misbinds +# the arguments and EngineCore initialization fails for all vLLM workers. The fix +# moves `*args` to the front and forwards cleanly. Drop this once a vllm-omni +# release includes the change. +# https://github.com/vllm-project/vllm-omni/commit/17cf60a63d240608653c4532084a4c00d6f02216 +VLLM_OMNI_CHERRY_PICK_COMMIT="17cf60a63d240608653c4532084a4c00d6f02216" + +omni_site="$(python3 -c 'import importlib.util, os; print(os.path.dirname(os.path.dirname(importlib.util.find_spec("vllm_omni").origin)))')" +full_patch="$(mktemp /tmp/vllm-omni-commit.XXXXXX.patch)" +cherry_pick_patch="$(mktemp /tmp/vllm-omni-cherry-pick.XXXXXX.patch)" + +curl -fsSL \ + "https://github.com/vllm-project/vllm-omni/commit/${VLLM_OMNI_CHERRY_PICK_COMMIT}.patch" \ + -o "${full_patch}" +# Keep only the vllm_omni/request.py hunk; the commit's new test file is not part +# of the installed wheel. +awk '/^diff --git a\/vllm_omni\/request.py/{f=1} f' "${full_patch}" > "${cherry_pick_patch}" + +if [ ! -s "${cherry_pick_patch}" ]; then + echo "ERROR: could not extract request.py hunk from vllm-omni commit ${VLLM_OMNI_CHERRY_PICK_COMMIT}" >&2 + exit 1 +fi + +if patch -p1 -d "${omni_site}" --forward --dry-run < "${cherry_pick_patch}" >/dev/null 2>&1; then + patch -p1 -d "${omni_site}" --forward < "${cherry_pick_patch}" + echo "Applied vllm-omni cherry-pick ${VLLM_OMNI_CHERRY_PICK_COMMIT}" +elif patch -p1 -d "${omni_site}" --reverse --dry-run < "${cherry_pick_patch}" >/dev/null 2>&1; then + echo "vllm-omni cherry-pick ${VLLM_OMNI_CHERRY_PICK_COMMIT} already present; skipping" +else + echo "ERROR: vllm-omni cherry-pick ${VLLM_OMNI_CHERRY_PICK_COMMIT} does not apply cleanly to the installed package" >&2 + exit 1 +fi + +rm -f "${full_patch}" "${cherry_pick_patch}" diff --git a/examples/backends/vllm/launch/agg_omni_realtime.sh b/examples/backends/vllm/launch/agg_omni_realtime.sh new file mode 100755 index 000000000000..9caefb3c5697 --- /dev/null +++ b/examples/backends/vllm/launch/agg_omni_realtime.sh @@ -0,0 +1,62 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -e +trap 'echo Cleaning up...; kill 0' EXIT + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../../common/launch_utils.sh" + +MODEL="Qwen/Qwen3-Omni-30B-A3B-Instruct" + +# Parse command line arguments +EXTRA_ARGS=() +while [[ $# -gt 0 ]]; do + case $1 in + --model) + MODEL="$2" + shift 2 + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +GPU_MEM_ARGS=$(build_vllm_gpu_mem_args) +print_launch_banner --no-curl "Launching vLLM-Omni Realtime (1 GPU)" "$MODEL" "$HTTP_PORT" +print_curl_footer < tuple[np.ndarray, int]: + """Read audio as float32 + sample rate via soundfile, or stdlib wave (PCM16).""" + try: + import soundfile as sf + + audio, sr = sf.read(path, dtype="float32", always_2d=True) + return audio, sr + except ImportError: + with wave.open(path, "rb") as wf: + if wf.getsampwidth() != 2: + raise SystemExit( + "soundfile not installed and input is not 16-bit PCM WAV; " + "install soundfile or pre-convert to PCM16 " + "(ffmpeg -i in.wav -ac 1 -ar 16000 -sample_fmt s16 out.wav)" + ) + sr = wf.getframerate() + frames = np.frombuffer(wf.readframes(wf.getnframes()), dtype=" str: + """Download the sample clip from GitHub to a temp file; return its path.""" + dest = os.path.join(tempfile.gettempdir(), "realtime_omni_" + os.path.basename(url)) + if not os.path.exists(dest): + print(f"[client] no --input-audio given; fetching sample from {url}") + urllib.request.urlretrieve(url, dest) # noqa: S310 - fixed https GitHub URL + else: + print(f"[client] using cached sample {dest}") + return dest + + +def _load_pcm16_16k(path: str | None, default_url: str) -> bytes: + """Load an audio file as 16 kHz mono PCM16 bytes, fetching a sample if needed.""" + if path is None: + path = _fetch_default_audio(default_url) + elif not os.path.isfile(path): + raise SystemExit( + f"--input-audio file not found: {path!r}. Pass a real audio file, " + "or omit --input-audio to fetch the vLLM-Omni sample clip." + ) + + audio, sr = _read_audio(path) + if audio.ndim > 1: + audio = audio.mean(axis=1) # downmix to mono + if sr != INPUT_SAMPLE_RATE: + # Linear resample to 16 kHz (adequate for a demo client). + duration = audio.shape[0] / sr + tgt_len = int(duration * INPUT_SAMPLE_RATE) + xp = np.linspace(0.0, duration, num=audio.shape[0], endpoint=False) + x = np.linspace(0.0, duration, num=tgt_len, endpoint=False) + audio = np.interp(x, xp, audio).astype(np.float32) + pcm16 = np.clip(audio, -1.0, 1.0) + pcm16 = (pcm16 * 32767.0).astype(" None: + with wave.open(path, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(pcm16_bytes) + + +async def run(args: argparse.Namespace) -> int: + pcm16 = _load_pcm16_16k(args.input_audio, args.input_audio_url) + chunk_bytes = max(INPUT_SAMPLE_RATE * 2 // 1000 * args.chunk_ms, 2) + + os.makedirs(args.output_dir, exist_ok=True) + + # Each response.output_audio.delta is written to its own chunk_NNNN.wav in + # the output folder as it arrives, and also appended to this buffer so the + # whole response can be written as a single concatenated response.wav. + audio_out = bytearray() + audio_delta_count = 0 + transcript = [] + response_id = None + status = None + + async with aiohttp.ClientSession() as session: + async with session.ws_connect(args.url, max_msg_size=64 * 1024 * 1024) as ws: + print(f"[client] connected to {args.url}") + + # 1) Select the model (the frontend picks the engine on + # session.update) and request output modalities. Asking for + # "audio" drives the Omni talker; "text" yields transcript only. + session_block = {"type": "realtime", "model": args.model} + if args.output_modalities: + session_block["output_modalities"] = args.output_modalities + await ws.send_str( + json.dumps({"type": "session.update", "session": session_block}) + ) + + # 2) Stream the audio in chunks, then commit to trigger generation. + for i in range(0, len(pcm16), chunk_bytes): + await ws.send_str( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode( + pcm16[i : i + chunk_bytes] + ).decode(), + } + ) + ) + print(f"[client] sent {i + chunk_bytes} bytes of audio\n") + await asyncio.sleep(0.05) + await ws.send_str(json.dumps({"type": "input_audio_buffer.commit"})) + print("[client] commit audio\n") + + # 3) Print every server event until the response completes. + while True: + msg = await asyncio.wait_for(ws.receive(), timeout=args.timeout) + if msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED): + print(f"[client] socket closed: {msg.data!r} {msg.extra!r}") + break + if msg.type is not aiohttp.WSMsgType.TEXT: + continue + event = json.loads(msg.data) + etype = event.get("type") + + if etype == "response.output_audio.delta": + delta = base64.b64decode(event.get("delta", "")) + if delta: + audio_delta_count += 1 + audio_out.extend(delta) + chunk_path = os.path.join( + args.output_dir, f"chunk_{audio_delta_count:04d}.wav" + ) + _write_wav(chunk_path, delta, args.output_sample_rate) + print( + f"<- response.output_audio.delta ({len(delta)} bytes) " + f"-> {os.path.basename(chunk_path)}" + ) + elif etype == "response.output_audio_transcript.delta": + transcript.append(event.get("delta", "")) + print(f"<- transcript.delta: {event.get('delta')!r}") + elif etype == "response.created": + response_id = event["response"]["id"] + print(f"<- response.created (id={response_id})") + elif etype == "response.done": + status = event["response"]["status"] + print(f"<- response.done (status={status})") + break + elif etype == "error": + print(f"<- ERROR: {json.dumps(event.get('error'), indent=2)}") + break + else: + print(f"<- {etype}") + + print("\n[client] === summary ===") + print(f" response_id : {response_id}") + print(f" status : {status}") + print(f" transcript : {''.join(transcript)!r}") + print( + f" audio : {audio_delta_count} delta(s) joined -> " + f"{len(audio_out)} bytes ({len(audio_out) // 2} samples)" + ) + if audio_out: + # Concatenate every audio delta into one WAV alongside the chunk files. + concat_path = os.path.join(args.output_dir, "response.wav") + _write_wav(concat_path, bytes(audio_out), args.output_sample_rate) + print( + f" saved audio : {audio_delta_count} chunk file(s) + " + f"{os.path.basename(concat_path)} in {args.output_dir}/ " + f"@ {args.output_sample_rate} Hz" + ) + else: + print(" saved audio : none (no audio modality / no audio returned)") + return 0 if status == "completed" else 1 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--url", default="ws://localhost:8000/v1/realtime") + parser.add_argument("--model", required=True, help="served model name") + parser.add_argument( + "--input-audio", default=None, help="any soundfile-readable audio" + ) + parser.add_argument( + "--input-audio-url", + default=DEFAULT_AUDIO_URL, + help="sample audio URL fetched when --input-audio is omitted", + ) + parser.add_argument( + "--output-modalities", + nargs="+", + choices=["audio", "text"], + default=["audio"], + help="output modalities requested via session.update " + "('audio' drives the Omni talker). Pass --output-modalities text " + "for transcript only, or 'text audio' for both.", + ) + parser.add_argument( + "--output-dir", + default="realtime_output", + help="folder to write per-delta chunk_NNNN.wav files and the " + "concatenated response.wav into", + ) + parser.add_argument( + "--output-sample-rate", + type=int, + default=24000, + help="sample rate to save the response audio at (Omni talker default 24kHz)", + ) + parser.add_argument( + "--chunk-ms", type=int, default=100, help="append chunk size in ms" + ) + parser.add_argument( + "--timeout", type=float, default=120.0, help="per-frame recv timeout" + ) + args = parser.parse_args() + sys.exit(asyncio.run(run(args))) + + +if __name__ == "__main__": + main() diff --git a/tests/frontend/realtime_omni_mock_worker.py b/tests/frontend/realtime_omni_mock_worker.py new file mode 100644 index 000000000000..f907cc1f6b8e --- /dev/null +++ b/tests/frontend/realtime_omni_mock_worker.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Realtime Omni worker driven by a mock vLLM-Omni engine, for the bridge e2e. + +Serves the real ``RealtimeOmniHandler`` (the production translation layer) but +backs it with a fake AsyncOmni that echoes appended audio back as +``OmniRequestOutput``-shaped frames. This exercises the full bidirectional +bridge — frontend ``/v1/realtime`` -> PushRouter -> Python engine -> handler — +without a GPU or model download, so it runs in the same file-discovery e2e +shape as ``realtime_echo_worker.py``. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import numpy as np +import uvloop +from vllm_omni.engine.mm_outputs import MultimodalPayload + +from dynamo.llm import ModelInput, ModelType, WorkerType, register_model +from dynamo.runtime import DistributedRuntime +from dynamo.vllm.omni.realtime_handler import RealtimeOmniHandler +from tests.frontend.test_realtime_omni_bridge import ( + ENDPOINT_PATH, + MOCK_TRANSCRIPT, + MODEL_NAME, +) + + +async def _passthrough_factory(audio_stream, input_stream): + """Stand-in for ``OpenAIServingRealtime.transcribe_realtime``. + + The real factory buffers audio into model prompts; the mock engine only + needs the raw float32 waveforms, so we yield each audio chunk straight + through. ``input_stream`` (the talker token-feedback queue) is unused here. + """ + async for waveform in audio_stream: + yield waveform + + +class _MockAsyncOmni: + """Fake AsyncOmni: drains the streaming audio input, then echoes it back. + + Yields a stage-0 text frame (transcript) followed by the accumulated audio + as a single multimodal output, matching the fields RealtimeOmniHandler reads + off a real ``OmniRequestOutput`` (``stage_id``, ``outputs[].text``, and a + ``MultimodalPayload`` whose ``tensors['audio']`` holds the waveform). + """ + + default_sampling_params_list: list = [] + + async def generate( + self, *, prompt, request_id, sampling_params_list=None, output_modalities=None + ): + chunks = [chunk async for chunk in prompt] + full = np.concatenate(chunks) if chunks else np.zeros(1, dtype=np.float32) + yield SimpleNamespace( + stage_id=0, + outputs=[SimpleNamespace(text=MOCK_TRANSCRIPT, token_ids=[1])], + prompt_token_ids=[0], + multimodal_output=MultimodalPayload(), + ) + yield SimpleNamespace( + stage_id=1, + outputs=[], + multimodal_output=MultimodalPayload( + tensors={"audio": full}, metadata={"sr": 16000} + ), + ) + + +async def main() -> None: + runtime = DistributedRuntime(asyncio.get_running_loop(), "file", "tcp") + endpoint = runtime.endpoint(ENDPOINT_PATH) + handler = RealtimeOmniHandler( + engine_client=_MockAsyncOmni(), + model_name=MODEL_NAME, + streaming_input_factory=_passthrough_factory, + ) + await register_model( + ModelInput.Text, + ModelType.Realtime, + endpoint, + MODEL_NAME, + model_name=MODEL_NAME, + worker_type=WorkerType.Aggregated, + ) + await endpoint.serve_bidirectional_endpoint(handler.generate) + + +if __name__ == "__main__": + uvloop.run(main()) diff --git a/tests/frontend/test_realtime_omni_bridge.py b/tests/frontend/test_realtime_omni_bridge.py new file mode 100644 index 000000000000..51cc999c4943 --- /dev/null +++ b/tests/frontend/test_realtime_omni_bridge.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +End-to-end realtime WebSocket test for the vLLM-Omni realtime bridge. + +A launched ``dynamo.frontend`` discovers a mock-Omni realtime worker (the real +``RealtimeOmniHandler`` backed by a fake AsyncOmni that echoes audio) and +installs a typed realtime PushRouter to it. A WebSocket client connects to +``/v1/realtime``, drives OpenAI Realtime client events, and asserts the +spec-shaped server events come back — exercising the full bridge without a GPU +or model download. + +Discovery uses the file backend (``DYN_FILE_KV``) and the tcp request plane, so +the two processes coordinate without etcd or nats. Mirrors +``test_realtime_python_bridge.py``. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import logging + +import aiohttp +import numpy as np +import pytest +import requests + +from tests.utils.managed_process import DynamoFrontendProcess, ManagedProcess +from tests.utils.port_utils import ServicePorts + +logger = logging.getLogger(__name__) + +# Shared with the worker module (realtime_omni_mock_worker.py imports these). +MODEL_NAME = "omni-realtime-mock" +ENDPOINT_PATH = "test_omni_ws_e2e.realtime.generate" +MOCK_TRANSCRIPT = "mock omni transcript" + +pytestmark = [ + pytest.mark.pre_merge, + pytest.mark.integration, + pytest.mark.vllm, + pytest.mark.multimodal, + pytest.mark.gpu_0, +] + + +class RealtimeOmniMockWorkerProcess(ManagedProcess): + """Launch the mock-Omni realtime worker; ready once the frontend lists it.""" + + def __init__(self, request, *, frontend_port: int) -> None: + super().__init__( + command=["python3", "-m", "tests.frontend.realtime_omni_mock_worker"], + health_check_urls=[ + (f"http://localhost:{frontend_port}/v1/models", self._model_listed) + ], + timeout=60, + display_output=True, + terminate_all_matching_process_names=False, + straggler_commands=["-m tests.frontend.realtime_omni_mock_worker"], + log_dir=f"{request.node.name}_realtime_omni_worker", + ) + + @staticmethod + def _model_listed(response: requests.Response) -> bool: + try: + if response.status_code != 200: + return False + data = response.json() + except (ValueError, KeyError): + return False + return any(model.get("id") == MODEL_NAME for model in data.get("data", [])) + + +@pytest.fixture(scope="function") +def realtime_omni_frontend( + request, file_storage_backend, dynamo_dynamic_ports: ServicePorts +): + """Launch the frontend + mock-Omni worker; yield the frontend port once discovered. + + Uses file-based discovery (the ``file_storage_backend`` fixture sets + ``DYN_FILE_KV``), the tcp request plane, and an explicit zmq event plane, so + the two processes coordinate without etcd or nats and an ambient + ``NATS_SERVER`` never forces a connection. Mirrors ``test_prompt_embeds.py``. + """ + _ = file_storage_backend # sets DYN_FILE_KV for both subprocesses + frontend_port = dynamo_dynamic_ports.frontend_port + with DynamoFrontendProcess( + request, + frontend_port=frontend_port, + extra_args=[ + "--discovery-backend", + "file", + "--request-plane", + "tcp", + "--event-plane", + "zmq", + ], + terminate_all_matching_process_names=False, + ): + logger.info("Frontend started on port %s", frontend_port) + with RealtimeOmniMockWorkerProcess(request, frontend_port=frontend_port): + logger.info("Mock-Omni realtime worker registered %s", MODEL_NAME) + yield frontend_port + + +async def _recv_json(ws: aiohttp.ClientWebSocketResponse, timeout_s: float) -> dict: + msg = await asyncio.wait_for(ws.receive(), timeout=timeout_s) + if msg.type is not aiohttp.WSMsgType.TEXT: + raise AssertionError(f"unexpected websocket frame: {msg.type!r} {msg.data!r}") + return json.loads(msg.data) + + +async def _drain_until( + ws: aiohttp.ClientWebSocketResponse, expected_type: str, timeout_s: float = 5.0 +) -> dict: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout_s + while loop.time() < deadline: + remaining = deadline - loop.time() + event = await _recv_json(ws, max(remaining, 0.01)) + if event.get("type") == expected_type: + return event + raise AssertionError( + f"timed out waiting for a {expected_type!r} frame on the websocket" + ) + + +async def _audio_round_trip(port: int) -> None: + async with aiohttp.ClientSession() as session: + async with session.ws_connect(f"ws://127.0.0.1:{port}/v1/realtime") as ws: + await _drain_until(ws, "session.created") + + await ws.send_str( + json.dumps( + { + "type": "session.update", + "session": {"type": "realtime", "model": MODEL_NAME}, + } + ) + ) + await _drain_until(ws, "session.updated") + + # Send a short PCM16 ramp; the mock engine echoes it back as audio. + pcm16 = np.linspace(-8000, 8000, 128, dtype=np.int16).tobytes() + await ws.send_str( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(pcm16).decode("utf-8"), + } + ) + ) + await ws.send_str(json.dumps({"type": "input_audio_buffer.commit"})) + + response_id: str | None = None + audio_b64_parts: list[str] = [] + transcript_parts: list[str] = [] + saw_audio_done = False + response_done_status: str | None = None + + loop = asyncio.get_event_loop() + deadline = loop.time() + 10.0 + while response_done_status is None: + remaining = deadline - loop.time() + if remaining <= 0: + raise AssertionError( + "timed out before response.done; " + f"audio_parts={len(audio_b64_parts)}, " + f"saw_audio_done={saw_audio_done}" + ) + event = await _recv_json(ws, remaining) + etype = event.get("type") + if etype == "response.created": + response_id = event["response"]["id"] + elif etype == "response.output_audio_transcript.delta": + transcript_parts.append(event["delta"]) + assert event["response_id"] == response_id, event + elif etype == "response.output_audio.delta": + audio_b64_parts.append(event["delta"]) + assert event["response_id"] == response_id, event + elif etype == "response.output_audio.done": + saw_audio_done = True + assert event["response_id"] == response_id, event + elif etype == "response.done": + response_done_status = event["response"]["status"] + assert event["response"]["id"] == response_id, event + else: + raise AssertionError(f"unexpected event type {etype!r}: {event}") + + assert response_id is not None + assert saw_audio_done, "engine should emit response.output_audio.done" + assert response_done_status == "completed", response_done_status + assert "".join(transcript_parts) == MOCK_TRANSCRIPT, transcript_parts + + # Concatenated audio deltas decode back to the input ramp (echo). + out_bytes = b"".join(base64.b64decode(p) for p in audio_b64_parts) + in_f32 = np.frombuffer(pcm16, dtype=np.int16).astype(np.float32) / 32768.0 + out_f32 = ( + np.frombuffer(out_bytes, dtype=np.int16).astype(np.float32) / 32767.0 + ) + assert out_f32.shape == in_f32.shape, (out_f32.shape, in_f32.shape) + assert np.allclose(out_f32, in_f32, atol=2e-4) + + +@pytest.mark.timeout(120) +def test_websocket_audio_round_trip(realtime_omni_frontend) -> None: + """Appended audio echoes back as the full spec response envelope.""" + asyncio.run(_audio_round_trip(realtime_omni_frontend)) From 8b63aaae605838b070380ebd837e85b6f09cb41c Mon Sep 17 00:00:00 2001 From: "Zhuangcheng(Jesse) Gu" <40918450+Chokoyo@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:33:58 -0700 Subject: [PATCH 018/320] perf(sglang): add video embedding cache benchmark workload (#10000) Signed-off-by: Zhuangcheng(Jesse) Gu --- benchmarks/multimodal/jsonl/README.md | 65 ++++++- benchmarks/multimodal/jsonl/args.py | 82 +++++++- .../multimodal/jsonl/generate_videos.py | 178 ++++++++++++++++++ benchmarks/multimodal/jsonl/main.py | 56 ++++++ benchmarks/multimodal/sweep/README.md | 1 + .../embedding_cache/sglang_e_pd.yaml | 38 ++++ .../benchmarks/multimodal/jsonl/test_main.py | 130 +++++++++++++ 7 files changed, 540 insertions(+), 10 deletions(-) create mode 100644 benchmarks/multimodal/jsonl/generate_videos.py create mode 100644 benchmarks/multimodal/sweep/experiments/embedding_cache/sglang_e_pd.yaml diff --git a/benchmarks/multimodal/jsonl/README.md b/benchmarks/multimodal/jsonl/README.md index 0b9376e5a1ab..6fc190373a81 100644 --- a/benchmarks/multimodal/jsonl/README.md +++ b/benchmarks/multimodal/jsonl/README.md @@ -1,6 +1,6 @@ # Multimodal JSONL Request Generator -Generates `.jsonl` benchmark files for [aiperf](https://github.com/ai-dynamo/aiperf) with single-turn multimodal requests (text + images). +Generates `.jsonl` benchmark files for [aiperf](https://github.com/ai-dynamo/aiperf) with multimodal requests (text + images or text + videos). ## Key concept: image pool reuse @@ -25,6 +25,51 @@ wget http://images.cocodataset.org/annotations/image_info_test2017.zip unzip image_info_test2017.zip ``` +## Video modes + +Video workloads use the same pool-reuse idea as images, but emit aiperf's +`videos` field instead of `images`. `--videos-pool` is the number of unique +videos to sample from, and `--videos-per-request` controls how many video slots +each request gets. A smaller pool relative to total video slots produces more +cross-request video reuse. + +Video modes generate deterministic local synthetic MP4 files under +`--synthetic-video-dir`. The generated content is derived from `--seed`, the +video index, and the synthetic video parameters, so rerunning the same command +on the same machine reuses equivalent clips. This is the simplest mode for +measuring embedding-cache reuse from repeated local video inputs. + +`video-single-turn` measures cross-request reuse from a shared video pool. + +To generate the video inputs referenced by +`benchmarks/multimodal/sweep/experiments/embedding_cache/sglang_e_pd.yaml`: + +```bash +python benchmarks/multimodal/jsonl/main.py video-single-turn \ + -n 100 \ + --videos-per-request 1 \ + --videos-pool 20 \ + --user-text-tokens 300 \ + --synthetic-video-dir /tmp/bench_videos \ + -o benchmarks/multimodal/jsonl/100req_1vid_20pool_300word_local.jsonl + +python benchmarks/multimodal/jsonl/main.py video-single-turn \ + -n 100 \ + --videos-per-request 1 \ + --videos-pool 80 \ + --user-text-tokens 300 \ + --synthetic-video-dir /tmp/bench_videos \ + -o benchmarks/multimodal/jsonl/100req_1vid_80pool_300word_local.jsonl + +python benchmarks/multimodal/jsonl/main.py video-single-turn \ + -n 100 \ + --videos-per-request 2 \ + --videos-pool 160 \ + --user-text-tokens 300 \ + --synthetic-video-dir /tmp/bench_videos \ + -o benchmarks/multimodal/jsonl/100req_2vid_160pool_300word_local.jsonl +``` + ## Usage ```bash @@ -39,9 +84,25 @@ python main.py -n 200 --images-pool 100 # More images per request python main.py -n 100 --images-per-request 20 --images-pool 500 + +# Video workload with repeated clips +python main.py video-single-turn \ + -n 100 \ + --videos-per-request 1 \ + --videos-pool 20 \ + --synthetic-video-dir /tmp/bench_videos \ + --seed 1 + +# Multiple videos per request with the same pool-reuse semantics +python main.py video-single-turn \ + -n 100 \ + --videos-per-request 2 \ + --videos-pool 160 + ``` -Output filename encodes the parameters, e.g. `500req_3img_200pool_300word_http.jsonl`. +Output filename encodes the parameters, e.g. `500req_3img_200pool_300word_http.jsonl` +or `100req_1vid_20pool_300word_local.jsonl`. ## Running with aiperf diff --git a/benchmarks/multimodal/jsonl/args.py b/benchmarks/multimodal/jsonl/args.py index 1028acbb074e..7b812b68b0ef 100644 --- a/benchmarks/multimodal/jsonl/args.py +++ b/benchmarks/multimodal/jsonl/args.py @@ -8,6 +8,7 @@ from pathlib import Path DEFAULT_IMAGES_PER_REQUEST = 3 +DEFAULT_VIDEOS_PER_REQUEST = 1 USER_TEXT_TOKENS = 300 COCO_ANNOTATIONS = Path(__file__).parent / "annotations" / "image_info_test2017.json" @@ -41,13 +42,6 @@ def _common_parser() -> argparse.ArgumentParser: default=None, help="Random seed for reproducible generation (default: time-based)", ) - p.add_argument( - "--uuid", - action=argparse.BooleanOptionalAction, - default=False, - help="Emit `image_uuids` parallel to `images` in each JSONL row (default: False). " - "Pass --uuid to enable for aiperf --mm-cache-mode {uuid-only,uuid-and-strip} runs.", - ) return p @@ -80,12 +74,53 @@ def _image_parser() -> argparse.ArgumentParser: default=COCO_ANNOTATIONS, help=f"Path to COCO image_info JSON for --image-mode http (default: {COCO_ANNOTATIONS})", ) + p.add_argument( + "--uuid", + action=argparse.BooleanOptionalAction, + default=False, + help="Emit `image_uuids` parallel to `images` in each JSONL row (default: False). " + "Pass --uuid to enable for aiperf --mm-cache-mode {uuid-only,uuid-and-strip} runs.", + ) + return p + + +def _video_parser() -> argparse.ArgumentParser: + """Args for synthetic video workloads.""" + p = argparse.ArgumentParser(add_help=False) + p.add_argument( + "--synthetic-video-dir", + type=Path, + default=Path("/tmp/bench_videos"), + help="Directory for generated synthetic MP4 videos " + "(default: /tmp/bench_videos)", + ) + p.add_argument( + "--synthetic-video-size", + type=_positive_int, + nargs=2, + default=[320, 240], + metavar=("WIDTH", "HEIGHT"), + help="Size of generated synthetic MP4 videos in pixels (default: 320 240)", + ) + p.add_argument( + "--synthetic-video-fps", + type=_positive_int, + default=8, + help="Frames per second for generated synthetic MP4 videos (default: 8)", + ) + p.add_argument( + "--synthetic-video-seconds", + type=_positive_int, + default=4, + help="Duration for generated synthetic MP4 videos (default: 4)", + ) return p def parse_args(description: str = "") -> argparse.Namespace: common = _common_parser() image = _image_parser() + video = _video_parser() parser = argparse.ArgumentParser( description=description, @@ -146,10 +181,41 @@ def parse_args(description: str = "") -> argparse.Namespace: "with window_size-1 overlap between consecutive turns (default: 5)", ) + # --- video-single-turn --- + vst = sub.add_parser( + "video-single-turn", + parents=[common, video], + help="Independent requests with random video sampling", + ) + vst.add_argument( + "-n", + "--num-requests", + type=int, + default=200, + help="Number of requests to generate (default: 200)", + ) + vst.add_argument( + "--videos-per-request", + type=int, + default=DEFAULT_VIDEOS_PER_REQUEST, + help=f"Number of videos per request (default: {DEFAULT_VIDEOS_PER_REQUEST})", + ) + vst.add_argument( + "--videos-pool", + type=int, + default=None, + help="Unique videos in pool. Smaller pool = more cross-request reuse. " + "Default: num_requests * videos_per_request (all unique).", + ) + # Default to single-turn when no subcommand given, but let top-level # `-h`/`--help` flow through the main parser so users see both # subcommands and the module description. - known_strategies = {"single-turn", "sliding-window"} + known_strategies = { + "single-turn", + "sliding-window", + "video-single-turn", + } argv = sys.argv[1:] help_requested = bool(argv) and argv[0] in {"-h", "--help"} if not help_requested and (not argv or argv[0] not in known_strategies): diff --git a/benchmarks/multimodal/jsonl/generate_videos.py b/benchmarks/multimodal/jsonl/generate_videos.py new file mode 100644 index 000000000000..7f49a7a91910 --- /dev/null +++ b/benchmarks/multimodal/jsonl/generate_videos.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for generating and sampling video pools.""" + +import hashlib +import random +from pathlib import Path + +import numpy as np + + +def _synthetic_video_key( + seed: int, + video_idx: int, + width: int, + height: int, + fps: int, + seconds: int, +) -> str: + return f"seed{seed}_idx{video_idx:04d}_{width}x{height}_{fps}fps_{seconds}s" + + +def _derive_synthetic_seed(key: str) -> int: + digest = hashlib.sha256(key.encode()).digest() + return int.from_bytes(digest[:8], byteorder="big", signed=False) + + +def _write_synthetic_video( + path: Path, + seed: int, + width: int, + height: int, + fps: int, + seconds: int, +) -> None: + try: + import imageio.v2 as imageio + except ImportError as exc: + raise RuntimeError( + "Synthetic video generation requires imageio. Install imageio and " + "imageio-ffmpeg." + ) from exc + + frame_count = fps * seconds + rng = np.random.default_rng(seed) + base = rng.integers(0, 256, size=3, dtype=np.uint8) + accent = rng.integers(0, 256, size=3, dtype=np.uint8) + rect_w = max(4, width // 4) + rect_h = max(4, height // 4) + span_x = max(1, width - rect_w + 1) + span_y = max(1, height - rect_h + 1) + offset_x = int(rng.integers(0, span_x)) + offset_y = int(rng.integers(0, span_y)) + speed_x = int(rng.integers(1, max(2, width // 8))) + speed_y = int(rng.integers(1, max(2, height // 8))) + xx = np.arange(width, dtype=np.uint16)[None, :] + yy = np.arange(height, dtype=np.uint16)[:, None] + + with imageio.get_writer( + str(path), + fps=fps, + codec="libx264", + macro_block_size=None, + ffmpeg_params=[ + "-pix_fmt", + "yuv420p", + "-map_metadata", + "-1", + "-threads", + "1", + ], + ) as writer: + for frame_idx in range(frame_count): + frame = np.empty((height, width, 3), dtype=np.uint8) + frame[:, :, 0] = ((xx + int(base[0]) + frame_idx * speed_x) % 256).astype( + np.uint8 + ) + frame[:, :, 1] = ((yy + int(base[1]) + frame_idx * speed_y) % 256).astype( + np.uint8 + ) + frame[:, :, 2] = ( + (xx // 2 + yy // 2 + int(base[2]) + frame_idx * 7) % 256 + ).astype(np.uint8) + + x = (offset_x + frame_idx * speed_x) % span_x + y = (offset_y + frame_idx * speed_y) % span_y + frame[y : y + rect_h, x : x + rect_w, :] = accent + writer.append_data(frame) + + +def generate_synthetic_video_pool( + pool_size: int, + video_dir: Path, + video_size: tuple[int, int], + fps: int, + seconds: int, + seed: int, +) -> list[str]: + """Generate pool_size deterministic local MP4 files and return their paths.""" + width, height = video_size + if width <= 0 or height <= 0: + raise ValueError(f"synthetic video size must be positive, got {video_size}") + if width % 2 or height % 2: + raise ValueError( + f"synthetic video dimensions must be even for yuv420p, got {video_size}" + ) + if fps <= 0: + raise ValueError(f"synthetic video fps must be positive, got {fps}") + if seconds <= 0: + raise ValueError(f"synthetic video seconds must be positive, got {seconds}") + + video_dir.mkdir(parents=True, exist_ok=True) + pool: list[str] = [] + for idx in range(pool_size): + key = _synthetic_video_key(seed, idx, width, height, fps, seconds) + path = video_dir / f"synthetic_{key}.mp4" + if not path.exists() or path.stat().st_size == 0: + video_seed = _derive_synthetic_seed(key) + _write_synthetic_video(path, video_seed, width, height, fps, seconds) + pool.append(str(path.resolve())) + + print( + f" {pool_size} synthetic {width}x{height} videos " + f"({fps} fps, {seconds}s) saved to {video_dir}" + ) + return pool + + +def sample_video_slots( + py_rng: random.Random, + pool: list[str], + num_requests: int, + videos_per_request: int, +) -> list[str]: + """Sample video slots from a fixed pool, no duplicates within each request. + + Every video in the pool is guaranteed to appear at least once. + """ + pool_size = len(pool) + total_slots = num_requests * videos_per_request + if pool_size < videos_per_request: + raise ValueError( + f"videos-pool ({pool_size}) must be >= " + f"videos-per-request ({videos_per_request})" + ) + if total_slots < pool_size: + raise ValueError( + f"total slots ({num_requests}x{videos_per_request}={total_slots}) < " + f"videos-pool ({pool_size}). Increase --num-requests or " + f"--videos-per-request, or reduce --videos-pool." + ) + + # Round-robin every pool video into requests so each appears at least once + shuffled = list(pool) + py_rng.shuffle(shuffled) + requests: list[list[str]] = [[] for _ in range(num_requests)] + for i, video in enumerate(shuffled): + requests[i % num_requests].append(video) + + # Fill remaining slots with random pool samples (no intra-request duplicates) + for req in requests: + remaining = videos_per_request - len(req) + if remaining > 0: + used = set(req) + available = [video for video in pool if video not in used] + req.extend(py_rng.sample(available, remaining)) + py_rng.shuffle(req) + + slot_refs = [video for req in requests for video in req] + num_unique = len(set(slot_refs)) + print( + f"Generated {total_slots} video slots from pool of {pool_size}: " + f"{num_unique} unique in use, " + f"{total_slots - num_unique} duplicate references " + f"({(total_slots - num_unique) / total_slots:.1%} reuse)" + ) + return slot_refs diff --git a/benchmarks/multimodal/jsonl/main.py b/benchmarks/multimodal/jsonl/main.py index 4751771cf876..f4d8c1fa8c02 100644 --- a/benchmarks/multimodal/jsonl/main.py +++ b/benchmarks/multimodal/jsonl/main.py @@ -6,11 +6,13 @@ Strategies: single-turn Independent requests with random image sampling (default) sliding-window Causal sessions with sliding-window image overlap + video-single-turn Independent requests with random video sampling Usage: python main.py -n 200 --images-pool 100 python main.py single-turn --image-mode http python main.py sliding-window --num-users 10 --turns-per-user 20 --window-size 5 + python main.py video-single-turn -n 200 --videos-pool 40 """ import argparse @@ -28,6 +30,7 @@ sample_slots, ) from generate_input_text import generate_filler +from generate_videos import generate_synthetic_video_pool, sample_video_slots def _make_pool( @@ -43,6 +46,20 @@ def _make_pool( ) +def _make_video_pool( + args: argparse.Namespace, + pool_size: int, +) -> list[str]: + return generate_synthetic_video_pool( + pool_size=pool_size, + video_dir=args.synthetic_video_dir, + video_size=tuple(args.synthetic_video_size), + fps=args.synthetic_video_fps, + seconds=args.synthetic_video_seconds, + seed=args.seed, + ) + + def run_single_turn( args: argparse.Namespace, np_rng: np.random.Generator, @@ -116,9 +133,47 @@ def run_sliding_window( print(f"Wrote {total_requests} requests ({num_users} sessions) to {output_path}") +def run_video_single_turn( + args: argparse.Namespace, + _np_rng: np.random.Generator, + py_rng: random.Random, +) -> None: + num_requests: int = args.num_requests + videos_per_request: int = args.videos_per_request + video_pool: int = args.videos_pool or (num_requests * videos_per_request) + + total_slots = num_requests * videos_per_request + if video_pool > total_slots: + raise ValueError( + f"total slots ({num_requests}x{videos_per_request}={total_slots}) < " + f"videos-pool ({video_pool}). Increase --num-requests or " + f"--videos-per-request, or reduce --videos-pool." + ) + + pool = _make_video_pool(args, video_pool) + slot_refs = sample_video_slots(py_rng, pool, num_requests, videos_per_request) + + output_filename = ( + f"{num_requests}req_{videos_per_request}vid_{video_pool}pool_" + f"{args.user_text_tokens}word_local.jsonl" + ) + output_path = args.output or (Path(__file__).parent / output_filename) + + with open(output_path, "w") as f: + for i in range(num_requests): + user_text = generate_filler(py_rng, args.user_text_tokens) + start = i * videos_per_request + videos = slot_refs[start : start + videos_per_request] + row: dict = {"text": user_text, "videos": videos} + f.write(json.dumps(row, separators=(",", ":")) + "\n") + + print(f"Wrote {num_requests} video requests to {output_path}") + + STRATEGIES = { "single-turn": run_single_turn, "sliding-window": run_sliding_window, + "video-single-turn": run_video_single_turn, } @@ -129,6 +184,7 @@ def main() -> None: args.seed if args.seed is not None else int(time.time() * 1000) % (2**32) ) print(f"Using seed: {seed}") + args.seed = seed np_rng = np.random.default_rng(seed) py_rng = random.Random(seed) diff --git a/benchmarks/multimodal/sweep/README.md b/benchmarks/multimodal/sweep/README.md index 3e0e4691b96a..b16551626d18 100644 --- a/benchmarks/multimodal/sweep/README.md +++ b/benchmarks/multimodal/sweep/README.md @@ -123,3 +123,4 @@ Given the config above with two input files and two configs (`cache-off`, | Embedding cache (vLLM serve) | `experiments/embedding_cache/vllm_serve.yaml` | Single-node vLLM | | Embedding cache (vLLM E+PD) | `experiments/embedding_cache/vllm_e_pd.yaml` | Disaggregated vLLM E+PD | | Embedding cache (TRT-LLM E+PD) | `experiments/embedding_cache/trtllm_e_pd.yaml` | Disaggregated TRT-LLM E+PD | +| Embedding cache (SGLang E+PD) | `experiments/embedding_cache/sglang_e_pd.yaml` | Disaggregated SGLang video understanding | diff --git a/benchmarks/multimodal/sweep/experiments/embedding_cache/sglang_e_pd.yaml b/benchmarks/multimodal/sweep/experiments/embedding_cache/sglang_e_pd.yaml new file mode 100644 index 000000000000..46faad98c8db --- /dev/null +++ b/benchmarks/multimodal/sweep/experiments/embedding_cache/sglang_e_pd.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# SGLang EPD embedding-cache comparison. +# +# Usage: +# python -m benchmarks.multimodal.sweep \ +# --config benchmarks/multimodal/sweep/experiments/embedding_cache/sglang_e_pd.yaml + +model: Qwen/Qwen3-VL-2B-Instruct +concurrencies: [1, 2, 4, 8] +osl: 128 +warmup_count: 2 +port: 8000 +timeout: 1200 +output_dir: benchmarks/multimodal/sweep/results/sglang_e_pd + +input_files: + - benchmarks/multimodal/jsonl/100req_1vid_20pool_300word_local.jsonl + - benchmarks/multimodal/jsonl/100req_1vid_80pool_300word_local.jsonl + - benchmarks/multimodal/jsonl/100req_2vid_160pool_300word_local.jsonl + +configs: + - label: cache-off + workflow: examples/backends/sglang/launch/multimodal_epd.sh + extra_args: + - --chat-template + - qwen2-vl + - --multimodal-embedding-cache-capacity-gb + - "0" + + - label: cache-on + workflow: examples/backends/sglang/launch/multimodal_epd.sh + extra_args: + - --chat-template + - qwen2-vl + - --multimodal-embedding-cache-capacity-gb + - "10" diff --git a/tests/benchmarks/multimodal/jsonl/test_main.py b/tests/benchmarks/multimodal/jsonl/test_main.py index f208a8f7c003..8c55afd312f3 100644 --- a/tests/benchmarks/multimodal/jsonl/test_main.py +++ b/tests/benchmarks/multimodal/jsonl/test_main.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import json import sys from pathlib import Path @@ -16,6 +17,7 @@ JSONL_DIR = Path(__file__).resolve().parents[4] / "benchmarks" / "multimodal" / "jsonl" sys.path.insert(0, str(JSONL_DIR)) +from generate_videos import generate_synthetic_video_pool # noqa: E402 from main import main # noqa: E402 pytestmark = [pytest.mark.unit, pytest.mark.pre_merge, pytest.mark.gpu_0] @@ -31,6 +33,23 @@ def _run_main(tmp_path: Path, argv: list[str]) -> list[dict]: return [json.loads(line) for line in f if line.strip()] +def _write_test_video( + path: Path, + seed: int, + width: int, + height: int, + fps: int, + seconds: int, +) -> None: + # Keep JSONL benchmark unit tests independent of CI ffmpeg codec availability. + path.write_bytes( + ( + f"synthetic-video\nseed={seed}\n" + f"size={width}x{height}\nfps={fps}\nseconds={seconds}\n" + ).encode() + ) + + class TestSingleTurnDefault: """single-turn is the default when no subcommand is given.""" @@ -134,3 +153,114 @@ def test_image_overlap(self, tmp_path: Path) -> None: assert ( prev[1:] == curr[:-1] ), f"Turn {i} and {i + 1} should share 3/4 images" + + +class TestVideoSingleTurn: + """video-single-turn mirrors image pool reuse across requests.""" + + def test_video_pool_reuse_with_multiple_videos_per_request( + self, tmp_path: Path + ) -> None: + with patch("generate_videos._write_synthetic_video", _write_test_video): + lines = _run_main( + tmp_path, + [ + "video-single-turn", + "-n", + "6", + "--videos-per-request", + "2", + "--videos-pool", + "3", + "--seed", + "11", + "--synthetic-video-dir", + str(tmp_path / "clips"), + "--synthetic-video-size", + "32", + "32", + "--synthetic-video-fps", + "2", + "--synthetic-video-seconds", + "1", + "-o", + str(tmp_path / "videos.jsonl"), + ], + ) + + assert len(lines) == 6 + all_videos = [] + for line in lines: + assert "text" in line + assert "session_id" not in line + assert len(line["videos"]) == 2 + assert len(set(line["videos"])) == 2 + all_videos.extend(line["videos"]) + + assert len(all_videos) == 12 + assert len(set(all_videos)) == 3 + + +class TestSyntheticVideo: + """Synthetic video mode generates reusable local MP4 inputs.""" + + def test_synthetic_video_generation_is_reproducible_in_place( + self, tmp_path: Path + ) -> None: + video_dir = tmp_path / "clips" + kwargs = dict( + pool_size=1, + video_dir=video_dir, + video_size=(32, 32), + fps=2, + seconds=1, + seed=123, + ) + + with patch("generate_videos._write_synthetic_video", _write_test_video): + pool = generate_synthetic_video_pool(**kwargs) + first_digest = hashlib.sha256(Path(pool[0]).read_bytes()).hexdigest() + Path(pool[0]).unlink() + + pool = generate_synthetic_video_pool(**kwargs) + second_digest = hashlib.sha256(Path(pool[0]).read_bytes()).hexdigest() + + assert first_digest == second_digest + + def test_video_single_turn_generates_local_synthetic_pool( + self, tmp_path: Path + ) -> None: + video_dir = tmp_path / "clips" + with patch("generate_videos._write_synthetic_video", _write_test_video): + lines = _run_main( + tmp_path, + [ + "video-single-turn", + "-n", + "4", + "--videos-per-request", + "1", + "--videos-pool", + "2", + "--synthetic-video-dir", + str(video_dir), + "--synthetic-video-size", + "32", + "32", + "--synthetic-video-fps", + "2", + "--synthetic-video-seconds", + "1", + "--seed", + "17", + "-o", + str(tmp_path / "synthetic.jsonl"), + ], + ) + + refs = [row["videos"][0] for row in lines] + assert len(lines) == 4 + assert len(set(refs)) == 2 + for ref in refs: + assert Path(ref).is_file() + assert Path(ref).suffix == ".mp4" From ec45adf2a3f7b719fcba5d5c7c2cf32747b4e030 Mon Sep 17 00:00:00 2001 From: Yongming Ding Date: Tue, 30 Jun 2026 18:09:28 -0700 Subject: [PATCH 019/320] feat(mocker): add explicit max_model_len support to vLLM mocker (#11069) Signed-off-by: Yongming Ding --- components/src/dynamo/mocker/args.py | 17 ++ components/src/dynamo/mocker/config.py | 5 +- .../dynamo/mocker/tests/unit/test_config.py | 68 ++++++++ components/src/dynamo/replay/main.py | 2 +- lib/bindings/python/rust/llm/replay.rs | 9 +- lib/bindings/python/src/dynamo/_core.pyi | 4 + lib/llm/src/mocker.rs | 71 ++++++++- lib/mocker/src/common/protocols.rs | 40 ++++- lib/mocker/src/common/sequence.rs | 8 +- lib/mocker/src/replay/collector.rs | 38 +++-- lib/mocker/src/replay/offline/agg.rs | 49 +++++- lib/mocker/src/replay/offline/disagg.rs | 18 ++- lib/mocker/src/replay/offline/disagg_tests.rs | 21 +++ lib/mocker/src/replay/offline/single.rs | 43 ++++- lib/mocker/src/scheduler/vllm/core.rs | 94 +++++++---- lib/mocker/src/scheduler/vllm/policy.rs | 38 ++++- lib/mocker/src/scheduler/vllm/policy/tests.rs | 148 +++++++++++++++++- lib/mocker/src/scheduler/vllm/tests.rs | 46 ++++++ 18 files changed, 648 insertions(+), 71 deletions(-) diff --git a/components/src/dynamo/mocker/args.py b/components/src/dynamo/mocker/args.py index 3c2008eda2aa..7e3159ec5242 100644 --- a/components/src/dynamo/mocker/args.py +++ b/components/src/dynamo/mocker/args.py @@ -18,6 +18,16 @@ logger = logging.getLogger(__name__) +def positive_int(value: str) -> int: + try: + parsed = int(value) + except ValueError as error: + raise argparse.ArgumentTypeError(str(error)) from error + if parsed <= 0: + raise argparse.ArgumentTypeError(f"must be positive, got {parsed}") + return parsed + + def non_negative_int(value: str) -> int: try: parsed = int(value) @@ -211,6 +221,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="Token block size for KV cache blocks. When unset, the default " "depends on engine: vLLM 64, SGLang 1, TRTLLM 32.", ) + parser.add_argument( + "--max-model-len", + type=positive_int, + default=None, + help="Maximum vLLM sequence length, including prompt and generated tokens. " + "When omitted, no model-length limit is enforced.", + ) parser.add_argument( "--max-num-seqs", type=int, diff --git a/components/src/dynamo/mocker/config.py b/components/src/dynamo/mocker/config.py index 7f639fcae82f..5e69a990f1a6 100644 --- a/components/src/dynamo/mocker/config.py +++ b/components/src/dynamo/mocker/config.py @@ -243,6 +243,7 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: aic_moe_ep_size = getattr(args, "aic_moe_ep_size", None) aic_attention_dp_size = getattr(args, "aic_attention_dp_size", None) engine_type = getattr(args, "engine_type", None) or "vllm" + max_model_len = getattr(args, "max_model_len", None) num_gpu_blocks = _resolve_num_gpu_blocks( explicit_num_gpu_blocks=getattr(args, "num_gpu_blocks", None), engine_type=engine_type, @@ -267,6 +268,7 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: engine_type=engine_type, num_gpu_blocks=num_gpu_blocks, block_size=getattr(args, "block_size", 0) or 0, + max_model_len=max_model_len, max_num_seqs=getattr(args, "max_num_seqs", _DEFAULT_MAX_NUM_SEQS), max_num_batched_tokens=getattr( args, "max_num_batched_tokens", _DEFAULT_MAX_NUM_BATCHED_TOKENS @@ -349,8 +351,7 @@ def build_runtime_config( engine_args: MockEngineArgs, ) -> tuple[int, ModelRuntimeConfig]: rc = ModelRuntimeConfig() - # Mocker does not enforce a model context limit. - rc.context_length = 0 + rc.context_length = engine_args.max_model_len or 0 rc.total_kv_blocks = engine_args.num_gpu_blocks rc.max_num_seqs = engine_args.max_num_seqs if rc.max_num_seqs is None: diff --git a/components/src/dynamo/mocker/tests/unit/test_config.py b/components/src/dynamo/mocker/tests/unit/test_config.py index 4f690775b1d2..c3406c958dd3 100644 --- a/components/src/dynamo/mocker/tests/unit/test_config.py +++ b/components/src/dynamo/mocker/tests/unit/test_config.py @@ -35,6 +35,7 @@ def make_args(**overrides): "engine_type": "vllm", "num_gpu_blocks": None, "block_size": None, + "max_model_len": None, "max_num_seqs": 256, "max_num_batched_tokens": 8192, "enable_prefix_caching": True, @@ -397,6 +398,73 @@ def test_mocker_cli_accepts_mtp_configuration(): assert args.aic_mtp_seed == 99 +def test_mocker_cli_accepts_max_model_len(): + args = parse_args(["--max-model-len", "32768"]) + + engine_args = CONFIG.build_mocker_engine_args(args) + _, runtime_config = CONFIG.build_runtime_config(engine_args) + + assert engine_args.max_model_len == 32768 + assert runtime_config.context_length == 32768 + + +@pytest.mark.parametrize("value", ["0", "-1"]) +def test_mocker_cli_rejects_non_positive_max_model_len(value): + with pytest.raises(SystemExit): + parse_args(["--max-model-len", value]) + + +def test_build_mocker_engine_args_keeps_max_model_len_explicit_only(): + engine_args = CONFIG.build_mocker_engine_args( + make_args(model_path="/models/mock", num_gpu_blocks=4096) + ) + + assert engine_args.max_model_len is None + + +def test_build_mocker_engine_args_preserves_explicit_max_model_len(): + engine_args = CONFIG.build_mocker_engine_args( + make_args( + model_path="/models/mock", + max_model_len=32768, + num_gpu_blocks=4096, + ) + ) + + assert engine_args.max_model_len == 32768 + + +def test_replay_engine_args_keeps_max_model_len_explicit_only(): + import dynamo.replay.main as replay_main + + engine_args = replay_main._load_engine_args( + json.dumps( + { + "num_gpu_blocks": 4096, + "aic_model_path": "/models/mock", + } + ) + ) + + assert engine_args.max_model_len is None + + +def test_replay_engine_args_preserves_explicit_max_model_len(): + import dynamo.replay.main as replay_main + + engine_args = replay_main._load_engine_args( + json.dumps( + { + "num_gpu_blocks": 4096, + "max_model_len": 32768, + "aic_model_path": "/models/mock", + } + ) + ) + + assert engine_args.max_model_len == 32768 + + def test_replay_engine_args_compute_kv_bytes_for_g3_before_validation(monkeypatch): import dynamo.replay.main as replay_main diff --git a/components/src/dynamo/replay/main.py b/components/src/dynamo/replay/main.py index 0053b30c7e5d..fb7d2d883173 100644 --- a/components/src/dynamo/replay/main.py +++ b/components/src/dynamo/replay/main.py @@ -295,7 +295,7 @@ def _engine_caps(args: MockEngineArgs) -> EngineCapabilities: num_gpu=1, max_num_batched_tokens=args.max_num_batched_tokens, max_num_seqs=args.max_num_seqs, - context_length=max_kv_tokens if max_kv_tokens > 0 else None, + context_length=args.max_model_len, max_kv_tokens=max_kv_tokens if max_kv_tokens > 0 else None, speculative_nextn=args.aic_nextn, ) diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index 7d4a62e84d58..cfdb300768d3 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -171,7 +171,7 @@ impl MockEngineArgs { #[pymethods] impl MockEngineArgs { #[new] - #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=None, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, aic_moe_tp_size=None, aic_moe_ep_size=None, aic_attention_dp_size=None, aic_nextn=None, aic_nextn_accept_rates=None, aic_mtp_seed=42, aic_gemm_dtype=None, aic_moe_dtype=None, aic_fmha_dtype=None, aic_kv_cache_dtype=None, aic_comm_dtype=None, gpu_memory_utilization=None, mem_fraction_static=None, free_gpu_memory_fraction=None, enable_local_indexer=false, bootstrap_port=None, handoff_session_timeout_ms=300000, kv_bytes_per_token=None, kv_transfer_bandwidth=None, kv_transfer_timing_mode="full_prompt", reasoning=None, response_replay_trace_path=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None, trtllm=None, num_g2_blocks=None, num_g3_blocks=None, offload_batch_size=None, bandwidth_g1_to_g2_gbps=None, bandwidth_g2_to_g1_gbps=None, bandwidth_g2_to_g3_gbps=None, bandwidth_g3_to_g2_gbps=None, enable_g4_storage=false, bandwidth_g2_to_g4_gbps=None, bandwidth_g4_to_g2_gbps=None))] + #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=None, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, aic_moe_tp_size=None, aic_moe_ep_size=None, aic_attention_dp_size=None, aic_nextn=None, aic_nextn_accept_rates=None, aic_mtp_seed=42, aic_gemm_dtype=None, aic_moe_dtype=None, aic_fmha_dtype=None, aic_kv_cache_dtype=None, aic_comm_dtype=None, gpu_memory_utilization=None, mem_fraction_static=None, free_gpu_memory_fraction=None, enable_local_indexer=false, bootstrap_port=None, handoff_session_timeout_ms=300000, kv_bytes_per_token=None, kv_transfer_bandwidth=None, kv_transfer_timing_mode="full_prompt", reasoning=None, response_replay_trace_path=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None, trtllm=None, num_g2_blocks=None, num_g3_blocks=None, offload_batch_size=None, bandwidth_g1_to_g2_gbps=None, bandwidth_g2_to_g1_gbps=None, bandwidth_g2_to_g3_gbps=None, bandwidth_g3_to_g2_gbps=None, enable_g4_storage=false, bandwidth_g2_to_g4_gbps=None, bandwidth_g4_to_g2_gbps=None, max_model_len=None))] #[allow(clippy::too_many_arguments)] fn new( engine_type: &str, @@ -230,6 +230,7 @@ impl MockEngineArgs { enable_g4_storage: bool, bandwidth_g2_to_g4_gbps: Option, bandwidth_g4_to_g2_gbps: Option, + max_model_len: Option, ) -> PyResult { let engine_type = parse_mocker_engine_type(engine_type)?; let worker_type = parse_worker_type(worker_type)?; @@ -248,6 +249,7 @@ impl MockEngineArgs { let mut builder = RsMockEngineArgs::builder() .engine_type(engine_type) .block_size(block_size) + .max_model_len(max_model_len) .max_num_seqs(max_num_seqs) .max_num_batched_tokens(max_num_batched_tokens) .enable_prefix_caching(enable_prefix_caching) @@ -365,6 +367,11 @@ impl MockEngineArgs { self.inner.num_gpu_blocks } + #[getter] + fn max_model_len(&self) -> Option { + self.inner.max_model_len + } + #[getter] fn max_num_seqs(&self) -> Option { self.inner.max_num_seqs diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 3248e7d7a714..9aadc3ae2958 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1965,6 +1965,7 @@ class MockEngineArgs: enable_g4_storage: bool = False, bandwidth_g2_to_g4_gbps: Optional[float] = None, bandwidth_g4_to_g2_gbps: Optional[float] = None, + max_model_len: Optional[int] = None, ) -> None: ... @@ -1983,6 +1984,9 @@ class MockEngineArgs: @num_gpu_blocks.setter def num_gpu_blocks(self, value: int) -> None: ... + @property + def max_model_len(self) -> Optional[int]: ... + @property def max_num_seqs(self) -> Optional[int]: ... diff --git a/lib/llm/src/mocker.rs b/lib/llm/src/mocker.rs index 1a557ed3dd21..9727f83656ba 100644 --- a/lib/llm/src/mocker.rs +++ b/lib/llm/src/mocker.rs @@ -791,6 +791,12 @@ impl AsyncEngine, ManyOut, Error> let max_output_tokens = planned_output_token_ids .as_ref() .map_or(requested_max_output_tokens, Vec::len); + let effective_max_output_tokens = + self.engine_args + .max_model_len + .map_or(max_output_tokens, |max_model_len| { + max_output_tokens.min(max_model_len.saturating_sub(request.token_ids.len())) + }); let native_timing = self .native_metrics .request_timing(&request.model, dp_rank, is_prefill, request_start) @@ -1013,14 +1019,14 @@ impl AsyncEngine, ManyOut, Error> break; }; - // A terminally rejected request never ran (its footprint - // exceeds the KV pool): emit no token and do not complete the - // bootstrap room — surface the rejection and end the stream - // before any token/prefill bookkeeping. + // A terminally rejected request never ran because it violated + // a worker admission limit. Emit no token and do not complete + // the bootstrap room; surface the rejection before any + // token/prefill bookkeeping. if signal.rejected { handoff_cancel.cancel(); let _ = stream_tx.send(LLMEngineOutput::error( - "request rejected: KV footprint exceeds pool capacity".to_string(), + "request rejected: request exceeds worker admission limits".to_string(), )); break; } @@ -1043,7 +1049,7 @@ impl AsyncEngine, ManyOut, Error> ..Default::default() }; - if signal.completed && token_count < max_output_tokens { + if signal.completed && token_count < effective_max_output_tokens { let _ = stream_tx.send(LLMEngineOutput::error("Completion signal received before max tokens reached".to_string())); break; } @@ -1239,6 +1245,22 @@ mod tests { .unwrap() } + fn decode_request(prompt_tokens: usize, max_tokens: u32) -> PreprocessedRequest { + PreprocessedRequest::builder() + .model("mock".to_string()) + .token_ids(vec![1; prompt_tokens]) + .stop_conditions(StopConditions { + max_tokens: Some(max_tokens), + ..Default::default() + }) + .sampling_options(SamplingOptions::default()) + .output_options(OutputOptions::default()) + .eos_token_ids(vec![]) + .annotations(vec![]) + .build() + .unwrap() + } + #[tokio::test(start_paused = true)] async fn no_bootstrap_prefill_delays_terminal_finish_once() { let args = MockEngineArgs::builder() @@ -1284,6 +1306,43 @@ mod tests { assert!(stream.next().await.is_none()); } + #[tokio::test] + async fn context_capped_completion_maps_to_length() { + let args = MockEngineArgs::builder() + .max_model_len(Some(4)) + .build() + .unwrap(); + let engine = MockEngine::new(args); + let (request_tx, mut request_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.request_senders.set(vec![request_tx]).unwrap(); + + let mut stream = engine + .generate(SingleIn::new(decode_request(3, 4))) + .await + .unwrap(); + let request = request_rx.recv().await.unwrap(); + assert_eq!(request.max_output_tokens, 4); + let request_id = request.uuid.unwrap(); + engine + .active_requests + .get(&request_id) + .unwrap() + .send(OutputSignal { + uuid: request_id, + token_id: Some(42), + completed: true, + rejected: false, + handoff_delay_ms: None, + }) + .unwrap(); + + let token = stream.next().await.unwrap(); + assert_eq!(token.token_ids.len(), 1); + assert!(token.finish_reason.is_none()); + assert_eq!(stream.next().await.unwrap(), LLMEngineOutput::length()); + assert!(stream.next().await.is_none()); + } + #[test] fn unbounded_sequence_limit_uses_finite_multi_handoff_capacity() { let args = MockEngineArgs::builder() diff --git a/lib/mocker/src/common/protocols.rs b/lib/mocker/src/common/protocols.rs index 3af8300294dc..ff82a5b24706 100644 --- a/lib/mocker/src/common/protocols.rs +++ b/lib/mocker/src/common/protocols.rs @@ -519,6 +519,7 @@ struct MockEngineArgsSerde { engine_type: OptionalConfigValue, num_gpu_blocks: OptionalConfigValue, block_size: OptionalConfigValue, + max_model_len: OptionalConfigValue, max_num_seqs: OptionalConfigValue, max_num_batched_tokens: OptionalConfigValue, enable_prefix_caching: OptionalConfigValue, @@ -612,6 +613,12 @@ pub struct MockEngineArgs { #[builder(default = "0")] pub block_size: usize, + /// Optional vLLM sequence-length limit, including prompt and generated + /// tokens. Requests with no room to generate are rejected before admission. + #[builder(default = "None")] + #[validate(range(min = 1))] + pub max_model_len: Option, + // This was 1024 in the past but reverted back to 256 #[builder(default = Some(256))] #[validate(range(min = 1))] @@ -927,6 +934,16 @@ fn validate_mock_engine_args(args: &MockEngineArgs) -> Result<(), ValidationErro "num_g3_blocks requires num_g2_blocks because mocker stages G3 through G2".to_string(), )); } + + if args.max_model_len.is_some() && args.engine_type != EngineType::Vllm { + return Err(mock_engine_args_validation_error( + "max_model_len_requires_vllm", + format!( + "max_model_len is supported only for engine_type=vllm, got engine_type={:?}", + args.engine_type + ), + )); + } if args.enable_g4_storage && args.num_g2_blocks.is_none() { return Err(mock_engine_args_validation_error( "g4_requires_g2", @@ -1015,6 +1032,9 @@ impl TryFrom for MockEngineArgs { if let Some(block_size) = compat.block_size.into_non_null("block_size")? { builder = builder.block_size(block_size); } + if let Some(max_model_len) = compat.max_model_len.into_nullable() { + builder = builder.max_model_len(max_model_len); + } if let Some(max_num_seqs) = compat.max_num_seqs.into_nullable() { builder = builder.max_num_seqs(max_num_seqs); } @@ -1428,6 +1448,7 @@ mod tests { fn test_mock_engine_args_json_round_trip_preserves_worker_type_and_nulls() { let args = MockEngineArgs::builder() .worker_type(WorkerType::Decode) + .max_model_len(Some(32768)) .max_num_seqs(None) .max_num_batched_tokens(None) .reasoning(None) @@ -1437,7 +1458,7 @@ mod tests { .normalized() .unwrap(); - let payload = serde_json::json!({ + let mut payload = serde_json::json!({ "engine_type": "vllm", "num_gpu_blocks": args.num_gpu_blocks, "block_size": args.block_size, @@ -1480,10 +1501,12 @@ mod tests { "sglang": args.sglang, "has_perf_model": true, }); + payload["max_model_len"] = serde_json::json!(args.max_model_len); let restored = MockEngineArgs::from_json_str(&payload.to_string()).unwrap(); assert_eq!(restored.worker_type, WorkerType::Decode); + assert_eq!(restored.max_model_len, Some(32768)); assert_eq!(restored.max_num_seqs, None); assert_eq!(restored.max_num_batched_tokens, None); assert_eq!( @@ -1690,6 +1713,21 @@ mod tests { .expect("in-range aic_nextn should validate"); } + #[test] + fn test_normalized_rejects_zero_max_model_len() { + let error = MockEngineArgs::builder() + .max_model_len(Some(0)) + .build() + .unwrap() + .normalized() + .unwrap_err(); + + assert!( + error.to_string().contains("max_model_len"), + "unexpected error: {error}", + ); + } + #[test] fn test_mtp_defaults_and_json_round_trip() { let args = MockEngineArgs::builder() diff --git a/lib/mocker/src/common/sequence.rs b/lib/mocker/src/common/sequence.rs index 104470cb10c8..aac2b4c5a501 100644 --- a/lib/mocker/src/common/sequence.rs +++ b/lib/mocker/src/common/sequence.rs @@ -372,10 +372,16 @@ impl ActiveSequence { } // Free all blocks when we reach max tokens - signals.extend(self.free_signal_for_tokens(self.len())); + signals.extend(self.terminal_signals()); (token, signals) } + /// Release the full sequence footprint after an independent terminal + /// condition, such as the model context-length limit, is reached. + pub(crate) fn terminal_signals(&self) -> Vec { + self.free_signal_for_tokens(self.len()) + } + fn free_signal_for_tokens(&self, active_tokens: usize) -> Vec { let active_blocks = active_tokens .div_ceil(self.block_size) diff --git a/lib/mocker/src/replay/collector.rs b/lib/mocker/src/replay/collector.rs index eb3c1dadd598..d65f81b55f84 100644 --- a/lib/mocker/src/replay/collector.rs +++ b/lib/mocker/src/replay/collector.rs @@ -325,7 +325,7 @@ struct TraceRequestStats { first_admit_ms: Option, token_times_ms: Vec, input_length: usize, - output_length: usize, + requested_output_length: usize, reused_input_tokens: usize, first_admission_reused_input_tokens: usize, /// Index of the prefill worker that handled this request, if any. @@ -397,6 +397,9 @@ pub struct PerRequestRecord { /// AIPerf's `inter_token_latency` field — one scalar per request. pub itl_ms: Option, pub input_length: usize, + /// Number of output tokens requested by the workload trace. + pub requested_output_length: usize, + /// Number of output tokens actually emitted by the mock engine. pub output_length: usize, pub reused_input_tokens: usize, pub prefill_worker_idx: Option, @@ -421,6 +424,7 @@ pub(crate) struct TraceRequestStatsSnapshot { pub first_token_ms: Option, pub last_token_ms: Option, pub input_length: usize, + pub requested_output_length: usize, pub output_length: usize, pub reused_input_tokens: usize, pub first_admission_reused_input_tokens: usize, @@ -512,6 +516,10 @@ impl TraceRequestStats { self.token_times_ms.last().copied() } + fn actual_output_length(&self) -> usize { + self.token_times_ms.len() + } + fn mean_tpot_ms(&self) -> Option { let num_gaps = self.token_times_ms.len().saturating_sub(1); if num_gaps == 0 { @@ -580,16 +588,16 @@ impl TraceCollector { uuid: Uuid, arrival_time_ms: f64, input_length: usize, - output_length: usize, + requested_output_length: usize, ) { self.requests.insert( uuid, TraceRequestStats { arrival_time_ms, first_admit_ms: None, - token_times_ms: Vec::with_capacity(output_length), + token_times_ms: Vec::with_capacity(requested_output_length), input_length, - output_length, + requested_output_length, reused_input_tokens: 0, prefill_worker_idx: None, decode_worker_idx: None, @@ -757,6 +765,12 @@ impl TraceCollector { Some((ttft_ms, mean_itl_ms)) } + pub(crate) fn actual_output_length(&self, uuid: Uuid) -> Option { + self.requests + .get(&uuid) + .map(TraceRequestStats::actual_output_length) + } + pub(crate) fn finish(self) -> TraceSimulationReport { // Build per-request records before we move `self.requests` into the // summary aggregation below. Gated on `capture_per_request` — the @@ -805,7 +819,8 @@ impl TraceCollector { completed_requests += 1; total_input_tokens += stats.input_length; - total_output_tokens += stats.output_length; + let output_length = stats.actual_output_length(); + total_output_tokens += output_length; total_reused_tokens += stats.reused_input_tokens; total_first_admission_reused_tokens += stats.first_admission_reused_input_tokens; duration_ms = duration_ms.max(last_token_ms); @@ -816,9 +831,9 @@ impl TraceCollector { e2e_latencies.push(e2e_ms); // Goodput classification (aiperf avg-ITL; see SlaThresholds::is_good). - if sla.is_set() && sla.is_good(ttft_ms, e2e_ms, stats.output_length) { + if sla.is_set() && sla.is_good(ttft_ms, e2e_ms, output_length) { goodput_requests += 1; - goodput_output_tokens += stats.output_length; + goodput_output_tokens += output_length; } if let Some(ttst_ms) = stats.ttst_ms() { @@ -937,7 +952,8 @@ impl TraceCollector { e2e_latency_ms: last_token_ms.map(|time| (time - stats.arrival_time_ms).max(0.0)), itl_ms: stats.mean_tpot_ms(), input_length: stats.input_length, - output_length: stats.output_length, + requested_output_length: stats.requested_output_length, + output_length: stats.actual_output_length(), reused_input_tokens: detail .prefill_reused_input_tokens .unwrap_or(stats.reused_input_tokens), @@ -976,7 +992,8 @@ impl TraceCollector { first_token_ms: stats.first_token_ms(), last_token_ms: stats.last_token_ms(), input_length: stats.input_length, - output_length: stats.output_length, + requested_output_length: stats.requested_output_length, + output_length: stats.actual_output_length(), reused_input_tokens: stats.reused_input_tokens, first_admission_reused_input_tokens: stats.first_admission_reused_input_tokens, }) @@ -992,7 +1009,8 @@ impl TraceCollector { first_token_ms: stats.first_token_ms(), last_token_ms: stats.last_token_ms(), input_length: stats.input_length, - output_length: stats.output_length, + requested_output_length: stats.requested_output_length, + output_length: stats.actual_output_length(), reused_input_tokens: stats.reused_input_tokens, first_admission_reused_input_tokens: stats.first_admission_reused_input_tokens, }) diff --git a/lib/mocker/src/replay/offline/agg.rs b/lib/mocker/src/replay/offline/agg.rs index 892d68e6a67e..80fdb73cae7d 100644 --- a/lib/mocker/src/replay/offline/agg.rs +++ b/lib/mocker/src/replay/offline/agg.rs @@ -508,9 +508,19 @@ impl AggRuntime { // traffic deltas (they still free their slot and advance below). if !signal.rejected { let latencies = self.collector.request_latencies(signal.uuid); + let actual_output_tokens = self + .collector + .actual_output_length(signal.uuid) + .ok_or_else(|| { + anyhow::anyhow!( + "offline replay missing collector state for {}", + signal.uuid + ) + })?; + debug_assert!(actual_output_tokens <= removed_state.output_tokens); self.traffic.on_request( removed_state.input_tokens, - removed_state.output_tokens, + actual_output_tokens, latencies, ); } @@ -2957,6 +2967,43 @@ mod tests { ); } + #[test] + fn test_drain_traffic_uses_context_capped_output_length() { + let args = MockEngineArgs::builder() + .block_size(4) + .num_gpu_blocks(32) + .max_model_len(Some(8)) + .max_num_batched_tokens(Some(16)) + .max_num_seqs(Some(4)) + .enable_prefix_caching(false) + .speedup_ratio(0.0) + .build() + .unwrap(); + let requests = VecDeque::from([DirectRequest { + tokens: vec![1; 7], + max_output_tokens: 4, + uuid: Some(Uuid::from_u128(1)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.0), + ..Default::default() + }]); + let mut rt = AggRuntime::new( + &args, + None, + None, + requests, + 1, + ReplayMode::Trace, + ReplayRouterMode::RoundRobin, + ) + .unwrap(); + + assert!(rt.advance_to(1000.0).unwrap()); + let stats = rt.drain_traffic(); + assert_eq!(stats.num_req, 1); + assert_eq!(stats.avg_osl, 1.0); + } + #[test] fn test_apply_scaling_without_startup_is_immediate() { let args = fast_router_args(); // no startup_time diff --git a/lib/mocker/src/replay/offline/disagg.rs b/lib/mocker/src/replay/offline/disagg.rs index 0e06e8cdc2ca..a5e15e83825d 100644 --- a/lib/mocker/src/replay/offline/disagg.rs +++ b/lib/mocker/src/replay/offline/disagg.rs @@ -1420,13 +1420,21 @@ impl DisaggRuntime { // latency — keep it out of the planner-facing traffic deltas (mirror the // aggregated path). It still frees its slot, advances, and is marked done. if !signal.rejected { - let state = self.state(signal.uuid)?; - let original = state.original_request()?; - let input_tokens = original.tokens.len(); - let output_tokens = original.max_output_tokens; + let (input_tokens, requested_output_tokens) = { + let state = self.state(signal.uuid)?; + let original = state.original_request()?; + (original.tokens.len(), original.max_output_tokens) + }; + let actual_output_tokens = self + .collector + .actual_output_length(signal.uuid) + .ok_or_else(|| { + anyhow!("offline replay missing collector state for {}", signal.uuid) + })?; + debug_assert!(actual_output_tokens <= requested_output_tokens); let latencies = self.collector.request_latencies(signal.uuid); self.traffic - .on_request(input_tokens, output_tokens, latencies); + .on_request(input_tokens, actual_output_tokens, latencies); } let terminal_status = if signal.rejected { ReplayTerminalStatus::Rejected diff --git a/lib/mocker/src/replay/offline/disagg_tests.rs b/lib/mocker/src/replay/offline/disagg_tests.rs index fcbfdbab56d7..e4ebc989c99f 100644 --- a/lib/mocker/src/replay/offline/disagg_tests.rs +++ b/lib/mocker/src/replay/offline/disagg_tests.rs @@ -1215,6 +1215,27 @@ fn test_advance_to_moves_clock_across_idle_gap() { assert!((stats.duration_s - 0.5).abs() < 1e-9); } +#[test] +fn test_disagg_traffic_uses_context_capped_output_length() { + let mut config = disagg_config(); + config.prefill_args.max_model_len = Some(8); + config.decode_args.max_model_len = Some(8); + let mut runtime = DisaggRuntime::new( + &config, + None, + None, + VecDeque::from([request(1, 7, 4, 0.0)]), + ReplayMode::Trace, + ReplayRouterMode::RoundRobin, + ) + .unwrap(); + + assert!(runtime.advance_to(1000.0).unwrap()); + let stats = runtime.drain_traffic(); + assert_eq!(stats.num_req, 1); + assert_eq!(stats.avg_osl, 1.0); +} + /// Setting `max_sim_time_ms` causes `run()` to break before scheduled /// arrivals past the cap. This test verifies the cap operates on /// **simulated** time (`now_ms`), not real wall-clock time: with diff --git a/lib/mocker/src/replay/offline/single.rs b/lib/mocker/src/replay/offline/single.rs index eefef4f1588d..e16e90c46287 100644 --- a/lib/mocker/src/replay/offline/single.rs +++ b/lib/mocker/src/replay/offline/single.rs @@ -331,7 +331,7 @@ mod tests { use super::*; use crate::common::protocols::EngineType; use crate::loadgen::{SessionTrace, Trace, TurnTrace}; - use crate::replay::{TraceRequestStatsSnapshot, TraceSimulationReport}; + use crate::replay::{ReplayTerminalStatus, TraceRequestStatsSnapshot, TraceSimulationReport}; use rstest::rstest; use std::collections::{HashMap, VecDeque}; use uuid::Uuid; @@ -1007,6 +1007,47 @@ mod tests { assert_eq!(report.request_counts.total_output_tokens, 4); } + #[test] + fn max_model_len_reports_actual_output_and_preserves_requested_output() { + let args = MockEngineArgs::builder() + .block_size(4) + .num_gpu_blocks(4) + .max_model_len(Some(8)) + .max_num_batched_tokens(Some(16)) + .max_num_seqs(Some(4)) + .enable_prefix_caching(false) + .enable_chunked_prefill(true) + .speedup_ratio(0.0) + .build() + .unwrap(); + let uuid = Uuid::from_u128(3); + let report = simulate_concurrency_single( + args, + vec![DirectRequest { + tokens: vec![1; 7], + max_output_tokens: 4, + uuid: Some(uuid), + dp_rank: 0, + arrival_timestamp_ms: Some(0.0), + ..Default::default() + }], + 1, + true, + None, + crate::replay::SlaThresholds::default(), + ) + .unwrap(); + + assert_eq!(report.request_counts.completed_requests, 1); + assert_eq!(report.request_counts.total_output_tokens, 1); + assert_eq!(report.per_request.len(), 1); + let record = &report.per_request[0]; + assert_eq!(record.uuid, uuid.to_string()); + assert_eq!(record.requested_output_length, 4); + assert_eq!(record.output_length, 1); + assert_eq!(record.terminal_status, ReplayTerminalStatus::Completed); + } + fn cap_request(uuid: u128, arrival_ms: f64) -> DirectRequest { DirectRequest { tokens: vec![1; 4], diff --git a/lib/mocker/src/scheduler/vllm/core.rs b/lib/mocker/src/scheduler/vllm/core.rs index 6602e3305339..5732e5e9e883 100644 --- a/lib/mocker/src/scheduler/vllm/core.rs +++ b/lib/mocker/src/scheduler/vllm/core.rs @@ -1387,26 +1387,37 @@ impl VllmCore { .iter() .filter_map(|running_uuid| self.state.requests.get(running_uuid)) .map(|request| &request.sequence); - let prompt_is_prebuilt = request.prompt_is_prebuilt(); - match admission.stage_for(prompt_is_prebuilt) { - AdmissionStage::Materialized => AdmissionDecision::Admit { - prefill_cost: PrefillCost { - new_blocks: 0, - new_tokens: 0, - cached_tokens: request.sequence.num_input_tokens(), - active_cached_tokens: request.sequence.num_input_tokens(), + if policy::should_reject_for_model_len( + scheduling_policy, + &request.sequence, + self.args.max_model_len, + ) { + AdmissionDecision::Reject + } else { + let prompt_is_prebuilt = request.prompt_is_prebuilt(); + match admission.stage_for(prompt_is_prebuilt) { + AdmissionStage::Materialized => AdmissionDecision::Admit { + prefill_cost: PrefillCost { + new_blocks: 0, + new_tokens: 0, + cached_tokens: request.sequence.num_input_tokens(), + active_cached_tokens: request.sequence.num_input_tokens(), + }, }, - }, - AdmissionStage::PendingDestinationHead => break, - AdmissionStage::FreshKv => policy::decide_waiting_admission( - scheduling_policy, - &request.sequence, - request.status == RequestStatus::Waiting, - running_seqs, - self.args.num_gpu_blocks, - self.args.block_size, - &self.kv_manager, - ), + AdmissionStage::PendingDestinationHead => break, + AdmissionStage::FreshKv => { + let is_fresh = request.status == RequestStatus::Waiting; + policy::decide_waiting_admission( + scheduling_policy, + &request.sequence, + is_fresh, + running_seqs, + self.args.num_gpu_blocks, + self.args.block_size, + &self.kv_manager, + ) + } + } } }; let prefill_cost = match decision { @@ -1418,8 +1429,14 @@ impl VllmCore { tracing::warn!( %uuid, ?scheduling_policy, + prompt_tokens = self + .state + .requests + .get(&uuid) + .map(|request| request.sequence.num_input_tokens()), + max_model_len = self.args.max_model_len, num_gpu_blocks = self.args.num_gpu_blocks, - "rejecting request whose admission footprint exceeds the entire KV pool" + "rejecting request that exceeds a worker admission limit" ); rejected_uuids.push(uuid); self.drop_request(uuid); @@ -1893,7 +1910,7 @@ impl VllmCore { continue; }; if request.num_computed_tokens < request.sequence.len() - || request.sequence.generated_tokens() >= request.sequence.max_output_tokens() + || policy::generation_complete(&request.sequence, self.args.max_model_len) { continue; } @@ -1941,8 +1958,11 @@ impl VllmCore { let Some(sequence) = self.state.running_sequence_mut(uuid) else { break; }; - let (token_id, signals) = sequence.generate_token(); - completed = sequence.generated_tokens() >= sequence.max_output_tokens(); + let (token_id, mut signals) = sequence.generate_token(); + completed = policy::generation_complete(sequence, self.args.max_model_len); + if completed && sequence.generated_tokens() < sequence.max_output_tokens() { + signals.extend(sequence.terminal_signals()); + } let effects = if completed { split_terminal_effects(signals) } else { @@ -2072,10 +2092,10 @@ impl VllmCore { .iter() .filter_map(|uuid| self.state.requests.get(uuid)) .map(|request| { - let remaining = request - .sequence - .max_output_tokens() - .saturating_sub(request.sequence.generated_tokens()); + let remaining = policy::remaining_generation_tokens( + &request.sequence, + self.args.max_model_len, + ); let burst = max_burst.min(remaining); let current_blocks = request.sequence.len().div_ceil(self.args.block_size); let target_blocks = @@ -2130,7 +2150,7 @@ impl VllmCore { continue; }; if request.num_computed_tokens == request.sequence.len() - && request.sequence.generated_tokens() < request.sequence.max_output_tokens() + && !policy::generation_complete(&request.sequence, self.args.max_model_len) { ready.push(uuid); } @@ -2179,10 +2199,10 @@ impl VllmCore { .requests .get(uuid) .expect("ready request must remain active"); - let remaining = request - .sequence - .max_output_tokens() - .saturating_sub(request.sequence.generated_tokens()); + let remaining = policy::remaining_generation_tokens( + &request.sequence, + self.args.max_model_len, + ); let burst = if self.args.worker_type == WorkerType::Prefill { remaining.min(1) } else { @@ -2205,9 +2225,15 @@ impl VllmCore { .requests .get_mut(&uuid) .expect("sampled request must remain active"); - let (token_id, signals) = request.sequence.generate_token(); + let (token_id, mut signals) = request.sequence.generate_token(); let is_complete = - request.sequence.generated_tokens() >= request.sequence.max_output_tokens(); + policy::generation_complete(&request.sequence, self.args.max_model_len); + if is_complete + && request.sequence.generated_tokens() + < request.sequence.max_output_tokens() + { + signals.extend(request.sequence.terminal_signals()); + } (token_id, signals, is_complete) }; let effects = if is_complete { diff --git a/lib/mocker/src/scheduler/vllm/policy.rs b/lib/mocker/src/scheduler/vllm/policy.rs index 3992ccef375c..f75df464fba1 100644 --- a/lib/mocker/src/scheduler/vllm/policy.rs +++ b/lib/mocker/src/scheduler/vllm/policy.rs @@ -3,8 +3,7 @@ //! Engine-specific policy for the shared vLLM/TRT-LLM scheduler core. -use crate::common::protocols::PrefillCost; -use crate::common::protocols::SchedulingPolicy; +use crate::common::protocols::{PrefillCost, SchedulingPolicy}; use crate::common::sequence::ActiveSequence; use crate::kv_manager::KvManager; @@ -15,6 +14,34 @@ pub(super) enum AdmissionDecision { Reject, } +pub(super) fn should_reject_for_model_len( + policy: SchedulingPolicy, + sequence: &ActiveSequence, + max_model_len: Option, +) -> bool { + policy == SchedulingPolicy::Vllm + && max_model_len.is_some_and(|limit| sequence.num_input_tokens() >= limit) +} + +/// Number of additional tokens the request may generate before reaching +/// either its requested output length or the model sequence-length limit. +pub(super) fn remaining_generation_tokens( + sequence: &ActiveSequence, + max_model_len: Option, +) -> usize { + let requested_remaining = sequence + .max_output_tokens() + .saturating_sub(sequence.generated_tokens()); + let context_remaining = max_model_len + .map(|limit| limit.saturating_sub(sequence.len())) + .unwrap_or(usize::MAX); + requested_remaining.min(context_remaining) +} + +pub(super) fn generation_complete(sequence: &ActiveSequence, max_model_len: Option) -> bool { + remaining_generation_tokens(sequence, max_model_len) == 0 +} + /// Decide whether the FIFO head can enter the shared scheduler core. /// /// vLLM reserves only the current known sequence. TRT-LLM @@ -32,11 +59,8 @@ pub(super) fn decide_waiting_admission<'a>( if is_fresh { match policy { SchedulingPolicy::Vllm => { - // TODO: Carry vLLM's max_model_len explicitly. Upstream bounds prompt - // plus generated tokens by max_model_len and sizes KV for one - // max-length sequence. Until the mocker models that value, total - // worker KV is only a proxy for the one-time fresh-sequence admission - // cap; it does not cap future output length. + // Total worker KV remains a fallback one-time admission cap + // when max_model_len is unset or larger than the KV pool. if sequence.current_known_blocks() > num_gpu_blocks { return AdmissionDecision::Reject; } diff --git a/lib/mocker/src/scheduler/vllm/policy/tests.rs b/lib/mocker/src/scheduler/vllm/policy/tests.rs index d87d59f572a8..3ae6296c9e9c 100644 --- a/lib/mocker/src/scheduler/vllm/policy/tests.rs +++ b/lib/mocker/src/scheduler/vllm/policy/tests.rs @@ -16,7 +16,7 @@ use crate::kv_manager::KvManager; use crate::kv_manager::kvbm_backend::G1Acquire; use crate::scheduler::vllm::{RequestStatus, VllmCore}; -use super::{AdmissionDecision, decide_waiting_admission}; +use super::{AdmissionDecision, decide_waiting_admission, should_reject_for_model_len}; mod vllm { use super::*; @@ -82,6 +82,152 @@ mod vllm { assert!(matches!(decision, AdmissionDecision::Reject)); } + #[test] + fn rejects_prompt_at_max_model_len() { + let sequence = ActiveSequence::new((0..8).collect(), 1, Some(4), false, false); + + assert!(should_reject_for_model_len( + SchedulingPolicy::Vllm, + &sequence, + Some(8) + )); + } + + #[test] + fn rejects_prompt_above_max_model_len() { + let sequence = ActiveSequence::new((0..9).collect(), 1, Some(4), false, false); + + assert!(should_reject_for_model_len( + SchedulingPolicy::Vllm, + &sequence, + Some(8) + )); + } + + #[test] + fn trtllm_does_not_apply_vllm_max_model_len() { + let sequence = ActiveSequence::new((0..9).collect(), 1, Some(4), false, false); + + assert!(!should_reject_for_model_len( + SchedulingPolicy::TrtllmGuaranteedNoEvict, + &sequence, + Some(8) + )); + } + + #[test] + fn core_rejects_prompt_above_max_model_len() { + let args = MockEngineArgs::builder() + .engine_type(EngineType::Vllm) + .block_size(4) + .num_gpu_blocks(4) + .max_model_len(Some(8)) + .max_num_batched_tokens(Some(16)) + .max_num_seqs(Some(4)) + .enable_chunked_prefill(true) + .enable_prefix_caching(false) + .speedup_ratio(0.0) + .build() + .unwrap(); + let mut core = VllmCore::new(args); + let uuid = Uuid::from_u128(1); + core.receive(DirectRequest { + tokens: (0..9).collect(), + max_output_tokens: 1, + uuid: Some(uuid), + dp_rank: 0, + ..Default::default() + }); + + let mut collector = crate::replay::TraceCollector::default(); + let pass = core.execute_pass(&mut collector, 0.0); + + assert!( + pass.output_signals + .iter() + .any(|signal| signal.uuid == uuid && signal.completed && signal.rejected) + ); + assert!(!core.state().requests.contains_key(&uuid)); + } + + #[test] + fn core_completes_at_max_model_len_without_rejecting() { + let args = MockEngineArgs::builder() + .engine_type(EngineType::Vllm) + .block_size(4) + .num_gpu_blocks(4) + .max_model_len(Some(8)) + .max_num_batched_tokens(Some(16)) + .max_num_seqs(Some(4)) + .enable_chunked_prefill(true) + .enable_prefix_caching(false) + .speedup_ratio(0.0) + .build() + .unwrap(); + let mut core = VllmCore::new(args); + let uuid = Uuid::from_u128(2); + core.receive(DirectRequest { + tokens: (0..7).collect(), + max_output_tokens: 4, + uuid: Some(uuid), + dp_rank: 0, + ..Default::default() + }); + + let mut collector = crate::replay::TraceCollector::default(); + let pass = core.execute_pass(&mut collector, 0.0); + + assert_eq!(pass.output_signals.len(), 1); + let terminal = &pass.output_signals[0]; + assert_eq!(terminal.uuid, uuid); + assert!(terminal.token_id.is_some()); + assert!(terminal.completed); + assert!(!terminal.rejected); + assert!(!core.state().requests.contains_key(&uuid)); + } + + #[test] + fn speculative_decode_does_not_burst_past_max_model_len() { + let args = MockEngineArgs::builder() + .engine_type(EngineType::Vllm) + .block_size(4) + .num_gpu_blocks(4) + .max_model_len(Some(8)) + .max_num_batched_tokens(Some(16)) + .max_num_seqs(Some(4)) + .enable_chunked_prefill(true) + .enable_prefix_caching(false) + .speedup_ratio(0.0) + .aic_nextn(Some(2)) + .aic_nextn_accept_rates(Some("1,1".to_string())) + .build() + .unwrap(); + let mut core = VllmCore::new(args); + let uuid = Uuid::from_u128(3); + core.receive(DirectRequest { + tokens: (0..5).collect(), + max_output_tokens: 8, + uuid: Some(uuid), + dp_rank: 0, + ..Default::default() + }); + + let mut collector = crate::replay::TraceCollector::default(); + let pass = core.execute_pass(&mut collector, 0.0); + + assert_eq!(pass.output_signals.len(), 3); + assert!( + pass.output_signals + .iter() + .take(2) + .all(|signal| !signal.completed) + ); + let terminal = pass.output_signals.last().unwrap(); + assert!(terminal.completed); + assert!(!terminal.rejected); + assert!(!core.state().requests.contains_key(&uuid)); + } + #[test] fn discounts_active_cached_prefix() { let mut manager = kv_manager(3); diff --git a/lib/mocker/src/scheduler/vllm/tests.rs b/lib/mocker/src/scheduler/vllm/tests.rs index ff6f9d7efb1a..336215c31450 100644 --- a/lib/mocker/src/scheduler/vllm/tests.rs +++ b/lib/mocker/src/scheduler/vllm/tests.rs @@ -255,6 +255,52 @@ mod destination_lifecycle { assert!(later.iter().all(|hash| !activation.contains(hash))); } + #[test] + fn materialized_prompt_above_max_model_len_is_rejected() { + let args = MockEngineArgs::builder() + .block_size(4) + .num_gpu_blocks(12) + .max_model_len(Some(8)) + .max_num_batched_tokens(Some(16)) + .max_num_seqs(Some(1)) + .enable_chunked_prefill(true) + .enable_prefix_caching(true) + .worker_type(WorkerType::Decode) + .speedup_ratio(0.0) + .build() + .unwrap(); + let mut core = VllmCore::new(args); + let handoff_id = HandoffId::from(Uuid::from_u128(30_001)); + let uuid = Uuid::from_u128(30_002); + + assert!(matches!( + core.apply_command(SchedulerCommand::ReserveDestination { + handoff_id, + request: request(uuid, vec![1; 9], 1), + }) + .unwrap(), + SchedulerCommandResult::DestinationAccepted { request_id } if request_id == uuid + )); + assert_eq!( + core.apply_command(SchedulerCommand::ActivateDestination { handoff_id }) + .unwrap(), + SchedulerCommandResult::Applied + ); + + let pass = execute(&mut core, 0.0); + assert!(matches!( + pass.output_signals.as_slice(), + [OutputSignal { + uuid: signal_uuid, + token_id: None, + completed: true, + rejected: true, + .. + }] if *signal_uuid == uuid + )); + assert!(!core.state().requests.contains_key(&uuid)); + } + fn drive_source_to_hold(core: &mut VllmCore, handoff_id: HandoffId, req: DirectRequest) { assert!(matches!( core.apply_command(SchedulerCommand::SubmitHandoffPrefill { From c790d15849eafc02cb2e7bffee214a77ff6d0286 Mon Sep 17 00:00:00 2001 From: Qi Wang Date: Wed, 1 Jul 2026 13:03:46 +0900 Subject: [PATCH 020/320] feat(multimodal): CustomEncoder ABC + mixed-embeds assembler + --custom-encoder-class (#10910) Co-authored-by: Claude Opus 4.8 (1M context) --- components/src/dynamo/vllm/backend_args.py | 74 ++++++++ .../dynamo/vllm/multimodal_utils/__init__.py | 8 + .../vllm/multimodal_utils/embed_assembler.py | 152 ++++++++++++++++ .../vision_encoder_backend.py | 169 ++++++++++++++++++ .../test_vllm_custom_encoder_flow.py | 78 ++++++++ .../test_vllm_embed_assembler.py | 110 ++++++++++++ .../test_vllm_vision_encoder_backend.py | 101 +++++++++++ .../dynamo/vllm/tests/test_backend_args.py | 95 ++++++++++ tests/frontend/test_realtime_omni_bridge.py | 13 ++ 9 files changed, 800 insertions(+) create mode 100644 components/src/dynamo/vllm/multimodal_utils/embed_assembler.py create mode 100644 components/src/dynamo/vllm/multimodal_utils/vision_encoder_backend.py create mode 100644 components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_custom_encoder_flow.py create mode 100644 components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_embed_assembler.py create mode 100644 components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_vision_encoder_backend.py diff --git a/components/src/dynamo/vllm/backend_args.py b/components/src/dynamo/vllm/backend_args.py index 2ad5a3915c19..bdc477e604bc 100644 --- a/components/src/dynamo/vllm/backend_args.py +++ b/components/src/dynamo/vllm/backend_args.py @@ -145,6 +145,21 @@ def add_arguments(self, parser) -> None: add_frontend_decoding_arg(g, env_prefix="VLLM") + add_argument( + g, + flag_name="--custom-encoder-class", + env_var="DYN_CUSTOM_ENCODER_CLASS", + default=None, + help=( + "Dotted module.ClassName path to a VisionEncoderBackend subclass. " + "When set, the aggregated worker wraps it in the in-process " + "AsyncVisionEncoder and runs encoder.encode(image_urls) for each " + "multimodal request, bypassing vLLM's built-in multimodal " + "processing. --model is passed verbatim to the backend's build(). " + "Example: 'my_package.encoders.MyEncoder'." + ), + ) + add_argument( g, flag_name="--embedding-transfer-mode", @@ -299,6 +314,9 @@ class DynamoVllmConfig(ConfigBase): ] # resolved to enum in validate() embedding_worker: bool = False + # CustomEncoder (image-only embeddings; worker assembles mixed prompt) + custom_encoder_class: Optional[str] = None + # Headless mode for multi-node TP/PP headless: bool = False @@ -324,6 +342,7 @@ def validate(self) -> None: self._validate_multimodal_role_exclusivity() self._validate_multimodal_requires_flag() self._validate_embedding_worker_exclusivity() + self._validate_custom_encoder() def _resolve_embedding_transfer_mode(self) -> None: """Resolve embedding_transfer_mode from string to enum.""" @@ -488,6 +507,61 @@ def _validate_multimodal_requires_flag(self) -> None: "Use --enable-multimodal when enabling any multimodal component" ) + def _validate_custom_encoder(self) -> None: + """Validate the aggregated CustomEncoder configuration. + + The encoder runs in-process in a single aggregated worker on the + token-in/token-out path and produces image embeds for the mixed + EmbedsPrompt, so it is a multimodal, aggregated-only, token-mode + component. Enforce those here (fail fast) instead of silently bypassing + the multimodal gate at request time, no-op'ing in a decode worker that + never reaches the custom-encoder branch, or loading the encoder in + --use-vllm-tokenizer text mode where it is never invoked. + """ + if not self.custom_encoder_class: + return + if ( + self.multimodal_worker + or self.multimodal_encode_worker + or self.multimodal_decode_worker + ): + raise ValueError( + "--custom-encoder-class is incompatible with the legacy multimodal " + "role flags (--multimodal-worker / --multimodal-encode-worker / " + "--multimodal-decode-worker): the custom encoder is its own " + "aggregated multimodal path and bypasses vLLM's built-in " + "multimodal processing." + ) + if not self.enable_multimodal: + raise ValueError( + "--custom-encoder-class requires --enable-multimodal " + "(the custom encoder is a multimodal component)." + ) + if self.use_vllm_tokenizer: + raise ValueError( + "--custom-encoder-class is incompatible with --use-vllm-tokenizer: " + "the custom encoder is wired into the token-in/token-out path, " + "which --use-vllm-tokenizer bypasses (text mode), so the encoder " + "would load but never run." + ) + if self.frontend_decoding: + raise ValueError( + "--custom-encoder-class is incompatible with --frontend-decoding: " + "the custom encoder consumes image URLs, but frontend decoding " + "pre-decodes images to tensors the encoder cannot accept." + ) + if self.disaggregation_mode != DisaggregationMode.AGGREGATED: + mode = ( + self.disaggregation_mode.value + if isinstance(self.disaggregation_mode, DisaggregationMode) + else self.disaggregation_mode + ) + raise ValueError( + f"--custom-encoder-class is only supported with " + f"--disaggregation-mode=agg (got {mode}). The custom encoder " + "runs in-process in a single aggregated worker." + ) + def _validate_embedding_worker_exclusivity(self) -> None: """Embedding worker is aggregated-only and exclusive of multimodal roles.""" if not self.embedding_worker: diff --git a/components/src/dynamo/vllm/multimodal_utils/__init__.py b/components/src/dynamo/vllm/multimodal_utils/__init__.py index 4e44b4dff896..8463a193703c 100644 --- a/components/src/dynamo/vllm/multimodal_utils/__init__.py +++ b/components/src/dynamo/vllm/multimodal_utils/__init__.py @@ -3,6 +3,7 @@ from dynamo.common.multimodal.image_loader import ImageLoader from dynamo.vllm.multimodal_utils.chat_message_utils import extract_user_text +from dynamo.vllm.multimodal_utils.embed_assembler import build_mixed_embeds from dynamo.vllm.multimodal_utils.encode_utils import ( encode_image_embeddings, get_embedding_hash, @@ -23,11 +24,18 @@ PatchedTokensPrompt, vLLMMultimodalRequest, ) +from dynamo.vllm.multimodal_utils.vision_encoder_backend import ( + Preprocessed, + VisionEncoderBackend, +) __all__ = [ + "build_mixed_embeds", "encode_image_embeddings", "extract_user_text", "get_encoder_components", + "Preprocessed", + "VisionEncoderBackend", "ImageLoader", "ModelFamily", "construct_mm_data", diff --git a/components/src/dynamo/vllm/multimodal_utils/embed_assembler.py b/components/src/dynamo/vllm/multimodal_utils/embed_assembler.py new file mode 100644 index 000000000000..221fed17aeef --- /dev/null +++ b/components/src/dynamo/vllm/multimodal_utils/embed_assembler.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mixed token-ids/embeds assembly for the aggregated CustomEncoder path. + +The encoder returns only the visual token embeddings; this module builds the +inputs for vLLM's mixed ``EmbedsPrompt`` mode (``prompt_token_ids`` + +``prompt_is_token_ids`` + ``prompt_embeds``): + + prompt_token_ids = [ text ... ... text ] + prompt_is_token_ids = [ True ... False False False ... True ] + prompt_embeds = [ zeros ... e0 e1 e2 ... zeros ] (seq_len, hidden) + +One image occupies a **contiguous run** of ``False`` positions — the single +placeholder token is expanded to the encoder tensor's row count (3 here), and +that image's embeds (``e0,e1,e2``) fill exactly those rows. vLLM embeds the +``True`` (text) positions itself with the model's real embedding table and +substitutes each ``False`` (image) row from ``prompt_embeds`` in the forward +pass. Dynamo therefore only fills the image rows — text rows stay zero (they are +overwritten) and no LM embedding weight is needed on the Dynamo side. + +The contract is **one placeholder token per image**: each occurrence of the +placeholder token in ``prompt_token_ids`` is one image slot, matched +positionally to the encoder tensors, and the single placeholder is expanded to +the tensor's row count so the encoder dictates the span length (mirroring +vLLM's own placeholder expansion); this keeps a mismatch between the tokenizer's +placeholder count and the encoder's visual-token count from raising. The chat +template therefore emits exactly one placeholder token per image and needs no +separator between consecutive images. +""" + +from __future__ import annotations + +import logging + +import torch + +logger = logging.getLogger(__name__) + + +def build_mixed_embeds( + token_ids: list[int], + img_tensors: list[torch.Tensor], + placeholder_id: int, +) -> tuple[torch.Tensor, list[int], list[bool]]: + """Build the mixed token-ids/embeds inputs for an aggregated request. + + Args: + token_ids: The full prompt token IDs (text + one placeholder token per + image). + img_tensors: Per-image visual token tensors, each ``(n_tokens, hidden)``, + in prompt order. + placeholder_id: The token ID marking image positions. + + Returns: + ``(prompt_embeds, prompt_token_ids, prompt_is_token_ids)`` where + ``prompt_embeds`` is a CPU ``(seq_len, hidden)`` tensor (text rows zero, + image rows from ``img_tensors``), ``prompt_token_ids`` has each + placeholder token expanded to its tensor's row count, and + ``prompt_is_token_ids[i]`` is ``False`` at image positions. + + Raises: + ValueError: if ``img_tensors`` is empty, the number of placeholder + tokens does not equal the number of image tensors, or the tensors + are not 2D with a consistent hidden dim. + """ + if not img_tensors: + raise ValueError("img_tensors must not be empty") + + positions = [i for i, tid in enumerate(token_ids) if tid == placeholder_id] + if len(positions) != len(img_tensors): + raise ValueError( + f"placeholder tokens ({len(positions)}) != image tensors " + f"({len(img_tensors)}) for placeholder token {placeholder_id} " + f"in sequence of length {len(token_ids)}" + ) + + # Check tensor 0 is 2D before reading its hidden dim, so a 1D encoder output + # raises a clear ValueError here instead of an opaque IndexError on shape[1]. + if img_tensors[0].dim() != 2: + raise ValueError( + f"image tensor 0 has shape {tuple(img_tensors[0].shape)}; expected " + "2D (n_tokens, hidden)" + ) + hidden = img_tensors[0].shape[1] + dtype = img_tensors[0].dtype + # Validate shapes before scattering so a bad encoder output raises a clear + # ValueError here (caught by the caller) instead of an opaque RuntimeError + # from the row-copy below on a width mismatch. + for i, tensor in enumerate(img_tensors): + if tensor.dim() != 2 or tensor.shape[1] != hidden: + raise ValueError( + f"image tensor {i} has shape {tuple(tensor.shape)}; expected " + f"2D with hidden dim {hidden} (from image tensor 0)" + ) + # A (0, hidden) tensor passes the 2D/hidden checks but would erase the + # image's placeholder token entirely, silently dropping the image from + # the prompt. An encoder returning no visual tokens for an image is a + # bug — fail loudly instead. + if tensor.shape[0] == 0: + raise ValueError( + f"image tensor {i} has 0 rows (shape {tuple(tensor.shape)}); the " + "encoder returned no visual tokens for an image" + ) + # forward_batch must fence + copy to CPU before returning, so the scatter + # below is a plain assignment into the CPU prompt_embeds buffer. Fail loud + # here instead of an opaque cross-device error on the row-copy. + if tensor.device.type != "cpu": + raise ValueError( + f"image tensor {i} is on {tensor.device}; forward_batch must " + "return CPU tensors" + ) + + # Build the token-id / mask layout and record where each image block lands, + # then scatter the image rows into one pre-zeroed (seq_len, hidden) buffer. + # Text rows stay zero (vLLM overwrites them via the model's embedding table), + # so there is no need to allocate per-text-segment zero tensors and concat. + out_token_ids: list[int] = [] + is_token_ids: list[bool] = [] + image_slots: list[tuple[int, torch.Tensor]] = [] # (row_start, tensor) + + def _emit_text(text_ids: list[int]) -> None: + if not text_ids: + return + out_token_ids.extend(text_ids) + is_token_ids.extend([True] * len(text_ids)) + + cursor = 0 + for pos, tensor in zip(positions, img_tensors): + _emit_text(token_ids[cursor:pos]) + n = tensor.shape[0] + image_slots.append((len(out_token_ids), tensor)) + out_token_ids.extend([placeholder_id] * n) + is_token_ids.extend([False] * n) + cursor = pos + 1 + _emit_text(token_ids[cursor:]) + + seq_len = len(out_token_ids) + # CPU tensor: vLLM's renderer forces prompt_embeds to CPU anyway. + prompt_embeds = torch.zeros(seq_len, hidden, dtype=dtype) + for row_start, tensor in image_slots: + n = tensor.shape[0] + prompt_embeds[row_start : row_start + n] = tensor + + logger.debug( + "[custom_embeds] images=%d seq_len=%d hidden=%d dtype=%s", + len(img_tensors), + seq_len, + hidden, + dtype, + ) + return prompt_embeds, out_token_ids, is_token_ids diff --git a/components/src/dynamo/vllm/multimodal_utils/vision_encoder_backend.py b/components/src/dynamo/vllm/multimodal_utils/vision_encoder_backend.py new file mode 100644 index 000000000000..193c2bb085c5 --- /dev/null +++ b/components/src/dynamo/vllm/multimodal_utils/vision_encoder_backend.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The author-written contract for a pluggable in-process vision encoder. + +``VisionEncoderBackend`` is the **single surface an encoder author implements**. +It is a pure policy + compute backend: no threads, no futures, no event loop. +Dynamo owns all the *driving* — the dedicated actor thread, cross-request +coalescing, the embeds splice, and the lifecycle — via ``ThreadedMicroBatcher`` +(the generic cross-request batcher) and ``AsyncVisionEncoder`` (the async +request-API glue). This module defines only the contract those drivers call. + +The encoder runs in the **same process** as the aggregated vLLM worker (no +separate encode worker, no NIXL transfer): it turns image inputs into the +visual-token embeddings for each image, and Dynamo splices those embeds into a +mixed ``EmbedsPrompt`` at the placeholder positions (see +``embed_assembler.build_mixed_embeds``) for a text-only LM. + +Division of labour (author vs. Dynamo): + +- ``build(model_id)`` — **actor thread, once.** Load weights / tokenizer; warm up + to peak; if ``buckets`` is set (once CUDA-graph batching is supported), capture + one CUDA graph per rung here so it is bound to the thread that later replays it + in ``forward_batch``. Pick the device yourself (``"cuda"`` / the current device). +- ``preprocess(raw) -> Preprocessed{item, cost}`` — **off the actor thread, + concurrent.** Deterministic, thread-safe, CUDA-free (fetch / resize / patchify + on CPU/pinned memory). ``cost`` is a **scalar** — how much the item adds toward + ``max_batch_cost`` (e.g. its visual-token count). Raise to reject a bad input — + it fails only that image, before any GPU work. **Off by default:** override + ``preprocess`` *and* set ``preprocess_concurrency > 0`` together to enable this + pool. With the defaults (identity passthrough, ``preprocess_concurrency = 0``) + there is no preprocess phase — raws go straight to ``forward_batch``. +- ``forward_batch(items, target_bucket=None) -> list[torch.Tensor]`` — **actor + thread, serialized.** ``items`` are a cost-bounded batch (summed ``cost`` within + the budget). Fence (stream event + sync) and **copy outputs to CPU** before + returning, so results are safe to consume from another thread and splice + directly. Returns one ``(n_visual_tokens, lm_hidden_dim)`` **CPU** tensor per + item, in input order. ``target_bucket`` is reserved for CUDA-graph batching, + once supported (the ladder rung to pad to); it is ``None`` until then. +- ``close()`` — actor thread, on teardown. Release any thread-affine resources. + +Attributes read **once at setup** (never per-request): + +- ``image_token_id`` — the token id marking image positions in the prompt; + **hardcode it for your model** (e.g. ``151655`` for Qwen3-VL's ``<|image_pad|>``). + Dynamo uses it to locate each image span for the splice. +- ``max_batch_cost`` — the scalar dispatch ceiling the batcher packs up to; a + *chosen* budget (a token budget when ``cost`` is a token count). ``None`` (the + default) ⇒ **pass-through**: no cap (the author owns sizing). +- ``buckets`` — sorted graph ladder, forward-compatible (unused until CUDA-graph + batching is supported). ``None``/empty ⇒ eager. +- ``preprocess_concurrency`` — size of the off-thread pool Dynamo runs + ``preprocess`` on. ``0`` (the **default**) ⇒ no preprocess phase: raws go + straight to ``forward_batch``. Set ``> 0`` (with an overridden ``preprocess``) + for off-loop fetch / resize / patchify. + +Batching is **one-dimensional**: Dynamo packs by scalar ``cost`` up to +``max_batch_cost`` and never inspects item shape — the author owns any +shape/padding concerns inside ``forward_batch``. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Generic, List, Optional, Sequence, TypeVar + +import torch + +RawT = TypeVar("RawT") # raw input the author preprocesses (e.g. an image URL) +ItemT = TypeVar("ItemT") # opaque payload preprocess() hands to forward_batch() + + +@dataclass(frozen=True) +class Preprocessed(Generic[ItemT]): + """The result of ``preprocess(raw)``: an opaque item plus its batching cost. + + ``cost`` is computed **once, off the actor thread**, so the batcher never + evaluates model policy (it stays torch-free) and packs purely by this scalar. + + Attributes: + item: Opaque payload passed verbatim to ``forward_batch``. + cost: Scalar size of this item (``>= 1``); packs toward ``max_batch_cost``. + Read only in **budgeted mode** (``max_batch_cost`` set). In + **pass-through mode** (``max_batch_cost`` is ``None``) the batcher + never reads it, so a pass-through author can leave it at the default + ``1``. + """ + + item: ItemT + cost: int = 1 + + +class VisionEncoderBackend(ABC, Generic[RawT, ItemT]): + """Author-written, in-process vision encoder contract. + + A pure policy + compute backend — no threads, no futures. Dynamo drives it + on a dedicated actor thread (``ThreadedMicroBatcher``) and exposes the async + request API (``AsyncVisionEncoder``). Subclasses implement ``build`` and + ``forward_batch`` and set ``image_token_id``; ``preprocess`` (default identity + passthrough), ``max_batch_cost``, ``buckets``, and ``preprocess_concurrency`` + are overridden only as needed. + """ + + #: Image placeholder token id — **hardcode it for your model** (e.g. ``151655`` + #: for Qwen3-VL's ``<|image_pad|>``; resolve it from your tokenizer offline if + #: unsure). Dynamo uses it to locate each image span for the splice. Declared + #: without a default so a backend that forgets to set it fails fast at startup. + image_token_id: int + + #: Scalar dispatch ceiling: the batcher packs items up to this summed ``cost`` + #: per ``forward_batch`` call. ``None`` (the default) ⇒ **pass-through**: no cap + #: — every drained item in one iteration is handed to a single ``forward_batch`` + #: (the author owns sizing; ``cost`` is ignored). + max_batch_cost: Optional[int] = None + + #: Sorted graph ladder (the captured rungs), **forward-compatible** — unused + #: until CUDA-graph batching is supported. ``None``/empty ⇒ eager. + buckets: Optional[Sequence[int]] = None + + #: Off-loop preprocess pool size Dynamo runs ``preprocess`` on. ``0`` (the + #: **default**) ⇒ **no preprocess phase**: raws go straight to ``forward_batch`` + #: (``raw`` is the item; do any prep there). Set ``> 0`` (with an overridden + #: ``preprocess``) to fetch / resize / patchify off the actor thread. Whether an + #: encoder needs off-loop prep is a property of the encoder, so it lives here; + #: the driver takes an optional override for tuning. + preprocess_concurrency: int = 0 + + # ---- subclass contract ------------------------------------------------- + + @abstractmethod + def build(self, model_id: str) -> None: + """Load weights / tokenizer, warm up, capture graphs (actor thread, once). + + Any CUDA graph captured here is bound to the thread that later replays it. + Pick the device yourself (``"cuda"`` / the current device). All CUDA init + happens here. + """ + ... + + def preprocess(self, raw: RawT) -> Preprocessed[ItemT]: + """Turn a raw input into a ``Preprocessed`` item (off the actor thread). + + The default is an **identity passthrough** (``raw`` is the item, ``cost`` + ``1``), so by default there is no preprocessing. Override it for off-loop + fetch + HF processing **and** set ``preprocess_concurrency > 0`` to run it + on the pool — it must then be deterministic, thread-safe, and CUDA-free. + Raise to reject a bad input — it fails only that image, before submit. + """ + return Preprocessed(item=raw) # type: ignore[arg-type] # ItemT == RawT + + @abstractmethod + def forward_batch( + self, items: List[ItemT], target_bucket: Optional[int] = None + ) -> List[torch.Tensor]: + """Encode one cost-bounded batch (actor thread); one tensor per item, in order. + + Fence (stream event + sync) and **copy outputs to CPU** before returning, + so results are safe to consume from another thread and splice directly. + Return one ``(n_visual_tokens, lm_hidden_dim)`` **CPU** tensor per item, in + input order. ``target_bucket`` is reserved for CUDA-graph batching, once + supported (the ladder rung to pad to), and is ``None`` until then. + """ + ... + + def close(self) -> None: + """Release thread-affine resources on teardown (actor thread). No-op by + default; override to free graphs / pools / weights.""" + return None diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_custom_encoder_flow.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_custom_encoder_flow.py new file mode 100644 index 000000000000..7ab647f1000f --- /dev/null +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_custom_encoder_flow.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Interaction test: a VisionEncoderBackend feeding build_mixed_embeds. + +The unit tests cover each side with the other mocked out — the backend test never +touches the assembler, and the assembler test hand-builds tensors instead of a +backend — so neither pins the seam. This walks the whole contract for one request +(build -> preprocess -> forward_batch -> build_mixed_embeds) and doubles as a +worked example of how the APIs fit: forward_batch emits one CPU +(n_tokens, hidden) tensor per image, and build_mixed_embeds splices them at the +one-placeholder-token-per-image positions. +""" + +import pytest +import torch + +from dynamo.vllm.multimodal_utils.embed_assembler import build_mixed_embeds +from dynamo.vllm.multimodal_utils.vision_encoder_backend import ( + Preprocessed, + VisionEncoderBackend, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.multimodal, +] + +_HIDDEN = 4 +_IMG = 151655 # the toy backend's hardcoded image placeholder token id + + +class _ToyEncoder(VisionEncoderBackend): + """Stand-in backend: one visual token per input char, distinct per-image embeds.""" + + image_token_id = _IMG + + def build(self, model_id): + ... + + def preprocess(self, raw): + n = len(raw) + return Preprocessed(item=(raw, n), cost=n) + + def forward_batch(self, items, target_bucket=None): + # Distinct non-zero fill per image (1.0, 2.0, ...) so the splice is + # unambiguous against the zero-filled text rows. CPU, per the contract. + return [ + torch.full((n, _HIDDEN), float(i + 1)) for i, (_, n) in enumerate(items) + ] + + +def test_backend_and_assembler_work_together(): + """Worked example of the whole contract for one two-image request.""" + enc = _ToyEncoder() + enc.build("toy-model") + + # Caller flow: preprocess each raw off-thread, then one batched forward. + pre = [enc.preprocess(r) for r in ("ab", "cde")] + assert [p.cost for p in pre] == [2, 3] + img_tensors = enc.forward_batch([p.item for p in pre]) + assert [tuple(t.shape) for t in img_tensors] == [(2, _HIDDEN), (3, _HIDDEN)] + + # Prompt: 10 11 12 — one placeholder token per image. + prompt_embeds, out_ids, is_token = build_mixed_embeds( + [10, _IMG, 11, 12, _IMG], img_tensors, enc.image_token_id + ) + + # Each placeholder is expanded to its image's row count (2 and 3). + assert out_ids == [10, _IMG, _IMG, 11, 12, _IMG, _IMG, _IMG] + assert is_token == [True, False, False, True, True, False, False, False] + # Image spans carry the encoder embeds; text rows stay zero. + assert torch.equal(prompt_embeds[1:3], torch.full((2, _HIDDEN), 1.0)) + assert torch.equal(prompt_embeds[5:8], torch.full((3, _HIDDEN), 2.0)) + assert torch.equal(prompt_embeds[0], torch.zeros(_HIDDEN)) diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_embed_assembler.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_embed_assembler.py new file mode 100644 index 000000000000..667bb6f53a07 --- /dev/null +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_embed_assembler.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for dynamo.vllm.multimodal_utils.embed_assembler. + +build_mixed_embeds assembles the mixed token-ids/embeds EmbedsPrompt for the +aggregated CustomEncoder path: each placeholder token is expanded to its encoder +tensor's row count, image rows carry the encoder embeddings, and text rows stay +zero (vLLM fills them from the model's embedding table). These tests pin that +layout, the per-image expansion (including back-to-back images), and the input +validation that surfaces a bad encoder output as a clear ValueError. +""" + +import pytest +import torch + +from dynamo.vllm.multimodal_utils.embed_assembler import build_mixed_embeds + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.multimodal, +] + +_HIDDEN = 8 +_PLACEHOLDER = 999 + + +def test_build_mixed_embeds_multi_image_token_expand(): + """Each placeholder token is one image and is expanded to its encoder + tensor's row count; back-to-back placeholders (adjacent images, no + separator) yield one block per image.""" + # img_a, then text, then img_b and img_c back-to-back (no separator). + token_ids = [1, 2, _PLACEHOLDER, 3, _PLACEHOLDER, _PLACEHOLDER, 4, 5] + img_a = torch.ones(2, _HIDDEN, dtype=torch.float16) + img_b = torch.ones(3, _HIDDEN, dtype=torch.float16) * 2.0 + img_c = torch.ones(1, _HIDDEN, dtype=torch.float16) * 3.0 + + embeds, out_ids, is_tok = build_mixed_embeds( + token_ids, [img_a, img_b, img_c], _PLACEHOLDER + ) + + # Layout: [1,2] + img_a(2) + [3] + img_b(3) + img_c(1) + [4,5] -> 11 rows. + assert embeds.shape == (11, _HIDDEN) + assert embeds.dtype == torch.float16 + assert embeds.device.type == "cpu" + assert len(out_ids) == 11 and len(is_tok) == 11 + assert out_ids == [1, 2, 999, 999, 3, 999, 999, 999, 999, 4, 5] + assert is_tok == [ + True, + True, + False, + False, + True, + False, + False, + False, + False, + True, + True, + ] + # Each image's rows carry its encoder values; text rows stay zero. + assert torch.all(embeds[2:4] == 1.0) # img_a + assert torch.all(embeds[5:8] == 2.0) # img_b + assert torch.all(embeds[8] == 3.0) # img_c (adjacent to img_b, no separator) + assert torch.all(embeds[:2] == 0) + assert torch.all(embeds[4] == 0) + assert torch.all(embeds[9:] == 0) + + +@pytest.mark.parametrize( + "token_ids, n_tensors", + [ + pytest.param([1, _PLACEHOLDER, 2], 2, id="more_tensors_than_placeholders"), + pytest.param([1, 2, 3], 1, id="no_placeholders_but_tensors"), + ], +) +def test_build_mixed_embeds_raises_on_placeholder_tensor_mismatch(token_ids, n_tensors): + """A placeholder-token count that differs from the image-tensor count is a + caller error and must raise ValueError, not silently mis-scatter.""" + tensors = [torch.ones(1, _HIDDEN, dtype=torch.float16)] * n_tensors + with pytest.raises(ValueError): + build_mixed_embeds(token_ids, tensors, _PLACEHOLDER) + + +def test_build_mixed_embeds_raises_on_empty_tensors(): + with pytest.raises(ValueError): + build_mixed_embeds([1, _PLACEHOLDER, 2], [], _PLACEHOLDER) + + +def test_build_mixed_embeds_raises_on_bad_tensor_shape(): + """A 1D encoder tensor (missing the hidden dim) must raise before the row + copy, not surface as an opaque RuntimeError.""" + with pytest.raises(ValueError): + build_mixed_embeds( + [1, _PLACEHOLDER, 2], [torch.ones(4, dtype=torch.float16)], _PLACEHOLDER + ) + + +def test_build_mixed_embeds_raises_on_empty_rows(): + """A (0, hidden) encoder tensor passes the 2D/hidden checks but would erase + the image's placeholder run, silently dropping the image — must raise.""" + with pytest.raises(ValueError, match="0 rows"): + build_mixed_embeds( + [1, _PLACEHOLDER, 2], + [torch.empty(0, _HIDDEN, dtype=torch.float16)], + _PLACEHOLDER, + ) diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_vision_encoder_backend.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_vision_encoder_backend.py new file mode 100644 index 000000000000..be08946dbf44 --- /dev/null +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_vision_encoder_backend.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for dynamo.vllm.multimodal_utils.vision_encoder_backend. + +Pin the author-facing contract surface: the ``Preprocessed`` carrier (item + +scalar cost, no bucket_key), the hardcoded ``image_token_id`` attribute, the +no-device ``build`` signature, and that the ABC cannot be instantiated without +the required methods. +""" + +import pytest +import torch + +from dynamo.vllm.multimodal_utils.vision_encoder_backend import ( + Preprocessed, + VisionEncoderBackend, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.multimodal, +] + + +class _MinimalBackend(VisionEncoderBackend): + """Smallest concrete backend — hardcodes the image token id.""" + + image_token_id = 151655 + + def build(self, model_id): + ... + + def preprocess(self, raw): + return Preprocessed(item=raw, cost=1) + + def forward_batch(self, items, target_bucket=None): + return [torch.zeros(1, 1) for _ in items] + + +class _PassthroughBackend(VisionEncoderBackend): + """Backend that does NOT override preprocess — exercises the identity default + (no preprocess phase; raws go straight to forward_batch).""" + + image_token_id = 151655 + + def build(self, model_id): + ... + + def forward_batch(self, items, target_bucket=None): + return [torch.zeros(1, 1) for _ in items] + + +def test_preprocessed_is_frozen(): + p = Preprocessed(item="x", cost=3) + assert (p.item, p.cost) == ("x", 3) + with pytest.raises(Exception): # FrozenInstanceError + p.cost = 4 # type: ignore[misc] + + +def test_preprocessed_cost_defaults_to_1(): + # A pass-through author can omit cost entirely. + assert Preprocessed(item="x").cost == 1 + + +def test_preprocessed_has_no_bucket_key(): + # Batching is one-dimensional (scalar cost only) — there is no bucket_key. + assert not hasattr(Preprocessed(item="x"), "bucket_key") + + +def test_abc_cannot_be_instantiated(): + # build + forward_batch are still abstract (preprocess is not). + with pytest.raises(TypeError): + VisionEncoderBackend() # type: ignore[abstract] + + +def test_preprocess_defaults_to_identity_passthrough(): + # Not overriding preprocess ⇒ raw IS the item, cost 1 (no preprocess phase). + p = _PassthroughBackend().preprocess("http://img") + assert isinstance(p, Preprocessed) + assert (p.item, p.cost) == ("http://img", 1) + + +def test_preprocess_concurrency_defaults_to_0(): + # No pool by default — authors opt in by overriding preprocess + setting > 0. + assert _PassthroughBackend().preprocess_concurrency == 0 + assert _MinimalBackend().preprocess_concurrency == 0 + + +def test_default_attrs_and_close_noop(): + e = _MinimalBackend() + # Defaults from the ABC: eager (no ladder) + pass-through (no cost cap) + no + # preprocess pool. + assert e.buckets is None + assert e.max_batch_cost is None + assert e.preprocess_concurrency == 0 + assert e.image_token_id == 151655 + assert e.close() is None diff --git a/components/src/dynamo/vllm/tests/test_backend_args.py b/components/src/dynamo/vllm/tests/test_backend_args.py index e569f10b3d3d..151bd6564b3e 100644 --- a/components/src/dynamo/vllm/tests/test_backend_args.py +++ b/components/src/dynamo/vllm/tests/test_backend_args.py @@ -18,6 +18,7 @@ pytest.mark.vllm, pytest.mark.pre_merge, pytest.mark.gpu_0, + pytest.mark.multimodal, ] @@ -40,6 +41,8 @@ def create_config() -> DynamoVllmConfig: config.enable_multimodal = False config.embedding_worker = False config.benchmark_mode = None + config.use_vllm_tokenizer = False + config.frontend_decoding = False return config @@ -187,3 +190,95 @@ def test_no_op_when_embedding_worker_disabled(self): config.embedding_worker = False config.benchmark_mode = "agg" config._validate_embedding_worker_exclusivity() + + +class TestValidateCustomEncoder: + """--custom-encoder-class is an in-process, aggregated-only multimodal + component, so validation must require --enable-multimodal and reject any + non-aggregated disaggregation mode (where the custom-encoder branch is + never reached) up front. + """ + + def test_requires_enable_multimodal(self): + # Without the gate the custom encoder processes images while multimodal + # is disabled, bypassing the normal multimodal enable check. + config = create_config() + config.custom_encoder_class = "my_pkg.MyEncoder" + config.disaggregation_mode = DisaggregationMode.AGGREGATED + config.enable_multimodal = False + with pytest.raises(ValueError, match="enable-multimodal"): + config._validate_custom_encoder() + + @pytest.mark.parametrize( + "mode", + [ + DisaggregationMode.PREFILL, + DisaggregationMode.DECODE, + DisaggregationMode.ENCODE, + ], + ) + def test_non_aggregated_mode_rejected(self, mode): + config = create_config() + config.custom_encoder_class = "my_pkg.MyEncoder" + config.enable_multimodal = True + config.disaggregation_mode = mode + with pytest.raises(ValueError, match="agg"): + config._validate_custom_encoder() + + def test_use_vllm_tokenizer_rejected(self): + # --use-vllm-tokenizer routes to text mode, which never invokes the + # custom encoder, so the encoder would load but sit unused. Reject it. + config = create_config() + config.custom_encoder_class = "my_pkg.MyEncoder" + config.enable_multimodal = True + config.disaggregation_mode = DisaggregationMode.AGGREGATED + config.use_vllm_tokenizer = True + with pytest.raises(ValueError, match="use-vllm-tokenizer"): + config._validate_custom_encoder() + + @pytest.mark.parametrize( + "role_flag", + [ + "multimodal_worker", + "multimodal_encode_worker", + "multimodal_decode_worker", + ], + ) + def test_legacy_multimodal_role_rejected(self, role_flag): + # The custom encoder is its own aggregated multimodal path; combining it + # with a legacy multimodal role flag sets up two conflicting multimodal + # paths (and --multimodal-worker resolves to agg, slipping past the + # disaggregation-mode check), so reject the combination up front. + config = create_config() + config.custom_encoder_class = "my_pkg.MyEncoder" + config.enable_multimodal = True + config.disaggregation_mode = DisaggregationMode.AGGREGATED + setattr(config, role_flag, True) + with pytest.raises(ValueError, match="legacy multimodal role flags"): + config._validate_custom_encoder() + + def test_frontend_decoding_rejected(self): + # --frontend-decoding pre-decodes images to tensors; the custom encoder + # consumes URLs, so the decoded inputs would fail extraction. Reject it. + config = create_config() + config.custom_encoder_class = "my_pkg.MyEncoder" + config.enable_multimodal = True + config.disaggregation_mode = DisaggregationMode.AGGREGATED + config.frontend_decoding = True + with pytest.raises(ValueError, match="frontend-decoding"): + config._validate_custom_encoder() + + def test_accepted_when_agg_and_multimodal(self): + config = create_config() + config.custom_encoder_class = "my_pkg.MyEncoder" + config.enable_multimodal = True + config.disaggregation_mode = DisaggregationMode.AGGREGATED + # Must not raise. + config._validate_custom_encoder() + + def test_no_op_when_unset(self): + # No custom encoder → validator must not touch unrelated configs. + config = create_config() + config.custom_encoder_class = None + config.enable_multimodal = False + config._validate_custom_encoder() diff --git a/tests/frontend/test_realtime_omni_bridge.py b/tests/frontend/test_realtime_omni_bridge.py index 51cc999c4943..a58abfeb841a 100644 --- a/tests/frontend/test_realtime_omni_bridge.py +++ b/tests/frontend/test_realtime_omni_bridge.py @@ -38,6 +38,19 @@ ENDPOINT_PATH = "test_omni_ws_e2e.realtime.generate" MOCK_TRANSCRIPT = "mock omni transcript" +# The vllm-runtime image bumped to vLLM 0.24.0 (PR #11076) but still pins +# vLLM-Omni 0.23.0rc1, whose vllm_omni/platforms/__init__.py imports +# `supports_xccl` from vllm.utils.torch_utils — a symbol 0.24.0 removed. So +# `import vllm_omni` fails image-wide and this test's mock worker (which +# hard-imports vLLM-Omni) crashes on startup. Skip the whole module up front — +# no point importing vLLM-Omni when we know it can't load. +# TODO: remove this skip once vLLM-Omni is bumped to a vLLM-0.24-compatible release. +pytest.skip( + "vLLM-Omni 0.23.0rc1 is incompatible with the image's vLLM 0.24.0 " + "(missing supports_xccl); re-enable when the vLLM-Omni pin is realigned.", + allow_module_level=True, +) + pytestmark = [ pytest.mark.pre_merge, pytest.mark.integration, From 364cc8aa543d97f0563e17de3069d98ea051f33f Mon Sep 17 00:00:00 2001 From: Yan Ru Pei Date: Tue, 30 Jun 2026 22:23:06 -0700 Subject: [PATCH 021/320] perf(kv-router): skip shape lock for internal CRTC nodes (#11102) Signed-off-by: PeaBrane --- .../matches.rs | 3 + .../concurrent_radix_tree_compressed/node.rs | 17 ++ .../concurrent_radix_tree_compressed/store.rs | 12 ++ .../concurrent_radix_tree_compressed/tests.rs | 152 ++++++++++++++++++ .../concurrent_radix_tree_compressed/types.rs | 2 + 5 files changed, 186 insertions(+) diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/matches.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/matches.rs index 057c6fc04c8d..2459874c7648 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/matches.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/matches.rs @@ -121,6 +121,9 @@ impl ConcurrentRadixTreeCompressed { } } + // NOTE(perf): Pre-reserving the output maps in these survivor-recording + // helpers did not produce a repeatable throughput improvement. Re-profile + // before adding eager capacity here. fn record_surviving_details(details: &mut MatchDetails, walk_result: &MatchWalkResult) { for worker in &walk_result.active { details diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/node.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/node.rs index 604eceaedb3b..52c0fc8113e8 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/node.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/node.rs @@ -34,6 +34,8 @@ fn record_last_matched_hash( #[derive(Debug)] pub(super) struct Node { shape_gate: RwLock<()>, + /// NOTE(concurrency): This is a post-commit validation token, not a seqlock. + /// Node state and children do not share one immutable publication boundary. shape_version: AtomicU64, /// Sticky logical-internal marker. Once true, this node is treated as /// internal even if cleanup removes all physical children later. @@ -103,6 +105,9 @@ impl Node { children: FxHashMap, ) -> Self { let internal = !children.is_empty(); + // NOTE(perf): Reducing child-map sharding substantially lowered memory + // usage but regressed throughput. Treat custom sharding as an explicit + // memory tradeoff rather than a throughput optimization. let children_map = DashMap::with_hasher(FxBuildHasher); for (key, child) in children { children_map.insert(key, child); @@ -121,6 +126,9 @@ impl Node { &self, plan: impl FnOnce(&NodeState, &NodeChildren, u64) -> R, ) -> Option { + // NOTE(perf): Replacing these shape-gated reads with state-only snapshots + // was neutral or regressive, and profiling did not identify the RwLock + // as a hotspot. Re-profile before removing this shape read. let _gate = self.shape_gate.read(); let shape_version = self.shape_version.load(Ordering::Acquire); let state = self.state.read(); @@ -404,6 +412,8 @@ impl Node { blocks: &[KvCacheStoredBlockData], ) -> ParentEdgeAction { match plan.action { + // NOTE(perf): Removing this validation did not produce a repeatable + // benefit and regressed scaled cumulative workloads. ParentEdgePlanAction::InsertFromParent => self .validate_shape_read(plan.shape_version, || { ParentEdgeAction::InsertFromParent(None) @@ -416,6 +426,9 @@ impl Node { } }) .unwrap_or(ParentEdgeAction::Stale), + // NOTE(perf): An additional sticky-internal rejection before this + // commit did not improve throughput. The check inside the gate + // closes the split race. ParentEdgePlanAction::ReuseSuffixAndExtendLeaf { append_start } => self .apply_edge_shape_update(plan.shape_version, |state, _children| { if !self.internal.load(Ordering::Acquire) { @@ -515,6 +528,10 @@ impl Node { blocks: &[KvCacheStoredBlockData], shape_version: u64, ) -> Option { + if self.internal.load(Ordering::Acquire) { + return Some(false); + } + self.apply_edge_shape_update(shape_version, |state, _children| { if self.internal.load(Ordering::Acquire) || blocks.is_empty() || state.edge.is_empty() { return (false, false); diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/store.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/store.rs index 526af04948a7..aa8906fca43c 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/store.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/store.rs @@ -572,4 +572,16 @@ impl ConcurrentRadixTreeCompressed { Ok(StoreInsertOutcome { duplicate_store }) } + + #[cfg(test)] + pub(super) fn insert_blocks_from_for_test( + &self, + lookup: &mut FxHashMap, + worker: WorkerWithDpRank, + parent: &SharedNode, + seed_hash: ExternalSequenceBlockHash, + blocks: &[KvCacheStoredBlockData], + ) -> Result { + self.insert_blocks_from(lookup, worker, parent, false, Some(seed_hash), blocks) + } } diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs index c276907aaa77..20622dc018d4 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs @@ -20,6 +20,13 @@ fn direct_lookup() -> DirectLookup { FxHashMap::default() } +fn stored_data(event: RouterEvent) -> KvCacheStoreData { + match event.event.data { + KvCacheEventData::Stored(op) => op, + _ => unreachable!("expected a store event"), + } +} + fn worker_lookup_len(lookup: &DirectLookup, worker: WorkerWithDpRank) -> Option { lookup.get(&worker).map(|worker_lookup| worker_lookup.len()) } @@ -207,6 +214,151 @@ mod race_tests { assert_direct_score(&index, &[1, 2, 3, 7, 8], worker1, 3); assert_direct_score(&index, &[1, 2, 3, 7, 8], worker2, 5); } + + #[test] + fn stale_scan_cannot_commit_after_split() { + let index = ConcurrentRadixTreeCompressed::new(); + let worker1 = worker(1); + let worker2 = worker(2); + let worker3 = worker(3); + let mut lookup1 = direct_lookup(); + let mut lookup2 = direct_lookup(); + let mut lookup3 = direct_lookup(); + + apply_direct( + &index, + &mut lookup1, + make_store_event(1, &[1, 2, 3, 4, 5, 6]), + ); + apply_direct( + &index, + &mut lookup2, + make_store_event(2, &[1, 2, 3, 4, 5, 6]), + ); + + let node = index + .root + .child_snapshot(LocalBlockHash(1)) + .expect("root child should exist"); + let blocks = stored_data(make_store_event(3, &[1, 2, 3, 4, 5, 6])).blocks; + let stale_scan = node.scan_store_prefix(&blocks); + + apply_direct( + &index, + &mut lookup2, + make_store_event_with_parent(2, &[1, 2, 3], &[7]), + ); + + assert!( + node.promote_to_full_with_version(worker3, stale_scan.shape_version) + .is_none() + ); + + apply_direct( + &index, + &mut lookup3, + make_store_event(3, &[1, 2, 3, 4, 5, 6]), + ); + + assert_eq!( + index.edge_topology_for_test(), + vec![edge_topology( + &[1, 2, 3], + vec![ + edge_topology(&[4, 5, 6], vec![]), + edge_topology(&[7], vec![]) + ], + )], + ); + assert_direct_score(&index, &[1, 2, 3, 4, 5, 6], worker1, 6); + assert_direct_score(&index, &[1, 2, 3, 4, 5, 6], worker2, 6); + assert_direct_score(&index, &[1, 2, 3, 4, 5, 6], worker3, 6); + } + + #[test] + fn tail_parent_split_before_child_lookup_repairs_to_suffix() { + let index = ConcurrentRadixTreeCompressed::new(); + let worker1 = worker(1); + let worker2 = worker(2); + let worker3 = worker(3); + let worker4 = worker(4); + let mut lookup1 = direct_lookup(); + let mut lookup2 = direct_lookup(); + let mut lookup3 = direct_lookup(); + let mut lookup4 = direct_lookup(); + + for (worker_id, lookup) in [ + (1, &mut lookup1), + (2, &mut lookup2), + (3, &mut lookup3), + (4, &mut lookup4), + ] { + apply_direct(&index, lookup, make_store_event(worker_id, &[1, 2, 3, 4])); + } + apply_direct( + &index, + &mut lookup1, + make_store_event_with_parent(1, &[1, 2, 3, 4], &[5, 6]), + ); + apply_direct( + &index, + &mut lookup2, + make_store_event_with_parent(2, &[1, 2, 3, 4], &[7, 8]), + ); + + let stale_parent = index + .root + .child_snapshot(LocalBlockHash(1)) + .expect("root child should exist"); + let continuation = stored_data(make_store_event_with_parent(3, &[1, 2, 3, 4], &[9])); + let parent_hash = continuation.parent_hash.expect("continuation has a parent"); + let plan = stale_parent + .plan_store_parent_edge(parent_hash, &continuation.blocks) + .expect("tail parent should be present before the split"); + assert!(matches!( + plan.action, + ParentEdgePlanAction::InsertFromParent + )); + + apply_direct( + &index, + &mut lookup4, + make_store_event_with_parent(4, &[1, 2], &[10]), + ); + + index + .insert_blocks_from_for_test( + &mut lookup3, + worker3, + &stale_parent, + parent_hash, + &continuation.blocks, + ) + .unwrap(); + + assert_eq!( + index.edge_topology_for_test(), + vec![edge_topology( + &[1, 2], + vec![ + edge_topology( + &[3, 4], + vec![ + edge_topology(&[5, 6], vec![]), + edge_topology(&[7, 8], vec![]), + edge_topology(&[9], vec![]), + ], + ), + edge_topology(&[10], vec![]), + ], + )], + ); + assert_direct_score(&index, &[1, 2, 3, 4, 5, 6], worker1, 6); + assert_direct_score(&index, &[1, 2, 3, 4, 7, 8], worker2, 6); + assert_direct_score(&index, &[1, 2, 3, 4, 9], worker3, 5); + assert_direct_score(&index, &[1, 2, 9], worker3, 2); + assert_direct_score(&index, &[1, 2, 10], worker4, 3); + } } mod remove { diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/types.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/types.rs index c453298bffaa..f975ddb6bf7c 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/types.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/types.rs @@ -19,6 +19,8 @@ pub(super) type SharedNode = Arc; pub(super) type WorkerLookup = FxHashMap; pub(super) struct MatchWalkResult { + // NOTE(perf): Replacing this set with a Vec did not improve throughput. Keep + // uniqueness by construction unless a new profile justifies changing it. pub(super) active: FxHashSet, pub(super) matched_depth: u32, pub(super) prev_edge_last_hash: Option, From 5245c0fa6a2b0e11456e7a990836756233d8b011 Mon Sep 17 00:00:00 2001 From: Tushar Sharma Date: Tue, 30 Jun 2026 23:28:45 -0700 Subject: [PATCH 022/320] fix(test): stop etcd_ha frontend from killing all python processes (#10890) Co-authored-by: Claude Opus 4.8 (1M context) --- tests/fault_tolerance/etcd_ha/utils.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/fault_tolerance/etcd_ha/utils.py b/tests/fault_tolerance/etcd_ha/utils.py index 547c4493d0b6..c4d0fc00247a 100644 --- a/tests/fault_tolerance/etcd_ha/utils.py +++ b/tests/fault_tolerance/etcd_ha/utils.py @@ -37,16 +37,21 @@ def __init__( "DYN_LOG": "debug", "ETCD_ENDPOINTS": ",".join(etcd_endpoints), } - # WARNING: terminate_all_matching_process_names=True is NOT pytest-xdist safe! - # DANGER: Kills ALL dynamo-frontend processes system-wide, including other parallel tests. - # For parallel-safe alternative, use terminate_all_matching_process_names=False. - # See tests/kvbm_integration/common.py:llm_server_kvbm for example. - # TODO: Switch to terminate_all_matching_process_names=False with dynamic ports + # terminate_all_matching_process_names=False is required here: the frontend + # is launched as `python -m dynamo.frontend`, so its _command_name is the + # generic "python". With True, ManagedProcess.__enter__ would SIGTERM/SIGKILL + # every "python" process on the host on startup — including this test's own + # etcd cluster and vLLM worker (and sibling framework jobs). That orphans + # etcd ("connection refused"), so the HA failover the test waits for never + # completes and the test hangs until the 600s pytest-timeout fires. + # Each test manages its frontend via a `with` block (cleaned up on exit) and + # runs in its own container on the fixed FRONTEND_PORT, so the system-wide + # name-based sweep is both unnecessary and unsafe here. super().__init__( request, router_mode="round-robin", extra_env=extra_env, - terminate_all_matching_process_names=True, # TODO: Change to False + terminate_all_matching_process_names=False, ) From 076ce9d82e2f31a5b5e1d9d29b41041b27b3a128 Mon Sep 17 00:00:00 2001 From: MatejKosec Date: Wed, 1 Jul 2026 11:22:26 +0200 Subject: [PATCH 023/320] fix(toolcalling): synthesize tool_calls finish_reason when stream lacks one (#11045) Signed-off-by: Matej Kosec --- .../src/protocols/openai/chat_completions.rs | 46 ++- .../protocols/openai/chat_completions/jail.rs | 382 +++++++++++++++++- .../openai/chat_completions/tool_parser_v2.rs | 304 ++++++++++++-- .../openai/responses/stream_converter.rs | 23 ++ lib/llm/tests/common/http_harness.rs | 1 + lib/llm/tests/responses_http_replay.rs | 6 +- lib/llm/tests/test_jail.rs | 179 ++++---- 7 files changed, 791 insertions(+), 150 deletions(-) diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index e07af11084a5..e021ca8b9759 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; -use dynamo_runtime::protocols::annotated::AnnotationsProvider; +use dynamo_runtime::protocols::annotated::{Annotated, AnnotationsProvider}; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; use validator::Validate; @@ -30,8 +30,9 @@ pub use delta::DeltaGenerator; use dynamo_parsers::tool_calling::{ToolCallResponse, ToolCallResponseChunk}; use dynamo_protocols::types::{ - ChatCompletionMessageToolCall, ChatCompletionMessageToolCallChunk, FunctionCall, - FunctionCallStream, FunctionType, + ChatChoiceStream, ChatCompletionMessageContent, ChatCompletionMessageToolCall, + ChatCompletionMessageToolCallChunk, ChatCompletionStreamResponseDelta, FinishReason, + FunctionCall, FunctionCallStream, FunctionType, }; /// Map a parser-native [`ToolCallResponse`] onto the protocol/wire @@ -239,6 +240,45 @@ pub struct NvCreateChatCompletionStreamResponse { pub llm_metrics: Option, } +/// Build one synthetic stream choice from an existing response template. +/// +/// Both streaming tool-call paths use this constructor when an engine omits a +/// terminal choice. Accounting data belongs only on the usage chunk and must +/// not be copied onto the synthetic choice. +pub(super) fn stream_choice_chunk_from_template( + template: &NvCreateChatCompletionStreamResponse, + index: u32, + content: Option, + tool_calls: Option>, + finish_reason: Option, +) -> Annotated { + let mut response = template.clone(); + response.inner.usage = None; + response.llm_metrics = None; + #[allow(deprecated)] + let choice = ChatChoiceStream { + index, + delta: ChatCompletionStreamResponseDelta { + role: None, + content, + tool_calls, + function_call: None, + refusal: None, + reasoning_content: None, + }, + finish_reason, + logprobs: None, + }; + response.inner.choices = vec![choice]; + Annotated { + data: Some(response), + id: None, + event: None, + comment: None, + error: None, + } +} + /// Implements `NvExtProvider` for `NvCreateChatCompletionRequest`, /// providing access to NVIDIA-specific extensions. impl NvExtProvider for NvCreateChatCompletionRequest { diff --git a/lib/llm/src/protocols/openai/chat_completions/jail.rs b/lib/llm/src/protocols/openai/chat_completions/jail.rs index 543c19c694c7..1a6aa1697d31 100644 --- a/lib/llm/src/protocols/openai/chat_completions/jail.rs +++ b/lib/llm/src/protocols/openai/chat_completions/jail.rs @@ -17,12 +17,12 @@ use dynamo_parsers::tool_calling::{ }; use dynamo_runtime::protocols::annotated::Annotated; use futures::{Stream, StreamExt}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use uuid::Uuid; use crate::utils::{MarkerMatcher, MatchResult}; -use super::NvCreateChatCompletionStreamResponse; +use super::{NvCreateChatCompletionStreamResponse, stream_choice_chunk_from_template}; fn is_harmony_parser(parser: Option<&str>) -> bool { parser == Some("harmony") @@ -1461,17 +1461,37 @@ impl JailedStream { where S: Stream> + Send + 'static, { + let _ = named_tool_active; + let _ = &jail_mode; stream! { tokio::pin!(input_stream); let mut has_tool_calls_per_choice: HashMap = HashMap::new(); + // Choices that already received a finish_reason during the stream — used by + // the backstop below to avoid synthesizing a duplicate. + let mut terminated: HashSet = HashSet::new(); + // Last response, kept (with choices cleared) as a template for a synthesized + // finish_reason chunk when the stream ended without one. + let mut template: Option = None; + // Choices for which this post-processor has already emitted a synthetic + // terminal chunk. Tracking this per choice allows a later tool-call choice + // to terminate even if an earlier empty-choices chunk emitted nothing. + let mut synthesized: HashSet = HashSet::new(); while let Some(mut response) = input_stream.next().await { - // Track if any choice emitted tool calls + // Track if any choice emitted tool calls, and which already terminated. if let Some(ref data) = response.data { for choice in &data.inner.choices { if choice.delta.tool_calls.is_some() { has_tool_calls_per_choice.insert(choice.index, true); } + if choice.finish_reason.is_some() { + terminated.insert(choice.index); + } + } + { + let mut t = data.clone(); + t.inner.choices.clear(); + template = Some(t); } } @@ -1487,7 +1507,6 @@ impl JailedStream { // choice, finish_reason MUST be "tool_calls" — regardless of // whether tool_choice was "auto", "required", or a named // function. - let _ = named_tool_active; match &jail_mode { JailMode::MarkerBased => { if has_tool_calls { @@ -1506,8 +1525,68 @@ impl JailedStream { } } + // OpenAI stream ordering: the terminal finish_reason chunk must precede + // the usage-only chunk. When a chunk with no choices arrives (the + // frontend's compliance usage chunk, or any other empty-choices chunk) + // and tool-call choices are still missing a finish_reason, synthesize + // their terminal `ToolCalls` chunks *before* yielding this one. + let is_empty_choices = response + .data + .as_ref() + .is_some_and(|d| d.inner.choices.is_empty()); + if is_empty_choices && let Some(template) = &template { + let mut indices: Vec<_> = has_tool_calls_per_choice + .iter() + .filter_map(|(index, has)| { + (*has && !terminated.contains(index) && !synthesized.contains(index)) + .then_some(*index) + }) + .collect(); + indices.sort_unstable(); + for index in indices { + yield stream_choice_chunk_from_template( + template, + index, + None, + None, + Some(FinishReason::ToolCalls), + ); + synthesized.insert(index); + } + } + yield response; } + + // Backstop: the stream ended without a finish_reason AND without an + // empty-choices/usage chunk to anchor the synthesized terminal chunks + // before (e.g. the engine dropped the terminal signal and the frontend + // never emitted a usage chunk). Emit one trailing `ToolCalls` chunk per + // tool-call choice that never received a finish_reason. Strict OpenAI + // clients wait for a non-null finish_reason before considering a tool call + // complete; without this they hang until their client-side timeout. + // Choices that never emitted tool calls are left alone — there + // is no signal to invent a finish_reason from for text-only output. + if let Some(template) = template { + let mut indices: Vec<_> = has_tool_calls_per_choice + .iter() + .filter_map(|(index, has)| { + (*has && !terminated.contains(index) && !synthesized.contains(index)) + .then_some(*index) + }) + .collect(); + indices.sort_unstable(); + for index in indices { + yield stream_choice_chunk_from_template( + &template, + index, + None, + None, + Some(FinishReason::ToolCalls), + ); + synthesized.insert(index); + } + } } } } @@ -1775,6 +1854,96 @@ mod tests { } } + /// A usage-only chunk (empty `choices`, a usage object) — the frontend's + /// OpenAI-compliance terminal chunk. Used to test that the synthesized + /// `ToolCalls` finish_reason chunk is emitted *before* this one. + fn usage_only_chunk() -> Annotated { + Annotated { + data: Some(NvCreateChatCompletionStreamResponse { + inner: CreateChatCompletionStreamResponse { + id: "id-42".to_string(), + object: "chat.completion.chunk".to_string(), + created: 0, + model: "test-model".to_string(), + choices: vec![], + usage: Some(dynamo_protocols::types::CompletionUsage { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + prompt_tokens_details: None, + completion_tokens_details: None, + }), + service_tier: None, + system_fingerprint: None, + }, + nvext: None, + llm_metrics: Some(crate::protocols::common::metrics::LLMMetricAnnotation { + input_tokens: 10, + output_tokens: 5, + chunk_tokens: 0, + cached_tokens: None, + prefill_worker_id: None, + prefill_dp_rank: None, + prefill_worker_type: None, + decode_worker_id: None, + decode_dp_rank: None, + decode_worker_type: None, + tokenize_latency: None, + detokenize_total_latency: None, + detokenize_count: None, + }), + }), + id: None, + event: None, + comment: None, + error: None, + } + } + + /// Build one data chunk whose choices have already emitted tool-call deltas. + fn tool_call_choices_chunk(indices: &[u32]) -> Annotated { + let mut chunk = text_chunk(""); + let data = chunk.data.as_mut().expect("tool-call response data"); + #[allow(deprecated)] + { + data.inner.choices = indices + .iter() + .map(|index| ChatChoiceStream { + index: *index, + delta: ChatCompletionStreamResponseDelta { + role: Some(Role::Assistant), + content: None, + tool_calls: Some(vec![ChatCompletionMessageToolCallChunk { + index: 0, + id: Some(format!("call-{index}")), + r#type: Some(FunctionType::Function), + function: Some(FunctionCallStream { + name: Some(format!("tool_{index}")), + arguments: Some("{}".to_string()), + }), + }]), + function_call: None, + refusal: None, + reasoning_content: None, + }, + finish_reason: None, + logprobs: None, + }) + .collect(); + } + chunk + } + + fn heartbeat() -> Annotated { + Annotated { + data: None, + id: None, + event: None, + comment: Some(vec!["heartbeat".to_string()]), + error: None, + } + } + /// Collect all emitted tool calls from the jailed stream output fn collect_tool_calls( responses: &[Annotated], @@ -2079,4 +2248,209 @@ mod tests { all_text ); } + + /// The last `finish_reason` a client sees across the stream. `None` if the + /// stream never carried one (the missing-finish-reason hang condition). + fn final_finish_reason( + responses: &[Annotated], + ) -> Option { + responses + .iter() + .filter_map(|r| r.data.as_ref()) + .flat_map(|d| d.inner.choices.iter()) + .filter_map(|c| c.finish_reason) + .next_back() + } + + // Missing-finish-reason regression: when the engine emits a complete tool call + // but the stream ends without any finish_reason chunk (speculative decoding + // folded EOS into content, or the terminal signal was dropped), a strict + // OpenAI client waits for a non-null finish_reason and hangs until its timeout. + // The jail path's finalize() emits the tool call with the absent upstream + // finish_reason; fix_finish_reason's end-of-stream path must synthesize + // `ToolCalls` so the client gets a terminal signal. + #[tokio::test] + async fn jail_synthesizes_tool_calls_finish_reason_when_stream_lacks_one() { + let jail = JailedStream::builder().tool_call_parser("hermes").build(); + + let chunks = vec![text_chunk( + "\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"SF\"}}\n", + )]; + + let input_stream = Box::pin(stream::iter(chunks)); + let output_stream = jail.apply_with_finish_reason(input_stream); + + let responses: Vec<_> = output_stream.collect().await; + let tool_calls = collect_tool_calls(&responses); + assert!( + !tool_calls.is_empty(), + "expected the hermes tool call to be parsed: {tool_calls:?}" + ); + assert_eq!(tool_calls[0].0, "get_weather"); + assert_eq!( + final_finish_reason(&responses), + Some(FinishReason::ToolCalls), + "backstop must synthesize ToolCalls when the stream ended without a finish_reason" + ); + } + + // Text-only corollary: a text-only stream that ends without a finish_reason + // must not get a synthesized one. There is no signal to invent a + // finish_reason from when no tool call was emitted. + #[tokio::test] + async fn jail_does_not_synthesize_finish_reason_for_text_only_stream() { + let jail = JailedStream::builder().tool_call_parser("hermes").build(); + + let chunks = vec![text_chunk("hello world"), text_chunk("")]; + + let input_stream = Box::pin(stream::iter(chunks)); + let output_stream = jail.apply_with_finish_reason(input_stream); + + let responses: Vec<_> = output_stream.collect().await; + let tool_calls = collect_tool_calls(&responses); + assert!( + tool_calls.is_empty(), + "no tool calls expected: {tool_calls:?}" + ); + assert_eq!( + final_finish_reason(&responses), + None, + "text-only stream with no upstream finish_reason must not get a synthetic one" + ); + } + + // Usage-ordering regression: a tool call is followed by a usage-only chunk, + // with no finish_reason chunk from the engine. The + // synthesized `ToolCalls` terminal chunk must be emitted *before* the + // usage-only chunk (OpenAI stream ordering — the terminal finish_reason + // precedes usage). This mirrors the production stream ordering. + #[tokio::test] + async fn jail_synthesizes_tool_calls_before_usage_only_chunk() { + let jail = JailedStream::builder().tool_call_parser("hermes").build(); + + let chunks = vec![ + heartbeat(), + text_chunk( + "\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"SF\"}}\n", + ), + usage_only_chunk(), + ]; + + let input_stream = Box::pin(stream::iter(chunks)); + let output_stream = jail.apply_with_finish_reason(input_stream); + + let responses: Vec<_> = output_stream.collect().await; + assert_eq!( + responses + .first() + .and_then(|response| response.comment.clone()), + Some(vec!["heartbeat".to_string()]), + "leading non-data annotations must pass through unchanged" + ); + let tool_calls = collect_tool_calls(&responses); + assert!( + !tool_calls.is_empty(), + "expected the hermes tool call: {tool_calls:?}" + ); + assert_eq!(tool_calls[0].0, "get_weather"); + assert_eq!( + final_finish_reason(&responses), + Some(FinishReason::ToolCalls), + "a synthesized ToolCalls terminal chunk must be present" + ); + + // The ToolCalls terminal chunk must precede the usage-only chunk. + let finish_pos = responses.iter().position(|r| { + r.data.as_ref().is_some_and(|d| { + d.inner + .choices + .iter() + .any(|c| c.finish_reason == Some(FinishReason::ToolCalls)) + }) + }); + let usage_pos = responses.iter().position(|r| { + r.data + .as_ref() + .is_some_and(|d| d.inner.usage.is_some() && d.inner.choices.is_empty()) + }); + let finish_pos = finish_pos.expect("no ToolCalls chunk emitted"); + let usage_pos = usage_pos.expect("no usage-only chunk in output"); + assert!( + finish_pos < usage_pos, + "ToolCalls chunk at {finish_pos} must precede the usage chunk at {usage_pos}" + ); + let finish_data = responses[finish_pos] + .data + .as_ref() + .expect("ToolCalls chunk has no response data"); + assert!( + finish_data.inner.usage.is_none(), + "synthesized ToolCalls chunk must not repeat usage data" + ); + assert!( + finish_data.llm_metrics.is_none(), + "synthesized ToolCalls chunk must not repeat LLM metrics" + ); + } + + // An empty-choices chunk can precede tool deltas (for example, a metadata + // response). It must not disable later synthesis. When several choices then + // emit tool calls, their terminal chunks must be ordered by choice index. + #[tokio::test] + async fn jail_synthesizes_late_tool_choices_in_index_order() { + let chunks = vec![ + usage_only_chunk(), + tool_call_choices_chunk(&[2, 0, 1]), + usage_only_chunk(), + ]; + + let responses: Vec<_> = + JailedStream::fix_finish_reason(stream::iter(chunks), JailMode::MarkerBased, false) + .collect() + .await; + + let usage_positions: Vec<_> = responses + .iter() + .enumerate() + .filter_map(|(position, response)| { + response + .data + .as_ref() + .is_some_and(|data| data.inner.choices.is_empty() && data.inner.usage.is_some()) + .then_some(position) + }) + .collect(); + assert_eq!( + usage_positions.len(), + 2, + "both empty-choices chunks must pass through" + ); + + let terminals: Vec<_> = responses + .iter() + .enumerate() + .flat_map(|(position, response)| { + response.data.iter().flat_map(move |data| { + data.inner.choices.iter().filter_map(move |choice| { + (choice.finish_reason == Some(FinishReason::ToolCalls)) + .then_some((position, choice.index)) + }) + }) + }) + .collect(); + assert_eq!( + terminals + .iter() + .map(|(_, index)| *index) + .collect::>(), + vec![0, 1, 2], + "synthetic terminal chunks must be deterministic" + ); + assert!( + terminals.iter().all(|(position, _)| { + usage_positions[0] < *position && *position < usage_positions[1] + }), + "terminal chunks must follow the early empty response and precede the final usage response" + ); + } } diff --git a/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs b/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs index 9ce7c12a02ff..4cdf482b8e50 100644 --- a/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs +++ b/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs @@ -33,7 +33,7 @@ use dynamo_parsers::tool_calling::{ }; use dynamo_parsers_v2::{Tool as ToolV2, ToolCallDelta, ToolParser, create_tool_parser_for_family}; -use super::NvCreateChatCompletionStreamResponse; +use super::{NvCreateChatCompletionStreamResponse, stream_choice_chunk_from_template}; /// Tool-call families with a `dynamo-parsers-v2` parser wired into both the batch and /// the streaming path. Must stay a subset of the families @@ -163,6 +163,62 @@ impl ChoiceState { } } +/// Finish every choice that has not received an upstream finish reason. This is +/// called before a usage-only chunk when one exists, with EOF as a fallback. +fn finish_unterminated_choices( + states: &mut HashMap, + finished: &mut HashSet, + tool_emitted: &mut HashSet, + template: &NvCreateChatCompletionStreamResponse, +) -> Vec> { + let mut indices: Vec<_> = states + .keys() + .copied() + .filter(|index| !finished.contains(index)) + .collect(); + indices.sort_unstable(); + + let mut responses = Vec::new(); + for index in indices { + finished.insert(index); + let state = states + .get_mut(&index) + .expect("choice index came from parser state map"); + let result = match state.parser.finish() { + Ok(result) => result, + Err(error) => { + tracing::warn!(error = %error, choice_index = index, "v2 stream finish failed"); + dynamo_parsers_v2::ToolParseResult::default() + } + }; + let tool_calls = state.emit_chunks(result.calls); + if tool_calls.is_some() { + tool_emitted.insert(index); + } + // A choice that produced tool calls during the stream must terminate + // with `ToolCalls` even when the backend never sent a finish_reason. + // Text-only output without an upstream finish reason stays `None`. + let finish_reason = if tool_emitted.contains(&index) { + Some(FinishReason::ToolCalls) + } else { + None + }; + let content = (!result.normal_text.is_empty()) + .then_some(ChatCompletionMessageContent::Text(result.normal_text)); + if content.is_none() && tool_calls.is_none() && finish_reason.is_none() { + continue; + } + responses.push(stream_choice_chunk_from_template( + template, + index, + content, + tool_calls, + finish_reason, + )); + } + responses +} + /// Streaming path: replace the jail with the `family` v2 parser. Each upstream text /// delta is pushed into the parser; the parser's `normal_text` becomes the emitted /// content and its tool-call deltas become OpenAI tool-call chunks. The jail is never @@ -213,6 +269,7 @@ where t.inner.choices.clear(); template = Some(t); } + let is_empty_choices = chat_response.inner.choices.is_empty(); for choice in chat_response.inner.choices.iter_mut() { let state = states.entry(choice.index).or_insert_with(|| { @@ -282,47 +339,36 @@ where } } + // OpenAI stream ordering requires a terminal finish_reason before the + // usage-only chunk. Finish every unterminated choice before yielding an + // empty-choices response; EOF below remains the fallback when no such + // response arrives. + if is_empty_choices && let Some(template) = &template { + for terminal in finish_unterminated_choices( + &mut states, + &mut finished, + &mut tool_emitted, + template, + ) { + yield terminal; + } + } + yield response; } // Backstop: the stream ended without a finish_reason for some choice. Flush - // each unfinished parser; emit a trailing chunk only when it yields output. - if let Some(template) = template { - for (index, state) in states.iter_mut() { - if finished.contains(index) { - continue; - } - let Ok(result) = state.parser.finish() else { - continue; - }; - let tool_calls = state.emit_chunks(result.calls); - if result.normal_text.is_empty() && tool_calls.is_none() { - continue; - } - let mut response = template.clone(); - #[allow(deprecated)] - let choice = dynamo_protocols::types::ChatChoiceStream { - index: *index, - delta: dynamo_protocols::types::ChatCompletionStreamResponseDelta { - role: None, - content: (!result.normal_text.is_empty()) - .then_some(ChatCompletionMessageContent::Text(result.normal_text)), - tool_calls, - function_call: None, - refusal: None, - reasoning_content: None, - }, - finish_reason: None, - logprobs: None, - }; - response.inner.choices = vec![choice]; - yield Annotated { - data: Some(response), - id: None, - event: None, - comment: None, - error: None, - }; + // each unfinished parser; emit a trailing chunk when the flush yields output + // or when the choice already emitted tool calls and still needs a terminal + // `ToolCalls` reason. + if let Some(template) = &template { + for terminal in finish_unterminated_choices( + &mut states, + &mut finished, + &mut tool_emitted, + template, + ) { + yield terminal; } } } @@ -332,10 +378,26 @@ where mod tests { use super::*; use dynamo_protocols::types::{ - ChatChoiceStream, ChatCompletionStreamResponseDelta, FinishReason, Role, + ChatChoiceStream, ChatCompletionStreamResponseDelta, CompletionUsage, FinishReason, Role, }; use futures::stream; + struct FinishErrorParser; + + impl ToolParser for FinishErrorParser { + fn create(_tools: &[ToolV2]) -> anyhow::Result> { + Ok(Box::new(Self)) + } + + fn push(&mut self, _chunk: &str) -> anyhow::Result { + Ok(dynamo_parsers_v2::ToolParseResult::default()) + } + + fn finish(&mut self) -> anyhow::Result { + anyhow::bail!("intentional finish failure") + } + } + const QWEN3_GET_WEATHER: &str = "\n\n\nParis\n\n\n"; // DeepSeek-V4 DSML: one get_weather(location="NYC") call. The `｜` glyphs are the @@ -379,6 +441,35 @@ mod tests { } } + fn usage_chunk() -> Annotated { + let mut chunk = chunk("", false); + let data = chunk.data.as_mut().expect("usage chunk response data"); + data.inner.choices.clear(); + data.inner.usage = Some(CompletionUsage { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + prompt_tokens_details: None, + completion_tokens_details: None, + }); + data.llm_metrics = Some(crate::protocols::common::metrics::LLMMetricAnnotation { + input_tokens: 10, + output_tokens: 5, + chunk_tokens: 0, + cached_tokens: None, + prefill_worker_id: None, + prefill_dp_rank: None, + prefill_worker_type: None, + decode_worker_id: None, + decode_dp_rank: None, + decode_worker_type: None, + tokenize_latency: None, + detokenize_total_latency: None, + detokenize_count: None, + }); + chunk + } + /// Reassemble the streamed tool-call deltas into (name, arguments) per index and /// collect all emitted content, mirroring how an OpenAI client reconstructs a /// streamed tool call. @@ -535,4 +626,139 @@ mod tests { "finish_reason must flip Stop->ToolCalls when tool calls are emitted" ); } + + // Missing-finish-reason regression: the stream emits a complete tool call but + // ends without any finish_reason chunk (e.g. speculative decoding folded EOS + // into content, or the engine dropped the terminal signal). A strict OpenAI + // client waits for a non-null finish_reason before considering the tool call + // complete; the end-of-stream backstop must synthesize `ToolCalls` so the + // client doesn't hang until its timeout. + #[tokio::test] + async fn qwen3_bypass_synthesizes_tool_calls_when_stream_lacks_finish_reason() { + // Same call as the incremental test, but the final chunk carries NO + // finish_reason — the stream simply ends after the tool markup. + let mut chunks: Vec<_> = QWEN3_GET_WEATHER + .as_bytes() + .chunks(8) + .map(|b| chunk(std::str::from_utf8(b).unwrap(), false)) + .collect(); + // A usage-only chunk arrives without any terminating choice. + chunks.push(usage_chunk()); + + let out: Vec<_> = apply_stream(stream::iter(chunks), None, "qwen3_coder".to_string()) + .collect::>() + .await; + + let (calls, _content) = reassemble(&out); + assert_eq!(calls.len(), 1, "expected exactly one tool call: {calls:?}"); + assert_eq!(calls[0].0, "get_weather"); + let args: serde_json::Value = serde_json::from_str(&calls[0].1).unwrap(); + assert_eq!(args["location"], "Paris"); + assert_eq!( + final_finish_reason(&out), + Some(FinishReason::ToolCalls), + "backstop must synthesize ToolCalls when the stream ended without a finish_reason" + ); + let finish_positions: Vec<_> = out + .iter() + .enumerate() + .filter_map(|(position, response)| { + response.data.as_ref().and_then(|data| { + data.inner + .choices + .iter() + .any(|choice| choice.finish_reason == Some(FinishReason::ToolCalls)) + .then_some(position) + }) + }) + .collect(); + assert_eq!( + finish_positions.len(), + 1, + "expected exactly one synthesized finish chunk" + ); + let usage_position = + out.iter() + .position(|response| { + response.data.as_ref().is_some_and(|data| { + data.inner.choices.is_empty() && data.inner.usage.is_some() + }) + }) + .expect("usage-only response"); + assert!( + finish_positions[0] < usage_position, + "synthesized finish chunk must precede usage" + ); + let terminal = out[finish_positions[0]] + .data + .as_ref() + .expect("synthesized terminal response"); + assert!( + terminal.inner.usage.is_none(), + "synthesized terminal chunk must not repeat usage" + ); + assert!( + terminal.llm_metrics.is_none(), + "synthesized terminal chunk must not repeat LLM metrics" + ); + } + + // Text-only corollary: when the stream ends without a finish_reason and no + // tool call was emitted, the backstop must not invent a finish_reason. There + // is no signal to synthesize one from. A trailing content chunk may be + // emitted, but its finish_reason stays None. + #[tokio::test] + async fn qwen3_bypass_does_not_synthesize_finish_reason_for_text_only_stream() { + let chunks = vec![chunk("hello world", false), chunk("", false)]; + + let out: Vec<_> = apply_stream(stream::iter(chunks), None, "qwen3_coder".to_string()) + .collect::>() + .await; + + let (calls, _content) = reassemble(&out); + assert!(calls.is_empty(), "no tool calls expected: {calls:?}"); + assert_eq!( + final_finish_reason(&out), + None, + "text-only stream with no upstream finish_reason must not get a synthetic one" + ); + } + + #[test] + fn finish_error_still_terminates_a_choice_that_emitted_tools() { + let mut states = HashMap::from([( + 3, + ChoiceState { + parser: Box::new(FinishErrorParser), + opened: HashSet::new(), + }, + )]); + let mut finished = HashSet::new(); + let mut tool_emitted = HashSet::from([3]); + let template = usage_chunk().data.expect("usage response data"); + + let responses = + finish_unterminated_choices(&mut states, &mut finished, &mut tool_emitted, &template); + + assert_eq!( + responses.len(), + 1, + "the choice still needs a terminal chunk" + ); + let response = responses[0].data.as_ref().expect("terminal response data"); + assert!( + response.inner.usage.is_none(), + "terminal chunk must not repeat usage" + ); + assert!( + response.llm_metrics.is_none(), + "terminal chunk must not repeat LLM metrics" + ); + assert_eq!(response.inner.choices.len(), 1); + assert_eq!(response.inner.choices[0].index, 3); + assert_eq!( + response.inner.choices[0].finish_reason, + Some(FinishReason::ToolCalls) + ); + } } diff --git a/lib/llm/src/protocols/openai/responses/stream_converter.rs b/lib/llm/src/protocols/openai/responses/stream_converter.rs index 2619009802d7..4a668bd2446a 100644 --- a/lib/llm/src/protocols/openai/responses/stream_converter.rs +++ b/lib/llm/src/protocols/openai/responses/stream_converter.rs @@ -1106,6 +1106,29 @@ mod tests { assert!(end_types.contains(&"response.completed".to_string())); } + #[test] + fn test_function_call_finish_reason_closes_tool_call() { + let mut conv = ResponseStreamConverter::new("test-model".into(), default_params()); + let _ = conv.emit_start_events(); + + let _ = conv.process_chunk(&tool_call_chunk( + 0, + Some("call-1"), + Some("get_weather"), + Some("{\"city\":\"SF\"}"), + )); + + let finish_types = + event_types(&conv.process_chunk(&finish_chunk(FinishReason::FunctionCall))); + assert_eq!( + finish_types, + vec![ + "response.function_call_arguments.done".to_string(), + "response.output_item.done".to_string(), + ] + ); + } + #[test] fn test_identity_only_tool_call_is_emitted_and_finished() { let mut conv = ResponseStreamConverter::new("test-model".into(), default_params()); diff --git a/lib/llm/tests/common/http_harness.rs b/lib/llm/tests/common/http_harness.rs index 9c187b616d49..76f3eb2f4a5c 100644 --- a/lib/llm/tests/common/http_harness.rs +++ b/lib/llm/tests/common/http_harness.rs @@ -139,6 +139,7 @@ pub async fn load_sse_fixture(path: impl AsRef) -> Result