test(vllm): multi-worker embedding routing tests (DIS-2098)

components/src/dynamo/vllm/health_check.py

@@ -113,6 +113,45 @@ def to_dict(self) -> dict[str, Any]:
         return _layer_probe_marker(super().to_dict())
 
 
+class VllmEmbeddingHealthCheckPayload(HealthCheckPayload):
+    """
+    vLLM-specific health check payload for pooling/embedding workers.
+
+    Embedding workers run an ``AsyncLLM`` in pooling mode and serve the
+    ``EmbeddingWorkerHandler.generate`` entry point, which expects the
+    OpenAI ``/v1/embeddings`` request shape -- ``{model, input}`` -- not
+    the token-id/sampling-options shape the chat health check uses.
+    Sending the chat payload through the embedding handler raises
+    "missing required 'input' field" and the canary stays unhealthy
+    forever, so the runtime's ``/health`` never flips to 200 and any
+    caller that waits on it (``health_check_workers=True`` in the test
+    harness, K8s readiness probes, etc.) races against worker startup.
+
+    The probe runs a real pooling forward pass with a short input, so
+    "healthy" actually means "the engine produced an embedding," not
+    just "the process is up." Cost is one small batch every probe
+    cycle.
+    """
+
+    def __init__(self, model_name: Optional[str] = None):
+        """
+        Args:
+            model_name: served model name to put on ``request["model"]``.
+                Optional -- ``EmbeddingWorkerHandler.generate`` falls
+                back to ``config.served_model_name`` when the field is
+                absent. Passing it keeps the probe self-describing in
+                worker logs; omit when the caller doesn't have a
+                specific name to advertise.
+        """
+        self.default_payload: dict[str, Any] = {"input": "probe"}
+        if model_name is not None:
+            self.default_payload["model"] = model_name
+        super().__init__()
+
+    def to_dict(self) -> dict[str, Any]:
+        return _layer_probe_marker(super().to_dict())
+
+
 class VllmPrefillHealthCheckPayload(HealthCheckPayload):
     """
     vLLM-specific health check payload for prefill workers in disaggregated mode.

components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py

@@ -15,6 +15,7 @@
 
 from dynamo.health_check import HEALTH_CHECK_KEY
 from dynamo.vllm.health_check import (
+    VllmEmbeddingHealthCheckPayload,
     VllmHealthCheckPayload,
     VllmOmniHealthCheckPayload,
     VllmPrefillHealthCheckPayload,
@@ -27,20 +28,24 @@
     pytest.mark.pre_merge,
 ]
 
-PAYLOAD_CLASSES = [
-    VllmHealthCheckPayload,
-    VllmPrefillHealthCheckPayload,
-    VllmOmniHealthCheckPayload,
+PAYLOAD_FACTORIES = [
+    pytest.param(VllmHealthCheckPayload, id="VllmHealthCheckPayload"),
+    pytest.param(VllmPrefillHealthCheckPayload, id="VllmPrefillHealthCheckPayload"),
+    pytest.param(VllmOmniHealthCheckPayload, id="VllmOmniHealthCheckPayload"),
+    pytest.param(
+        VllmEmbeddingHealthCheckPayload,
+        id="VllmEmbeddingHealthCheckPayload",
+    ),
 ]
 
 
-@pytest.mark.parametrize("cls", PAYLOAD_CLASSES, ids=lambda c: c.__name__)
-def test_payload_has_marker(cls):
-    assert cls().to_dict()[HEALTH_CHECK_KEY] is True
+@pytest.mark.parametrize("make", PAYLOAD_FACTORIES)
+def test_payload_has_marker(make):
+    assert make().to_dict().get(HEALTH_CHECK_KEY) is True
 
 
-@pytest.mark.parametrize("cls", PAYLOAD_CLASSES, ids=lambda c: c.__name__)
-def test_env_override_preserves_marker(monkeypatch, cls):
+@pytest.mark.parametrize("make", PAYLOAD_FACTORIES)
+def test_env_override_preserves_marker(monkeypatch, make):
     """DYN_HEALTH_CHECK_PAYLOAD must not drop the canary marker."""
     monkeypatch.setenv(
         "DYN_HEALTH_CHECK_PAYLOAD",
@@ -52,4 +57,35 @@ def test_env_override_preserves_marker(monkeypatch, cls):
             }
         ),
     )
-    assert cls().to_dict()[HEALTH_CHECK_KEY] is True
+    assert make().to_dict().get(HEALTH_CHECK_KEY) is True
+
+
+def test_embedding_payload_shape_matches_handler_contract():
+    """``VllmEmbeddingHealthCheckPayload`` must produce the
+    ``{model, input}`` shape that ``EmbeddingWorkerHandler.generate``
+    expects. The chat payload (``token_ids`` + ``sampling_options``)
+    would be rejected by that handler with "missing required 'input'
+    field" and leave the canary stuck unhealthy forever -- regression
+    pin for the bug PR #9765's canary readiness fix targets."""
+    payload = VllmEmbeddingHealthCheckPayload(
+        model_name="Qwen/Qwen3-Embedding-0.6B"
+    ).to_dict()
+    assert payload.get("model") == "Qwen/Qwen3-Embedding-0.6B"
+    assert payload.get("input") == "probe"
+    assert payload.get(HEALTH_CHECK_KEY) is True
+    # Must NOT carry chat-shaped keys -- those would make the embedding
+    # handler's request shape validation fail.
+    assert "token_ids" not in payload
+    assert "sampling_options" not in payload
+    assert "stop_conditions" not in payload
+
+
+def test_embedding_payload_omits_model_when_no_name():
+    """``model_name`` is optional. When omitted, the ``model`` key is
+    absent from the payload and the handler falls back to
+    ``config.served_model_name`` (see ``EmbeddingWorkerHandler.generate``).
+    """
+    payload = VllmEmbeddingHealthCheckPayload().to_dict()
+    assert "model" not in payload
+    assert payload.get("input") == "probe"
+    assert payload.get(HEALTH_CHECK_KEY) is True

components/src/dynamo/vllm/worker_factory.py

@@ -35,7 +35,11 @@
     PrefillWorkerHandler,
     get_dp_range_for_worker,
 )
-from .health_check import VllmHealthCheckPayload, VllmPrefillHealthCheckPayload
+from .health_check import (
+    VllmEmbeddingHealthCheckPayload,
+    VllmHealthCheckPayload,
+    VllmPrefillHealthCheckPayload,
+)
 from .multimodal_handlers import EncodeWorkerHandler
 from .publisher import StatLoggerFactory
 
@@ -282,12 +286,17 @@ async def _create_embedding_worker(
             shutdown_event=shutdown_event,
         )
 
+        embedding_health_check_payload = VllmEmbeddingHealthCheckPayload(
+            model_name=config.served_model_name or config.model
+        ).to_dict()
+
         logger.info("Starting to serve the embedding worker endpoint...")
         try:
             await asyncio.gather(
                 generate_endpoint.serve_endpoint(
                     handler.generate,
                     metrics_labels=[("model", config.model)],
+                    health_check_payload=embedding_health_check_payload,
                 ),
                 self.register_vllm_model(
                     ModelInput.Text,

examples/backends/vllm/launch/agg_embed_multiworker.sh

@@ -0,0 +1,131 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Multi-worker aggregated embedding serving.
+#
+# Spawns 1 frontend + 2 embedding workers, each on its own GPU and its own
+# DYN_SYSTEM_PORT so per-worker metrics can be scraped independently.
+#
+# Used by the multi-worker embedding tests to verify:
+#   1. Same-model load balancing — pass the same MODEL twice; the frontend
+#      should weighted-randomly distribute requests across both workers.
+#   2. Multi-model dispatch       — pass two different MODELs; the
+#      name-keyed router (lib/llm/src/discovery/model_manager.rs) should
+#      send each request only to the worker registered for that model.
+#
+# GPUs: 2
+#
+# Usage:
+#   agg_embed_multiworker.sh MODEL1 MODEL2 [EXTRA_DYNAMO_VLLM_ARGS...]
+#
+# EXTRA args (after the two model positions) are forwarded verbatim to
+# *both* dynamo.vllm worker processes. The current launch matches the
+# single-worker ``agg_embed.sh`` script (``--runner pooling``,
+# ``--dtype float32``, MEAN pooler config, ``--max-model-len 2048``,
+# ``--no-enable-prefix-caching``).
+
+set -e
+trap 'echo Cleaning up...; kill 0' EXIT
+
+SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
+source "$SCRIPT_DIR/../../../common/gpu_utils.sh"   # build_vllm_gpu_mem_args
+source "$SCRIPT_DIR/../../../common/launch_utils.sh" # print_launch_banner, wait_any_exit
+
+if [[ $# -lt 2 ]]; then
+    echo "Usage: $0 MODEL1 MODEL2 [extra dynamo.vllm args...]" >&2
+    echo "" >&2
+    echo "Examples:" >&2
+    echo "  # Same-model load balance test:" >&2
+    echo "  $0 Qwen/Qwen3-Embedding-0.6B Qwen/Qwen3-Embedding-0.6B" >&2
+    echo "" >&2
+    echo "  # Multi-model dispatch test:" >&2
+    echo "  $0 Qwen/Qwen3-Embedding-0.6B BAAI/bge-small-en-v1.5" >&2
+    exit 2
+fi
+
+MODEL1="$1"
+MODEL2="$2"
+shift 2
+EXTRA_ARGS=("$@")
+
+GPU_MEM_ARGS=$(build_vllm_gpu_mem_args)
+
+HTTP_PORT="${DYN_HTTP_PORT:-8000}"
+# Fall through to ``DYN_SYSTEM_PORT`` (the single-worker convention used by
+# sibling launch scripts like ``agg_embed.sh``) for worker 1 so callers
+# that only set the non-numbered env var still drive the first worker's
+# port. ``SYSTEM_PORT2`` stays numbered-only -- there's no single-worker
+# equivalent for it.
+SYSTEM_PORT1="${DYN_SYSTEM_PORT1:-${DYN_SYSTEM_PORT:-8081}}"
+SYSTEM_PORT2="${DYN_SYSTEM_PORT2:-8082}"
+
+print_launch_banner --no-curl "Launching Multi-Worker Embeddings (2 GPUs)" "${MODEL1} + ${MODEL2}" "$HTTP_PORT"
+
+print_curl_footer <<CURL
+  curl http://localhost:${HTTP_PORT}/v1/embeddings \\
+    -H 'Content-Type: application/json' \\
+    -d '{"model": "${MODEL1}", "input": "Hello, world!"}'
+
+  # Per-worker metrics:
+  curl http://localhost:${SYSTEM_PORT1}/metrics  # worker on GPU 0 (${MODEL1})
+  curl http://localhost:${SYSTEM_PORT2}/metrics  # worker on GPU 1 (${MODEL2})
+CURL
+
+# Frontend — same routing layer as the single-worker case; the name-keyed
+# DashMap lookup in lib/llm/src/discovery/model_manager.rs handles both the
+# same-model fan-out (weighted-random across matching workers) and the
+# multi-model dispatch (model field selects the set of eligible workers).
+python3 -m dynamo.frontend &
+
+# Tunable: see agg_embed.sh for the rationale.
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-2048}"
+
+# Common worker args. Mirrors agg_embed.sh — the embedding worker handler
+# is the same; only the per-worker model + GPU + system port differ.
+common_worker_args=(
+    --embedding-worker
+    --runner pooling
+    --dtype float32
+    --pooler-config '{"pooling_type": "MEAN", "use_activation": false}'
+    --max-model-len "$MAX_MODEL_LEN"
+    --no-enable-prefix-caching
+    --trust-remote-code
+)
+
+# Each worker registers under its OWN Dynamo NAMESPACE.
+#
+# Why namespaces and not just unique components or endpoints: the frontend
+# keys ``Model.worker_sets`` by ``(namespace, model_type)`` (see
+# ``worker_set_key`` in ``lib/llm/src/discovery/watcher.rs``), and
+# ``add_worker_set`` is an insert-overwrite on that key. Two workers
+# sharing a namespace -- regardless of how their endpoint paths differ --
+# both hash to the same ``ws_key``, and the second registration silently
+# replaces the first ``WorkerSet`` (along with its push_router). Only the
+# last-registered worker survives in the routing table; the other one
+# stays alive but is orphaned from the frontend's
+# ``select_worker_set_with`` selector. The original "one model per
+# endpoint" symptom was a sibling consequence of the same collision --
+# this resolves both by giving each worker its own namespace, so both
+# ``WorkerSet``s coexist and ``select_worker_set_with`` does its
+# weighted-random fan-out as designed.
+#
+# Endpoint format is ``namespace.component.endpoint`` (dots, not slashes).
+DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=${SYSTEM_PORT1} \
+CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \
+    --model "$MODEL1" \
+    --endpoint embed-worker-1.vllm.generate \
+    "${common_worker_args[@]}" \
+    $GPU_MEM_ARGS \
+    "${EXTRA_ARGS[@]}" &
+
+DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=${SYSTEM_PORT2} \
+CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm \
+    --model "$MODEL2" \
+    --endpoint embed-worker-2.vllm.generate \
+    "${common_worker_args[@]}" \
+    $GPU_MEM_ARGS \
+    "${EXTRA_ARGS[@]}" &
+
+# Exit on first worker failure; kill 0 in the EXIT trap tears down the rest
+wait_any_exit