regenerate.py 4.64 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
#!/usr/bin/env python3

"""
This script should use a very simple, functional programming style.
Avoid Jinja macros in favor of native Python functions.

Don't go overboard on code generation; use Python only to generate
content that can't be easily declared statically using CircleCI's YAML API.

Data declarations (e.g. the nested loops for defining the configuration matrix)
should be at the top of the file for easy updating.

See this comment for design rationale:
https://github.com/pytorch/vision/pull/1321#issuecomment-531033978
"""
16
17

import jinja2
18
import yaml
19
20
import os.path

21

22
def workflows(prefix='', filter_branch=None, upload=False, indentation=6, windows_latest_only=False):
23
24
    w = []
    for btype in ["wheel", "conda"]:
25
        for os_type in ["linux", "macos", "win"]:
26
            python_versions = ["3.6", "3.7", "3.8"]
27
28
29
            cu_versions = (["cpu", "cu92", "cu101", "cu102"] if os_type == "linux" or os_type == "win" else ["cpu"])
            for python_version in python_versions:
                for cu_version in cu_versions:
30
                    for unicode in ([False, True] if btype == "wheel" and python_version == "2.7" else [False]):
31
32
33
34
35
                        fb = filter_branch
                        if windows_latest_only and os_type == "win" and filter_branch is None and \
                            (python_version != python_versions[-1] or
                             (cu_version not in [cu_versions[0], cu_versions[-1]])):
                            fb = "master"
36
37
                        w += workflow_pair(
                            btype, os_type, python_version, cu_version,
38
                            unicode, prefix, upload, filter_branch=fb)
39
40
41
42

    return indent(indentation, w)


43
def workflow_pair(btype, os_type, python_version, cu_version, unicode, prefix='', upload=False, *, filter_branch=None):
44
45
46
47
48

    w = []
    unicode_suffix = "u" if unicode else ""
    base_workflow_name = f"{prefix}binary_{os_type}_{btype}_py{python_version}{unicode_suffix}_{cu_version}"

49
50
51
    w.append(generate_base_workflow(
        base_workflow_name, python_version, cu_version,
        unicode, os_type, btype, filter_branch=filter_branch))
52
53

    if upload:
54
        w.append(generate_upload_workflow(base_workflow_name, os_type, btype, cu_version, filter_branch=filter_branch))
55
56
57
58

    return w


59
60
61
62
63
64
65
66
67
68
69
70
71
72
manylinux_images = {
    "cu92": "pytorch/manylinux-cuda92",
    "cu101": "pytorch/manylinux-cuda101",
    "cu102": "pytorch/manylinux-cuda102",
}


def get_manylinux_image(cu_version):
    cu_suffix = "102"
    if cu_version.startswith('cu'):
        cu_suffix = cu_version[len('cu'):]
    return f"pytorch/manylinux-cuda{cu_suffix}"


73
74
def generate_base_workflow(base_workflow_name, python_version, cu_version,
                           unicode, os_type, btype, *, filter_branch=None):
75
76
77
78

    d = {
        "name": base_workflow_name,
        "python_version": python_version,
79
        "cu_version": cu_version,
80
81
    }

82
    if os_type != "win" and unicode:
83
84
        d["unicode_abi"] = '1'

85
86
    if os_type != "win":
        d["wheel_docker_image"] = get_manylinux_image(cu_version)
87

88
    if filter_branch is not None:
89
90
91
92
93
94
95
96
97
98
        d["filters"] = {
            "branches": {
                "only": filter_branch
            },
            "tags": {
                # Using a raw string here to avoid having to escape
                # anything
                "only": r"/v[0-9]+(\.[0-9]+)*-rc[0-9]+/"
            }
        }
99

100
101
    w = f"binary_{os_type}_{btype}_release" if os_type == "win" else f"binary_{os_type}_{btype}"
    return {w: d}
102
103


104
def generate_upload_workflow(base_workflow_name, os_type, btype, cu_version, *, filter_branch=None):
105
106
107
108
109
110
111
112
113
    d = {
        "name": f"{base_workflow_name}_upload",
        "context": "org-member",
        "requires": [base_workflow_name],
    }

    if btype == 'wheel':
        d["subfolder"] = "" if os_type == 'macos' else cu_version + "/"

114
    if filter_branch is not None:
115
116
117
118
119
120
121
122
123
124
        d["filters"] = {
            "branches": {
                "only": filter_branch
            },
            "tags": {
                # Using a raw string here to avoid having to escape
                # anything
                "only": r"/v[0-9]+(\.[0-9]+)*-rc[0-9]+/"
            }
        }
125

126
127
128
129
    return {f"binary_{btype}_upload": d}


def indent(indentation, data_list):
130
131
    return ("\n" + " " * indentation).join(
        yaml.dump(data_list, default_flow_style=False).splitlines())
132
133
134
135
136
137
138
139
140
141
142
143


if __name__ == "__main__":
    d = os.path.dirname(__file__)
    env = jinja2.Environment(
        loader=jinja2.FileSystemLoader(d),
        lstrip_blocks=True,
        autoescape=False,
    )

    with open(os.path.join(d, 'config.yml'), 'w') as f:
        f.write(env.get_template('config.yml.in').render(workflows=workflows))