feat(sglang): gate chat-shaped Prometheus collectors on embedding worker (DIS-2107)

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

@@ -460,18 +460,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.
@@ -498,7 +531,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