From 00220774a0cccb802e71d4e107d0ad6357cbe0fa Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 3 Sep 2026 09:31:15 -0700 Subject: [PATCH 1/7] fix(router): preserve per-rank DP capacity Signed-off-by: PeaBrane --- components/src/dynamo/sglang/capacity.py | 20 +- components/src/dynamo/sglang/publisher.py | 11 +- .../tests/test_sglang_local_dp_ranks.py | 26 +- .../sglang/tests/test_sglang_publisher.py | 77 ++++++ components/src/dynamo/vllm/publisher.py | 40 ++- components/src/dynamo/vllm/snapshot.py | 15 +- .../dynamo/vllm/tests/test_vllm_publisher.py | 75 +++++- .../dynamo/vllm/tests/test_vllm_snapshot.py | 50 ++++ .../vllm/tests/test_vllm_worker_factory.py | 14 +- components/src/dynamo/vllm/worker_factory.py | 52 ++-- lib/llm/src/discovery/worker_monitor.rs | 238 +++++++++++++++--- lib/llm/src/kv_router/publisher/tests.rs | 57 +++++ .../src/kv_router/publisher/worker_metrics.rs | 129 ++++++++-- lib/sidecar/sglang/src/engine.rs | 46 +++- lib/sidecar/vllm/src/model.rs | 28 ++- lib/sidecar/vllm/src/tests.rs | 46 +++- 16 files changed, 798 insertions(+), 126 deletions(-) create mode 100644 components/src/dynamo/vllm/tests/test_vllm_snapshot.py diff --git a/components/src/dynamo/sglang/capacity.py b/components/src/dynamo/sglang/capacity.py index 07dd7c20bbc8..e66abf7325f7 100644 --- a/components/src/dynamo/sglang/capacity.py +++ b/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 diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index 54892c47f468..38e839529227 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/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, ) diff --git a/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py b/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py index a4bd56ad50cc..788b20ee1adb 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py +++ b/components/src/dynamo/sglang/tests/test_sglang_local_dp_ranks.py @@ -3,7 +3,6 @@ from __future__ import annotations -import importlib.util from types import SimpleNamespace import pytest @@ -11,6 +10,7 @@ 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 diff --git a/components/src/dynamo/sglang/tests/test_sglang_publisher.py b/components/src/dynamo/sglang/tests/test_sglang_publisher.py index 7b105cbd257f..6aa1c0ddf2f4 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_publisher.py +++ b/components/src/dynamo/sglang/tests/test_sglang_publisher.py @@ -39,6 +39,17 @@ def test_get_local_dp_rank_range_defaults_to_rank_zero(): assert list(get_local_dp_rank_range(server_args)) == [0] +def test_get_local_dp_rank_range_includes_every_pure_dp_replica(): + server_args = SimpleNamespace( + dp_size=4, + enable_dp_attention=False, + nnodes=1, + node_rank=0, + ) + + assert list(get_local_dp_rank_range(server_args)) == [0, 1, 2, 3] + + def test_get_local_dp_rank_range_respects_multinode_dp_attention(): server_args = SimpleNamespace( dp_size=8, @@ -498,6 +509,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"}', diff --git a/components/src/dynamo/vllm/publisher.py b/components/src/dynamo/vllm/publisher.py index ff1113b307d7..8f697d15b67d 100644 --- a/components/src/dynamo/vllm/publisher.py +++ b/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,7 +145,7 @@ class StatLoggerFactory: def __init__( self, - endpoint: Endpoint, + endpoint: Optional[Endpoint], component_gauges: Optional[LLMBackendMetrics] = None, embedding_worker: bool = False, ) -> None: @@ -142,6 +153,7 @@ def __init__( 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 @@ -160,17 +172,25 @@ def create_stat_logger(self, dp_rank: int) -> StatLoggerBase: 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() diff --git a/components/src/dynamo/vllm/snapshot.py b/components/src/dynamo/vllm/snapshot.py index c5e636fa6162..2a7509edbf7d 100644 --- a/components/src/dynamo/vllm/snapshot.py +++ b/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, diff --git a/components/src/dynamo/vllm/tests/test_vllm_publisher.py b/components/src/dynamo/vllm/tests/test_vllm_publisher.py index 70d8a5716339..29354ed80f10 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_publisher.py +++ b/components/src/dynamo/vllm/tests/test_vllm_publisher.py @@ -12,7 +12,7 @@ """ from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock import pytest @@ -125,3 +125,76 @@ def _fake_publisher(*args, **kwargs): assert constructed[0]["endpoint"] is endpoint assert constructed[0]["dp_rank"] == 3 assert constructed[0]["component_gauges"] is component_gauges + + +def test_factory_initializes_every_dp_rank_logger(monkeypatch): + loggers = [] + + def _fake_publisher(*args, **kwargs): + logger = Mock(spec=DynamoStatLoggerPublisher) + loggers.append(logger) + return logger + + monkeypatch.setattr(publisher_mod, "DynamoStatLoggerPublisher", _fake_publisher) + + factory = StatLoggerFactory( + endpoint=SimpleNamespace(), component_gauges=SimpleNamespace() + ) + for dp_rank in range(3): + factory.create_stat_logger(dp_rank=dp_rank) + + factory.set_num_gpu_blocks_all(4096) + factory.init_publish() + + assert factory.created_logger is loggers[-1] + assert factory.created_loggers == dict(enumerate(loggers)) + for logger in loggers: + logger.set_num_gpu_block.assert_called_once_with(4096) + logger.init_publish.assert_called_once_with() + + +def test_factory_binds_deferred_endpoint_to_every_dp_rank_logger(monkeypatch): + loggers = [] + + def _fake_publisher(*args, **kwargs): + assert kwargs["endpoint"] is None + logger = Mock(spec=DynamoStatLoggerPublisher) + loggers.append(logger) + return logger + + monkeypatch.setattr(publisher_mod, "DynamoStatLoggerPublisher", _fake_publisher) + + factory = StatLoggerFactory(endpoint=None, component_gauges=SimpleNamespace()) + factory.create_stat_logger(dp_rank=0) + factory.create_stat_logger(dp_rank=1) + + endpoint = SimpleNamespace() + factory.bind_endpoint(endpoint) + + assert factory.endpoint is endpoint + for logger in loggers: + logger.bind_endpoint.assert_called_once_with(endpoint) + + +@pytest.mark.asyncio +async def test_deferred_logger_starts_with_fresh_metrics_state(monkeypatch): + publishers = [Mock(create_endpoint=AsyncMock()), Mock(create_endpoint=AsyncMock())] + monkeypatch.setattr( + publisher_mod, "WorkerMetricsPublisher", Mock(side_effect=publishers) + ) + + logger = DynamoStatLoggerPublisher( + endpoint=None, + component_gauges=SimpleNamespace(), + ) + logger.inner.publish(dp_rank=0, kv_used_blocks=7) + + endpoint = SimpleNamespace() + logger.bind_endpoint(endpoint) + assert logger.inner is publishers[1] + + assert logger._endpoint_task is not None + await logger._endpoint_task + publishers[0].publish.assert_called_once_with(dp_rank=0, kv_used_blocks=7) + publishers[0].create_endpoint.assert_not_called() + publishers[1].create_endpoint.assert_awaited_once_with(endpoint) diff --git a/components/src/dynamo/vllm/tests/test_vllm_snapshot.py b/components/src/dynamo/vllm/tests/test_vllm_snapshot.py new file mode 100644 index 000000000000..c5f36e814c32 --- /dev/null +++ b/components/src/dynamo/vllm/tests/test_vllm_snapshot.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from dynamo.vllm import snapshot +from dynamo.vllm.publisher import StatLoggerFactory + +pytestmark = [ + pytest.mark.unit, + pytest.mark.vllm, + pytest.mark.core, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + + +@pytest.mark.asyncio +async def test_prepare_snapshot_preserves_engine_stat_loggers(monkeypatch): + snapshot_config = Mock() + snapshot_config.run_lifecycle = AsyncMock(return_value=True) + monkeypatch.setattr(snapshot.SnapshotConfig, "from_env", lambda: snapshot_config) + monkeypatch.setattr(snapshot, "configure_snapshot_capture_env", Mock()) + + engine_client = Mock(checkpoint_prepare=Mock(), checkpoint_restore=Mock()) + engine_setup = ( + engine_client, + Mock(), + Mock(), + Mock(), + Mock(), + ) + setup_vllm_engine = Mock(return_value=engine_setup) + config = SimpleNamespace( + headless=False, + embedding_worker=False, + engine_args=SimpleNamespace(enable_sleep_mode=False), + ) + + controller = await snapshot.prepare_snapshot_engine(config, setup_vllm_engine) + + assert controller is not None + restored_engine, stat_logger_factory = controller.engine + assert restored_engine is engine_setup + assert isinstance(stat_logger_factory, StatLoggerFactory) + setup_vllm_engine.assert_called_once_with(config, stat_logger_factory) + assert config.engine_args.enable_sleep_mode is True diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py b/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py index 3e47d13be394..4754be114c4f 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py @@ -15,6 +15,7 @@ from dynamo.vllm.constants import DisaggregationMode from dynamo.vllm.worker_factory import ( EngineSetupResult, + SnapshotEngineSetupResult, WorkerFactory, _DecodeWorkerLifecycle, _wait_and_load_benchmark, @@ -849,11 +850,14 @@ async def test_passes_snapshot_engine(self, factory: WorkerFactory) -> None: runtime = Mock() shutdown_event = asyncio.Event() shutdown_endpoints: list = [] - snapshot_engine: EngineSetupResult = ( - Mock(), - Mock(), - Mock(), - "/tmp/prometheus", + snapshot_engine: SnapshotEngineSetupResult = ( + ( + Mock(), + Mock(), + Mock(), + "/tmp/prometheus", + Mock(), + ), Mock(), ) diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 2ff0fbd97b96..671c5d6c28e4 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -67,6 +67,7 @@ # have no KV cache / scheduler gauges, so setup_vllm_engine() skips the # LLMBackendMetrics registration there. EngineSetupResult = tuple[AsyncLLM, VllmConfig, Any, Any, Optional[LLMBackendMetrics]] +SnapshotEngineSetupResult = tuple[EngineSetupResult, StatLoggerFactory] def _benchmark_rank_path(base_path: Path, dp_rank: int) -> Path: @@ -644,7 +645,7 @@ async def create( config: Config, shutdown_event: asyncio.Event, shutdown_endpoints: list, - snapshot_engine: Optional[EngineSetupResult] = None, + snapshot_engine: Optional[SnapshotEngineSetupResult] = None, ) -> None: """Create the appropriate multimodal worker based on config flags.""" @@ -706,7 +707,7 @@ async def _create_realtime_worker( config: Config, shutdown_event: asyncio.Event, shutdown_endpoints: list, - snapshot_engine: Optional[EngineSetupResult] = None, + snapshot_engine: Optional[SnapshotEngineSetupResult] = None, ) -> None: """Initialize an aggregated vLLM realtime worker.""" del shutdown_event # Connection cancellation is carried by Dynamo Context. @@ -718,18 +719,16 @@ async def _create_realtime_worker( fpm_worker_id = str(generate_endpoint.connection_id()) if snapshot_engine is not None: + engine_setup, factory = snapshot_engine ( engine_client, vllm_config, _default_sampling_params, prometheus_temp_dir, - component_gauges, - ) = snapshot_engine + _component_gauges, + ) = engine_setup os.environ[ENV_FPM_WORKER_ID] = fpm_worker_id - factory = StatLoggerFactory( - endpoint=generate_endpoint, - component_gauges=component_gauges, - ) + factory.bind_endpoint(generate_endpoint) else: factory = StatLoggerFactory(endpoint=generate_endpoint) ( @@ -737,7 +736,7 @@ async def _create_realtime_worker( vllm_config, _default_sampling_params, prometheus_temp_dir, - component_gauges, + _component_gauges, ) = self.setup_vllm_engine( config, factory, @@ -1097,7 +1096,7 @@ async def _create_decode_worker( config: Config, shutdown_event: asyncio.Event, shutdown_endpoints: list, # mutated in place - snapshot_engine: Optional[EngineSetupResult] = None, + snapshot_engine: Optional[SnapshotEngineSetupResult] = None, ) -> None: """ Instantiate and serve @@ -1121,7 +1120,7 @@ async def _run_decode_worker( config: Config, shutdown_event: asyncio.Event, shutdown_endpoints: list, # mutated in place - snapshot_engine: Optional[EngineSetupResult], + snapshot_engine: Optional[SnapshotEngineSetupResult], lifecycle: _DecodeWorkerLifecycle, ) -> None: """Initialize and serve a decode worker.""" @@ -1173,19 +1172,16 @@ async def _run_decode_worker( # Use pre-created engine if provided (checkpoint mode), otherwise create new fpm_worker_id = str(generate_endpoint.connection_id()) if snapshot_engine is not None: + engine_setup, factory = snapshot_engine ( engine_client, vllm_config, default_sampling_params, prometheus_temp_dir, - component_gauges, - ) = snapshot_engine + _component_gauges, + ) = engine_setup os.environ[ENV_FPM_WORKER_ID] = fpm_worker_id - # Factory is created after unpack so component_gauges is available - factory = StatLoggerFactory( - endpoint=generate_endpoint, - component_gauges=component_gauges, - ) + factory.bind_endpoint(generate_endpoint) else: # Factory is created without component_gauges; setup_vllm_engine() will # create the gauges after setup_multiprocess_prometheus() and set them @@ -1198,7 +1194,7 @@ async def _run_decode_worker( vllm_config, default_sampling_params, prometheus_temp_dir, - component_gauges, + _component_gauges, ) = self.setup_vllm_engine(config, factory, fpm_worker_id=fpm_worker_id) lifecycle.engine_client = engine_client lifecycle.vllm_config = vllm_config @@ -1432,7 +1428,7 @@ async def _create_prefill_worker( config: Config, shutdown_event: asyncio.Event, shutdown_endpoints: list, - snapshot_engine: Optional[EngineSetupResult] = None, + snapshot_engine: Optional[SnapshotEngineSetupResult] = None, ) -> None: try: await self._run_prefill_worker( @@ -1452,7 +1448,7 @@ async def _run_prefill_worker( config: Config, shutdown_event: asyncio.Event, shutdown_endpoints: list, # mutated in place - snapshot_engine: Optional[EngineSetupResult] = None, + snapshot_engine: Optional[SnapshotEngineSetupResult] = None, ) -> None: """ Instantiate and serve @@ -1487,14 +1483,17 @@ async def _run_prefill_worker( # Use pre-created engine if provided (checkpoint mode), otherwise create new fpm_worker_id = str(generate_endpoint.connection_id()) + snapshot_factory: Optional[StatLoggerFactory] = None if snapshot_engine is not None: + engine_setup, snapshot_factory = snapshot_engine ( engine_client, vllm_config, default_sampling_params, prometheus_temp_dir, _component_gauges, - ) = snapshot_engine + ) = engine_setup + snapshot_factory.bind_endpoint(generate_endpoint) # TODO: The scheduler in the child process still has worker_id="" # because the engine was forked before the runtime existed. # Propagating the new ID to the child requires shared memory or @@ -1510,6 +1509,15 @@ async def _run_prefill_worker( ) = self.setup_vllm_engine(config, fpm_worker_id=fpm_worker_id) await configure_kv_event_block_size(engine_client, vllm_config) + if snapshot_factory is not None: + _, dp_size = get_dp_range_for_worker(vllm_config) + per_rank_num_gpu_blocks = per_rank_kv_blocks( + vllm_config.cache_config.num_gpu_blocks, + dp_size, + ) + snapshot_factory.set_num_gpu_blocks_all(per_rank_num_gpu_blocks or 0) + snapshot_factory.init_publish() + encode_worker_client = await self._maybe_get_encode_worker_client( runtime, config ) diff --git a/lib/llm/src/discovery/worker_monitor.rs b/lib/llm/src/discovery/worker_monitor.rs index b983fbbc99b0..2ce354dd2442 100644 --- a/lib/llm/src/discovery/worker_monitor.rs +++ b/lib/llm/src/discovery/worker_monitor.rs @@ -253,10 +253,72 @@ pub struct WorkerLoadState { pub active_prefill_tokens: HashMap, /// max_num_batched_tokens from runtime config (same for all dp_ranks) pub max_num_batched_tokens: HashMap, + /// The current router-visible ranks declared by this worker's runtime config. + /// `None` allows observations received before discovery to remain usable until + /// the runtime config arrives and establishes the authoritative rank set. + declared_dp_ranks: Option>, decode_overload_latches: HashMap, } impl WorkerLoadState { + fn reconcile_runtime_config( + &mut self, + dp_start: u32, + dp_size: u32, + total_kv_blocks: Option, + max_num_batched_tokens: Option, + active_decode_blocks_threshold: Option, + ) -> HashSet { + let dp_end = dp_start.saturating_add(dp_size); + let declared_dp_ranks: HashSet<_> = (dp_start..dp_end).collect(); + + self.active_decode_blocks + .retain(|dp_rank, _| declared_dp_ranks.contains(dp_rank)); + self.kv_used_blocks + .retain(|dp_rank, _| declared_dp_ranks.contains(dp_rank)); + self.active_prefill_tokens + .retain(|dp_rank, _| declared_dp_ranks.contains(dp_rank)); + + self.kv_total_blocks.clear(); + if let Some(total_blocks) = total_kv_blocks { + self.kv_total_blocks.extend( + declared_dp_ranks + .iter() + .map(|&dp_rank| (dp_rank, total_blocks)), + ); + } + + self.max_num_batched_tokens.clear(); + if let Some(max_batched) = max_num_batched_tokens { + self.max_num_batched_tokens.extend( + declared_dp_ranks + .iter() + .map(|&dp_rank| (dp_rank, max_batched)), + ); + } + + self.decode_overload_latches.clear(); + if let Some(threshold) = active_decode_blocks_threshold { + for &dp_rank in &declared_dp_ranks { + self.update_decode_overload_latch( + dp_rank, + self.active_decode_blocks.get(&dp_rank).copied(), + self.kv_used_blocks.get(&dp_rank).copied(), + threshold, + ); + } + } + + self.declared_dp_ranks = Some(declared_dp_ranks.clone()); + declared_dp_ranks + } + + fn accepts_dp_rank(&self, dp_rank: u32) -> bool { + self.declared_dp_ranks + .as_ref() + .is_none_or(|declared| declared.contains(&dp_rank)) + } + fn is_decode_signal_overloaded( used_blocks: u64, total_blocks: u64, @@ -343,10 +405,14 @@ impl WorkerLoadState { &mut self, observation: LoadObservation, active_decode_blocks_threshold: Option, - ) { + ) -> bool { let (worker, active_decode_blocks, active_prefill_tokens, kv_used_blocks) = observation.parts(); let dp_rank = worker.dp_rank; + if !self.accepts_dp_rank(dp_rank) { + return false; + } + if let Some(active_blocks) = active_decode_blocks { self.active_decode_blocks.insert(dp_rank, active_blocks); } @@ -364,6 +430,7 @@ impl WorkerLoadState { threshold, ); } + true } #[cfg(test)] @@ -404,15 +471,18 @@ impl WorkerLoadState { return false; } - // Get all dp_ranks we know about - let all_dp_ranks: std::collections::HashSet<_> = self - .active_decode_blocks - .keys() - .chain(self.kv_used_blocks.keys()) - .chain(self.decode_overload_latches.keys()) - .chain(self.active_prefill_tokens.keys()) - .copied() - .collect(); + // Once discovery has supplied the runtime config, its rank set is + // authoritative. An expected rank without a load observation is free, + // so one noisy rank cannot exclude the whole worker during startup. + let all_dp_ranks = self.declared_dp_ranks.clone().unwrap_or_else(|| { + self.active_decode_blocks + .keys() + .chain(self.kv_used_blocks.keys()) + .chain(self.decode_overload_latches.keys()) + .chain(self.active_prefill_tokens.keys()) + .copied() + .collect() + }); // If no dp_ranks known, not overloaded if all_dp_ranks.is_empty() { @@ -791,36 +861,37 @@ impl WorkerLoadMonitor for KvWorkerMonitor { overloaded_tracker.remove_workers(&removed_workers); client.clear_overloaded_instances_for_removed(&removed_workers); - // Update worker load states with runtime config values for all dp_ranks - // This ensures we track workers from MDCs even if they don't publish ActiveLoad + let cfg = thresholds.get(); + + // Reconcile worker state to the authoritative rank range from discovery. + // This also makes expected-but-unobserved ranks participate in the + // worker-level "all ranks overloaded" decision. for (lease_id, runtime_config) in runtime_configs.iter() { let mut state = worker_load_states.entry(*lease_id).or_default(); - let dp_start = runtime_config.data_parallel_start_rank; - let dp_end = dp_start + runtime_config.data_parallel_size; - - // Track dp_ranks for this worker (for cleanup when worker disappears) - let dp_ranks_set = known_worker_dp_ranks.entry(*lease_id).or_default(); - for dp_rank in dp_start..dp_end { - dp_ranks_set.insert(dp_rank); - } - - // Populate total_blocks for all dp_ranks (they share the same total) - if let Some(total_blocks) = runtime_config.total_kv_blocks { - for dp_rank in dp_start..dp_end { - state.kv_total_blocks.insert(dp_rank, total_blocks); - } - } + let declared_dp_ranks = state.reconcile_runtime_config( + dp_start, + runtime_config.data_parallel_size, + runtime_config.total_kv_blocks, + runtime_config.max_num_batched_tokens, + cfg.active_decode_blocks_threshold, + ); - // Populate max_num_batched_tokens for all dp_ranks - if let Some(max_batched) = runtime_config.max_num_batched_tokens { - for dp_rank in dp_start..dp_end { - state.max_num_batched_tokens.insert(dp_rank, max_batched); - } + if let Some(previous_dp_ranks) = known_worker_dp_ranks + .insert(*lease_id, declared_dp_ranks.clone()) + { + let removed_dp_ranks: Vec<_> = previous_dp_ranks + .difference(&declared_dp_ranks) + .copied() + .collect(); + cleanup_worker_metrics( + *lease_id, + &removed_dp_ranks, + source.metric_label(), + ); } } - let cfg = thresholds.get(); last_thresholds = cfg.clone(); let overloaded_workers = collect_overloaded_workers(&worker_load_states, &cfg); if overloaded_tracker.replace(overloaded_workers) { @@ -855,6 +926,18 @@ impl WorkerLoadMonitor for KvWorkerMonitor { ); continue; } + if worker_load_states + .get(&worker.worker_id) + .is_some_and(|state| !state.accepts_dp_rank(worker.dp_rank)) + { + tracing::debug!( + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + source = ?source, + "dropping load event outside the worker's current runtime-config rank range" + ); + continue; + } // Track known worker/dp_rank combinations for cleanup known_worker_dp_ranks @@ -936,6 +1019,18 @@ impl WorkerLoadMonitor for KvWorkerMonitor { ); continue; } + if worker_load_states + .get(&worker.worker_id) + .is_some_and(|state| !state.accepts_dp_rank(worker.dp_rank)) + { + tracing::debug!( + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + source = ?source, + "dropping scheduler load outside the worker's current runtime-config rank range" + ); + continue; + } known_worker_dp_ranks .entry(worker.worker_id) @@ -1236,6 +1331,83 @@ mod tests { assert!(state.is_overloaded(Some(0.6), Some(u64::MAX), Some(2.0))); } + #[test] + fn expected_but_unobserved_dp_rank_keeps_worker_available() { + let mut state = WorkerLoadState::default(); + state.reconcile_runtime_config(0, 2, Some(100), Some(1_000), Some(0.6)); + + state.update_from_active_load( + &ActiveLoad { + worker_id: 1, + dp_rank: 0, + active_decode_blocks: None, + active_prefill_tokens: None, + kv_used_blocks: Some(90), + }, + Some(0.6), + ); + assert!(!state.is_overloaded(Some(0.6), None, None)); + + state.update_from_active_load( + &ActiveLoad { + worker_id: 1, + dp_rank: 1, + active_decode_blocks: None, + active_prefill_tokens: None, + kv_used_blocks: Some(90), + }, + Some(0.6), + ); + assert!(state.is_overloaded(Some(0.6), None, None)); + } + + #[test] + fn runtime_config_update_reconciles_rank_range_and_optional_capacity() { + let mut state = WorkerLoadState::default(); + state.reconcile_runtime_config(2, 2, Some(100), Some(1_000), Some(0.6)); + + for dp_rank in 2..4 { + state.update_from_active_load( + &ActiveLoad { + worker_id: 1, + dp_rank, + active_decode_blocks: Some(90), + active_prefill_tokens: Some(900), + kv_used_blocks: Some(90), + }, + Some(0.6), + ); + } + assert!(state.is_overloaded(Some(0.6), None, Some(0.5))); + + let declared = state.reconcile_runtime_config(3, 1, None, None, Some(0.6)); + assert_eq!(declared, HashSet::from([3])); + assert!(!state.active_decode_blocks.contains_key(&2)); + assert!(!state.kv_used_blocks.contains_key(&2)); + assert!(!state.active_prefill_tokens.contains_key(&2)); + assert!(state.kv_total_blocks.is_empty()); + assert!(state.max_num_batched_tokens.is_empty()); + assert!(state.decode_overload_latches.is_empty()); + assert!(!state.is_overloaded(Some(0.6), None, Some(0.5))); + + assert!(!state.apply_load_observation( + LoadObservation::Remote(RemoteActiveLoadSnapshot { + worker: WorkerWithDpRank::new(1, 2), + active_decode_blocks: Some(100), + active_prefill_tokens: Some(1_000), + kv_used_blocks: Some(100), + }), + Some(0.6), + )); + + let declared = state.reconcile_runtime_config(4, 1, Some(100), Some(1_000), Some(0.6)); + assert_eq!(declared, HashSet::from([4])); + assert!(state.active_decode_blocks.is_empty()); + assert!(state.kv_used_blocks.is_empty()); + assert!(state.active_prefill_tokens.is_empty()); + assert!(!state.is_overloaded(Some(0.6), None, Some(0.5))); + } + #[test] fn decode_overload_latch_sets_overloaded_if_any_signal_is_overloaded() { let mut state = WorkerLoadState::default(); diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 8c1e7fc610c6..35ae0fa8cafb 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -2132,6 +2132,63 @@ mod test_event_dedup_filter { } } +#[cfg(test)] +mod worker_metrics_tests { + use std::time::Duration; + + use anyhow::Result; + use dynamo_kv_router::protocols::ActiveLoad; + + use super::super::worker_metrics::{WorkerMetricsPublisher, WorkerMetricsSink}; + + struct ChannelSink(tokio::sync::mpsc::UnboundedSender); + + #[async_trait::async_trait] + impl WorkerMetricsSink for ChannelSink { + async fn publish(&self, active_load: ActiveLoad) -> Result<()> { + self.0 + .send(active_load) + .map_err(|_| anyhow::anyhow!("metrics test channel closed")) + } + } + + #[tokio::test] + async fn publish_debounces_updates_independently_per_rank() { + let publisher = WorkerMetricsPublisher::new().unwrap(); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + publisher.start_metrics_publishing_with(ChannelSink(tx), 42); + + publisher.publish(Some(0), None, Some(100)).unwrap(); + publisher.publish(Some(1), None, Some(200)).unwrap(); + publisher.publish(Some(0), None, Some(300)).unwrap(); + + let mut published = Vec::new(); + for _ in 0..2 { + published.push( + tokio::time::timeout(Duration::from_millis(100), rx.recv()) + .await + .expect("timed out waiting for rank metrics") + .expect("metrics publishing task stopped"), + ); + } + published.sort_unstable_by_key(|load| load.dp_rank); + + assert_eq!(published[0].worker_id, 42); + assert_eq!(published[0].dp_rank, 0); + assert_eq!(published[0].kv_used_blocks, Some(300)); + assert_eq!(published[1].worker_id, 42); + assert_eq!(published[1].dp_rank, 1); + assert_eq!(published[1].kv_used_blocks, Some(200)); + + assert!( + tokio::time::timeout(Duration::from_millis(10), rx.recv()) + .await + .is_err(), + "same-rank updates should be coalesced" + ); + } +} + #[cfg(all(test, feature = "integration"))] mod test_integration_publisher { use super::*; diff --git a/lib/llm/src/kv_router/publisher/worker_metrics.rs b/lib/llm/src/kv_router/publisher/worker_metrics.rs index a38e9c5b5a47..1e6145257632 100644 --- a/lib/llm/src/kv_router/publisher/worker_metrics.rs +++ b/lib/llm/src/kv_router/publisher/worker_metrics.rs @@ -1,6 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; +use std::time::Duration; + use anyhow::Result; use dynamo_kv_router::protocols::{ActiveLoad, DpRank}; @@ -10,6 +13,8 @@ use dynamo_runtime::transports::event_plane::EventPublisher; use crate::kv_router::KV_METRICS_SUBJECT; +const PUBLISH_DEBOUNCE: Duration = Duration::from_millis(1); + #[derive(Debug, Clone, Default, PartialEq)] struct WorkerMetrics { dp_rank: DpRank, @@ -17,14 +22,86 @@ struct WorkerMetrics { kv_used_blocks: Option, } +struct PendingMetrics { + metrics: WorkerMetrics, + deadline: tokio::time::Instant, +} + +struct WorkerMetricsDebouncer { + debounce: Duration, + last_metrics: HashMap, + pending: HashMap, +} + +impl WorkerMetricsDebouncer { + fn new(debounce: Duration) -> Self { + Self { + debounce, + last_metrics: HashMap::new(), + pending: HashMap::new(), + } + } + + fn observe( + &mut self, + metrics_by_rank: &HashMap, + now: tokio::time::Instant, + ) { + for (&dp_rank, metrics) in metrics_by_rank { + if self.last_metrics.get(&dp_rank) == Some(metrics) { + continue; + } + + self.last_metrics.insert(dp_rank, metrics.clone()); + self.pending.insert( + dp_rank, + PendingMetrics { + metrics: metrics.clone(), + deadline: now + self.debounce, + }, + ); + } + } + + fn next_deadline(&self) -> Option { + self.pending.values().map(|pending| pending.deadline).min() + } + + fn take_due(&mut self, now: tokio::time::Instant) -> Vec { + let due_ranks = self + .pending + .iter() + .filter_map(|(&dp_rank, pending)| (pending.deadline <= now).then_some(dp_rank)) + .collect::>(); + + due_ranks + .into_iter() + .filter_map(|dp_rank| self.pending.remove(&dp_rank)) + .map(|pending| pending.metrics) + .collect() + } +} + +#[async_trait::async_trait] +pub(super) trait WorkerMetricsSink: Send + 'static { + async fn publish(&self, active_load: ActiveLoad) -> Result<()>; +} + +#[async_trait::async_trait] +impl WorkerMetricsSink for EventPublisher { + async fn publish(&self, active_load: ActiveLoad) -> Result<()> { + EventPublisher::publish(self, &active_load).await + } +} + pub struct WorkerMetricsPublisher { - tx: tokio::sync::watch::Sender, - rx: tokio::sync::watch::Receiver, + tx: tokio::sync::watch::Sender>, + rx: tokio::sync::watch::Receiver>, } impl WorkerMetricsPublisher { pub fn new() -> Result { - let (tx, rx) = tokio::sync::watch::channel(WorkerMetrics::default()); + let (tx, rx) = tokio::sync::watch::channel(HashMap::new()); Ok(Self { tx, rx }) } @@ -49,9 +126,15 @@ impl WorkerMetricsPublisher { metrics.active_decode_blocks, metrics.kv_used_blocks ); - self.tx - .send(metrics) - .map_err(|_| anyhow::anyhow!("metrics channel closed")) + self.tx.send_if_modified(|metrics_by_rank| { + if metrics_by_rank.get(&metrics.dp_rank) == Some(&metrics) { + return false; + } + + metrics_by_rank.insert(metrics.dp_rank, metrics); + true + }); + Ok(()) } pub async fn create_endpoint(&self, endpoint: Endpoint) -> Result<()> { @@ -62,12 +145,18 @@ impl WorkerMetricsPublisher { } pub(super) fn start_metrics_publishing(&self, event_publisher: EventPublisher, worker_id: u64) { + self.start_metrics_publishing_with(event_publisher, worker_id); + } + + pub(super) fn start_metrics_publishing_with(&self, sink: S, worker_id: u64) + where + S: WorkerMetricsSink, + { let metrics_rx = self.rx.clone(); tokio::spawn(async move { let mut rx = metrics_rx; - let mut last_metrics: Option = None; - let mut pending_publish: Option = None; + let mut debouncer = WorkerMetricsDebouncer::new(PUBLISH_DEBOUNCE); let publish_timer = tokio::time::sleep(tokio::time::Duration::ZERO); tokio::pin!(publish_timer); @@ -81,20 +170,14 @@ impl WorkerMetricsPublisher { break; } - let metrics = rx.borrow_and_update().clone(); - if last_metrics.as_ref() == Some(&metrics) { - continue; + let now = tokio::time::Instant::now(); + debouncer.observe(&rx.borrow_and_update(), now); + if let Some(deadline) = debouncer.next_deadline() { + publish_timer.as_mut().reset(deadline); } - - pending_publish = Some(metrics.clone()); - last_metrics = Some(metrics); - publish_timer.as_mut().reset( - tokio::time::Instant::now() - + tokio::time::Duration::from_millis(1) - ); } - _ = &mut publish_timer, if pending_publish.is_some() => { - if let Some(metrics) = pending_publish.take() { + _ = &mut publish_timer, if debouncer.next_deadline().is_some() => { + for metrics in debouncer.take_due(tokio::time::Instant::now()) { let active_load = ActiveLoad { worker_id, dp_rank: metrics.dp_rank, @@ -103,10 +186,14 @@ impl WorkerMetricsPublisher { kv_used_blocks: metrics.kv_used_blocks, }; - if let Err(e) = event_publisher.publish(&active_load).await { + if let Err(e) = sink.publish(active_load).await { tracing::warn!("Failed to publish metrics: {}", e); } } + + if let Some(deadline) = debouncer.next_deadline() { + publish_timer.as_mut().reset(deadline); + } } } } diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index 4496976cf300..dfe86fe4e116 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -790,9 +790,14 @@ fn build_engine_config( let dp_size = client::json_u32(&discovery.server_info, "dp_size") .unwrap_or(1) .max(1); + let enable_dp_attention = discovery + .server_info + .get("enable_dp_attention") + .and_then(Value::as_bool) + .unwrap_or(false); let max_num_seqs = client::json_u64(&discovery.server_info, "max_running_requests").map(|value| { - if dp_size > 1 { + if enable_dp_attention && dp_size > 1 { value / u64::from(dp_size) } else { value @@ -801,11 +806,6 @@ fn build_engine_config( let max_num_batched_tokens = client::json_u64(&discovery.server_info, "max_prefill_tokens").or(max_total_tokens); - let enable_dp_attention = discovery - .server_info - .get("enable_dp_attention") - .and_then(Value::as_bool) - .unwrap_or(false); let (data_parallel_start_rank, data_parallel_size) = if enable_dp_attention && dp_size > 1 { // Native gRPC is exposed by the rank-zero frontend for the complete // multi-node SGLang endpoint, so one sidecar registers every DP rank. @@ -952,6 +952,40 @@ mod tests { assert_eq!(registration.data_parallel_size, Some(16)); } + #[test] + fn pure_dp_preserves_per_rank_max_num_seqs() { + let config = build_engine_config( + &discovery(json!({ + "dp_size": 4, + "enable_dp_attention": false, + "max_running_requests": 256, + })), + DisaggregationMode::Decode, + None, + None, + ) + .unwrap(); + + assert_eq!(config.llm.unwrap().max_num_seqs, Some(256)); + } + + #[test] + fn attention_dp_normalizes_aggregate_max_num_seqs_per_rank() { + let config = build_engine_config( + &discovery(json!({ + "dp_size": 4, + "enable_dp_attention": true, + "max_running_requests": 256, + })), + DisaggregationMode::Decode, + None, + None, + ) + .unwrap(); + + assert_eq!(config.llm.unwrap().max_num_seqs, Some(64)); + } + #[test] fn dcp_registers_logical_kv_block_size() { let config = build_engine_config( diff --git a/lib/sidecar/vllm/src/model.rs b/lib/sidecar/vllm/src/model.rs index 2f4bc944e34c..e8c6383693a1 100644 --- a/lib/sidecar/vllm/src/model.rs +++ b/lib/sidecar/vllm/src/model.rs @@ -145,7 +145,7 @@ impl DiscoveredModel { llm: Some(LlmRegistration { context_length: nonzero(self.server.max_model_len), kv_cache_block_size: nonzero(self.server.kv_block_size), - total_kv_blocks: nonzero(self.server.total_kv_blocks), + total_kv_blocks: self.total_kv_blocks_per_rank(), max_num_seqs: nonzero(self.server.max_running_requests), max_num_batched_tokens: nonzero(self.server.max_batched_tokens), data_parallel_size: parallelism @@ -162,6 +162,32 @@ impl DiscoveredModel { .as_ref() .map_or(1, |parallelism| parallelism.data_parallel_size) } + + fn total_kv_blocks_per_rank(&self) -> Option { + let total_kv_blocks = nonzero(self.server.total_kv_blocks)?; + let data_parallel_size = u64::from(self.data_parallel_size()); + let per_rank = total_kv_blocks / data_parallel_size; + + if per_rank == 0 { + tracing::warn!( + total_kv_blocks, + data_parallel_size, + "vLLM reported fewer total KV blocks than DP ranks; publishing one block per rank" + ); + return Some(1); + } + + if total_kv_blocks % data_parallel_size != 0 { + tracing::warn!( + total_kv_blocks, + data_parallel_size, + per_rank, + "vLLM aggregate KV blocks are not divisible by DP ranks; publishing floor per-rank capacity" + ); + } + + Some(per_rank) + } } fn required(field: &str, value: String) -> Result { diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 0132636a3e83..1a2a754538bc 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -865,6 +865,50 @@ fn discovery_rejects_incompatible_model_metadata() { } } +#[test] +fn engine_config_normalizes_total_kv_blocks_per_dp_rank() { + for (data_parallel_size, aggregate_blocks, expected_per_rank_blocks) in + [(1, 2048, 2048), (2, 4096, 2048), (4, 8192, 2048)] + { + let mut server = server_info(); + server + .parallelism + .as_mut() + .expect("parallelism metadata") + .data_parallel_size = data_parallel_size; + server.total_kv_blocks = aggregate_blocks; + + let model = + DiscoveredModel::from_proto(model_info(), server).expect("valid discovery metadata"); + let registration = model.engine_config().llm.expect("LLM registration"); + + assert_eq!( + registration.total_kv_blocks, + Some(expected_per_rank_blocks), + "DP size {data_parallel_size}" + ); + } +} + +#[test] +fn engine_config_handles_inexact_aggregate_kv_capacity_conservatively() { + for (aggregate_blocks, expected_per_rank_blocks) in + [(0, None), (1, Some(1)), (4097, Some(2048))] + { + let mut server = server_info(); + server.total_kv_blocks = aggregate_blocks; + + let model = + DiscoveredModel::from_proto(model_info(), server).expect("valid discovery metadata"); + let registration = model.engine_config().llm.expect("LLM registration"); + + assert_eq!( + registration.total_kv_blocks, expected_per_rank_blocks, + "aggregate blocks {aggregate_blocks}" + ); + } +} + #[tokio::test] async fn startup_rejects_model_identity_change_after_bootstrap() { let server = FakeServer::start(FakeVllm::default()).await; @@ -902,7 +946,7 @@ async fn aggregated_generation_converts_request_stream_and_usage() { let registration = config.llm.expect("LLM registration"); assert_eq!(registration.context_length, Some(8192)); assert_eq!(registration.kv_cache_block_size, Some(16)); - assert_eq!(registration.total_kv_blocks, Some(4096)); + assert_eq!(registration.total_kv_blocks, Some(2048)); assert_eq!(registration.max_num_seqs, Some(128)); assert_eq!(registration.max_num_batched_tokens, Some(2048)); assert_eq!(registration.data_parallel_size, Some(2)); From 76762aa8cb6e8e0249de5bb70a83c6ef6df48fd0 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 3 Sep 2026 09:41:05 -0700 Subject: [PATCH 2/7] test(router): trim low-signal DP coverage Signed-off-by: PeaBrane --- .../src/dynamo/sglang/tests/test_sglang_publisher.py | 11 ----------- lib/sidecar/vllm/src/tests.rs | 6 ++---- 2 files changed, 2 insertions(+), 15 deletions(-) diff --git a/components/src/dynamo/sglang/tests/test_sglang_publisher.py b/components/src/dynamo/sglang/tests/test_sglang_publisher.py index 6aa1c0ddf2f4..9024ef5a5c30 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_publisher.py +++ b/components/src/dynamo/sglang/tests/test_sglang_publisher.py @@ -39,17 +39,6 @@ def test_get_local_dp_rank_range_defaults_to_rank_zero(): assert list(get_local_dp_rank_range(server_args)) == [0] -def test_get_local_dp_rank_range_includes_every_pure_dp_replica(): - server_args = SimpleNamespace( - dp_size=4, - enable_dp_attention=False, - nnodes=1, - node_rank=0, - ) - - assert list(get_local_dp_rank_range(server_args)) == [0, 1, 2, 3] - - def test_get_local_dp_rank_range_respects_multinode_dp_attention(): server_args = SimpleNamespace( dp_size=8, diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 1a2a754538bc..99f6071bf71b 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -891,10 +891,8 @@ fn engine_config_normalizes_total_kv_blocks_per_dp_rank() { } #[test] -fn engine_config_handles_inexact_aggregate_kv_capacity_conservatively() { - for (aggregate_blocks, expected_per_rank_blocks) in - [(0, None), (1, Some(1)), (4097, Some(2048))] - { +fn engine_config_handles_zero_and_inexact_aggregate_kv_capacity() { + for (aggregate_blocks, expected_per_rank_blocks) in [(0, None), (4097, Some(2048))] { let mut server = server_info(); server.total_kv_blocks = aggregate_blocks; From 8243d28bd969975964bdb59fe381f46adbdebee8 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 3 Sep 2026 09:50:39 -0700 Subject: [PATCH 3/7] fix(discovery): reject zero DP rank counts Signed-off-by: PeaBrane --- lib/llm/src/discovery/runtime_configs.rs | 27 +++++++++++++++++++++++ lib/llm/src/local_model/runtime_config.rs | 18 +++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/lib/llm/src/discovery/runtime_configs.rs b/lib/llm/src/discovery/runtime_configs.rs index 631a61972503..1bc885571f38 100644 --- a/lib/llm/src/discovery/runtime_configs.rs +++ b/lib/llm/src/discovery/runtime_configs.rs @@ -62,6 +62,13 @@ fn base_runtime_config_watch( if id.model_suffix.is_some() || card.lora.is_some() { continue; } + if card.runtime_config.data_parallel_size == 0 { + tracing::warn!( + instance_id = id.instance_id, + "Ignoring base model runtime config with zero data_parallel_size" + ); + continue; + } configs.insert(id.instance_id, card.runtime_config); } Ok(DiscoveryEvent::ModelTaintsUpdated(update)) => { @@ -269,6 +276,26 @@ mod tests { assert!(configs.borrow().is_empty()); } + #[tokio::test] + async fn zero_data_parallel_size_is_ignored_before_runtime_config_watch() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let stream: DiscoveryStream = + Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(rx)); + let mut configs = base_runtime_config_watch(stream, CancellationToken::new()); + let mut invalid = ModelDeploymentCard::default(); + invalid.runtime_config.data_parallel_size = 0; + let valid = ModelDeploymentCard::default(); + + tx.send(Ok(DiscoveryEvent::Added(model_instance(7, None, &invalid)))) + .unwrap(); + tx.send(Ok(DiscoveryEvent::Added(model_instance(8, None, &valid)))) + .unwrap(); + + configs.changed().await.unwrap(); + assert!(!configs.borrow().contains_key(&7)); + assert_eq!(configs.borrow().get(&8).unwrap().data_parallel_size, 1); + } + #[tokio::test] async fn scoped_taint_updates_replace_only_known_base_worker_taints() { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 6fcf6319b030..72afe02887db 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -555,6 +555,13 @@ fn validate_kv_transfer_domain(domain: &str) -> Result<(), ValidationError> { } fn validate_model_runtime_config(config: &ModelRuntimeConfig) -> Result<(), ValidationError> { + if config.data_parallel_size == 0 { + return Err(validation_error( + "invalid_data_parallel_size", + "data_parallel_size must be at least 1", + )); + } + if let Some(parser) = config .tool_call_parser .as_deref() @@ -1215,6 +1222,17 @@ mod tests { } } + #[test] + fn test_validate_config_rejects_zero_data_parallel_size() { + let config = ModelRuntimeConfig { + data_parallel_size: 0, + ..Default::default() + }; + + let error = config.validate_config().unwrap_err(); + assert!(error.contains("data_parallel_size must be at least 1")); + } + #[test] fn test_validate_config_rejects_invalid_topology_components() { for config in [ From 9004098cbb4452fd41303f6a958bb4b45d73e13d Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 3 Sep 2026 11:24:54 -0700 Subject: [PATCH 4/7] fix(router): address DP capacity review feedback Signed-off-by: PeaBrane --- components/src/dynamo/vllm/capacity.py | 4 ++++ lib/sidecar/sglang/src/engine.rs | 16 +++++++------- lib/sidecar/vllm/src/model.rs | 2 ++ lib/sidecar/vllm/src/tests.rs | 30 ++++++++++---------------- 4 files changed, 25 insertions(+), 27 deletions(-) diff --git a/components/src/dynamo/vllm/capacity.py b/components/src/dynamo/vllm/capacity.py index 92081e10f05b..93105cf66164 100644 --- a/components/src/dynamo/vllm/capacity.py +++ b/components/src/dynamo/vllm/capacity.py @@ -29,6 +29,10 @@ 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. + """ if total_kv_blocks is None: return None diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index dfe86fe4e116..d0ad6a0bd816 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -806,13 +806,10 @@ fn build_engine_config( let max_num_batched_tokens = client::json_u64(&discovery.server_info, "max_prefill_tokens").or(max_total_tokens); - let (data_parallel_start_rank, data_parallel_size) = if enable_dp_attention && dp_size > 1 { - // Native gRPC is exposed by the rank-zero frontend for the complete - // multi-node SGLang endpoint, so one sidecar registers every DP rank. - (Some(0), Some(dp_size)) - } else { - (Some(0), Some(1)) - }; + // Native gRPC is exposed by the rank-zero frontend for the complete + // SGLang endpoint, so one sidecar registers every pure- or attention-DP rank. + let data_parallel_start_rank = Some(0); + let data_parallel_size = Some(dp_size); if mode.is_prefill() && (bootstrap_host.is_none() || bootstrap_port.is_none()) { return Err(client::protocol_error( @@ -966,7 +963,10 @@ mod tests { ) .unwrap(); - assert_eq!(config.llm.unwrap().max_num_seqs, Some(256)); + let registration = config.llm.unwrap(); + assert_eq!(registration.max_num_seqs, Some(256)); + assert_eq!(registration.data_parallel_start_rank, Some(0)); + assert_eq!(registration.data_parallel_size, Some(4)); } #[test] diff --git a/lib/sidecar/vllm/src/model.rs b/lib/sidecar/vllm/src/model.rs index e8c6383693a1..7200732c2339 100644 --- a/lib/sidecar/vllm/src/model.rs +++ b/lib/sidecar/vllm/src/model.rs @@ -166,6 +166,8 @@ impl DiscoveredModel { fn total_kv_blocks_per_rank(&self) -> Option { let total_kv_blocks = nonzero(self.server.total_kv_blocks)?; let data_parallel_size = u64::from(self.data_parallel_size()); + // Control exposes only the aggregate across DP engines. This arithmetic-mean + // estimate assumes homogeneous ranks; exact division does not prove they are equal. let per_rank = total_kv_blocks / data_parallel_size; if per_rank == 0 { diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 99f6071bf71b..da912159ea49 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -867,27 +867,19 @@ fn discovery_rejects_incompatible_model_metadata() { #[test] fn engine_config_normalizes_total_kv_blocks_per_dp_rank() { - for (data_parallel_size, aggregate_blocks, expected_per_rank_blocks) in - [(1, 2048, 2048), (2, 4096, 2048), (4, 8192, 2048)] - { - let mut server = server_info(); - server - .parallelism - .as_mut() - .expect("parallelism metadata") - .data_parallel_size = data_parallel_size; - server.total_kv_blocks = aggregate_blocks; + let mut server = server_info(); + server + .parallelism + .as_mut() + .expect("parallelism metadata") + .data_parallel_size = 2; + server.total_kv_blocks = 4096; - let model = - DiscoveredModel::from_proto(model_info(), server).expect("valid discovery metadata"); - let registration = model.engine_config().llm.expect("LLM registration"); + let model = + DiscoveredModel::from_proto(model_info(), server).expect("valid discovery metadata"); + let registration = model.engine_config().llm.expect("LLM registration"); - assert_eq!( - registration.total_kv_blocks, - Some(expected_per_rank_blocks), - "DP size {data_parallel_size}" - ); - } + assert_eq!(registration.total_kv_blocks, Some(2048)); } #[test] From 349e827b57a939a22261ad426074b25c75d00139 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 3 Sep 2026 12:06:36 -0700 Subject: [PATCH 5/7] docs(router): record rank capacity follow-ups Signed-off-by: PeaBrane --- components/src/dynamo/planner/connectors/mdc.py | 3 +++ components/src/dynamo/sglang/register.py | 3 +++ components/src/dynamo/thunderagent_router/capacity.py | 3 +++ components/src/dynamo/vllm/capacity.py | 2 ++ lib/bindings/python/rust/backend.rs | 3 +++ lib/kv-router/src/scheduling/queue.rs | 3 +++ lib/kv-router/src/services/selection/types.rs | 3 +++ lib/llm/src/discovery/worker_monitor.rs | 4 ++++ lib/llm/src/kv_dc_relay/load.rs | 3 +++ lib/llm/src/local_model/runtime_config.rs | 9 +++++++-- lib/sidecar/vllm/src/model.rs | 2 ++ 11 files changed, 36 insertions(+), 2 deletions(-) diff --git a/components/src/dynamo/planner/connectors/mdc.py b/components/src/dynamo/planner/connectors/mdc.py index ea776c727266..b28b3892da11 100644 --- a/components/src/dynamo/planner/connectors/mdc.py +++ b/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, diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index 4c5893bc6b82..42e6c020ce55 100644 --- a/components/src/dynamo/sglang/register.py +++ b/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") diff --git a/components/src/dynamo/thunderagent_router/capacity.py b/components/src/dynamo/thunderagent_router/capacity.py index dc6f2682b38e..6fa8ddf5a67b 100644 --- a/components/src/dynamo/thunderagent_router/capacity.py +++ b/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", {}) ) diff --git a/components/src/dynamo/vllm/capacity.py b/components/src/dynamo/vllm/capacity.py index 93105cf66164..f0afd4657c84 100644 --- a/components/src/dynamo/vllm/capacity.py +++ b/components/src/dynamo/vllm/capacity.py @@ -32,6 +32,8 @@ def per_rank_kv_blocks( """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 diff --git a/lib/bindings/python/rust/backend.rs b/lib/bindings/python/rust/backend.rs index 8ba566c20b74..c66b07939d90 100644 --- a/lib/bindings/python/rust/backend.rs +++ b/lib/bindings/python/rust/backend.rs @@ -215,6 +215,9 @@ pub struct LlmRegistration { #[pymethods] impl LlmRegistration { + // TODO(rank-aware-kv-capacity): append any rank-capacity arguments so existing positional + // callers do not shift, and update the Python dataclass, duck-typed extraction, stub, and + // Rust-to-MDC copy as one compatibility boundary. #[new] #[pyo3(signature = ( context_length = None, diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index ccc48bce35ef..4f9cac14dbf0 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -1439,6 +1439,9 @@ impl< .max_num_batched_tokens() .unwrap_or(DEFAULT_MAX_BATCHED_TOKENS) as usize, + // TODO(rank-aware-kv-capacity): resolve the selected DP rank and preserve + // capacity quality. Estimated fallbacks must not authorize load-based + // admission/bypass decisions that require an exact or conservative bound. total_kv_blocks: config.total_kv_blocks().map(|blocks| blocks as usize), }; SelectedWorkerForRequest { diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index b9117e94d1b9..b0ceea00f0fd 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -338,6 +338,9 @@ pub struct WorkerPatchRequest { impl WorkerCatalogRecord { pub(super) fn apply_patch(&mut self, patch: WorkerPatchRequest) { + // TODO(rank-aware-kv-capacity): when the rank map is added, treat rank range, map, + // scalar fallback, and provenance as one replace-only snapshot. A legacy scalar/range + // patch must clear stale exact data rather than leave it winning lookup precedence. if patch.endpoint.is_some() { self.endpoint = patch.endpoint; } diff --git a/lib/llm/src/discovery/worker_monitor.rs b/lib/llm/src/discovery/worker_monitor.rs index 2ce354dd2442..3da7cc980184 100644 --- a/lib/llm/src/discovery/worker_monitor.rs +++ b/lib/llm/src/discovery/worker_monitor.rs @@ -281,6 +281,10 @@ impl WorkerLoadState { self.kv_total_blocks.clear(); if let Some(total_blocks) = total_kv_blocks { + // TODO(rank-aware-kv-capacity): resolve each rank from a validated advertisement and + // retain its provenance. Aggregate/representative estimates may support approximate + // routing, but must not trip this hard overload threshold. Exclusion remains + // worker-granular until the overloaded-worker contract itself becomes rank-aware. self.kv_total_blocks.extend( declared_dp_ranks .iter() diff --git a/lib/llm/src/kv_dc_relay/load.rs b/lib/llm/src/kv_dc_relay/load.rs index d71c2c837757..356055433374 100644 --- a/lib/llm/src/kv_dc_relay/load.rs +++ b/lib/llm/src/kv_dc_relay/load.rs @@ -180,6 +180,9 @@ fn load_ranks_from_configs( // registering process and uses zero as an unknown-capacity sentinel. Runtime // config does not carry backend identity, and zero is never a usable pressure // denominator for any backend, so normalize it fail-closed for every engine. + // TODO(rank-aware-kv-capacity): resolve exact/conservative capacity per rank and carry + // quality into the pool snapshot. Estimated fallbacks must not count as authoritative + // rank coverage merely because every rank received a scalar. let total_kv_blocks = config.total_kv_blocks.filter(|&total| total != 0); for dp_rank in config.data_parallel_start_rank..end { ranks.insert( diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 72afe02887db..97f525d116f3 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -199,8 +199,13 @@ pub struct ModelRuntimeConfig { #[serde(default, skip_serializing_if = "Option::is_none")] pub context_length: Option, - /// Physical KV-cache capacity for each router-visible data-parallel rank. - /// This is per rank, never the aggregate capacity of the worker process. + /// Compatibility KV-cache capacity applied to each router-visible data-parallel rank. + /// Some adapters derive this scalar from aggregate or representative-rank data. + /// + /// TODO(rank-aware-kv-capacity): Add an additive per-rank advertisement whose resolver + /// preserves exact/conservative/estimated provenance. Exact heterogeneous producers must + /// dual-write their minimum here for old readers; aggregate division stays an adapter-only + /// estimate and must not silently become a hard-admission denominator. pub total_kv_blocks: Option, pub max_num_seqs: Option, diff --git a/lib/sidecar/vllm/src/model.rs b/lib/sidecar/vllm/src/model.rs index 7200732c2339..85522a7f2b81 100644 --- a/lib/sidecar/vllm/src/model.rs +++ b/lib/sidecar/vllm/src/model.rs @@ -168,6 +168,8 @@ impl DiscoveredModel { let data_parallel_size = u64::from(self.data_parallel_size()); // Control exposes only the aggregate across DP engines. This arithmetic-mean // estimate assumes homogeneous ranks; exact division does not prove they are equal. + // TODO(rank-aware-kv-capacity): consume a per-rank Control response when available and + // publish it atomically; never relabel this quotient as exact for hard admission. let per_rank = total_kv_blocks / data_parallel_size; if per_rank == 0 { From 03d54f9c326af9ed7cf54e4956dae742d7d5877a Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 3 Sep 2026 15:54:04 -0700 Subject: [PATCH 6/7] refactor(router): remove redundant DP state Signed-off-by: PeaBrane --- components/src/dynamo/vllm/publisher.py | 2 -- .../dynamo/vllm/tests/test_vllm_publisher.py | 6 ----- lib/llm/src/discovery/worker_monitor.rs | 24 ++++++++++++------- 3 files changed, 15 insertions(+), 17 deletions(-) diff --git a/components/src/dynamo/vllm/publisher.py b/components/src/dynamo/vllm/publisher.py index 8f697d15b67d..85d2cd016de4 100644 --- a/components/src/dynamo/vllm/publisher.py +++ b/components/src/dynamo/vllm/publisher.py @@ -152,7 +152,6 @@ def __init__( 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: @@ -171,7 +170,6 @@ 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 diff --git a/components/src/dynamo/vllm/tests/test_vllm_publisher.py b/components/src/dynamo/vllm/tests/test_vllm_publisher.py index 29354ed80f10..7cbc105baecf 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_publisher.py +++ b/components/src/dynamo/vllm/tests/test_vllm_publisher.py @@ -59,11 +59,6 @@ def _explode(*_a, **_kw): logger = factory.create_stat_logger(dp_rank=0) assert isinstance(logger, NoopStatLogger) - # Embedding factory never tracks a created chat logger, so the - # downstream ``init_publish`` / ``set_num_gpu_blocks_all`` calls in - # the chat path are safe no-ops if anyone ever wires them on the - # embedding branch by mistake. - assert factory.created_logger is None def test_noop_stat_logger_record_is_safe_with_none_stats(): @@ -146,7 +141,6 @@ def _fake_publisher(*args, **kwargs): factory.set_num_gpu_blocks_all(4096) factory.init_publish() - assert factory.created_logger is loggers[-1] assert factory.created_loggers == dict(enumerate(loggers)) for logger in loggers: logger.set_num_gpu_block.assert_called_once_with(4096) diff --git a/lib/llm/src/discovery/worker_monitor.rs b/lib/llm/src/discovery/worker_monitor.rs index 3da7cc980184..6856256dec7e 100644 --- a/lib/llm/src/discovery/worker_monitor.rs +++ b/lib/llm/src/discovery/worker_monitor.rs @@ -478,15 +478,21 @@ impl WorkerLoadState { // Once discovery has supplied the runtime config, its rank set is // authoritative. An expected rank without a load observation is free, // so one noisy rank cannot exclude the whole worker during startup. - let all_dp_ranks = self.declared_dp_ranks.clone().unwrap_or_else(|| { - self.active_decode_blocks - .keys() - .chain(self.kv_used_blocks.keys()) - .chain(self.decode_overload_latches.keys()) - .chain(self.active_prefill_tokens.keys()) - .copied() - .collect() - }); + let fallback_dp_ranks; + let all_dp_ranks = match &self.declared_dp_ranks { + Some(declared_dp_ranks) => declared_dp_ranks, + None => { + fallback_dp_ranks = self + .active_decode_blocks + .keys() + .chain(self.kv_used_blocks.keys()) + .chain(self.decode_overload_latches.keys()) + .chain(self.active_prefill_tokens.keys()) + .copied() + .collect(); + &fallback_dp_ranks + } + }; // If no dp_ranks known, not overloaded if all_dp_ranks.is_empty() { From 6fecc488894373e847104e0b7e5579ce26d90c89 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 3 Sep 2026 19:08:24 -0700 Subject: [PATCH 7/7] fix(router): bound advertised DP rank ranges Signed-off-by: PeaBrane --- lib/llm/src/discovery/runtime_configs.rs | 34 +++++++++---- lib/llm/src/discovery/worker_monitor.rs | 29 ++++++----- lib/llm/src/local_model/runtime_config.rs | 59 ++++++++++++++++++++--- 3 files changed, 93 insertions(+), 29 deletions(-) diff --git a/lib/llm/src/discovery/runtime_configs.rs b/lib/llm/src/discovery/runtime_configs.rs index 1bc885571f38..1ae77bfc9e25 100644 --- a/lib/llm/src/discovery/runtime_configs.rs +++ b/lib/llm/src/discovery/runtime_configs.rs @@ -62,14 +62,16 @@ fn base_runtime_config_watch( if id.model_suffix.is_some() || card.lora.is_some() { continue; } - if card.runtime_config.data_parallel_size == 0 { + if let Err(error) = card.runtime_config.data_parallel_rank_range() { tracing::warn!( instance_id = id.instance_id, - "Ignoring base model runtime config with zero data_parallel_size" + %error, + "Ignoring base model runtime config with invalid data-parallel rank range" ); - continue; + configs.remove(&id.instance_id); + } else { + configs.insert(id.instance_id, card.runtime_config); } - configs.insert(id.instance_id, card.runtime_config); } Ok(DiscoveryEvent::ModelTaintsUpdated(update)) => { if update.id.model_suffix.is_some() { @@ -277,23 +279,35 @@ mod tests { } #[tokio::test] - async fn zero_data_parallel_size_is_ignored_before_runtime_config_watch() { + async fn invalid_data_parallel_ranges_are_ignored_before_runtime_config_watch() { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); let stream: DiscoveryStream = Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(rx)); let mut configs = base_runtime_config_watch(stream, CancellationToken::new()); - let mut invalid = ModelDeploymentCard::default(); - invalid.runtime_config.data_parallel_size = 0; + let mut zero_size = ModelDeploymentCard::default(); + zero_size.runtime_config.data_parallel_size = 0; + let mut oversized = ModelDeploymentCard::default(); + oversized.runtime_config.data_parallel_size = 4097; + let mut overflowing = ModelDeploymentCard::default(); + overflowing.runtime_config.data_parallel_start_rank = u32::MAX; let valid = ModelDeploymentCard::default(); - tx.send(Ok(DiscoveryEvent::Added(model_instance(7, None, &invalid)))) + for (instance_id, card) in [(6, &zero_size), (7, &oversized), (8, &overflowing)] { + tx.send(Ok(DiscoveryEvent::Added(model_instance( + instance_id, + None, + card, + )))) .unwrap(); - tx.send(Ok(DiscoveryEvent::Added(model_instance(8, None, &valid)))) + } + tx.send(Ok(DiscoveryEvent::Added(model_instance(9, None, &valid)))) .unwrap(); configs.changed().await.unwrap(); + assert!(!configs.borrow().contains_key(&6)); assert!(!configs.borrow().contains_key(&7)); - assert_eq!(configs.borrow().get(&8).unwrap().data_parallel_size, 1); + assert!(!configs.borrow().contains_key(&8)); + assert_eq!(configs.borrow().get(&9).unwrap().data_parallel_size, 1); } #[tokio::test] diff --git a/lib/llm/src/discovery/worker_monitor.rs b/lib/llm/src/discovery/worker_monitor.rs index 6856256dec7e..12f51570fb93 100644 --- a/lib/llm/src/discovery/worker_monitor.rs +++ b/lib/llm/src/discovery/worker_monitor.rs @@ -263,14 +263,12 @@ pub struct WorkerLoadState { impl WorkerLoadState { fn reconcile_runtime_config( &mut self, - dp_start: u32, - dp_size: u32, + dp_ranks: std::ops::Range, total_kv_blocks: Option, max_num_batched_tokens: Option, active_decode_blocks_threshold: Option, ) -> HashSet { - let dp_end = dp_start.saturating_add(dp_size); - let declared_dp_ranks: HashSet<_> = (dp_start..dp_end).collect(); + let declared_dp_ranks: HashSet<_> = dp_ranks.collect(); self.active_decode_blocks .retain(|dp_rank, _| declared_dp_ranks.contains(dp_rank)); @@ -878,10 +876,19 @@ impl WorkerLoadMonitor for KvWorkerMonitor { // worker-level "all ranks overloaded" decision. for (lease_id, runtime_config) in runtime_configs.iter() { let mut state = worker_load_states.entry(*lease_id).or_default(); - let dp_start = runtime_config.data_parallel_start_rank; + let dp_ranks = match runtime_config.data_parallel_rank_range() { + Ok(dp_ranks) => dp_ranks, + Err(error) => { + tracing::warn!( + worker_id = *lease_id, + %error, + "ignoring runtime config with an invalid data-parallel rank range" + ); + continue; + } + }; let declared_dp_ranks = state.reconcile_runtime_config( - dp_start, - runtime_config.data_parallel_size, + dp_ranks, runtime_config.total_kv_blocks, runtime_config.max_num_batched_tokens, cfg.active_decode_blocks_threshold, @@ -1344,7 +1351,7 @@ mod tests { #[test] fn expected_but_unobserved_dp_rank_keeps_worker_available() { let mut state = WorkerLoadState::default(); - state.reconcile_runtime_config(0, 2, Some(100), Some(1_000), Some(0.6)); + state.reconcile_runtime_config(0..2, Some(100), Some(1_000), Some(0.6)); state.update_from_active_load( &ActiveLoad { @@ -1374,7 +1381,7 @@ mod tests { #[test] fn runtime_config_update_reconciles_rank_range_and_optional_capacity() { let mut state = WorkerLoadState::default(); - state.reconcile_runtime_config(2, 2, Some(100), Some(1_000), Some(0.6)); + state.reconcile_runtime_config(2..4, Some(100), Some(1_000), Some(0.6)); for dp_rank in 2..4 { state.update_from_active_load( @@ -1390,7 +1397,7 @@ mod tests { } assert!(state.is_overloaded(Some(0.6), None, Some(0.5))); - let declared = state.reconcile_runtime_config(3, 1, None, None, Some(0.6)); + let declared = state.reconcile_runtime_config(3..4, None, None, Some(0.6)); assert_eq!(declared, HashSet::from([3])); assert!(!state.active_decode_blocks.contains_key(&2)); assert!(!state.kv_used_blocks.contains_key(&2)); @@ -1410,7 +1417,7 @@ mod tests { Some(0.6), )); - let declared = state.reconcile_runtime_config(4, 1, Some(100), Some(1_000), Some(0.6)); + let declared = state.reconcile_runtime_config(4..5, Some(100), Some(1_000), Some(0.6)); assert_eq!(declared, HashSet::from([4])); assert!(state.active_decode_blocks.is_empty()); assert!(state.kv_used_blocks.is_empty()); diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 97f525d116f3..72904a6b1872 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -4,6 +4,7 @@ use std::{ borrow::Cow, collections::{HashMap, HashSet}, + ops::Range, str::FromStr, }; @@ -32,6 +33,9 @@ pub const TOPOLOGY_TAINT_PREFIX: &str = "dynamo.topology/"; /// Runtime-data key for an engine-published token-overflow contract. pub const TOKEN_BUDGET_RUNTIME_KEY: &str = "token_budget"; +/// Resource-safety bound for rank ranges advertised by one worker. +pub(crate) const MAX_DATA_PARALLEL_RANKS_PER_WORKER: u32 = 4096; + /// Runtime-data key indicating that a backend expects tool structural tags to /// exclude reasoning and manages grammar activation around reasoning itself. /// @@ -566,6 +570,9 @@ fn validate_model_runtime_config(config: &ModelRuntimeConfig) -> Result<(), Vali "data_parallel_size must be at least 1", )); } + config + .data_parallel_rank_range() + .map_err(|error| validation_error("invalid_data_parallel_rank_range", error))?; if let Some(parser) = config .tool_call_parser @@ -640,6 +647,23 @@ impl ModelRuntimeConfig { self.validate().map_err(|error| error.to_string()) } + pub(crate) fn data_parallel_rank_range(&self) -> Result, String> { + if self.data_parallel_size == 0 { + return Err("data_parallel_size must be at least 1".to_string()); + } + if self.data_parallel_size > MAX_DATA_PARALLEL_RANKS_PER_WORKER { + return Err(format!( + "data_parallel_size {} exceeds the supported maximum {}", + self.data_parallel_size, MAX_DATA_PARALLEL_RANKS_PER_WORKER + )); + } + let end = self + .data_parallel_start_rank + .checked_add(self.data_parallel_size) + .ok_or_else(|| "data-parallel rank range overflows u32".to_string())?; + Ok(self.data_parallel_start_rank..end) + } + pub fn set_engine_specific(&mut self, key: &str, value: T) -> anyhow::Result<()> { self.runtime_data .insert(key.to_string(), serde_json::to_value(value)?); @@ -1228,14 +1252,33 @@ mod tests { } #[test] - fn test_validate_config_rejects_zero_data_parallel_size() { - let config = ModelRuntimeConfig { - data_parallel_size: 0, - ..Default::default() - }; - - let error = config.validate_config().unwrap_err(); - assert!(error.contains("data_parallel_size must be at least 1")); + fn test_validate_config_rejects_invalid_data_parallel_ranges() { + for (config, expected_error) in [ + ( + ModelRuntimeConfig { + data_parallel_size: 0, + ..Default::default() + }, + "data_parallel_size must be at least 1", + ), + ( + ModelRuntimeConfig { + data_parallel_size: MAX_DATA_PARALLEL_RANKS_PER_WORKER + 1, + ..Default::default() + }, + "exceeds the supported maximum", + ), + ( + ModelRuntimeConfig { + data_parallel_start_rank: u32::MAX, + ..Default::default() + }, + "data-parallel rank range overflows u32", + ), + ] { + let error = config.validate_config().unwrap_err(); + assert!(error.contains(expected_error), "{error}"); + } } #[test]