regenerate.py 4.9 KB
Newer Older
facebook-github-bot's avatar
facebook-github-bot committed
1
#!/usr/bin/env python3
2
# Copyright (c) Meta Platforms, Inc. and affiliates.
Patrick Labatut's avatar
Patrick Labatut committed
3
4
5
6
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
facebook-github-bot's avatar
facebook-github-bot committed
7
8
9
10
11
12

"""
This script is adapted from the torchvision one.
"""

import os.path
13

facebook-github-bot's avatar
facebook-github-bot committed
14
15
import jinja2
import yaml
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
16
from packaging import version
facebook-github-bot's avatar
facebook-github-bot committed
17
18


Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
19
20
21
# The CUDA versions which have pytorch conda packages available for linux for each
# version of pytorch.
CONDA_CUDA_VERSIONS = {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
22
23
    "1.12.0": ["cu113", "cu116"],
    "1.12.1": ["cu113", "cu116"],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
24
    "1.13.0": ["cu116", "cu117"],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
25
    "1.13.1": ["cu116", "cu117"],
26
    "2.0.0": ["cu117", "cu118"],
27
    "2.0.1": ["cu117", "cu118"],
28
    "2.1.0": ["cu118", "cu121"],
29
30
31
    "2.1.1": ["cu118", "cu121"],
    "2.1.2": ["cu118", "cu121"],
    "2.2.0": ["cu118", "cu121"],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
32
33
34
}


Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
35
def conda_docker_image_for_cuda(cuda_version):
36
37
38
    if len(cuda_version) != 5:
        raise ValueError("Unknown cuda version")
    return "pytorch/conda-builder:cuda" + cuda_version[2:]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
39
40


41
def pytorch_versions_for_python(python_version):
42
    if python_version in ["3.8", "3.9"]:
43
        return list(CONDA_CUDA_VERSIONS)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
44
45
46
47
48
49
    if python_version == "3.10":
        return [
            i
            for i in CONDA_CUDA_VERSIONS
            if version.Version(i) >= version.Version("1.11.0")
        ]
50
51
52
53
54
55
    if python_version == "3.11":
        return [
            i
            for i in CONDA_CUDA_VERSIONS
            if version.Version(i) >= version.Version("2.1.0")
        ]
56
57
58
59
60
61
    if python_version == "3.12":
        return [
            i
            for i in CONDA_CUDA_VERSIONS
            if version.Version(i) >= version.Version("2.2.0")
        ]
62
63


facebook-github-bot's avatar
facebook-github-bot committed
64
65
66
def workflows(prefix="", filter_branch=None, upload=False, indentation=6):
    w = []
    for btype in ["conda"]:
67
        for python_version in ["3.8", "3.9", "3.10", "3.11", "3.12"]:
68
            for pytorch_version in pytorch_versions_for_python(python_version):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
69
70
71
72
73
74
75
76
77
78
                for cu_version in CONDA_CUDA_VERSIONS[pytorch_version]:
                    w += workflow_pair(
                        btype=btype,
                        python_version=python_version,
                        pytorch_version=pytorch_version,
                        cu_version=cu_version,
                        prefix=prefix,
                        upload=upload,
                        filter_branch=filter_branch,
                    )
facebook-github-bot's avatar
facebook-github-bot committed
79
80
81
82
83

    return indent(indentation, w)


def workflow_pair(
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
84
85
86
87
88
89
90
91
    *,
    btype,
    python_version,
    pytorch_version,
    cu_version,
    prefix="",
    upload=False,
    filter_branch,
facebook-github-bot's avatar
facebook-github-bot committed
92
93
94
):

    w = []
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
95
96
97
    py = python_version.replace(".", "")
    pyt = pytorch_version.replace(".", "")
    base_workflow_name = f"{prefix}linux_{btype}_py{py}_{cu_version}_pyt{pyt}"
facebook-github-bot's avatar
facebook-github-bot committed
98
99
100
101
102

    w.append(
        generate_base_workflow(
            base_workflow_name=base_workflow_name,
            python_version=python_version,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
103
            pytorch_version=pytorch_version,
facebook-github-bot's avatar
facebook-github-bot committed
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
            cu_version=cu_version,
            btype=btype,
            filter_branch=filter_branch,
        )
    )

    if upload:
        w.append(
            generate_upload_workflow(
                base_workflow_name=base_workflow_name,
                btype=btype,
                cu_version=cu_version,
                filter_branch=filter_branch,
            )
        )

    return w


def generate_base_workflow(
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
124
125
126
127
128
129
130
    *,
    base_workflow_name,
    python_version,
    cu_version,
    pytorch_version,
    btype,
    filter_branch=None,
facebook-github-bot's avatar
facebook-github-bot committed
131
132
133
134
135
136
):

    d = {
        "name": base_workflow_name,
        "python_version": python_version,
        "cu_version": cu_version,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
137
        "pytorch_version": pytorch_version,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
138
        "context": "DOCKERHUB_TOKEN",
facebook-github-bot's avatar
facebook-github-bot committed
139
140
    }

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
141
142
143
144
    conda_docker_image = conda_docker_image_for_cuda(cu_version)
    if conda_docker_image is not None:
        d["conda_docker_image"] = conda_docker_image

facebook-github-bot's avatar
facebook-github-bot committed
145
146
147
148
149
150
    if filter_branch is not None:
        d["filters"] = {"branches": {"only": filter_branch}}

    return {f"binary_linux_{btype}": d}


151
def generate_upload_workflow(*, base_workflow_name, btype, cu_version, filter_branch):
facebook-github-bot's avatar
facebook-github-bot committed
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
    d = {
        "name": f"{base_workflow_name}_upload",
        "context": "org-member",
        "requires": [base_workflow_name],
    }

    if btype == "wheel":
        d["subfolder"] = cu_version + "/"

    if filter_branch is not None:
        d["filters"] = {"branches": {"only": filter_branch}}

    return {f"binary_{btype}_upload": d}


def indent(indentation, data_list):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
168
169
    if len(data_list) == 0:
        return ""
facebook-github-bot's avatar
facebook-github-bot committed
170
171
172
173
174
175
176
177
    return ("\n" + " " * indentation).join(
        yaml.dump(data_list, default_flow_style=False).splitlines()
    )


if __name__ == "__main__":
    d = os.path.dirname(__file__)
    env = jinja2.Environment(
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
178
179
180
181
        loader=jinja2.FileSystemLoader(d),
        lstrip_blocks=True,
        autoescape=False,
        keep_trailing_newline=True,
facebook-github-bot's avatar
facebook-github-bot committed
182
183
184
185
    )

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