diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py index f497ab54be..fa41f2d3d9 100644 --- a/tools/scripts/generate_binary_build_matrix.py +++ b/tools/scripts/generate_binary_build_matrix.py @@ -203,7 +203,7 @@ def get_wheel_install_command(os: str, channel: str, gpu_arch_type: str, gpu_arc whl_install_command = f"{WHL_INSTALL_BASE} --pre {PACKAGES_TO_INSTALL_WHL}" if channel == "nightly" else f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL}" return f"{whl_install_command} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}" -def generate_conda_matrix(os: str, channel: str, with_cuda: str, limit_pr_builds: bool) -> List[Dict[str, str]]: +def generate_conda_matrix(os: str, channel: str, with_cuda: str, with_rocm: str, limit_pr_builds: bool) -> List[Dict[str, str]]: ret: List[Dict[str, str]] = [] arches = ["cpu"] python_versions = list(mod.PYTHON_ARCHES) @@ -247,6 +247,7 @@ def generate_libtorch_matrix( os: str, channel: str, with_cuda: str, + with_rocm: str, limit_pr_builds: str, abi_versions: Optional[List[str]] = None, arches: Optional[List[str]] = None, @@ -265,10 +266,14 @@ def generate_libtorch_matrix( if with_cuda == ENABLE: if os == "linux": arches += mod.CUDA_ARCHES - arches += mod.ROCM_ARCHES elif os == "windows": arches += mod.CUDA_ARCHES + if with_rocm == ENABLE: + if os == "linux": + arches += mod.ROCM_ARCHES + + if abi_versions is None: if os == "windows": abi_versions = [RELEASE, DEBUG] @@ -336,6 +341,7 @@ def generate_wheels_matrix( os: str, channel: str, with_cuda: str, + with_rocm: str, limit_pr_builds: bool, arches: Optional[List[str]] = None, python_versions: Optional[List[str]] = None, @@ -358,10 +364,14 @@ def generate_wheels_matrix( if with_cuda == ENABLE: upload_to_base_bucket = "no" if os == "linux": - arches += mod.CUDA_ARCHES + mod.ROCM_ARCHES + arches += mod.CUDA_ARCHES elif os == "windows": arches += mod.CUDA_ARCHES + if with_rocm == ENABLE: + if os == "linux": + arches += mod.ROCM_ARCHES + if limit_pr_builds: python_versions = [ python_versions[0] ] @@ -404,6 +414,7 @@ def generate_build_matrix( operating_system: str, channel: str, with_cuda: str, + with_rocm: str, limit_pr_builds: str) -> Dict[str, List[Dict[str, str]]]: includes = [] @@ -420,6 +431,7 @@ def generate_build_matrix( GENERATING_FUNCTIONS_BY_PACKAGE_TYPE[package](operating_system, channel, with_cuda, + with_rocm, limit_pr_builds == "true") ) @@ -454,6 +466,13 @@ def main(args) -> None: choices=[ENABLE, DISABLE], default=os.getenv("WITH_CUDA", ENABLE), ) + parser.add_argument( + "--with-rocm", + help="Build with Rocm?", + type=str, + choices=[ENABLE, DISABLE], + default=os.getenv("WITH_ROCM", ENABLE), + ) # By default this is false for this script but expectation is that the caller # workflow will default this to be true most of the time, where a pull # request is synchronized and does not contain the label "ciflow/binaries/all" @@ -466,7 +485,6 @@ def main(args) -> None: ) - options = parser.parse_args(args) build_matrix = generate_build_matrix( @@ -474,6 +492,7 @@ def main(args) -> None: options.operating_system, options.channel, options.with_cuda, + options.with_rocm, options.limit_pr_builds) print(json.dumps(build_matrix)) diff --git a/tools/tests/assets/build_matrix_linux_wheel_cuda_norocm.json b/tools/tests/assets/build_matrix_linux_wheel_cuda_norocm.json new file mode 100644 index 0000000000..4c59fa0edf --- /dev/null +++ b/tools/tests/assets/build_matrix_linux_wheel_cuda_norocm.json @@ -0,0 +1 @@ +{"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.0.1"}, {"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.0.1"}, {"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.0.1"}, {"python_version": "3.9", "gpu_arch_type": "cpu", "gpu_arch_version": "", "desired_cuda": "cpu", "container_image": "pytorch/manylinux-builder:cpu", "package_type": "manywheel", "build_name": "manywheel-py3_9-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.0.1"}, {"python_version": "3.9", "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_9-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.0.1"}, {"python_version": "3.9", "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_9-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.0.1"}, {"python_version": "3.10", "gpu_arch_type": "cpu", "gpu_arch_version": "", "desired_cuda": "cpu", "container_image": "pytorch/manylinux-builder:cpu", "package_type": "manywheel", "build_name": "manywheel-py3_10-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.0.1"}, {"python_version": "3.10", "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_10-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.0.1"}, {"python_version": "3.10", "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_10-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.0.1"}, {"python_version": "3.11", "gpu_arch_type": "cpu", "gpu_arch_version": "", "desired_cuda": "cpu", "container_image": "pytorch/manylinux-builder:cpu", "package_type": "manywheel", "build_name": "manywheel-py3_11-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.0.1"}, {"python_version": "3.11", "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_11-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.0.1"}, {"python_version": "3.11", "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_11-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.0.1"}]} diff --git a/tools/tests/test_generate_binary_build_matrix.py b/tools/tests/test_generate_binary_build_matrix.py index adcc74982c..aaf656f555 100644 --- a/tools/tests/test_generate_binary_build_matrix.py +++ b/tools/tests/test_generate_binary_build_matrix.py @@ -10,13 +10,14 @@ class GenerateBuildMatrixTest(TestCase): def matrix_compare_helper( - self, package_type, operating_system, cuda, reference_output_file + self, package_type, operating_system, cuda, rocm, reference_output_file ): out = generate_build_matrix( package_type, operating_system, "nightly", "enable" if cuda else "disable", + "enable" if rocm else "disable", "false", ) @@ -33,6 +34,7 @@ def test_linux_wheel_cuda(self): package_type="wheel", operating_system="linux", cuda=True, + rocm=True, reference_output_file="build_matrix_linux_wheel_cuda.json", ) @@ -41,6 +43,7 @@ def test_linux_conda_cuda(self): package_type="conda", operating_system="linux", cuda=True, + rocm=True, reference_output_file="build_matrix_linux_conda_cuda.json", ) @@ -49,6 +52,7 @@ def test_macos_wheel(self): package_type="wheel", operating_system="macos", cuda=False, + rocm=False, reference_output_file="build_matrix_macos_wheel.json", ) @@ -57,6 +61,7 @@ def test_macos_conda(self): package_type="conda", operating_system="macos", cuda=False, + rocm=False, reference_output_file="build_matrix_macos_conda.json", ) @@ -65,6 +70,7 @@ def test_windows_wheel_cuda(self): package_type="wheel", operating_system="windows", cuda=True, + rocm=True, reference_output_file="build_matrix_windows_wheel_cuda.json", ) @@ -73,9 +79,19 @@ def test_windows_conda_cuda(self): package_type="conda", operating_system="windows", cuda=True, + rocm=True, reference_output_file="build_matrix_windows_conda_cuda.json", ) + def test_linux_wheel_cuda_norocm(self): + self.matrix_compare_helper( + package_type="wheel", + operating_system="linux", + cuda=True, + rocm=False, + reference_output_file="build_matrix_linux_wheel_cuda_norocm.json", + ) + if __name__ == "__main__": main()