From 7331f9d9913841887ca796ae437e73cda2aa53e0 Mon Sep 17 00:00:00 2001 From: gitover22 Date: Thu, 10 Sep 2026 18:13:29 +0800 Subject: [PATCH] fix(vllm): publish KV gauges from prefill workers Signed-off-by: gitover22 --- components/src/dynamo/vllm/publisher.py | 2 +- .../dynamo/vllm/tests/test_vllm_publisher.py | 66 +++++++++++++++++-- .../vllm/tests/test_vllm_worker_factory.py | 44 +++++++++++-- components/src/dynamo/vllm/worker_factory.py | 23 ++++--- tests/serve/test_vllm.py | 8 ++- 5 files changed, 119 insertions(+), 24 deletions(-) diff --git a/components/src/dynamo/vllm/publisher.py b/components/src/dynamo/vllm/publisher.py index 85d2cd016de4..da7317e8653a 100644 --- a/components/src/dynamo/vllm/publisher.py +++ b/components/src/dynamo/vllm/publisher.py @@ -89,7 +89,7 @@ def record( def init_publish(self) -> None: self.inner.publish(self.dp_rank, kv_used_blocks=0) dp_rank_str = str(self.dp_rank) - self.component_gauges.set_total_blocks(dp_rank_str, 0) + self.component_gauges.set_total_blocks(dp_rank_str, self.num_gpu_block) self.component_gauges.set_gpu_cache_usage(dp_rank_str, 0.0) def log_engine_initialized(self) -> None: diff --git a/components/src/dynamo/vllm/tests/test_vllm_publisher.py b/components/src/dynamo/vllm/tests/test_vllm_publisher.py index 7cbc105baecf..138a11fc68ec 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_publisher.py +++ b/components/src/dynamo/vllm/tests/test_vllm_publisher.py @@ -12,11 +12,14 @@ """ from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, Mock, call import pytest +from prometheus_client import CollectorRegistry import dynamo.vllm.publisher as publisher_mod +from dynamo import prometheus_names +from dynamo.common.utils.prometheus import LLMBackendMetrics from dynamo.vllm.publisher import ( DynamoStatLoggerPublisher, NoopStatLogger, @@ -135,18 +138,29 @@ def _fake_publisher(*args, **kwargs): factory = StatLoggerFactory( endpoint=SimpleNamespace(), component_gauges=SimpleNamespace() ) - for dp_rank in range(3): + dp_ranks = (2, 4, 7) + for dp_rank in dp_ranks: factory.create_stat_logger(dp_rank=dp_rank) factory.set_num_gpu_blocks_all(4096) factory.init_publish() - assert factory.created_loggers == dict(enumerate(loggers)) + assert factory.created_loggers == dict(zip(dp_ranks, loggers, strict=True)) for logger in loggers: logger.set_num_gpu_block.assert_called_once_with(4096) logger.init_publish.assert_called_once_with() +def test_factory_initialization_without_loggers_is_a_noop(): + """Disabled vLLM stat logging leaves the factory without loggers.""" + factory = StatLoggerFactory(endpoint=SimpleNamespace()) + + factory.set_num_gpu_blocks_all(4096) + factory.init_publish() + + assert factory.created_loggers == {} + + def test_factory_binds_deferred_endpoint_to_every_dp_rank_logger(monkeypatch): loggers = [] @@ -177,11 +191,19 @@ async def test_deferred_logger_starts_with_fresh_metrics_state(monkeypatch): publisher_mod, "WorkerMetricsPublisher", Mock(side_effect=publishers) ) + registry = CollectorRegistry() + component_gauges = LLMBackendMetrics( + registry=registry, + model_name="test-model", + component_name="prefill", + ) + logger = DynamoStatLoggerPublisher( endpoint=None, - component_gauges=SimpleNamespace(), + dp_rank=5, + component_gauges=component_gauges, ) - logger.inner.publish(dp_rank=0, kv_used_blocks=7) + logger.inner.publish(dp_rank=5, kv_used_blocks=7) endpoint = SimpleNamespace() logger.bind_endpoint(endpoint) @@ -189,6 +211,38 @@ async def test_deferred_logger_starts_with_fresh_metrics_state(monkeypatch): 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].publish.assert_called_once_with(dp_rank=5, kv_used_blocks=7) publishers[0].create_endpoint.assert_not_called() publishers[1].create_endpoint.assert_awaited_once_with(endpoint) + + metric_labels = { + prometheus_names.labels.MODEL: "test-model", + prometheus_names.labels.COMPONENT: "prefill", + prometheus_names.labels.DP_RANK: "5", + } + total_blocks_name = ( + f"{prometheus_names.name_prefix.COMPONENT}_" + f"{prometheus_names.kvstats.TOTAL_BLOCKS}" + ) + cache_usage_name = ( + f"{prometheus_names.name_prefix.COMPONENT}_" + f"{prometheus_names.kvstats.GPU_CACHE_USAGE_PERCENT}" + ) + + logger.set_num_gpu_block(400) + logger.init_publish() + assert registry.get_sample_value(total_blocks_name, metric_labels) == 400 + assert registry.get_sample_value(cache_usage_name, metric_labels) == 0.0 + + logger.record(SimpleNamespace(kv_cache_usage=0.25), None) + assert registry.get_sample_value(total_blocks_name, metric_labels) == 400 + assert registry.get_sample_value(cache_usage_name, metric_labels) == 0.25 + + logger.record(SimpleNamespace(kv_cache_usage=0.0), None) + assert registry.get_sample_value(total_blocks_name, metric_labels) == 400 + assert registry.get_sample_value(cache_usage_name, metric_labels) == 0.0 + assert publishers[1].publish.call_args_list == [ + call(5, kv_used_blocks=0), + call(5, kv_used_blocks=100), + call(5, kv_used_blocks=0), + ] 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 a9e9019b7412..2bfeab3044d4 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py @@ -1114,21 +1114,32 @@ async def _noop(*_args, **_kwargs) -> None: @pytest.mark.asyncio @pytest.mark.parametrize("lora_enabled", [True, False]) -async def test_prefill_serves_lora_lifecycle_endpoints_when_enabled( +@pytest.mark.parametrize("snapshot_mode", [True, False]) +async def test_prefill_initializes_metrics_and_serves_lora_lifecycle( monkeypatch: pytest.MonkeyPatch, + tmp_path, lora_enabled: bool, + snapshot_mode: bool, ) -> None: engine_client = Mock() - vllm_config = Mock(additional_config={}) + vllm_config = Mock( + additional_config={}, + cache_config=SimpleNamespace(num_gpu_blocks=96), + ) engine_tuple: EngineSetupResult = ( engine_client, vllm_config, Mock(), - "/tmp/prom", + str(tmp_path / "prometheus"), Mock(), ) + stat_logger = Mock() + snapshot_engine: SnapshotEngineSetupResult | None = ( + (engine_tuple, stat_logger) if snapshot_mode else None + ) + setup_vllm_engine = Mock(return_value=engine_tuple) factory = WorkerFactory( - setup_vllm_engine_fn=Mock(return_value=engine_tuple), + setup_vllm_engine_fn=setup_vllm_engine, setup_kv_event_publisher_fn=Mock(return_value=None), register_vllm_model_fn=AsyncMock(), setup_fpm_relay_fn=Mock(return_value=None), @@ -1155,6 +1166,13 @@ async def _noop(*_args, **_kwargs) -> None: monkeypatch.setattr( "dynamo.vllm.worker_factory.configure_kv_event_block_size", _noop ) + monkeypatch.setattr( + "dynamo.vllm.worker_factory.get_dp_range_for_worker", lambda _config: (3, 2) + ) + stat_logger_factory = Mock(return_value=stat_logger) + monkeypatch.setattr( + "dynamo.vllm.worker_factory.StatLoggerFactory", stat_logger_factory + ) endpoints: dict[str, Mock] = {} @@ -1187,8 +1205,26 @@ def _endpoint(name: str) -> Mock: config, asyncio.Event(), shutdown_endpoints, + snapshot_engine=snapshot_engine, ) + if snapshot_mode: + stat_logger_factory.assert_not_called() + setup_vllm_engine.assert_not_called() + stat_logger.bind_endpoint.assert_called_once_with( + endpoints["dyn.prefill.generate"] + ) + else: + stat_logger_factory.assert_called_once_with( + endpoint=endpoints["dyn.prefill.generate"] + ) + setup_vllm_engine.assert_called_once_with( + config, stat_logger, fpm_worker_id="cid" + ) + stat_logger.bind_endpoint.assert_not_called() + stat_logger.set_num_gpu_blocks_all.assert_called_once_with(48) + stat_logger.init_publish.assert_called_once_with() + lifecycle_names = { "dyn.prefill.load_lora", "dyn.prefill.unload_lora", diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 0324f49b3db2..e5c47f4b98d7 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -1557,9 +1557,8 @@ 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_setup, factory = snapshot_engine ( engine_client, vllm_config, @@ -1567,30 +1566,30 @@ async def _run_prefill_worker( prometheus_temp_dir, _component_gauges, ) = engine_setup - snapshot_factory.bind_endpoint(generate_endpoint) + 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 # a restart of the EngineCore process. os.environ[ENV_FPM_WORKER_ID] = fpm_worker_id else: + factory = StatLoggerFactory(endpoint=generate_endpoint) ( engine_client, vllm_config, default_sampling_params, prometheus_temp_dir, _component_gauges, - ) = self.setup_vllm_engine(config, fpm_worker_id=fpm_worker_id) + ) = self.setup_vllm_engine(config, factory, 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() + _, 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, + ) + factory.set_num_gpu_blocks_all(per_rank_num_gpu_blocks or 0) + factory.init_publish() encode_worker_client = await self._maybe_get_encode_worker_client( runtime, config diff --git a/tests/serve/test_vllm.py b/tests/serve/test_vllm.py index b9cc9c3a2b9c..e173a2302a09 100644 --- a/tests/serve/test_vllm.py +++ b/tests/serve/test_vllm.py @@ -365,11 +365,17 @@ class VLLMConfig(EngineConfig): pytest.mark.core, pytest.mark.gpu_2, pytest.mark.pre_merge, - ], # TODO: profile to get max_vram and timeout + pytest.mark.timeout(720), # 3x ~240s local runtime + ], # TODO: profile to get max_vram model="Qwen/Qwen3-0.6B", request_payloads=[ chat_payload_default(), completion_payload_default(), + metric_payload_default( + min_num_requests=1, + backend="vllm", + port=DefaultPort.SYSTEM2.value, + ), ], ), "disaggregated_same_gpu": VLLMConfig(