Skip to content

Commit

Permalink
Option to Generate Matrix without ROCM (#4360)
Browse files Browse the repository at this point in the history 
Will add tests once #4386 is
merged.
  • Loading branch information
@osalpekar
osalpekar authored Jul 13, 2023
1 parent abf394b commit 63d1ae3
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 5 deletions.
27 changes: 23 additions & 4 deletions tools/scripts/generate_binary_build_matrix.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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]
Expand Down Expand Up @@ -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,
Expand All @@ -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] ]

Expand Down Expand Up @@ -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 = []

Expand All @@ -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")
)

Expand Down Expand Up @@ -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"
Expand All @@ -466,14 +485,14 @@ def main(args) -> None:
)



options = parser.parse_args(args)

build_matrix = generate_build_matrix(
options.package_type,
options.operating_system,
options.channel,
options.with_cuda,
options.with_rocm,
options.limit_pr_builds)

print(json.dumps(build_matrix))
Expand Down
1 change: 1 addition & 0 deletions tools/tests/assets/build_matrix_linux_wheel_cuda_norocm.json
View file
Original file line number Diff line number Diff line change
@@ -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"}]}
18 changes: 17 additions & 1 deletion tools/tests/test_generate_binary_build_matrix.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)

Expand All @@ -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",
)

Expand All @@ -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",
)

Expand All @@ -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",
)

Expand All @@ -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",
)

Expand All @@ -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",
)

Expand All @@ -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()

0 comments on commit 63d1ae3

Please sign in to comment.