From ebb6d3e25133083f0f4ac992a2d08b4bad742eb8 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Thu, 21 May 2026 11:57:48 -0500 Subject: [PATCH 01/12] feat(sglang): gate chat-shaped Prometheus collectors on embedding worker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the first half of DIS-2107. The SGLang publisher previously registered chat-shaped collectors (SGLang's multiprocess ``sglang:*`` metrics, the Dynamo ``LLMBackendMetrics`` gauges ``total_blocks`` / ``gpu_cache_usage`` / ``model_load_time``, the ``DynamoSglangPublisher`` with its ZMQ scheduler pull + KV-event + FPM-relay wiring) on every worker, including embedding workers where they emit zeros forever. Operators monitoring an embedding fleet saw noise instead of signal. ``setup_sgl_metrics`` now short-circuits when ``config.dynamo_args.embedding_worker`` is set: no chat-shaped constructors run. It still returns a uniform ``(publisher_or_None, asyncio.Task, metrics_labels)`` shape so the embedding init path can keep the same ``await + cancel()`` cleanup sequence as the chat path. The corresponding ``dynamo_embedding_batch_size`` / ``dynamo_embedding_input_tokens`` collectors that DO belong on an embedding worker are registered separately by ``init_embedding.py`` via ``init_embedding_metrics`` (see PR #9753). Tests ----- ``test_sglang_publisher.py`` gains two new cases: - ``test_setup_sgl_metrics_skips_chat_pipeline_for_embedding_worker`` — patches ``setup_prometheus_registry``, ``register_engine_metrics_callback``, ``LLMBackendMetrics``, and ``DynamoSglangPublisher`` to ``raise`` if invoked. Setting ``embedding_worker=True`` must not trigger any of them, the return tuple must have ``publisher is None``, the task must be cancellable, and the model label still propagates. - ``test_setup_sgl_metrics_returns_publisher_for_chat_worker`` — sibling check that the chat path is unchanged: all chat-shaped constructors fire and a publisher comes back. vLLM publisher gating is intentionally NOT in this PR — vLLM has no embedding-worker code path on ``main`` today. The mirror change lands once the vLLM-side embedding worker (DIS-2092) merges and the vLLM-side embedding observation PR (DIS-2094 part 3) is on deck. Refs DIS-2107. Signed-off-by: Tzu-Ling (cherry picked from commit e1b83e80b470b2a8283c089789a1d97a078e700c) --- components/src/dynamo/sglang/publisher.py | 34 +++- .../sglang/tests/test_sglang_publisher.py | 153 ++++++++++++++++++ 2 files changed, 184 insertions(+), 3 deletions(-) diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index 2cc0eb40440c..cd389292e4f2 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/components/src/dynamo/sglang/publisher.py @@ -449,9 +449,25 @@ async def setup_sgl_metrics( config: Config, generate_endpoint: Endpoint, kv_worker_id: Optional[int] = None, -) -> tuple[DynamoSglangPublisher, asyncio.Task, list[tuple[str, str]]]: +) -> tuple[Optional[DynamoSglangPublisher], asyncio.Task, list[tuple[str, str]]]: """Create publisher, initialize metrics, and start the metrics publishing loop. + For chat/decode workers (the default), this registers SGLang's + multiprocess ``sglang:*`` metrics, the Dynamo ``LLMBackendMetrics`` + chat-shaped gauges (KV total_blocks, gpu_cache_usage, model_load_time), + and starts a ``DynamoSglangPublisher`` that pulls scheduler metrics + over ZMQ and (optionally) forwards KV events / FPM stats. + + For **embedding workers** (``config.dynamo_args.embedding_worker``), + the chat-shaped pipeline is **skipped entirely**: pooling engines + have no KV cache, no prefill/decode phase, and no scheduler metrics + worth collecting, so every metric in that pipeline would emit zeros + forever. The function returns ``(None, , metrics_labels)`` + so callers can keep the same ``await setup_sgl_metrics(...)`` shape + and ``metrics_task.cancel()`` cleanup. Embedding-shaped metrics are + registered separately by ``init_embedding.py`` via + ``init_embedding_metrics`` (see DIS-2094 part 1). + Args: engine: The SGLang engine instance. config: SGLang configuration including server args. @@ -459,8 +475,21 @@ async def setup_sgl_metrics( kv_worker_id: Optional worker identity for KV event attribution. Returns: - Tuple of (publisher instance, running asyncio task, metrics labels). + Tuple of (publisher instance or None, asyncio task, metrics labels). """ + metrics_labels = [("model", engine.server_args.served_model_name)] + + if getattr(config.dynamo_args, "embedding_worker", False): + logging.info( + "Embedding worker: skipping chat-shaped Prometheus + KV-event " + "wiring (no KV cache, no prefill/decode, no scheduler metrics). " + "Embedding-shaped metrics are registered separately." + ) + # Hold a never-completing task so callers can ``cancel()`` + ``await`` + # it uniformly in their finally blocks, matching the chat-worker shape. + task = asyncio.create_task(asyncio.Event().wait()) + return None, task, metrics_labels + # Register SGLang multiprocess metrics only when --enable-metrics was passed. # SGLang only calls set_prometheus_multiproc_dir() when enable_metrics=True, # so MultiProcessCollector will crash without it. @@ -487,7 +516,6 @@ async def setup_sgl_metrics( component_name=config.dynamo_args.component, ) - metrics_labels = [("model", engine.server_args.served_model_name)] publisher = DynamoSglangPublisher( engine, config, diff --git a/components/src/dynamo/sglang/tests/test_sglang_publisher.py b/components/src/dynamo/sglang/tests/test_sglang_publisher.py index 8b1141cd3121..a7740653830d 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_publisher.py +++ b/components/src/dynamo/sglang/tests/test_sglang_publisher.py @@ -15,6 +15,7 @@ get_local_dp_rank_range, handle_non_leader_node, set_forward_pass_metrics_worker_id, + setup_sgl_metrics, ) pytestmark = [ @@ -476,3 +477,155 @@ def shutdown(self): assert calls[0]["worker_id"] == 0 publisher.cleanup() + + +# ---- DIS-2107: per-worker metric gating ---- + + +@pytest.mark.asyncio +async def test_setup_sgl_metrics_skips_chat_pipeline_for_embedding_worker(monkeypatch): + """``setup_sgl_metrics`` short-circuits for embedding workers. + + Chat-shaped collectors (``sglang:*`` multiproc metrics, the Dynamo + ``LLMBackendMetrics`` gauges, ``DynamoSglangPublisher`` itself with + its KV-events / FPM relay wiring) emit zeros forever on a pooling + engine. Verify they are not constructed when + ``config.dynamo_args.embedding_worker`` is True, while preserving + the ``(publisher, task, metrics_labels)`` return shape so the + embedding init path can keep its uniform cleanup. + """ + calls: dict[str, int] = {} + + def _track(name): + def _wrapped(*_a, **_kw): + calls[name] = calls.get(name, 0) + 1 + raise AssertionError( + f"setup_sgl_metrics should not call {name} on the embedding-worker path" + ) + + return _wrapped + + monkeypatch.setattr(publisher_mod, "setup_prometheus_registry", _track("setup_prometheus_registry")) + monkeypatch.setattr( + publisher_mod, "register_engine_metrics_callback", _track("register_engine_metrics_callback") + ) + monkeypatch.setattr(publisher_mod, "LLMBackendMetrics", _track("LLMBackendMetrics")) + monkeypatch.setattr(publisher_mod, "DynamoSglangPublisher", _track("DynamoSglangPublisher")) + + engine = SimpleNamespace( + server_args=SimpleNamespace( + served_model_name="Qwen/Qwen3-Embedding-4B", + enable_metrics=True, # would normally enable chat-shaped sglang:* metrics + node_rank=0, + ) + ) + config = SimpleNamespace( + dynamo_args=SimpleNamespace(embedding_worker=True), + server_args=engine.server_args, + ) + generate_endpoint = SimpleNamespace() + + publisher, task, metrics_labels = await setup_sgl_metrics( + engine, config, generate_endpoint + ) + + try: + assert publisher is None + assert metrics_labels == [("model", "Qwen/Qwen3-Embedding-4B")] + assert isinstance(task, asyncio.Task) + # Task is intentionally a never-completing waiter so callers can + # ``cancel()`` + ``await`` uniformly in their finally blocks. + assert not task.done() + finally: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Hard assertion: NONE of the chat-shaped constructors fired. + assert calls == {} + + +@pytest.mark.asyncio +async def test_setup_sgl_metrics_returns_publisher_for_chat_worker(monkeypatch): + """The chat-worker path still constructs the publisher. + + Sibling to the embedding-worker test above: gives confidence the + gating doesn't accidentally short-circuit the default code path. + """ + constructed: dict[str, int] = {} + + def _count(name): + def _wrapped(*_a, **_kw): + constructed[name] = constructed.get(name, 0) + 1 + return SimpleNamespace() + + return _wrapped + + monkeypatch.setattr(publisher_mod, "setup_prometheus_registry", _count("setup_prometheus_registry")) + monkeypatch.setattr( + publisher_mod, "register_engine_metrics_callback", _count("register_engine_metrics_callback") + ) + monkeypatch.setattr(publisher_mod, "LLMBackendMetrics", _count("LLMBackendMetrics")) + + # Replace the publisher constructor with one that returns a stub whose + # methods exist but no-op, so we can keep the test free of real ZMQ / NATS. + class _StubPublisher: + def __init__(self, *_a, **_kw): + constructed["DynamoSglangPublisher"] = constructed.get("DynamoSglangPublisher", 0) + 1 + self.metrics_publisher = SimpleNamespace( + create_endpoint=lambda _ep: _async_noop() + ) + + def init_engine_metrics_publish(self): + pass + + def init_kv_event_publish(self): + pass + + def init_fpm_relay(self): + pass + + async def run(self): + await asyncio.Event().wait() + + async def _async_noop(): + return None + + monkeypatch.setattr(publisher_mod, "DynamoSglangPublisher", _StubPublisher) + + engine = SimpleNamespace( + server_args=SimpleNamespace( + served_model_name="Qwen/Qwen3-0.6B", + enable_metrics=False, # skip setup_prometheus_registry but still run chat path + node_rank=0, + ) + ) + config = SimpleNamespace( + dynamo_args=SimpleNamespace( + embedding_worker=False, + component="sglang-decode", + ), + server_args=engine.server_args, + ) + generate_endpoint = SimpleNamespace() + + publisher, task, _labels = await setup_sgl_metrics( + engine, config, generate_endpoint + ) + + try: + assert publisher is not None + # All three chat-shaped constructors fired. + assert constructed.get("register_engine_metrics_callback", 0) == 1 + assert constructed.get("LLMBackendMetrics", 0) == 1 + assert constructed.get("DynamoSglangPublisher", 0) == 1 + # setup_prometheus_registry was gated off by enable_metrics=False. + assert "setup_prometheus_registry" not in constructed + finally: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass From 1cc2976175be691cffabea8c6507a078886eaa2d Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Thu, 21 May 2026 12:15:22 -0500 Subject: [PATCH 02/12] fix(sglang/publisher): mypy + black + scrub internal ticket refs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses three findings on PR #9830: 1. mypy: ``setup_sgl_metrics`` now declares ``Optional[DynamoSglangPublisher]`` in its return tuple, but the chat-worker call sites (init_llm.py x2, init_diffusion.py x1) pass the publisher to functions that expect the non-Optional type. Add ``assert publisher is not None`` at those three call sites — those code paths are reached only when ``embedding_worker=False``, so the assertion is correct by construction and narrows the type for mypy. 2. mypy: the embedding-worker no-op task used ``asyncio.create_task(asyncio.Event().wait())``. ``Event.wait()`` returns ``Literal[True]`` when set, which mypy then conflicts with the chat path's ``publisher.run() -> None`` return at the second ``create_task`` call. Wrap the wait in a local ``async _idle() -> None`` so both branches feed ``Coroutine[Any, Any, None]`` into ``create_task``. 3. Repository policy: scrub internal Linear ticket IDs from source. Removed ``DIS-2094`` from the publisher docstring and ``DIS-2107`` from the test-file section header. 4. Black: re-format the two files mentioned by pre-commit. No behavioral changes. Signed-off-by: Tzu-Ling (cherry picked from commit 2029e201d086b9fafe3fcfcf6239174806506d8b) --- .../src/dynamo/sglang/init_diffusion.py | 3 +++ components/src/dynamo/sglang/init_llm.py | 6 +++++ components/src/dynamo/sglang/publisher.py | 8 ++++-- .../sglang/tests/test_sglang_publisher.py | 26 ++++++++++++++----- 4 files changed, 34 insertions(+), 9 deletions(-) diff --git a/components/src/dynamo/sglang/init_diffusion.py b/components/src/dynamo/sglang/init_diffusion.py index 32f990efafb5..552e24cc0221 100644 --- a/components/src/dynamo/sglang/init_diffusion.py +++ b/components/src/dynamo/sglang/init_diffusion.py @@ -67,6 +67,9 @@ async def init_llm_diffusion( publisher, metrics_task, metrics_labels = await setup_sgl_metrics( engine, config, generate_endpoint ) + # ``setup_sgl_metrics`` only returns ``None`` for embedding workers, + # which take a different init path entirely. Narrow for mypy. + assert publisher is not None, "setup_sgl_metrics returned None on chat path" if server_args.node_rank >= 1: await handle_non_leader_node(engine, publisher, metrics_task) diff --git a/components/src/dynamo/sglang/init_llm.py b/components/src/dynamo/sglang/init_llm.py index 06608094c166..bdb068571b90 100644 --- a/components/src/dynamo/sglang/init_llm.py +++ b/components/src/dynamo/sglang/init_llm.py @@ -93,6 +93,9 @@ async def init_decode( publisher, metrics_task, metrics_labels = await setup_sgl_metrics( engine, config, generate_endpoint ) + # ``setup_sgl_metrics`` only returns ``None`` for embedding workers, + # which take a different init path entirely. Narrow for mypy. + assert publisher is not None, "setup_sgl_metrics returned None on chat path" publisher.component_gauges.set_model_load_time(load_time) logging.debug(f"SGLang model load time: {load_time:.2f}s") @@ -238,6 +241,9 @@ async def init_prefill( publisher, metrics_task, metrics_labels = await setup_sgl_metrics( engine, config, generate_endpoint ) + # ``setup_sgl_metrics`` only returns ``None`` for embedding workers, + # which take a different init path entirely. Narrow for mypy. + assert publisher is not None, "setup_sgl_metrics returned None on chat path" publisher.component_gauges.set_model_load_time(load_time) diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index cd389292e4f2..da958c049a6e 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/components/src/dynamo/sglang/publisher.py @@ -466,7 +466,7 @@ async def setup_sgl_metrics( so callers can keep the same ``await setup_sgl_metrics(...)`` shape and ``metrics_task.cancel()`` cleanup. Embedding-shaped metrics are registered separately by ``init_embedding.py`` via - ``init_embedding_metrics`` (see DIS-2094 part 1). + ``init_embedding_metrics``. Args: engine: The SGLang engine instance. @@ -485,9 +485,13 @@ async def setup_sgl_metrics( "wiring (no KV cache, no prefill/decode, no scheduler metrics). " "Embedding-shaped metrics are registered separately." ) + # Hold a never-completing task so callers can ``cancel()`` + ``await`` # it uniformly in their finally blocks, matching the chat-worker shape. - task = asyncio.create_task(asyncio.Event().wait()) + async def _idle() -> None: + await asyncio.Event().wait() + + task = asyncio.create_task(_idle()) return None, task, metrics_labels # Register SGLang multiprocess metrics only when --enable-metrics was passed. diff --git a/components/src/dynamo/sglang/tests/test_sglang_publisher.py b/components/src/dynamo/sglang/tests/test_sglang_publisher.py index a7740653830d..034445b0da1e 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_publisher.py +++ b/components/src/dynamo/sglang/tests/test_sglang_publisher.py @@ -479,7 +479,7 @@ def shutdown(self): publisher.cleanup() -# ---- DIS-2107: per-worker metric gating ---- +# ---- per-worker metric gating (embedding vs chat) ---- @pytest.mark.asyncio @@ -505,12 +505,18 @@ def _wrapped(*_a, **_kw): return _wrapped - monkeypatch.setattr(publisher_mod, "setup_prometheus_registry", _track("setup_prometheus_registry")) monkeypatch.setattr( - publisher_mod, "register_engine_metrics_callback", _track("register_engine_metrics_callback") + publisher_mod, "setup_prometheus_registry", _track("setup_prometheus_registry") + ) + monkeypatch.setattr( + publisher_mod, + "register_engine_metrics_callback", + _track("register_engine_metrics_callback"), ) monkeypatch.setattr(publisher_mod, "LLMBackendMetrics", _track("LLMBackendMetrics")) - monkeypatch.setattr(publisher_mod, "DynamoSglangPublisher", _track("DynamoSglangPublisher")) + monkeypatch.setattr( + publisher_mod, "DynamoSglangPublisher", _track("DynamoSglangPublisher") + ) engine = SimpleNamespace( server_args=SimpleNamespace( @@ -563,9 +569,13 @@ def _wrapped(*_a, **_kw): return _wrapped - monkeypatch.setattr(publisher_mod, "setup_prometheus_registry", _count("setup_prometheus_registry")) monkeypatch.setattr( - publisher_mod, "register_engine_metrics_callback", _count("register_engine_metrics_callback") + publisher_mod, "setup_prometheus_registry", _count("setup_prometheus_registry") + ) + monkeypatch.setattr( + publisher_mod, + "register_engine_metrics_callback", + _count("register_engine_metrics_callback"), ) monkeypatch.setattr(publisher_mod, "LLMBackendMetrics", _count("LLMBackendMetrics")) @@ -573,7 +583,9 @@ def _wrapped(*_a, **_kw): # methods exist but no-op, so we can keep the test free of real ZMQ / NATS. class _StubPublisher: def __init__(self, *_a, **_kw): - constructed["DynamoSglangPublisher"] = constructed.get("DynamoSglangPublisher", 0) + 1 + constructed["DynamoSglangPublisher"] = ( + constructed.get("DynamoSglangPublisher", 0) + 1 + ) self.metrics_publisher = SimpleNamespace( create_endpoint=lambda _ep: _async_noop() ) From 15c13c47ce5857fd799317ebdde9852d07062e7f Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Wed, 20 May 2026 10:24:33 -0500 Subject: [PATCH 03/12] test(vllm): add multi-worker embedding routing tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds end-to-end + unit coverage for the two routing properties of the embedding worker pool: 1. Same-model load balancing — two workers serving the same embedding model weighted-randomly absorb a burst of /v1/embeddings traffic. The test asserts both workers' dynamo_component_requests_total counters move > 0 over a 20-request burst. 2. Multi-model dispatch — two workers serving two different models. Each model's burst must land only on its registered worker; the wrong-model worker observes zero delta. Verifies the name-keyed get_embeddings_engine(model) -> select_worker_set_with() path specifically for embedding traffic. (Chat-completions already exercises the same router code; this is explicit embedding-side coverage.) Per-burst deltas, not absolute counts ===================================== In the multi-model test, burst A leaves worker A's counter > 0; burst B's check would falsely pass if it only required "B's worker has > 0 traffic". The new payload snapshots /metrics at the start of each burst and asserts the delta through the end of the burst, which makes "wrong-model traffic stayed out of worker A during burst B" actually expressible. Pieces ====== - examples/backends/vllm/launch/agg_embed_multiworker.sh: frontend + N embedding workers (one per GPU, distinct DYN_SYSTEM_PORT each). Takes two model args. Mirrors agg_embed.sh's pooler config. - tests/utils/payloads.py: new EmbeddingMultiWorkerDispatchPayload. Takes a list of system ports (resolved from DefaultPort.SYSTEM{1,2} placeholders by the harness) and a set of INDICES (port mapping happens at runtime, so absolute port numbers can't be compared). Snapshots on first request, deltas on last. - tests/serve/test_vllm.py: two pytest.mark.gpu_2 / pre_merge test functions: test_embedding_multi_worker_same_model_load_balance, test_embedding_multi_worker_multi_model_dispatch. - tests/utils/test_embedding_dispatch_payload.py: unit tests for the payload's baseline+delta logic. Monkeypatches requests.get with a canned per-port /metrics sequence; no GPU needed. Covers the six branches that matter: load-balance pass/fail, dispatch pass/fail, min-total-delta floor, baseline-snapshot timing. Signed-off-by: Tzu-Ling (cherry picked from commit 7fc1a38ca251cfe5e86fd9a95921d6de9605a864) --- .../vllm/launch/agg_embed_multiworker.sh | 106 ++++++++ tests/serve/test_vllm.py | 209 ++++++++++++++++ tests/utils/payloads.py | 130 ++++++++++ .../utils/test_embedding_dispatch_payload.py | 231 ++++++++++++++++++ 4 files changed, 676 insertions(+) create mode 100755 examples/backends/vllm/launch/agg_embed_multiworker.sh create mode 100644 tests/utils/test_embedding_dispatch_payload.py diff --git a/examples/backends/vllm/launch/agg_embed_multiworker.sh b/examples/backends/vllm/launch/agg_embed_multiworker.sh new file mode 100755 index 000000000000..0bb1848862c8 --- /dev/null +++ b/examples/backends/vllm/launch/agg_embed_multiworker.sh @@ -0,0 +1,106 @@ +#!/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}" +SYSTEM_PORT1="${DYN_SYSTEM_PORT1:-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 < EmbeddingPayload: + """One quick embedding request used as a smoke check before the burst.""" + return EmbeddingPayload( + body={"model": model, "input": "warmup"}, + expected_response=["Generated 1 embeddings with dimension"], + expected_log=[], + repeat_count=1, + ) + + +def _embedding_dispatch_burst( + *, + model: str, + repeat_count: int, + expected_worker_indices_with_delta: set[int], + min_total_delta: int, +) -> EmbeddingMultiWorkerDispatchPayload: + """One burst payload that drives the dispatch assertion. + + ``system_ports`` is fixed at ``[SYSTEM_PORT1, SYSTEM_PORT2]`` — the + harness remaps those placeholders to the per-test dynamic ports. + Dispatch expectations are expressed as INDICES into that list (index 0 + = first worker = GPU 0, index 1 = second worker = GPU 1). + """ + return EmbeddingMultiWorkerDispatchPayload( + body={"model": model, "input": "Hello, world!"}, + expected_response=["Generated 1 embeddings with dimension"], + expected_log=[], + repeat_count=repeat_count, + system_ports=[DefaultPort.SYSTEM1.value, DefaultPort.SYSTEM2.value], + expected_worker_indices_with_delta=expected_worker_indices_with_delta, + min_total_delta=min_total_delta, + ) + + +# Same model on both GPUs — verifies weighted-random selection in +# `select_worker_set_with` fans out across both registered workers. +_EMBED_SAME_MODEL = "Qwen/Qwen3-Embedding-0.6B" + +# Multi-model setup: each model is served by exactly one worker, so a +# request whose `model` field names model A must never reach model B's +# worker. BGE-small-en is intentionally small (33M params, fits alongside +# Qwen3-Embedding-0.6B on stock CI nodes). +_EMBED_MODEL_A = "Qwen/Qwen3-Embedding-0.6B" +_EMBED_MODEL_B = "BAAI/bge-small-en-v1.5" + + +@pytest.mark.vllm +@pytest.mark.core +@pytest.mark.e2e +@pytest.mark.gpu_2 +@pytest.mark.model(_EMBED_SAME_MODEL) +@pytest.mark.profiled_vram_gib(5.0) # per GPU; mirrors single-worker embedding_agg +@pytest.mark.requested_vllm_kv_cache_bytes(559_693_824) +@pytest.mark.timeout( + 420 +) # 2x cold-load vs single-worker embedding_agg (2 GPUs in parallel) +@pytest.mark.pre_merge +@pytest.mark.parametrize("num_system_ports", [2], indirect=True) +def test_embedding_multi_worker_same_model_load_balance( + request, + runtime_services_dynamic_ports, + dynamo_dynamic_ports, + num_system_ports, + predownload_models, +): + """Two workers serving the same model: a burst of requests should be + weighted-randomly distributed so both workers' /metrics counters > 0. + + Burst size is deliberately small (20) so pure chance of all-to-one- + worker is negligible (≈ 1 in 2^19) while keeping test runtime tight + on 2-GPU CI nodes. + """ + assert num_system_ports >= 2, "Requires SYSTEM_PORT1 + SYSTEM_PORT2" + + # 20 repeats inside the burst; the payload uses repeat 1 as its + # baseline snapshot and asserts the delta across repeats 2..20. + burst = _embedding_dispatch_burst( + model=_EMBED_SAME_MODEL, + repeat_count=20, + # Both workers (indices 0 and 1) should see delta > 0. + expected_worker_indices_with_delta={0, 1}, + # 19 post-baseline requests; loose lower bound absorbs any frontend + # health probes that the worker happens to count. + min_total_delta=15, + ) + + config = VLLMConfig( + name="embedding_multi_worker_same_model", + directory=vllm_dir, + script_name="agg_embed_multiworker.sh", + script_args=[_EMBED_SAME_MODEL, _EMBED_SAME_MODEL], + marks=[], # markers at function level + model=_EMBED_SAME_MODEL, + timeout=420, + # Poll each DYN_SYSTEM_PORT*/health endpoint before sending + # traffic so we don't race the second worker's model + # registration with the frontend. Without this, the first burst + # can fire while only one worker has registered its model, + # which manifests as HTTP 404 on /v1/embeddings. + health_check_workers=True, + request_payloads=[ + _embedding_warmup_payload(_EMBED_SAME_MODEL), + burst, + ], + ) + + config = dataclasses.replace( + config, frontend_port=dynamo_dynamic_ports.frontend_port + ) + run_serve_deployment(config, request, ports=dynamo_dynamic_ports) + + +@pytest.mark.vllm +@pytest.mark.core +@pytest.mark.e2e +@pytest.mark.gpu_2 +@pytest.mark.model(_EMBED_MODEL_A) +@pytest.mark.model(_EMBED_MODEL_B) +@pytest.mark.profiled_vram_gib(5.0) # Qwen3-Embed (0.6B) is the larger of the two +@pytest.mark.requested_vllm_kv_cache_bytes(559_693_824) +@pytest.mark.timeout(420) +@pytest.mark.pre_merge +@pytest.mark.parametrize("num_system_ports", [2], indirect=True) +def test_embedding_multi_worker_multi_model_dispatch( + request, + runtime_services_dynamic_ports, + dynamo_dynamic_ports, + num_system_ports, + predownload_models, +): + """Two workers, two different models: requests for model A must reach + only worker A; symmetric for model B. Verifies name-keyed dispatch in + ``get_embeddings_engine`` for embedding traffic. + """ + assert num_system_ports >= 2, "Requires SYSTEM_PORT1 + SYSTEM_PORT2" + + # Worker A → SYSTEM_PORT1 (GPU 0, model A, payload index 0) + # Worker B → SYSTEM_PORT2 (GPU 1, model B, payload index 1) + # + # Each burst takes its own baseline snapshot and checks the DELTA + # over its repeats — so burst_b's check is independent of burst_a's + # absolute count, and "wrong-model traffic stays out" can actually + # be expressed (no delta on the wrong worker during this burst). + burst_a = _embedding_dispatch_burst( + model=_EMBED_MODEL_A, + repeat_count=10, + expected_worker_indices_with_delta={0}, # only worker A + min_total_delta=5, + ) + burst_b = _embedding_dispatch_burst( + model=_EMBED_MODEL_B, + repeat_count=10, + expected_worker_indices_with_delta={1}, # only worker B + min_total_delta=5, + ) + + config = VLLMConfig( + name="embedding_multi_worker_multi_model", + directory=vllm_dir, + script_name="agg_embed_multiworker.sh", + script_args=[_EMBED_MODEL_A, _EMBED_MODEL_B], + marks=[], # markers at function level + # ``model`` here is just metadata for the test runner; the real + # per-request model is set in each payload's body. + model=_EMBED_MODEL_A, + # BGE-small-en-v1.5's architecture caps at ``max_position_embeddings=512`` + # — applying the script's default ``MAX_MODEL_LEN=2048`` to it crashes + # the second worker at engine init. Drop the cap to BGE's native max; + # Qwen3-Embedding-0.6B happily accepts the lower cap. + env={"MAX_MODEL_LEN": "512"}, + timeout=420, + # Poll each DYN_SYSTEM_PORT*/health endpoint before sending + # traffic. Qwen3-Embedding-0.6B (~600M params) loads ~25s + # slower than BGE-small (~33M); without the per-worker wait, + # the first burst can fire while only BGE's name is + # registered with the frontend, which manifests as HTTP 404 + # for the Qwen3 model. + health_check_workers=True, + request_payloads=[ + _embedding_warmup_payload(_EMBED_MODEL_A), + _embedding_warmup_payload(_EMBED_MODEL_B), + burst_a, + burst_b, + ], + ) + + config = dataclasses.replace( + config, frontend_port=dynamo_dynamic_ports.frontend_port + ) + run_serve_deployment(config, request, ports=dynamo_dynamic_ports) diff --git a/tests/utils/payloads.py b/tests/utils/payloads.py index a72e13270d68..2b20fa3d2ab8 100644 --- a/tests/utils/payloads.py +++ b/tests/utils/payloads.py @@ -962,6 +962,136 @@ def response_handler(self, response: Any) -> str: return EmbeddingPayload.extract_embeddings(response) +@dataclass +class EmbeddingMultiWorkerDispatchPayload(BasePayload): + """Send ``repeat_count`` embedding requests to the frontend, capturing a + per-worker ``/metrics`` snapshot on the FIRST iteration and on the LAST + iteration. Assert the delta — i.e. requests attributed during *this* + burst — matches an expected per-worker pattern. + + Diff semantics are important because the same fixture may run multiple + bursts back-to-back (e.g. one per model in a multi-model dispatch + test): the prior burst leaves a worker's absolute counter > 0, so only + deltas can prove the second burst did not also reach that worker. + + Two routing properties of the embedding worker pool are checked + through this payload: + + 1. **Same-model load balancing** — when multiple workers serve the + same embedding model, ``select_worker_set_with()`` does + weighted-random selection across them, so both workers should see + ``>0`` delta. Set ``expected_workers_with_delta={port1, port2}``. + + 2. **Multi-model dispatch** — when workers serve different models, + the name-keyed ``get_embeddings_engine(model)`` lookup must route + only to the worker registered for the requested model. Set + ``expected_workers_with_delta={port_of_requested_model}`` so the + check asserts the wrong-model worker observed exactly 0 delta + during this burst. + + The per-worker counter sampled is + ``dynamo_component_requests_total`` — the same counter exercised by + ``MetricsPayload`` for single-worker tests. The dispatch assertion + fires once after the last repeat. + """ + + endpoint: str = "/v1/embeddings" + + # Indices into ``self.system_ports`` (inherited from BasePayload, with + # DefaultPort.SYSTEM{1,2} entries remapped to per-test dynamic ports + # by the harness). Each indexed worker must have observed >0 delta + # during this burst; other workers must observe exactly 0 delta. + # + # Index-based on purpose: the actual port values are not known at + # config-construction time because the harness assigns dynamic ports. + expected_worker_indices_with_delta: set[int] = field(default_factory=set) + + # Lower bound on the SUM of per-worker deltas across all workers. Use + # ``repeat_count`` for an exact-match expectation in clean fixtures. + # Defaults to 0 (predicate-only — workers_with_delta must match). + min_total_delta: int = 0 + + # Settle delay applied around each /metrics scrape so the worker has + # time to flush the most recent counter increment. + settle_seconds: float = 1.0 + + # Internal: iteration counter and baseline snapshot. Both are mutated + # in validate(); the dataclass machinery treats them as normal + # attributes once the instance is constructed. + _calls_seen: int = 0 + _baseline: dict[int, float] = field(default_factory=dict) + + def with_model(self, model): + # Embedding body is set externally; this is a no-op. + return self + + def response_handler(self, response: Any) -> str: + # Validate the per-iteration response shape so an HTML/JSON error + # surfaces immediately. The *dispatch* assertion happens in + # validate() on the final repeat. + return EmbeddingPayload.extract_embeddings(response) + + def _scrape(self) -> dict[int, float]: + prefix = prometheus_names.name_prefix.COMPONENT + counter_name = f"{prefix}_{prometheus_names.work_handler.REQUESTS_TOTAL}" + + counts: dict[int, float] = {} + for port in self.system_ports: + r = requests.get( + f"http://{self.host}:{port}/metrics", + timeout=self.timeout, + ) + r.raise_for_status() + counts[port] = sum_metric_samples(r.text, counter_name) + return counts + + def validate(self, response: Any, content: str) -> None: + """First repeat: snapshot baseline counts. Last repeat: scrape + again, compute deltas, and assert the dispatch pattern. + """ + self._calls_seen += 1 + + if self._calls_seen == 1: + # Snapshot AFTER the first request; this is fine because the + # baseline is subtracted out below — we only care that the + # delta from this point forward matches expectations. + time.sleep(self.settle_seconds) + self._baseline = self._scrape() + logger.info("Baseline per-worker counts: %s", self._baseline) + return + + if self._calls_seen < self.repeat_count: + return + + # Last repeat — compute delta. + time.sleep(self.settle_seconds) + final = self._scrape() + delta = {p: final[p] - self._baseline.get(p, 0.0) for p in final} + logger.info( + "Per-worker delta (final - baseline) over %d requests: %s", + self.repeat_count - 1, + delta, + ) + + workers_with_delta_idx = { + i for i, port in enumerate(self.system_ports) if delta.get(port, 0) > 0 + } + assert workers_with_delta_idx == self.expected_worker_indices_with_delta, ( + f"Expected worker indices with delta " + f"{self.expected_worker_indices_with_delta}, got " + f"{workers_with_delta_idx}. Per-worker delta: {delta} " + f"(baseline={self._baseline}, final={final}, " + f"system_ports={self.system_ports})" + ) + + if self.min_total_delta > 0: + total_delta = sum(delta.values()) + assert total_delta >= self.min_total_delta, ( + f"Expected at least {self.min_total_delta} total delta across " + f"workers, got {int(total_delta)}. Per-worker delta: {delta}" + ) + + @dataclass class MetricCheck: """Definition of a metric validation check""" diff --git a/tests/utils/test_embedding_dispatch_payload.py b/tests/utils/test_embedding_dispatch_payload.py new file mode 100644 index 000000000000..a78f22cba74a --- /dev/null +++ b/tests/utils/test_embedding_dispatch_payload.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for ``EmbeddingMultiWorkerDispatchPayload``. + +These tests exercise the baseline-snapshot / final-snapshot diff logic +that backs the multi-worker dispatch checks — without needing a GPU or +a real vLLM worker. We monkeypatch ``requests.get`` to return canned +``/metrics`` responses that mimic the +``dynamo_component_requests_total`` counter rising as requests land. +""" + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from tests.utils import payloads as payloads_mod +from tests.utils.payloads import EmbeddingMultiWorkerDispatchPayload + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, +] + + +def _metrics_text(value: float) -> str: + """Minimal Prometheus exposition with the counter we care about.""" + return ( + "# HELP dynamo_component_requests_total Total requests handled\n" + "# TYPE dynamo_component_requests_total counter\n" + f'dynamo_component_requests_total{{model="m"}} {value}\n' + ) + + +class _FakeMetricsServer: + """Sequence of canned ``/metrics`` responses keyed by port. + + Construct with the per-port value sequence. Each call to ``get(port)`` + advances and returns the next value. Calling more times than entries + were registered raises ``IndexError`` so over-scraping is caught. + """ + + def __init__(self, port_sequences: dict[int, list[float]]) -> None: + self._port_sequences = port_sequences + self._port_index: dict[int, int] = dict.fromkeys(port_sequences, 0) + + def get(self, port: int) -> str: + idx = self._port_index[port] + value = self._port_sequences[port][idx] + self._port_index[port] = idx + 1 + return _metrics_text(value) + + +@pytest.fixture +def patch_requests_get(monkeypatch): + """Return a callable that installs a fake ``requests.get`` driven by + a ``_FakeMetricsServer``. + """ + + def install(server: _FakeMetricsServer): + def fake_get(url: str, timeout: float = 0) -> Any: + # url looks like ``http://localhost:/metrics`` + port = int(url.split(":")[-1].split("/")[0]) + text = server.get(port) + resp = MagicMock() + resp.raise_for_status.return_value = None + resp.text = text + return resp + + monkeypatch.setattr(payloads_mod.requests, "get", fake_get) + + return install + + +def _mock_response() -> Any: + """Return a MagicMock that looks like a successful /v1/embeddings response. + + Only ``raise_for_status`` and ``json`` are touched by + ``EmbeddingPayload.extract_embeddings`` (which is what + ``response_handler`` delegates to). + """ + resp = MagicMock() + resp.raise_for_status.return_value = None + resp.json.return_value = { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}], + } + return resp + + +def _make_payload( + *, + ports: tuple[int, int] = (8081, 8082), + expected_indices: set[int], + repeat_count: int = 3, + min_total_delta: int = 0, +) -> EmbeddingMultiWorkerDispatchPayload: + return EmbeddingMultiWorkerDispatchPayload( + body={"model": "m", "input": "x"}, + expected_response=["Generated 1 embeddings with dimension"], + expected_log=[], + repeat_count=repeat_count, + host="localhost", + system_ports=list(ports), + expected_worker_indices_with_delta=expected_indices, + min_total_delta=min_total_delta, + settle_seconds=0.0, + ) + + +def _drive_payload( + payload: EmbeddingMultiWorkerDispatchPayload, + *, + response_text: str = "Generated 1 embeddings with dimension 3", +) -> None: + """Invoke validate() ``payload.repeat_count`` times, mirroring what the + test harness loop does for each request iteration. + """ + for _ in range(payload.repeat_count): + payload.validate(_mock_response(), response_text) + + +# ── Tests ────────────────────────────────────────────────────────────────── + + +def test_same_model_load_balance_passes_when_both_workers_increment( + patch_requests_get, +): + """Both workers' counters rise → expected_indices={0, 1} passes.""" + # baseline scrape (repeat 1), then final scrape (repeat 3); repeat 2 is a no-op + patch_requests_get( + _FakeMetricsServer( + { + 8081: [5.0, 12.0], # +7 on worker A + 8082: [3.0, 11.0], # +8 on worker B + } + ) + ) + payload = _make_payload(expected_indices={0, 1}) + _drive_payload(payload) + + +def test_same_model_load_balance_fails_when_only_one_worker_increments( + patch_requests_get, +): + """One worker idle while the other absorbs all traffic → assert fires.""" + patch_requests_get( + _FakeMetricsServer( + { + 8081: [5.0, 25.0], # +20 on A + 8082: [3.0, 3.0], # +0 on B (idle) + } + ) + ) + payload = _make_payload(expected_indices={0, 1}) + with pytest.raises(AssertionError, match="Expected worker indices"): + _drive_payload(payload) + + +def test_multi_model_dispatch_passes_when_only_target_worker_increments( + patch_requests_get, +): + """Model-A traffic should appear only on worker A (index 0).""" + patch_requests_get( + _FakeMetricsServer( + { + 8081: [5.0, 15.0], # +10 on A + 8082: [3.0, 3.0], # +0 on B — correct! + } + ) + ) + payload = _make_payload(expected_indices={0}, min_total_delta=10) + _drive_payload(payload) + + +def test_multi_model_dispatch_fails_when_wrong_worker_gets_traffic( + patch_requests_get, +): + """If model-A traffic leaks onto worker B, the index check catches it.""" + patch_requests_get( + _FakeMetricsServer( + { + 8081: [5.0, 15.0], # +10 on A + 8082: [3.0, 4.0], # +1 on B — leak! + } + ) + ) + payload = _make_payload(expected_indices={0}) + with pytest.raises(AssertionError, match="Expected worker indices"): + _drive_payload(payload) + + +def test_min_total_delta_lower_bound(patch_requests_get): + """min_total_delta enforces a per-burst floor on summed delta.""" + patch_requests_get( + _FakeMetricsServer( + { + 8081: [0.0, 2.0], # +2 + 8082: [0.0, 1.0], # +1 → sum = 3 + } + ) + ) + payload = _make_payload( + expected_indices={0, 1}, + min_total_delta=10, # floor 10, actual 3 — should fail + ) + with pytest.raises(AssertionError, match="total delta"): + _drive_payload(payload) + + +def test_baseline_snapshot_taken_on_first_repeat(patch_requests_get): + """Baseline is taken after the first request — earlier traffic on the + workers should not be subtracted from this burst's delta. + """ + # Worker A starts already at 100 (e.g. a prior burst's leftover), then + # increments to 105 by the end of this burst. Worker B stays flat at 50. + patch_requests_get( + _FakeMetricsServer( + { + 8081: [100.0, 105.0], # +5 + 8082: [50.0, 50.0], # +0 + } + ) + ) + payload = _make_payload( + expected_indices={0}, # only A should show delta + min_total_delta=5, + ) + _drive_payload(payload) From 6cff5126f641be17dde0131659d222a7ca170ddf Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Thu, 21 May 2026 11:46:52 -0500 Subject: [PATCH 04/12] fix(test): use delayed_start instead of health_check_workers for multi-worker embed tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous fix used ``health_check_workers=True`` to wait for both embedding workers before sending requests, but the Dynamo ``/health`` endpoint on an embedding worker stays at HTTP 503 because the embedding worker handler in PR #9713 doesn't call ``set_health_status(True)`` after registering its model. The harness polled ``/health`` for the full pytest-timeout (420s) and the test failed with "Timeout (>420.0s) from pytest-timeout". Switch to ``delayed_start=90`` on both multi-worker configs. 90s is comfortably above the observed ~30s per-worker load time and gives both workers room to register their model names with the frontend before the first burst fires. Long-term fix is for the embedding worker handler to flip the readiness flag once the model is registered — tracked separately. Signed-off-by: Tzu-Ling (cherry picked from commit cdd2ab27793b040641abaace9767302605fa096f) --- tests/serve/test_vllm.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/tests/serve/test_vllm.py b/tests/serve/test_vllm.py index d24e8a4f3d77..1af9f26ad242 100644 --- a/tests/serve/test_vllm.py +++ b/tests/serve/test_vllm.py @@ -1006,12 +1006,15 @@ def test_embedding_multi_worker_same_model_load_balance( marks=[], # markers at function level model=_EMBED_SAME_MODEL, timeout=420, - # Poll each DYN_SYSTEM_PORT*/health endpoint before sending - # traffic so we don't race the second worker's model - # registration with the frontend. Without this, the first burst - # can fire while only one worker has registered its model, - # which manifests as HTTP 404 on /v1/embeddings. - health_check_workers=True, + # Crude readiness gate: wait long enough for both Qwen3-Embedding + # workers to load + register before sending traffic. The + # framework's ``health_check_workers=True`` path polls + # ``/health`` on each ``DYN_SYSTEM_PORT*``, but the embedding + # worker handler in PR #9713 doesn't call + # ``set_health_status(True)`` so ``/health`` stays at 503 + # indefinitely — using ``delayed_start`` instead. 90s is + # comfortably above the observed ~30s per-worker load time. + delayed_start=90, request_payloads=[ _embedding_warmup_payload(_EMBED_SAME_MODEL), burst, @@ -1083,13 +1086,12 @@ def test_embedding_multi_worker_multi_model_dispatch( # Qwen3-Embedding-0.6B happily accepts the lower cap. env={"MAX_MODEL_LEN": "512"}, timeout=420, - # Poll each DYN_SYSTEM_PORT*/health endpoint before sending - # traffic. Qwen3-Embedding-0.6B (~600M params) loads ~25s - # slower than BGE-small (~33M); without the per-worker wait, - # the first burst can fire while only BGE's name is - # registered with the frontend, which manifests as HTTP 404 - # for the Qwen3 model. - health_check_workers=True, + # Crude readiness gate. ``health_check_workers=True`` would race + # against ``/health`` which stays at HTTP 503 on the embedding + # worker (see same-model test for the rationale). 90s is enough + # for both Qwen3-Embedding-0.6B and BGE-small to load and + # register their model names with the frontend. + delayed_start=90, request_payloads=[ _embedding_warmup_payload(_EMBED_MODEL_A), _embedding_warmup_payload(_EMBED_MODEL_B), From b5e316a262c2b34679b0dbac30c97e0624e2ff97 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Thu, 21 May 2026 23:10:54 -0500 Subject: [PATCH 05/12] fix(test): give each multi-worker embedding worker a unique --endpoint Dynamo enforces one model per endpoint path (namespace/component/endpoint). The previous multi-worker script let both workers default to the same ``dynamo/backend/generate`` endpoint; that works for the same-model load-balance test (both workers register the same model name) but the multi-model dispatch test fails with: Failed to serve embedding worker endpoint: Cannot register model 'Qwen/Qwen3-Embedding-0.6B' on endpoint 'dynamo/backend/generate': a different model 'BAAI/bge-small-en-v1.5' is already registered there Give worker 1 ``--endpoint dynamo/embed-worker-1/generate`` and worker 2 ``--endpoint dynamo/embed-worker-2/generate``. The frontend's name-keyed ``get_embeddings_engine(model)`` dispatch still routes correctly by the ``model`` field regardless of which component path hosts each model. Signed-off-by: Tzu-Ling (cherry picked from commit c85091bf6862cf235ad391b121a5e83ea1ee37f2) --- examples/backends/vllm/launch/agg_embed_multiworker.sh | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/examples/backends/vllm/launch/agg_embed_multiworker.sh b/examples/backends/vllm/launch/agg_embed_multiworker.sh index 0bb1848862c8..0e20d9aabe6d 100755 --- a/examples/backends/vllm/launch/agg_embed_multiworker.sh +++ b/examples/backends/vllm/launch/agg_embed_multiworker.sh @@ -88,9 +88,16 @@ common_worker_args=( --trust-remote-code ) +# Each worker registers at its OWN ``--endpoint`` path. Dynamo enforces one +# model per endpoint, so without unique endpoints the second worker fails +# to register with: "Cannot register model 'X' on endpoint Y: a different +# model 'Z' is already registered there". The frontend's model-keyed +# dispatch (``get_embeddings_engine(model)``) still routes correctly by the +# ``model`` field regardless of which component path hosts each model. DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=${SYSTEM_PORT1} \ CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \ --model "$MODEL1" \ + --endpoint dynamo/embed-worker-1/generate \ "${common_worker_args[@]}" \ $GPU_MEM_ARGS \ "${EXTRA_ARGS[@]}" & @@ -98,6 +105,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \ DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=${SYSTEM_PORT2} \ CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm \ --model "$MODEL2" \ + --endpoint dynamo/embed-worker-2/generate \ "${common_worker_args[@]}" \ $GPU_MEM_ARGS \ "${EXTRA_ARGS[@]}" & From 3bd48f259a65bad9f6e169cc7bd6c18df0f97e76 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Fri, 22 May 2026 11:20:40 -0500 Subject: [PATCH 06/12] fix(test): correct --endpoint separator (dots, not slashes) Dynamo expects namespace.component.endpoint (dot-separated); the previous commit used slashes and the worker died at startup with: ValueError: Invalid endpoint format: 'dynamo/embed-worker-1/generate'. Expected 'dyn://namespace.component.endpoint' or 'namespace.component.endpoint'. Signed-off-by: Tzu-Ling (cherry picked from commit 051ecbbdf7c0ed5b69c40e0c1f70e0f9965b836a) --- examples/backends/vllm/launch/agg_embed_multiworker.sh | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/backends/vllm/launch/agg_embed_multiworker.sh b/examples/backends/vllm/launch/agg_embed_multiworker.sh index 0e20d9aabe6d..29b66b115890 100755 --- a/examples/backends/vllm/launch/agg_embed_multiworker.sh +++ b/examples/backends/vllm/launch/agg_embed_multiworker.sh @@ -94,10 +94,12 @@ common_worker_args=( # model 'Z' is already registered there". The frontend's model-keyed # dispatch (``get_embeddings_engine(model)``) still routes correctly by the # ``model`` field regardless of which component path hosts each model. +# +# 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 dynamo/embed-worker-1/generate \ + --endpoint dynamo.embed-worker-1.generate \ "${common_worker_args[@]}" \ $GPU_MEM_ARGS \ "${EXTRA_ARGS[@]}" & @@ -105,7 +107,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \ DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=${SYSTEM_PORT2} \ CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm \ --model "$MODEL2" \ - --endpoint dynamo/embed-worker-2/generate \ + --endpoint dynamo.embed-worker-2.generate \ "${common_worker_args[@]}" \ $GPU_MEM_ARGS \ "${EXTRA_ARGS[@]}" & From 8e5c894dc4762b117e56d7591ca97207d824c821 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Fri, 22 May 2026 15:20:41 -0500 Subject: [PATCH 07/12] fix(vllm): canary-based readiness for embedding workers The multi-worker embedding tests have been flaking on multi-GPU CI because there is no real readiness signal -- ``delayed_start=90`` is just a fixed sleep, and the per-worker cold-load time on the AMD GPU pool varies enough that one worker can still be loading when the test starts firing requests. The router only sees workers that have finished registering, so every burst request lands on whichever worker happened to be ready first. The recent failure mode is the clear signature: ``Expected worker indices with delta {0, 1}, got {1}``, with worker 0's per-port counter sitting at zero across both baseline AND burst snapshots -- the test fixture never reached worker 0 either, because worker 0 was not yet in the registry. Fix --- Mirror the chat worker pattern: give the embedding worker a real canary payload that the runtime probes ``serve_endpoint``'s registered handler with. Once the probe runs a successful embedding forward pass, the runtime flips ``/health`` to 200, and the test harness's ``health_check_workers=True`` path can be used as the real readiness gate instead of a wall-clock sleep. - ``components/src/dynamo/vllm/health_check.py`` -- new ``VllmEmbeddingHealthCheckPayload`` returns ``{"model": , "input": "probe", "_HEALTH_CHECK": True}``. Mirrors ``VllmHealthCheckPayload`` / ``VllmPrefillHealthCheckPayload`` shape so the runtime treats it the same way (canary registration, periodic probe, /health gate). - ``components/src/dynamo/vllm/worker_factory.py::_create_embedding_worker`` -- builds the payload and passes it through ``serve_endpoint(health_check_payload=...)`` alongside the ``register_vllm_model`` task. The chat-shaped payload would be rejected by ``EmbeddingWorkerHandler.generate`` ("missing required 'input' field"); embedding-shaped payload goes through cleanly and runs a real pooling pass on a tiny input. - ``tests/serve/test_vllm.py`` -- ``test_embedding_multi_worker_same_model_load_balance`` and ``test_embedding_multi_worker_multi_model_dispatch`` both switch from ``delayed_start=90`` to ``health_check_workers=True``. The framework now polls each worker's ``/health`` until it returns 200 before firing the warmup or burst payloads. No more wall-clock guessing. Probe cost ---------- The canary runs one real ``engine.encode("probe")`` per probe cycle on each worker. For Qwen3-Embedding-0.6B that is a single ~10ms pooling pass; for BGE-small-en-v1.5 (33M params) the pass is effectively free. Trades that minor steady-state GPU cycle for a deterministic readiness signal. Signed-off-by: Tzu-Ling (cherry picked from commit f1a6263216cb8d3d9f99f218bac09d15c0f6ff39) --- components/src/dynamo/vllm/health_check.py | 39 ++++++++++++++++++++ components/src/dynamo/vllm/worker_factory.py | 19 +++++++++- tests/serve/test_vllm.py | 28 +++++++------- 3 files changed, 70 insertions(+), 16 deletions(-) diff --git a/components/src/dynamo/vllm/health_check.py b/components/src/dynamo/vllm/health_check.py index 775b3973b63e..e06cae9b50ab 100644 --- a/components/src/dynamo/vllm/health_check.py +++ b/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: str): + """ + Args: + model_name: served model name the handler will accept on + ``request["model"]``. The handler also falls back to + ``config.served_model_name`` when the field is absent, + so this is technically optional, but providing it keeps + the probe self-describing in worker logs. + """ + self.default_payload = { + "model": model_name, + "input": "probe", + } + 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. diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index b08d3a3a9b37..39c389bf6ebf 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/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,25 @@ async def _create_embedding_worker( shutdown_event=shutdown_event, ) + # Canary payload for the runtime's periodic /health check. The + # chat-path payload (token_ids + sampling_options) is rejected by + # the embedding handler ("missing required 'input' field") so + # without an embedding-shaped probe the worker's /health stays + # at 503 forever -- which makes K8s readiness probes and the + # test harness's health_check_workers=True path race against + # startup. A real pooling pass through "probe" is cheap and + # actually verifies the engine works end-to-end. + 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, diff --git a/tests/serve/test_vllm.py b/tests/serve/test_vllm.py index 1af9f26ad242..7934af1bed12 100644 --- a/tests/serve/test_vllm.py +++ b/tests/serve/test_vllm.py @@ -1006,15 +1006,13 @@ def test_embedding_multi_worker_same_model_load_balance( marks=[], # markers at function level model=_EMBED_SAME_MODEL, timeout=420, - # Crude readiness gate: wait long enough for both Qwen3-Embedding - # workers to load + register before sending traffic. The - # framework's ``health_check_workers=True`` path polls - # ``/health`` on each ``DYN_SYSTEM_PORT*``, but the embedding - # worker handler in PR #9713 doesn't call - # ``set_health_status(True)`` so ``/health`` stays at 503 - # indefinitely — using ``delayed_start`` instead. 90s is - # comfortably above the observed ~30s per-worker load time. - delayed_start=90, + # Poll each worker's ``/health`` until 200 before sending + # traffic. The embedding worker registers a canary + # ``VllmEmbeddingHealthCheckPayload`` with ``serve_endpoint``, + # so the runtime drives /health to 200 once the engine produces + # a real embedding -- not a fixed-time sleep that races against + # variable per-worker cold-load latency. + health_check_workers=True, request_payloads=[ _embedding_warmup_payload(_EMBED_SAME_MODEL), burst, @@ -1086,12 +1084,12 @@ def test_embedding_multi_worker_multi_model_dispatch( # Qwen3-Embedding-0.6B happily accepts the lower cap. env={"MAX_MODEL_LEN": "512"}, timeout=420, - # Crude readiness gate. ``health_check_workers=True`` would race - # against ``/health`` which stays at HTTP 503 on the embedding - # worker (see same-model test for the rationale). 90s is enough - # for both Qwen3-Embedding-0.6B and BGE-small to load and - # register their model names with the frontend. - delayed_start=90, + # Poll each worker's ``/health`` until 200 before sending + # traffic -- same canary mechanism as the same-model test. + # Both workers register an embedding-shaped probe via + # ``VllmEmbeddingHealthCheckPayload``, so the framework only + # advances once each engine has produced a real embedding. + health_check_workers=True, request_payloads=[ _embedding_warmup_payload(_EMBED_MODEL_A), _embedding_warmup_payload(_EMBED_MODEL_B), From 5c01882f9c261163434186792cb593a173cb44bc Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Fri, 22 May 2026 17:34:53 -0500 Subject: [PATCH 08/12] fix(test): use distinct namespaces per multi-worker embedding worker The multi-worker embedding tests have been flaking on multi-GPU CI with all burst requests landing on a single worker, even when both workers are alive and reachable. After the canary readiness fix ensured both workers reach ``/health = 200`` deterministically, the underlying cause surfaced cleanly: the previous "unique endpoint" workaround did not give the workers truly independent routing. Root cause ---------- The frontend's discovery layer keys ``Model.worker_sets`` by ``(namespace, model_type)`` (see ``worker_set_key`` in ``lib/llm/src/discovery/watcher.rs``) -- the endpoint component is NOT part of the key. ``add_worker_set`` on that key is an insert-overwrite (``DashMap::insert``), so two workers sharing a namespace -- no matter 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). The first worker stays alive on its system port and the canary keeps it "healthy," but it is orphaned from the frontend's ``select_worker_set_with`` selector: 100% of routed traffic lands on whichever worker registered last. The same collision was the ORIGINAL symptom that drove the unique- endpoint workaround -- "Cannot register model 'X' on endpoint Y: a different model 'Z' is already registered there." Splitting the endpoint name only hid that symptom; the underlying WorkerSet collision was still there. Fix --- Give each worker its own Dynamo namespace (``embed-worker-1.vllm.generate`` and ``embed-worker-2.vllm.generate``). Different namespaces produce different ``ws_key``s, both WorkerSets coexist in ``Model.worker_sets``, and ``select_worker_set_with`` does its weighted-random fan-out across them as designed -- which is the routing path the test was authored to verify. The frontend discovers models across namespaces without configuration; ``python3 -m dynamo.frontend`` watches all namespaces by default, so no frontend change is needed. Signed-off-by: Tzu-Ling (cherry picked from commit 34b496c400f93c38477295358b175e008e34f137) --- .../vllm/launch/agg_embed_multiworker.sh | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/examples/backends/vllm/launch/agg_embed_multiworker.sh b/examples/backends/vllm/launch/agg_embed_multiworker.sh index 29b66b115890..9ab944af2be9 100755 --- a/examples/backends/vllm/launch/agg_embed_multiworker.sh +++ b/examples/backends/vllm/launch/agg_embed_multiworker.sh @@ -88,18 +88,28 @@ common_worker_args=( --trust-remote-code ) -# Each worker registers at its OWN ``--endpoint`` path. Dynamo enforces one -# model per endpoint, so without unique endpoints the second worker fails -# to register with: "Cannot register model 'X' on endpoint Y: a different -# model 'Z' is already registered there". The frontend's model-keyed -# dispatch (``get_embeddings_engine(model)``) still routes correctly by the -# ``model`` field regardless of which component path hosts each model. +# 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 dynamo.embed-worker-1.generate \ + --endpoint embed-worker-1.vllm.generate \ "${common_worker_args[@]}" \ $GPU_MEM_ARGS \ "${EXTRA_ARGS[@]}" & @@ -107,7 +117,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \ DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=${SYSTEM_PORT2} \ CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm \ --model "$MODEL2" \ - --endpoint dynamo.embed-worker-2.generate \ + --endpoint embed-worker-2.vllm.generate \ "${common_worker_args[@]}" \ $GPU_MEM_ARGS \ "${EXTRA_ARGS[@]}" & From b82e1b4302594b44761eb54082d96ced019ed13a Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Tue, 26 May 2026 19:51:50 -0500 Subject: [PATCH 09/12] chore(test): accept DYN_SYSTEM_PORT fallback for worker 1 Sibling launch scripts (``agg_embed.sh``, ``agg.sh``, etc.) use ``DYN_SYSTEM_PORT`` as the single-worker convention; multi-worker scripts use the numbered ``DYN_SYSTEM_PORT1`` / ``DYN_SYSTEM_PORT2``. Make the multiworker embedding launcher accept either form for worker 1 by falling through ``DYN_SYSTEM_PORT1 -> DYN_SYSTEM_PORT -> 8081`` so callers that only set the non-numbered env var still drive worker 1's port. ``SYSTEM_PORT2`` stays numbered-only -- there's no single-worker equivalent. Per CodeRabbit feedback on PR #9765. Signed-off-by: Tzu-Ling (cherry picked from commit baf7777d5d50166b5bb80927fd04ab3e5fe390ad) --- examples/backends/vllm/launch/agg_embed_multiworker.sh | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/backends/vllm/launch/agg_embed_multiworker.sh b/examples/backends/vllm/launch/agg_embed_multiworker.sh index 9ab944af2be9..6cf8e1d3c756 100755 --- a/examples/backends/vllm/launch/agg_embed_multiworker.sh +++ b/examples/backends/vllm/launch/agg_embed_multiworker.sh @@ -52,7 +52,12 @@ EXTRA_ARGS=("$@") GPU_MEM_ARGS=$(build_vllm_gpu_mem_args) HTTP_PORT="${DYN_HTTP_PORT:-8000}" -SYSTEM_PORT1="${DYN_SYSTEM_PORT1:-8081}" +# 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" From 97135bb2c93e2cc8abefa086c2bf6db9ba98cf40 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Wed, 27 May 2026 10:35:09 -0500 Subject: [PATCH 10/12] fix(embed-test): enable canary in serve test + add embedding payload unit test Three follow-ups from @tmonty12's review on PR #9765: 1. ``components/src/dynamo/vllm/worker_factory.py`` -- drop the inline comment block above the ``VllmEmbeddingHealthCheckPayload`` construction. It was duplicating the class docstring almost verbatim. The construction line is now self-explanatory; readers who want the "why" go straight to the class. 2. ``tests/serve/test_vllm.py`` -- set ``DYN_HEALTH_CHECK_ENABLED=true`` on both multi-worker embedding test configs. Without that env var the runtime's canary never actually runs and ``/health`` returns 200 immediately after endpoint registration, so ``health_check_workers=True`` gates on a constant-true signal and we race startup just like the old ``delayed_start`` path. The whole point of the canary readiness fix is to gate on "engine produced an embedding" -- both knobs (flag + embedding-shaped probe payload) must be on. 3. ``components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py`` -- add ``VllmEmbeddingHealthCheckPayload`` to the marker / env-override parametrize and a new ``test_embedding_payload_shape_matches_handler_contract`` regression pin that asserts the payload produces ``{model, input}`` (the shape ``EmbeddingWorkerHandler.generate`` accepts) and explicitly NOT ``token_ids`` / ``sampling_options`` / ``stop_conditions`` (the chat shape that would trip the embedding handler's input validation). The parametrize list switches from ``PAYLOAD_CLASSES`` to ``PAYLOAD_FACTORIES`` because ``VllmEmbeddingHealthCheckPayload`` requires a ``model_name`` argument; the other three subclasses stay no-arg-constructable so they're passed directly. Signed-off-by: Tzu-Ling (cherry picked from commit e35a2d34303ca0a690b7ecc79a9b4226822bde40) --- .../tests/test_vllm_health_check_payloads.py | 46 +++++++++++++++---- components/src/dynamo/vllm/worker_factory.py | 8 ---- tests/serve/test_vllm.py | 26 ++++++----- 3 files changed, 50 insertions(+), 30 deletions(-) diff --git a/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py b/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py index 925061c93c6a..f556780e74ab 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py +++ b/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,27 @@ pytest.mark.pre_merge, ] -PAYLOAD_CLASSES = [ - VllmHealthCheckPayload, - VllmPrefillHealthCheckPayload, - VllmOmniHealthCheckPayload, +# Each entry is a no-arg factory because ``VllmEmbeddingHealthCheckPayload`` +# requires a ``model_name`` -- the other three subclasses are constructed +# with no args. +PAYLOAD_FACTORIES = [ + pytest.param(VllmHealthCheckPayload, id="VllmHealthCheckPayload"), + pytest.param(VllmPrefillHealthCheckPayload, id="VllmPrefillHealthCheckPayload"), + pytest.param(VllmOmniHealthCheckPayload, id="VllmOmniHealthCheckPayload"), + pytest.param( + lambda: VllmEmbeddingHealthCheckPayload(model_name="test-model"), + 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()[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 +60,22 @@ def test_env_override_preserves_marker(monkeypatch, cls): } ), ) - assert cls().to_dict()[HEALTH_CHECK_KEY] is True + assert make().to_dict()[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["model"] == "Qwen/Qwen3-Embedding-0.6B" + assert payload["input"] == "probe" + assert payload[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 diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 39c389bf6ebf..de19b6de4e87 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -286,14 +286,6 @@ async def _create_embedding_worker( shutdown_event=shutdown_event, ) - # Canary payload for the runtime's periodic /health check. The - # chat-path payload (token_ids + sampling_options) is rejected by - # the embedding handler ("missing required 'input' field") so - # without an embedding-shaped probe the worker's /health stays - # at 503 forever -- which makes K8s readiness probes and the - # test harness's health_check_workers=True path race against - # startup. A real pooling pass through "probe" is cheap and - # actually verifies the engine works end-to-end. embedding_health_check_payload = VllmEmbeddingHealthCheckPayload( model_name=config.served_model_name or config.model ).to_dict() diff --git a/tests/serve/test_vllm.py b/tests/serve/test_vllm.py index 7934af1bed12..e606df0a65d6 100644 --- a/tests/serve/test_vllm.py +++ b/tests/serve/test_vllm.py @@ -1006,13 +1006,16 @@ def test_embedding_multi_worker_same_model_load_balance( marks=[], # markers at function level model=_EMBED_SAME_MODEL, timeout=420, - # Poll each worker's ``/health`` until 200 before sending - # traffic. The embedding worker registers a canary - # ``VllmEmbeddingHealthCheckPayload`` with ``serve_endpoint``, - # so the runtime drives /health to 200 once the engine produces - # a real embedding -- not a fixed-time sleep that races against - # variable per-worker cold-load latency. + # ``DYN_HEALTH_CHECK_ENABLED=true`` flips the runtime's canary + # on. Without it ``/health`` returns 200 the moment the endpoint + # is registered (before the engine has produced anything), so + # ``health_check_workers=True`` would gate on a constant-true + # signal and we'd race startup just like the old ``delayed_start`` + # path. Setting the flag plus the embedding-shaped probe payload + # in ``_create_embedding_worker`` is what actually makes + # readiness mean "engine produced an embedding". health_check_workers=True, + env={"DYN_HEALTH_CHECK_ENABLED": "true"}, request_payloads=[ _embedding_warmup_payload(_EMBED_SAME_MODEL), burst, @@ -1082,13 +1085,12 @@ def test_embedding_multi_worker_multi_model_dispatch( # — applying the script's default ``MAX_MODEL_LEN=2048`` to it crashes # the second worker at engine init. Drop the cap to BGE's native max; # Qwen3-Embedding-0.6B happily accepts the lower cap. - env={"MAX_MODEL_LEN": "512"}, + # ``DYN_HEALTH_CHECK_ENABLED=true`` is required for + # ``health_check_workers=True`` below to gate on the canary + # rather than on endpoint registration (see same-model test for + # the full rationale). + env={"MAX_MODEL_LEN": "512", "DYN_HEALTH_CHECK_ENABLED": "true"}, timeout=420, - # Poll each worker's ``/health`` until 200 before sending - # traffic -- same canary mechanism as the same-model test. - # Both workers register an embedding-shaped probe via - # ``VllmEmbeddingHealthCheckPayload``, so the framework only - # advances once each engine has produced a real embedding. health_check_workers=True, request_payloads=[ _embedding_warmup_payload(_EMBED_MODEL_A), From 01291547d81a81d0f496cfca32b1ea7b7278b5e4 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Wed, 27 May 2026 11:46:48 -0500 Subject: [PATCH 11/12] chore(test): black-23 line wrap on long VllmEmbeddingHealthCheckPayload call CI's pre-commit black 23.1.0 wants the constructor call split across multiple lines; my local black 26 considered the one-liner fine. Matching CI to unblock the pre-commit job on PR #9765. Signed-off-by: Tzu-Ling (cherry picked from commit 822085184253a965e03341b5f548df55964caddb) --- .../src/dynamo/vllm/tests/test_vllm_health_check_payloads.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py b/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py index f556780e74ab..62b9820baeb0 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py +++ b/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py @@ -70,7 +70,9 @@ def test_embedding_payload_shape_matches_handler_contract(): 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() + payload = VllmEmbeddingHealthCheckPayload( + model_name="Qwen/Qwen3-Embedding-0.6B" + ).to_dict() assert payload["model"] == "Qwen/Qwen3-Embedding-0.6B" assert payload["input"] == "probe" assert payload[HEALTH_CHECK_KEY] is True From 8bf12812887945a277ac76db49c3878a718e10bf Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Thu, 28 May 2026 19:45:59 -0500 Subject: [PATCH 12/12] chore(vllm): make VllmEmbeddingHealthCheckPayload.model_name optional Per @biswapanda's review on PR #9765: 1. ``health_check.py`` -- ``model_name`` is now ``Optional[str] = None``. When the caller doesn't pass a model name, the ``model`` key is omitted from the payload entirely and ``EmbeddingWorkerHandler.generate`` falls back to ``config.served_model_name`` (which the handler already does on the missing-field path). Existing callers that pass a name still get the same self-describing log behavior. 2. ``test_vllm_health_check_payloads.py``: - Switch ``payload[key]`` lookups to ``payload.get(key)`` so a missing key produces a clean ``None != expected`` assertion instead of a ``KeyError`` before the assertion runs. - The embedding factory entry can now use the class directly (no lambda) since the constructor takes no required args. - Add a sibling regression test (``test_embedding_payload_omits_model_when_no_name``) pinning the no-arg behavior: the ``model`` key is absent and the rest of the shape is intact. Signed-off-by: Tzu-Ling (cherry picked from commit e9e8da53f0bfa27166e08d7cb8343026c6690598) --- components/src/dynamo/vllm/health_check.py | 20 +++++++------- .../tests/test_vllm_health_check_payloads.py | 26 ++++++++++++------- 2 files changed, 27 insertions(+), 19 deletions(-) diff --git a/components/src/dynamo/vllm/health_check.py b/components/src/dynamo/vllm/health_check.py index e06cae9b50ab..9ef929e4f808 100644 --- a/components/src/dynamo/vllm/health_check.py +++ b/components/src/dynamo/vllm/health_check.py @@ -133,19 +133,19 @@ class VllmEmbeddingHealthCheckPayload(HealthCheckPayload): cycle. """ - def __init__(self, model_name: str): + def __init__(self, model_name: Optional[str] = None): """ Args: - model_name: served model name the handler will accept on - ``request["model"]``. The handler also falls back to - ``config.served_model_name`` when the field is absent, - so this is technically optional, but providing it keeps - the probe self-describing in worker logs. + 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 = { - "model": model_name, - "input": "probe", - } + 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]: diff --git a/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py b/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py index 62b9820baeb0..01e46f46e8f2 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py +++ b/components/src/dynamo/vllm/tests/test_vllm_health_check_payloads.py @@ -28,15 +28,12 @@ pytest.mark.pre_merge, ] -# Each entry is a no-arg factory because ``VllmEmbeddingHealthCheckPayload`` -# requires a ``model_name`` -- the other three subclasses are constructed -# with no args. PAYLOAD_FACTORIES = [ pytest.param(VllmHealthCheckPayload, id="VllmHealthCheckPayload"), pytest.param(VllmPrefillHealthCheckPayload, id="VllmPrefillHealthCheckPayload"), pytest.param(VllmOmniHealthCheckPayload, id="VllmOmniHealthCheckPayload"), pytest.param( - lambda: VllmEmbeddingHealthCheckPayload(model_name="test-model"), + VllmEmbeddingHealthCheckPayload, id="VllmEmbeddingHealthCheckPayload", ), ] @@ -44,7 +41,7 @@ @pytest.mark.parametrize("make", PAYLOAD_FACTORIES) def test_payload_has_marker(make): - assert make().to_dict()[HEALTH_CHECK_KEY] is True + assert make().to_dict().get(HEALTH_CHECK_KEY) is True @pytest.mark.parametrize("make", PAYLOAD_FACTORIES) @@ -60,7 +57,7 @@ def test_env_override_preserves_marker(monkeypatch, make): } ), ) - assert make().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(): @@ -73,11 +70,22 @@ def test_embedding_payload_shape_matches_handler_contract(): payload = VllmEmbeddingHealthCheckPayload( model_name="Qwen/Qwen3-Embedding-0.6B" ).to_dict() - assert payload["model"] == "Qwen/Qwen3-Embedding-0.6B" - assert payload["input"] == "probe" - assert payload[HEALTH_CHECK_KEY] is True + 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