Use matrix generate script for docker release workflows (#115949)

Enable both supported CUDA version builds for docker release. Rather then building only 1 version.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115949
Approved by: https://github.com/huydhn

.github/scripts/generate_binary_build_matrix.py

@@ -16,6 +16,12 @@ from typing import Dict, List, Optional, Tuple
 CUDA_ARCHES = ["11.8", "12.1"]
 
 
+CUDA_ARCHES_FULL_VERSION = {"11.8": "11.8.0", "12.1": "12.1.1"}
+
+
+CUDA_ARCHES_CUDNN_VERSION = {"11.8": "8", "12.1": "8"}
+
+
 ROCM_ARCHES = ["5.6", "5.7"]
 
 
@@ -24,6 +30,7 @@ CPU_CXX11_ABI_ARCH = ["cpu-cxx11-abi"]
 
 CPU_AARCH64_ARCH = ["cpu-aarch64"]
 
+
 PYTORCH_EXTRA_INSTALL_REQUIREMENTS = {
     "11.8": (
         "nvidia-cuda-nvrtc-cu11==11.8.89; platform_system == 'Linux' and platform_machine == 'x86_64' | "  # noqa: B950

.github/scripts/generate_docker_release_matrix.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+"""Generates a matrix for docker releases through github actions
+
+Will output a condensed version of the matrix. Will include fllowing:
+    * CUDA version short
+    * CUDA full verison
+    * CUDNN version short
+    * Image type either runtime or devel
+    * Platform linux/arm64,linux/amd64
+
+"""
+
+import json
+from typing import Dict, List
+
+import generate_binary_build_matrix
+
+DOCKER_IMAGE_TYPES = ["runtime", "devel"]
+
+
+def generate_docker_matrix() -> Dict[str, List[Dict[str, str]]]:
+    ret: List[Dict[str, str]] = []
+    for cuda, version in generate_binary_build_matrix.CUDA_ARCHES_FULL_VERSION.items():
+        for image in DOCKER_IMAGE_TYPES:
+            ret.append(
+                {
+                    "cuda": cuda,
+                    "cuda_full_version": version,
+                    "cudnn_version": generate_binary_build_matrix.CUDA_ARCHES_CUDNN_VERSION[
+                        cuda
+                    ],
+                    "image_type": image,
+                    "platform": "linux/arm64,linux/amd64",
+                }
+            )
+    return {"include": ret}
+
+
+if __name__ == "__main__":
+    build_matrix = generate_docker_matrix()
+    print(json.dumps(build_matrix))

.github/workflows/docker-release.yml

@@ -29,22 +29,38 @@ env:
   WITH_PUSH: ${{ github.event_name == 'push' && (github.event.ref == 'refs/heads/nightly' || startsWith(github.event.ref, 'refs/tags/v')) }}
 
 jobs:
+  generate-matrix:
+    if: github.repository_owner == 'pytorch'
+    runs-on: [self-hosted, linux.large]
+    outputs:
+      matrix: ${{ steps.generate-matrix.outputs.matrix }}
+    steps:
+      - name: Checkout PyTorch
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+        with:
+          fetch-depth: 1
+          submodules: true
+      - name: Get docker release matrix
+        id: generate-matrix
+        run: |
+          MATRIX_BLOB="$(python3 .github/scripts/generate_docker_release_matrix.py)"
+          echo "${MATRIX_BLOB}"
+          echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}"
+
   build:
     if: ${{ github.repository == 'pytorch/pytorch' }}
     runs-on: [self-hosted, linux.2xlarge]
     environment: ${{ (github.ref == 'refs/heads/nightly' || startsWith(github.event.ref, 'refs/tags/v')) && 'docker-build' || '' }}
     timeout-minutes: 240
+    needs: generate-matrix
     strategy:
-      matrix:
+      matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix) }}
-        include:
+      fail-fast: false
-          # nvidia specific images don't exist for arm64 so only build the runtime image
-          - image_type: runtime
-            platform: linux/arm64,linux/amd64
-          - image_type: devel
-            platform: linux/amd64
     env:
       BUILD_IMAGE_TYPE: ${{ matrix.image_type }}
       BUILD_PLATFORMS: ${{ matrix.platform }}
+      CUDA_VERSION: ${{ matrix.cuda_full_version }}
+      CUDNN_VERSION: ${{ matrix.cudnn_version }}
     steps:
       - name: Setup SSH (Click me for login details)
         uses: pytorch/test-infra/.github/actions/setup-ssh@main

Dockerfile

@@ -63,7 +63,7 @@ RUN --mount=type=cache,target=/opt/ccache \
 
 FROM conda as conda-installs
 ARG PYTHON_VERSION=3.8
-ARG CUDA_VERSION=11.7
+ARG CUDA_VERSION=12.1
 ARG CUDA_CHANNEL=nvidia
 ARG INSTALL_CHANNEL=pytorch-nightly
 # Automatically set by buildx

docker.Makefile

@@ -8,8 +8,8 @@ $(warning WARNING: No docker user found using results from whoami)
 DOCKER_ORG                = $(shell whoami)
 endif
 
-CUDA_VERSION              = 12.1.1
+CUDA_VERSION             ?= 12.1.1
-CUDNN_VERSION             = 8
+CUDNN_VERSION            ?= 8
 BASE_RUNTIME              = ubuntu:22.04
 BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu22.04
 CMAKE_VARS               ?=