forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathgenerate_binary_build_matrix.py
More file actions
86 lines (73 loc) · 1.96 KB
/
Copy pathgenerate_binary_build_matrix.py
File metadata and controls
86 lines (73 loc) · 1.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
#!/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()