Add option to generate matrix for getting started page (#6971)

atalman · web-flow · commit bb05b7916e63 · 2025-08-05T12:53:46.000-04:00
We will be removing torchaudio from getting started page. However would still like to include it in validation. Test: ``` python tools/scripts/generate_binary_build_matrix.py --getting-started true {"include": [{"python_version": "3.9", "gpu_arch_type": "cpu", "gpu_arch_version": "", "desired_cuda": "cpu", "container_image": "pytorch/manylinux2_28-builder:cpu", "package_type": "manywheel", "build_name": "manywheel-py3_9-cpu", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.1", "use_split_build": false}, {"python_version": "3.9", "gpu_arch_type": "cuda", "gpu_arch_version": "12.6", "desired_cuda": "cu126", "container_image": "pytorch/manylinux2_28-builder:cuda12.6", "package_type": "manywheel", "build_name": "manywheel-py3_9-cuda12_6", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu126", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.1", "use_split_build": false}, {"python_version": "3.9", "gpu_arch_type": "cuda", "gpu_arch_version": "12.8", "desired_cuda": "cu128", "container_image": "pytorch/manylinux2_28-builder:cuda12.8", "package_type": "manywheel", "build_name": "manywheel-py3_9-cuda12_8", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.1", "use_split_build": false}, {"python_version": "3.9", "gpu_arch_type": "cuda", "gpu_arch_version": "12.9", "desired_cuda": "cu129", "container_image": "pytorch/manylinux2_28-builder:cuda12.9", "package_type": "manywheel", "build_name": "manywheel-py3_9-cuda12_9", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu129", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.1", "use_split_build": false}, {"python_version": "3.9", "gpu_arch_type": "rocm", "gpu_arch_version": "6.3", "desired_cuda": "rocm6.3", "container_image": "pytorch/manylinux2_28-builder:rocm6.3", "package_type": "manywheel", "build_name": "manywheel-py3_9-rocm6_3", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.3", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.1", "use_split_build": false}, {"python_version": "3.9", "gpu_arch_type": "rocm", "gpu_arch_version": "6.4", "desired_cuda": "rocm6.4", "container_image": "pytorch/manylinux2_28-builder:rocm6.4", "package_type": "manywheel", "build_name": "manywheel-py3_9-rocm6_4", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.4", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.1", "use_split_build": false}, {"python_version": "3.9", "gpu_arch_type": "xpu", "gpu_arch_version": "", "desired_cuda": "xpu", "container_image": "pytorch/manylinux2_28-builder:xpu", "package_type": "manywheel", "build_name": "manywheel-py3_9-xpu", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/xpu", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.1", "use_split_build": false}]} ```
diff --git a/.github/workflows/generate_binary_build_matrix.yml b/.github/workflows/generate_binary_build_matrix.yml
@@ -59,6 +59,11 @@ on:
         required: false
         type: boolean
         default: false
+      getting-started:
+        description: "Release matrix for getting started page"
+        required: false
+        type: boolean
+        default: false
 
     outputs:
       matrix:
@@ -98,6 +103,7 @@ jobs:
           USE_ONLY_DL_PYTORCH_ORG: ${{ inputs.use-only-dl-pytorch-org }}
           BUILD_PYTHON_ONLY: ${{ inputs.build-python-only }}
           USE_SPLIT_BUILD: ${{ inputs.use_split_build }}
+          GETTING_STARTED: ${{ inputs.getting-started }}
           PYTHON_VERSIONS: ${{ inputs.python-versions }}
         run: |
           set -eou pipefail
diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py
@@ -99,6 +99,7 @@
 MACOS_M1_RUNNER = "macos-m1-stable"
 
 PACKAGES_TO_INSTALL_WHL = "torch torchvision torchaudio"
+PACKAGES_TO_INSTALL_GETTING_STARTED_WHL = "torch torchvision"
 PACKAGES_TO_INSTALL_WHL_WIN_ARM64 = "torch"
 WHL_INSTALL_BASE = "pip3 install"
 DOWNLOAD_URL_BASE = "https://download.pytorch.org"
@@ -274,10 +275,16 @@ def get_wheel_install_command(
     python_version: str,
     use_only_dl_pytorch_org: bool,
     use_split_build: bool = False,
+    getting_started: bool = False,
 ) -> str:
+    PACKAGES_TO_INSTALL = (
+        PACKAGES_TO_INSTALL_GETTING_STARTED_WHL
+        if getting_started
+        else PACKAGES_TO_INSTALL_WHL
+    )
     if use_split_build:
         if (gpu_arch_version in CUDA_ARCHES) and (os == LINUX) and (channel == NIGHTLY):
-            return f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}_pypi_pkg"  # noqa: E501
+            return f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}_pypi_pkg"  # noqa: E501
         else:
             raise ValueError(
                 "Split build is not supported for this configuration. It is only supported for CUDA 11.8, 12.4, 12.6 on Linux nightly builds."  # noqa: E501
@@ -291,7 +298,7 @@ def get_wheel_install_command(
             or (os == LINUX_AARCH64)
         )
     ):
-        return f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL}"
+        return f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL}"
     else:
         whl_install_command = ""
         if os == WINDOWS_ARM64:
@@ -300,9 +307,9 @@ def get_wheel_install_command(
                 f"{WHL_INSTALL_BASE} --pre {PACKAGES_TO_INSTALL_WHL_WIN_ARM64}"  # noqa: E501
             )
         elif channel == "nightly":
-            whl_install_command = f"{WHL_INSTALL_BASE} --pre {PACKAGES_TO_INSTALL_WHL}"
+            whl_install_command = f"{WHL_INSTALL_BASE} --pre {PACKAGES_TO_INSTALL}"
         else:
-            whl_install_command = f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL}"
+            whl_install_command = f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL}"
         return f"{whl_install_command} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}"  # noqa: E501
 
 
@@ -409,6 +416,7 @@ def generate_wheels_matrix(
     limit_pr_builds: bool,
     use_only_dl_pytorch_org: bool,
     use_split_build: bool = False,
+    getting_started: bool = False,
     python_versions: Optional[List[str]] = None,
     arches: Optional[List[str]] = None,
 ) -> List[Dict[str, str]]:
@@ -483,6 +491,8 @@ def generate_wheels_matrix(
                     desired_cuda,
                     python_version,
                     use_only_dl_pytorch_org,
+                    use_split_build,
+                    getting_started,
                 ),
                 "channel": channel,
                 "upload_to_base_bucket": upload_to_base_bucket,
@@ -526,6 +536,7 @@ def generate_build_matrix(
     use_only_dl_pytorch_org: str,
     build_python_only: str,
     use_split_build: str = "false",
+    getting_started: str = "false",
     python_versions: Optional[List[str]] = None,
 ) -> Dict[str, List[Dict[str, str]]]:
     includes = []
@@ -550,6 +561,7 @@ def generate_build_matrix(
                     limit_pr_builds == "true",
                     use_only_dl_pytorch_org == "true",
                     use_split_build == "true",
+                    getting_started == "true",
                     python_versions,
                 )
             )
@@ -645,6 +657,14 @@ def main(args: List[str]) -> None:
         default=os.getenv("USE_SPLIT_BUILD", DISABLE),
     )
 
+    parser.add_argument(
+        "--getting-started",
+        help="Matrix for getting started page",
+        type=str,
+        choices=["true", "false"],
+        default=os.getenv("GETTING_STARTED", "false"),
+    )
+
     parser.add_argument(
         "--python-versions",
         help="Only build the select JSON-encoded list of python versions",
@@ -674,6 +694,7 @@ def main(args: List[str]) -> None:
         options.use_only_dl_pytorch_org,
         options.build_python_only,
         options.use_split_build,
+        options.getting_started,
         python_versions,
     )
 
