diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 3845d6297a..5dd37a6f7e 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -4,10 +4,11 @@ import asyncio import time from contextlib import ExitStack -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, call import pytest +import vllm.v1.engine.async_llm as async_llm_module from vllm import SamplingParams from vllm.assets.image import ImageAsset from vllm.config import VllmConfig @@ -59,6 +60,49 @@ VISION_PROMPT = { } +def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): + vllm_config = MagicMock() + vllm_config.observability_config.otlp_traces_endpoint = None + vllm_config.scheduler_config.stream_interval = 1 + vllm_config.profiler_config.profiler = "cuda" + vllm_config.profiler_config.ignore_frontend = False + + renderer = MagicMock() + engine_core = MagicMock() + engine_core.profile_async = AsyncMock() + monkeypatch.setattr( + async_llm_module, "maybe_register_config_serialize_by_value", MagicMock() + ) + monkeypatch.setattr( + async_llm_module, + "load_stat_logger_plugin_factories", + MagicMock(return_value=[]), + ) + monkeypatch.setattr( + async_llm_module, + "renderer_from_config", + MagicMock(return_value=renderer), + ) + monkeypatch.setattr(async_llm_module, "InputProcessor", MagicMock()) + monkeypatch.setattr(async_llm_module, "OutputProcessor", MagicMock()) + monkeypatch.setattr( + async_llm_module.EngineCoreClient, + "make_async_mp_client", + MagicMock(return_value=engine_core), + ) + + engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) + + async def profile(): + await engine.start_profile() + await engine.stop_profile() + + asyncio.run(profile()) + + assert engine.profiler is None + engine_core.profile_async.assert_has_awaits([call(True, None), call(False)]) + + async def generate( engine: AsyncLLM, request_id: str, diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index cc132afd97..91618c4093 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -7,7 +7,7 @@ import time import warnings from collections.abc import AsyncGenerator, Iterable, Mapping from copy import copy -from typing import Any, Optional +from typing import Any import vllm.envs as envs from vllm import TokensPrompt @@ -91,7 +91,7 @@ class AsyncLLM(EngineClient): client_addresses: dict[str, Any] | None = None, client_count: int = 1, client_index: int = 0, - profiler: Optional[TorchProfilerWrapper] = None, # type: ignore # noqa + profiler: TorchProfilerWrapper | None = None, ) -> None: """ Create an AsyncLLM. @@ -190,6 +190,7 @@ class AsyncLLM(EngineClient): except RuntimeError: pass + self.profiler = profiler if ( vllm_config.profiler_config.profiler == "torch" and not vllm_config.profiler_config.ignore_frontend