chore(e2e): validation snapshot — combined #9830 + #9765 (do not merge)

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)

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)
 

components/src/dynamo/sglang/publisher.py

@@ -449,18 +449,51 @@ 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, <noop task>, 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``.
+
     Args:
         engine: The SGLang engine instance.
         config: SGLang configuration including server args.
         generate_endpoint: The Dynamo endpoint for generation requests.
         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.
+        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.
     # SGLang only calls set_prometheus_multiproc_dir() when enable_metrics=True,
     # so MultiProcessCollector will crash without it.
@@ -487,7 +520,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,

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,167 @@ def shutdown(self):
 
     assert calls[0]["worker_id"] == 0
     publisher.cleanup()
+
+
+# ---- per-worker metric gating (embedding vs chat) ----
+
+
+@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

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