[ROCm][Bugfix] Fix `use_v2_model_runner` inside Ray driver thread (#45998)
Signed-off-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Micah Williamson committed
Jun 23, 2026 at 19:41 UTC
e2bdc24612ab0b7bf7a1bc67c955fd244e8660c4
2 files changed
+23
-4
.buildkite/test_areas/distributed.yaml
+13
-1
@@ -210,7 +210,7 @@ steps:
210
- TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
211
- pytest -v -s -x lora/test_mixtral.py
212
213
-- label: Distributed Tests (2 GPUs)(H100)
213
+- label: Distributed Tests (2xH100-2xMI300)
214
key: distributed-tests-2-gpus-h100
215
timeout_in_minutes: 15
216
device: h100
@@ -224,6 +224,18 @@ steps:
224
- pytest -v -s tests/v1/distributed/test_dbo.py
225
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
226
- pytest -v -s tests/distributed/test_packed_tensor.py
227
+ mirror:
228
+ amd:
229
+ device: mi300_2
230
+ timeout_in_minutes: 180
231
+ depends_on:
232
+ - image-build-amd
233
+ commands:
234
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
235
+ - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048
236
+ - pytest -v -s tests/v1/distributed/test_dbo.py
237
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
238
+ - pytest -v -s tests/distributed/test_packed_tensor.py
239
240
- label: Distributed Tests (2 GPUs)(B200)
241
key: distributed-tests-2-gpus-b200
vllm/triton_utils/importing.py
+10
-3
@@ -7,6 +7,7 @@ from importlib.metadata import version
7
from importlib.util import find_spec
8
9
from vllm.logger import init_logger
10
+from vllm.platforms import current_platform
11
from vllm.utils.math_utils import cdiv
12
13
logger = init_logger(__name__)
@@ -27,10 +28,16 @@ if HAS_TRITON:
28
]
29
30
# Check if we're in a distributed environment where CUDA_VISIBLE_DEVICES
30
- # might be temporarily empty (e.g., Ray sets it to "" during actor init)
31
- cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
31
+ # or HIP_VISIBLE_DEVICES might be temporarily empty (e.g., Ray sets it to ""
32
+ # during actor init)
33
+ visible_devices_env = (
34
+ "HIP_VISIBLE_DEVICES"
35
+ if current_platform.is_rocm()
36
+ else "CUDA_VISIBLE_DEVICES"
37
+ )
38
+ visible_devices = os.environ.get(visible_devices_env)
39
is_distributed_env = (
33
- cuda_visible_devices is not None and len(cuda_visible_devices.strip()) == 0
40
+ visible_devices is not None and len(visible_devices.strip()) == 0
41
)
42
43
# Apply lenient driver check for distributed environments