[ROCm][Bugfix] Fix profiler in TheRock image (#56190)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Sep 10, 2026 at 11:39 UTC
48cb12c184e9b753c5de1c8ecab7953ded127b87
1 file changed
+74
-11
vllm/env_override.py
+74
-11
@@ -5,21 +5,28 @@ import importlib.util
5
import os
6
7
8
-def _get_torch_cuda_version():
9
- """Peripheral function to _maybe_set_cuda_compatibility_path().
8
+def _get_torch_root():
9
+ """Locate the installed torch package without importing it."""
10
+ spec = importlib.util.find_spec("torch")
11
+ if not spec:
12
+ return None
13
+ if spec.origin:
14
+ return os.path.dirname(spec.origin)
15
+ if spec.submodule_search_locations:
16
+ return spec.submodule_search_locations[0]
17
+ return None
18
+
19
+
20
+def _get_torch_version_attr(attr):
21
+ """Read an attribute of torch.version without importing torch.
22
+
23
PyTorch version must not be determined by importing directly
24
because it will trigger the CUDA initialization, losing the
25
chance to set the LD_LIBRARY_PATH beforehand.
26
"""
27
try:
15
- spec = importlib.util.find_spec("torch")
16
- if not spec:
17
- return None
18
- if spec.origin:
19
- torch_root = os.path.dirname(spec.origin)
20
- elif spec.submodule_search_locations:
21
- torch_root = spec.submodule_search_locations[0]
22
- else:
28
+ torch_root = _get_torch_root()
29
+ if not torch_root:
30
return None
31
version_path = os.path.join(torch_root, "version.py")
32
if not os.path.exists(version_path):
@@ -31,11 +38,16 @@ def _get_torch_cuda_version():
38
module = importlib.util.module_from_spec(ver_spec)
39
# Avoid registering in sys.modules to not confuse future imports
40
ver_spec.loader.exec_module(module)
34
- return getattr(module, "cuda", None)
41
+ return getattr(module, attr, None)
42
except Exception:
43
return None
44
45
46
+def _get_torch_cuda_version():
47
+ """Peripheral function to _maybe_set_cuda_compatibility_path()."""
48
+ return _get_torch_version_attr("cuda")
49
+
50
+
51
def _maybe_set_cuda_compatibility_path():
52
"""Set LD_LIBRARY_PATH for CUDA forward compatibility if enabled.
53
@@ -82,7 +94,58 @@ def _maybe_set_cuda_compatibility_path():
94
os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(new_paths)
95
96
97
+def _maybe_promote_torch_symbols_for_rocm():
98
+ """Put libtorch_cpu.so in the global symbol scope on ROCm, for GPU profiling.
99
+
100
+ Must run before 'import torch'. libkineto advertises itself to
101
+ rocprofiler-sdk by exporting rocprofiler_configure from libtorch_cpu.so,
102
+ and rocprofiler-sdk looks for that symbol exactly once, from a lazy
103
+ initializer that runs *during* 'import torch'. CPython loads extension
104
+ modules RTLD_LOCAL, so unless something has already placed the symbol in
105
+ the global scope the lookup misses and no client is ever registered.
106
+
107
+ The failure is silent and total: no queue interception, so roctracer
108
+ yields no GPU records, while torch.profiler still writes a complete
109
+ CPU-only trace and reports success.
110
+
111
+ Best effort. If the library cannot be loaded early we leave the process
112
+ exactly as it would have been, losing only GPU tracing.
113
+ """
114
+ if _get_torch_version_attr("hip") is None:
115
+ return
116
+ try:
117
+ import ctypes
118
+ import glob
119
+
120
+ torch_root = _get_torch_root()
121
+ if not torch_root:
122
+ return
123
+ lib = os.path.join(torch_root, "lib", "libtorch_cpu.so")
124
+ if not os.path.exists(lib):
125
+ return
126
+ try:
127
+ ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
128
+ except OSError:
129
+ # Some ROCm wheel layouts leave a stale RUNPATH in libtorch_cpu.so,
130
+ # so a dependency that resolves once torch has set things up does
131
+ # not resolve for a standalone load. Satisfy it by search, then
132
+ # retry; give up quietly if it is not where we expect.
133
+ site_root = os.path.dirname(torch_root)
134
+ deps = glob.glob(
135
+ os.path.join(
136
+ site_root, "_rocm_sdk_*", "lib", "*", "lib", "librocm-openblas.so.*"
137
+ )
138
+ )
139
+ if not deps:
140
+ return
141
+ ctypes.CDLL(deps[0], mode=ctypes.RTLD_LOCAL)
142
+ ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
143
+ except Exception:
144
+ return
145
+
146
+
147
_maybe_set_cuda_compatibility_path()
148
+_maybe_promote_torch_symbols_for_rocm()
149
150
import torch
151