| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import time |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | from vllm.v1.engine import EngineCoreOutputs |
| 8 | from vllm.v1.engine.core import EngineCore |
| 9 | from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats |
| 10 | |
| 11 | |
| 12 | class FakeEngineCore: |
| 13 | def _make_iteration_details_stats( |
| 14 | self, iteration_details: SchedulerIterationDetails |
| 15 | ) -> SchedulerStats: |
| 16 | return SchedulerStats(iteration_details=iteration_details) |
| 17 | |
| 18 | |
| 19 | def make_iteration_details() -> SchedulerIterationDetails: |
| 20 | return SchedulerIterationDetails( |
| 21 | iteration_index=1, |
| 22 | num_ctx_requests=2, |
| 23 | num_ctx_tokens=3, |
| 24 | num_generation_requests=4, |
| 25 | num_generation_tokens=5, |
| 26 | elapsed_ms=6.7, |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | def make_fake_engine(log_stats: bool = True) -> SimpleNamespace: |
| 31 | return SimpleNamespace( |
| 32 | log_stats=log_stats, |
| 33 | vllm_config=SimpleNamespace( |
| 34 | observability_config=SimpleNamespace( |
| 35 | enable_logging_iteration_details=True, |
| 36 | ) |
| 37 | ), |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | def test_capture_iteration_details_disabled_without_log_stats(): |
| 42 | engine = make_fake_engine(log_stats=False) |
| 43 | |
| 44 | with EngineCore.capture_iteration_details(engine, None) as iteration_details: |
| 45 | assert iteration_details is None |
| 46 | |
| 47 | assert not hasattr(engine, "_iteration_index") |
| 48 | |
| 49 | |
| 50 | def test_capture_iteration_details_fills_elapsed_time(): |
| 51 | engine = make_fake_engine() |
| 52 | |
| 53 | with EngineCore.capture_iteration_details(engine, None) as iteration_details: |
| 54 | assert iteration_details is not None |
| 55 | assert iteration_details.elapsed_ms == 0.0 |
| 56 | assert iteration_details.is_dummy |
| 57 | time.sleep(0.001) |
| 58 | |
| 59 | assert iteration_details is not None |
| 60 | assert iteration_details.elapsed_ms > 0.0 |
| 61 | assert engine._iteration_index == 1 |
| 62 | |
| 63 | |
| 64 | def test_attach_iteration_details_uses_existing_output(): |
| 65 | iteration_details = make_iteration_details() |
| 66 | outputs = { |
| 67 | 2: EngineCoreOutputs(scheduler_stats=SchedulerStats()), |
| 68 | 1: EngineCoreOutputs(scheduler_stats=SchedulerStats()), |
| 69 | } |
| 70 | |
| 71 | EngineCore._attach_iteration_details(FakeEngineCore(), outputs, iteration_details) |
| 72 | |
| 73 | assert 0 not in outputs |
| 74 | assert outputs[2].scheduler_stats is not None |
| 75 | assert outputs[2].scheduler_stats.iteration_details == iteration_details |
| 76 | assert outputs[1].scheduler_stats is not None |
| 77 | assert outputs[1].scheduler_stats.iteration_details is None |
| 78 | |
| 79 | |
| 80 | def test_attach_iteration_details_falls_back_to_client_zero_without_outputs(): |
| 81 | iteration_details = make_iteration_details() |
| 82 | outputs: dict[int, EngineCoreOutputs] = {} |
| 83 | |
| 84 | EngineCore._attach_iteration_details(FakeEngineCore(), outputs, iteration_details) |
| 85 | |
| 86 | assert set(outputs) == {0} |
| 87 | assert outputs[0].scheduler_stats is not None |
| 88 | assert outputs[0].scheduler_stats.iteration_details == iteration_details |