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