Merge branch 'main' into main

Author: zvictor

.changeset/forty-windows-shop.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Added the ability to retry runs that fail with an Out Of Memory (OOM) error on a larger machine.

apps/webapp/app/v3/failedTaskRun.server.ts

@@ -5,14 +5,14 @@ import {
   TaskRunExecutionRetry,
   TaskRunFailedExecutionResult,
 } from "@trigger.dev/core/v3";
+import type { Prisma, TaskRun } from "@trigger.dev/database";
+import * as semver from "semver";
 import { logger } from "~/services/logger.server";
+import { sharedQueueTasks } from "./marqs/sharedQueueConsumer.server";
 import { BaseService } from "./services/baseService.server";
-import { isFailableRunStatus, isFinalAttemptStatus } from "./taskStatus";
-import type { Prisma, TaskRun } from "@trigger.dev/database";
 import { CompleteAttemptService } from "./services/completeAttempt.server";
 import { CreateTaskRunAttemptService } from "./services/createTaskRunAttempt.server";
-import { sharedQueueTasks } from "./marqs/sharedQueueConsumer.server";
-import * as semver from "semver";
+import { isFailableRunStatus, isFinalAttemptStatus } from "./taskStatus";
 
 const FailedTaskRunRetryGetPayload = {
   select: {
@@ -180,13 +180,52 @@ export class FailedTaskRunRetryHelper extends BaseService {
     }
   }
 
-  static async getExecutionRetry({
+  static getExecutionRetry({
     run,
     execution,
   }: {
     run: TaskRunWithWorker;
     execution: TaskRunExecution;
-  }): Promise<TaskRunExecutionRetry | undefined> {
+  }): TaskRunExecutionRetry | undefined {
+    try {
+      const retryConfig = FailedTaskRunRetryHelper.getRetryConfig({ run, execution });
+      if (!retryConfig) {
+        return;
+      }
+
+      const delay = calculateNextRetryDelay(retryConfig, execution.attempt.number);
+
+      if (!delay) {
+        logger.debug("[FailedTaskRunRetryHelper] No more retries", {
+          run,
+          execution,
+        });
+
+        return;
+      }
+
+      return {
+        timestamp: Date.now() + delay,
+        delay,
+      };
+    } catch (error) {
+      logger.error("[FailedTaskRunRetryHelper] Failed to get execution retry", {
+        run,
+        execution,
+        error,
+      });
+
+      return;
+    }
+  }
+
+  static getRetryConfig({
+    run,
+    execution,
+  }: {
+    run: TaskRunWithWorker;
+    execution: TaskRunExecution;
+  }): RetryOptions | undefined {
     try {
       const retryConfig = run.lockedBy?.retryConfig;
 
@@ -247,21 +286,7 @@ export class FailedTaskRunRetryHelper extends BaseService {
         return;
       }
 
-      const delay = calculateNextRetryDelay(parsedRetryConfig.data, execution.attempt.number);
-
-      if (!delay) {
-        logger.debug("[FailedTaskRunRetryHelper] No more retries", {
-          run,
-          execution,
-        });
-
-        return;
-      }
-
-      return {
-        timestamp: Date.now() + delay,
-        delay,
-      };
+      return parsedRetryConfig.data;
     } catch (error) {
       logger.error("[FailedTaskRunRetryHelper] Failed to get execution retry", {
         run,

apps/webapp/app/v3/marqs/index.server.ts

@@ -638,7 +638,8 @@ export class MarQS {
   }
 
   /**
-   * Negative acknowledge a message, which will requeue the message
+   * Negative acknowledge a message, which will requeue the message.
+   * Returns whether it went back into the queue or not.
    */
   public async nackMessage(
     messageId: string,
@@ -657,7 +658,7 @@ export class MarQS {
             updates,
             service: this.name,
           });
-          return;
+          return false;
         }
 
         const nackCount = await this.#getNackCount(messageId);
@@ -676,7 +677,7 @@ export class MarQS {
 
           // If we have reached the maximum nack count, we will ack the message
           await this.acknowledgeMessage(messageId, "maximum nack count reached");
-          return;
+          return false;
         }
 
         span.setAttributes({
@@ -705,6 +706,8 @@ export class MarQS {
         });
 
         await this.options.subscriber?.messageNacked(message);
+
+        return true;
       },
       {
         kind: SpanKind.CONSUMER,

apps/webapp/app/v3/services/completeAttempt.server.ts

@@ -1,13 +1,16 @@
 import { Attributes } from "@opentelemetry/api";
 import {
   TaskRunContext,
+  TaskRunError,
   TaskRunErrorCodes,
   TaskRunExecution,
   TaskRunExecutionResult,
   TaskRunExecutionRetry,
   TaskRunFailedExecutionResult,
   TaskRunSuccessfulExecutionResult,
+  exceptionEventEnhancer,
   flattenAttributes,
+  internalErrorFromUnexpectedExit,
   sanitizeError,
   shouldRetryError,
   taskRunErrorEnhancer,
@@ -233,7 +236,7 @@ export class CompleteAttemptService extends BaseService {
 
     if (!executionRetry && shouldInfer) {
       executionRetryInferred = true;
-      executionRetry = await FailedTaskRunRetryHelper.getExecutionRetry({
+      executionRetry = FailedTaskRunRetryHelper.getExecutionRetry({
         run: {
           ...taskRunAttempt.taskRun,
           lockedBy: taskRunAttempt.backgroundWorkerTask,
@@ -243,7 +246,47 @@ export class CompleteAttemptService extends BaseService {
       });
     }
 
-    const retriableError = shouldRetryError(taskRunErrorEnhancer(completion.error));
+    let retriableError = shouldRetryError(taskRunErrorEnhancer(completion.error));
+    let isOOMRetry = false;
+
+    //OOM errors should retry (if an OOM machine is specified)
+    if (isOOMError(completion.error)) {
+      const retryConfig = FailedTaskRunRetryHelper.getRetryConfig({
+        run: {
+          ...taskRunAttempt.taskRun,
+          lockedBy: taskRunAttempt.backgroundWorkerTask,
+          lockedToVersion: taskRunAttempt.backgroundWorker,
+        },
+        execution,
+      });
+
+      if (
+        retryConfig?.outOfMemory?.machine &&
+        retryConfig.outOfMemory.machine !== taskRunAttempt.taskRun.machinePreset
+      ) {
+        //we will retry
+        isOOMRetry = true;
+        retriableError = true;
+        executionRetry = FailedTaskRunRetryHelper.getExecutionRetry({
+          run: {
+            ...taskRunAttempt.taskRun,
+            lockedBy: taskRunAttempt.backgroundWorkerTask,
+            lockedToVersion: taskRunAttempt.backgroundWorker,
+          },
+          execution,
+        });
+
+        //update the machine on the run
+        await this._prisma.taskRun.update({
+          where: {
+            id: taskRunAttempt.taskRunId,
+          },
+          data: {
+            machinePreset: retryConfig.outOfMemory.machine,
+          },
+        });
+      }
+    }
 
     if (
       retriableError &&
@@ -257,6 +300,7 @@ export class CompleteAttemptService extends BaseService {
         taskRunAttempt,
         environment,
         checkpoint,
+        forceRequeue: isOOMRetry,
       });
     }
 
@@ -378,12 +422,14 @@ export class CompleteAttemptService extends BaseService {
     executionRetryInferred,
     checkpointEventId,
     supportsLazyAttempts,
+    forceRequeue = false,
   }: {
     run: TaskRun;
     executionRetry: TaskRunExecutionRetry;
     executionRetryInferred: boolean;
     checkpointEventId?: string;
     supportsLazyAttempts: boolean;
+    forceRequeue?: boolean;
   }) {
     const retryViaQueue = () => {
       logger.debug("[CompleteAttemptService] Enqueuing retry attempt", { runId: run.id });
@@ -434,6 +480,12 @@ export class CompleteAttemptService extends BaseService {
       return;
     }
 
+    if (forceRequeue) {
+      logger.debug("[CompleteAttemptService] Forcing retry via queue", { runId: run.id });
+      await retryViaQueue();
+      return;
+    }
+
     // Workers that never checkpoint between attempts will exit after completing their current attempt if the retry delay exceeds the threshold
     if (
       !this.opts.supportsRetryCheckpoints &&
@@ -466,13 +518,15 @@ export class CompleteAttemptService extends BaseService {
     taskRunAttempt,
     environment,
     checkpoint,
+    forceRequeue = false,
   }: {
     execution: TaskRunExecution;
     executionRetry: TaskRunExecutionRetry;
     executionRetryInferred: boolean;
     taskRunAttempt: NonNullable<FoundAttempt>;
     environment: AuthenticatedEnvironment;
     checkpoint?: CheckpointData;
+    forceRequeue?: boolean;
   }) {
     const retryAt = new Date(executionRetry.timestamp);
 
@@ -533,6 +587,7 @@ export class CompleteAttemptService extends BaseService {
       executionRetry,
       supportsLazyAttempts: taskRunAttempt.backgroundWorker.supportsLazyAttempts,
       executionRetryInferred,
+      forceRequeue,
     });
 
     return "RETRIED";
@@ -634,3 +689,12 @@ async function findAttempt(prismaClient: PrismaClientOrTransaction, friendlyId:
     },
   });
 }
+
+function isOOMError(error: TaskRunError) {
+  if (error.type !== "INTERNAL_ERROR") return false;
+  if (error.code === "TASK_PROCESS_OOM_KILLED" || error.code === "TASK_PROCESS_MAYBE_OOM_KILLED") {
+    return true;
+  }
+
+  return false;
+}