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 15, 06:00 UTC
[Bugfix] Fix faulty triton importing logic when using Ray for DP (#19734)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jun 18, 2025 at 05:59 UTC
b447624ee399019740626f6217a05c00178ed17d
1 file changed
+17
-1
vllm/triton_utils/importing.py
+17
-1
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import os
5
import types
6
from importlib.util import find_spec
7
@@ -23,7 +24,22 @@ if HAS_TRITON:
24
x.driver for x in backends.values()
25
if x.driver and x.driver.is_active()
26
]
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
43
logger.info(
44
"Triton is installed but %d active driver(s) found "
45
"(expected 1). Disabling Triton to prevent runtime errors.",