From 95048475516c09f909101791fdbfd3846f99f71e Mon Sep 17 00:00:00 2001 From: atalman Date: Wed, 10 Apr 2024 08:27:02 -0700 Subject: [PATCH] Add docker image to docker release matrix --- tools/scripts/generate_binary_build_matrix.py | 1 + tools/scripts/generate_docker_release_matrix.py | 12 ++++++++++++ 2 files changed, 13 insertions(+) diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py index 590c88ec8f..d2140a2038 100644 --- a/tools/scripts/generate_binary_build_matrix.py +++ b/tools/scripts/generate_binary_build_matrix.py @@ -63,6 +63,7 @@ ROCM = "rocm" +CURRENT_NIGHTLY_VERSION = "2.4.0" CURRENT_CANDIDATE_VERSION = "2.3.0" CURRENT_STABLE_VERSION = "2.2.2" mod.CURRENT_VERSION = CURRENT_STABLE_VERSION diff --git a/tools/scripts/generate_docker_release_matrix.py b/tools/scripts/generate_docker_release_matrix.py index b08a42f0ab..570c7e44f3 100644 --- a/tools/scripts/generate_docker_release_matrix.py +++ b/tools/scripts/generate_docker_release_matrix.py @@ -16,6 +16,7 @@ import sys import argparse from typing import Dict, List +from datetime import datetime import generate_binary_build_matrix @@ -27,6 +28,16 @@ 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] + + prefix = "ghcr.io/pytorch/pytorch" + docker_image_version = "" + if channel == "release": + docker_image_version = f"{prefix}:{generate_binary_build_matrix.CURRENT_STABLE_VERSION}" + elif channel == "test": + docker_image_version = f"{prefix}-test:{generate_binary_build_matrix.CURRENT_CANDIDATE_VERSION}" + else: + docker_image_version = f"{prefix}-nightly:{generate_binary_build_matrix.CURRENT_NIGHTLY_VERSION}.dev{datetime.today().strftime('%Y%m%d')}" + for image in DOCKER_IMAGE_TYPES: ret.append( { @@ -34,6 +45,7 @@ def generate_docker_matrix(channel: str) -> Dict[str, List[Dict[str, str]]]: "cuda_full_version": version["cuda"], "cudnn_version": version["cudnn"], "image_type": image, + "docker": f"{docker_image_version}-cuda{cuda}-cudnn{version['cudnn']}-{image}", "platform": "linux/arm64,linux/amd64", } )