| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copy |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from tests.plugins.vllm_add_dummy_stat_logger.dummy_stat_logger.dummy_stat_logger import ( # noqa E501 |
| 8 | DummyStatLogger, |
| 9 | ) |
| 10 | from tests.utils import wait_for_memory_to_settle |
| 11 | from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM |
| 12 | from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger |
| 13 | |
| 14 | |
| 15 | @pytest.fixture |
| 16 | def log_stats_enabled_engine_args(): |
| 17 | """Shared fixture providing common AsyncEngineArgs configuration |
| 18 | used across multiple tests. |
| 19 | """ |
| 20 | return AsyncEngineArgs( |
| 21 | model="distilbert/distilgpt2", |
| 22 | dtype="half", |
| 23 | disable_log_stats=False, |
| 24 | enforce_eager=True, |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | @pytest.mark.asyncio |
| 29 | async def test_async_llm_replace_default_loggers(log_stats_enabled_engine_args): |
| 30 | """RayPrometheusStatLogger should replace the default PrometheusStatLogger.""" |
| 31 | engine = AsyncLLM.from_engine_args( |
| 32 | log_stats_enabled_engine_args, stat_loggers=[RayPrometheusStatLogger] |
| 33 | ) |
| 34 | try: |
| 35 | assert isinstance( |
| 36 | engine.logger_manager.stat_loggers[0], RayPrometheusStatLogger |
| 37 | ) |
| 38 | finally: |
| 39 | engine.shutdown() |
| 40 | wait_for_memory_to_settle( |
| 41 | threshold_ratio=1.0 - log_stats_enabled_engine_args.gpu_memory_utilization |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | @pytest.mark.asyncio |
| 46 | async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args): |
| 47 | """It's still possible to use custom stat loggers exclusively by passing |
| 48 | disable_log_stats=True in addition to a list of custom stat loggers. |
| 49 | """ |
| 50 | # Create engine_args with disable_log_stats=True for this test |
| 51 | disabled_log_engine_args = copy.deepcopy(log_stats_enabled_engine_args) |
| 52 | disabled_log_engine_args.disable_log_stats = True |
| 53 | |
| 54 | # Disable default loggers; pass custom stat logger to the constructor |
| 55 | engine = AsyncLLM.from_engine_args( |
| 56 | disabled_log_engine_args, stat_loggers=[DummyStatLogger] |
| 57 | ) |
| 58 | |
| 59 | try: |
| 60 | assert len(engine.logger_manager.stat_loggers) == 2 |
| 61 | assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1 |
| 62 | assert isinstance( |
| 63 | engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0], |
| 64 | DummyStatLogger, |
| 65 | ) |
| 66 | |
| 67 | # log_stats is still True, since custom stat loggers are used |
| 68 | assert engine.log_stats |
| 69 | finally: |
| 70 | engine.shutdown() |
| 71 | wait_for_memory_to_settle( |
| 72 | threshold_ratio=1.0 - disabled_log_engine_args.gpu_memory_utilization |
| 73 | ) |