[Bugfix] Fix cuda profiler missing bug (#55237)
Wei Zhao committed
Sep 6, 2026 at 23:15 UTC
c3ec0d29f54cdd450576d819b1e5229ce6ef080e
2 files changed
+48
-3
tests/v1/engine/test_async_llm.py
+45
-1
@@ -4,10 +4,11 @@
4
import asyncio
5
import time
6
from contextlib import ExitStack
7
-from unittest.mock import MagicMock
7
+from unittest.mock import AsyncMock, MagicMock, call
8
9
import pytest
10
11
+import vllm.v1.engine.async_llm as async_llm_module
12
from vllm import SamplingParams
13
from vllm.assets.image import ImageAsset
14
from vllm.config import VllmConfig
@@ -59,6 +60,49 @@ VISION_PROMPT = {
60
}
61
62
63
+def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch):
64
+ vllm_config = MagicMock()
65
+ vllm_config.observability_config.otlp_traces_endpoint = None
66
+ vllm_config.scheduler_config.stream_interval = 1
67
+ vllm_config.profiler_config.profiler = "cuda"
68
+ vllm_config.profiler_config.ignore_frontend = False
69
+
70
+ renderer = MagicMock()
71
+ engine_core = MagicMock()
72
+ engine_core.profile_async = AsyncMock()
73
+ monkeypatch.setattr(
74
+ async_llm_module, "maybe_register_config_serialize_by_value", MagicMock()
75
+ )
76
+ monkeypatch.setattr(
77
+ async_llm_module,
78
+ "load_stat_logger_plugin_factories",
79
+ MagicMock(return_value=[]),
80
+ )
81
+ monkeypatch.setattr(
82
+ async_llm_module,
83
+ "renderer_from_config",
84
+ MagicMock(return_value=renderer),
85
+ )
86
+ monkeypatch.setattr(async_llm_module, "InputProcessor", MagicMock())
87
+ monkeypatch.setattr(async_llm_module, "OutputProcessor", MagicMock())
88
+ monkeypatch.setattr(
89
+ async_llm_module.EngineCoreClient,
90
+ "make_async_mp_client",
91
+ MagicMock(return_value=engine_core),
92
+ )
93
+
94
+ engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False)
95
+
96
+ async def profile():
97
+ await engine.start_profile()
98
+ await engine.stop_profile()
99
+
100
+ asyncio.run(profile())
101
+
102
+ assert engine.profiler is None
103
+ engine_core.profile_async.assert_has_awaits([call(True, None), call(False)])
104
+
105
+
106
async def generate(
107
engine: AsyncLLM,
108
request_id: str,
vllm/v1/engine/async_llm.py
+3
-2
@@ -7,7 +7,7 @@ import time
7
import warnings
8
from collections.abc import AsyncGenerator, Iterable, Mapping
9
from copy import copy
10
-from typing import Any, Optional
10
+from typing import Any
11
12
import vllm.envs as envs
13
from vllm import TokensPrompt
@@ -91,7 +91,7 @@ class AsyncLLM(EngineClient):
91
client_addresses: dict[str, Any] | None = None,
92
client_count: int = 1,
93
client_index: int = 0,
94
- profiler: Optional[TorchProfilerWrapper] = None, # type: ignore # noqa
94
+ profiler: TorchProfilerWrapper | None = None,
95
) -> None:
96
"""
97
Create an AsyncLLM.
@@ -190,6 +190,7 @@ class AsyncLLM(EngineClient):
190
except RuntimeError:
191
pass
192
193
+ self.profiler = profiler
194
if (
195
vllm_config.profiler_config.profiler == "torch"
196
and not vllm_config.profiler_config.ignore_frontend