diff --git a/tests/v1/metrics/test_engine_logger_apis.py b/tests/v1/metrics/test_engine_logger_apis.py index e6a4d0a2a2e8..603790f87db5 100644 --- a/tests/v1/metrics/test_engine_logger_apis.py +++ b/tests/v1/metrics/test_engine_logger_apis.py @@ -5,7 +5,9 @@ import pytest from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM +from vllm.v1.metrics.loggers import LoggingStatLogger from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger +from vllm.v1.metrics.stats import IterationStats class DummyStatLogger: @@ -31,6 +33,16 @@ def log_engine_initialized(self): self.engine_initialized = True +class DummyLoggingStatLogger(LoggingStatLogger): + """ + A dummy logging stat logger for testing purposes. + Implemented the record and log APIs + """ + + def get_num_preempted_reqs(self) -> int: + return self.num_preempted_reqs + + @pytest.fixture def log_stats_enabled_engine_args(): """ @@ -62,7 +74,7 @@ async def test_async_llm_replace_default_loggers( @pytest.mark.asyncio async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args): """ - It's still possible to use custom stat loggers exclusively by passing + It's still possible to use custom stat loggers exclusively by passing disable_log_stats=True in addition to a list of custom stat loggers. """ # Create engine_args with disable_log_stats=True for this test @@ -81,3 +93,41 @@ async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args): assert engine.log_stats engine.shutdown() + + +@pytest.mark.asyncio +async def test_logger_iteration_stats(log_stats_enabled_engine_args): + """ + """ + # Create engine_args with disable_log_stats=True for this test + disabled_log_engine_args = copy.deepcopy(log_stats_enabled_engine_args) + disabled_log_engine_args.disable_log_stats = True + + # Disable default loggers; pass custom stat logger to the constructor + engine = AsyncLLM.from_engine_args(disabled_log_engine_args, + stat_loggers=[DummyLoggingStatLogger]) + + dummy_logger = engine.logger_manager.per_engine_logger_dict[0][0] + + assert len(engine.logger_manager.per_engine_logger_dict[0]) == 1 + assert isinstance(dummy_logger, DummyLoggingStatLogger) + + stats_1 = IterationStats() + stats_1.num_preempted_reqs = 1 + stats_1.num_generation_tokens = 10 + stats_1.num_prompt_tokens = 100 + + stats_2 = IterationStats() + stats_2.num_preempted_reqs = 2 + stats_2.num_generation_tokens = 20 + stats_2.num_prompt_tokens = 200 + + # Expect the record will update the local iteration stats correctly + dummy_logger.record(scheduler_stats=None, iteration_stats=stats_1) + dummy_logger.record(scheduler_stats=None, iteration_stats=stats_2) + + assert dummy_logger.num_preempted_reqs == 3 + assert dummy_logger.num_generation_tokens == 30 + assert dummy_logger.num_prompt_tokens == 300 + + engine.shutdown() diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index d68d111c67ca..9d44254268f7 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -72,11 +72,13 @@ def _reset(self, now): # Tracked stats over current local logging interval. self.num_prompt_tokens: int = 0 self.num_generation_tokens: int = 0 + self.num_preempted_reqs: int = 0 def _track_iteration_stats(self, iteration_stats: IterationStats): # Save tracked stats for token counters. self.num_prompt_tokens += iteration_stats.num_prompt_tokens self.num_generation_tokens += iteration_stats.num_generation_tokens + self.num_preempted_reqs += iteration_stats.num_preempted_reqs def _get_throughput(self, tracked_stats: int, now: float) -> float: # Compute summary metrics for tracked stats