[Bugfix][Profiler] Fix API server crash on double /stop_profile (#51678)
Signed-off-by: Alberto Massidda <amassidda@nvidia.com> Co-authored-by: Alberto Massidda <amassidda@nvidia.com>
Alberto Massidda committed
Sep 1, 2026 at 22:05 UTC
55178f2d09309c42bad093484c4c1d5df340ab14
1 file changed
+8
-15
vllm/v1/engine/async_llm.py
+8
-15
@@ -7,9 +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
11
-
12
-import torch
10
+from typing import Any, Optional
11
12
import vllm.envs as envs
13
from vllm import TokensPrompt
@@ -34,6 +32,7 @@ from vllm.lora.request import LoRARequest
32
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
33
from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput
34
from vllm.pooling_params import PoolingParams
35
+from vllm.profiler.wrapper import TorchProfilerWrapper
36
from vllm.renderers import renderer_from_config
37
from vllm.renderers.inputs.preprocess import extract_prompt_components
38
from vllm.sampling_params import RequestOutputKind, SamplingParams
@@ -92,6 +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
95
) -> None:
96
"""
97
Create an AsyncLLM.
@@ -195,19 +195,12 @@ class AsyncLLM(EngineClient):
195
profiler_dir,
196
)
197
worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
198
- self.profiler = torch.profiler.profile(
199
- activities=[
200
- torch.profiler.ProfilerActivity.CPU,
201
- ],
202
- with_stack=vllm_config.profiler_config.torch_profiler_with_stack,
203
- on_trace_ready=torch.profiler.tensorboard_trace_handler(
204
- profiler_dir,
205
- worker_name=worker_name,
206
- use_gzip=vllm_config.profiler_config.torch_profiler_use_gzip,
207
- ),
198
+ self.profiler = TorchProfilerWrapper(
199
+ vllm_config.profiler_config,
200
+ worker_name=worker_name,
201
+ local_rank=0,
202
+ activities=["CPU"],
203
)
209
- else:
210
- self.profiler = None
204
205
@classmethod
206
def from_vllm_config(