From 8af9fe63a38a05435085e28034a09356acb9c692 Mon Sep 17 00:00:00 2001 From: Bubu Date: Wed, 10 Jun 2026 21:01:05 +0800 Subject: [PATCH] fix: persist Windows ROCm BNB version (#6048) * fix: persist Windows ROCm BNB version * style: apply kwarg spacing hook * fix: avoid persisting caller ROCm overrides * fix: redetect managed BNB ROCm defaults * style: apply ROCm guard test formatting --------- Co-authored-by: Daniel Han --- studio/backend/core/training/worker.py | 13 +- studio/backend/main.py | 10 +- studio/install_python_stack.py | 94 +++++++++- tests/studio/install/test_rocm_support.py | 211 ++++++++++++++++++++++ 4 files changed, 315 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index f7fbaa0b69..b6ec057231 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2026,9 +2026,13 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB # wheel may ship a DLL whose suffix doesn't match. Detect the actual - # DLL name and override; "72" is a safe fallback. Callers may - # pre-set the var to override. - if "BNB_ROCM_VERSION" not in os.environ: + # DLL name and override; "72" is a safe fallback. Values seeded by + # the installer are redetectable defaults, while caller overrides + # remain authoritative. + if ( + "BNB_ROCM_VERSION" not in os.environ + or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize" + ): _bnb_rocm_ver = None try: import glob as _glob @@ -2053,8 +2057,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> _bnb_rocm_ver = max(_all_vers, key = lambda v: int(v)) except Exception: pass - _bnb_rocm_ver = _bnb_rocm_ver or "72" + _bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver + os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected" logger.info( "Windows ROCm: set BNB_ROCM_VERSION=%s " "(detected from installed BNB wheel; " diff --git a/studio/backend/main.py b/studio/backend/main.py index 38605a5009..74a7dbc8c5 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -86,7 +86,12 @@ if sys.platform == "win32": # found") without this. Detect the shipped DLL and fall back to "72" (mirrors # worker.py). Gate on the rocm bnb DLL / HIP_PATH rather than torch.version.hip # to avoid importing torch on every Windows host. - if "BNB_ROCM_VERSION" not in os.environ: + # Values seeded by the installer's sitecustomize.py are redetectable + # defaults; explicit caller values remain authoritative. + if ( + "BNB_ROCM_VERSION" not in os.environ + or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize" + ): import glob as _glob import logging as _logging @@ -118,8 +123,9 @@ if sys.platform == "win32": ) # rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72") if _found_rocm_bnb or _hip_env: - _bnb_rocm_ver_final = _bnb_rocm_ver or "72" + _bnb_rocm_ver_final = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final + os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected" _logging.getLogger(__name__).info( "Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)", _bnb_rocm_ver_final, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e742d85eed..eaab3f5559 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -18,6 +18,7 @@ import re import shutil import subprocess import sys +import sysconfig import tempfile import urllib.request from pathlib import Path @@ -489,6 +490,73 @@ def _detect_bnb_rocm_dll_ver() -> str | None: return max(all_vers, key = lambda v: int(v)) if all_vers else None +_BNB_ROCM_SITECUSTOMIZE_BEGIN = "# BEGIN Unsloth BNB_ROCM_VERSION" +_BNB_ROCM_SITECUSTOMIZE_END = "# END Unsloth BNB_ROCM_VERSION" +_BNB_ROCM_VERSION_SOURCE_ENV = "UNSLOTH_BNB_ROCM_VERSION_SOURCE" +_BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE = "sitecustomize" +_BNB_ROCM_VERSION_SOURCE_DETECTED = "detected" + + +def _persist_bnb_rocm_version(version: str) -> bool: + """Persist BNB_ROCM_VERSION for future Python processes in this venv.""" + version = str(version).strip() + if not version: + return False + + site_packages = sysconfig.get_path("purelib") + if not site_packages: + return False + + sitecustomize_path = Path(site_packages) / "sitecustomize.py" + block = ( + f"{_BNB_ROCM_SITECUSTOMIZE_BEGIN}\n" + "import os as _unsloth_os\n" + "_unsloth_existing_bnb_rocm = _unsloth_os.environ.get('BNB_ROCM_VERSION')\n" + f"_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', {version!r})\n" + "if _unsloth_existing_bnb_rocm is None and " + f"_unsloth_os.environ.get('BNB_ROCM_VERSION') == {version!r}:\n" + " _unsloth_os.environ.setdefault(" + f"{_BNB_ROCM_VERSION_SOURCE_ENV!r}, " + f"{_BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE!r})\n" + "del _unsloth_existing_bnb_rocm\n" + f"{_BNB_ROCM_SITECUSTOMIZE_END}\n" + ) + + try: + sitecustomize_path.parent.mkdir(parents = True, exist_ok = True) + existing = ( + sitecustomize_path.read_text(encoding = "utf-8") if sitecustomize_path.exists() else "" + ) + # Strip all managed regions, including one whose END marker was lost to + # an interrupted write, then append exactly one fresh block. + pattern = re.compile( + rf"{re.escape(_BNB_ROCM_SITECUSTOMIZE_BEGIN)}.*?" + rf"(?:{re.escape(_BNB_ROCM_SITECUSTOMIZE_END)}\n?|\Z)", + re.DOTALL, + ) + remainder = pattern.sub("", existing) + separator = "" if not remainder or remainder.endswith("\n") else "\n" + updated = f"{remainder}{separator}{block}" + tmp_path = sitecustomize_path.with_name( + f"{sitecustomize_path.name}.unsloth-tmp{os.getpid()}" + ) + try: + tmp_path.write_text(updated, encoding = "utf-8") + if sitecustomize_path.exists(): + shutil.copymode(sitecustomize_path, tmp_path) + os.replace(tmp_path, sitecustomize_path) + finally: + tmp_path.unlink(missing_ok = True) + except (OSError, UnicodeDecodeError) as exc: + print( + f" Warning: could not persist BNB_ROCM_VERSION={version} " + f"to {sitecustomize_path}: {exc}" + ) + return False + + return True + + def _has_rocm_gpu() -> bool: """Return True only if an actual AMD GPU is visible (not just ROCm tools installed).""" for cmd, check_fn in ( @@ -636,15 +704,27 @@ def _install_bnb_windows_rocm() -> bool: ) if not _ok: return False - # After install: detect the actual ROCm DLL suffix in the wheel and set - # BNB_ROCM_VERSION so bitsandbytes loads the correct DLL regardless of what - # torch.version.hip reports. The wheel may ship an older suffix (e.g. "72") - # while torch reports a newer HIP version (e.g. 7.13); the override stops - # bitsandbytes from failing on a non-existent DLL. The worker subprocess - # inherits this env var. Fall back to "72" if detection fails (no-op/dry-run). - if "BNB_ROCM_VERSION" not in os.environ: + # After install: detect the actual ROCm DLL suffix shipped in the wheel and + # set BNB_ROCM_VERSION so bitsandbytes loads the correct DLL regardless of + # what torch.version.hip reports. The wheel may ship an older suffix (e.g. + # "72") while torch reports a newer HIP version (e.g. 7.13); the env var + # override ensures bitsandbytes does not fail looking for a non-existent DLL. + # The worker subprocess inherits this env var automatically. + # Fall back to "72" if detection fails (e.g. install was a no-op / dry-run). + _env_ver = os.environ.get("BNB_ROCM_VERSION") + _env_is_persisted_default = ( + os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE + ) + _persist_detected_version = False + if _env_ver and not _env_is_persisted_default: + _ver = _env_ver + else: _ver = _detect_bnb_rocm_dll_ver() or "72" os.environ["BNB_ROCM_VERSION"] = _ver + os.environ[_BNB_ROCM_VERSION_SOURCE_ENV] = _BNB_ROCM_VERSION_SOURCE_DETECTED + _persist_detected_version = True + if _persist_detected_version: + _persist_bnb_rocm_version(_ver) # Make hipInfo.exe (shipped into the venv Scripts dir by the AMD torch # wheel) resolvable via PATH for this process and every child python the # installer spawns (import checks, precompile): bitsandbytes runs diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index c761cd090f..02f52e7700 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1709,6 +1709,17 @@ class TestGfxArchNameFallback: class TestInstallBnbWindowsRocm: """Verify AMD Windows BNB wheel install helper.""" + @pytest.fixture(autouse = True) + def _isolate_sitecustomize_persistence(self, monkeypatch, request): + """Keep helper tests from writing to the active interpreter site-packages.""" + if request.node.name.startswith("test_persist"): + return + monkeypatch.setattr( + stack_mod, + "_persist_bnb_rocm_version", + lambda version: True, + ) + def test_calls_pip_install_try_with_win_amd64_url(self): """Should call pip_install_try with the win_amd64 wheel URL via plain pip.""" with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip: @@ -1762,6 +1773,7 @@ class TestInstallBnbWindowsRocm: """BNB_ROCM_VERSION is set from the DLL detected after install.""" with patch.dict(os.environ, {}, clear = False): os.environ.pop("BNB_ROCM_VERSION", None) + os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"): stack_mod._install_bnb_windows_rocm() @@ -1771,6 +1783,7 @@ class TestInstallBnbWindowsRocm: """If AMD ships a newer DLL (e.g. rocm713.dll), that version is used.""" with patch.dict(os.environ, {}, clear = False): os.environ.pop("BNB_ROCM_VERSION", None) + os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"): stack_mod._install_bnb_windows_rocm() @@ -1780,6 +1793,7 @@ class TestInstallBnbWindowsRocm: """Falls back to '72' when DLL detection returns None.""" with patch.dict(os.environ, {}, clear = False): os.environ.pop("BNB_ROCM_VERSION", None) + os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None): stack_mod._install_bnb_windows_rocm() @@ -1788,10 +1802,207 @@ class TestInstallBnbWindowsRocm: def test_does_not_override_existing_bnb_rocm_version(self): """An explicit BNB_ROCM_VERSION in the caller's env must not be clobbered.""" with patch.dict(os.environ, {"BNB_ROCM_VERSION": "60"}): + os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): stack_mod._install_bnb_windows_rocm() assert os.environ.get("BNB_ROCM_VERSION") == "60" + def test_does_not_persist_existing_bnb_rocm_version(self): + """A caller override must not become the venv's managed default.""" + with patch.dict(os.environ, {"BNB_ROCM_VERSION": "60"}): + os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver") as mock_detect: + with patch.object( + stack_mod, "_persist_bnb_rocm_version", return_value = True + ) as mock_persist: + stack_mod._install_bnb_windows_rocm() + + assert os.environ.get("BNB_ROCM_VERSION") == "60" + mock_detect.assert_not_called() + mock_persist.assert_not_called() + + def test_redetects_when_bnb_rocm_version_came_from_sitecustomize(self): + """Persisted defaults should not mask a newer DLL suffix after reinstall.""" + with patch.dict( + os.environ, + { + "BNB_ROCM_VERSION": "72", + stack_mod._BNB_ROCM_VERSION_SOURCE_ENV: ( + stack_mod._BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE + ), + }, + ): + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"): + with patch.object( + stack_mod, "_persist_bnb_rocm_version", return_value = True + ) as mock_persist: + stack_mod._install_bnb_windows_rocm() + + assert os.environ.get("BNB_ROCM_VERSION") == "713" + assert ( + os.environ.get(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV) + == stack_mod._BNB_ROCM_VERSION_SOURCE_DETECTED + ) + mock_persist.assert_called_once_with("713") + + def test_persists_bnb_rocm_version_for_direct_venv_python(self, tmp_path): + """BNB_ROCM_VERSION must apply to a fresh Python process in the venv.""" + site_packages = tmp_path / "site-packages" + + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"): + with patch.object( + stack_mod.sysconfig, "get_path", return_value = str(site_packages) + ): + stack_mod._install_bnb_windows_rocm() + + sitecustomize = site_packages / "sitecustomize.py" + source = sitecustomize.read_text(encoding = "utf-8") + assert "BNB_ROCM_VERSION" in source + assert stack_mod._BNB_ROCM_VERSION_SOURCE_ENV in source + assert "'72'" in source + + probe_env = os.environ.copy() + probe_env.pop("BNB_ROCM_VERSION", None) + probe_env.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) + probe_env["PYTHONPATH"] = str(site_packages) + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import os; " + "print(os.environ.get('BNB_ROCM_VERSION', ''), " + "os.environ.get('UNSLOTH_BNB_ROCM_VERSION_SOURCE', ''))" + ), + ], + env = probe_env, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + text = True, + check = True, + ) + assert result.stdout.strip() == "72 sitecustomize" + + def test_persist_bnb_rocm_version_replaces_existing_managed_block(self, tmp_path): + """Updating sitecustomize.py must not duplicate the managed BNB block.""" + site_packages = tmp_path / "site-packages" + site_packages.mkdir() + sitecustomize = site_packages / "sitecustomize.py" + sitecustomize.write_text( + "EXISTING = True\n" + "# BEGIN Unsloth BNB_ROCM_VERSION\n" + "import os as _unsloth_os\n" + "_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', '72')\n" + "# END Unsloth BNB_ROCM_VERSION\n", + encoding = "utf-8", + ) + + with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + assert stack_mod._persist_bnb_rocm_version("713") is True + + source = sitecustomize.read_text(encoding = "utf-8") + assert source.count("# BEGIN Unsloth BNB_ROCM_VERSION") == 1 + assert "EXISTING = True" in source + assert "'713'" in source + assert "'72'" not in source + + def test_persist_bnb_rocm_version_handles_non_utf8_sitecustomize(self, tmp_path): + """A legacy non-UTF-8 sitecustomize.py should not abort installation.""" + site_packages = tmp_path / "site-packages" + site_packages.mkdir() + sitecustomize = site_packages / "sitecustomize.py" + sitecustomize.write_bytes(b"\xff\xfe\x00") + + with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + assert stack_mod._persist_bnb_rocm_version("72") is False + + def test_persist_bnb_rocm_version_repairs_truncated_block(self, tmp_path): + """A managed block missing its END marker is replaced, not duplicated.""" + site_packages = tmp_path / "site-packages" + site_packages.mkdir() + sitecustomize = site_packages / "sitecustomize.py" + sitecustomize.write_text( + "EXISTING = True\n" + "# BEGIN Unsloth BNB_ROCM_VERSION\n" + "import os as _unsloth_os\n" + "_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', '72')\n", + encoding = "utf-8", + ) + + with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + assert stack_mod._persist_bnb_rocm_version("713") is True + + source = sitecustomize.read_text(encoding = "utf-8") + assert source.count("# BEGIN Unsloth BNB_ROCM_VERSION") == 1 + assert source.count("# END Unsloth BNB_ROCM_VERSION") == 1 + assert "EXISTING = True" in source + assert "'713'" in source + assert "'72'" not in source + + def test_persist_bnb_rocm_version_dedupes_duplicate_blocks(self, tmp_path): + """Multiple managed blocks collapse to one while preserving user content.""" + site_packages = tmp_path / "site-packages" + site_packages.mkdir() + sitecustomize = site_packages / "sitecustomize.py" + block = ( + "# BEGIN Unsloth BNB_ROCM_VERSION\n" + "import os as _unsloth_os\n" + "_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', '72')\n" + "# END Unsloth BNB_ROCM_VERSION\n" + ) + sitecustomize.write_text(block + "USER_MID = 1\n" + block, encoding = "utf-8") + + with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + assert stack_mod._persist_bnb_rocm_version("713") is True + + source = sitecustomize.read_text(encoding = "utf-8") + assert source.count("# BEGIN Unsloth BNB_ROCM_VERSION") == 1 + assert source.count("# END Unsloth BNB_ROCM_VERSION") == 1 + assert "USER_MID = 1" in source + assert "'713'" in source + assert "'72'" not in source + + def test_persist_bnb_rocm_version_atomic_no_leftover_tmp(self, tmp_path): + """The write-then-rename path must not leave its temp file behind.""" + site_packages = tmp_path / "site-packages" + site_packages.mkdir() + + with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + assert stack_mod._persist_bnb_rocm_version("72") is True + + leftovers = [p.name for p in site_packages.iterdir() if "unsloth-tmp" in p.name] + assert leftovers == [] + assert (site_packages / "sitecustomize.py").exists() + + +class TestRuntimeBnbRocmSourceGuards: + """Runtime entrypoints redetect managed defaults but keep caller overrides.""" + + _MAIN_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py" + _TRAINING_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" + + def test_main_gate_redetects_persisted_default(self): + source = self._MAIN_PATH.read_text(encoding = "utf-8") + assert 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' in source + assert 'os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"' in source + + def test_worker_gate_redetects_persisted_default(self): + source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8") + assert 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' in source + assert 'os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"' in source + + def test_fallback_prefers_seeded_value_over_hardcoded_72(self): + """A failed redetect must not downgrade a persisted suffix to '72'.""" + for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH): + source = path.read_text(encoding = "utf-8") + assert 'os.environ.get("BNB_ROCM_VERSION") or "72"' in source, path.name + class TestDetectBnbRocmDllVer: """Unit tests for _detect_bnb_rocm_dll_ver()."""