pytorch

1
#!/usr/bin/env python3
2

3
import argparse
4
import fnmatch
5
import pathlib
6
import subprocess
7
import textwrap
8

9
from typing import Any, Dict, List
10

11
import yaml
12

13

14
REPO_ROOT = pathlib.Path(__file__).parent.parent.parent
15
CONFIG_YML = REPO_ROOT / ".circleci" / "config.yml"
16
WORKFLOWS_DIR = REPO_ROOT / ".github" / "workflows"
17

18

19
WORKFLOWS_TO_CHECK = [
20
    "binary_builds",
21
    "build",
22
    "master_build",
23
    # These are formatted slightly differently, skip them
24
    # "scheduled-ci",
25
    # "debuggable-scheduled-ci",
26
    # "slow-gradcheck-scheduled-ci",
27
    # "promote",
28
]
29

30

31
def add_job(
32
    workflows: Dict[str, Any],
33
    workflow_name: str,
34
    type: str,
35
    job: Dict[str, Any],
36
    past_jobs: Dict[str, Any],
37
) -> None:
38
    """
39
    Add job 'job' under 'type' and 'workflow_name' to 'workflow' in place. Also
40
    add any dependencies (they must already be in 'past_jobs')
41
    """
42
    if workflow_name not in workflows:
43
        workflows[workflow_name] = {"when": "always", "jobs": []}
44

45
    requires = job.get("requires", None)
46
    if requires is not None:
47
        for requirement in requires:
48
            dependency = past_jobs[requirement]
49
            add_job(
50
                workflows,
51
                dependency["workflow_name"],
52
                dependency["type"],
53
                dependency["job"],
54
                past_jobs,
55
            )
56

57
    workflows[workflow_name]["jobs"].append({type: job})
58

59

60
def get_filtered_circleci_config(
61
    workflows: Dict[str, Any], relevant_jobs: List[str]
62
) -> Dict[str, Any]:
63
    """
64
    Given an existing CircleCI config, remove every job that's not listed in
65
    'relevant_jobs'
66
    """
67
    new_workflows: Dict[str, Any] = {}
68
    past_jobs: Dict[str, Any] = {}
69
    for workflow_name, workflow in workflows.items():
70
        if workflow_name not in WORKFLOWS_TO_CHECK:
71
            # Don't care about this workflow, skip it entirely
72
            continue
73

74
        for job_dict in workflow["jobs"]:
75
            for type, job in job_dict.items():
76
                if "name" not in job:
77
                    # Job doesn't have a name so it can't be handled
78
                    print("Skipping", type)
79
                else:
80
                    if job["name"] in relevant_jobs:
81
                        # Found a job that was specified at the CLI, add it to
82
                        # the new result
83
                        add_job(new_workflows, workflow_name, type, job, past_jobs)
84

85
                    # Record the job in case it's needed as a dependency later
86
                    past_jobs[job["name"]] = {
87
                        "workflow_name": workflow_name,
88
                        "type": type,
89
                        "job": job,
90
                    }
91

92
    return new_workflows
93

94

95
def commit_ci(files: List[str], message: str) -> None:
96
    # Check that there are no other modified files than the ones edited by this
97
    # tool
98
    stdout = subprocess.run(
99
        ["git", "status", "--porcelain"], stdout=subprocess.PIPE
100
    ).stdout.decode()
101
    for line in stdout.split("\n"):
102
        if line == "":
103
            continue
104
        if line[0] != " ":
105
            raise RuntimeError(
106
                f"Refusing to commit while other changes are already staged: {line}"
107
            )
108

109
    # Make the commit
110
    subprocess.run(["git", "add"] + files)
111
    subprocess.run(["git", "commit", "-m", message])
112

113

114
if __name__ == "__main__":
115
    parser = argparse.ArgumentParser(
116
        description="make .circleci/config.yml only have a specific set of jobs and delete GitHub actions"
117
    )
118
    parser.add_argument("--job", action="append", help="job name", default=[])
119
    parser.add_argument(
120
        "--filter-gha", help="keep only these github actions (glob match)", default=""
121
    )
122
    parser.add_argument(
123
        "--make-commit",
124
        action="store_true",
125
        help="add change to git with to a do-not-merge commit",
126
    )
127
    args = parser.parse_args()
128

129
    touched_files = [CONFIG_YML]
130
    with open(CONFIG_YML) as f:
131
        config_yml = yaml.safe_load(f.read())
132

133
    config_yml["workflows"] = get_filtered_circleci_config(
134
        config_yml["workflows"], args.job
135
    )
136

137
    with open(CONFIG_YML, "w") as f:
138
        yaml.dump(config_yml, f)
139

140
    if args.filter_gha:
141
        for relative_file in WORKFLOWS_DIR.iterdir():
142
            path = REPO_ROOT.joinpath(relative_file)
143
            if not fnmatch.fnmatch(path.name, args.filter_gha):
144
                touched_files.append(path)
145
                path.resolve().unlink()
146

147
    if args.make_commit:
148
        jobs_str = "\n".join([f" * {job}" for job in args.job])
149
        message = textwrap.dedent(
150
            f"""
151
        [skip ci][do not merge] Edit config.yml to filter specific jobs
152

153
        Filter CircleCI to only run:
154
        {jobs_str}
155

156
        See [Run Specific CI Jobs](https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md#run-specific-ci-jobs) for details.
157
        """
158
        ).strip()
159
        commit_ci([str(f.relative_to(REPO_ROOT)) for f in touched_files], message)
160