pytorch

1
#!/usr/bin/env python3
2

3
import contextlib
4
import os
5
import re
6
import subprocess
7
import sys
8
from typing import Any, Generator
9

10
from github_utils import gh_post_pr_comment as gh_post_comment
11
from gitutils import get_git_remote_name, get_git_repo_dir, GitRepo
12
from trymerge import GitHubPR
13

14
SAME_SHA_ERROR = (
15
    "\n```\nAborting rebase because rebasing the branch resulted in the same sha as the target branch.\n"
16
    + "This usually happens because the PR has already been merged.  Please rebase locally and push.\n```"
17
)
18

19

20
def parse_args() -> Any:
21
    from argparse import ArgumentParser
22

23
    parser = ArgumentParser("Rebase PR into branch")
24
    parser.add_argument("--dry-run", action="store_true")
25
    parser.add_argument("--branch", type=str)
26
    parser.add_argument("pr_num", type=int)
27
    return parser.parse_args()
28

29

30
def post_already_uptodate(
31
    pr: GitHubPR, repo: GitRepo, onto_branch: str, dry_run: bool
32
) -> None:
33
    msg = f"Tried to rebase and push PR #{pr.pr_num}, but it was already up to date."
34
    def_branch = pr.default_branch()
35
    def_branch_fcn = f"refs/remotes/{repo.remote}/{def_branch}"
36
    if onto_branch != def_branch_fcn and repo.rev_parse(
37
        def_branch_fcn
38
    ) != repo.rev_parse(onto_branch):
39
        def_branch_url = f"https://github.com/{pr.org}/{pr.project}/tree/{def_branch}"
40
        msg += f" Try rebasing against [{def_branch}]({def_branch_url}) by issuing:"
41
        msg += f"\n`@pytorchbot rebase -b {def_branch}`"
42

43
    gh_post_comment(
44
        pr.org,
45
        pr.project,
46
        pr.pr_num,
47
        msg,
48
        dry_run=dry_run,
49
    )
50

51

52
def rebase_onto(
53
    pr: GitHubPR, repo: GitRepo, onto_branch: str, dry_run: bool = False
54
) -> bool:
55
    branch = f"pull/{pr.pr_num}/head"
56
    remote_url = f"https://github.com/{pr.info['headRepository']['nameWithOwner']}.git"
57
    refspec = f"{branch}:{pr.head_ref()}"
58

59
    repo.fetch(branch, branch)
60
    repo._run_git("rebase", onto_branch, branch)
61

62
    if repo.rev_parse(branch) == repo.rev_parse(onto_branch):
63
        raise Exception(SAME_SHA_ERROR)
64

65
    if dry_run:
66
        push_result = repo._run_git("push", "--dry-run", "-f", remote_url, refspec)
67
    else:
68
        push_result = repo._run_git("push", "-f", remote_url, refspec)
69
    if "Everything up-to-date" in push_result:
70
        post_already_uptodate(pr, repo, onto_branch, dry_run)
71
        return False
72
    else:
73
        gh_post_comment(
74
            pr.org,
75
            pr.project,
76
            pr.pr_num,
77
            f"Successfully rebased `{pr.head_ref()}` onto `{onto_branch}`, please pull locally "
78
            + f"before adding more changes (for example, via `git checkout {pr.head_ref()} && "
79
            + "git pull --rebase`)",
80
            dry_run=dry_run,
81
        )
82
        return True
83

84

85
def rebase_ghstack_onto(
86
    pr: GitHubPR, repo: GitRepo, onto_branch: str, dry_run: bool = False
87
) -> bool:
88
    if (
89
        subprocess.run(
90
            [sys.executable, "-m", "ghstack", "--help"],
91
            capture_output=True,
92
            check=False,
93
        ).returncode
94
        != 0
95
    ):
96
        subprocess.run([sys.executable, "-m", "pip", "install", "ghstack"], check=True)
97
    orig_ref = f"{re.sub(r'/head$', '/orig', pr.head_ref())}"
98

99
    repo.fetch(orig_ref, orig_ref)
100
    repo._run_git("rebase", onto_branch, orig_ref)
101

102
    if repo.rev_parse(orig_ref) == repo.rev_parse(onto_branch):
103
        raise Exception(SAME_SHA_ERROR)
104

105
    # steal the identity of the committer of the commit on the orig branch
106
    email = repo._run_git("log", orig_ref, "--pretty=format:%ae", "-1")
107
    name = repo._run_git("log", orig_ref, "--pretty=format:%an", "-1")
108
    repo._run_git("config", "--global", "user.email", email)
109
    repo._run_git("config", "--global", "user.name", name)
110

111
    os.environ["OAUTH_TOKEN"] = os.environ["GITHUB_TOKEN"]
112
    with open(".ghstackrc", "w+") as f:
113
        f.write(
114
            "[ghstack]\n"
115
            + "github_url=github.com\n"
116
            + "github_username=pytorchmergebot\n"
117
            + "remote_name=origin"
118
        )
119

120
    if dry_run:
121
        print("Don't know how to dry-run ghstack")
122
        return False
123
    else:
124
        ghstack_result = subprocess.run(["ghstack"], capture_output=True, check=True)
125
        push_result = ghstack_result.stdout.decode("utf-8")
126
        print(push_result)
127
        if ghstack_result.returncode != 0:
128
            print(ghstack_result.stderr.decode("utf-8"))
129
            raise Exception(f"\n```{push_result}```")
130
        # The contents of a successful push result should look like:
131
        # Summary of changes (ghstack 0.6.0)
132

133
        #  - Updated https://github.com/clee2000/random-testing/pull/2
134
        #  - Updated https://github.com/clee2000/random-testing/pull/1
135

136
        # Facebook employees can import your changes by running
137
        # (on a Facebook machine):
138

139
        #     ghimport -s https://github.com/clee2000/random-testing/pull/2
140

141
        # If you want to work on this diff stack on another machine:
142

143
        #     ghstack checkout https://github.com/clee2000/random-testing/pull/2
144
        org, project = repo.gh_owner_and_name()
145
        for line in push_result.splitlines():
146
            if "Updated" in line:
147
                pr_num = int(line.split("/")[-1])
148
                if pr_num != pr.pr_num:
149
                    gh_post_comment(
150
                        pr.org,
151
                        pr.project,
152
                        pr_num,
153
                        f"Rebased `{orig_ref}` onto `{onto_branch}` because #{pr.pr_num} was rebased, "
154
                        "please pull locally before adding more changes (for example, via `ghstack "
155
                        + f"checkout https://github.com/{org}/{project}/pull/{pr_num}`)",
156
                        dry_run=dry_run,
157
                    )
158
                else:
159
                    gh_post_comment(
160
                        pr.org,
161
                        pr.project,
162
                        pr_num,
163
                        f"Successfully rebased `{orig_ref}` onto `{onto_branch}`, please pull locally "
164
                        + "before adding more changes (for example, via `ghstack "
165
                        + f"checkout https://github.com/{org}/{project}/pull/{pr.pr_num}`)",
166
                        dry_run=dry_run,
167
                    )
168

169
        if (
170
            f"Skipped https://github.com/{org}/{project}/pull/{pr.pr_num}"
171
            in push_result
172
        ):
173
            post_already_uptodate(pr, repo, onto_branch, dry_run)
174
            return False
175
        return True
176

177

178
def additional_rebase_failure_info(e: Exception) -> str:
179
    if re.search(
180
        r"remote: Permission to .* denied to .*\.\nfatal: unable to access", str(e)
181
    ):
182
        return (
183
            "\nThis is likely because the author did not allow edits from maintainers on the PR or because the "
184
            "repo has additional permissions settings that mergebot does not qualify."
185
        )
186
    return ""
187

188

189
@contextlib.contextmanager
190
def git_config_guard(repo: GitRepo) -> Generator[None, None, None]:
191
    """Restores user.name and user.email global properties after context is finished"""
192
    user_email = repo._run_git("config", "user.email")
193
    user_name = repo._run_git("config", "user.name")
194
    try:
195
        yield
196
    finally:
197
        if user_email:
198
            repo._run_git("config", "--global", "user.email", user_email)
199
        if user_name:
200
            repo._run_git("config", "--global", "user.name", user_name)
201

202

203
def main() -> None:
204
    args = parse_args()
205
    repo = GitRepo(get_git_repo_dir(), get_git_remote_name(), debug=True)
206
    org, project = repo.gh_owner_and_name()
207

208
    pr = GitHubPR(org, project, args.pr_num)
209
    onto_branch = args.branch if args.branch else pr.default_branch()
210
    onto_branch = f"refs/remotes/{repo.remote}/{onto_branch}"
211
    onto_branch_url = (
212
        f"https://github.com/{org}/{project}/commit/{repo.rev_parse(onto_branch)}"
213
    )
214

215
    msg = f"@pytorchbot started a rebase job onto [{onto_branch}]({onto_branch_url})."
216
    msg += f" Check the current status [here]({os.getenv('GH_RUN_URL')})"
217
    gh_post_comment(org, project, args.pr_num, msg, dry_run=args.dry_run)
218

219
    if pr.is_closed():
220
        gh_post_comment(
221
            org,
222
            project,
223
            args.pr_num,
224
            f"PR #{args.pr_num} is closed, won't rebase",
225
            dry_run=args.dry_run,
226
        )
227
        return
228

229
    try:
230
        if pr.is_ghstack_pr():
231
            with git_config_guard(repo):
232
                rc = rebase_ghstack_onto(pr, repo, onto_branch, dry_run=args.dry_run)
233
        else:
234
            rc = rebase_onto(pr, repo, onto_branch, dry_run=args.dry_run)
235
        sys.exit(0 if rc else 1)
236

237
    except Exception as e:
238
        msg = f"Rebase failed due to {e}"
239
        msg += additional_rebase_failure_info(e)
240
        run_url = os.getenv("GH_RUN_URL")
241
        if run_url is not None:
242
            msg += f"\nRaised by {run_url}"
243
        gh_post_comment(org, project, args.pr_num, msg, dry_run=args.dry_run)
244

245

246
if __name__ == "__main__":
247
    main()
248