Fix for OOM retrying

Author: matt-aitken

internal-packages/run-engine/src/engine/retrying.ts

@@ -6,6 +6,7 @@ import {
   TaskRunExecutionRetry,
   taskRunErrorEnhancer,
   sanitizeError,
+  calculateNextRetryDelay,
 } from "@trigger.dev/core/v3";
 import { PrismaClientOrTransaction, TaskRunStatus } from "@trigger.dev/database";
 import { MAX_TASK_RUN_ATTEMPTS } from "./consts";
@@ -47,26 +48,33 @@ export async function retryOutcomeFromCompletion(
 
   const sanitizedError = sanitizeError(error);
 
-  // No retry settings
-  if (!retrySettings) {
-    return { outcome: "fail_run", sanitizedError };
-  }
-
   // OOM error (retry on a larger machine or fail)
   if (isOOMRunError(error)) {
-    const newMachine = await retryOOMOnMachine(prisma, runId);
-    if (!newMachine) {
+    const oomResult = await retryOOMOnMachine(prisma, runId);
+    if (!oomResult) {
+      return { outcome: "fail_run", sanitizedError, wasOOMError: true };
+    }
+
+    const delay = calculateNextRetryDelay(oomResult.retrySettings, attemptNumber ?? 1);
+
+    if (!delay) {
+      //no more retries left
       return { outcome: "fail_run", sanitizedError, wasOOMError: true };
     }
 
     return {
       outcome: "retry",
       method: "queue",
-      settings: retrySettings,
-      machine: newMachine,
+      machine: oomResult.machine,
+      settings: { timestamp: Date.now() + delay, delay },
     };
   }
 
+  // No retry settings
+  if (!retrySettings) {
+    return { outcome: "fail_run", sanitizedError };
+  }
+
   // Not a retriable error: fail
   const retriableError = shouldRetryError(taskRunErrorEnhancer(error));
   if (!retriableError) {
@@ -112,7 +120,7 @@ export async function retryOutcomeFromCompletion(
 async function retryOOMOnMachine(
   prisma: PrismaClientOrTransaction,
   runId: string
-): Promise<string | undefined> {
+): Promise<{ machine: string; retrySettings: RetryOptions } | undefined> {
   try {
     const run = await prisma.taskRun.findFirst({
       where: {
@@ -153,7 +161,7 @@ async function retryOOMOnMachine(
       return;
     }
 
-    return retryMachine;
+    return { machine: retryMachine, retrySettings: parsedRetryConfig.data };
   } catch (error) {
     console.error("[FailedTaskRunRetryHelper] Failed to get execution retry", {
       runId,

internal-packages/run-engine/src/engine/tests/attemptFailures.test.ts

@@ -8,6 +8,7 @@ import { trace } from "@opentelemetry/api";
 import { expect } from "vitest";
 import { EventBusEventArgs } from "../eventBus.js";
 import { RunEngine } from "../index.js";
+import { setTimeout } from "node:timers/promises";
 
 describe("RunEngine attempt failures", () => {
   containerTest(
@@ -479,10 +480,6 @@ describe("RunEngine attempt failures", () => {
           ok: false,
           id: dequeued[0].run.id,
           error,
-          retry: {
-            timestamp: Date.now(),
-            delay: 0,
-          },
         },
       });
 
@@ -603,10 +600,6 @@ describe("RunEngine attempt failures", () => {
             ok: false,
             id: dequeued[0].run.id,
             error,
-            retry: {
-              timestamp: Date.now(),
-              delay: 0,
-            },
           },
         });
 
@@ -714,6 +707,8 @@ describe("RunEngine attempt failures", () => {
 
         //create background worker
         await setupBackgroundWorker(prisma, authenticatedEnvironment, taskIdentifier, undefined, {
+          maxTimeoutInMs: 10,
+          maxAttempts: 10,
           outOfMemory: {
             machine: "small-2x",
           },
@@ -768,10 +763,6 @@ describe("RunEngine attempt failures", () => {
             ok: false,
             id: dequeued[0].run.id,
             error,
-            retry: {
-              timestamp: Date.now(),
-              delay: 0,
-            },
           },
         });
 
@@ -787,12 +778,16 @@ describe("RunEngine attempt failures", () => {
         expect(executionData.run.attemptNumber).toBe(1);
         expect(executionData.run.status).toBe("RETRYING_AFTER_FAILURE");
 
+        //wait for 1s
+        await setTimeout(1_000);
+
         //dequeue again
         const dequeued2 = await engine.dequeueFromMasterQueue({
           consumerId: "test_12345",
           masterQueue: run.masterQueue,
           maxRunCount: 10,
         });
+        expect(dequeued2.length).toBe(1);
 
         //create second attempt
         const attemptResult2 = await engine.startRunAttempt({