diff --git a/setup.py b/setup.py index b6c1f5bc8ac3f..7dfcec7f9f0c5 100644 --- a/setup.py +++ b/setup.py @@ -324,21 +324,26 @@ def run(self) -> None: def _is_hpu() -> bool: - is_hpu_available = True + # if VLLM_TARGET_DEVICE env var was set explicitly, skip HPU autodetection + if os.getenv("VLLM_TARGET_DEVICE", None) == VLLM_TARGET_DEVICE: + return VLLM_TARGET_DEVICE == "hpu" + + # if VLLM_TARGET_DEVICE was not set explicitly, check if hl-smi succeeds, + # and if it doesn't, check if habanalabs driver is loaded + is_hpu_available = False try: - subprocess.run(["hl-smi"], capture_output=True, check=True) + out = subprocess.run(["hl-smi"], capture_output=True, check=True) + is_hpu_available = out.returncode == 0 except (FileNotFoundError, PermissionError, subprocess.CalledProcessError): - if not os.path.exists('/dev/accel/accel0') and not os.path.exists( - '/dev/accel/accel_controlD0'): - # last resort... + if sys.platform.startswith("linux"): try: output = subprocess.check_output( 'lsmod | grep habanalabs | wc -l', shell=True) is_hpu_available = int(output) > 0 except (ValueError, FileNotFoundError, PermissionError, subprocess.CalledProcessError): - is_hpu_available = False - return is_hpu_available or VLLM_TARGET_DEVICE == "hpu" + pass + return is_hpu_available def _no_device() -> bool: