Ship a single ROCm version to the nightly getting-started matrix

The getting-started page renders one ROCm choice. gen_quick_start_module.py
picks it with max() over version strings, and "7.14" < "7.2" lexicographically,
so sending both made pytorch.org advertise ROCm 7.2 as the nightly ROCm after
7.14 landed.

Send only the newest, chosen with a version-aware key. Scoped to
getting-started + nightly: the binary build matrix is untouched, so nightly
still builds every ROCm in ROCM_ARCHES_DICT.

Author: atalman

tools/scripts/generate_binary_build_matrix.py

@@ -131,6 +131,10 @@
 DISABLE = "disable"
 
 
+def parse_version(version: str) -> Tuple[int, ...]:
+    return tuple(int(part) for part in version.split("."))
+
+
 def arch_type(arch_version: str) -> str:
     if arch_version in CUDA_ARCHES:
         return CUDA
@@ -182,6 +186,8 @@ def initialize_globals(
 
     CUDA_ARCHES = CUDA_ARCHES_DICT[channel]
     ROCM_ARCHES = ROCM_ARCHES_DICT[channel]
+    if getting_started and channel == NIGHTLY:
+        ROCM_ARCHES = [max(ROCM_ARCHES, key=parse_version)]
     if build_python_only:
         # Only select the oldest version of python if building a python only package
         PYTHON_ARCHES = [PYTHON_ARCHES_DICT[channel][0]]

tools/tests/test_generate_binary_build_matrix.py

@@ -4,7 +4,11 @@
 import sys
 from unittest import main, TestCase
 
-from tools.scripts.generate_binary_build_matrix import generate_build_matrix
+from tools.scripts.generate_binary_build_matrix import (
+    generate_build_matrix,
+    parse_version,
+    ROCM_ARCHES_DICT,
+)
 
 
 ASSETS_DIR = "tools/tests/assets"
@@ -200,6 +204,46 @@ def test_torch_only_install_command_for_torch_only_arches(self):
                 self.assertNotIn("torchvision", entry["installation"])
                 self.assertIn("torch", entry["installation"])
 
+    def _rocm_versions(self, channel: str, getting_started: str) -> set:
+        out = generate_build_matrix(
+            "wheel",
+            "linux",
+            channel,
+            "enable",
+            "enable",
+            "enable",
+            "enable",
+            "false",
+            "false",
+            "disable",
+            getting_started,
+            None,
+        )
+        return {
+            entry["gpu_arch_version"]
+            for entry in out["include"]
+            if entry["gpu_arch_type"] == "rocm"
+        }
+
+    def test_parse_version_orders_double_digit_minors(self):
+        self.assertGreater(parse_version("7.14"), parse_version("7.2"))
+        self.assertEqual(
+            max(["7.2", "7.14"], key=parse_version),
+            "7.14",
+        )
+
+    def test_getting_started_nightly_ships_one_rocm(self):
+        versions = self._rocm_versions("nightly", "true")
+        self.assertEqual(len(versions), 1)
+        self.assertEqual(
+            versions, {max(ROCM_ARCHES_DICT["nightly"], key=parse_version)}
+        )
+
+    def test_nightly_builds_keep_every_rocm(self):
+        versions = self._rocm_versions("nightly", "false")
+        self.assertEqual(versions, set(ROCM_ARCHES_DICT["nightly"]))
+        self.assertGreater(len(versions), 1)
+
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Test generate build matrix")