remove checkpoints between attempts

Author: nicktrn

apps/coordinator/src/checkpointer.ts

@@ -17,6 +17,7 @@ type CheckpointAndPushOptions = {
   projectRef: string;
   deploymentVersion: string;
   shouldHeartbeat?: boolean;
+  attemptNumber?: number;
 };
 
 type CheckpointAndPushResult =
@@ -258,6 +259,7 @@ export class Checkpointer {
     leaveRunning = true, // This mirrors kubernetes behaviour more accurately
     projectRef,
     deploymentVersion,
+    attemptNumber,
   }: CheckpointAndPushOptions): Promise<CheckpointAndPushResult> {
     this.#logger.log("Checkpointing with backoff", {
       runId,
@@ -297,6 +299,7 @@ export class Checkpointer {
           leaveRunning,
           projectRef,
           deploymentVersion,
+          attemptNumber,
         });
 
         if (result.success) {
@@ -359,6 +362,7 @@ export class Checkpointer {
     leaveRunning = true, // This mirrors kubernetes behaviour more accurately
     projectRef,
     deploymentVersion,
+    attemptNumber,
   }: CheckpointAndPushOptions): Promise<CheckpointAndPushResult> {
     await this.init();
 
@@ -367,6 +371,7 @@ export class Checkpointer {
       leaveRunning,
       projectRef,
       deploymentVersion,
+      attemptNumber,
     };
 
     if (!this.#dockerMode && !this.#canCheckpoint) {
@@ -417,7 +422,7 @@ export class Checkpointer {
 
       this.#logger.log("Checkpointing:", { options });
 
-      const containterName = this.#getRunContainerName(runId);
+      const containterName = this.#getRunContainerName(runId, attemptNumber);
 
       // Create checkpoint (docker)
       if (this.#dockerMode) {
@@ -581,7 +586,7 @@ export class Checkpointer {
     return this.#failedCheckpoints.has(runId);
   }
 
-  #getRunContainerName(suffix: string) {
-    return `task-run-${suffix}`;
+  #getRunContainerName(suffix: string, attemptNumber?: number) {
+    return `task-run-${suffix}${attemptNumber && attemptNumber > 1 ? `-att${attemptNumber}` : ""}`;
   }
 }

apps/coordinator/src/index.ts

@@ -294,6 +294,7 @@ class TaskCoordinator {
           setSocketDataFromHeader("projectRef", "x-trigger-project-ref");
           setSocketDataFromHeader("runId", "x-trigger-run-id");
           setSocketDataFromHeader("attemptFriendlyId", "x-trigger-attempt-friendly-id", false);
+          setSocketDataFromHeader("attemptNumber", "x-trigger-attempt-number", false);
           setSocketDataFromHeader("envId", "x-trigger-env-id");
           setSocketDataFromHeader("deploymentId", "x-trigger-deployment-id");
           setSocketDataFromHeader("deploymentVersion", "x-trigger-deployment-version");
@@ -310,6 +311,10 @@ class TaskCoordinator {
       onConnection: async (socket, handler, sender) => {
         const logger = new SimpleLogger(`[prod-worker][${socket.id}]`);
 
+        const getAttemptNumber = () => {
+          return socket.data.attemptNumber ? parseInt(socket.data.attemptNumber) : undefined;
+        };
+
         const crashRun = async (error: { name: string; message: string; stack?: string }) => {
           try {
             this.#platformSocket?.send("RUN_CRASHED", {
@@ -385,6 +390,10 @@ class TaskCoordinator {
           socket.data.attemptFriendlyId = attemptFriendlyId;
         };
 
+        const updateAttemptNumber = (attemptNumber: string | number) => {
+          socket.data.attemptNumber = String(attemptNumber);
+        };
+
         this.#platformSocket?.send("LOG", {
           metadata: socket.data,
           text: "connected",
@@ -434,6 +443,7 @@ class TaskCoordinator {
             });
 
             updateAttemptFriendlyId(executionAck.payload.execution.attempt.id);
+            updateAttemptNumber(executionAck.payload.execution.attempt.number);
           } catch (error) {
             logger.error("Error", { error });
 
@@ -509,11 +519,17 @@ class TaskCoordinator {
 
           updateAttemptFriendlyId(message.attemptFriendlyId);
 
-          this.#platformSocket?.send("READY_FOR_RESUME", message);
+          if (message.version === "v2") {
+            updateAttemptNumber(message.attemptNumber);
+          }
+
+          this.#platformSocket?.send("READY_FOR_RESUME", { ...message, version: "v1" });
         });
 
         // MARK: RUN COMPLETED
-        socket.on("TASK_RUN_COMPLETED", async ({ completion, execution }, callback) => {
+        socket.on("TASK_RUN_COMPLETED", async (message, callback) => {
+          const { completion, execution } = message;
+
           logger.log("completed task", { completionId: completion.id });
 
           // Cancel all in-progress checkpoints (if any)
@@ -522,8 +538,10 @@ class TaskCoordinator {
           await chaosMonkey.call({ throwErrors: false });
 
           const completeWithoutCheckpoint = (shouldExit: boolean) => {
+            const supportsRetryCheckpoints = message.version === "v1";
+
             this.#platformSocket?.send("TASK_RUN_COMPLETED", {
-              version: "v1",
+              version: supportsRetryCheckpoints ? "v1" : "v2",
               execution,
               completion,
             });
@@ -553,6 +571,11 @@ class TaskCoordinator {
             return;
           }
 
+          if (message.version === "v2") {
+            completeWithoutCheckpoint(true);
+            return;
+          }
+
           const { canCheckpoint, willSimulate } = await this.#checkpointer.init();
 
           const willCheckpointAndRestore = canCheckpoint || willSimulate;
@@ -685,6 +708,7 @@ class TaskCoordinator {
             runId: socket.data.runId,
             projectRef: socket.data.projectRef,
             deploymentVersion: socket.data.deploymentVersion,
+            attemptNumber: getAttemptNumber(),
           });
 
           if (!checkpoint) {
@@ -756,6 +780,7 @@ class TaskCoordinator {
             runId: socket.data.runId,
             projectRef: socket.data.projectRef,
             deploymentVersion: socket.data.deploymentVersion,
+            attemptNumber: getAttemptNumber(),
           });
 
           if (!checkpoint) {
@@ -825,6 +850,7 @@ class TaskCoordinator {
             runId: socket.data.runId,
             projectRef: socket.data.projectRef,
             deploymentVersion: socket.data.deploymentVersion,
+            attemptNumber: getAttemptNumber(),
           });
 
           if (!checkpoint) {
@@ -909,6 +935,7 @@ class TaskCoordinator {
           }
 
           updateAttemptFriendlyId(createAttempt.executionPayload.execution.attempt.id);
+          updateAttemptNumber(createAttempt.executionPayload.execution.attempt.number);
 
           callback({
             success: true,
@@ -928,6 +955,10 @@ class TaskCoordinator {
           if (message.attemptFriendlyId) {
             updateAttemptFriendlyId(message.attemptFriendlyId);
           }
+
+          if (message.attemptNumber) {
+            updateAttemptNumber(message.attemptNumber);
+          }
         });
       },
       onDisconnect: async (socket, handler, sender, logger) => {

apps/docker-provider/src/index.ts

@@ -109,7 +109,7 @@ class DockerTaskOperations implements TaskOperations {
   async create(opts: TaskOperationsCreateOptions) {
     await this.init();
 
-    const containerName = this.#getRunContainerName(opts.runId);
+    const containerName = this.#getRunContainerName(opts.runId, opts.nextAttemptNumber);
 
     const runArgs = [
       "run",
@@ -150,7 +150,7 @@ class DockerTaskOperations implements TaskOperations {
   async restore(opts: TaskOperationsRestoreOptions) {
     await this.init();
 
-    const containerName = this.#getRunContainerName(opts.runId);
+    const containerName = this.#getRunContainerName(opts.runId, opts.attemptNumber);
 
     if (!this.#canCheckpoint || this.opts.forceSimulate) {
       logger.log("Simulating restore");
@@ -195,8 +195,8 @@ class DockerTaskOperations implements TaskOperations {
     return `task-index-${suffix}`;
   }
 
-  #getRunContainerName(suffix: string) {
-    return `task-run-${suffix}`;
+  #getRunContainerName(suffix: string, attemptNumber?: number) {
+    return `task-run-${suffix}${attemptNumber && attemptNumber > 1 ? `-att${attemptNumber}` : ""}`;
   }
 
   async #sendPostStart(containerName: string): Promise<void> {

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

@@ -1,11 +1,9 @@
 import { TaskRunFailedExecutionResult } from "@trigger.dev/core/v3";
-import { TaskRunStatus } from "@trigger.dev/database";
 import { logger } from "~/services/logger.server";
 import { createExceptionPropertiesFromError, eventRepository } from "./eventRepository.server";
 import { BaseService } from "./services/baseService.server";
 import { FinalizeTaskRunService } from "./services/finalizeTaskRun.server";
-
-const FAILABLE_TASK_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "PENDING", "WAITING_FOR_DEPLOY"];
+import { FAILABLE_RUN_STATUSES } from "./taskStatus";
 
 export class FailedTaskRunService extends BaseService {
   public async call(anyRunId: string, completion: TaskRunFailedExecutionResult) {
@@ -27,7 +25,7 @@ export class FailedTaskRunService extends BaseService {
       return;
     }
 
-    if (!FAILABLE_TASK_RUN_STATUSES.includes(taskRun.status)) {
+    if (!FAILABLE_RUN_STATUSES.includes(taskRun.status)) {
       logger.error("[FailedTaskRunService] Task run is not in a failable state", {
         taskRun,
         completion,

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

@@ -128,6 +128,7 @@ function createCoordinatorNamespace(io: Server) {
           completion: message.completion,
           execution: message.execution,
           checkpoint: message.checkpoint,
+          supportsRetryCheckpoints: message.version === "v1",
         });
       },
       TASK_RUN_FAILED_TO_RUN: async (message) => {

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

@@ -54,6 +54,7 @@ export const SharedQueueMessageBody = z.discriminatedUnion("type", [
     type: z.literal("EXECUTE"),
     taskIdentifier: z.string(),
     checkpointEventId: z.string().optional(),
+    retryCheckpointsDisabled: z.boolean().optional(),
   }),
   WithTraceContext.extend({
     type: z.literal("RESUME"),
@@ -294,12 +295,9 @@ export class SharedQueueConsumer {
 
         const retryingFromCheckpoint = !!messageBody.data.checkpointEventId;
 
-        const EXECUTABLE_RUN_STATUSES: {
-          fromCheckpoint: TaskRunStatus[];
-          withoutCheckpoint: TaskRunStatus[];
-        } = {
-          fromCheckpoint: ["WAITING_TO_RESUME"],
-          withoutCheckpoint: ["PENDING", "RETRYING_AFTER_FAILURE"],
+        const EXECUTABLE_RUN_STATUSES = {
+          fromCheckpoint: ["WAITING_TO_RESUME"] satisfies TaskRunStatus[],
+          withoutCheckpoint: ["PENDING", "RETRYING_AFTER_FAILURE"] satisfies TaskRunStatus[],
         };
 
         if (
@@ -474,7 +472,10 @@ export class SharedQueueConsumer {
           ? lockedTaskRun.attempts[0].number + 1
           : 1;
 
-        const isRetry = lockedTaskRun.status === "WAITING_TO_RESUME" && nextAttemptNumber > 1;
+        const isRetry =
+          nextAttemptNumber > 1 &&
+          (lockedTaskRun.status === "WAITING_TO_RESUME" ||
+            lockedTaskRun.status === "RETRYING_AFTER_FAILURE");
 
         try {
           if (messageBody.data.checkpointEventId) {
@@ -515,11 +516,13 @@ export class SharedQueueConsumer {
             }
           }
 
-          if (isRetry) {
+          if (isRetry && !messageBody.data.retryCheckpointsDisabled) {
             socketIo.coordinatorNamespace.emit("READY_FOR_RETRY", {
               version: "v1",
               runId: lockedTaskRun.id,
             });
+
+            // Retries for workers with disabled retry checkpoints will be handled just like normal attempts
           } else {
             const machineConfig = lockedTaskRun.lockedBy?.machineConfig;
             const machine = machinePresetFromConfig(machineConfig ?? {});
@@ -531,6 +534,7 @@ export class SharedQueueConsumer {
                 image: deployment.imageReference,
                 version: deployment.version,
                 machine,
+                nextAttemptNumber,
                 // identifiers
                 id: "placeholder", // TODO: Remove this completely in a future release
                 envId: lockedTaskRun.runtimeEnvironment.id,