diff --git a/studio/backend/main.py b/studio/backend/main.py index 30ffb6d353..488ce936ff 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -17,6 +17,10 @@ os.environ["PYTHONWARNINGS"] = "ignore" # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. if sys.platform == "win32": + # Retained at module scope -- os.add_dll_directory returns a handle that + # removes the search-path entry when garbage collected. + _ROCM_DLL_HANDLES: list = [] + def _add_rocm_dll_dirs() -> None: candidates = [] # 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer @@ -40,7 +44,7 @@ if sys.platform == "win32": for _d in candidates: if os.path.isdir(_d): try: - os.add_dll_directory(_d) + _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d)) except (OSError, AttributeError): pass @@ -54,10 +58,18 @@ if sys.platform == "win32": # this the server process crashes with "Configured ROCm binary not found". # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION # before any import that pulls in bitsandbytes (mirrors worker.py logic). - # Guard: only set on ROCm hosts (HIP_PATH/ROCM_PATH present) -- setting - # BNB_ROCM_VERSION on a Windows CUDA machine makes bitsandbytes look for a - # ROCm DLL that doesn't exist and fail to initialise the CUDA backend. - _is_rocm_host = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")) + # Gate on the active torch runtime, not env-var presence -- HIP_PATH / + # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a + # CUDA torch wheel, and setting BNB_ROCM_VERSION there makes bitsandbytes + # look for a ROCm DLL that doesn't exist and crash the CUDA backend. + _is_rocm_host = False + if os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"): + try: + import torch as _torch_probe + _is_rocm_host = bool(getattr(_torch_probe.version, "hip", None)) + del _torch_probe + except Exception: + pass if _is_rocm_host and "BNB_ROCM_VERSION" not in os.environ: import glob as _glob diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 6929e5eab7..297cd126da 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -259,7 +259,10 @@ def _detect_windows_gfx_arch() -> str | None: return None text = result.stdout.decode(errors = "replace") m = re.search(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text) - return m.group(1).strip() if m else None + # Lowercase the captured token -- some hipinfo builds emit "Gfx1151" + # which would miss the lowercase keys in _GFX_TO_AMD_INDEX_ARCH and + # silently fall back to CPU torch. + return m.group(1).strip().lower() if m else None except Exception: return None @@ -384,8 +387,8 @@ def _detect_amd_gfx_codes() -> list[str]: _rocm_windows_torch_installed: bool = False -def _install_bnb_windows_rocm() -> None: - """Install the AMD Windows BNB prerelease wheel. +def _install_bnb_windows_rocm() -> bool: + """Install the AMD Windows BNB prerelease wheel. Returns True on success. The continuous-release wheel is intentionally mismatched: the filename encodes version 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the @@ -395,11 +398,11 @@ def _install_bnb_windows_rocm() -> None: """ _bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64") if _bnb_win_url is None: - return + return False _prev = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1" try: - pip_install_try( + _ok = pip_install_try( "bitsandbytes (AMD Windows, pre-release main)", "--force-reinstall", "--no-cache-dir", @@ -412,6 +415,8 @@ def _install_bnb_windows_rocm() -> None: os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) else: os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = _prev + if not _ok: + return False # After install: detect the actual ROCm DLL suffix from the wheel so any # post-install BNB import in this process loads the correct DLL. # The worker subprocess does the same detection independently (worker.py §1f). @@ -419,6 +424,7 @@ def _install_bnb_windows_rocm() -> None: if "BNB_ROCM_VERSION" not in os.environ: _ver = _detect_bnb_rocm_dll_ver() or "72" os.environ["BNB_ROCM_VERSION"] = _ver + return True def _ensure_rocm_torch() -> None: @@ -431,14 +437,38 @@ def _ensure_rocm_torch() -> None: Uses pip_install() to respect uv, constraints, and --python targeting. """ global _rocm_windows_torch_installed - # setup.ps1 sets this when it already installed AMD wheels; skip the probe. + # setup.ps1 sets this when it already installed AMD wheels; skip the probe + # only when torch is actually importable as ROCm. If the venv was wiped + # between runs, the stale env-var would suppress a needed reinstall. if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": - _rocm_windows_torch_installed = True - # setup.ps1 already installed ROCm torch, but we still need to install - # the AMD Windows BNB wheel here — the PyPI bitsandbytes wheel ships - # only CUDA DLLs and will fail to load on ROCm (no libbitsandbytes_rocm72.dll). - _install_bnb_windows_rocm() - return + _torch_ok = False + try: + _probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import torch; " + "hip=getattr(torch.version,'hip','') or ''; " + "import sys; " + "sys.exit(0 if (hip or 'rocm' in torch.__version__.lower()) else 1)" + ), + ], + stdout = subprocess.DEVNULL, + stderr = subprocess.DEVNULL, + timeout = 30, + ) + _torch_ok = _probe.returncode == 0 + except (OSError, subprocess.TimeoutExpired): + pass + if _torch_ok: + _rocm_windows_torch_installed = True + # setup.ps1 already installed ROCm torch, but we still need to install + # the AMD Windows BNB wheel here -- the PyPI bitsandbytes wheel ships + # only CUDA DLLs and will fail to load on ROCm. + _install_bnb_windows_rocm() + return + # torch was wiped between runs; fall through to the full install path if IS_MACOS: return @@ -488,11 +518,13 @@ def _ensure_rocm_torch() -> None: "torchaudio", constrain = False, ) - # Always install AMD Windows bitsandbytes — the PyPI wheel ships only + # Always install AMD Windows bitsandbytes -- the PyPI wheel ships only # CUDA DLLs and will fail to load on ROCm. Install even when torch was # already a ROCm build so that `studio update` repairs a broken bnb. - _install_bnb_windows_rocm() - _rocm_windows_torch_installed = True + # Only flip the success flag when the install actually succeeds; otherwise + # the post-install "manual install may be required" warning is suppressed. + if _install_bnb_windows_rocm(): + _rocm_windows_torch_installed = True return # ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ── diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 7011abc451..035a9ea3c6 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1788,11 +1788,19 @@ class TestDetectBnbRocmDllVer: class TestRocmTorchInstalledEnvVar: """Verify UNSLOTH_ROCM_TORCH_INSTALLED=1 skips main install but still installs BNB.""" + @staticmethod + def _ok_torch_probe(*a, **kw): + # subprocess.run probe returns 0 when torch imports as ROCm + rv = MagicMock() + rv.returncode = 0 + return rv + @patch.object(stack_mod, "_install_bnb_windows_rocm") @patch.object(stack_mod, "pip_install") def test_env_var_skips_main_pip_install(self, mock_pip, mock_bnb): """UNSLOTH_ROCM_TORCH_INSTALLED=1 should not trigger torch pip_install.""" - with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}): + with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \ + patch.object(stack_mod.subprocess, "run", side_effect=self._ok_torch_probe): stack_mod._ensure_rocm_torch() mock_pip.assert_not_called() @@ -1800,7 +1808,8 @@ class TestRocmTorchInstalledEnvVar: @patch.object(stack_mod, "pip_install") def test_env_var_calls_bnb_install(self, mock_pip, mock_bnb): """UNSLOTH_ROCM_TORCH_INSTALLED=1 should still call _install_bnb_windows_rocm.""" - with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}): + with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \ + patch.object(stack_mod.subprocess, "run", side_effect=self._ok_torch_probe): stack_mod._ensure_rocm_torch() mock_bnb.assert_called_once() @@ -1809,10 +1818,28 @@ class TestRocmTorchInstalledEnvVar: def test_env_var_sets_rocm_windows_flag(self, mock_pip, mock_bnb): """UNSLOTH_ROCM_TORCH_INSTALLED=1 should set _rocm_windows_torch_installed.""" stack_mod._rocm_windows_torch_installed = False - with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}): + with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \ + patch.object(stack_mod.subprocess, "run", side_effect=self._ok_torch_probe): stack_mod._ensure_rocm_torch() assert stack_mod._rocm_windows_torch_installed is True + @patch.object(stack_mod, "_install_bnb_windows_rocm") + @patch.object(stack_mod, "pip_install") + def test_env_var_falls_through_when_torch_missing(self, mock_pip, mock_bnb): + """If the venv was wiped between runs, the stale env-var must not suppress reinstall.""" + stack_mod._rocm_windows_torch_installed = False + def _bad_probe(*a, **kw): + rv = MagicMock() + rv.returncode = 1 + return rv + with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \ + patch.object(stack_mod.subprocess, "run", side_effect=_bad_probe), \ + patch.object(stack_mod, "IS_WINDOWS", False), \ + patch.object(stack_mod, "IS_MACOS", True): + stack_mod._ensure_rocm_torch() + # macOS branch is the next exit -- but the point is the early-return did NOT fire. + mock_bnb.assert_not_called() + # ============================================================================= # TEST: worker.py -- Windows ROCm patches (source-level checks)