diff --git a/studio/backend/main.py b/studio/backend/main.py index 942cf8b2fc..609702ca2f 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -47,6 +47,32 @@ if sys.platform == "win32": _add_rocm_dll_dirs() del _add_rocm_dll_dirs + # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─ + # bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm.dll + # where comes from torch.version.hip (e.g. "7.13..." → "713"). + # The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without + # 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). + if "BNB_ROCM_VERSION" not in os.environ: + import glob as _glob + _bnb_rocm_ver = None + try: + import importlib.util as _ilu + _bnb_spec = _ilu.find_spec("bitsandbytes") + if _bnb_spec and _bnb_spec.origin: + _pkg_dir = os.path.dirname(_bnb_spec.origin) + _dlls = _glob.glob(os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")) + import re as _re_bnb + for _dll in sorted(_dlls): + _m = _re_bnb.search(r"libbitsandbytes_rocm(\d+)\.dll", _dll) + if _m: + _bnb_rocm_ver = _m.group(1) + break + except Exception: + pass + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver or "72" + # 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) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index d0fe3822b5..0968248679 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -946,10 +946,23 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non def _determine_attention_impl_for_gpu_estimate(config) -> str: import copy as _copy - # torch.distributed is incomplete on Windows ROCm — it ships without the - # process-group helpers (is_initialized, is_available, etc.). - # resolve_attention_implementation (unsloth) calls is_initialized() - # unconditionally, so patch any missing attrs before importing it. + # torch.distributed is incomplete on Windows ROCm — torch._C is a C + # extension (not a package), so Python cannot import the submodule + # torch._C._distributed_c10d that torch.distributed depends on. + # Inject an empty stub into sys.modules BEFORE importing torch.distributed + # so the import succeeds, then patch the missing process-group helpers. + import sys as _sys + import types as _types + + if _sys.platform == "win32": + for _c10d_name in ( + "torch._C._distributed_c10d", + "torch._C._distributed_autograd", + "torch._C._distributed_rpc", + ): + if _c10d_name not in _sys.modules: + _sys.modules[_c10d_name] = _types.ModuleType(_c10d_name) + try: import torch.distributed as _td diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 822514eeb1..b615d6d0d0 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2348,5 +2348,73 @@ class TestSetupShGccInstallDir: assert "gcc install dir" in source or "GCC_INSTALL_DIR" in source +# ============================================================================= +# TEST: main.py -- BNB_ROCM_VERSION server startup + distributed stubs +# ============================================================================= + +_MAIN_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py" +_HARDWARE_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + + +class TestServerStartupRocmFixes: + """Verify main.py sets BNB_ROCM_VERSION before any bitsandbytes import and + hardware.py injects torch._C._distributed_c10d stubs before torch.distributed.""" + + # ── BNB_ROCM_VERSION in server process ──────────────────────────────────── + + def test_main_py_sets_bnb_rocm_version(self): + """main.py must set BNB_ROCM_VERSION in the server process before imports.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert "BNB_ROCM_VERSION" in source + + def test_main_py_bnb_detection_scoped_to_win32(self): + """main.py BNB_ROCM_VERSION logic must be inside the win32 platform guard.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + win32_idx = source.find('sys.platform == "win32"') + bnb_idx = source.find("BNB_ROCM_VERSION") + assert win32_idx != -1 and bnb_idx != -1 + assert win32_idx < bnb_idx + + def test_main_py_bnb_dll_detection_uses_glob(self): + """main.py must scan for libbitsandbytes_rocm*.dll to find the right version.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert "libbitsandbytes_rocm" in source + + def test_main_py_bnb_falls_back_to_72(self): + """main.py must fall back to BNB_ROCM_VERSION='72' when no DLL is found.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert '"72"' in source or "'72'" in source + + def test_main_py_bnb_only_set_when_not_already_in_env(self): + """main.py must not override an existing BNB_ROCM_VERSION env var.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert '"BNB_ROCM_VERSION" not in os.environ' in source + + # ── torch._C._distributed_c10d stubs in hardware.py ────────────────────── + + def test_hardware_py_injects_distributed_c10d_stub(self): + """hardware.py must inject torch._C._distributed_c10d into sys.modules.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "_distributed_c10d" in source + + def test_hardware_py_stub_injected_before_distributed_import(self): + """The sys.modules stub must be injected BEFORE import torch.distributed.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + c10d_idx = source.find("_distributed_c10d") + dist_idx = source.find("import torch.distributed") + assert c10d_idx != -1 and dist_idx != -1 + assert c10d_idx < dist_idx + + def test_hardware_py_stub_uses_types_moduletype(self): + """hardware.py must create the stub with types.ModuleType.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "ModuleType" in source + + def test_hardware_py_stub_scoped_to_win32(self): + """hardware.py distributed stub injection must be gated on win32.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert 'platform == "win32"' in source or "win32" in source + + if __name__ == "__main__": pytest.main([__file__, "-v"])