diff --git a/.github/workflows/generate_docker_release_matrix.yml b/.github/workflows/generate_docker_release_matrix.yml new file mode 100644 index 0000000000..5a01233ab1 --- /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.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..b08a42f0ab --- /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() -> 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() + + build_matrix = generate_docker_matrix(options.channel) + print(json.dumps(build_matrix)) + +if __name__ == "__main__": + main()