mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
1b7a4f9cde
Summary:
Signed-off-by: Eli Uriegas <eliuriegas@fb.com>
Fixes #{issue number}
Pull Request resolved: https://github.com/pytorch/pytorch/pull/50633
Reviewed By: samestep
Differential Revision: D26083492
Pulled By: seemethere
fbshipit-source-id: c133671b9cf5074539133ee79fca5c680793a85d
87 lines
2.0 KiB
Python
87 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
|
|
"""Generates a matrix to be utilized through github actions
|
|
|
|
Will output a condensed version of the matrix if on a pull request that only
|
|
includes the latest version of python we support built on three different
|
|
architectures:
|
|
* CPU
|
|
* Latest CUDA
|
|
* Latest ROCM
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import itertools
|
|
|
|
CUDA_ARCHES = [
|
|
"10.1",
|
|
"10.2",
|
|
"11.0"
|
|
]
|
|
|
|
ROCM_ARCHES = [
|
|
"3.10",
|
|
"4.0"
|
|
]
|
|
|
|
FULL_ARCHES = [
|
|
"cpu",
|
|
*CUDA_ARCHES,
|
|
*ROCM_ARCHES
|
|
]
|
|
|
|
CONTAINER_IMAGES = {
|
|
**{
|
|
# TODO: Re-do manylinux CUDA image tagging scheme to be similar to
|
|
# ROCM so we don't have to do this replacement
|
|
gpu_arch: f"pytorch/manylinux-cuda{gpu_arch.replace('.', '')}"
|
|
for gpu_arch in CUDA_ARCHES
|
|
},
|
|
**{
|
|
gpu_arch: f"pytorch/manylinux-rocm:{gpu_arch}"
|
|
for gpu_arch in ROCM_ARCHES
|
|
},
|
|
"cpu": "pytorch/manylinux-cpu"
|
|
}
|
|
|
|
FULL_PYTHON_VERSIONS = [
|
|
"3.6",
|
|
"3.7",
|
|
"3.8",
|
|
"3.9",
|
|
]
|
|
|
|
|
|
def is_pull_request():
|
|
return os.environ.get("GITHUB_HEAD_REF")
|
|
|
|
def generate_matrix():
|
|
python_versions = FULL_PYTHON_VERSIONS
|
|
arches = FULL_ARCHES
|
|
if is_pull_request():
|
|
python_versions = [python_versions[-1]]
|
|
arches = ["cpu", CUDA_ARCHES[-1], ROCM_ARCHES[-1]]
|
|
matrix = []
|
|
for item in itertools.product(python_versions, arches):
|
|
python_version, arch_version = item
|
|
# Not my favorite code here
|
|
gpu_arch_type = "cuda"
|
|
if "rocm" in CONTAINER_IMAGES[arch_version]:
|
|
gpu_arch_type = "rocm"
|
|
elif "cpu" in CONTAINER_IMAGES[arch_version]:
|
|
gpu_arch_type = "cpu"
|
|
matrix.append({
|
|
"python_version": python_version,
|
|
"gpu_arch_type": gpu_arch_type,
|
|
"gpu_arch_version": arch_version,
|
|
"container_image": CONTAINER_IMAGES[arch_version]
|
|
})
|
|
return json.dumps({"include": matrix})
|
|
|
|
def main():
|
|
print(generate_matrix())
|
|
|
|
if __name__ == "__main__":
|
|
main()
|