[Bugfix] Don't attempt to use triton if no driver is active (#19561)
This commit is contained in:
@@ -12,6 +12,36 @@ HAS_TRITON = (
|
|||||||
find_spec("triton") is not None
|
find_spec("triton") is not None
|
||||||
or find_spec("pytorch-triton-xpu") is not None # Not compatible
|
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:
|
if not HAS_TRITON:
|
||||||
logger.info("Triton not installed or not compatible; certain GPU-related"
|
logger.info("Triton not installed or not compatible; certain GPU-related"
|
||||||
|
|||||||
Reference in New Issue
Block a user