[Bugfix] Don't attempt to use triton if no driver is active (#19561)

This commit is contained in:
Konrad Zawora 2025-06-14 21:30:54 +02:00 committed by GitHub
parent bc956b38d0
commit 861a0a0a39
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -12,6 +12,36 @@ HAS_TRITON = (
find_spec("triton") is not None
or find_spec("pytorch-triton-xpu") is not None # Not compatible
)
if HAS_TRITON:
try:
from triton.backends import backends
# It's generally expected that x.driver exists and has
# an is_active method.
# The `x.driver and` check adds a small layer of safety.
active_drivers = [
x.driver for x in backends.values()
if x.driver and x.driver.is_active()
]
if len(active_drivers) != 1:
logger.info(
"Triton is installed but %d active driver(s) found "
"(expected 1). Disabling Triton to prevent runtime errors.",
len(active_drivers))
HAS_TRITON = False
except ImportError:
# This can occur if Triton is partially installed or triton.backends
# is missing.
logger.warning(
"Triton is installed, but `triton.backends` could not be imported. "
"Disabling Triton.")
HAS_TRITON = False
except Exception as e:
# Catch any other unexpected errors during the check.
logger.warning(
"An unexpected error occurred while checking Triton active drivers:"
" %s. Disabling Triton.", e)
HAS_TRITON = False
if not HAS_TRITON:
logger.info("Triton not installed or not compatible; certain GPU-related"