diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py index 8ebea945fc7..dcfa328c187 100644 --- a/.github/scripts/generate_binary_build_matrix.py +++ b/.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 diff --git a/.github/scripts/generate_docker_release_matrix.py b/.github/scripts/generate_docker_release_matrix.py new file mode 100644 index 00000000000..4dfa016be86 --- /dev/null +++ b/.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)) diff --git a/.github/workflows/docker-release.yml b/.github/workflows/docker-release.yml index 2b487038182..c678fccb104 100644 --- a/.github/workflows/docker-release.yml +++ b/.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: - 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 }} + 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 diff --git a/Dockerfile b/Dockerfile index 62a5b839a1b..a2c4bef9659 100644 --- a/Dockerfile +++ b/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 diff --git a/docker.Makefile b/docker.Makefile index 184648ccf4d..63d4126f0df 100644 --- a/docker.Makefile +++ b/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:22.04 BASE_DEVEL = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu22.04 CMAKE_VARS ?=