Simplify .github/scripts/generate_ci_workflows.py (#58491)

Summary:
This PR simplifies `.github/scripts/generate_ci_workflows.py` by using the same strategy as https://github.com/pytorch/pytorch/issues/54344, representing workflows as plain data to avoid duplicating the definition of the `generate_workflow_file` function. This will make the script easier to maintain if/when that function is modified and/or more workflow types are added.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/58491

Test Plan:
The Lint job in CI; specifically:
```
make generate-gha-workflows
mypy --config mypy-strict.ini
```

Reviewed By: malfet, seemethere

Differential Revision: D28511918

Pulled By: samestep

fbshipit-source-id: aaf415a954d938a29aee7c9367c9bc2b9f44bb01

.github/scripts/generate_ci_workflows.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 from pathlib import Path
+from typing import Any, Dict
 
 import jinja2
 
@@ -9,77 +10,63 @@ DOCKER_REGISTRY = "308535385114.dkr.ecr.us-east-1.amazonaws.com"
 GITHUB_DIR = Path(__file__).parent.parent
 
 
-class PyTorchWindowsWorkflow:
+# it would be nice to statically specify that build_environment must be
+# present, but currently Python has no easy way to do that
+# https://github.com/python/mypy/issues/4617
+PyTorchWorkflow = Dict[str, Any]
+
+
+def PyTorchWindowsWorkflow(
+    *,
+    build_environment: str,
+    on_pull_request: bool = False
+) -> PyTorchWorkflow:
     CPU_TEST_RUNNER = "windows.4xlarge"
     CUDA_TEST_RUNNER = "windows.8xlarge.nvidia.gpu"
+    return {
+        "build_environment": build_environment,
+        "test_runner_type": (
+            CUDA_TEST_RUNNER
+            if "cuda" in build_environment
+            else CPU_TEST_RUNNER
+        ),
+        "on_pull_request": on_pull_request,
+    }
 
-    def __init__(
-            self,
-            build_environment: str,
-            on_pull_request: bool = False
-    ):
-        self.build_environment = build_environment
-        self.test_runner_type = self.CPU_TEST_RUNNER
-        self.on_pull_request = on_pull_request
-        if "cuda" in build_environment:
-            self.test_runner_type = self.CUDA_TEST_RUNNER
 
-    def generate_workflow_file(
+def PyTorchLinuxWorkflow(
-        self, workflow_template: jinja2.Template,
+    *,
-    ) -> Path:
+    build_environment: str,
-        output_file_path = GITHUB_DIR.joinpath(
+    docker_image_base: str,
-            f"workflows/{self.build_environment}.yml"
+    on_pull_request: bool = False,
-        )
+    enable_doc_jobs: bool = False,
-        with open(output_file_path, "w") as output_file:
+) -> PyTorchWorkflow:
-            output_file.writelines(["# @generated DO NOT EDIT MANUALLY\n"])
-            output_file.write(
-                workflow_template.render(
-                    build_environment=self.build_environment,
-                    test_runner_type=self.test_runner_type,
-                    on_pull_request=self.on_pull_request
-                )
-            )
-            output_file.write('\n')
-        return output_file_path
-
-class PyTorchLinuxWorkflow:
     CPU_TEST_RUNNER = "linux.2xlarge"
     CUDA_TEST_RUNNER = "linux.8xlarge.nvidia.gpu"
+    return {
+        "build_environment": build_environment,
+        "docker_image_base": docker_image_base,
+        "test_runner_type": (
+            CUDA_TEST_RUNNER
+            if "cuda" in build_environment
+            else CPU_TEST_RUNNER
+        ),
+        "on_pull_request": on_pull_request,
+        "enable_doc_jobs": enable_doc_jobs,
+    }
 
-    def __init__(
-            self,
-            build_environment: str,
-            docker_image_base: str,
-            on_pull_request: bool = False,
-            enable_doc_jobs: bool = False,
-    ):
-        self.build_environment = build_environment
-        self.docker_image_base = docker_image_base
-        self.test_runner_type = self.CPU_TEST_RUNNER
-        self.on_pull_request = on_pull_request
-        self.enable_doc_jobs = enable_doc_jobs
-        if "cuda" in build_environment:
-            self.test_runner_type = self.CUDA_TEST_RUNNER
 
-    def generate_workflow_file(
+def generate_workflow_file(
-        self, workflow_template: jinja2.Template,
+    *,
-    ) -> Path:
+    workflow: PyTorchWorkflow,
-        output_file_path = GITHUB_DIR.joinpath(
+    workflow_template: jinja2.Template,
-            f"workflows/{self.build_environment}.yml"
+) -> Path:
-        )
+    output_file_path = GITHUB_DIR / f"workflows/{workflow['build_environment']}.yml"
-        with open(output_file_path, "w") as output_file:
+    with open(output_file_path, "w") as output_file:
-            output_file.writelines(["# @generated DO NOT EDIT MANUALLY\n"])
+        output_file.writelines(["# @generated DO NOT EDIT MANUALLY\n"])
-            output_file.write(
+        output_file.write(workflow_template.render(**workflow))
-                workflow_template.render(
+        output_file.write("\n")
-                    build_environment=self.build_environment,
+    return output_file_path
-                    docker_image_base=self.docker_image_base,
-                    test_runner_type=self.test_runner_type,
-                    enable_doc_jobs=self.enable_doc_jobs,
-                    on_pull_request=self.on_pull_request,
-                )
-            )
-            output_file.write('\n')
-        return output_file_path
 
 
 WINDOWS_WORKFLOWS = [
@@ -196,4 +183,4 @@ if __name__ == "__main__":
     ]
     for template, workflows in template_and_workflows:
         for workflow in workflows:
-            print(workflow.generate_workflow_file(workflow_template=template))
+            print(generate_workflow_file(workflow=workflow, workflow_template=template))

mypy-strict.ini

@@ -37,6 +37,7 @@ strict_equality = True
 
 files =
     .github/scripts/generate_binary_build_matrix.py,
+    .github/scripts/generate_ci_workflows.py,
     .github/scripts/parse_ref.py,
     benchmarks/instruction_counts,
     tools/actions_local_runner.py,