From 13095a53bb60866cf032152bf2f9204f1d2f4dcd Mon Sep 17 00:00:00 2001
From: driazati <driazati@users.noreply.github.com>
Date: Wed, 3 Aug 2022 11:23:58 -0700
Subject: [PATCH] [ci][tvmbot] Enable re-run for GitHub Actions

---
 .github/workflows/main.yml     |  7 +++-
 .github/workflows/tvmbot.yml   |  3 +-
 tests/scripts/github_tvmbot.py | 76 ++++++++++++++++++++++++++--------
 3 files changed, 64 insertions(+), 22 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 55fe5f1441cb3..077fbb7ce68a1 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -36,12 +36,15 @@ concurrency:
 
 jobs:
   MacOS:
-    if: ${{ github.repository == 'apache/tvm' }}
-    runs-on: macOS-latest
+    if: ${{ github.repository == 'driazati/tvm' }}
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
         with:
           submodules: 'recursive'
+      - name: fail
+        run: |
+          exit 1
       - name: Set up environment
         uses: ./.github/actions/setup
       - name: Conda Build
diff --git a/.github/workflows/tvmbot.yml b/.github/workflows/tvmbot.yml
index 792977f92ee5d..349cb4d2cca95 100644
--- a/.github/workflows/tvmbot.yml
+++ b/.github/workflows/tvmbot.yml
@@ -1,7 +1,6 @@
 
 name: tvm-bot
 on:
-  status:
   pull_request_review:
     types:
       - submitted
@@ -21,7 +20,7 @@ jobs:
       issues: write
       pull-requests: write
       statuses: write
-    if: ${{ github.event.issue.pull_request && github.repository == 'apache/tvm' }}
+    if: ${{ github.event.issue.pull_request && github.repository == 'driazati/tvm' }}
     runs-on: ubuntu-20.04
     steps:
       - uses: actions/checkout@v2
diff --git a/tests/scripts/github_tvmbot.py b/tests/scripts/github_tvmbot.py
index e83318e18e513..7ef7119035260 100755
--- a/tests/scripts/github_tvmbot.py
+++ b/tests/scripts/github_tvmbot.py
@@ -106,6 +106,7 @@ def to_json_str(obj: Any) -> str:
                     nodes {
                       ... on CheckRun {
                         name
+                        databaseId
                         checkSuite {
                           workflowRun {
                             workflow {
@@ -503,6 +504,16 @@ def rerun_jenkins_ci(self) -> None:
         else:
             post(url, auth=("tvm-bot", TVM_BOT_JENKINS_TOKEN))
 
+    def rerun_github_actions(self) -> None:
+        job_ids = []
+        for item in self.head_commit()["statusCheckRollup"]["contexts"]["nodes"]:
+            if "checkSuite" in item:
+                job_ids.append(item["databaseId"])
+
+        logging.info(f"Rerunning GitHub Actions jobs with IDs: {job_ids}")
+        for job_id in job_ids:
+            self.github.post(f"actions/jobs/{job_id}/rerun", data={})
+
     def comment_failure(self, msg: str, exception: Exception):
         if not self.dry_run:
             exception_msg = traceback.format_exc()
@@ -514,6 +525,40 @@ def comment_failure(self, msg: str, exception: Exception):
         return exception
 
 
+def check_author(pr, triggering_comment, args):
+    comment_author = triggering_comment["user"]["login"]
+    if pr.author() == comment_author:
+        logging.info("Comment user is PR author, continuing")
+        return True
+    return False
+
+
+def check_collaborator(pr, triggering_comment, args):
+    logging.info("Checking collaborators")
+    # Get the list of collaborators for the repo filtered by the comment
+    # author
+    if args.testing_collaborators_json:
+        collaborators = json.loads(args.testing_collaborators_json)
+    else:
+        collaborators = pr.search_collaborator(triggering_comment["user"]["login"])
+    logging.info(f"Found collaborators: {collaborators}")
+
+    return len(collaborators) > 0
+
+
+def check_anyone(pr, triggering_comment, args):
+    return True
+
+
+AUTH_CHECKS = {
+    "anyone": check_anyone,
+    "collaborators": check_collaborator,
+    "author": check_author,
+}
+# Stash the keys so they're accessible from the values
+AUTH_CHECKS = {k: (k, v) for k, v in AUTH_CHECKS.items()}
+
+
 class Merge:
     triggers = [
         "merge",
@@ -521,6 +566,8 @@ class Merge:
         "merge this pr",
     ]
 
+    auth = [AUTH_CHECKS["collaborators"], AUTH_CHECKS["author"]]
+
     @staticmethod
     def run(pr: PR):
         info = None
@@ -548,9 +595,15 @@ class Rerun:
         "run ci",
     ]
 
+    auth = [AUTH_CHECKS["anyone"]]
+
     @staticmethod
     def run(pr: PR):
-        pr.rerun_jenkins_ci()
+        try:
+            # pr.rerun_jenkins_ci()
+            pr.rerun_github_actions()
+        except Exception as e:
+            pr.comment_failure("Failed to re-run CI", e)
 
 
 if __name__ == "__main__":
@@ -618,24 +671,11 @@ def run(pr: PR):
     # Acknowledge the comment with a react
     pr.plus_one(comment)
 
-    # Check the comment author
-    comment_author = comment["user"]["login"]
-    if pr.author() == comment_author:
-        logging.info("Comment user is PR author, continuing")
-    else:
-        logging.info("Comment is not from PR author, checking collaborators")
-        # Get the list of collaborators for the repo filtered by the comment
-        # author
-        if args.testing_collaborators_json:
-            collaborators = json.loads(args.testing_collaborators_json)
-        else:
-            collaborators = pr.search_collaborator(comment_author)
-        logging.info(f"Found collaborators: {collaborators}")
-
-        if len(collaborators) > 0:
-            logging.info("Comment is from collaborator")
+    for name, check in command_to_run.auth:
+        if check(pr, comment, args):
+            logging.info(f"Passed auth check '{name}', continuing")
         else:
-            logging.info("Comment is not from from PR author or collaborator, quitting")
+            logging.info(f"Failed auth check '{name}', quitting")
             exit(0)
 
     state = pr.state()