From 0ca3978e8e559ebd503c49367ed7cefd9059b717 Mon Sep 17 00:00:00 2001 From: atalman Date: Fri, 5 Apr 2024 12:49:19 -0700 Subject: [PATCH 1/3] Add docker release matrix workflow --- .../generate_docker_release_matrix.yml | 50 ++++++++++++++++ tools/scripts/generate_binary_build_matrix.py | 5 ++ .../scripts/generate_docker_release_matrix.py | 58 +++++++++++++++++++ 3 files changed, 113 insertions(+) create mode 100644 .github/workflows/generate_docker_release_matrix.yml create mode 100644 tools/scripts/generate_docker_release_matrix.py diff --git a/.github/workflows/generate_docker_release_matrix.yml b/.github/workflows/generate_docker_release_matrix.yml new file mode 100644 index 0000000000..f796216e89 --- /dev/null +++ b/.github/workflows/generate_docker_release_matrix.yml @@ -0,0 +1,50 @@ +name: Generates the docker release matrix + +on: + workflow_call: + inputs: + channel: + description: "Channel to use (nightly, test, release, all)" + default: "" + type: string + test-infra-repository: + description: "Test infra repository to use" + default: "pytorch/test-infra" + type: string + test-infra-ref: + description: "Test infra reference to use" + default: "main" + type: string + outputs: + matrix: + description: "Generated build matrix" + value: ${{ jobs.generate.outputs.matrix }} + +jobs: + generate: + outputs: + matrix: ${{ steps.generate.outputs.matrix }} + runs-on: ubuntu-latest + steps: + - uses: actions/setup-python@v4 + with: + python-version: '3.10' + - name: Checkout test-infra repository + uses: actions/checkout@v3 + with: + repository: ${{ inputs.test-infra-repository }} + ref: ${{ inputs.test-infra-ref }} + - uses: ./.github/actions/set-channel + - name: Generate docker release matrix + id: generate + env: + CHANNEL: ${{ inputs.channel != '' && inputs.channel || env.CHANNEL }} + run: | + set -eou pipefail + MATRIX_BLOB="$(python3 tools/scripts/generate_docker_release_matrix.py)" + echo "${MATRIX_BLOB}" + echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}" + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.package-type }}-${{ inputs.os }}-${{ inputs.test-infra-repository }}-${{ inputs.test-infra-ref }} + cancel-in-progress: true diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py index 0f9f115eee..590c88ec8f 100644 --- a/tools/scripts/generate_binary_build_matrix.py +++ b/tools/scripts/generate_binary_build_matrix.py @@ -36,6 +36,11 @@ "release": ["5.6", "5.7"], } +CUDA_CUDDN_VERSIONS = { + "11.8": { "cuda": "11.8.0", "cudnn": "8" }, + "12.1": { "cuda": "12.1.1", "cudnn": "8"}, +} + PACKAGE_TYPES = ["wheel", "conda", "libtorch"] PRE_CXX11_ABI = "pre-cxx11" CXX11_ABI = "cxx11-abi" diff --git a/tools/scripts/generate_docker_release_matrix.py b/tools/scripts/generate_docker_release_matrix.py new file mode 100644 index 0000000000..82691eaf1a --- /dev/null +++ b/tools/scripts/generate_docker_release_matrix.py @@ -0,0 +1,58 @@ +#!/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 version + * CUDNN version short + * Image type either runtime or devel + * Platform linux/arm64,linux/amd64 + +""" + +import json +import os +import sys +import argparse +from typing import Dict, List + +import generate_binary_build_matrix + +DOCKER_IMAGE_TYPES = ["runtime", "devel"] + + +def generate_docker_matrix(channel: str) -> Dict[str, List[Dict[str, str]]]: + + ret: List[Dict[str, str]] = [] + for cuda in generate_binary_build_matrix.CUDA_ARCHES_DICT[channel]: + version = generate_binary_build_matrix.CUDA_CUDDN_VERSIONS[cuda] + for image in DOCKER_IMAGE_TYPES: + ret.append( + { + "cuda": cuda, + "cuda_full_version": version["cuda"], + "cudnn_version": version["cudnn"], + "image_type": image, + "platform": "linux/arm64,linux/amd64", + } + ) + return {"include": ret} + + +def main(args) -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--channel", + help="Channel to use, default nightly", + type=str, + choices=["nightly", "test", "release", "all"], + default=os.getenv("CHANNEL", "nightly"), + ) + options = parser.parse_args(args) + + build_matrix = generate_docker_matrix(options.channel) + print(json.dumps(build_matrix)) + +if __name__ == "__main__": + main(sys.argv[1:]) From 8f82307790b489351f20fcf09a38f4e973a1cfa8 Mon Sep 17 00:00:00 2001 From: atalman Date: Fri, 5 Apr 2024 13:00:40 -0700 Subject: [PATCH 2/3] lint --- .github/workflows/generate_docker_release_matrix.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/generate_docker_release_matrix.yml b/.github/workflows/generate_docker_release_matrix.yml index f796216e89..5a01233ab1 100644 --- a/.github/workflows/generate_docker_release_matrix.yml +++ b/.github/workflows/generate_docker_release_matrix.yml @@ -46,5 +46,5 @@ jobs: echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}" concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.package-type }}-${{ inputs.os }}-${{ inputs.test-infra-repository }}-${{ inputs.test-infra-ref }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.test-infra-repository }}-${{ inputs.test-infra-ref }} cancel-in-progress: true From 6506b147bd4d8887828f0d51d0b561c915c2d72d Mon Sep 17 00:00:00 2001 From: atalman Date: Fri, 5 Apr 2024 13:51:33 -0700 Subject: [PATCH 3/3] fix --- tools/scripts/generate_docker_release_matrix.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tools/scripts/generate_docker_release_matrix.py b/tools/scripts/generate_docker_release_matrix.py index 82691eaf1a..b08a42f0ab 100644 --- a/tools/scripts/generate_docker_release_matrix.py +++ b/tools/scripts/generate_docker_release_matrix.py @@ -40,7 +40,7 @@ def generate_docker_matrix(channel: str) -> Dict[str, List[Dict[str, str]]]: return {"include": ret} -def main(args) -> None: +def main() -> None: parser = argparse.ArgumentParser() parser.add_argument( "--channel", @@ -49,10 +49,10 @@ def main(args) -> None: choices=["nightly", "test", "release", "all"], default=os.getenv("CHANNEL", "nightly"), ) - options = parser.parse_args(args) + options = parser.parse_args() build_matrix = generate_docker_matrix(options.channel) print(json.dumps(build_matrix)) if __name__ == "__main__": - main(sys.argv[1:]) + main()