From 1e3b1c97ad45e76e9644c01a13803c0d12d0e690 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 15:50:14 +0000 Subject: [PATCH] install_python_stack.py: keep the ROCm 7.1 repair on the 2.11 line The rocm7.1 leaf fell through to _ROCM_TORCH_PKG_SPECS["_default"], which caps the trio below 2.11. Now that install.sh leaves a rocm7.1 leaf on the widened default range, a fresh install resolves torch 2.11.0+rocm7.1 while the later dependency pass force-reinstalled 2.10.0+rocm7.1 over it, so any repair or `studio update` silently downgraded the environment. download.pytorch.org/whl/rocm7.1 serves a paired 2.11 trio, verified with uv pip compile --no-deps against that index: install.sh default range -> torch 2.11.0+rocm7.1, torchvision 0.26.0+rocm7.1, torchaudio 2.11.0+rocm7.1 _default repair spec -> torch 2.10.0+rocm7.1, torchvision 0.25.0+rocm7.1, torchaudio 2.10.0+rocm7.1 Give rocm7.1 its own entry carrying install.sh's default range rather than the rocm7.2 tuple: only the _grouped_mm arches take the hard 2.11 floor, so _ROCM_KNOWN_TORCH211_VERSIONS stays {(7, 2)}. _default keeps its literal <2.11 ceiling because rocm7.0 and older genuinely top out below it (rocm7.0 at 2.10.0, rocm6.4 and rocm6.3 at 2.9.1, rocm6.2 at 2.5.1), and the stale index comments are corrected to match what those indexes serve today. --- studio/install_python_stack.py | 20 ++++++++++++++---- tests/studio/install/test_rocm_support.py | 25 ++++++++++++++++++++++- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index fb450a4b54..57c647d11a 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -58,8 +58,8 @@ PLATFORM_LACKS_TORCHCODEC_WHEEL = ( # Detected ROCm (major, minor) -> best PyTorch wheel tag, checked newest-first (>=). _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { (7, 2): "rocm7.2", # torch 2.11.0 - (7, 1): "rocm7.1", # torch 2.10.0 - (7, 0): "rocm7.0", + (7, 1): "rocm7.1", # torch 2.11.0 + (7, 0): "rocm7.0", # torch 2.10.0 (6, 4): "rocm6.4", (6, 3): "rocm6.3", (6, 2): "rocm6.2", @@ -97,14 +97,26 @@ _ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset( # rocmX.Y indexes KNOWN to ship torch 2.11; never floor an unknown newer rocm speculatively. _ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)}) -# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x). +# Per-tag pip specs for the repair/update path; must land on the same wheels a fresh +# install.sh run would pick, otherwise `studio update` silently downgrades the venv. _ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { + # Floored at 2.11 (the _grouped_mm bug), matching install.sh's rocm7.2|gfx* case. "rocm7.2": ( "torch>=2.11.0,<2.12.0", "torchvision>=0.26.0,<0.27.0", "torchaudio>=2.11.0,<2.12.0", ), - # rocm7.1 and earlier: torch 2.x below 2.11 + # rocm7.1 also serves a paired 2.11 trio (torch 2.11.0 / torchvision 0.26.0 / + # torchaudio 2.11.0), so it takes install.sh's widened DEFAULT range rather than + # the 2.11 floor: no _grouped_mm floor applies here, but capping at <2.11 would + # force-reinstall 2.10 over the 2.11 a fresh install just resolved. + "rocm7.1": ( + "torch>=2.4,<2.12.0", + "torchvision>=0.19,<0.27.0", + "torchaudio>=2.4,<2.12.0", + ), + # rocm7.0 and earlier genuinely top out below 2.11 (rocm7.0: torch 2.10.0, + # rocm6.4/6.3: 2.9.1, rocm6.2: 2.5.1), so the old ceiling stays literal. "_default": ( "torch>=2.4,<2.11.0", "torchvision>=0.19,<0.26.0", diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4f65857a3e..f8d2d4e374 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3457,12 +3457,35 @@ class TestRocmTorchPkgSpecs: assert "2.11" in torch_spec def test_default_caps_below_211(self): - """Default spec (rocm7.1 and earlier) should cap below 2.11.""" + """Default spec (rocm7.0 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_rocm71_repair_matches_install_sh_default_range(self): + """rocm7.1 serves a paired 2.11 trio, so the repair path must not cap at <2.11. + + install.sh leaves a rocm7.1 leaf on its default trio (torch>=2.4,<2.12.0 / + torchvision>=0.19,<0.27.0 / torchaudio>=2.4,<2.12.0), which resolves + torch 2.11.0+rocm7.1 on that index. Falling back to _default here would + force-reinstall 2.10.0+rocm7.1 over it on the next `studio update`. + """ + specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("rocm7.1") + assert specs is not None, "rocm7.1 must have its own repair spec" + assert specs == ( + "torch>=2.4,<2.12.0", + "torchvision>=0.19,<0.27.0", + "torchaudio>=2.4,<2.12.0", + ) + # Not the rocm7.2 spec: no 2.11 floor applies to rocm7.1. + assert specs != stack_mod._ROCM_TORCH_PKG_SPECS["rocm7.2"] + + def test_rocm71_is_not_a_known_211_floor_version(self): + """The widened rocm7.1 range must NOT promote it to a floored 2.11 line.""" + assert (7, 1) not in stack_mod._ROCM_KNOWN_TORCH211_VERSIONS + assert (7, 2) in stack_mod._ROCM_KNOWN_TORCH211_VERSIONS + 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():