From 796d69f30da4cc220729e6b8ce3ca5e6fee9e2c9 Mon Sep 17 00:00:00 2001 From: PaliC <> Date: Tue, 25 Jun 2024 16:54:44 -0700 Subject: [PATCH] [Split Build] limit split build matrix generation to cuda nightly jobs Validation this works ``` sahanp@sahanp-mbp ~/test-infra [2]> python tools/scripts/generate_binary_build_matrix.py --use-split-build true (base) { "include": [ { "python_version": "3.8", "gpu_arch_type": "cpu", "gpu_arch_version": "", "desired_cuda": "cpu", "container_image": "pytorch/manylinux-builder:cpu", "package_type": "manywheel", "build_name": "manywheel-py3_8-cpu", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": false }, { "python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "11.8", "desired_cuda": "cu118", "container_image": "pytorch/manylinux-builder:cuda11.8", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda11_8", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": false }, { "python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "11.8", "desired_cuda": "cu118", "container_image": "pytorch/manylinux-builder:cuda11.8", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda11_8-split", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": true }, { "python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "12.1", "desired_cuda": "cu121", "container_image": "pytorch/manylinux-builder:cuda12.1", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda12_1", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": false }, { "python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "12.1", "desired_cuda": "cu121", "container_image": "pytorch/manylinux-builder:cuda12.1", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda12_1-split", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": true }, { "python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "12.4", "desired_cuda": "cu124", "container_image": "pytorch/manylinux-builder:cuda12.4", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda12_4", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu124", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": false }, { "python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "12.4", "desired_cuda": "cu124", "container_image": "pytorch/manylinux-builder:cuda12.4", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda12_4-split", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu124", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": true }, { "python_version": "3.8", "gpu_arch_type": "rocm", "gpu_arch_version": "6.0", "desired_cuda": "rocm6.0", "container_image": "pytorch/manylinux-builder:rocm6.0", "package_type": "manywheel", "build_name": "manywheel-py3_8-rocm6_0", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.0", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": false }, { "python_version": "3.8", "gpu_arch_type": "rocm", "gpu_arch_version": "6.1", "desired_cuda": "rocm6.1", "container_image": "pytorch/manylinux-builder:rocm6.1", "package_type": "manywheel", "build_name": "manywheel-py3_8-rocm6_1", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.1", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.3.1", "use_split_build": false } ] } ``` --- tools/scripts/generate_binary_build_matrix.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py index 70465f3152..960544accd 100644 --- a/tools/scripts/generate_binary_build_matrix.py +++ b/tools/scripts/generate_binary_build_matrix.py @@ -545,7 +545,7 @@ def generate_wheels_matrix( "use_split_build": False, } ret.append(entry) - if use_split_build: + if use_split_build and (gpu_arch_version in CUDA_ARCHES) and (os == LINUX) and (channel == NIGHTLY): entry = entry.copy() entry["build_name"] = f"{package_type}-py{python_version}-{gpu_arch_type}{gpu_arch_version}-split".replace( ".", "_"