"""Tests for AMD ROCm support across install pathways. Verifies that ROCm detection and installation logic works correctly WITHOUT breaking existing CUDA, CPU, macOS, and Windows pathways. All tests use mocks -- no AMD hardware required. """ import importlib.util import json import os import subprocess import sys from pathlib import Path from unittest.mock import MagicMock, mock_open, patch, PropertyMock import pytest # ── Load modules under test ────────────────────────────────────────────────── PACKAGE_ROOT = Path(__file__).resolve().parents[3] # install_llama_prebuilt.py _PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" _PREBUILT_SPEC = importlib.util.spec_from_file_location( "studio_install_llama_prebuilt", _PREBUILT_PATH ) assert _PREBUILT_SPEC is not None and _PREBUILT_SPEC.loader is not None prebuilt_mod = importlib.util.module_from_spec(_PREBUILT_SPEC) sys.modules[_PREBUILT_SPEC.name] = prebuilt_mod _PREBUILT_SPEC.loader.exec_module(prebuilt_mod) HostInfo = prebuilt_mod.HostInfo AssetChoice = prebuilt_mod.AssetChoice PrebuiltFallback = prebuilt_mod.PrebuiltFallback resolve_upstream_asset_choice = prebuilt_mod.resolve_upstream_asset_choice runtime_patterns_for_choice = prebuilt_mod.runtime_patterns_for_choice _apply_host_overrides = prebuilt_mod._apply_host_overrides _normalize_forwarded_gfx = prebuilt_mod._normalize_forwarded_gfx # install_python_stack.py _STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" _STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH) assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None stack_mod = importlib.util.module_from_spec(_STACK_SPEC) sys.modules[_STACK_SPEC.name] = stack_mod _STACK_SPEC.loader.exec_module(stack_mod) _detect_rocm_version = stack_mod._detect_rocm_version _ensure_rocm_torch = stack_mod._ensure_rocm_torch _has_rocm_gpu = stack_mod._has_rocm_gpu _has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu _ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX _windows_rocm_index_url = stack_mod._windows_rocm_index_url _detect_windows_gfx_arch = stack_mod._detect_windows_gfx_arch _install_bnb_windows_rocm = stack_mod._install_bnb_windows_rocm def _extract_sh_function_body(source: str, name: str) -> str: """Return the body of a shell function from `source` by brace matching. Used by structural tests that need to assert ordering of helper calls inside a specific function rather than across the whole install.sh file. """ needle = f"{name}() {{" start = source.find(needle) if start < 0: return "" depth = 0 i = start + len(needle) - 1 # land on the opening brace n = len(source) while i < n: ch = source[i] if ch == "{": depth += 1 elif ch == "}": depth -= 1 if depth == 0: return source[start : i + 1] i += 1 return source[start:] # ── Helper: build HostInfo for different scenarios ────────────────────────── def nvidia_host(**overrides) -> HostInfo: """NVIDIA Linux x86_64 host.""" defaults = dict( system = "Linux", machine = "x86_64", is_windows = False, is_linux = True, is_macos = False, is_x86_64 = True, is_arm64 = False, nvidia_smi = "/usr/bin/nvidia-smi", driver_cuda_version = (12, 6), compute_caps = ["89"], visible_cuda_devices = None, has_physical_nvidia = True, has_usable_nvidia = True, has_rocm = False, ) defaults.update(overrides) return HostInfo(**defaults) def rocm_host(**overrides) -> HostInfo: """AMD ROCm Linux x86_64 host (no NVIDIA).""" defaults = dict( system = "Linux", machine = "x86_64", is_windows = False, is_linux = True, is_macos = False, is_x86_64 = True, is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], visible_cuda_devices = None, has_physical_nvidia = False, has_usable_nvidia = False, has_rocm = True, ) defaults.update(overrides) return HostInfo(**defaults) def cpu_host(**overrides) -> HostInfo: """CPU-only Linux x86_64 host.""" defaults = dict( system = "Linux", machine = "x86_64", is_windows = False, is_linux = True, is_macos = False, is_x86_64 = True, is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], visible_cuda_devices = None, has_physical_nvidia = False, has_usable_nvidia = False, has_rocm = False, ) defaults.update(overrides) return HostInfo(**defaults) def macos_host(**overrides) -> HostInfo: """macOS arm64 host.""" defaults = dict( system = "Darwin", machine = "arm64", is_windows = False, is_linux = False, is_macos = True, is_x86_64 = False, is_arm64 = True, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], visible_cuda_devices = None, has_physical_nvidia = False, has_usable_nvidia = False, has_rocm = False, ) defaults.update(overrides) return HostInfo(**defaults) def windows_host(**overrides) -> HostInfo: """Windows x86_64 host.""" defaults = dict( system = "Windows", machine = "amd64", is_windows = True, is_linux = False, is_macos = False, is_x86_64 = True, is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], visible_cuda_devices = None, has_physical_nvidia = False, has_usable_nvidia = False, has_rocm = False, ) defaults.update(overrides) return HostInfo(**defaults) def windows_rocm_host(**overrides) -> HostInfo: """Windows x86_64 host with ROCm.""" defaults = dict( system = "Windows", machine = "amd64", is_windows = True, is_linux = False, is_macos = False, is_x86_64 = True, is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], visible_cuda_devices = None, has_physical_nvidia = False, has_usable_nvidia = False, has_rocm = True, ) defaults.update(overrides) return HostInfo(**defaults) # ── Upstream asset fixture ─────────────────────────────────────────────────── LLAMA_TAG = "b8508" UPSTREAM_ASSETS = { f"llama-{LLAMA_TAG}-bin-ubuntu-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-linux-cpu.tar.gz", f"llama-{LLAMA_TAG}-bin-ubuntu-rocm-7.2-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-linux-rocm.tar.gz", f"llama-{LLAMA_TAG}-bin-win-cpu-x64.zip": f"https://example.com/{LLAMA_TAG}-win-cpu.zip", f"llama-{LLAMA_TAG}-bin-win-cuda-12.4-x64.zip": f"https://example.com/{LLAMA_TAG}-win-cuda.zip", f"llama-{LLAMA_TAG}-bin-win-hip-radeon-x64.zip": f"https://example.com/{LLAMA_TAG}-win-hip.zip", f"llama-{LLAMA_TAG}-bin-macos-arm64.tar.gz": f"https://example.com/{LLAMA_TAG}-macos-arm64.tar.gz", f"llama-{LLAMA_TAG}-bin-macos-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-macos-x64.tar.gz", } # TEST: install_llama_prebuilt.py -- resolve_upstream_asset_choice class TestResolveUpstreamAssetChoice: """Verify that the asset selection logic picks the right binary for each platform.""" @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_nvidia_linux_gets_cpu_asset(self, mock_assets): """NVIDIA host should NOT hit the ROCm path -- gets CPU asset (CUDA handled elsewhere).""" host = nvidia_host() choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "linux-cpu" assert "ubuntu-x64" in choice.name assert "rocm" not in choice.name @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_rocm_linux_gets_rocm_prebuilt(self, mock_assets): """AMD ROCm Linux host should get the ROCm prebuilt.""" host = rocm_host() choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "linux-rocm" assert "rocm" in choice.name @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_cpu_linux_gets_cpu_asset(self, mock_assets): """CPU-only Linux host should get CPU asset.""" host = cpu_host() choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "linux-cpu" assert "ubuntu-x64" in choice.name @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_macos_arm64_gets_macos_asset(self, mock_assets): """macOS arm64 host should get macOS asset.""" host = macos_host() choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "macos-arm64" assert "macos-arm64" in choice.name @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_windows_cpu_gets_cpu_asset(self, mock_assets): """Windows CPU-only host should get Windows CPU asset.""" host = windows_host() choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "windows-cpu" assert "win-cpu" in choice.name @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_windows_rocm_gets_hip_asset(self, mock_assets): """Windows ROCm host should get Windows HIP asset.""" host = windows_rocm_host() choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "windows-hip" assert "hip" in choice.name @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_mixed_nvidia_rocm_prefers_nvidia(self, mock_assets): """Host with both NVIDIA and ROCm should use NVIDIA (CPU path here, CUDA elsewhere).""" host = nvidia_host(has_rocm = True) choice = resolve_upstream_asset_choice(host, LLAMA_TAG) # NVIDIA hosts go through the normal path (CUDA handled by resolve_linux_cuda_choice) assert choice.install_kind == "linux-cpu" assert "rocm" not in choice.name @patch.object(prebuilt_mod, "github_release_assets") def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets): """AMD ROCm host should fall back to source build when no ROCm prebuilt exists.""" # Remove the ROCm asset from available assets assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k} mock_assets.return_value = assets_without_rocm host = rocm_host() with pytest.raises(PrebuiltFallback, match = "ROCm detected"): resolve_upstream_asset_choice(host, LLAMA_TAG) @patch.object(prebuilt_mod, "github_release_assets") def test_windows_rocm_no_hip_falls_to_cpu(self, mock_assets): """Windows+ROCm with HIP prebuilt missing should fall through to CPU.""" assets_no_hip = {k: v for k, v in UPSTREAM_ASSETS.items() if "hip" not in k} mock_assets.return_value = assets_no_hip host = windows_rocm_host() choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "windows-cpu" @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_macos_rocm_impossible_has_rocm_false(self, mock_assets): """macOS host should never have has_rocm=True in practice; verify it gets macOS asset.""" host = macos_host(has_rocm = True) choice = resolve_upstream_asset_choice(host, LLAMA_TAG) assert choice.install_kind == "macos-arm64" @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) def test_linux_aarch64_rocm_gets_prebuilt_fallback(self, mock_assets): """Linux aarch64 with ROCm -- no x86_64 match, should raise PrebuiltFallback.""" host = rocm_host(machine = "aarch64", is_x86_64 = False, is_arm64 = True) with pytest.raises(PrebuiltFallback): resolve_upstream_asset_choice(host, LLAMA_TAG) # TEST: install_llama_prebuilt.py -- runtime_patterns_for_choice class TestRuntimePatterns: """Verify runtime file patterns for all install kinds.""" def test_linux_cpu_patterns(self): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cpu" ) patterns = runtime_patterns_for_choice(choice) assert "llama-server" in patterns assert "llama-quantize" in patterns # Broad lib*.so* covers libllama, libggml, libmtmd, libggml-cpu-*, # plus the libllama--impl.so split that ggml-org/llama.cpp # #23462 introduced between b9279 and b9283. assert "lib*.so*" in patterns def test_linux_cuda_patterns(self): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cuda" ) patterns = runtime_patterns_for_choice(choice) # libggml-cuda.so is matched by lib*.so* now. assert "lib*.so*" in patterns def test_linux_rocm_patterns(self): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-rocm" ) patterns = runtime_patterns_for_choice(choice) # libggml-hip.so is matched by lib*.so* now. assert "lib*.so*" in patterns assert "llama-server" in patterns def test_windows_hip_patterns(self): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "windows-hip", ) patterns = runtime_patterns_for_choice(choice) # Narrowed from "*.exe" to the two binaries Studio actually # invokes, mirroring the Linux/macOS pattern style. assert "llama-server.exe" in patterns assert "llama-quantize.exe" in patterns assert "*.dll" in patterns def test_macos_patterns(self): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "macos-arm64", ) patterns = runtime_patterns_for_choice(choice) assert "lib*.dylib" in patterns def test_diffusion_visual_server_kept(self): # The DiffusionGemma visual-server ships in the prebuilt bundle and must # survive the prune so Studio can serve DiffusionGemma GGUFs natively. for kind, name in ( ("linux-cuda", "llama-diffusion-gemma-visual-server"), ("macos-arm64", "llama-diffusion-gemma-visual-server"), ("windows-cuda", "llama-diffusion-gemma-visual-server.exe"), ): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = kind ) assert name in runtime_patterns_for_choice(choice) # TEST: install_llama_prebuilt.py -- HostInfo.has_rocm field class TestHostInfoRocm: """Verify has_rocm field does not affect other HostInfo behavior.""" def test_has_rocm_default_false(self): host = HostInfo( system = "Linux", machine = "x86_64", is_windows = False, is_linux = True, is_macos = False, is_x86_64 = True, is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], visible_cuda_devices = None, has_physical_nvidia = False, has_usable_nvidia = False, ) assert host.has_rocm is False def test_has_rocm_explicit_true(self): host = rocm_host() assert host.has_rocm is True def test_nvidia_host_no_rocm(self): host = nvidia_host() assert host.has_rocm is False assert host.has_usable_nvidia is True def test_detect_host_has_rocm_detection_logic(self): """detect_host() should have ROCm GPU detection logic.""" import inspect source = inspect.getsource(prebuilt_mod.detect_host) # Must probe for actual GPU, not just tool presence assert "rocminfo" in source or "amd-smi" in source def test_detect_host_windows_rocm_detection(self): """detect_host() source should have Windows-specific ROCm GPU detection.""" import inspect source = inspect.getsource(prebuilt_mod.detect_host) assert "hipinfo" in source or "amd-smi" in source # TEST: install_python_stack.py -- _detect_rocm_version class TestDetectRocmVersion: """Verify ROCm version detection from various sources.""" def test_no_rocm_returns_none(self, tmp_path): """No ROCm installed should return None.""" with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): with patch("shutil.which", return_value = None): result = _detect_rocm_version() assert result is None def test_version_from_file(self, tmp_path): """Reads version from /opt/rocm/.info/version.""" info_dir = tmp_path / ".info" info_dir.mkdir() (info_dir / "version").write_text("7.1.0-12345\n") with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): result = _detect_rocm_version() assert result == (7, 1) def test_version_62(self, tmp_path): """Reads ROCm 6.2 version.""" info_dir = tmp_path / ".info" info_dir.mkdir() (info_dir / "version").write_text("6.2.0\n") with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): result = _detect_rocm_version() assert result == (6, 2) def test_hipconfig_fallback(self, tmp_path): """Falls back to hipconfig --version when file not found.""" with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = b"6.3.21234.2\n" with patch("shutil.which", return_value = "/usr/bin/hipconfig"): with patch("subprocess.run", return_value = mock_result): result = _detect_rocm_version() assert result == (6, 3) def test_dpkg_fallback_without_hipconfig(self, tmp_path): """dpkg rocm-core fallback works when amd-smi and hipconfig are absent (regression: a shadowing local re import raised UnboundLocalError).""" def which(cmd): return "/usr/bin/dpkg-query" if cmd == "dpkg-query" else None mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = "1:6.3.0-1\n" with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): with patch("shutil.which", side_effect = which): with patch("subprocess.run", return_value = mock_result): assert _detect_rocm_version() == (6, 3) def test_empty_version_file(self, tmp_path): """Empty version file should return None.""" info_dir = tmp_path / ".info" info_dir.mkdir() (info_dir / "version").write_text("") with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): with patch("shutil.which", return_value = None): result = _detect_rocm_version() assert result is None def test_version_with_epoch_prefix(self, tmp_path): """Debian epoch prefix (2:6.2.0) -- version file has no epoch, so should parse.""" info_dir = tmp_path / ".info" info_dir.mkdir() # Version files don't typically have epoch prefix, but lib/rocm_version might (info_dir / "version").write_text("6.2.0\n") with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): result = _detect_rocm_version() assert result == (6, 2) def test_multiple_version_sources_first_wins(self, tmp_path): """When both .info/version and lib/rocm_version exist, first found wins.""" info_dir = tmp_path / ".info" info_dir.mkdir() (info_dir / "version").write_text("7.1.0\n") lib_dir = tmp_path / "lib" lib_dir.mkdir() (lib_dir / "rocm_version").write_text("6.3.0\n") with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): result = _detect_rocm_version() assert result == (7, 1) # .info/version checked first def test_hipconfig_multiline_output(self, tmp_path): """hipconfig with multi-line output -- should use first line.""" with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = b"6.3.21234.2\nSome extra info\n" with patch("shutil.which", return_value = "/usr/bin/hipconfig"): with patch("subprocess.run", return_value = mock_result): result = _detect_rocm_version() assert result == (6, 3) def test_hipconfig_timeout(self, tmp_path): """hipconfig that times out should return None.""" with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): with patch("shutil.which", return_value = "/usr/bin/hipconfig"): with patch( "subprocess.run", side_effect = subprocess.TimeoutExpired("hipconfig", 5), ): result = _detect_rocm_version() assert result is None # TEST: install_python_stack.py -- _ensure_rocm_torch class TestEnsureRocmTorch: """Verify ROCm torch reinstall logic.""" @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) def test_no_rocm_skips(self, mock_nvidia, mock_pip): """No ROCm toolchain should skip entirely.""" # _detect_windows_gfx_arch pinned to None: on a real AMD test host its # WMI name fallback would otherwise answer and defeat the "no ROCm # anywhere" premise of this test. with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None): with patch("os.path.isdir", return_value = False): with patch("shutil.which", return_value = None): _ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """If torch already has CUDA, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"12.6\n" # CUDA version string with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """If torch already has HIP, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"7.1.12345\n" # HIP version string with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) def test_cpu_torch_gets_rocm_reinstall( self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try ): """CPU-only torch on ROCm host should trigger reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"\n" # empty = no GPU backend with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() # Should install torch via pip_install and bitsandbytes via pip_install_try. assert mock_pip.call_count == 1 assert "rocm7.1" in str(mock_pip.call_args_list[0]) assert mock_pip_try.call_count >= 1 assert "bitsandbytes" in str(mock_pip_try.call_args_list[0]) assert mock_pip_try.call_args.kwargs["force_pip"] is True @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 3)) def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """ROCm 6.3 should select rocm6.3 tag.""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"\n" with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() torch_call = mock_pip.call_args_list[0] assert "rocm6.3" in str(torch_call) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (5, 0)) def test_old_rocm_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """ROCm version too old (below 6.0) should skip.""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"\n" with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = None) def test_version_unreadable_prints_warning( self, mock_ver, mock_gpu, mock_nvidia, mock_pip, capsys ): """ROCm detected but version unreadable should print warning and skip.""" with patch("os.path.isdir", return_value = True): _ensure_rocm_torch() mock_pip.assert_not_called() captured = capsys.readouterr() assert "unreadable" in captured.out @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) def test_rocm_72_selects_72_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """ROCm 7.2 should select rocm7.2 tag (now in mapping with torch 2.11.0).""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"\n" with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() torch_call = mock_pip.call_args_list[0] assert "rocm7.2" in str(torch_call) @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) def test_probe_timeout_triggers_reinstall( self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try ): """Probe subprocess timeout should not crash; should proceed to reinstall.""" with patch("os.path.isdir", return_value = True): with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)): _ensure_rocm_torch() # If probe times out, the function should treat torch as unusable and reinstall # both torch (via pip_install) and bitsandbytes (via pip_install_try). assert mock_pip.call_count == 1 assert "rocm7.1" in str(mock_pip.call_args_list[0]) assert mock_pip_try.call_count >= 1 assert mock_pip_try.call_args.kwargs["force_pip"] is True @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = False) def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip): """ROCm tools present but no actual AMD GPU should skip entirely.""" # Pin the Windows arch probe to None: on a real AMD test host the WMI # name fallback would otherwise answer and defeat the "no actual GPU" # premise (the Linux path under test uses _has_rocm_gpu, mocked False). with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None): with patch("os.path.isdir", return_value = True): _ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip): """UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe.""" with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}): # Reload _TORCH_BACKEND from the patched environment. with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"): _ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip): """UNSLOTH_TORCH_BACKEND=cpu must short-circuit before any GPU probe.""" with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cpu"}): with patch.object(stack_mod, "_TORCH_BACKEND", "cpu"): _ensure_rocm_torch() mock_pip.assert_not_called() # TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard class TestHasRocmGpuKfdVendorGuard: """Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes. These tests are source-level: they verify the regex and logic present in the _has_rocm_gpu implementation rather than running the sysfs traversal (which requires Linux path conventions). """ def _src(self) -> str: """Return the source of _has_rocm_gpu from install_python_stack.py.""" import inspect return inspect.getsource(stack_mod._has_rocm_gpu) def test_vendor_id_check_present(self): """_has_rocm_gpu sysfs fallback must check vendor_id 4098 (AMD 0x1002).""" src = self._src() assert "vendor_id" in src, ( "_has_rocm_gpu KFD sysfs fallback must read the properties file " "to check vendor_id and exclude NVIDIA KFD nodes" ) assert "4098" in src, ( "_has_rocm_gpu must require AMD vendor_id 4098 (0x1002) in the " "KFD node properties to avoid false positives on NVIDIA systems" ) def test_vendor_regex_pattern_anchored(self): """The vendor_id regex must use a word boundary to avoid partial matches.""" import re as _re src = self._src() # The pattern should have a word boundary before and after the number # so "vendor_id 41098" doesn't match "vendor_id 4098". assert ( _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src ), "_has_rocm_gpu vendor_id check should use word boundary anchors" def test_sysfs_fallback_guarded_by_non_win32(self): """KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32').""" src = self._src() assert "win32" in src, "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check" def test_cpu_node_excluded(self): """gpu_id == '0' must be excluded (CPU topology nodes).""" src = self._src() assert ( '!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src ), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)" def test_install_sh_has_vendor_check(self): """_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_amd_rocm_gpu()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] assert "vendor_id" in func_body, "_has_amd_rocm_gpu sysfs fallback must check vendor_id" assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)" def test_has_rocm_gpu_returns_false_when_nvidia_present(self): """_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True. This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs produce a false positive, an NVIDIA GPU always wins. """ with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True): with patch("shutil.which", return_value = "/usr/bin/rocminfo"): # Simulate rocminfo claiming an AMD GPU is present mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = "Name: gfx1100\n" with patch("subprocess.run", return_value = mock_result): assert not stack_mod._has_rocm_gpu(), ( "_has_rocm_gpu must return False when NVIDIA GPU is detected, " "regardless of what rocminfo reports" ) def test_install_sh_has_rocm_gpu_nvidia_guard(self): """_has_amd_rocm_gpu in install.sh must call _has_usable_nvidia_gpu and return 1 if true.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_amd_rocm_gpu()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] assert ( "_has_usable_nvidia_gpu" in func_body ), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts" assert ( "return 1" in func_body ), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected" def test_has_usable_nvidia_gpu_proc_fallback_present(self): """`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback.""" import inspect src = inspect.getsource(stack_mod._has_usable_nvidia_gpu) assert "/proc/driver/nvidia" in src, ( "_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus when " "nvidia-smi subprocess fails, to handle PATH gaps and driver init races" ) def test_install_sh_has_usable_nvidia_gpu_proc_fallback(self): """_has_usable_nvidia_gpu in install.sh must also have a /proc/driver/nvidia fallback.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_usable_nvidia_gpu()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] assert "/proc/driver/nvidia" in func_body, ( "_has_usable_nvidia_gpu in install.sh must fall back to " "/proc/driver/nvidia/gpus when nvidia-smi fails" ) # TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping class TestRocmTorchIndex: """Verify the ROCm version -> torch index tag mapping.""" def test_mapping_is_sorted_descending(self): """Keys should be in descending order for the next() iteration to work.""" keys = list(_ROCM_TORCH_INDEX.keys()) assert keys == sorted(keys, reverse = True) def test_rocm_72_in_mapping(self): """ROCm 7.2 should be in the active mapping (torch 2.11.0 now supported).""" assert (7, 2) in _ROCM_TORCH_INDEX assert _ROCM_TORCH_INDEX[(7, 2)] == "rocm7.2" def test_rocm_71_maps_correctly(self): assert _ROCM_TORCH_INDEX[(7, 1)] == "rocm7.1" def test_rocm_63_maps_correctly(self): assert _ROCM_TORCH_INDEX[(6, 3)] == "rocm6.3" def test_rocm_60_maps_correctly(self): assert _ROCM_TORCH_INDEX[(6, 0)] == "rocm6.0" def test_all_tags_use_download_pytorch(self): """All tags should be for download.pytorch.org, not repo.radeon.com.""" for tag in _ROCM_TORCH_INDEX.values(): assert tag.startswith("rocm") assert "radeon" not in tag def test_newer_rocm_selects_best_match(self): """ROCm 7.2 (now in map) should select rocm7.2 directly.""" ver = (7, 2) tag = next( ( t for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) if ver >= (maj, mn) ), None, ) assert tag == "rocm7.2" def test_rocm_64_selects_64(self): ver = (6, 4) tag = next( ( t for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) if ver >= (maj, mn) ), None, ) assert tag == "rocm6.4" # TEST: hardware.py -- IS_ROCM flag and detect_hardware class TestHardwareRocmFlag: """Verify IS_ROCM flag behavior without importing the full hardware module.""" def test_hardware_py_has_is_rocm(self): """hardware.py should define IS_ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") assert "IS_ROCM: bool" in source and "False" in source def test_hardware_py_sets_is_rocm_on_hip(self): """detect_hardware() should set IS_ROCM when torch.version.hip is set.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") assert 'torch.version, "hip"' in source or "torch.version.hip" in source def test_hardware_py_still_returns_cuda_for_rocm(self): """DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") # Ensure ROCM is NOT a DeviceType member enum_section = source.split("class DeviceType")[1].split("\n\n")[0] assert "ROCM" not in enum_section def test_hardware_py_has_rocm_in_package_versions(self): """get_package_versions() should include 'rocm' key.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") assert '"rocm"' in source def test_hardware_py_device_type_cuda_references_intact(self): """All existing DeviceType.CUDA references should still be present.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") # Key functions that must still reference DeviceType.CUDA assert "DeviceType.CUDA" in source assert "DEVICE = DeviceType.CUDA" in source def test_is_rocm_exported_from_init(self): """IS_ROCM should be exported from hardware __init__.py.""" init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" source = init_path.read_text(encoding = "utf-8") assert "IS_ROCM" in source def test_is_rocm_in_all_list(self): """IS_ROCM should be in __all__ list in __init__.py.""" init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" source = init_path.read_text(encoding = "utf-8") # Extract __all__ section assert '"IS_ROCM"' in source def test_get_package_versions_returns_rocm_key(self): """get_package_versions() source should return both 'cuda' and 'rocm' keys.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") # Find the get_package_versions function body func_start = source.find("def get_package_versions") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert '"cuda"' in func_body assert '"rocm"' in func_body def test_distributed_stubs_cover_is_torchelastic_launched(self): """_determine_attention_impl_for_gpu_estimate must stub is_torchelastic_launched. resolve_attention_implementation calls is_torchelastic_launched() on Windows ROCm where torch.distributed ships without that helper, causing a warning: 'module torch.distributed has no attribute is_torchelastic_launched'. """ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") assert "is_torchelastic_launched" in source def test_distributed_stubs_cover_core_helpers(self): """_determine_attention_impl_for_gpu_estimate must stub the four core distributed helpers.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") for attr in ("is_initialized", "is_available", "get_rank", "get_world_size"): assert attr in source, f"distributed stub for '{attr}' missing from hardware.py" # TEST: tokenizer_utils.py -- error message class TestTokenizerErrorMessage: """Verify the AMD error message is updated.""" def test_no_old_amd_message(self): """Old 'We do not support AMD' message should be gone.""" tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py" source = tu_path.read_text(encoding = "utf-8") assert "We do not support AMD" not in source def test_new_message_has_docs_link(self): """New message should point to Unsloth AMD docs.""" tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py" source = tu_path.read_text(encoding = "utf-8") assert "docs.unsloth.ai" in source or "No GPU detected" in source # TEST: install.sh -- structural checks class TestInstallShStructure: """Verify install.sh structural properties without running it.""" def test_no_here_strings(self): """install.sh must not use the bash-only `<<<` here-string operator. `<<<` inside a quoted literal (e.g. a marker in a printf) is just data, not a here-string, so strip quoted spans first: this still catches a real `cmd <<< word` (outside quotes) without false-positiving on data. """ import re sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") # <<< is bash-only; breaks dash for i, line in enumerate(source.splitlines(), 1): stripped = line.lstrip() if stripped.startswith("#"): continue # Remove quoted string literals so `<<<` inside them is ignored; # a genuine here-string operator lives outside any quotes. unquoted = re.sub(r"'[^']*'", "", line) unquoted = re.sub(r'"[^"]*"', "", unquoted) assert "<<<" not in unquoted, f"install.sh:{i} uses non-POSIX <<< here-string" def test_rocm_detection_present(self): """install.sh should have ROCm detection in get_torch_index_url.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert "amd-smi" in source assert "rocm" in source.lower() def test_cuda_precedence(self): """ROCm detection should only run when nvidia-smi is absent. install.sh defines _has_amd_rocm_gpu and _has_usable_nvidia_gpu helpers near each other (file-position order has no semantic meaning), so check the runtime ordering inside get_torch_index_url instead: NVIDIA branch runs first and the AMD/ROCm branch only fires inside the `if [ -z "$_smi" ]` block. """ sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") body = _extract_sh_function_body(source, "get_torch_index_url") nvidia_call = body.find("_has_usable_nvidia_gpu") # Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to # handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu # returns true via /proc/driver/nvidia/gpus. no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]') if no_nvidia_branch < 0: no_nvidia_branch = body.find('if [ -z "$_smi" ]') rocm_call = body.find("_has_amd_rocm_gpu") assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu" assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch" assert ( rocm_call > no_nvidia_branch ), "ROCm detection should sit inside the 'no NVIDIA' branch" assert ( nvidia_call < no_nvidia_branch ), "NVIDIA detection should run before the no-NVIDIA branch" def test_bitsandbytes_amd_install(self): """install.sh should install bitsandbytes for AMD when ROCm detected.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert "bitsandbytes" in source assert "rocm*)" in source # case pattern for ROCm URLs def test_cpu_hint_mentions_amd(self): """CPU-only hint should mention AMD ROCm.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert "ROCm" in source def test_rocm72_supported_future_capped(self): """ROCm 7.2 should pass through directly; 7.3+ falls back to rocm7.2.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert 'echo "$_base/rocm7.2"' in source # fallback for unknown future versions # Allowlisted versions should pass through directly assert "rocm6.*" in source assert "rocm7.0" in source assert "rocm7.1" in source assert "rocm7.2" in source def test_rocm_tag_validation_guard_exists(self): """install.sh should validate _rocm_tag with a case guard.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert "rocm[1-9]*.[0-9]*)" in source assert '_rocm_tag=""' in source # rejection path def test_dpkg_epoch_handling(self): """install.sh should strip Debian epoch prefix from dpkg-query output.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert "sed 's/^[0-9]*://' " in source or "sed 's/^[0-9]*://'" in source def test_no_double_bracket_in_rocm_block(self): """ROCm detection block should not use [[ ]] (bash-only, not POSIX). Note: [[:space:]], [[:digit:]] etc. are valid POSIX character classes, not bash [[ ]].""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] import re for i, line in enumerate(func_body.splitlines(), 1): stripped = line.lstrip() if stripped.startswith("#"): continue # Remove POSIX character classes [[:foo:]] before checking for [[ ]] cleaned = re.sub(r"\[\[:[a-z]+:\]\]", "", line) assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [[" def test_no_arithmetic_expansion_in_rocm_block(self): """ROCm detection block should not use (( )) (bash-only).""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] for i, line in enumerate(func_body.splitlines(), 1): stripped = line.lstrip() if stripped.startswith("#"): continue assert ( "((" not in line or "))" not in line or "$(()" in line ), f"get_torch_index_url line {i} may use non-POSIX (( ))" def test_macos_returns_cpu_before_rocm_check(self): """macOS should return CPU immediately (before any ROCm check).""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") func_body = source[func_start:] darwin_pos = func_body.find("Darwin") rocm_pos = func_body.find("amd-smi") assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection" def test_unsloth_torch_backend_exported_after_get_torch_index_url(self): """install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set. This lets install_python_stack.py skip ROCm torch operations on CUDA and CPU hosts without re-running GPU detection in a subprocess. """ sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)") backend_pos = source.find("UNSLOTH_TORCH_BACKEND") assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh" assert ( backend_pos > torch_url_pos ), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved" # Verify all three cases are covered assert '"cuda"' in source[backend_pos : backend_pos + 500] assert '"rocm"' in source[backend_pos : backend_pos + 500] assert '"cpu"' in source[backend_pos : backend_pos + 500] # Must be exported so subprocesses (setup.sh, install_python_stack.py) see it assert "export UNSLOTH_TORCH_BACKEND" in source def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self): """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098. NVIDIA open kernel module (560+) registers KFD nodes with vendor_id 4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0 (true) on NVIDIA-only hosts that have the nvidia-open driver, causing get_torch_index_url to select a ROCm wheel index. """ sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_amd_rocm_gpu()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] assert ( "vendor_id" in func_body ), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes" assert ( "4098" in func_body ), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)" def test_kfd_awk_resets_state_per_file(self): """KFD sysfs awk must reset gpu/amd state per file (FNR==1). Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU (gpu_id > 0, vendor_id 4318) can produce a false positive: node 0 sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1 before vendor_id 4318 is seen on node 1. """ sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_amd_rocm_gpu()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] assert "FNR==1" in func_body, ( "_has_amd_rocm_gpu KFD awk must reset state per file with FNR==1 " "to avoid false positives on Ryzen+NVIDIA hosts with multiple KFD nodes" ) def test_get_torch_index_url_uses_nvidia_detected_flag(self): """get_torch_index_url must track NVIDIA detection independently of _smi. When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback but nvidia-smi is not on PATH, _smi stays empty. Without a separate _nvidia_detected flag, the function falls into the AMD/CPU branch even though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA. """ sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] assert "_nvidia_detected" in func_body, ( "get_torch_index_url must use a _nvidia_detected flag (separate from " "_smi) so that proc-only NVIDIA detection still selects CUDA wheels" ) # The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on # _smi being empty. assert ( '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body ), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1" # TEST: Live regression on current host (NVIDIA B200 expected) class TestLiveRegression: """Live checks that run on the actual host -- skip if no NVIDIA GPU.""" def test_get_torch_index_url_returns_cuda_on_nvidia(self): """On an NVIDIA machine, get_torch_index_url should return a CUDA URL.""" import shutil if not shutil.which("nvidia-smi"): pytest.skip("No nvidia-smi available") # Skip if nvidia-smi exists but does not actually list a GPU on this # host (containers occasionally ship the binary without a driver). check = subprocess.run( [ "bash", "-c", "nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{f=1} END{exit !f}'", ], capture_output = True, ) if check.returncode != 0: pytest.skip("nvidia-smi is on PATH but no GPU is listed") sh_path = PACKAGE_ROOT / "install.sh" # get_torch_index_url calls _has_usable_nvidia_gpu and # _has_amd_rocm_gpu, so all three function definitions must be # in scope when we eval the extract. extract_cmd = ( f"sed -n '/^_has_amd_rocm_gpu()/,/^}}$/p; " f"/^_has_usable_nvidia_gpu()/,/^}}$/p; " f"/^get_torch_index_url()/,/^}}$/p' '{sh_path}'" ) result = subprocess.run( ["bash", "-c", f'eval "$({extract_cmd})"; get_torch_index_url'], capture_output = True, text = True, timeout = 30, ) if result.returncode != 0: pytest.skip("Could not extract get_torch_index_url for live test") url = result.stdout.strip() assert "cu1" in url or "cuda" in url.lower(), f"Expected CUDA URL, got: {url}" # TEST: worker.py -- ROCm Mamba/SSM source build path # Load worker.py module _WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" _EXPORT_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py" # The torchao Windows-ROCm stub was de-duplicated out of the export/training # workers into a shared module; both workers now call into it. _TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py" # The wheel-probe subprocess was hoisted out of worker.py into wheel_utils # during the wheel-resolver refactor; the probe script literal lives there. _WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py" class TestWorkerRocmMambaSsm: """Verify worker.py Mamba/SSM install logic on ROCm.""" def test_probe_returns_hip_version_field(self): """The wheel probe should include hip_version, and worker.py should consume it.""" assert "hip_version" in _WHEEL_UTILS_PATH.read_text(encoding = "utf-8") assert "hip_version" in _WORKER_PATH.read_text(encoding = "utf-8") def test_probe_script_has_getattr_hip(self): """Probe script should use getattr for torch.version.hip (safe on CUDA).""" source = _WHEEL_UTILS_PATH.read_text(encoding = "utf-8") assert "getattr(torch.version, 'hip', None)" in source def test_direct_wheel_url_returns_none_without_cuda_major(self, monkeypatch): """_direct_wheel_url should return None when cuda_major is empty (ROCm).""" # Load module for function access _worker_spec = importlib.util.spec_from_file_location("test_worker", _WORKER_PATH) assert _worker_spec is not None and _worker_spec.loader is not None worker_mod = importlib.util.module_from_spec(_worker_spec) # Stub worker.py's imports via monkeypatch so the fakes (notably a # non-package "utils") are undone and don't break later tests that # import the real utils.* package. loggers_mock = MagicMock() loggers_mock.get_logger = MagicMock(return_value = MagicMock()) monkeypatch.setitem(sys.modules, "structlog", MagicMock()) monkeypatch.setitem(sys.modules, "loggers", loggers_mock) monkeypatch.setitem(sys.modules, "utils", MagicMock()) monkeypatch.setitem(sys.modules, "utils.hardware", MagicMock()) try: _worker_spec.loader.exec_module(worker_mod) except Exception: pytest.skip("Could not load worker module in test environment") env_rocm = { "python_tag": "cp312", "torch_mm": "2.6", "cuda_major": "", "hip_version": "7.1.12345", "cxx11abi": "TRUE", } result = worker_mod._direct_wheel_url( filename_prefix = "causal_conv1d", package_version = "1.6.1", release_tag = "v1.6.1.post4", release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download", env = env_rocm, ) assert result is None def test_hipcc_check_exists_in_source(self): """worker.py should check for hipcc before ROCm source builds.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "hipcc" in source def test_rocm_source_build_status_message(self): """worker.py should send a specific status for ROCm source compilation.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "Compiling" in source and "from source for ROCm" in source def test_rocm_build_failure_message(self): """worker.py should send a clear error on ROCm build failure.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "Failed to compile" in source and "for ROCm" in source def test_timeout_on_install(self): """worker.py should have a timeout on pip install subprocess.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "TimeoutExpired" in source assert "timeout" in source # TEST: amd.py -- AMD GPU monitoring class TestAmdGpuMonitoring: """Verify amd.py module structure and mock behavior.""" def test_amd_py_exists(self): """amd.py should exist in the hardware directory.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" assert amd_path.exists() def test_amd_py_has_required_functions(self): """amd.py should export the same function signatures as nvidia.py.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" source = amd_path.read_text(encoding = "utf-8") assert "def get_physical_gpu_count" in source assert "def get_primary_gpu_utilization" in source assert "def get_visible_gpu_utilization" in source def test_amd_smi_json_parsing(self, monkeypatch): """Verify _extract_gpu_metrics parses amd-smi JSON correctly.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" _amd_spec = importlib.util.spec_from_file_location("test_amd", amd_path) assert _amd_spec is not None and _amd_spec.loader is not None amd_mod = importlib.util.module_from_spec(_amd_spec) loggers_mock = MagicMock() loggers_mock.get_logger = MagicMock(return_value = MagicMock()) monkeypatch.setitem(sys.modules, "loggers", loggers_mock) try: _amd_spec.loader.exec_module(amd_mod) except Exception: pytest.skip("Could not load amd module in test environment") # Simulate amd-smi metric JSON output gpu_data = { "usage": {"gfx_activity": "85"}, "temperature": {"edge": "72"}, "power": { "current_socket_power": "200.5", "power_cap": "300", }, "vram": { "vram_used": 8192, # MB "vram_total": 16384, # MB }, } metrics = amd_mod._extract_gpu_metrics(gpu_data) assert metrics["gpu_utilization_pct"] == 85.0 assert metrics["temperature_c"] == 72.0 assert metrics["power_draw_w"] == 200.5 assert metrics["power_limit_w"] == 300.0 assert metrics["vram_used_gb"] == round(8192 / 1024, 2) assert metrics["vram_total_gb"] == round(16384 / 1024, 2) assert metrics["vram_utilization_pct"] is not None assert metrics["power_utilization_pct"] is not None def test_amd_primary_gpu_with_mock(self, monkeypatch): """get_primary_gpu_utilization returns correct dict with mocked amd-smi.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" _amd_spec = importlib.util.spec_from_file_location("test_amd2", amd_path) assert _amd_spec is not None and _amd_spec.loader is not None amd_mod = importlib.util.module_from_spec(_amd_spec) loggers_mock = MagicMock() loggers_mock.get_logger = MagicMock(return_value = MagicMock()) monkeypatch.setitem(sys.modules, "loggers", loggers_mock) try: _amd_spec.loader.exec_module(amd_mod) except Exception: pytest.skip("Could not load amd module") # _first_visible_amd_gpu_id() short-circuits to None when any of # HIP / ROCR / CUDA_VISIBLE_DEVICES is "" or "-1". CI often sets # CUDA_VISIBLE_DEVICES="", so the test must not inherit that. for var in ( "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES", ): monkeypatch.delenv(var, raising = False) # amd-smi is gated off on Windows w/o a HIP SDK; this test mocks it as # available, so opt in so the gate allows it on every platform. monkeypatch.setenv("UNSLOTH_ENABLE_AMD_SMI", "1") mock_json = json.dumps( [ { "usage": {"gfx_activity": "50"}, "temperature": {"edge": "65"}, "power": {"current_socket_power": "150", "power_cap": "250"}, "vram": {"vram_used": 4096, "vram_total": 16384}, } ] ) mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = mock_json # The premise is "amd-smi exists and answers": the absence guard # which()-checks before spawning, so hosts without a real amd-smi # (Linux CI, driver-only Windows) need which mocked too. with patch.object(amd_mod.shutil, "which", return_value = "/usr/bin/amd-smi"): with patch.object(subprocess, "run", return_value = mock_result): result = amd_mod.get_primary_gpu_utilization() assert result["available"] is True assert result["gpu_utilization_pct"] == 50.0 assert result["temperature_c"] == 65.0 def test_amd_smi_not_found_returns_unavailable(self, monkeypatch): """get_primary_gpu_utilization returns available=False when amd-smi is missing.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" _amd_spec = importlib.util.spec_from_file_location("test_amd3", amd_path) assert _amd_spec is not None and _amd_spec.loader is not None amd_mod = importlib.util.module_from_spec(_amd_spec) loggers_mock = MagicMock() loggers_mock.get_logger = MagicMock(return_value = MagicMock()) monkeypatch.setitem(sys.modules, "loggers", loggers_mock) try: _amd_spec.loader.exec_module(amd_mod) except Exception: pytest.skip("Could not load amd module") # Opt in so the call reaches subprocess.run (gated off on Windows w/o a # HIP SDK); testing the OSError handling here. with ( patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}), patch.object(subprocess, "run", side_effect = OSError("amd-smi not found")), ): result = amd_mod.get_primary_gpu_utilization() assert result["available"] is False def test_amd_timeout_returns_unavailable(self, monkeypatch): """get_primary_gpu_utilization handles timeout gracefully.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" _amd_spec = importlib.util.spec_from_file_location("test_amd4", amd_path) assert _amd_spec is not None and _amd_spec.loader is not None amd_mod = importlib.util.module_from_spec(_amd_spec) loggers_mock = MagicMock() loggers_mock.get_logger = MagicMock(return_value = MagicMock()) monkeypatch.setitem(sys.modules, "loggers", loggers_mock) try: _amd_spec.loader.exec_module(amd_mod) except Exception: pytest.skip("Could not load amd module") # Opt in so the call reaches subprocess.run (gated off on Windows w/o a # HIP SDK); testing the timeout handling here. with ( patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}), patch.object( subprocess, "run", side_effect = subprocess.TimeoutExpired("amd-smi", 5), ), ): result = amd_mod.get_primary_gpu_utilization() assert result["available"] is False # TEST: hardware.py -- IS_ROCM branching to amd.py class TestHardwareAmdBranching: """Verify hardware.py branches to amd.py when IS_ROCM is True.""" def test_hardware_imports_amd_module(self): """hardware.py should import from amd module when IS_ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): """get_gpu_utilization should dispatch to amd.py via _smi_query when IS_ROCM, and the dispatcher itself must check IS_ROCM and import the amd backend.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert '_smi_query("get_primary_gpu_utilization"' in func_body smi = source[ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) ] assert "IS_ROCM" in smi assert "from . import amd" in smi def test_hardware_branches_on_is_rocm_for_visible(self): """get_visible_gpu_utilization should dispatch to amd.py via _smi_query when IS_ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_visible_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] # The dispatcher call may wrap onto multiple lines; allow whitespace # between the open paren and the literal func name argument. import re as _re assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body) smi = source[ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) ] assert "IS_ROCM" in smi assert "from . import amd" in smi def test_hardware_branches_on_is_rocm_for_physical_count(self): """get_physical_gpu_count should try amd.py when IS_ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_physical_gpu_count") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "IS_ROCM" in func_body assert "from . import amd" in func_body # TEST: hardware.py -- apply_gpu_ids ROCm fallback (issue #5180) class TestApplyGpuIdsRocmFallback: """Verify apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when IS_ROCM is still False (worker subprocess before detect_hardware runs).""" def test_apply_gpu_ids_falls_back_to_torch_version_hip(self): """apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert 'getattr(_torch.version, "hip", None)' in func_body def test_apply_gpu_ids_sets_hip_but_not_rocr_visible_devices(self): """apply_gpu_ids should set HIP_VISIBLE_DEVICES but leave ROCR_VISIBLE_DEVICES inherited. ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU indices. Overwriting it breaks multi-GPU ROCm systems (see issue #6118). """ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert 'os.environ["HIP_VISIBLE_DEVICES"] = value' in func_body assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' not in func_body def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self): """torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "import torch as _torch" in func_body assert "except Exception" in func_body # TEST: install_python_stack.py -- Windows AMD warning class TestWindowsRocmWarning: """Verify Windows AMD GPU detection and warning message.""" def test_windows_amd_warning_in_source(self): """install_python_stack.py should warn Windows AMD users.""" source = _STACK_PATH.read_text(encoding = "utf-8") assert "AMD GPU detected" in source def test_windows_amd_warning_checks_hipinfo_or_amdsmi(self): """Warning should check for hipinfo or amd-smi.""" source = _STACK_PATH.read_text(encoding = "utf-8") assert "hipinfo" in source assert "amd-smi" in source def test_windows_amd_warning_has_docs_link(self): """Warning should include AMD docs link.""" source = _STACK_PATH.read_text(encoding = "utf-8") assert "docs.unsloth.ai/get-started/install-and-update/amd" in source # TEST: unsloth/kernels/utils.py -- is_rdna() expansion class TestIsRdnaExpansion: """Verify is_rdna() covers RDNA2, RDNA3, RDNA3.5, RDNA4 architectures.""" def test_is_rdna_source_has_rdna2(self): """is_rdna() should include RDNA2 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1030" in func_body assert "gfx1031" in func_body assert "gfx1032" in func_body assert "gfx1033" in func_body assert "gfx1034" in func_body assert "gfx1035" in func_body assert "gfx1036" in func_body def test_is_rdna_source_has_rdna3(self): """is_rdna() should include RDNA3 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1100" in func_body assert "gfx1101" in func_body assert "gfx1102" in func_body assert "gfx1103" in func_body def test_is_rdna_source_has_rdna35(self): """is_rdna() should include RDNA3.5 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1150" in func_body assert "gfx1151" in func_body assert "gfx1152" in func_body def test_is_rdna_source_has_rdna4(self): """is_rdna() should include RDNA4 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1200" in func_body assert "gfx1201" in func_body def test_is_cdna_not_changed(self): """is_cdna() should remain unchanged (no RDNA architectures added).""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_cdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx940" in func_body assert "gfx941" in func_body assert "gfx942" in func_body assert "gfx950" in func_body # RDNA architectures should NOT be in is_cdna assert "gfx1030" not in func_body assert "gfx1100" not in func_body # TEST: install_python_stack.py -- _windows_rocm_index_url arch mapping class TestWindowsRocmIndexUrl: """Verify GPU arch → AMD pip index URL mapping.""" def test_gfx1200_maps_to_gfx120x_all(self): url = stack_mod._windows_rocm_index_url("gfx1200") assert url is not None assert "gfx120X-all" in url def test_gfx1201_maps_to_gfx120x_all(self): url = stack_mod._windows_rocm_index_url("gfx1201") assert url is not None assert "gfx120X-all" in url def test_gfx1151_maps_to_gfx1151(self): url = stack_mod._windows_rocm_index_url("gfx1151") assert url is not None assert "gfx1151" in url def test_gfx1150_maps_to_gfx1150(self): url = stack_mod._windows_rocm_index_url("gfx1150") assert url is not None assert "gfx1150" in url def test_gfx1100_maps_to_gfx110x_all(self): url = stack_mod._windows_rocm_index_url("gfx1100") assert url is not None assert "gfx110X-all" in url def test_unknown_arch_returns_none(self): assert stack_mod._windows_rocm_index_url("gfx9999") is None def test_none_arch_returns_none(self): assert stack_mod._windows_rocm_index_url(None) is None def test_url_ends_with_slash(self): """AMD pip index URLs must end with / for --index-url compatibility.""" url = stack_mod._windows_rocm_index_url("gfx1200") assert url is not None assert url.endswith("/") def test_base_url_uses_repo_amd_com_by_default(self): url = stack_mod._windows_rocm_index_url("gfx1200") assert url is not None assert "repo.amd.com" in url def test_mirror_env_var_overrides_base(self, monkeypatch): monkeypatch.setenv("UNSLOTH_ROCM_WINDOWS_MIRROR", "https://my-mirror.example.com/rocm/whl") # Reload module-level constant by calling helper directly url = stack_mod._windows_rocm_index_url("gfx1200") # The env var is read at module load time for _ROCM_WINDOWS_INDEX_BASE, # so just verify the helper itself doesn't error. assert url is not None # TEST: install_python_stack.py -- _detect_windows_gfx_arch class TestDetectWindowsGfxArch: """Verify hipinfo parsing for GPU arch detection on Windows.""" def test_returns_none_when_hipinfo_not_on_path(self): # Also neutralise the venv-hipInfo and WMI-name fallbacks: this test # pins "no probe source available -> None", and the suite may run on a # real AMD host where WMI would legitimately answer. with patch("shutil.which", return_value = None): with patch("os.path.isfile", return_value = False): with patch("subprocess.run", side_effect = FileNotFoundError): result = stack_mod._detect_windows_gfx_arch() assert result is None def test_parses_gcnarchname_from_hipinfo_output(self): mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n" with patch("shutil.which", return_value = "/usr/bin/hipinfo"): with patch("subprocess.run", return_value = mock_result): result = stack_mod._detect_windows_gfx_arch() assert result == "gfx1200" def test_returns_none_on_nonzero_returncode(self): mock_result = MagicMock() mock_result.returncode = 1 mock_result.stdout = b"gcnArchName : gfx1200\n" with patch("shutil.which", return_value = "/usr/bin/hipinfo"): with patch("subprocess.run", return_value = mock_result): result = stack_mod._detect_windows_gfx_arch() assert result is None def test_returns_none_when_no_gcnarchname_in_output(self): # hipinfo answers but without a gcnArchName line. Route only the # hipinfo/amd-smi probes to the mock; the WMI fallback must get # nothing (FileNotFoundError) -- otherwise the mocked device name # would legitimately resolve via the name table. mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = b"deviceName : SomeUnknownDevice\n" def _run(cmd, **kwargs): if cmd and "powershell" in str(cmd[0]).lower(): raise FileNotFoundError(cmd[0]) return mock_result with patch("shutil.which", return_value = "/usr/bin/hipinfo"): with patch("subprocess.run", side_effect = _run): result = stack_mod._detect_windows_gfx_arch() assert result is None def test_returns_none_on_timeout(self): with patch("shutil.which", return_value = "/usr/bin/hipinfo"): with patch( "subprocess.run", side_effect = subprocess.TimeoutExpired("hipinfo", 10), ): result = stack_mod._detect_windows_gfx_arch() assert result is None def test_strips_whitespace_from_arch(self): mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = b" gcnArchName : gfx1201 \n" with patch("shutil.which", return_value = "/usr/bin/hipinfo"): with patch("subprocess.run", return_value = mock_result): result = stack_mod._detect_windows_gfx_arch() assert result == "gfx1201" # TEST: install_python_stack.py -- GPU-name / WMI fallback (no amd-smi, no hipinfo) class TestGfxArchNameFallback: """amd-smi does not exist on Windows (neither Adrenalin consistently nor the HIP SDK ship a CLI) and driver-only hosts lack hipinfo too. The detection chain must still resolve the arch from the GPU marketing name (WMI), mirroring setup.ps1's $nameArchTable.""" @pytest.mark.parametrize( "name, expected", [ ("AMD Radeon(TM) 8060S Graphics", "gfx1151"), ("AMD Ryzen AI MAX+ 395 w/ Radeon 8060S", "gfx1151"), ("AMD Radeon(TM) 890M", "gfx1150"), ("AMD Ryzen AI 9 HX 370 w/ Radeon 890M", "gfx1150"), ("AMD Radeon RX 9070 XT", "gfx1201"), ("AMD Radeon RX 9070", "gfx1200"), ("AMD Radeon RX 7700S", "gfx1102"), # (?!S) lookahead must not hit gfx1100 ("AMD Radeon RX 7700 XT", "gfx1100"), ("AMD Radeon(TM) 780M", "gfx1103"), ("NVIDIA GeForce RTX 4090", None), ("Microsoft Basic Display Adapter", None), ("", None), ], ) def test_name_to_arch_mapping(self, name, expected): assert stack_mod._gfx_arch_from_gpu_name(name) == expected def test_wmi_fallback_resolves_arch_without_any_tools(self): """hipinfo absent everywhere + amd-smi absent -> WMI name fallback.""" ps_result = MagicMock() ps_result.returncode = 0 ps_result.stdout = b"AMD Radeon(TM) 8060S Graphics\r\nMicrosoft Basic Display Adapter\r\n" def _run(cmd, **kwargs): if cmd and "powershell.exe" in str(cmd[0]).lower(): return ps_result raise FileNotFoundError(cmd[0]) with patch.dict(os.environ, {}, clear = False): for _v in ( "HIP_PATH", "ROCM_PATH", "UNSLOTH_ROCM_GFX_ARCH", "UNSLOTH_ENABLE_AMD_SMI", ): os.environ.pop(_v, None) with patch("shutil.which", return_value = None): with patch("os.path.isfile", return_value = False): with patch("subprocess.run", side_effect = _run): result = stack_mod._detect_windows_gfx_arch() assert result == "gfx1151" def test_wmi_fallback_returns_none_for_non_amd_hosts(self): ps_result = MagicMock() ps_result.returncode = 0 ps_result.stdout = b"NVIDIA GeForce RTX 4090\r\n" def _run(cmd, **kwargs): if cmd and "powershell.exe" in str(cmd[0]).lower(): return ps_result raise FileNotFoundError(cmd[0]) with patch.dict(os.environ, {}, clear = False): for _v in ("HIP_PATH", "ROCM_PATH", "UNSLOTH_ROCM_GFX_ARCH"): os.environ.pop(_v, None) with patch("shutil.which", return_value = None): with patch("os.path.isfile", return_value = False): with patch("subprocess.run", side_effect = _run): result = stack_mod._detect_windows_gfx_arch() assert result is None def test_stack_probes_venv_hipinfo(self): """The venv Scripts dir hipInfo.exe (shipped by AMD torch wheels) must be a probe candidate so `studio update` works on driver-only hosts.""" source = _STACK_PATH.read_text(encoding = "utf-8") assert 'os.path.join(os.path.dirname(sys.executable), "hipInfo.exe")' in source def test_prebuilt_resolve_exe_probes_venv_dir(self): """install_llama_prebuilt's _resolve_exe must include the venv Scripts candidate for the same driver-only standalone-rerun scenario.""" source = _PREBUILT_PATH.read_text(encoding = "utf-8") assert "_venv_candidate" in source def test_runtime_monitor_guards_amd_smi_absence(self): """amd.py must which()-check amd-smi before spawning so absence disables the poller in one step (no FileNotFoundError strikes).""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" source = amd_path.read_text(encoding = "utf-8") assert 'shutil.which("amd-smi") is None' in source # TEST: install_python_stack.py -- _install_bnb_windows_rocm 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: stack_mod._install_bnb_windows_rocm() assert mock_pip.call_count == 1 call_args = str(mock_pip.call_args_list[0]) assert "bitsandbytes" in call_args assert "win_amd64" in call_args # Must force plain pip (uv mangles the bitsandbytes wheel) -- see # https://unsloth.ai/docs/get-started/install/amd/amd-hackathon assert mock_pip.call_args.kwargs.get("force_pip") is True def test_forces_plain_pip_not_uv(self): """The bnb wheel must be installed with plain pip, never uv.""" with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip: stack_mod._install_bnb_windows_rocm() assert mock_pip.call_args.kwargs.get("force_pip") is True def test_does_not_touch_uv_skip_env_var(self): """The UV_SKIP_WHEEL_FILENAME_CHECK hack is gone; the env must be untouched.""" observed = {} def _capture(*args, **kwargs): observed["during"] = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") return True with patch.dict(os.environ, {}, clear = False): os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) with patch.object(stack_mod, "pip_install_try", side_effect = _capture): stack_mod._install_bnb_windows_rocm() assert observed.get("during") is None assert "UV_SKIP_WHEEL_FILENAME_CHECK" not in os.environ def test_returns_false_on_pip_failure(self): """A failed pip_install_try must surface as a False return, not BNB_ROCM_VERSION.""" with patch.dict(os.environ, {}, clear = False): os.environ.pop("BNB_ROCM_VERSION", None) with patch.object(stack_mod, "pip_install_try", return_value = False): result = stack_mod._install_bnb_windows_rocm() assert result is False assert "BNB_ROCM_VERSION" not in os.environ def test_no_op_when_win_amd64_url_missing(self): """Should be silent no-op if win_amd64 key absent from _BNB_ROCM_PRERELEASE_URLS.""" with patch.object(stack_mod, "_BNB_ROCM_PRERELEASE_URLS", {}): with patch.object(stack_mod, "pip_install_try") as mock_pip: stack_mod._install_bnb_windows_rocm() mock_pip.assert_not_called() def test_sets_bnb_rocm_version_from_detected_dll(self): """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() assert os.environ.get("BNB_ROCM_VERSION") == "72" def test_sets_bnb_rocm_version_from_newer_dll(self): """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() assert os.environ.get("BNB_ROCM_VERSION") == "713" def test_falls_back_to_72_when_detection_fails(self): """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() assert os.environ.get("BNB_ROCM_VERSION") == "72" 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 ( '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source ), path.name def test_main_requires_found_rocm_dll(self): """HIP_PATH/ROCM_PATH alone (HIP SDK on a CUDA/CPU box) must not force a ROCm backend onto a non-ROCm bitsandbytes.""" source = self._MAIN_PATH.read_text(encoding = "utf-8") assert "if _found_rocm_bnb:" in source assert "_hip_env" not in source def test_worker_requires_found_rocm_dll(self): """No DLL found: the worker must not write any override or touch the seeded marker (later import fixes must still see sitecustomize).""" source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8") assert "if _found_rocm_bnb:" in source class TestDetectBnbRocmDllVer: """Unit tests for _detect_bnb_rocm_dll_ver().""" def test_returns_none_when_bnb_not_installed(self): """Returns None if bitsandbytes is not importable.""" import importlib.util with patch.object(importlib.util, "find_spec", return_value = None): assert stack_mod._detect_bnb_rocm_dll_ver() is None def test_detects_rocm72_dll(self, tmp_path): """Returns '72' when libbitsandbytes_rocm72.dll is present.""" (tmp_path / "libbitsandbytes_rocm72.dll").write_text("") mock_spec = MagicMock() mock_spec.submodule_search_locations = [str(tmp_path)] import importlib.util with patch.object(importlib.util, "find_spec", return_value = mock_spec): assert stack_mod._detect_bnb_rocm_dll_ver() == "72" def test_detects_rocm713_dll(self, tmp_path): """Returns '713' when libbitsandbytes_rocm713.dll is present.""" (tmp_path / "libbitsandbytes_rocm713.dll").write_text("") mock_spec = MagicMock() mock_spec.submodule_search_locations = [str(tmp_path)] import importlib.util with patch.object(importlib.util, "find_spec", return_value = mock_spec): assert stack_mod._detect_bnb_rocm_dll_ver() == "713" def test_returns_none_when_only_cuda_dlls(self, tmp_path): """Returns None when only CUDA DLLs are present (no ROCm DLL).""" (tmp_path / "libbitsandbytes_cuda121.dll").write_text("") mock_spec = MagicMock() mock_spec.submodule_search_locations = [str(tmp_path)] import importlib.util with patch.object(importlib.util, "find_spec", return_value = mock_spec): assert stack_mod._detect_bnb_rocm_dll_ver() is None def test_picks_highest_suffix_when_multiple_dlls(self, tmp_path): """Returns the highest numeric suffix when multiple ROCm DLL variants exist. Filesystem glob order is not guaranteed, so the function must not stop at the first match — it must always return the highest one. """ (tmp_path / "libbitsandbytes_rocm72.dll").write_text("") (tmp_path / "libbitsandbytes_rocm713.dll").write_text("") mock_spec = MagicMock() mock_spec.submodule_search_locations = [str(tmp_path)] import importlib.util with patch.object(importlib.util, "find_spec", return_value = mock_spec): assert stack_mod._detect_bnb_rocm_dll_ver() == "713" # TEST: install_python_stack.py -- UNSLOTH_ROCM_TORCH_INSTALLED early-return path 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"}), patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe), ): stack_mod._ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "_install_bnb_windows_rocm") @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"}), patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe), ): stack_mod._ensure_rocm_torch() mock_bnb.assert_called_once() @patch.object(stack_mod, "_install_bnb_windows_rocm") @patch.object(stack_mod, "pip_install") 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"}), 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) class TestWorkerWindowsRocmPatches: """Verify worker.py contains the required Windows ROCm runtime patches.""" def test_grouped_mm_dispatch_patch_present(self): """worker.py must register a _grouped_mm CUDA dispatch override.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert '_gm_lib.impl("_grouped_mm"' in source def test_grouped_mm_patch_targets_cuda_dispatch_key(self): """The dispatch override must target the CUDA key (not CompositeImplicitAutograd).""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert '"_grouped_mm", _grouped_mm_safe_impl, "CUDA"' in source def test_grouped_mm_lib_kept_alive(self): """The Library object must be stored to prevent GC clearing the registration.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "_WINDOWS_ROCM_GROUPED_MM_LIB" in source def test_grouped_mm_handles_offs_grouped_case(self): """_grouped_mm fallback must handle the grouped (offs!=None) variant.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "offs_list" in source assert "offs.tolist()" in source def test_worker_calls_shared_torchao_stub(self): """worker.py must invoke the shared torchao stub entrypoint.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "install_torchao_windows_rocm_stub()" in source def test_export_worker_calls_shared_torchao_stub(self): """export/worker.py must invoke the same shared torchao stub entrypoint.""" source = _EXPORT_WORKER_PATH.read_text(encoding = "utf-8") assert "install_torchao_windows_rocm_stub()" in source def test_torchao_stub_uses_stub_type_meta(self): """Torchao stub must use _StubTypeMeta so isinstance() returns False not TypeError.""" source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "_StubTypeMeta" in source def test_stub_type_meta_has_instancecheck(self): """_StubTypeMeta must define __instancecheck__ returning False.""" source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "__instancecheck__" in source def test_stub_subpackage_finder_registered(self): """_StubSubpackageFinder must be appended to sys.meta_path.""" source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "sys.meta_path.append(_StubSubpackageFinder())" in source def test_torchao_key_submodules_pre_stubbed(self): """Key torchao submodules (dtypes, quantization) must be pre-stubbed.""" source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8") assert "torchao.dtypes" in source assert "torchao.quantization" in source def test_torchdynamo_disabled_on_windows_rocm(self): """worker.py should disable dynamo on Windows ROCm as belt-and-suspenders.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "TORCHDYNAMO_DISABLE" in source def test_bnb_rocm_version_set_on_windows_rocm(self): """worker.py must set BNB_ROCM_VERSION in the Windows ROCm section. BNB auto-detects HIP version from torch.version.hip, which can mismatch the DLL suffix in the AMD prerelease wheel. The worker must detect the actual DLL suffix and override BNB's auto-detection before ML imports. """ source = _WORKER_PATH.read_text(encoding = "utf-8") # Env var must be set assert "BNB_ROCM_VERSION" in source # Detection helper must be used assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source # Falls back to the seeded value, never a blind "72" (which would # force a ROCm backend onto a non-ROCm bitsandbytes wheel) assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION")' in source def test_bnb_rocm_version_set_before_ml_imports(self): """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports.""" source = _WORKER_PATH.read_text(encoding = "utf-8") idx_bnb = source.find("BNB_ROCM_VERSION") # Use the specific section-2 marker that appears in the worker process # entry-point function (not the trainer helper which has its own "# ── 2."). idx_sec2 = source.find("# ── 2. Now import ML libraries") assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py" assert idx_sec2 != -1, "'# ── 2. Now import ML libraries' marker not found in worker.py" assert idx_bnb < idx_sec2, ( "BNB_ROCM_VERSION must be set before section 2 ML imports " f"(found at {idx_bnb}, section 2 at {idx_sec2})" ) def test_grouped_mm_patch_guarded_by_windows_and_hip_check(self): """_grouped_mm patch must only apply on Windows + HIP torch.""" source = _WORKER_PATH.read_text(encoding = "utf-8") # Must check sys.platform == "win32" assert 'sys.platform == "win32"' in source # Must gate on HIP version — code uses getattr chain: "version" and "hip" assert '"version"' in source and '"hip"' in source def test_hip_ver_at_least_helper_defined(self): """_hip_ver_at_least helper must be defined inside the Windows ROCm block.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "def _hip_ver_at_least(major: int, minor: int)" in source def test_grouped_mm_patch_gated_on_hip_lt_713(self): """_grouped_mm patch must be skipped on HIP >= 7.13 (AMD fixed the bug in ROCm 7.13).""" source = _WORKER_PATH.read_text(encoding = "utf-8") # The guard must call _hip_ver_at_least with exactly (7, 13) assert "_hip_ver_at_least(7, 13)" in source # The patch must be inside the `if not` branch (negated guard) assert "if not _hip_ver_at_least(7, 13):" in source def test_grouped_mm_hip_713_skip_message_present(self): """worker.py must log a message when skipping the patch on HIP >= 7.13.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "HIP >= 7.13" in source assert "7.13" in source def test_grouped_mm_patch_else_branch_present(self): """An else branch must follow the _hip_ver_at_least gate (skip path for 7.13+).""" source = _WORKER_PATH.read_text(encoding = "utf-8") # There must be an else: after the if not _hip_ver_at_least(7, 13): block gate_idx = source.find("if not _hip_ver_at_least(7, 13):") assert gate_idx != -1, "Version gate not found in worker.py" # The else: branch must appear after the gate else_idx = source.find("else:", gate_idx) assert else_idx != -1, "else: branch after _hip_ver_at_least gate not found" def test_hip_ver_at_least_handles_amd_version_format(self): """_hip_ver_at_least must split on '.' and compare only major.minor (handles '7.13.99004').""" source = _WORKER_PATH.read_text(encoding = "utf-8") # Must split the version string and take the first two parts assert 'split(".")[:2]' in source or ".split('.')[:2]" in source # TEST: install_python_stack.py -- _ROCM_TORCH_PKG_SPECS mapping class TestRocmTorchPkgSpecs: """Verify per-tag torch version specs are correct.""" def test_rocm72_has_torch_211(self): """rocm7.2 should specify torch 2.11.x.""" specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("rocm7.2") assert specs is not None torch_spec = specs[0] assert "2.11" in torch_spec def test_default_caps_below_211(self): """Default spec (rocm7.1 and earlier) should cap below 2.11.""" specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("_default") assert specs is not None torch_spec = specs[0] assert "<2.11" in torch_spec def test_specs_have_torch_vision_audio(self): """Each entry should be a 3-tuple: torch, torchvision, torchaudio.""" for tag, specs in stack_mod._ROCM_TORCH_PKG_SPECS.items(): assert len(specs) == 3, f"{tag}: expected (torch, torchvision, torchaudio)" assert "torch" in specs[0] assert "torchvision" in specs[1] assert "torchaudio" in specs[2] def test_gfx_to_amd_index_covers_rdna4(self): """_GFX_TO_AMD_INDEX_ARCH must cover gfx1200 and gfx1201 (RDNA 4).""" mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH assert mapping.get("gfx1200") == "gfx120X-all" assert mapping.get("gfx1201") == "gfx120X-all" def test_gfx_to_amd_index_covers_strix_halo(self): """_GFX_TO_AMD_INDEX_ARCH must cover gfx1151 and gfx1150 (RDNA 3.5).""" mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH assert mapping.get("gfx1151") == "gfx1151" assert mapping.get("gfx1150") == "gfx1150" def test_gfx_to_amd_index_covers_rdna3(self): """_GFX_TO_AMD_INDEX_ARCH must cover gfx1100-gfx1103 (RDNA 3).""" mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH for arch in ("gfx1100", "gfx1101", "gfx1102", "gfx1103"): assert mapping.get(arch) == "gfx110X-all", f"{arch} missing from mapping" # TEST: setup.ps1 / install.ps1 -- Strix Halo gfx arch detection _SETUP_PS1_PATH = PACKAGE_ROOT / "studio" / "setup.ps1" _INSTALL_PS1_PATH = PACKAGE_ROOT / "install.ps1" class TestStrixHaloGfxArchDetection: """Verify that setup.ps1 and install.ps1 have robust gfx arch detection for Strix Halo / iGPU users who only have the HIP runtime (no hipinfo).""" def test_amd_smi_static_asic_attempted_in_setup(self): """setup.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "static --asic" in source def test_amd_smi_static_asic_attempted_in_install(self): """install.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "static --asic" in source def test_env_var_override_in_setup(self): """setup.ps1 must honour UNSLOTH_ROCM_GFX_ARCH as a manual arch override.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "UNSLOTH_ROCM_GFX_ARCH" in source def test_env_var_override_in_install(self): """install.ps1 must honour UNSLOTH_ROCM_GFX_ARCH as a manual arch override.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "UNSLOTH_ROCM_GFX_ARCH" in source def test_name_arch_table_covers_strix_halo_in_setup(self): """setup.ps1 name→arch table must map 890M / Strix Halo to gfx1151.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "gfx1151" in source assert "890M" in source or "Strix Halo" in source def test_name_arch_table_covers_strix_halo_in_install(self): """install.ps1 name→arch table must map 890M / Strix Halo to gfx1151.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "gfx1151" in source assert "890M" in source or "Strix Halo" in source def test_name_arch_table_covers_strix_point_in_setup(self): """setup.ps1 name→arch table must map 880M / Strix Point to gfx1150.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "gfx1150" in source assert "880M" in source or "Strix Point" in source def test_name_arch_table_covers_strix_point_in_install(self): """install.ps1 name→arch table must map 880M / Strix Point to gfx1150.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "gfx1150" in source assert "880M" in source or "Strix Point" in source def test_name_arch_table_covers_rdna3_phoenix_in_setup(self): """setup.ps1 name→arch table must map 780M / Phoenix to gfx1103.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "gfx1103" in source assert "780M" in source or "Phoenix" in source def test_wmi_does_not_set_hasrocm_in_setup(self): """WMI block in setup.ps1 must NOT set $HasROCm = $true (no runtime confirmation).""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") # Find the WMI block and confirm HasROCm is not set inside it wmi_idx = source.find("Win32_VideoController") assert wmi_idx != -1, "WMI block not found in setup.ps1" # The nearest HasROCm = $true must not appear between the WMI block # and the closing brace of that if-block. We check by confirming # $HasROCm = $true does NOT appear within 300 chars of the WMI call. wmi_context = source[wmi_idx : wmi_idx + 300] assert "$HasROCm = $true" not in wmi_context def test_gfx_arch_regex_parses_from_amd_smi_output(self): """Both files must use the gfx\\d+[a-z]? regex to parse arch from amd-smi output.""" for path in (_SETUP_PS1_PATH, _INSTALL_PS1_PATH): source = path.read_text(encoding = "utf-8") # The regex pattern used to match gfx arches assert ( "gfx\\d+" in source or r"gfx\d+" in source ), f"gfx arch regex not found in {path.name}" # TEST: HIP SDK tool path resolution via HIP_PATH / ROCM_PATH env vars class TestHipSdkEnvPathResolution: """Verify that both install scripts resolve hipinfo/hipconfig via HIP_PATH and ROCM_PATH when the tools are not on $PATH, and emit explicit warnings.""" # ── hipinfo resolution ──────────────────────────────────────────────────── def test_setup_checks_hip_path_for_hipinfo(self): """setup.ps1 must reference HIP_PATH when resolving hipinfo.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "HIP_PATH" in source assert "hipinfo" in source def test_install_checks_hip_path_for_hipinfo(self): """install.ps1 must reference HIP_PATH when resolving hipinfo.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "HIP_PATH" in source assert "hipinfo" in source def test_setup_checks_rocm_path_as_hipinfo_fallback(self): """setup.ps1 must also check ROCM_PATH as a secondary hipinfo fallback.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "ROCM_PATH" in source # Confirm the fallback pattern: HIP_PATH ?? ROCM_PATH (or equivalent elseif) assert "ROCM_PATH" in source and "HIP_PATH" in source def test_install_checks_rocm_path_as_hipinfo_fallback(self): """install.ps1 must also check ROCM_PATH as a secondary hipinfo fallback.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "ROCM_PATH" in source assert "ROCM_PATH" in source and "HIP_PATH" in source def test_setup_resolves_hipinfo_via_bin_subdir(self): """setup.ps1 must join the env var root with 'bin\\hipinfo.exe'.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert r"bin\hipinfo.exe" in source def test_install_resolves_hipinfo_via_bin_subdir(self): """install.ps1 must join the env var root with 'bin\\hipinfo.exe'.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert r"bin\hipinfo.exe" in source # ── hipinfo not-on-PATH warning ─────────────────────────────────────────── def test_setup_warns_when_hipinfo_not_on_path(self): """setup.ps1 must warn when hipinfo is found via env var but not on PATH.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "hipinfo not on PATH" in source def test_install_warns_when_hipinfo_not_on_path(self): """install.ps1 must warn when hipinfo is found via env var but not on PATH.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "hipinfo not on PATH" in source # ── warn when HIP_PATH set but exe missing ──────────────────────────────── def test_setup_warns_when_hip_path_set_but_exe_missing(self): """setup.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") # The warning must mention that the SDK install may be incomplete assert "incomplete" in source or "not found at" in source def test_install_warns_when_hip_path_set_but_exe_missing(self): """install.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "incomplete" in source or "not found at" in source # ── hipinfo runtime error warning ───────────────────────────────────────── def test_setup_warns_on_hipinfo_nonzero_exit(self): """setup.ps1 must warn when hipinfo runs but returns a non-zero exit code.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "HIP runtime error" in source or "runtime error" in source.lower() def test_install_warns_on_hipinfo_nonzero_exit(self): """install.ps1 must warn when hipinfo runs but returns a non-zero exit code.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "HIP runtime error" in source or "runtime error" in source.lower() # ── hipconfig resolution ────────────────────────────────────────────────── def test_setup_resolves_hipconfig_via_bin_subdir(self): """setup.ps1 must also fall back to HIP_PATH/bin/hipconfig.exe for version detection.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert r"bin\hipconfig.exe" in source def test_install_resolves_hipconfig_via_bin_subdir(self): """install.ps1 must also fall back to HIP_PATH/bin/hipconfig.exe for version detection.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert r"bin\hipconfig.exe" in source def test_setup_warns_when_hipconfig_not_on_path(self): """setup.ps1 must warn when hipconfig is found via env var but not on PATH.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "hipconfig not on PATH" in source def test_install_warns_when_hipconfig_not_on_path(self): """install.ps1 must warn when hipconfig is found via env var but not on PATH.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "hipconfig not on PATH" in source # ── PATH fix hint ───────────────────────────────────────────────────────── def test_setup_provides_path_fix_hint(self): """setup.ps1 must tell the user how to add the HIP bin dir to PATH.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") # Should mention adding to PATH or SetEnvironmentVariable assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source) def test_install_provides_path_fix_hint(self): """install.ps1 must tell the user how to add the HIP bin dir to PATH.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source) # TEST: HIP SDK detected substep -- path + hipconfig version shown in terminal class TestHipSdkDetectedSubstep: """Verify that both scripts print HIP SDK path and full hipconfig version as substeps under the gpu step when AMD ROCm is successfully detected.""" def test_setup_prints_hip_sdk_path_substep(self): """setup.ps1 must print an 'HIP SDK:' substep showing the resolved path.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "HIP SDK:" in source def test_install_prints_hip_sdk_path_substep(self): """install.ps1 must print an 'HIP SDK:' substep showing the resolved path.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "HIP SDK:" in source def test_setup_shows_hipconfig_full_version(self): """setup.ps1 must capture and display the full hipconfig version string.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "ROCmVersionFull" in source or "hipconfig:" in source def test_install_shows_hipconfig_full_version(self): """install.ps1 must capture and display the full hipconfig version string.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "ROCmVersionFull" in source or "hipconfig:" in source def test_setup_captures_full_version_not_just_major_minor(self): """setup.ps1 must store the raw hipconfig output line, not just major.minor.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "ROCmVersionFull" in source def test_install_captures_full_version_not_just_major_minor(self): """install.ps1 must store the raw hipconfig output line, not just major.minor.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "ROCmVersionFull" in source def test_setup_uses_hip_path_or_rocm_path_for_sdk_display(self): """setup.ps1 HIP SDK path substep must check HIP_PATH then ROCM_PATH.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "HIP_PATH" in source and "ROCM_PATH" in source def test_install_uses_hip_path_or_rocm_path_for_sdk_display(self): """install.ps1 HIP SDK path substep must check HIP_PATH then ROCM_PATH.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "HIP_PATH" in source and "ROCM_PATH" in source def test_setup_rocm_step_uses_full_version(self): """setup.ps1 'rocm' step label must prefer the full version string.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "ROCmVersionFull" in source and "rocm" in source # TEST: install.sh -- Strix Halo rocm7.1 → rocm7.2 override _INSTALL_SH_PATH = PACKAGE_ROOT / "install.sh" _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh" class TestStrixRocm71Override: """Verify install.sh skips Radeon repo and routes to AMD arch-specific index for gfx1151/gfx1150 when ROCm 7.1 would otherwise be selected (known _grouped_mm segfault). AMD's repo.amd.com/rocm/whl/gfx1151/ serves torch 2.11+rocm7.13 which has the real fix.""" def test_strix_gfx_detection_in_install_sh(self): """install.sh must detect gfx1151 and gfx1150 for the override.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "gfx1151" in source and "gfx1150" in source def test_rocm71_override_to_amd_arch_index_in_install_sh(self): """install.sh must override TORCH_INDEX_URL to AMD arch-specific index for Strix.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") # The override must route to AMD's arch-specific index (repo.amd.com/rocm/whl) assert "repo.amd.com/rocm/whl" in source assert "_strix_gfx" in source # The URL must incorporate the detected gfx arch so gfx1151 → .../gfx1151/ strix_idx = source.find("_amd_strix_base") assert strix_idx != -1 ctx = source[strix_idx : strix_idx + 500] assert "_strix_gfx" in ctx def test_radeon_repo_bypassed_for_strix_in_install_sh(self): """install.sh must set _amd_gpu_radeon=false when Strix + ROCm 7.1 detected.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "_amd_gpu_radeon=false" in source def test_strix_override_warns_with_moe_utils_reference(self): """install.sh must emit a [WARN] mentioning the moe_utils segfault.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "moe_utils" in source or "_grouped_mm" in source def test_strix_override_only_fires_on_rocm71(self): """install.sh must scope the Strix override to rocm7.1 only (not rocm7.2+).""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") strix_idx = source.find("_strix_gfx") assert strix_idx != -1 # Look back for the rocm7.1 pattern within 600 chars before _strix_gfx context_before = source[max(0, strix_idx - 2400) : strix_idx] assert "rocm7.1" in context_before def test_torch_constraint_updated_for_strix_amd_index(self): """install.sh must set TORCH_CONSTRAINT>=2.11 when routing Strix to AMD index.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "TORCH_CONSTRAINT" in source and "2.11" in source def test_amd_rocm_mirror_env_var_respected(self): """install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "UNSLOTH_AMD_ROCM_MIRROR" in source def test_tauri_family_recognises_amd_arch_url(self): """_tauri_torch_index_family must return a rocm* family for AMD arch-specific URLs.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") # The function must have a case branch for repo.amd.com/rocm/whl/gfx* URLs assert "rocm/whl/gfx" in source # TEST: setup.sh -- gcc-install-dir fix for Ubuntu 24.04 + ROCm 7.x clang-20 class TestSetupShGccInstallDir: """Verify setup.sh applies the --gcc-install-dir flag when building llama.cpp with HIP on Ubuntu 24.04+ to work around ROCm 7.x clang-20 header path bug.""" def test_gcc_install_dir_search_loop_present(self): """setup.sh must iterate gcc versions 14→11 to find one with C++ headers.""" source = _SETUP_SH_PATH.read_text(encoding = "utf-8") assert "_GCC_INSTALL_DIR" in source assert "/usr/lib/gcc/x86_64-linux-gnu" in source def test_gcc_install_dir_checks_include_dir(self): """setup.sh must check that the gcc dir has an 'include' subdirectory.""" source = _SETUP_SH_PATH.read_text(encoding = "utf-8") assert "include" in source and "_GCC_INSTALL_DIR" in source def test_gcc_install_dir_appended_to_cmake_hip_flags(self): """setup.sh must pass --gcc-install-dir via CMAKE_HIP_FLAGS.""" source = _SETUP_SH_PATH.read_text(encoding = "utf-8") assert "CMAKE_HIP_FLAGS" in source assert "gcc-install-dir" in source def test_gcc_install_dir_only_applied_in_hip_build_block(self): """The --gcc-install-dir fix must only apply in the HIP/ROCm build branch.""" source = _SETUP_SH_PATH.read_text(encoding = "utf-8") # GGML_HIP=ON must appear before gcc-install-dir in the source hip_idx = source.find("GGML_HIP=ON") gcc_idx = source.find("gcc-install-dir") assert hip_idx != -1 and gcc_idx != -1 assert hip_idx < gcc_idx def test_gcc_install_dir_logs_substep(self): """setup.sh must print a substep when the gcc install dir is resolved.""" source = _SETUP_SH_PATH.read_text(encoding = "utf-8") 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 # ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ──────────────── # bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via subprocess PATH # at import time. The AMD torch wheel ships hipInfo.exe in the venv # Scripts dir, which is on PATH only for activated venvs -- Studio and the # installer launch python directly, so without the prepend every bnb # import logs "Could not detect ROCm GPU architecture: [WinError 2]". def test_main_py_prepends_hipinfo_dir_to_path(self): """main.py must make hipInfo.exe resolvable before bnb imports.""" source = _MAIN_PY_PATH.read_text(encoding = "utf-8") assert "hipInfo.exe" in source # The prepend must come before the BNB_ROCM_VERSION block (both run # pre-import; order documents that bnb sees the fixed PATH). assert source.find("hipInfo.exe") < source.find("BNB_ROCM_VERSION") def test_main_py_hipinfo_prepend_gated_on_file_presence(self): """Only AMD ROCm wheels ship hipInfo.exe; NVIDIA/CPU hosts must be untouched, so the prepend must check the file exists first.""" source = _MAIN_PY_PATH.read_text(encoding = "utf-8") assert 'os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe"))' in source def test_worker_py_prepends_hipinfo_dir_to_path(self): """worker.py must mirror the prepend for standalone-spawned workers.""" source = _WORKER_PATH.read_text(encoding = "utf-8") assert "hipInfo.exe" in source def test_install_stack_prepends_hipinfo_dir_to_path(self): """install_python_stack.py must prepend so the installer's child import checks inherit a PATH where bnb's probe succeeds.""" source = _STACK_PATH.read_text(encoding = "utf-8") assert "hipInfo.exe" 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 def test_hardware_py_stub_exposes_fake_process_group(self): """hardware.py stub must set FakeProcessGroup so torch.distributed doesn't raise AttributeError.""" source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") assert "FakeProcessGroup" in source def test_hardware_py_stub_exposes_process_group(self): """hardware.py stub must set ProcessGroup on the c10d stub.""" source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") assert "ProcessGroup" in source def test_hardware_py_stub_uses_setattr_for_symbols(self): """hardware.py must use setattr to populate stub symbols dynamically.""" source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") assert "setattr" in source def test_hardware_py_stub_all_c10d_siblings_covered(self): """hardware.py must stub all three torch._C._distributed_* submodules.""" source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") assert "_distributed_c10d" in source assert "_distributed_autograd" in source assert "_distributed_rpc" in source # TEST: install.ps1 / setup.ps1 -- HipSdkInstalled flag (SDK found, device inaccessible) class TestHipSdkInstalledButDeviceInaccessible: """Verify that when hipinfo is found but exits non-zero (device not ROCm-accessible), both scripts distinguish this from 'HIP SDK not found' and emit the correct message.""" def test_install_ps1_has_hip_sdk_installed_flag(self): """install.ps1 must track HipSdkInstalled separately from HasROCm.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "HipSdkInstalled" in source def test_setup_ps1_has_hip_sdk_installed_flag(self): """setup.ps1 must track HipSdkInstalled separately from HasROCm.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "HipSdkInstalled" in source def test_install_ps1_sets_flag_when_hipinfo_binary_found(self): """install.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") # HipSdkInstalled must be assigned inside the hipinfoExe block hipinfo_block_idx = source.find("if ($hipinfoExe)") sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx) assert hipinfo_block_idx != -1 and sdk_flag_idx != -1 assert sdk_flag_idx > hipinfo_block_idx def test_setup_ps1_sets_flag_when_hipinfo_binary_found(self): """setup.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") hipinfo_block_idx = source.find("if ($hipinfoExe)") sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx) assert hipinfo_block_idx != -1 and sdk_flag_idx != -1 assert sdk_flag_idx > hipinfo_block_idx def test_install_ps1_version_capture_runs_when_sdk_installed(self): """install.ps1 must capture hipconfig version when HipSdkInstalled even if HasROCm is false.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "HasROCm -or $HipSdkInstalled" in source or "$HipSdkInstalled" in source def test_setup_ps1_version_capture_runs_when_sdk_installed(self): """setup.ps1 must capture hipconfig version when HipSdkInstalled even if HasROCm is false.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "HasROCm -or $HipSdkInstalled" in source or "$HipSdkInstalled" in source def test_install_ps1_distinct_message_for_sdk_found_but_device_inaccessible(self): """install.ps1 must show 'not ROCm-accessible' message (not 'HIP SDK not found') when SDK present.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "not ROCm-accessible" in source def test_setup_ps1_distinct_message_for_sdk_found_but_device_inaccessible(self): """setup.ps1 must show 'not ROCm-accessible' message (not 'HIP SDK not found') when SDK present.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "not ROCm-accessible" in source def test_install_ps1_driver_guidance_in_sdk_found_branch(self): """install.ps1 must tell user this is a driver issue, not an SDK issue.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "driver issue" in source def test_setup_ps1_driver_guidance_in_sdk_found_branch(self): """setup.ps1 must tell user this is a driver issue, not an SDK issue.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "driver issue" in source def test_install_ps1_cpu_hint_distinguishes_driver_vs_no_sdk(self): """install.ps1 CPU-only hint must say 'GPU not ROCm-accessible' not 'require the HIP SDK' when SDK found.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") assert "GPU not ROCm-accessible" in source # TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 hand their resolved gfx arch # to install_llama_prebuilt.py so the lemonade HIP prebuilt is selected even when # the installer's own hipinfo/amd-smi probe cannot report it. _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh" class TestNormalizeForwardedGfx: """A forwarded gfx string is reduced to a single clean gfx token.""" def test_plain_token(self): assert _normalize_forwarded_gfx("gfx1151") == "gfx1151" def test_uppercase_normalized(self): assert _normalize_forwarded_gfx("GFX1151") == "gfx1151" def test_extracts_from_noise(self): assert _normalize_forwarded_gfx("gcnArchName: gfx942") == "gfx942" def test_malformed_is_ignored(self): assert _normalize_forwarded_gfx("not-a-gpu") is None def test_empty_and_none(self): assert _normalize_forwarded_gfx("") is None assert _normalize_forwarded_gfx(None) is None class TestApplyHostOverrides: """Forwarded ROCm detection is folded into the host profile correctly.""" def test_forwarded_gfx_fills_empty_probe(self): # amd-smi-only / name-inferred host: installer probe found no gfx. host = rocm_host(rocm_gfx_target = None) out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151") assert out.has_rocm is True assert out.rocm_gfx_target == "gfx1151" def test_forwarded_gfx_implies_rocm(self): # A CPU-looking host with a forwarded gfx is an AMD host. out = _apply_host_overrides(cpu_host(), override_rocm_gfx = "gfx1200") assert out.has_rocm is True assert out.rocm_gfx_target == "gfx1200" def test_forwarded_gfx_is_authoritative(self): # setup already applied visible-device selection; its value wins. host = rocm_host(rocm_gfx_target = "gfx1100") out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151") assert out.rocm_gfx_target == "gfx1151" def test_has_rocm_only_keeps_probe_gfx(self): out = _apply_host_overrides(cpu_host(), override_has_rocm = True) assert out.has_rocm is True assert out.rocm_gfx_target is None def test_malformed_forwarded_gfx_falls_back_to_has_rocm(self): out = _apply_host_overrides(cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk") assert out.has_rocm is True assert out.rocm_gfx_target is None def test_no_overrides_leaves_host_unchanged(self): host = nvidia_host() assert _apply_host_overrides(host) is host class TestRocmGfxForwarding: """setup.sh / setup.ps1 forward their resolved gfx; the installer accepts it.""" def test_installer_exposes_rocm_gfx_arg(self): source = _PREBUILT_PATH.read_text(encoding = "utf-8") assert '"--rocm-gfx"' in source # Defaults to the env override so a standalone run still works. assert 'os.environ.get("UNSLOTH_ROCM_GFX_ARCH")' in source def test_setup_sh_forwards_rocm_gfx(self): source = _SETUP_SH_PATH.read_text(encoding = "utf-8") assert "--rocm-gfx" in source assert '"$_setup_gfx"' in source def test_setup_sh_forwards_has_rocm(self): # When AMD is detected but gfx resolution fails, setup.sh must still # forward --has-rocm so the installer knows ROCm is present. source = _SETUP_SH_PATH.read_text(encoding = "utf-8") assert "--has-rocm" in source assert "_setup_amd_detected" in source def test_setup_ps1_forwards_rocm_gfx(self): source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "--rocm-gfx" in source assert "$script:ROCmGfxArch" in source # TEST: WSL ROCDXG fixes -- drop-in persistence + system-HIP-before-bundle _INSTALL_SH_PATH = PACKAGE_ROOT / "install.sh" _LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" class TestWslSystemRocmLibDirs: """install_llama_prebuilt._wsl_system_rocm_lib_dirs: a strict no-op off a ROCDXG WSL host; returns the system ROCm lib dir there so binary_env can load the WSL-capable HIP runtime before a prebuilt's bundled one.""" def test_empty_without_dev_dxg(self): with patch("os.path.exists", return_value = False): assert prebuilt_mod._wsl_system_rocm_lib_dirs() == [] def test_empty_on_bare_metal_linux(self): # /dev/dxg present but /proc/version is not a WSL kernel. with patch("os.path.exists", lambda p: p == "/dev/dxg"): with patch( "builtins.open", mock_open(read_data = "Linux version 6.8.0-generic"), ): assert prebuilt_mod._wsl_system_rocm_lib_dirs() == [] def test_returns_system_lib_on_wsl_with_librocdxg(self): # Normalize separators: os.path.join uses "\" on the Windows test host # though the target path is Linux. def _exists(p): p = str(p).replace("\\", "/") return p in ("/dev/dxg", "/opt/rocm/lib/librocdxg.so") with patch("os.path.exists", _exists): with patch( "builtins.open", mock_open(read_data = "Linux version 5.15.0-microsoft-standard-WSL2"), ): assert prebuilt_mod._wsl_system_rocm_lib_dirs() == ["/opt/rocm/lib"] def test_empty_on_wsl_without_librocdxg(self): # WSL kernel + /dev/dxg but no librocdxg -> not a ROCDXG ROCm install. with patch("os.path.exists", lambda p: p == "/dev/dxg"): with patch( "builtins.open", mock_open(read_data = "microsoft-standard-WSL2"), ): assert prebuilt_mod._wsl_system_rocm_lib_dirs() == [] class TestBinaryEnvWslOrdering: """binary_env must put the system ROCm lib dir AHEAD of the prebuilt's bundle dir on a ROCDXG WSL host, and set HSA_ENABLE_DXG_DETECTION; no-op on bare-metal Linux.""" @staticmethod def _linux_host(): return HostInfo( system = "Linux", machine = "x86_64", is_windows = False, is_linux = True, is_macos = False, is_x86_64 = True, is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], visible_cuda_devices = None, has_physical_nvidia = False, has_usable_nvidia = False, has_rocm = True, ) def test_wsl_prepends_system_rocm_and_sets_hsa(self, tmp_path): binary = tmp_path / "bundle" / "llama-server" binary.parent.mkdir(parents = True) binary.write_text("") # binary_env's dedupe_existing_dirs drops non-existent dirs, so use a # real dir to stand in for the system ROCm lib path. sys_rocm = tmp_path / "sysrocm" sys_rocm.mkdir() with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = [str(sys_rocm)]): with patch.dict(os.environ, {}, clear = True): env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) ld = env["LD_LIBRARY_PATH"].split(os.pathsep) # Compare resolved paths (dedupe_existing_dirs calls Path.resolve()). ld_resolved = [str(Path(p).resolve()) for p in ld] assert ld_resolved[0] == str(sys_rocm.resolve()) assert str(binary.parent.resolve()) in ld_resolved assert ld_resolved.index(str(sys_rocm.resolve())) < ld_resolved.index( str(binary.parent.resolve()) ) assert env.get("HSA_ENABLE_DXG_DETECTION") == "1" def test_bare_metal_linux_unchanged(self, tmp_path): binary = tmp_path / "bundle" / "llama-server" binary.parent.mkdir(parents = True) binary.write_text("") with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []): with patch.dict(os.environ, {}, clear = True): env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) ld = env["LD_LIBRARY_PATH"].split(os.pathsep) assert ld[0] == str(binary.parent) # bundle dir first, as before assert "HSA_ENABLE_DXG_DETECTION" not in env class TestInstallShDropinPersistence: """install.sh must persist the ROCm-on-WSL drop-in even when rocminfo already enumerates the GPU (via the transient probe env), so a reinstall over an existing /opt/rocm doesn't leave login shells without the env.""" def test_has_persist_helper(self): source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "_persist_rocm_wsl_dropin()" in source def test_gate5_early_return_persists_dropin(self): """The rocminfo-already-works early return must call the persist helper BEFORE returning.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") # Find the rocminfo gfx1151 gate and assert the persist call precedes # its `return 0`. gate = source.find("Name:[[:space:]]*gfx1151") assert gate != -1 window = source[gate : gate + 900] assert "_persist_rocm_wsl_dropin" in window assert window.find("_persist_rocm_wsl_dropin") < window.find("return 0") def test_persist_helper_gated_on_librocdxg(self): source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") body_start = source.find("_persist_rocm_wsl_dropin()") body = source[body_start : body_start + 1200] assert "librocdxg.so" in body assert "profile.d/unsloth-rocm-wsl.sh" in body class TestLlamaCppRuntimeWslOrdering: """The serve-time launcher must mirror binary_env: system HIP before the bundle dir on WSL, so a prebuilt that passed install validation runs.""" def test_has_wsl_helper(self): source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") assert "_wsl_system_rocm_lib_dirs" in source def test_prepends_before_binary_dir(self): source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") # The Linux lib_dirs build must add WSL rocm dirs before binary_dir. idx_helper = source.find("for _wsl_rocm in _wsl_system_rocm_lib_dirs()") idx_binary = source.find("lib_dirs.append(binary_dir)") assert idx_helper != -1 and idx_binary != -1 assert idx_helper < idx_binary if __name__ == "__main__": pytest.main([__file__, "-v"])