• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2
3import os
4import sys
5from dataclasses import asdict, dataclass, field
6from pathlib import Path
7from typing import Dict, Iterable, List, Literal, Set
8from typing_extensions import TypedDict  # Python 3.11+
9
10import generate_binary_build_matrix  # type: ignore[import]
11import jinja2
12
13
14Arch = Literal["windows", "linux", "macos"]
15
16GITHUB_DIR = Path(__file__).resolve().parent.parent
17
18LABEL_CIFLOW_TRUNK = "ciflow/trunk"
19LABEL_CIFLOW_UNSTABLE = "ciflow/unstable"
20LABEL_CIFLOW_BINARIES = "ciflow/binaries"
21LABEL_CIFLOW_PERIODIC = "ciflow/periodic"
22LABEL_CIFLOW_BINARIES_LIBTORCH = "ciflow/binaries_libtorch"
23LABEL_CIFLOW_BINARIES_CONDA = "ciflow/binaries_conda"
24LABEL_CIFLOW_BINARIES_WHEEL = "ciflow/binaries_wheel"
25
26
27@dataclass
28class 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
44class Config(TypedDict):
45    num_shards: int
46    runner: str
47
48
49@dataclass
50class 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    use_split_build: bool = False
65
66    def __post_init__(self) -> None:
67        if self.abi_version:
68            self.build_environment = (
69                f"{self.os}-binary-{self.package_type}-{self.abi_version}"
70            )
71        else:
72            self.build_environment = f"{self.os}-binary-{self.package_type}"
73        if self.use_split_build:
74            # added to distinguish concurrency groups
75            self.build_environment += "-split"
76
77    def generate_workflow_file(self, workflow_template: jinja2.Template) -> None:
78        output_file_path = (
79            GITHUB_DIR
80            / f"workflows/generated-{self.build_environment}-{self.branches}.yml"
81        )
82        if self.use_split_build:
83            output_file_path = (
84                GITHUB_DIR
85                / f"workflows/generated-{self.build_environment}-{self.branches}"
86            )
87        with open(output_file_path, "w") as output_file:
88            GENERATED = "generated"  # Note that please keep the variable GENERATED otherwise phabricator will hide the whole file
89            output_file.writelines([f"# @{GENERATED} DO NOT EDIT MANUALLY\n"])
90            try:
91                content = workflow_template.render(asdict(self))
92            except Exception as e:
93                print(f"Failed on template: {workflow_template}", file=sys.stderr)
94                raise e
95            output_file.write(content)
96            if content[-1] != "\n":
97                output_file.write("\n")
98        print(output_file_path)
99
100
101class OperatingSystem:
102    LINUX = "linux"
103    WINDOWS = "windows"
104    MACOS = "macos"
105    MACOS_ARM64 = "macos-arm64"
106    LINUX_AARCH64 = "linux-aarch64"
107    LINUX_S390X = "linux-s390x"
108
109
110LINUX_BINARY_BUILD_WORFKLOWS = [
111    BinaryBuildWorkflow(
112        os=OperatingSystem.LINUX,
113        package_type="manywheel",
114        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
115            OperatingSystem.LINUX
116        ),
117        ciflow_config=CIFlowConfig(
118            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
119            isolated_workflow=True,
120        ),
121    ),
122    BinaryBuildWorkflow(
123        os=OperatingSystem.LINUX,
124        package_type="manywheel",
125        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
126            OperatingSystem.LINUX,
127            use_split_build=True,
128            arches=["11.8", "12.1", "12.4", "cpu"],
129        ),
130        ciflow_config=CIFlowConfig(
131            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
132            isolated_workflow=True,
133        ),
134        use_split_build=True,
135    ),
136    BinaryBuildWorkflow(
137        os=OperatingSystem.LINUX,
138        package_type="conda",
139        build_configs=generate_binary_build_matrix.generate_conda_matrix(
140            OperatingSystem.LINUX
141        ),
142        ciflow_config=CIFlowConfig(
143            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_CONDA},
144            isolated_workflow=True,
145        ),
146    ),
147    BinaryBuildWorkflow(
148        os=OperatingSystem.LINUX,
149        package_type="libtorch",
150        abi_version=generate_binary_build_matrix.CXX11_ABI,
151        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
152            OperatingSystem.LINUX,
153            generate_binary_build_matrix.CXX11_ABI,
154            libtorch_variants=["shared-with-deps"],
155        ),
156        ciflow_config=CIFlowConfig(
157            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
158            isolated_workflow=True,
159        ),
160    ),
161    BinaryBuildWorkflow(
162        os=OperatingSystem.LINUX,
163        package_type="libtorch",
164        abi_version=generate_binary_build_matrix.PRE_CXX11_ABI,
165        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
166            OperatingSystem.LINUX,
167            generate_binary_build_matrix.PRE_CXX11_ABI,
168            libtorch_variants=["shared-with-deps"],
169        ),
170        ciflow_config=CIFlowConfig(
171            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
172            isolated_workflow=True,
173        ),
174    ),
175]
176
177LINUX_BINARY_SMOKE_WORKFLOWS = [
178    BinaryBuildWorkflow(
179        os=OperatingSystem.LINUX,
180        package_type="manywheel",
181        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
182            OperatingSystem.LINUX,
183            arches=["11.8", "12.1", "12.4"],
184            python_versions=["3.9"],
185        ),
186        branches="main",
187    ),
188    BinaryBuildWorkflow(
189        os=OperatingSystem.LINUX,
190        package_type="manywheel",
191        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
192            OperatingSystem.LINUX,
193            arches=["11.8", "12.1", "12.4"],
194            python_versions=["3.9"],
195            use_split_build=True,
196        ),
197        ciflow_config=CIFlowConfig(
198            labels={LABEL_CIFLOW_PERIODIC},
199        ),
200        branches="main",
201        use_split_build=True,
202    ),
203    BinaryBuildWorkflow(
204        os=OperatingSystem.LINUX,
205        package_type="libtorch",
206        abi_version=generate_binary_build_matrix.CXX11_ABI,
207        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
208            OperatingSystem.LINUX,
209            generate_binary_build_matrix.CXX11_ABI,
210            arches=["cpu"],
211            libtorch_variants=["shared-with-deps"],
212        ),
213        branches="main",
214    ),
215    BinaryBuildWorkflow(
216        os=OperatingSystem.LINUX,
217        package_type="libtorch",
218        abi_version=generate_binary_build_matrix.PRE_CXX11_ABI,
219        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
220            OperatingSystem.LINUX,
221            generate_binary_build_matrix.PRE_CXX11_ABI,
222            arches=["cpu"],
223            libtorch_variants=["shared-with-deps"],
224        ),
225        branches="main",
226    ),
227]
228
229WINDOWS_BINARY_BUILD_WORKFLOWS = [
230    BinaryBuildWorkflow(
231        os=OperatingSystem.WINDOWS,
232        package_type="wheel",
233        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
234            OperatingSystem.WINDOWS
235        ),
236        ciflow_config=CIFlowConfig(
237            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
238            isolated_workflow=True,
239        ),
240    ),
241    BinaryBuildWorkflow(
242        os=OperatingSystem.WINDOWS,
243        package_type="conda",
244        build_configs=generate_binary_build_matrix.generate_conda_matrix(
245            OperatingSystem.WINDOWS
246        ),
247        ciflow_config=CIFlowConfig(
248            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_CONDA},
249            isolated_workflow=True,
250        ),
251    ),
252    BinaryBuildWorkflow(
253        os=OperatingSystem.WINDOWS,
254        package_type="libtorch",
255        abi_version=generate_binary_build_matrix.RELEASE,
256        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
257            OperatingSystem.WINDOWS,
258            generate_binary_build_matrix.RELEASE,
259            libtorch_variants=["shared-with-deps"],
260        ),
261        ciflow_config=CIFlowConfig(
262            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
263            isolated_workflow=True,
264        ),
265    ),
266    BinaryBuildWorkflow(
267        os=OperatingSystem.WINDOWS,
268        package_type="libtorch",
269        abi_version=generate_binary_build_matrix.DEBUG,
270        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
271            OperatingSystem.WINDOWS,
272            generate_binary_build_matrix.DEBUG,
273            libtorch_variants=["shared-with-deps"],
274        ),
275        ciflow_config=CIFlowConfig(
276            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
277            isolated_workflow=True,
278        ),
279    ),
280]
281
282WINDOWS_BINARY_SMOKE_WORKFLOWS = [
283    BinaryBuildWorkflow(
284        os=OperatingSystem.WINDOWS,
285        package_type="libtorch",
286        abi_version=generate_binary_build_matrix.RELEASE,
287        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
288            OperatingSystem.WINDOWS,
289            generate_binary_build_matrix.RELEASE,
290            arches=["cpu"],
291            libtorch_variants=["shared-with-deps"],
292        ),
293        branches="main",
294        ciflow_config=CIFlowConfig(
295            isolated_workflow=True,
296        ),
297    ),
298    BinaryBuildWorkflow(
299        os=OperatingSystem.WINDOWS,
300        package_type="libtorch",
301        abi_version=generate_binary_build_matrix.DEBUG,
302        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
303            OperatingSystem.WINDOWS,
304            generate_binary_build_matrix.DEBUG,
305            arches=["cpu"],
306            libtorch_variants=["shared-with-deps"],
307        ),
308        branches="main",
309        ciflow_config=CIFlowConfig(
310            isolated_workflow=True,
311        ),
312    ),
313]
314
315MACOS_BINARY_BUILD_WORKFLOWS = [
316    BinaryBuildWorkflow(
317        os=OperatingSystem.MACOS_ARM64,
318        package_type="libtorch",
319        abi_version=generate_binary_build_matrix.CXX11_ABI,
320        build_configs=generate_binary_build_matrix.generate_libtorch_matrix(
321            OperatingSystem.MACOS,
322            generate_binary_build_matrix.CXX11_ABI,
323            libtorch_variants=["shared-with-deps"],
324        ),
325        cross_compile_arm64=False,
326        macos_runner="macos-14-xlarge",
327        ciflow_config=CIFlowConfig(
328            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_LIBTORCH},
329            isolated_workflow=True,
330        ),
331    ),
332    BinaryBuildWorkflow(
333        os=OperatingSystem.MACOS_ARM64,
334        package_type="wheel",
335        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
336            OperatingSystem.MACOS_ARM64
337        ),
338        cross_compile_arm64=False,
339        macos_runner="macos-14-xlarge",
340        ciflow_config=CIFlowConfig(
341            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
342            isolated_workflow=True,
343        ),
344    ),
345    BinaryBuildWorkflow(
346        os=OperatingSystem.MACOS_ARM64,
347        package_type="conda",
348        cross_compile_arm64=False,
349        macos_runner="macos-14-xlarge",
350        build_configs=generate_binary_build_matrix.generate_conda_matrix(
351            OperatingSystem.MACOS_ARM64
352        ),
353        ciflow_config=CIFlowConfig(
354            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_CONDA},
355            isolated_workflow=True,
356        ),
357    ),
358]
359
360AARCH64_BINARY_BUILD_WORKFLOWS = [
361    BinaryBuildWorkflow(
362        os=OperatingSystem.LINUX_AARCH64,
363        package_type="manywheel",
364        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
365            OperatingSystem.LINUX_AARCH64
366        ),
367        ciflow_config=CIFlowConfig(
368            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
369            isolated_workflow=True,
370        ),
371    ),
372]
373
374S390X_BINARY_BUILD_WORKFLOWS = [
375    BinaryBuildWorkflow(
376        os=OperatingSystem.LINUX_S390X,
377        package_type="manywheel",
378        build_configs=generate_binary_build_matrix.generate_wheels_matrix(
379            OperatingSystem.LINUX_S390X
380        ),
381        ciflow_config=CIFlowConfig(
382            labels={LABEL_CIFLOW_BINARIES, LABEL_CIFLOW_BINARIES_WHEEL},
383            isolated_workflow=True,
384        ),
385    ),
386]
387
388
389def main() -> None:
390    jinja_env = jinja2.Environment(
391        variable_start_string="!{{",
392        loader=jinja2.FileSystemLoader(str(GITHUB_DIR.joinpath("templates"))),
393        undefined=jinja2.StrictUndefined,
394    )
395
396    # not ported yet
397    template_and_workflows = [
398        (
399            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
400            LINUX_BINARY_BUILD_WORFKLOWS,
401        ),
402        (
403            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
404            AARCH64_BINARY_BUILD_WORKFLOWS,
405        ),
406        (
407            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
408            S390X_BINARY_BUILD_WORKFLOWS,
409        ),
410        (
411            jinja_env.get_template("linux_binary_build_workflow.yml.j2"),
412            LINUX_BINARY_SMOKE_WORKFLOWS,
413        ),
414        (
415            jinja_env.get_template("windows_binary_build_workflow.yml.j2"),
416            WINDOWS_BINARY_BUILD_WORKFLOWS,
417        ),
418        (
419            jinja_env.get_template("windows_binary_build_workflow.yml.j2"),
420            WINDOWS_BINARY_SMOKE_WORKFLOWS,
421        ),
422        (
423            jinja_env.get_template("macos_binary_build_workflow.yml.j2"),
424            MACOS_BINARY_BUILD_WORKFLOWS,
425        ),
426    ]
427    # Delete the existing generated files first, this should align with .gitattributes file description.
428    existing_workflows = GITHUB_DIR.glob("workflows/generated-*")
429    for w in existing_workflows:
430        try:
431            os.remove(w)
432        except Exception as e:
433            print(f"Error occurred when deleting file {w}: {e}")
434
435    for template, workflows in template_and_workflows:
436        # added Iterable check to appease the mypy gods
437        if not isinstance(workflows, Iterable):
438            raise Exception(  # noqa: TRY002
439                f"How is workflows not iterable? {workflows}"
440            )  # noqa: TRY002
441        for workflow in workflows:
442            workflow.generate_workflow_file(workflow_template=template)
443
444
445if __name__ == "__main__":
446    main()
447