Fix ROCm/gfx pin case normalization, ROCm-tag requirement, and CUDA-leaf classification

Normalize torch-index leaves to lowercase before the gfx*/rocm*/cu* allowlist
matches so the canonical gfx120X-all (capital X) gets the torch 2.11 floor in
install.sh (leaf, flavor and repairable helpers). Require an installed +rocm
local tag before a rocmX.Y or non-2.11 gfx pin is judged satisfied in
setup.ps1 Get-RocmPinStaleTags and the Python _rocm_pin_family_mismatch, so an
untagged CPU/CUDA wheel never leaves the pin unapplied. Classify a leaf as CUDA
only via ^cu[0-9]: the Python _TORCH_BACKEND derivation now uses
_is_cuda_family_leaf, and install.sh brands cuda only on cu[0-9]* (unset on an
unknown /current /custom mirror leaf) so the stack probes the GPU instead of
skipping ROCm repair. Add bash, Python and PowerShell tests for capital
gfx120X-all floor, current/custom not-cuda, and untagged-wheel ROCm pins.

studio/install_python_stack.py

@@ -1119,6 +1119,12 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
     _inst_rocm = re.search(r"\+rocm(\d+)\.(\d+)", installed_ver)
     _inst_ver = (int(_inst_rocm.group(1)), int(_inst_rocm.group(2))) if _inst_rocm else None
     _inst_is_perarch = re.search(r"\+rocm\d+\.\d+\.\d+", installed_ver) is not None
+    # A ROCm build MUST carry a +rocm local tag. An untagged CPU/CUDA wheel (no
+    # +rocm, e.g. "2.10.0" / "2.11.0") never satisfies a ROCm pin -- always a
+    # mismatch -- mirroring setup.ps1's Get-RocmPinStaleTags. (In practice
+    # _ensure_rocm_torch only calls this when has_hip_torch is True, but keep the
+    # pure function correct for any input so it stays in lockstep with the PS side.)
+    _inst_has_rocm = re.search(r"\+rocm", installed_ver) is not None
     # Whether the installed torch RELEASE (before "+") is 2.11+.
     _inst_rel = re.match(r"^(\d+)\.(\d+)", installed_ver)
     _inst_is_211 = (
@@ -1135,8 +1141,9 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
             # a generic rocm wheel or any pre-2.11 build IS a mismatch even at 2.11.
             return not (_inst_is_211 and _inst_is_perarch)
         # Non-2.11 gfx leaf: install path uses default <2.11 specs, so a correct
-        # <2.11 wheel must stay. Mismatch only when the installed torch is 2.11+.
-        return _inst_is_211
+        # <2.11 wheel must stay. An untagged (no +rocm) wheel never satisfies the
+        # pin -> mismatch; otherwise mismatch only when the installed torch is 2.11+.
+        return (not _inst_has_rocm) or _inst_is_211
 
     # rocmX.Y pin.
     _pin_is_211 = _pin_ver >= (7, 2) if _pin_ver is not None else False
@@ -1146,7 +1153,10 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
         # (7, 13) -> mismatch, which correctly reinstalls the generic wheel the
         # user pinned instead of leaving the per-arch one in place.
         return _pin_ver != _inst_ver
-    # rocm pin with an unreadable installed version: compare on the torch 2.11 line.
+    # rocm pin with an unreadable installed version: compare on the torch 2.11 line,
+    # but an untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> mismatch.
+    if not _inst_has_rocm:
+        return True
     return _pin_is_211 != _inst_is_211
 
 
@@ -1707,7 +1717,7 @@ def _ensure_rocm_torch() -> None:
                 None,
             )
         if tag is None:
-            print(f"   No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall")
+            print(f"   No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping torch reinstall")
         else:
             if _override_idx is None:
                 index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
@@ -1833,7 +1843,12 @@ if not _TORCH_BACKEND:
         _TORCH_BACKEND = "rocm"
     elif _idx_leaf == "cpu":
         _TORCH_BACKEND = "cpu"
-    elif _idx_leaf.startswith("cu"):
+    elif _is_cuda_family_leaf(_idx_leaf):
+        # Require a digit after "cu" (^cu[0-9]) so a full-override URL ending in
+        # /current or /custom is NOT branded CUDA. A wrong "cuda" backend makes
+        # _ensure_rocm_torch() return early on AMD hosts and leaves a CPU/wrong
+        # torch unrepaired; falling through here keeps _TORCH_BACKEND="" so the
+        # helpers probe the GPU instead.
         _TORCH_BACKEND = "cuda"
 
 

studio/setup.ps1

@@ -457,6 +457,9 @@ function Test-CudaFamilyLeaf {
 #     is stale; an already-installed per-arch wheel is NOT (no rebuild loop).
 #   * gfx pin NOT in the allowlist (gfx110X-all/gfx90a/gfx908) -> default <2.11
 #     specs, so a 2.10+rocm wheel is correct; only a 2.11+ build is stale.
+# A ROCm pin (gfx* or rocmX.Y) is satisfied ONLY by an installed wheel carrying a
+# +rocm local tag: an untagged CPU/CUDA wheel (e.g. 2.10.0 / 2.11.0) never
+# satisfies a ROCm pin, so it is reported stale and the pin is (re)applied.
 function Get-RocmPinStaleTags {
     param([string]$PinLeaf, [string]$TorchVersion)
     $_pinRocm = [regex]::Match($PinLeaf, '^rocm(\d+)\.(\d+)')
@@ -465,6 +468,9 @@ function Get-RocmPinStaleTags {
     $_instRocm = [regex]::Match($TorchVersion, '\+rocm(\d+)\.(\d+)')
     $_instVer = if ($_instRocm.Success) { "$($_instRocm.Groups[1].Value).$($_instRocm.Groups[2].Value)" } else { $null }
     $_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+')
+    # A ROCm build MUST carry a +rocm local tag. Without it the wheel is a CPU/CUDA
+    # build that cannot satisfy any ROCm pin, regardless of its release line.
+    $_instHasRocm = [regex]::IsMatch($TorchVersion, '\+rocm')
     $_instRel = [regex]::Match($TorchVersion, '^(\d+)\.(\d+)')
     $_instIs211 = $false
     if ($_instRel.Success) {
@@ -474,29 +480,36 @@ function Get-RocmPinStaleTags {
     if ($PinLeaf -like 'gfx*') {
         if (Test-RocmGfx211Leaf $PinLeaf) {
             # Expect the AMD per-arch (three-part) 2.11 wheel. Satisfied only when
-            # BOTH a 2.11 release AND a three-part rocm tag are installed.
+            # BOTH a 2.11 release AND a three-part rocm tag are installed (the
+            # three-part tag already implies +rocm).
             $installed = if ($_instIs211 -and $_instPerArch) { "rocm-perarch(torch>=2.11)" } else { "rocm-generic-or-old" }
             return @{ Expected = "rocm-perarch(torch>=2.11)"; Installed = $installed }
         }
-        # Non-2.11 gfx leaf: default <2.11 spec. Stale only when the build is 2.11+.
+        # Non-2.11 gfx leaf: default <2.11 spec. An untagged (no +rocm) wheel never
+        # satisfies the pin -> stale. Otherwise stale only when the build is 2.11+.
+        $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
         return @{
             Expected  = "rocm(torch<2.11)"
-            Installed = if ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
+            Installed = $installed
         }
     }
 
     # rocmX.Y pin.
     if ($_pinVer -and $_instVer) {
-        # Both rocm versions readable: exact comparison.
+        # Both rocm versions readable: exact comparison. A readable $_instVer already
+        # implies a +rocm tag, so no separate tag check is needed here.
         return @{ Expected = "rocm$_pinVer"; Installed = "rocm$_instVer" }
     }
     $_pinNeeds211 = $false
     if ($_pinRocm.Success) {
         $_pinNeeds211 = ([int]$_pinRocm.Groups[1].Value -gt 7) -or ([int]$_pinRocm.Groups[1].Value -eq 7 -and [int]$_pinRocm.Groups[2].Value -ge 2)
     }
+    # Fallback (installed rocm version unreadable): compare on the 2.11 line, but an
+    # untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> report it stale.
+    $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
     return @{
         Expected  = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
-        Installed = if ($_instIs211)   { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
+        Installed = $installed
     }
 }