Skip to content

Commit b447624

Browse files
authored
[Bugfix] Fix faulty triton importing logic when using Ray for DP (#19734)
Signed-off-by: mgoin <mgoin64@gmail.com>
1 parent cda9230 commit b447624

File tree

1 file changed

+17
-1
lines changed

1 file changed

+17
-1
lines changed

vllm/triton_utils/importing.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# SPDX-License-Identifier: Apache-2.0
22
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
33

4+
import os
45
import types
56
from importlib.util import find_spec
67

@@ -23,7 +24,22 @@
2324
x.driver for x in backends.values()
2425
if x.driver and x.driver.is_active()
2526
]
26-
if len(active_drivers) != 1:
27+
28+
# Check if we're in a distributed environment where CUDA_VISIBLE_DEVICES
29+
# might be temporarily empty (e.g., Ray sets it to "" during actor init)
30+
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
31+
is_distributed_env = (cuda_visible_devices is not None
32+
and len(cuda_visible_devices.strip()) == 0)
33+
34+
# Apply lenient driver check for distributed environments
35+
if is_distributed_env and len(active_drivers) == 0:
36+
# Allow 0 drivers in distributed environments - they may become
37+
# active later when CUDA context is properly initialized
38+
logger.debug(
39+
"Triton found 0 active drivers in distributed environment. "
40+
"This is expected during initialization.")
41+
elif not is_distributed_env and len(active_drivers) != 1:
42+
# Strict check for non-distributed environments
2743
logger.info(
2844
"Triton is installed but %d active driver(s) found "
2945
"(expected 1). Disabling Triton to prevent runtime errors.",

0 commit comments

Comments
 (0)