diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index 06abb347091..a125a2a37aa 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.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( - self, workflow_template: jinja2.Template, - ) -> Path: - output_file_path = GITHUB_DIR.joinpath( - f"workflows/{self.build_environment}.yml" - ) - with open(output_file_path, "w") as output_file: - 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: +def PyTorchLinuxWorkflow( + *, + build_environment: str, + docker_image_base: str, + on_pull_request: bool = False, + enable_doc_jobs: bool = False, +) -> PyTorchWorkflow: 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( - self, workflow_template: jinja2.Template, - ) -> Path: - output_file_path = GITHUB_DIR.joinpath( - f"workflows/{self.build_environment}.yml" - ) - with open(output_file_path, "w") as output_file: - output_file.writelines(["# @generated DO NOT EDIT MANUALLY\n"]) - output_file.write( - workflow_template.render( - build_environment=self.build_environment, - 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 +def generate_workflow_file( + *, + workflow: PyTorchWorkflow, + workflow_template: jinja2.Template, +) -> Path: + output_file_path = GITHUB_DIR / f"workflows/{workflow['build_environment']}.yml" + with open(output_file_path, "w") as output_file: + output_file.writelines(["# @generated DO NOT EDIT MANUALLY\n"]) + output_file.write(workflow_template.render(**workflow)) + 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)) diff --git a/mypy-strict.ini b/mypy-strict.ini index c46a482bda8..3ac5b6055f3 100644 --- a/mypy-strict.ini +++ b/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,