Skip to content

Commit 8b4021a

Browse files
mengchengTangmengchengTang
committed
vllm ascend profiler system support max_iteration and delay_iteration
Signed-off-by: mengchengTang <745274877@qq.com>
1 parent 49b9147 commit 8b4021a

5 files changed

Lines changed: 416 additions & 221 deletions

File tree

‎tests/ut/profiler/test_torch_npu_profiler.py‎

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
from unittest.mock import MagicMock, patch
2+
3+
from vllm.config import ProfilerConfig
4+
5+
from tests.ut.base import TestBase
6+
7+
8+
class TestTorchNPUProfilerWrapper(TestBase):
9+
def test_init_creates_underlying_profiler(self):
10+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
11+
12+
profiler_config = ProfilerConfig(
13+
profiler="torch",
14+
torch_profiler_dir="/path/to/traces",
15+
)
16+
mock_profiler = MagicMock()
17+
18+
with patch.object(TorchNPUProfilerWrapper, "_create_profiler", return_value=mock_profiler) as mock_create:
19+
wrapper = TorchNPUProfilerWrapper(profiler_config, "dp0_pp0_tp0_dcp0_ep0_rank0")
20+
21+
mock_create.assert_called_once_with(profiler_config, "dp0_pp0_tp0_dcp0_ep0_rank0")
22+
self.assertIs(wrapper.profiler, mock_profiler)
23+
24+
def test_start_stop_delegate_to_underlying_profiler(self):
25+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
26+
27+
profiler_config = ProfilerConfig(
28+
profiler="torch",
29+
torch_profiler_dir="/path/to/traces",
30+
)
31+
mock_profiler = MagicMock()
32+
33+
with patch.object(TorchNPUProfilerWrapper, "_create_profiler", return_value=mock_profiler):
34+
wrapper = TorchNPUProfilerWrapper(profiler_config, "trace_name")
35+
36+
wrapper._start()
37+
wrapper._stop()
38+
39+
mock_profiler.start.assert_called_once()
40+
mock_profiler.stop.assert_called_once()
41+
42+
@patch("vllm_ascend.profiler.torch_npu_profiler.envs_ascend")
43+
@patch("torch_npu.profiler._ExperimentalConfig")
44+
@patch("torch_npu.profiler.profile")
45+
@patch("torch_npu.profiler.tensorboard_trace_handler")
46+
@patch("torch_npu.profiler.ExportType")
47+
@patch("torch_npu.profiler.ProfilerLevel")
48+
@patch("torch_npu.profiler.AiCMetrics")
49+
@patch("torch_npu.profiler.ProfilerActivity")
50+
def test_create_profiler_enabled(
51+
self,
52+
mock_profiler_activity,
53+
mock_aic_metrics,
54+
mock_profiler_level,
55+
mock_export_type,
56+
mock_trace_handler,
57+
mock_profile,
58+
mock_experimental_config,
59+
mock_envs_ascend,
60+
):
61+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
62+
63+
mock_envs_ascend.MSMONITOR_USE_DAEMON = 0
64+
65+
profiler_config = ProfilerConfig(
66+
profiler="torch",
67+
torch_profiler_dir="/path/to/traces",
68+
torch_profiler_with_stack=True,
69+
torch_profiler_with_memory=True,
70+
)
71+
72+
mock_export_type.Text = "Text"
73+
mock_profiler_level.Level1 = "Level1"
74+
mock_aic_metrics.AiCoreNone = "AiCoreNone"
75+
mock_profiler_activity.CPU = "CPU"
76+
mock_profiler_activity.NPU = "NPU"
77+
78+
mock_trace_handler_instance = MagicMock()
79+
mock_trace_handler.return_value = mock_trace_handler_instance
80+
mock_profiler_instance = MagicMock()
81+
mock_profile.return_value = mock_profiler_instance
82+
83+
result = TorchNPUProfilerWrapper._create_profiler(
84+
profiler_config,
85+
"warmup_dp0_pp0_tp0_dcp0_ep0_rank0",
86+
)
87+
88+
mock_experimental_config.assert_called_once()
89+
config_kwargs = mock_experimental_config.call_args.kwargs
90+
expected_config = {
91+
"export_type": "Text",
92+
"profiler_level": "Level1",
93+
"msprof_tx": False,
94+
"aic_metrics": "AiCoreNone",
95+
"l2_cache": False,
96+
"op_attr": False,
97+
"data_simplification": True,
98+
"record_op_args": False,
99+
"gc_detect_threshold": None,
100+
}
101+
for key, expected_value in expected_config.items():
102+
self.assertEqual(config_kwargs[key], expected_value)
103+
104+
mock_trace_handler.assert_called_once_with(
105+
"/path/to/traces",
106+
worker_name="warmup_dp0_pp0_tp0_dcp0_ep0_rank0",
107+
)
108+
109+
mock_profile.assert_called_once()
110+
profile_kwargs = mock_profile.call_args.kwargs
111+
self.assertEqual(profile_kwargs["activities"], ["CPU", "NPU"])
112+
self.assertTrue(profile_kwargs["profile_memory"])
113+
self.assertEqual(profile_kwargs["with_modules"], True)
114+
self.assertEqual(profile_kwargs["on_trace_ready"], mock_trace_handler_instance)
115+
self.assertEqual(result, mock_profiler_instance)
116+
117+
def test_create_profiler_disabled(self):
118+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
119+
120+
profiler_config = ProfilerConfig(profiler=None, torch_profiler_dir="")
121+
122+
with self.assertRaises(RuntimeError) as cm:
123+
TorchNPUProfilerWrapper._create_profiler(profiler_config, "test_trace")
124+
125+
self.assertIn("Unrecognized profiler: None", str(cm.exception))
126+
127+
def test_create_profiler_empty_dir(self):
128+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
129+
130+
profiler_config = MagicMock()
131+
profiler_config.profiler = "torch"
132+
profiler_config.torch_profiler_dir = ""
133+
134+
with self.assertRaises(RuntimeError) as cm:
135+
TorchNPUProfilerWrapper._create_profiler(profiler_config, "test_trace")
136+
137+
self.assertIn("torch_profiler_dir cannot be empty", str(cm.exception))
138+
139+
@patch("vllm_ascend.profiler.torch_npu_profiler.envs_ascend")
140+
def test_create_profiler_raises_when_msmonitor_enabled(self, mock_envs_ascend):
141+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
142+
143+
mock_envs_ascend.MSMONITOR_USE_DAEMON = 1
144+
profiler_config = ProfilerConfig(
145+
profiler="torch",
146+
torch_profiler_dir="/path/to/traces",
147+
)
148+
149+
with self.assertRaises(RuntimeError) as cm:
150+
TorchNPUProfilerWrapper._create_profiler(profiler_config, "test_trace")
151+
152+
self.assertIn(
153+
"MSMONITOR_USE_DAEMON and torch profiler cannot be both enabled at the same time.",
154+
str(cm.exception),
155+
)
156+
157+
def test_profiler_step_returns_true(self):
158+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
159+
160+
profiler_config = ProfilerConfig(
161+
profiler="torch",
162+
torch_profiler_dir="/path/to/traces",
163+
)
164+
165+
with patch.object(TorchNPUProfilerWrapper, "_create_profiler", return_value=MagicMock()):
166+
wrapper = TorchNPUProfilerWrapper(profiler_config, "trace_name")
167+
168+
self.assertTrue(wrapper._profiler_step())
169+
170+
def test_step_calls_underlying_start_after_delay_iterations(self):
171+
"""Work matches vLLM WorkerProfiler: first N worker steps defer _start."""
172+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
173+
174+
profiler_config = ProfilerConfig(
175+
profiler="torch",
176+
torch_profiler_dir="/path/to/traces",
177+
delay_iterations=3,
178+
max_iterations=0,
179+
)
180+
mock_profiler = MagicMock()
181+
182+
with patch.object(TorchNPUProfilerWrapper, "_create_profiler", return_value=mock_profiler):
183+
wrapper = TorchNPUProfilerWrapper(profiler_config, "trace_name")
184+
185+
wrapper.start()
186+
mock_profiler.start.assert_not_called()
187+
188+
wrapper.step()
189+
wrapper.step()
190+
mock_profiler.start.assert_not_called()
191+
192+
wrapper.step()
193+
mock_profiler.start.assert_called_once()
194+
195+
def test_step_stops_underlying_profiler_after_max_iterations(self):
196+
"""Work matches vLLM WorkerProfiler: stop when recorded steps exceed max_iterations."""
197+
from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper
198+
199+
profiler_config = ProfilerConfig(
200+
profiler="torch",
201+
torch_profiler_dir="/path/to/traces",
202+
delay_iterations=0,
203+
max_iterations=1,
204+
)
205+
mock_profiler = MagicMock()
206+
207+
with patch.object(TorchNPUProfilerWrapper, "_create_profiler", return_value=mock_profiler):
208+
wrapper = TorchNPUProfilerWrapper(profiler_config, "trace_name")
209+
210+
wrapper.start()
211+
mock_profiler.start.assert_called_once()
212+
mock_profiler.stop.assert_not_called()
213+
214+
wrapper.step()
215+
mock_profiler.stop.assert_not_called()
216+
217+
wrapper.step()
218+
mock_profiler.stop.assert_called_once()

0 commit comments

Comments
 (0)