Support retryable_step_names in retryBot for per-repo retryable step config (#7876)

Allow repos to declare step names (e.g. "CUDA Compute Check") that
should always trigger a retry when they fail, regardless of whether test
steps also failed. This enables retrying Helion jobs on GPU health check
failures.

PR on helion side pytorch/helion#1808

Signed-off-by: Huy Do <huydhn@gmail.com>

Author: huydhn

torchci/lib/bot/retryBot.ts

@@ -72,7 +72,8 @@ async function retryCurrentWorkflow(
   defaultBranch: string,
   workflowName: string,
   workflowJobs: any[],
-  runId: number
+  runId: number,
+  retryableStepNames: string[] = []
 ) {
   const failedJobs = workflowJobs.filter((job) =>
     FAILURE_CONCLUSIONS.includes(job.conclusion!)
@@ -135,6 +136,22 @@ async function retryCurrentWorkflow(
       }
     }
 
+    // If a retryable step name failed, always retry (e.g. CUDA Compute Check
+    // indicates a bad runner, not a code problem)
+    if (retryableStepNames.length > 0) {
+      const hasRetryableStepFailure = job.steps?.some(
+        (step: any) =>
+          step.conclusion !== null &&
+          FAILURE_CONCLUSIONS.includes(step.conclusion) &&
+          retryableStepNames.some(
+            (name) => step.name.toLowerCase() === name.toLowerCase()
+          )
+      );
+      if (hasRetryableStepFailure) {
+        return true;
+      }
+    }
+
     // if no test steps failed, can rerun
     return !doesLookLikeUserFailure(job, (step) =>
       step.name.toLowerCase().includes("test")
@@ -181,6 +198,8 @@ function retryBot(app: Probot): void {
     const config: any = await tracker.loadConfig(ctx);
     const allowedWorkflowPrefixes: string[] | undefined =
       config != null ? config["retryable_workflows"] : undefined;
+    const retryableStepNames: string[] =
+      config != null ? config["retryable_step_names"] ?? [] : [];
 
     if (allowedWorkflowPrefixes === undefined) {
       return;
@@ -225,7 +244,8 @@ function retryBot(app: Probot): void {
         defaultBranch,
         workflowName,
         workflowJobs,
-        runId
+        runId,
+        retryableStepNames
       );
     }
 

torchci/test/retryBot.test.ts

@@ -440,6 +440,114 @@ describe("retry-bot", () => {
     handleScope(scope);
   });
 
+  test("rerun when a retryable step name fails", async () => {
+    mockIsPytorchbotSupportedOrg(true);
+    const event = requireDeepCopy("./fixtures/workflow_run.completed.json");
+    event.payload.workflow_run.name = "test";
+    const workflow_jobs = requireDeepCopy("./fixtures/workflow_jobs.json");
+    // Job fails on a retryable step (CUDA Compute Check)
+    workflow_jobs.jobs[4].conclusion = "failure";
+    workflow_jobs.jobs[4].steps = [
+      {
+        name: "CUDA Compute Check",
+        status: "completed",
+        conclusion: "failure",
+        number: 1,
+        started_at: "2022-10-06T18:09:54.000-07:00",
+        completed_at: "2022-10-06T18:09:54.000-07:00",
+      },
+    ];
+
+    const owner = event.payload.repository.owner.login;
+    const repo = event.payload.repository.name;
+    const attempt_number = event.payload.workflow_run.run_attempt;
+    const run_id = event.payload.workflow_run.id;
+
+    const scope = nock("https://api.github.com")
+      .get(
+        `/repos/${owner}/${repo}/actions/runs/${run_id}/attempts/${attempt_number}/jobs?page=1&per_page=100`
+      )
+      .reply(200, workflow_jobs)
+      .get(
+        `/repos/${owner}/${repo}/contents/${encodeURIComponent(
+          ".github/pytorch-probot.yml"
+        )}`
+      )
+      .reply(
+        200,
+        '{retryable_workflows: ["test", "benchmark"], retryable_step_names: ["CUDA Compute Check"]}'
+      )
+      .post(
+        `/repos/${owner}/${repo}/actions/jobs/${workflow_jobs.jobs[4].id}/rerun`
+      )
+      .reply(200);
+
+    const mock = jest.spyOn(clickhouse, "queryClickhouseSaved");
+    mock.mockImplementation(() => Promise.resolve([]));
+
+    await probot.receive(event);
+
+    handleScope(scope);
+  });
+
+  test("rerun when retryable step fails even if test step also failed", async () => {
+    mockIsPytorchbotSupportedOrg(true);
+    const event = requireDeepCopy("./fixtures/workflow_run.completed.json");
+    event.payload.workflow_run.name = "test";
+    const workflow_jobs = requireDeepCopy("./fixtures/workflow_jobs.json");
+    // Job fails on both a retryable step and a test step
+    workflow_jobs.jobs[4].conclusion = "failure";
+    workflow_jobs.jobs[4].steps = [
+      {
+        name: "CUDA Compute Check",
+        status: "completed",
+        conclusion: "failure",
+        number: 1,
+        started_at: "2022-10-06T18:09:54.000-07:00",
+        completed_at: "2022-10-06T18:09:54.000-07:00",
+      },
+      {
+        name: "test",
+        status: "completed",
+        conclusion: "failure",
+        number: 2,
+        started_at: "2022-10-06T18:10:54.000-07:00",
+        completed_at: "2022-10-06T18:10:54.000-07:00",
+      },
+    ];
+
+    const owner = event.payload.repository.owner.login;
+    const repo = event.payload.repository.name;
+    const attempt_number = event.payload.workflow_run.run_attempt;
+    const run_id = event.payload.workflow_run.id;
+
+    const scope = nock("https://api.github.com")
+      .get(
+        `/repos/${owner}/${repo}/actions/runs/${run_id}/attempts/${attempt_number}/jobs?page=1&per_page=100`
+      )
+      .reply(200, workflow_jobs)
+      .get(
+        `/repos/${owner}/${repo}/contents/${encodeURIComponent(
+          ".github/pytorch-probot.yml"
+        )}`
+      )
+      .reply(
+        200,
+        '{retryable_workflows: ["test", "benchmark"], retryable_step_names: ["CUDA Compute Check"]}'
+      )
+      .post(
+        `/repos/${owner}/${repo}/actions/jobs/${workflow_jobs.jobs[4].id}/rerun`
+      )
+      .reply(200);
+
+    const mock = jest.spyOn(clickhouse, "queryClickhouseSaved");
+    mock.mockImplementation(() => Promise.resolve([]));
+
+    await probot.receive(event);
+
+    handleScope(scope);
+  });
+
   test("dont re-run unless retryable_workflows is specified in .github/pytorch-probot.yml", async () => {
     mockIsPytorchbotSupportedOrg(true);
     const event = requireDeepCopy("./fixtures/workflow_run.completed.json");