Fix update logic for Dr. CI (#3921)

Make available Dr. CI update feature also for domains.

.github/workflows/update-drci-comments.yml

@@ -14,8 +14,16 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v3
 
-      - name: Retrieve rockset query results and update Dr. CI comments
+      - name: Retrieve rockset query results and update Dr. CI comments for the PyTorch repo
         run: |
           curl --request POST \
           --url 'https://www.torch-ci.com/api/drci/drci' \
-          --header 'Authorization: ${{ secrets.DRCI_BOT_KEY }}'
+          --header 'Authorization: ${{ secrets.DRCI_BOT_KEY }}' \
+          --data 'repo=pytorch'
+
+      - name: Retrieve rockset query results and update Dr. CI comments for the Vision repo
+        run: |
+          curl --request POST \
+          --url 'https://www.torch-ci.com/api/drci/drci' \
+          --header 'Authorization: ${{ secrets.DRCI_BOT_KEY }}' \
+          --data 'repo=vision'

torchci/lib/bot/pytorchBotHandler.ts

@@ -371,10 +371,10 @@ The explanation needs to be clear on why this is needed. Here are some good exam
 
   async handleDrCI() {
     await this.logger.log("Dr. CI");
-    const { ctx, prNum } = this;
+    const { ctx, prNum, repo } = this;
 
     await this.ackComment();
-    await updateDrciComments(ctx.octokit, prNum.toString());
+    await updateDrciComments(ctx.octokit, repo, prNum.toString());
   }
 
   async handlePytorchCommands(inputArgs: string) {

torchci/lib/fetchRecentWorkflows.ts

@@ -4,6 +4,7 @@ import rocksetVersions from "rockset/prodVersions.json";
 import { RecentWorkflowsData } from "./types";
 
 export async function fetchRecentWorkflows(
+  repo: string = "pytorch/pytorch",
   prNumber: string = "0",
   numMinutes: string = "30"
 ): Promise<RecentWorkflowsData[]> {
@@ -24,6 +25,11 @@ export async function fetchRecentWorkflows(
           type: "int",
           value: prNumber,
         },
+        {
+          name: "repo",
+          type: "string",
+          value: repo,
+        },
       ],
     }
   );

torchci/pages/api/drci/drci.ts

@@ -33,22 +33,30 @@ export interface FlakyRule {
   captures: string[];
 }
 
+export interface UpdateCommentBody {
+    repo: string;
+}
+
 export default async function handler(
     req: NextApiRequest,
     res: NextApiResponse<void>
 ) {
     const authorization = req.headers.authorization;
+
     if (authorization === process.env.DRCI_BOT_KEY) {
         const { prNumber } = req.query;
-        const octokit = await getOctokit(OWNER, REPO);
-        updateDrciComments(octokit, prNumber as string);
+        const { repo }: UpdateCommentBody = req.body;
+        const octokit = await getOctokit(OWNER, repo);
+        updateDrciComments(octokit, repo, prNumber as string);
+
         res.status(200).end();
     }
     res.status(403).end();
 }
 
-export async function updateDrciComments(octokit: Octokit, prNumber?: string) {
+export async function updateDrciComments(octokit: Octokit, repo: string = "pytorch", prNumber?: string) {
     const recentWorkflows: RecentWorkflowsData[] = await fetchRecentWorkflows(
+        `${OWNER}/${repo}`,
         prNumber,
         NUM_MINUTES + ""
     );
@@ -80,12 +88,12 @@ export async function updateDrciComments(octokit: Octokit, prNumber?: string) {
       const comment = formDrciComment(
         pr_info.pr_number,
         OWNER,
-        REPO,
+        repo,
         failureInfo,
         formDrciSevBody(sevs)
       );
 
-      await updateCommentWithWorkflow(octokit, pr_info, comment);
+      await updateCommentWithWorkflow(octokit, pr_info, comment, repo);
     });
 }
 
@@ -200,8 +208,7 @@ export function constructResultsComment(
 
         output += `\nAs of commit ${sha}:`;
         output += `\n:green_heart: Looks good so far! There are no failures yet. :green_heart:`;
-    }
-    else {
+    } else {
         output += someFailing;
         if (hasPending) {
             output += somePending;
@@ -319,9 +326,10 @@ export async function updateCommentWithWorkflow(
     octokit: Octokit,
     pr_info: PRandJobs,
     comment: string,
+    repo: string,
 ): Promise<void> {
     const { pr_number } = pr_info;
-    const { id, body } = await getDrciComment(octokit, OWNER, REPO, pr_number!);
+    const { id, body } = await getDrciComment(octokit, OWNER, repo, pr_number!);
 
     if (id === 0 || body === comment) {
         return;
@@ -330,7 +338,7 @@ export async function updateCommentWithWorkflow(
     await octokit.rest.issues.updateComment({
         body: comment,
         owner: OWNER,
-        repo: REPO,
+        repo: repo,
         comment_id: id,
     });
 }

torchci/rockset/commons/__sql/recent_pr_workflows_query.sql

@@ -13,7 +13,7 @@ WITH recent_shas as (
             )
             or :prNumber = p.number
         )
-        and p.base.repo.full_name = 'pytorch/pytorch'
+        and p.base.repo.full_name = :repo
 )
 SELECT
     w.id as workflow_id,

torchci/rockset/commons/recent_pr_workflows_query.lambda.json

@@ -10,6 +10,11 @@
       "name": "prNumber",
       "type": "int",
       "value": "0"
+    },
+    {
+      "name": "repo",
+      "type": "string",
+      "value": "pytorch/pytorch"
     }
   ],
   "description": "retrieve the CI check results in the past N minutes, as well as the other workflows from the same PRs that were previously run/pending"

torchci/rockset/prodVersions.json

@@ -15,7 +15,7 @@
     "issue_query": "276aea6b8186d808",
     "failure_samples_query": "18e7e696b4949f05",
     "num_commits_master": "08d299a1b7386dbb",
-    "recent_pr_workflows_query": "0deed4379aec49ce",
+    "recent_pr_workflows_query": "dc6152f7ab81a336",
     "reverted_prs_with_reason": "751f01cba16364f0",
     "unclassified": "1b31a2d8f4ab7230",
     "test_insights_overview": "c9be5cbdda1030af",