A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 18:15 UTC
Support custom URI schemes and trace handlers for profiler (#32393)
David Ramon Prados committed
Jan 22, 2026 at 12:45 UTC
3a63be0faa8e45401b9eddaafaca3deecac0a882
3 files changed
+74
-19
tests/v1/worker/test_gpu_profiler.py
+34
@@ -3,6 +3,7 @@
3
import pytest
4
5
from vllm.config import ProfilerConfig
6
+from vllm.config.profiler import _is_uri_path
7
from vllm.profiler.wrapper import WorkerProfiler
8
9
@@ -202,3 +203,36 @@ def test_mixed_delay_and_stop(default_profiler_config):
203
profiler.step()
204
205
assert profiler.start_call_count == 0
206
+
207
+
208
+class TestIsUriPath:
209
+ """Tests for the _is_uri_path helper function."""
210
+
211
+ @pytest.mark.parametrize(
212
+ "path,expected",
213
+ [
214
+ # Valid URI schemes - should return True
215
+ ("gs://bucket/path", True),
216
+ ("s3://bucket/path", True),
217
+ ("hdfs://cluster/path", True),
218
+ ("abfs://container/path", True),
219
+ ("http://example.com/path", True),
220
+ ("https://example.com/path", True),
221
+ # Local paths - should return False
222
+ ("/tmp/local/path", False),
223
+ ("./relative/path", False),
224
+ ("relative/path", False),
225
+ ("/absolute/path", False),
226
+ # Windows drive letters - should return False (single char scheme)
227
+ ("C://windows/path", False),
228
+ ("D://drive/path", False),
229
+ # Edge cases
230
+ ("", False),
231
+ ("no-scheme", False),
232
+ ("scheme-no-slashes:", False),
233
+ ("://no-scheme", False),
234
+ ],
235
+ )
236
+ def test_is_uri_path(self, path, expected):
237
+ """Test that _is_uri_path correctly identifies URI vs local paths."""
238
+ assert _is_uri_path(path) == expected
vllm/config/profiler.py
+19
-11
@@ -18,6 +18,20 @@ logger = init_logger(__name__)
18
ProfilerKind = Literal["torch", "cuda"]
19
20
21
+def _is_uri_path(path: str) -> bool:
22
+ """Check if path is a URI (scheme://...), excluding Windows drive letters.
23
+
24
+ Supports custom URI schemes like gs://, s3://, hdfs://, etc.
25
+ These paths should not be converted to absolute paths.
26
+ """
27
+ if "://" in path:
28
+ scheme = path.split("://")[0]
29
+ # Windows drive letters are single characters (e.g., C://)
30
+ # Valid URI schemes have more than one character
31
+ return len(scheme) > 1
32
+ return False
33
+
34
+
35
@config
36
@dataclass
37
class ProfilerConfig:
@@ -54,7 +68,7 @@ class ProfilerConfig:
68
Disabled by default."""
69
70
ignore_frontend: bool = False
57
- """If `True`, disables the front-end profiling of AsyncLLM when using the
71
+ """If `True`, disables the front-end profiling of AsyncLLM when using the
72
'torch' profiler. This is needed to reduce overhead when using delay/limit options,
73
since the front-end profiling does not track iterations and will capture the
74
entire range.
@@ -185,15 +199,9 @@ class ProfilerConfig:
199
if self.profiler == "torch" and not profiler_dir:
200
raise ValueError("torch_profiler_dir must be set when profiler is 'torch'")
201
188
- if profiler_dir:
189
- is_gs_path = (
190
- profiler_dir.startswith("gs://")
191
- and profiler_dir[5:]
192
- and profiler_dir[5] != "/"
193
- )
194
- if not is_gs_path:
195
- self.torch_profiler_dir = os.path.abspath(
196
- os.path.expanduser(profiler_dir)
197
- )
202
+ # Support any URI scheme (gs://, s3://, hdfs://, etc.)
203
+ # These paths should not be converted to absolute paths
204
+ if profiler_dir and not _is_uri_path(profiler_dir):
205
+ self.torch_profiler_dir = os.path.abspath(os.path.expanduser(profiler_dir))
206
207
return self
vllm/profiler/wrapper.py
+21
-8
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import ABC, abstractmethod
5
+from collections.abc import Callable
6
from contextlib import nullcontext
7
from typing import Literal
8
@@ -9,6 +10,7 @@ import torch
10
from typing_extensions import override
11
12
from vllm.config import ProfilerConfig
13
+from vllm.config.profiler import _is_uri_path
14
from vllm.logger import init_logger
15
16
logger = init_logger(__name__)
@@ -151,6 +153,7 @@ class TorchProfilerWrapper(WorkerProfiler):
153
worker_name: str,
154
local_rank: int,
155
activities: list[TorchProfilerActivity],
156
+ on_trace_ready: Callable[[torch.profiler.profile], None] | None = None,
157
) -> None:
158
super().__init__(profiler_config)
159
@@ -172,6 +175,17 @@ class TorchProfilerWrapper(WorkerProfiler):
175
profiler_config.torch_profiler_with_flops,
176
)
177
178
+ # Determine trace handler: use custom handler if provided,
179
+ # otherwise default to tensorboard trace handler
180
+ if on_trace_ready is not None:
181
+ trace_handler = on_trace_ready
182
+ else:
183
+ trace_handler = torch.profiler.tensorboard_trace_handler(
184
+ torch_profiler_trace_dir,
185
+ worker_name=worker_name,
186
+ use_gzip=profiler_config.torch_profiler_use_gzip,
187
+ )
188
+
189
self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
190
self.profiler = torch.profiler.profile(
191
activities=[TorchProfilerActivityMap[activity] for activity in activities],
@@ -179,11 +193,7 @@ class TorchProfilerWrapper(WorkerProfiler):
193
profile_memory=profiler_config.torch_profiler_with_memory,
194
with_stack=profiler_config.torch_profiler_with_stack,
195
with_flops=profiler_config.torch_profiler_with_flops,
182
- on_trace_ready=torch.profiler.tensorboard_trace_handler(
183
- torch_profiler_trace_dir,
184
- worker_name=worker_name,
185
- use_gzip=profiler_config.torch_profiler_use_gzip,
186
- ),
196
+ on_trace_ready=trace_handler,
197
)
198
199
@override
@@ -198,12 +208,15 @@ class TorchProfilerWrapper(WorkerProfiler):
208
rank = self.local_rank
209
if profiler_config.torch_profiler_dump_cuda_time_total:
210
profiler_dir = profiler_config.torch_profiler_dir
201
- profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
211
sort_key = "self_cuda_time_total"
212
table = self.profiler.key_averages().table(sort_by=sort_key)
213
205
- with open(profiler_out_file, "w") as f:
206
- print(table, file=f)
214
+ # Skip file write for URI paths (gs://, s3://, etc.)
215
+ # as standard file I/O doesn't work with URI schemes
216
+ if not _is_uri_path(profiler_dir):
217
+ profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
218
+ with open(profiler_out_file, "w") as f:
219
+ print(table, file=f)
220
221
# only print profiler results on rank 0
222
if rank == 0: