diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 1b789152e9..dc8d7152b7 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -15,6 +15,7 @@ from vllm.config import ( SpeculativeConfig, VllmConfig, ) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalKwargsItem, @@ -3990,6 +3991,87 @@ def test_delayed_kv_connector_free_keeps_scheduler_active(): assert not scheduler.has_finished_requests() +def test_scheduler_kv_connector_stats(): + """Test worker-side, scheduler-side, and combined KV connector stats.""" + + class GenericKVConnectorStats(KVConnectorStats): + def reset(self): + self.data = {} + + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: + self.data.update(other.data) + return self + + def reduce(self) -> dict[str, int | float]: + return {} + + def is_empty(self) -> bool: + return not self.data + + test_cases = ( + ({"worker": 1}, None, {"worker": 1}), + (None, {"scheduler": 2}, {"scheduler": 2}), + ({"worker": 1}, {"scheduler": 2}, {"worker": 1, "scheduler": 2}), + ) + + for worker_data, scheduler_data, expected_data in test_cases: + scheduler = create_scheduler() + worker_stats = ( + GenericKVConnectorStats(data=worker_data) if worker_data else None + ) + scheduler_stats = ( + GenericKVConnectorStats(data=scheduler_data) if scheduler_data else None + ) + scheduler.connector = Mock() + scheduler.connector.get_kv_connector_stats.return_value = ( + scheduler_stats if worker_stats is None else None + ) + scheduler.connector.take_events.return_value = [] + + def update_connector_output( + kv_connector_output: KVConnectorOutput, + scheduler=scheduler, + scheduler_stats=scheduler_stats, + ): + scheduler.connector.get_kv_connector_stats.return_value = scheduler_stats + + scheduler.connector.update_connector_output.side_effect = ( + update_connector_output + ) + + model_output = ModelRunnerOutput( + req_ids=["req_0"], + req_id_to_index={"req_0": 0}, + sampled_token_ids=[[123]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[None], + kv_connector_output=KVConnectorOutput(kv_connector_stats=worker_stats) + if worker_stats + else None, + ) + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=None, + num_scheduled_tokens={"req_0": 1}, + total_num_scheduled_tokens=1, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + + engine_core_outputs = scheduler.update_from_output( + scheduler_output, model_output + ) + + final_stats = next( + iter(engine_core_outputs.values()) + ).scheduler_stats.kv_connector_stats + assert final_stats == expected_data + + # ============================================================================== # Variable-length encoder cross-attention block allocation tests # ============================================================================== diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 6ac6b4318c..2d6fa834d2 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -58,6 +58,7 @@ class MockConnector(KVConnectorBase_V1): mock = MagicMock(spec_set=KVConnectorBase_V1) # Override just build_kv_connector_stats mock.build_kv_connector_stats = cls.build_kv_connector_stats + mock.get_kv_connector_stats.return_value = None return mock @classmethod @@ -93,6 +94,7 @@ class MockHMAConnector(KVConnectorBase_V1, SupportsHMA): def __new__(cls, *args, **kwargs): mock = MagicMock(spec_set=cls) + mock.get_kv_connector_stats.return_value = None return mock def start_load_kv(self, forward_context, **kwargs): @@ -368,7 +370,7 @@ def test_multi_example_connector_consistency(): def _ignore_event_collection(events: list[str]) -> list[str]: # Filter out per-step polling hooks that the scheduler calls repeatedly # and which are not meaningful state transitions for these assertions. - ignored = {"take_events", "has_pending_push_work"} + ignored = {"get_kv_connector_stats", "has_pending_push_work", "take_events"} return [event for event in events if event not in ignored] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 5ca90fd129..e215c698c4 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1412,13 +1412,6 @@ class Scheduler(SchedulerInterface): outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats: KVConnectorStats | None = ( - kv_connector_output.kv_connector_stats if kv_connector_output else None - ) - if kv_connector_stats and self.connector: - kv_stats = self.connector.get_kv_connector_stats() - if kv_stats: - kv_connector_stats = kv_connector_stats.aggregate(kv_stats) failed_kv_load_req_ids = None if kv_connector_output and kv_connector_output.invalid_block_ids: @@ -1665,6 +1658,23 @@ class Scheduler(SchedulerInterface): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if ( + scheduler_kv_connector_stats is not None + and not scheduler_kv_connector_stats.is_empty() + ): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + # collect KV cache events from KV cache manager events = self.kv_cache_manager.take_events()