Add functionality for rebasebot to rebase onto viable/strict branch (pytorch#78276)

Fixes ROCm#325
**Summary**: Currently, the pytorchbot only allows for rebasing to the master branch. These modifications add functionality for rebasing to the 'viable/strict' branch of pytorch/pytorch by adding a flag to the comment.
**Test Plan:** tested manually on personal fork ([ROCm#1](swang392#1)), and included a test case in test_tryrebase.py that checks if rebasing to viable/strict branch was successful.
Pull Request resolved: pytorch#78276
Approved by: https://github.com/clee2000, https://github.com/janeyx99

Author: swang392

.github/scripts/test_tryrebase.py

@@ -22,6 +22,21 @@ def test_rebase(self, mocked_post_comment: Any, mocked_run_git: Any, mocked_gql:
         mocked_run_git.assert_has_calls(calls)
         self.assertTrue("Successfully rebased `master` onto `master`" in mocked_post_comment.call_args[0][3])
 
+    @mock.patch('trymerge.gh_graphql', side_effect=mocked_gh_graphql)
+    @mock.patch('gitutils.GitRepo._run_git')
+    @mock.patch('tryrebase.gh_post_comment')
+    def test_rebase_to_stable(self, mocked_post_comment: Any, mocked_run_git: Any, mocked_gql: Any) -> None:
+        "Tests rebase to viable/strict successfully"
+        pr = GitHubPR("pytorch", "pytorch", 31093)
+        repo = GitRepo(get_git_repo_dir(), get_git_remote_name())
+        rebase_onto(pr, repo, False, True)
+        calls = [mock.call('fetch', 'origin', 'pull/31093/head:pull/31093/head'),
+                 mock.call('rebase', 'viable/strict', 'pull/31093/head'),
+                 mock.call('push', '-f', 'https://github.com/mingxiaoh/pytorch.git', 'pull/31093/head:master')]
+        mocked_run_git.assert_has_calls(calls)
+        self.assertTrue(
+            "Successfully rebased `master` onto `viable/strict`" in mocked_post_comment.call_args[0][3])
+
     @mock.patch('trymerge.gh_graphql', side_effect=mocked_gh_graphql)
     @mock.patch('gitutils.GitRepo._run_git', return_value="Everything up-to-date")
     @mock.patch('tryrebase.gh_post_comment')

.github/scripts/tryrebase.py

@@ -13,13 +13,14 @@ def parse_args() -> Any:
     from argparse import ArgumentParser
     parser = ArgumentParser("Rebase PR into branch")
     parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--stable", action="store_true")
     parser.add_argument("pr_num", type=int)
     return parser.parse_args()
 
 
-def rebase_onto(pr: GitHubPR, repo: GitRepo, dry_run: bool = False) -> None:
+def rebase_onto(pr: GitHubPR, repo: GitRepo, dry_run: bool = False, stable: bool = False) -> None:
     branch = f"pull/{pr.pr_num}/head"
-    onto_branch = pr.default_branch()
+    onto_branch = "viable/strict" if stable else pr.default_branch()
     remote_url = f"https://github.com/{pr.info['headRepository']['nameWithOwner']}.git"
     refspec = f"{branch}:{pr.head_ref()}"
 
@@ -39,11 +40,11 @@ def rebase_onto(pr: GitHubPR, repo: GitRepo, dry_run: bool = False) -> None:
                         "git pull --rebase`)", dry_run=dry_run)
 
 
-def rebase_ghstack_onto(pr: GitHubPR, repo: GitRepo, dry_run: bool = False) -> None:
+def rebase_ghstack_onto(pr: GitHubPR, repo: GitRepo, dry_run: bool = False, stable: bool = False) -> None:
     if subprocess.run([sys.executable, "-m", "ghstack", "--help"], capture_output=True).returncode != 0:
         subprocess.run([sys.executable, "-m", "pip", "install", "ghstack"])
     orig_ref = f"{re.sub(r'/head$', '/orig', pr.head_ref())}"
-    onto_branch = pr.default_branch()
+    onto_branch = "viable/strict" if stable else pr.default_branch()
 
     repo.fetch(orig_ref, orig_ref)
     repo._run_git("rebase", onto_branch, orig_ref)
@@ -102,9 +103,9 @@ def main() -> None:
 
     try:
         if pr.is_ghstack_pr():
-            rebase_ghstack_onto(pr, repo, dry_run=args.dry_run)
+            rebase_ghstack_onto(pr, repo, dry_run=args.dry_run, stable=args.stable)
             return
-        rebase_onto(pr, repo, dry_run=args.dry_run)
+        rebase_onto(pr, repo, dry_run=args.dry_run, stable=args.stable)
     except Exception as e:
         msg = f"Rebase failed due to {e}"
         run_url = os.getenv("GH_RUN_URL")

.github/workflows/tryrebase.yml

@@ -30,5 +30,11 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.MERGEBOT_TOKEN }}
           PR_NUM: ${{ github.event.client_payload.pr_num }}
           GH_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+          STABLE: ${{ github.event.client_payload.stable}}
         run: |
-          python3 .github/scripts/tryrebase.py "${PR_NUM}"
+          set -ex
+          if [ -n "${STABLE}" ]; then
+            python3 .github/scripts/tryrebase.py --stable "${PR_NUM}"
+          else
+            python3 .github/scripts/tryrebase.py "${PR_NUM}"
+          fi