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

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

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

* release 2.1.2 only changes, fix rebase

* lint

Author: atalman

.github/scripts/generate_binary_build_matrix.py

@@ -16,6 +16,12 @@
 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.5", "5.6"]
 
 

.github/scripts/generate_docker_release_matrix.py

@@ -0,0 +1,45 @@
+#!/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:
+            platform = (
+                "linux/arm64,linux/amd64" if image == "runtime" else "linux/amd64"
+            )
+            ret.append(
+                {
+                    "cuda": cuda,
+                    "cuda_full_version": version,
+                    "cudnn_version": generate_binary_build_matrix.CUDA_ARCHES_CUDNN_VERSION[
+                        cuda
+                    ],
+                    "image_type": image,
+                    "platform": platform,
+                }
+            )
+    return {"include": ret}
+
+
+if __name__ == "__main__":
+    build_matrix = generate_docker_matrix()
+    print(json.dumps(build_matrix))

.github/workflows/docker-release.yml

@@ -39,23 +39,39 @@ env:
   WITH_PUSH: ${{ inputs.channel == 'release' }}
 
 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/main' || startsWith(github.event.ref, 'refs/tags/v')) && 'docker-build' || '' }}
     timeout-minutes: 240
+    needs: generate-matrix
     strategy:
-      matrix:
-        include:
-          # 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
+      matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix) }}
+      fail-fast: false
     env:
       BUILD_IMAGE_TYPE: ${{ matrix.image_type }}
       BUILD_PLATFORMS: ${{ matrix.platform }}
       CHANNEL: ${{ inputs.channel }}
+      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@release/2.1

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
-CUDNN_VERSION             = 8
+CUDA_VERSION              ?= 12.1.1
+CUDNN_VERSION             ?= 8
 BASE_RUNTIME              = ubuntu:20.04
 BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu20.04
 CMAKE_VARS               ?=