pytorch

1
#!/usr/bin/env python3
2

3
import os
4
import sys
5
from dataclasses import asdict, dataclass, field
6
from pathlib import Path
7
from typing import Dict, Iterable, List, Literal, Set
8
from typing_extensions import TypedDict  # Python 3.11+
9

10
import generate_binary_build_matrix  # type: ignore[import]
11
import jinja2
12

13

14
Arch = Literal["windows", "linux", "macos"]
15

16
GITHUB_DIR = Path(__file__).resolve().parent.parent
17

18
LABEL_CIFLOW_TRUNK = "ciflow/trunk"
19
LABEL_CIFLOW_UNSTABLE = "ciflow/unstable"
20
LABEL_CIFLOW_BINARIES = "ciflow/binaries"
21
LABEL_CIFLOW_PERIODIC = "ciflow/periodic"
22
LABEL_CIFLOW_BINARIES_LIBTORCH = "ciflow/binaries_libtorch"
23
LABEL_CIFLOW_BINARIES_CONDA = "ciflow/binaries_conda"
24
LABEL_CIFLOW_BINARIES_WHEEL = "ciflow/binaries_wheel"
25

26

27
@dataclass
28
class CIFlowConfig:
29
    # For use to enable workflows to run on pytorch/pytorch-canary
30
    run_on_canary: bool = False
31
    labels: Set[str] = field(default_factory=set)
32
    # Certain jobs might not want to be part of the ciflow/[all,trunk] workflow
33
    isolated_workflow: bool = False
34
    unstable: bool = False
35

36
    def __post_init__(self) -> None:
37
        if not self.isolated_workflow:
38
            if LABEL_CIFLOW_PERIODIC not in self.labels:
39
                self.labels.add(
40
                    LABEL_CIFLOW_TRUNK if not self.unstable else LABEL_CIFLOW_UNSTABLE
41
                )
42

43

44
class Config(TypedDict):
45
    num_shards: int
46
    runner: str
47

48

49
@dataclass
50
class BinaryBuildWorkflow:
51
    os: str
52
    build_configs: List[Dict[str, str]]
53
    package_type: str
54

55
    # Optional fields
56
    build_environment: str = ""
57
    abi_version: str = ""
58
    ciflow_config: CIFlowConfig = field(default_factory=CIFlowConfig)
59
    is_scheduled: str = ""
60
    branches: str = "nightly"
61
    # Mainly for macos
62
    cross_compile_arm64: bool = False
63
    macos_runner: str = "macos-14-xlarge"
64

65
    def __post_init__(self) -> None:
66
        if self.abi_version:
67
            self.build_environment = (
68
                f"{self.os}-binary-{self.package_type}-{self.abi_version}"
69
            )
70
        else:
71
            self.build_environment = f"{self.os}-binary-{self.package_type}"
72

73
    def generate_workflow_file(self, workflow_template: jinja2.Template) -> None:
74
        output_file_path = (
75
            GITHUB_DIR
76
            / f"workflows/generated-{self.build_environment}-{self.branches}.yml"
77
        )
78
        with open(output_file_path, "w") as output_file:
79
            GENERATED = "generated"  # Note that please keep the variable GENERATED otherwise phabricator will hide the whole file
80
            output_file.writelines([f"# @{GENERATED} DO NOT EDIT MANUALLY\n"])
81
            try:
82
                content = workflow_template.render(asdict(self))
83
            except Exception as e:
84
                print(f"Failed on template: {workflow_template}", file=sys.stderr)
85
                raise e
86
            output_file.write(content)
87
            if content[-1] != "\n":
88
                output_file.write("\n")
89
        print(output_file_path)
90

91

92
class OperatingSystem:
93
    LINUX = "linux"
94
    WINDOWS = "windows"
95
    MACOS = "macos"
96
    MACOS_ARM64 = "macos-arm64"
97
    LINUX_AARCH64 = "linux-aarch64"
98
    LINUX_S390X = "linux-s390x"
99

100

101
LINUX_BINARY_BUILD_WORFKLOWS = [
102
    BinaryBuildWorkflow(
103
        os=OperatingSystem.LINUX,
104
        package_type="manywheel",
105
        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
106
            OperatingSystem.LINUX
107
        ),
108
        ciflow_config=CIFlowConfig(
109
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
110
            isolated_workflow=True,
111
        ),
112
    ),
113
    BinaryBuildWorkflow(
114
        os=OperatingSystem.LINUX,
115
        package_type="conda",
116
        build_configs=generate_binary_build_matrix.generate_conda_matrix(
117
            OperatingSystem.LINUX
118
        ),
119
        ciflow_config=CIFlowConfig(
120
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_CONDA},
121
            isolated_workflow=True,
122
        ),
123
    ),
124
    BinaryBuildWorkflow(
125
        os=OperatingSystem.LINUX,
126
        package_type="libtorch",
127
        abi_version=generate_binary_build_matrix.CXX11_ABI,
128
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
129
            OperatingSystem.LINUX,
130
            generate_binary_build_matrix.CXX11_ABI,
131
            libtorch_variants=["shared-with-deps"],
132
        ),
133
        ciflow_config=CIFlowConfig(
134
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
135
            isolated_workflow=True,
136
        ),
137
    ),
138
    BinaryBuildWorkflow(
139
        os=OperatingSystem.LINUX,
140
        package_type="libtorch",
141
        abi_version=generate_binary_build_matrix.PRE_CXX11_ABI,
142
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
143
            OperatingSystem.LINUX,
144
            generate_binary_build_matrix.PRE_CXX11_ABI,
145
            libtorch_variants=["shared-with-deps"],
146
        ),
147
        ciflow_config=CIFlowConfig(
148
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
149
            isolated_workflow=True,
150
        ),
151
    ),
152
]
153

154
LINUX_BINARY_SMOKE_WORKFLOWS = [
155
    BinaryBuildWorkflow(
156
        os=OperatingSystem.LINUX,
157
        package_type="manywheel",
158
        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
159
            OperatingSystem.LINUX,
160
            arches=["11.8", "12.1", "12.4"],
161
            python_versions=["3.9"],
162
        ),
163
        branches="main",
164
    ),
165
    BinaryBuildWorkflow(
166
        os=OperatingSystem.LINUX,
167
        package_type="libtorch",
168
        abi_version=generate_binary_build_matrix.CXX11_ABI,
169
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
170
            OperatingSystem.LINUX,
171
            generate_binary_build_matrix.CXX11_ABI,
172
            arches=["cpu"],
173
            libtorch_variants=["shared-with-deps"],
174
        ),
175
        branches="main",
176
    ),
177
    BinaryBuildWorkflow(
178
        os=OperatingSystem.LINUX,
179
        package_type="libtorch",
180
        abi_version=generate_binary_build_matrix.PRE_CXX11_ABI,
181
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
182
            OperatingSystem.LINUX,
183
            generate_binary_build_matrix.PRE_CXX11_ABI,
184
            arches=["cpu"],
185
            libtorch_variants=["shared-with-deps"],
186
        ),
187
        branches="main",
188
    ),
189
]
190

191
WINDOWS_BINARY_BUILD_WORKFLOWS = [
192
    BinaryBuildWorkflow(
193
        os=OperatingSystem.WINDOWS,
194
        package_type="wheel",
195
        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
196
            OperatingSystem.WINDOWS
197
        ),
198
        ciflow_config=CIFlowConfig(
199
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
200
            isolated_workflow=True,
201
        ),
202
    ),
203
    BinaryBuildWorkflow(
204
        os=OperatingSystem.WINDOWS,
205
        package_type="conda",
206
        build_configs=generate_binary_build_matrix.generate_conda_matrix(
207
            OperatingSystem.WINDOWS
208
        ),
209
        ciflow_config=CIFlowConfig(
210
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_CONDA},
211
            isolated_workflow=True,
212
        ),
213
    ),
214
    BinaryBuildWorkflow(
215
        os=OperatingSystem.WINDOWS,
216
        package_type="libtorch",
217
        abi_version=generate_binary_build_matrix.RELEASE,
218
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
219
            OperatingSystem.WINDOWS,
220
            generate_binary_build_matrix.RELEASE,
221
            libtorch_variants=["shared-with-deps"],
222
        ),
223
        ciflow_config=CIFlowConfig(
224
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
225
            isolated_workflow=True,
226
        ),
227
    ),
228
    BinaryBuildWorkflow(
229
        os=OperatingSystem.WINDOWS,
230
        package_type="libtorch",
231
        abi_version=generate_binary_build_matrix.DEBUG,
232
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
233
            OperatingSystem.WINDOWS,
234
            generate_binary_build_matrix.DEBUG,
235
            libtorch_variants=["shared-with-deps"],
236
        ),
237
        ciflow_config=CIFlowConfig(
238
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
239
            isolated_workflow=True,
240
        ),
241
    ),
242
]
243

244
WINDOWS_BINARY_SMOKE_WORKFLOWS = [
245
    BinaryBuildWorkflow(
246
        os=OperatingSystem.WINDOWS,
247
        package_type="libtorch",
248
        abi_version=generate_binary_build_matrix.RELEASE,
249
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
250
            OperatingSystem.WINDOWS,
251
            generate_binary_build_matrix.RELEASE,
252
            arches=["cpu"],
253
            libtorch_variants=["shared-with-deps"],
254
        ),
255
        branches="main",
256
        ciflow_config=CIFlowConfig(
257
            isolated_workflow=True,
258
        ),
259
    ),
260
    BinaryBuildWorkflow(
261
        os=OperatingSystem.WINDOWS,
262
        package_type="libtorch",
263
        abi_version=generate_binary_build_matrix.DEBUG,
264
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
265
            OperatingSystem.WINDOWS,
266
            generate_binary_build_matrix.DEBUG,
267
            arches=["cpu"],
268
            libtorch_variants=["shared-with-deps"],
269
        ),
270
        branches="main",
271
        ciflow_config=CIFlowConfig(
272
            isolated_workflow=True,
273
        ),
274
    ),
275
]
276

277
MACOS_BINARY_BUILD_WORKFLOWS = [
278
    BinaryBuildWorkflow(
279
        os=OperatingSystem.MACOS_ARM64,
280
        package_type="libtorch",
281
        abi_version=generate_binary_build_matrix.CXX11_ABI,
282
        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
283
            OperatingSystem.MACOS,
284
            generate_binary_build_matrix.CXX11_ABI,
285
            libtorch_variants=["shared-with-deps"],
286
        ),
287
        cross_compile_arm64=False,
288
        macos_runner="macos-14-xlarge",
289
        ciflow_config=CIFlowConfig(
290
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
291
            isolated_workflow=True,
292
        ),
293
    ),
294
    BinaryBuildWorkflow(
295
        os=OperatingSystem.MACOS_ARM64,
296
        package_type="wheel",
297
        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
298
            OperatingSystem.MACOS_ARM64
299
        ),
300
        cross_compile_arm64=False,
301
        macos_runner="macos-14-xlarge",
302
        ciflow_config=CIFlowConfig(
303
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
304
            isolated_workflow=True,
305
        ),
306
    ),
307
    BinaryBuildWorkflow(
308
        os=OperatingSystem.MACOS_ARM64,
309
        package_type="conda",
310
        cross_compile_arm64=False,
311
        macos_runner="macos-14-xlarge",
312
        build_configs=generate_binary_build_matrix.generate_conda_matrix(
313
            OperatingSystem.MACOS_ARM64
314
        ),
315
        ciflow_config=CIFlowConfig(
316
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_CONDA},
317
            isolated_workflow=True,
318
        ),
319
    ),
320
]
321

322
AARCH64_BINARY_BUILD_WORKFLOWS = [
323
    BinaryBuildWorkflow(
324
        os=OperatingSystem.LINUX_AARCH64,
325
        package_type="manywheel",
326
        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
327
            OperatingSystem.LINUX_AARCH64
328
        ),
329
        ciflow_config=CIFlowConfig(
330
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
331
            isolated_workflow=True,
332
        ),
333
    ),
334
]
335

336
S390X_BINARY_BUILD_WORKFLOWS = [
337
    BinaryBuildWorkflow(
338
        os=OperatingSystem.LINUX_S390X,
339
        package_type="manywheel",
340
        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
341
            OperatingSystem.LINUX_S390X
342
        ),
343
        ciflow_config=CIFlowConfig(
344
            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
345
            isolated_workflow=True,
346
        ),
347
    ),
348
]
349

350

351
def main() -> None:
352
    jinja_env = jinja2.Environment(
353
        variable_start_string="!{{",
354
        loader=jinja2.FileSystemLoader(str(GITHUB_DIR.joinpath("templates"))),
355
        undefined=jinja2.StrictUndefined,
356
    )
357

358
    # not ported yet
359
    template_and_workflows = [
360
        (
361
            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
362
            LINUX_BINARY_BUILD_WORFKLOWS,
363
        ),
364
        (
365
            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
366
            AARCH64_BINARY_BUILD_WORKFLOWS,
367
        ),
368
        (
369
            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
370
            S390X_BINARY_BUILD_WORKFLOWS,
371
        ),
372
        (
373
            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
374
            LINUX_BINARY_SMOKE_WORKFLOWS,
375
        ),
376
        (
377
            jinja_env.get_template("windows_binary_build_workflow.yml.j2"),
378
            WINDOWS_BINARY_BUILD_WORKFLOWS,
379
        ),
380
        (
381
            jinja_env.get_template("windows_binary_build_workflow.yml.j2"),
382
            WINDOWS_BINARY_SMOKE_WORKFLOWS,
383
        ),
384
        (
385
            jinja_env.get_template("macos_binary_build_workflow.yml.j2"),
386
            MACOS_BINARY_BUILD_WORKFLOWS,
387
        ),
388
    ]
389
    # Delete the existing generated files first, this should align with .gitattributes file description.
390
    existing_workflows = GITHUB_DIR.glob("workflows/generated-*")
391
    for w in existing_workflows:
392
        try:
393
            os.remove(w)
394
        except Exception as e:
395
            print(f"Error occurred when deleting file {w}: {e}")
396

397
    for template, workflows in template_and_workflows:
398
        # added Iterable check to appease the mypy gods
399
        if not isinstance(workflows, Iterable):
400
            raise Exception(  # noqa: TRY002
401
                f"How is workflows not iterable? {workflows}"
402
            )  # noqa: TRY002
403
        for workflow in workflows:
404
            workflow.generate_workflow_file(workflow_template=template)
405

406

407
if __name__ == "__main__":
408
    main()
409