From 6bbc6925872d167d08640be1a9eac4e37fb38c5a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 6 Jul 2026 04:22:13 +0000 Subject: [PATCH] install: keep the torch-index marker additive to flavor validation Three narrow fixes in the marker-based stale-venv detection: - setup.ps1: a matching marker no longer overwrites the detected installed flavor. The marker compare is now an additional rebuild trigger, so a stale wheel (torch swapped to a +cpu build while the marker still records a cuXXX pin) is still caught by the flavor check instead of being masked as up to date. - setup.ps1: a supported AMD arch carrying CPU torch is no longer marked stale and wiped. The downstream AMD Windows ROCm override upgrades CPU torch to ROCm in place, so wiping first would delete the venv and abort with "Virtual environment not found". Only a genuinely wrong CUDA wheel still rebuilds. - install.sh: the Radeon --find-links path records its repo.radeon.com base in the marker instead of the generic pytorch.org ROCm fallback index, so a later pin to that generic family correctly reinstalls rather than comparing equal. Mirrors install.ps1/setup.ps1, which already record the real AMD index. --- install.sh | 14 ++- studio/setup.ps1 | 50 ++++---- tests/python/test_cross_platform_parity.py | 132 ++++++++++----------- tests/studio/install/test_cuda_repair.py | 12 +- tests/studio/install/test_rocm_support.py | 84 ++++++------- 5 files changed, 155 insertions(+), 137 deletions(-) diff --git a/install.sh b/install.sh index f748aee08c..d5498e7f7a 100755 --- a/install.sh +++ b/install.sh @@ -2996,6 +2996,15 @@ elif [ -n "$TORCH_INDEX_URL" ]; then --index-url "$TORCH_INDEX_URL" else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." + # Record the ACTUAL wheel source for the torch-index marker: this + # path installs from repo.radeon.com via --find-links, not from + # $TORCH_INDEX_URL (the generic pytorch.org ROCm fallback index). + # Recording $TORCH_INDEX_URL here would make the marker claim the + # generic index was used, so a later pin to that same ROCm family + # would compare-equal and skip the reinstall even though the wheels + # came from a different source. Mirrors install.ps1/setup.ps1, which + # record $ROCmIndexUrl (the real AMD index) for their AMD path. + _TORCH_MARKER_INDEX_URL="$_RADEON_BASE_URL" # Pass explicit wheel URLs so the matched trio is # installed together. --find-links lets uv discover # the Radeon listing for any local lookup, and PyPI @@ -3150,8 +3159,11 @@ fi # update` (install_python_stack.py / setup.ps1) can detect a later pin change by an # exact string compare rather than the version-tag heuristic. Only when torch was # actually installed from a resolved index (skip --no-torch / no-URL fallback). +# Reflects the actual source: the Radeon --find-links path sets +# _TORCH_MARKER_INDEX_URL to its repo.radeon.com base; every other path falls back +# to $TORCH_INDEX_URL (the CUDA/CPU/ROCm/pinned index it installed from). if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then - _write_torch_index_marker "$VENV_DIR" "$TORCH_INDEX_URL" + _write_torch_index_marker "$VENV_DIR" "${_TORCH_MARKER_INDEX_URL:-$TORCH_INDEX_URL}" fi # ── Run studio setup ── diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 4e28af3d93..7b0d09255a 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2744,26 +2744,20 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode $_expectedKnown = $true if ($_pinnedIdx) { $_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx - # Torch-index marker: when the last install recorded an index, compare the - # pin against it EXACTLY. This is the only signal that catches a per-arch - # switch between two 2.11 gfx indexes (gfx1151 -> gfx120X-all -- both - # install a +rocm7.13.0 wheel, so the tag heuristic sees no difference) and - # a custom-URL change (/simple -> /current). $null = no usable marker -> - # fall back to the version-tag heuristic (old venv; backward compatible). + # Torch-index marker compare: an EXACT compare of the pin against the index + # the last install recorded. This is the only signal that catches a per-arch + # switch between two 2.11 gfx indexes (gfx1151 -> gfx120X-all -- both install + # a +rocm7.13.0 wheel, so the tag heuristic sees no difference) and a + # custom-URL change (/simple -> /current). It is an ADDITIONAL rebuild + # trigger, NOT a substitute for validating the installed flavor: a matching + # marker must not hide a stale wheel (e.g. torch later swapped to a +cpu + # build while the marker still records the cuXXX pin), so the flavor check + # below still runs for known cu*/cpu/rocm leaves. $null = no usable marker -> + # flavor heuristic only (old venv; backward compatible). $_markerMismatch = Test-MarkerPinMismatch -VenvDir $VenvDir -PinUrl $_pinnedIdx - if ($null -ne $_markerMismatch) { - # Drive the rebuild decision purely off the marker compare. - $_expectedKnown = $true - if ($_markerMismatch) { - $expectedTorchTag = "pinned:$_pinnedIdx" - $installedTorchTag = "marker-mismatch" - } else { - $expectedTorchTag = "pinned:$_pinnedIdx" - $installedTorchTag = "pinned:$_pinnedIdx" - } - } + if ($_markerMismatch -eq $true) { $shouldRebuild = $true } # cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds. - elseif ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { + if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { # Do NOT collapse a pinned ROCm/gfx leaf to a generic "rocm": that # would match any installed +rocm wheel and mask a pin change from # one ROCm family to another (e.g. rocm6.4 -> gfx1151, or rocm6.4 @@ -2784,9 +2778,9 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode $expectedTorchTag = $_pinLeaf } else { # Custom index whose final segment is not a torch flavor (e.g. a - # PEP 503 mirror ending in /simple) and no marker to compare against. - # We cannot infer the flavor, so trust the pinned URL and do not - # rebuild on a bogus tag comparison. + # PEP 503 mirror ending in /simple). We cannot infer the flavor from + # the wheel tag, so the marker compare above is the only signal; do + # not also rebuild on a bogus tag comparison here. $_expectedKnown = $false $expectedTorchTag = $installedTorchTag } @@ -2812,7 +2806,19 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode "gfx90a", "gfx908" # MI200 / MI100 ) if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) { - $expectedTorchTag = "rocm" + # A correct +rocm wheel is not stale. A CPU wheel on a supported AMD + # arch is NOT wiped here either: the AMD Windows ROCm override below + # (the `$CuTag -eq "cpu"` reinstall block, plus the CPU-torch force in + # the version fast-path) upgrades CPU torch to ROCm in place. Forcing a + # rebuild would delete the venv and then hit "Virtual environment not + # found", so an older CPU-only install would be lost instead of + # repaired. Expect "cpu" for that case and let the override upgrade it. + # A genuinely wrong CUDA wheel (cu*) still mismatches "rocm" -> rebuild. + if ($installedTorchTag -eq "cpu") { + $expectedTorchTag = "cpu" + } else { + $expectedTorchTag = "rocm" + } } else { $expectedTorchTag = "cpu" } diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index d3c1f02709..c3d101ee30 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -47,9 +47,9 @@ class TestNoTorchBackendAutoInInstallSh: def test_fallback_uses_torch_backend_auto(self): """The fallback branch should use --torch-backend=auto as recovery.""" text = INSTALL_SH.read_text(encoding = "utf-8") - assert ( - "GPU detection failed" in text - ), "install.sh should have a fallback branch for when GPU detection fails" + assert "GPU detection failed" in text, ( + "install.sh should have a fallback branch for when GPU detection fails" + ) class TestInstallShHasGpuDetection: @@ -57,15 +57,15 @@ class TestInstallShHasGpuDetection: def test_function_exists(self): text = INSTALL_SH.read_text(encoding = "utf-8") - assert ( - "get_torch_index_url()" in text - ), "install.sh is missing the get_torch_index_url() function" + assert "get_torch_index_url()" in text, ( + "install.sh is missing the get_torch_index_url() function" + ) def test_torch_index_url_assigned(self): text = INSTALL_SH.read_text(encoding = "utf-8") - assert ( - "TORCH_INDEX_URL=$(get_torch_index_url)" in text - ), "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()" + assert "TORCH_INDEX_URL=$(get_torch_index_url)" in text, ( + "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()" + ) class TestCudaMappingParity: @@ -133,15 +133,15 @@ class TestPyTorchMirrorEnvVar: def test_install_sh_has_mirror_var(self): text = INSTALL_SH.read_text(encoding = "utf-8") - assert ( - "UNSLOTH_PYTORCH_MIRROR" in text - ), "install.sh should reference UNSLOTH_PYTORCH_MIRROR" + assert "UNSLOTH_PYTORCH_MIRROR" in text, ( + "install.sh should reference UNSLOTH_PYTORCH_MIRROR" + ) def test_install_ps1_has_mirror_var(self): text = INSTALL_PS1.read_text(encoding = "utf-8") - assert ( - "UNSLOTH_PYTORCH_MIRROR" in text - ), "install.ps1 should reference UNSLOTH_PYTORCH_MIRROR" + assert "UNSLOTH_PYTORCH_MIRROR" in text, ( + "install.ps1 should reference UNSLOTH_PYTORCH_MIRROR" + ) class TestUvBytecodeCompileTimeout: @@ -167,21 +167,21 @@ class TestUvBytecodeCompileTimeout: def test_install_sh_preserves_timeout_override(self): text = INSTALL_SH.read_text(encoding = "utf-8") - assert ( - ': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text - ), "install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers" - assert ( - "export UV_COMPILE_BYTECODE_TIMEOUT" in text - ), "install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses" + assert ': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text, ( + "install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers" + ) + assert "export UV_COMPILE_BYTECODE_TIMEOUT" in text, ( + "install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses" + ) def test_install_ps1_preserves_timeout_override(self): text = INSTALL_PS1.read_text(encoding = "utf-8") - assert ( - "if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text - ), "install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides" - assert ( - '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text - ), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT" + assert "if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text, ( + "install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides" + ) + assert '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text, ( + "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT" + ) class TestTorchIndexOverrideParity: @@ -207,9 +207,9 @@ class TestTorchIndexOverrideParity: # The AMD ROCm reroute must be skipped when the index is explicitly pinned, # so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten. text = path.read_text(encoding = "utf-8") - assert ( - "TorchIndexPinned" in text - ), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag" + assert "TorchIndexPinned" in text, ( + f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag" + ) def test_cuda_pin_overrides_cvd_hide_gate(self): # A pinned cu* index skips ALL host-GPU probing (parity with install.sh's @@ -226,9 +226,9 @@ class TestTorchIndexOverrideParity: "_ensure_cuda_torch should compute a CUDA-pin flag so the pin can " "override the CVD hide gate" ) - assert re.search( - r"if not _cuda_pinned and _cvd is not None", body - ), "the CVD hide gate must be bypassed when a CUDA index is pinned" + assert re.search(r"if not _cuda_pinned and _cvd is not None", body), ( + "the CVD hide gate must be bypassed when a CUDA index is pinned" + ) def test_cpu_repair_pins_supported_torch_range(self): # The explicit-CPU repair must not install a bare torch trio: the /cpu @@ -281,12 +281,12 @@ class TestGfx211AllowlistParity: # install-spec path must reuse it, so the stale check and install spec can # never disagree again. text = SETUP_PS1.read_text(encoding = "utf-8") - assert ( - "function Test-RocmGfx211Leaf" in text - ), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper" - assert re.search( - r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower() - ), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist" + assert "function Test-RocmGfx211Leaf" in text, ( + "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper" + ) + assert re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()), ( + "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist" + ) assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, ( "setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not " "re-hardcode the allowlist (they must not diverge)" @@ -294,9 +294,9 @@ class TestGfx211AllowlistParity: def test_stack_py_allowlist(self): text = STACK_PY.read_text(encoding = "utf-8").lower() - assert ( - '"gfx120x-all", "gfx1151", "gfx1150"' in text - ), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed" + assert '"gfx120x-all", "gfx1151", "gfx1150"' in text, ( + "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed" + ) class TestCudaLeafDigitParity: @@ -307,32 +307,32 @@ class TestCudaLeafDigitParity: def test_stack_py_requires_cu_digit(self): text = STACK_PY.read_text(encoding = "utf-8") - assert re.search( - r'r"\^cu\[0-9\]"', text - ), "install_python_stack.py _is_cuda_family_leaf must match ^cu[0-9]" + assert re.search(r'r"\^cu\[0-9\]"', text), ( + "install_python_stack.py _is_cuda_family_leaf must match ^cu[0-9]" + ) def test_setup_ps1_requires_cu_digit(self): text = SETUP_PS1.read_text(encoding = "utf-8") - assert re.search( - r"'\^cu\[0-9\]'", text - ), "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9], not a bare cu* glob" + assert re.search(r"'\^cu\[0-9\]'", text), ( + "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9], not a bare cu* glob" + ) # The stale-venv branch must go through the digit-guarded helper. - assert ( - "Test-CudaFamilyLeaf $_pinLeaf" in text - ), "setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf" + assert "Test-CudaFamilyLeaf $_pinLeaf" in text, ( + "setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf" + ) def test_install_ps1_requires_cu_digit_in_gpu_branch(self): text = INSTALL_PS1.read_text(encoding = "utf-8") - assert re.search( - r"'\^cu\[0-9\]'", text - ), "install.ps1 Get-TauriGpuBranch must require a digit after cu" + assert re.search(r"'\^cu\[0-9\]'", text), ( + "install.ps1 Get-TauriGpuBranch must require a digit after cu" + ) def test_install_sh_requires_cu_digit_in_gpu_branch(self): text = INSTALL_SH.read_text(encoding = "utf-8") # The _tauri_gpu_branch cuda case must be cu[0-9]*, not a bare cu*. - assert re.search( - r"cu\[0-9\]\*\)\s*echo \"cuda\"", text - ), "install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*" + assert re.search(r"cu\[0-9\]\*\)\s*echo \"cuda\"", text), ( + "install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*" + ) def test_install_sh_backend_export_requires_cu_digit(self): text = INSTALL_SH.read_text(encoding = "utf-8") @@ -340,13 +340,13 @@ class TestCudaLeafDigitParity: # bare catch-all *) -> cuda would mis-brand /current, /custom mirror pins # as CUDA and make the stack skip ROCm repair on AMD hosts (comment #2's # bug via install.sh instead of standalone studio update). - assert re.search( - r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text - ), "install.sh backend export must brand cuda only on cu[0-9]*" + assert re.search(r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text), ( + "install.sh backend export must brand cuda only on cu[0-9]*" + ) # An unknown leaf must NOT commit a cuda backend (it unsets instead). - assert re.search( - r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text - ), "install.sh backend export must unset (not force cuda) on an unknown leaf" + assert re.search(r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text), ( + "install.sh backend export must unset (not force cuda) on an unknown leaf" + ) def test_install_sh_lowercases_backend_leaf(self): text = INSTALL_SH.read_text(encoding = "utf-8") @@ -417,9 +417,9 @@ class TestKnown211SetParity: def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self): text = INSTALL_SH.read_text(encoding = "utf-8") # The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves. - assert re.search( - r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text - ), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150" + assert re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text), ( + "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150" + ) # No speculative rocm7.3 anywhere. assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3" diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index 8f064f1a03..cb37a29d28 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -359,12 +359,12 @@ class TestTorchBackendDerivationFromPin: def test_source_uses_helper_not_bare_startswith(self): # Guard against a regression back to elif _idx_leaf.startswith("cu"). src = _STACK_PATH.read_text(encoding = "utf-8") - assert ( - "elif _is_cuda_family_leaf(_idx_leaf):" in src - ), "_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf" - assert ( - 'elif _idx_leaf.startswith("cu"):' not in src - ), "_TORCH_BACKEND derivation must not use a bare startswith('cu')" + assert "elif _is_cuda_family_leaf(_idx_leaf):" in src, ( + "_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf" + ) + assert 'elif _idx_leaf.startswith("cu"):' not in src, ( + "_TORCH_BACKEND derivation must not use a bare startswith('cu')" + ) # CUDA index ladder. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 10d198fd85..79c3c2d45f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -825,9 +825,9 @@ class TestEnsureRocmTorch: _args = [str(a) for a in _call.args] if "--index-url" in _args: _url = _args[_args.index("--index-url") + 1] - assert "rocm7.2" not in _url or "torch" not in " ".join( - _args - ), "torch must not be reinstalled when the pin already matches" + assert "rocm7.2" not in _url or "torch" not in " ".join(_args), ( + "torch must not be reinstalled when the pin already matches" + ) # A torch reinstall would pass torch>=... as a positional; assert none did. assert not any( any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list @@ -1294,9 +1294,9 @@ class TestHasRocmGpuKfdVendorGuard: src = self._src() # Word boundary 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" + 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').""" @@ -1306,9 +1306,9 @@ class TestHasRocmGpuKfdVendorGuard: 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)" + 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.""" @@ -1341,12 +1341,12 @@ class TestHasRocmGpuKfdVendorGuard: 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" + 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.""" @@ -1562,12 +1562,12 @@ class TestInstallShStructure: 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" + 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.""" @@ -1633,9 +1633,9 @@ class TestInstallShStructure: 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 (( ))" + 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).""" @@ -1654,9 +1654,9 @@ class TestInstallShStructure: 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" + assert backend_pos > torch_url_pos, ( + "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved" + ) 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] @@ -1670,12 +1670,12 @@ class TestInstallShStructure: 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)" + 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) to avoid Ryzen+NVIDIA false positives.""" @@ -1700,9 +1700,9 @@ class TestInstallShStructure: "get_torch_index_url must use a _nvidia_detected flag (separate from " "_smi) so that proc-only NVIDIA detection still selects CUDA wheels" ) - 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" + 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) @@ -2691,9 +2691,9 @@ class TestRuntimeBnbRocmSourceGuards: """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 + 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 @@ -3142,9 +3142,9 @@ class TestStrixHaloGfxArchDetection: """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") - assert ( - "gfx\\d+" in source or r"gfx\d+" in source - ), f"gfx arch regex not found in {path.name}" + 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