fix: register ROCm DLL directory before torch import on Windows

Python 3.8+ ignores PATH for extension DLL loading on Windows; amdhip64.dll
and other HIP runtime DLLs must be registered via os.add_dll_directory().
Without this, torch.cuda.is_available() always returns False on AMD ROCm
Windows even when HIP_PATH is correctly set in system environment variables.

Reads HIP_PATH / ROCM_PATH env vars first, then falls back to scanning
common ROCm install roots (C:\Program Files\AMD\ROCm, F:\ROCm, C:\ROCm).
This commit is contained in:
LeoBorcherding 2026-05-06 20:43:18 -05:00
commit f09a4240a1

View file

@ -12,6 +12,40 @@ from pathlib import Path as _Path
# Suppress annoying C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
# On Windows, Python 3.8+ uses a secure DLL search that ignores PATH for
# extension modules. torch's HIP backend (amdhip64.dll etc.) won't be found
# even if F:\ROCm\...\bin is in PATH unless we explicitly register the
# directory with os.add_dll_directory(). Do this before any torch import.
if sys.platform == "win32":
import ctypes as _ctypes
def _add_rocm_dll_dirs() -> None:
hip_path = os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")
candidates = []
if hip_path:
candidates.append(os.path.join(hip_path, "bin"))
# Also scan common install roots in case HIP_PATH is not set
for _root in (r"C:\Program Files\AMD\ROCm", r"F:\ROCm", r"C:\ROCm"):
try:
if os.path.isdir(_root):
for _ver in sorted(os.listdir(_root), reverse=True):
_bin = os.path.join(_root, _ver, "bin")
if os.path.isdir(_bin):
candidates.append(_bin)
break
except OSError:
pass
for _d in candidates:
if os.path.isdir(_d):
try:
os.add_dll_directory(_d)
except (OSError, AttributeError):
pass
_add_rocm_dll_dirs()
del _add_rocm_dll_dirs, _ctypes
# Ensure backend dir is on sys.path so _platform_compat is importable when
# main.py is launched directly (e.g. `uvicorn main:app`).
_backend_dir = str(_Path(__file__).parent)