fix(router): preserve per-rank DP capacity [DYN-4273]

components/src/dynamo/planner/connectors/mdc.py

@@ -145,6 +145,9 @@ def worker_info_from_mdc(
     if context_length is None:
         context_length = card.get("architectural_max_context_length")
 
+    # TODO(rank-aware-kv-capacity): propagate capacity provenance into WorkerInfo. Only an exact
+    # or conservative scalar is safe for Planner's scale-down feasibility check; an aggregate
+    # mean must remain explicitly approximate rather than masquerade as a worker minimum.
     return WorkerInfo(
         k8s_name=k8s_name,
         component_name=component_name,

components/src/dynamo/sglang/capacity.py

@@ -24,7 +24,10 @@ def local_dp_rank_bounds(server_args: Any) -> tuple[int, int]:
     nnodes = getattr(server_args, "nnodes", 1) or 1
     node_rank = getattr(server_args, "node_rank", 0) or 0
 
-    if enable_dp_attention and dp_size > 1:
+    if not enable_dp_attention:
+        return 0, dp_size
+
+    if dp_size > 1:
         local_dp_size = dp_size // nnodes if nnodes > 0 else dp_size
         start_dp_rank = node_rank * local_dp_size
         return start_dp_rank, start_dp_rank + local_dp_size
@@ -40,21 +43,15 @@ def publishes_kv_events(server_args: Any) -> bool:
     to one logical worker. That only yields a unique key per node while DP
     attention gives each node a distinct rank slice.
 
-    Without DP attention, ``local_dp_rank_bounds`` returns ``[0, 1)`` on every
-    node. Every node of a multinode gang would therefore advertise the same
-    ``(leader_worker_id, 0)`` source. The frontend marks that key ambiguous and
-    never activates the direct-ZMQ ingress.
-
-    Only the leader owns the single logical rank in TP-only mode. SGLang emits
-    radix-cache events from the rank-0 scheduler, so non-leader sockets have
-    nothing distinct to contribute.
+    Pure DP is single-node in SGLang, so its leader publishes every replica's
+    distinct rank. Only the leader owns the single logical rank in multinode
+    TP-only mode.
     """
     dp_size = getattr(server_args, "dp_size", 1) or 1
     enable_dp_attention = getattr(server_args, "enable_dp_attention", False)
     nnodes = getattr(server_args, "nnodes", 1) or 1
     node_rank = getattr(server_args, "node_rank", 0) or 0
 
-    # Mirrors the branch in local_dp_rank_bounds: per-node distinct slices.
     if enable_dp_attention and dp_size > 1:
         return True
 
@@ -72,7 +69,8 @@ def per_rank_max_running_requests(server_args: Any) -> int | None:
         return None
 
     dp_size = getattr(server_args, "dp_size", 1) or 1
-    if dp_size <= 1:
+    enable_dp_attention = getattr(server_args, "enable_dp_attention", False)
+    if dp_size <= 1 or not enable_dp_attention:
         return max_running_requests
 
     return max_running_requests // dp_size

components/src/dynamo/sglang/publisher.py

@@ -293,10 +293,11 @@ def init_engine_metrics_publish(self) -> None:
     def init_kv_event_publish(self) -> List[KvEventPublisher]:
         """Initialize KV event publisher(s) if configured.
 
-        For DP attention mode, creates one subscriber per LOCAL DP rank port.
-        Each SGLang scheduler in DP attention mode publishes to a unique port
-        (base_port + attn_dp_rank). In multi-node setups, each node's dynamo.sglang
-        instance subscribes only to the DP ranks running on that node.
+        Creates one subscriber per local KV-cache rank. Pure DP schedulers use
+        their DP replica rank while DP-attention schedulers use their attention
+        DP rank. Both publish to a unique port derived from the base endpoint.
+        In multi-node DP-attention setups, each node's dynamo.sglang instance
+        subscribes only to the ranks running on that node.
 
         Multi-node handling:
         - Each node runs dynamo.sglang alongside its local SGLang DP ranks
@@ -328,7 +329,7 @@ def init_kv_event_publish(self) -> List[KvEventPublisher]:
             dp_ranks = get_local_dp_rank_range(self.server_args)
             if len(dp_ranks) > 1:
                 logging.info(
-                    "DP attention mode: subscribing to local DP ranks [%d, %d)",
+                    "Subscribing to local DP ranks [%d, %d)",
                     dp_ranks.start,
                     dp_ranks.stop,
                 )

components/src/dynamo/sglang/register.py

@@ -514,6 +514,9 @@ async def get_runtime_config(
         return runtime_config
 
     try:
+        # TODO(rank-aware-kv-capacity): scheduler_infos[0] is only a representative rank.
+        # Collect every declared rank before the create-only MDC registration and publish one
+        # atomic rank-capacity snapshot; the card cannot be enriched after registration.
         scheduler_info = engine._scheduler_init_result.scheduler_infos[0]
         capacity = runtime_capacity(server_args, scheduler_info)
         max_total_tokens = scheduler_info.get("max_total_num_tokens")

components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py

@@ -3,14 +3,14 @@
 
 from __future__ import annotations
 
-import importlib.util
 from types import SimpleNamespace
 
 import pytest
 
 from dynamo.sglang.capacity import (
     local_dp_rank_bounds,
     model_card_dp_rank_bounds,
+    per_rank_max_running_requests,
     publishes_kv_events,
 )
 
@@ -19,10 +19,6 @@
     pytest.mark.sglang,
     pytest.mark.gpu_0,
     pytest.mark.pre_merge,
-    pytest.mark.skipif(
-        importlib.util.find_spec("sglang") is None,
-        reason="sglang not installed in this container",
-    ),
 ]
 
 
@@ -47,6 +43,10 @@ def test_single_node_publishes_kv_events():
     assert publishes_kv_events(_args()) is True
 
 
+def test_single_node_pure_dp_exposes_every_local_rank():
+    assert local_dp_rank_bounds(_args(dp_size=4)) == (0, 4)
+
+
 def test_multinode_without_dp_attention_publishes_only_from_leader():
     """TP-only multinode must advertise one source per logical worker."""
     leader = _args(nnodes=2, node_rank=0)
@@ -75,3 +75,19 @@ def test_dp_size_one_with_dp_attention_still_leader_only():
         publishes_kv_events(_args(enable_dp_attention=True, nnodes=2, node_rank=1))
         is False
     )
+
+
+def test_pure_dp_keeps_per_scheduler_max_running_requests():
+    server_args = _args(dp_size=4, max_running_requests=128)
+
+    assert per_rank_max_running_requests(server_args) == 128
+
+
+def test_dp_attention_splits_global_max_running_requests():
+    server_args = _args(
+        dp_size=4,
+        enable_dp_attention=True,
+        max_running_requests=128,
+    )
+
+    assert per_rank_max_running_requests(server_args) == 32

components/src/dynamo/sglang/tests/test_sglang_publisher.py

@@ -498,6 +498,72 @@ def __init__(self, **kwargs):
     assert {call["kv_block_size"] for call in calls} == {32}
 
 
+def test_init_kv_event_publish_subscribes_to_every_pure_dp_replica(monkeypatch):
+    calls = []
+
+    class FakeKvEventPublisher:
+        def __init__(self, **kwargs):
+            calls.append(kwargs)
+
+        def shutdown(self):
+            pass
+
+    monkeypatch.setattr(publisher_mod, "KvEventPublisher", FakeKvEventPublisher)
+    monkeypatch.setattr(
+        publisher_mod,
+        "get_zmq_socket",
+        lambda *args, **kwargs: SimpleNamespace(close=lambda linger=0: None),
+    )
+    monkeypatch.setattr(publisher_mod, "get_local_ip_auto", lambda: "127.0.0.1")
+    monkeypatch.setattr(
+        publisher_mod,
+        "ZmqEventPublisher",
+        SimpleNamespace(
+            offset_endpoint_port=staticmethod(
+                lambda base_ep, dp_rank: f"tcp://*:{5557 + dp_rank}"
+            )
+        ),
+    )
+
+    server_args = SimpleNamespace(
+        kv_events_config='{"endpoint": "tcp://*:5557"}',
+        page_size=16,
+        dcp_size=1,
+        dp_size=4,
+        enable_dp_attention=False,
+        nnodes=1,
+        node_rank=0,
+    )
+    config = SimpleNamespace(
+        server_args=server_args,
+        dynamo_args=SimpleNamespace(
+            enable_local_indexer=False,
+            kv_state_endpoint=None,
+            use_kv_events=True,
+        ),
+    )
+    publisher = DynamoSglangPublisher(
+        engine=SimpleNamespace(
+            port_args=SimpleNamespace(metrics_ipc_name="ipc://metrics")
+        ),
+        config=config,
+        generate_endpoint=SimpleNamespace(),
+        component_gauges=SimpleNamespace(),
+    )
+
+    publishers = publisher.init_kv_event_publish()
+
+    assert len(publishers) == 4
+    assert [call["dp_rank"] for call in calls] == [0, 1, 2, 3]
+    assert [call["zmq_endpoint"] for call in calls] == [
+        "tcp://127.0.0.1:5557",
+        "tcp://127.0.0.1:5558",
+        "tcp://127.0.0.1:5559",
+        "tcp://127.0.0.1:5560",
+    ]
+    publisher.cleanup()
+
+
 def test_init_kv_event_publish_uses_effective_kv_event_setting():
     server_args = SimpleNamespace(
         kv_events_config='{"publisher": "null", "endpoint": "tcp://*:5557"}',

components/src/dynamo/thunderagent_router/capacity.py

@@ -135,6 +135,9 @@ def _pool_tokens(body: dict, runtime_config: dict) -> Optional[int]:
         ):
             return None
         tokens = int(block_size) * int(total_blocks)
+        # TODO(rank-aware-kv-capacity): resolve device blocks per rank with provenance, and type
+        # native offload as per-rank versus shared before summing it. Do not fan a shared pool out
+        # to every DP rank or use an estimated device value as an exact admission budget.
         offloaded = get_native_offloading_capacity_tokens(
             runtime_config.get("runtime_data", {})
         )

components/src/dynamo/vllm/capacity.py

@@ -29,6 +29,12 @@ def publish_vllm_token_budget(runtime_config: Any, max_model_len: int | None) ->
 def per_rank_kv_blocks(
     total_kv_blocks: int | None, data_parallel_size: int
 ) -> int | None:
+    """Estimate rank capacity from vLLM's process-wide DP aggregate.
+
+    The arithmetic mean assumes homogeneous ranks; exact division does not prove equality.
+    TODO(rank-aware-kv-capacity): consume a per-rank Control response when vLLM exposes one,
+    then publish the rank vector atomically instead of upgrading this quotient to exact data.
+    """
     if total_kv_blocks is None:
         return None
 

components/src/dynamo/vllm/publisher.py

@@ -24,27 +24,38 @@ class DynamoStatLoggerPublisher(StatLoggerBase):
 
     def __init__(
         self,
-        endpoint: Endpoint,
+        endpoint: Optional[Endpoint],
         dp_rank: int = 0,
         component_gauges: Optional[LLMBackendMetrics] = None,
     ) -> None:
         self.inner = WorkerMetricsPublisher()
         self._endpoint = endpoint
+        self._endpoint_task: Optional[asyncio.Task[None]] = None
         self.dp_rank = dp_rank
         self.component_gauges = component_gauges or LLMBackendMetrics()
         self.num_gpu_block = 1
-        # Schedule async endpoint creation
-        self._endpoint_task = asyncio.create_task(self._create_endpoint())
+        if endpoint is not None:
+            self.bind_endpoint(endpoint)
 
-    async def _create_endpoint(self) -> None:
+    async def _create_endpoint(self, endpoint: Endpoint) -> None:
         """Create the NATS endpoint asynchronously."""
         try:
-            await self.inner.create_endpoint(self._endpoint)
+            await self.inner.create_endpoint(endpoint)
             logging.debug("vLLM metrics publisher endpoint created")
         except Exception:
             logging.exception("Failed to create vLLM metrics publisher endpoint")
             raise
 
+    def bind_endpoint(self, endpoint: Endpoint) -> None:
+        if self._endpoint_task is not None:
+            raise RuntimeError("vLLM metrics publisher endpoint is already bound")
+        if self._endpoint is None:
+            # Drop pre-restore samples so init_publish emits into the newly bound
+            # publisher instead of being deduplicated against snapshot state.
+            self.inner = WorkerMetricsPublisher()
+        self._endpoint = endpoint
+        self._endpoint_task = asyncio.create_task(self._create_endpoint(endpoint))
+
     # TODO: Remove this and pass as metadata through shared storage
     def set_num_gpu_block(self, num_blocks: int) -> None:
         self.num_gpu_block = num_blocks
@@ -134,14 +145,14 @@ class StatLoggerFactory:
 
     def __init__(
         self,
-        endpoint: Endpoint,
+        endpoint: Optional[Endpoint],
         component_gauges: Optional[LLMBackendMetrics] = None,
         embedding_worker: bool = False,
     ) -> None:
         self.endpoint = endpoint
         self.component_gauges = component_gauges
         self.embedding_worker = embedding_worker
-        self.created_logger: Optional[DynamoStatLoggerPublisher] = None
+        self.created_loggers: dict[int, DynamoStatLoggerPublisher] = {}
 
     def create_stat_logger(self, dp_rank: int) -> StatLoggerBase:
         # Embedding workers have no KV cache and no scheduler stats worth
@@ -159,18 +170,25 @@ def create_stat_logger(self, dp_rank: int) -> StatLoggerBase:
             dp_rank=dp_rank,
             component_gauges=self.component_gauges,
         )
-        self.created_logger = logger
+        self.created_loggers[dp_rank] = logger
 
         return logger
 
     def __call__(self, vllm_config: VllmConfig, dp_rank: int) -> StatLoggerBase:
         return self.create_stat_logger(dp_rank=dp_rank)
 
+    def bind_endpoint(self, endpoint: Endpoint) -> None:
+        if self.endpoint is not None:
+            raise RuntimeError("vLLM stat logger endpoint is already bound")
+        self.endpoint = endpoint
+        for logger in self.created_loggers.values():
+            logger.bind_endpoint(endpoint)
+
     # TODO Remove once we publish metadata to shared storage
     def set_num_gpu_blocks_all(self, num_blocks: int) -> None:
-        if self.created_logger:
-            self.created_logger.set_num_gpu_block(num_blocks)
+        for logger in self.created_loggers.values():
+            logger.set_num_gpu_block(num_blocks)
 
     def init_publish(self) -> None:
-        if self.created_logger:
-            self.created_logger.init_publish()
+        for logger in self.created_loggers.values():
+            logger.init_publish()

components/src/dynamo/vllm/snapshot.py

@@ -14,15 +14,16 @@
 
 from .args import Config
 from .handlers import VllmEnginePauseController
-from .worker_factory import EngineSetupResult
+from .publisher import StatLoggerFactory
+from .worker_factory import EngineSetupResult, SnapshotEngineSetupResult
 
 logger = logging.getLogger(__name__)
 
 
 async def prepare_snapshot_engine(
     config: Config,
-    setup_vllm_engine: Callable[[Config], EngineSetupResult],
-) -> EngineSnapshotController[EngineSetupResult] | None:
+    setup_vllm_engine: Callable[..., EngineSetupResult],
+) -> EngineSnapshotController[SnapshotEngineSetupResult] | None:
     snapshot_config = SnapshotConfig.from_env()
     if snapshot_config is None:
         return None
@@ -38,7 +39,11 @@ async def prepare_snapshot_engine(
     logger.info("Snapshot mode enabled (watcher-driven signals)")
     config.engine_args.enable_sleep_mode = True
 
-    engine = setup_vllm_engine(config)
+    stat_logger_factory = StatLoggerFactory(
+        endpoint=None,
+        embedding_worker=config.embedding_worker,
+    )
+    engine = setup_vllm_engine(config, stat_logger_factory)
     # Decide before the first pause: reaching this at pause time would raise
     # after sleep() had already released the engine's memory.
     checkpoint_hooks = all(
@@ -54,7 +59,7 @@ async def prepare_snapshot_engine(
 
     gc.collect()
     snapshot_controller = EngineSnapshotController(
-        engine=engine,
+        engine=(engine, stat_logger_factory),
         pause_controller=VllmEnginePauseController(
             engine[0],
             prepare_for_process_checkpoint=checkpoint_hooks,