atomic post-completion checkpoints

apps/coordinator/src/index.ts

@@ -151,10 +151,10 @@ class Checkpointer {
     );
 
     try {
+      this.#logger.log("Checkpointing:", { opts });
+
       // Create checkpoint (docker)
       if (this.#dockerMode) {
-        this.#logger.log("Checkpointing:", opts.podName);
-
         try {
           if (this.opts.forceSimulate || !this.#canCheckpoint) {
             this.#logger.log("Simulating checkpoint");
@@ -436,22 +436,26 @@ class TaskCoordinator {
         socket.on("TASK_RUN_COMPLETED", async ({ completion, execution }, callback) => {
           logger.log("completed task", { completionId: completion.id });
 
-          const sendCompletionToPlatform = () => {
-            this.#platformSocket?.send("TASK_RUN_COMPLETED", {
-              version: "v1",
-              execution,
-              completion,
-            });
+          type CheckpointData = {
+            docker: boolean;
+            location: string;
           };
 
           const confirmCompletion = ({
             didCheckpoint,
             shouldExit,
+            checkpoint,
           }: {
             didCheckpoint: boolean;
             shouldExit: boolean;
+            checkpoint?: CheckpointData;
           }) => {
-            sendCompletionToPlatform();
+            this.#platformSocket?.send("TASK_RUN_COMPLETED", {
+              version: "v1",
+              execution,
+              completion,
+              checkpoint,
+            });
             callback({ didCheckpoint, shouldExit });
           };
 
@@ -494,21 +498,7 @@ class TaskCoordinator {
             return;
           }
 
-          this.#platformSocket?.send("CHECKPOINT_CREATED", {
-            version: "v1",
-            attemptId: socket.data.attemptId,
-            docker: checkpoint.docker,
-            location: checkpoint.location,
-            reason: {
-              type: "RETRYING_AFTER_FAILURE",
-              attemptNumber: execution.attempt.number,
-              // TODO: attach completion data here
-            },
-          });
-
-          // TODO: replace this with
-          // callback({ didCheckpoint: true, shouldExit: false });
-          confirmCompletion({ didCheckpoint: true, shouldExit: false });
+          confirmCompletion({ didCheckpoint: true, shouldExit: false, checkpoint });
         });
 
         socket.on("WAIT_FOR_DURATION", async (message, callback) => {

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

@@ -67,7 +67,11 @@ function createCoordinatorNamespace(io: Server) {
       },
       TASK_RUN_COMPLETED: async (message) => {
         const completeAttempt = new CompleteAttemptService();
-        await completeAttempt.call(message.completion, message.execution);
+        await completeAttempt.call({
+          completion: message.completion,
+          execution: message.execution,
+          checkpoint: message.checkpoint,
+        });
       },
       TASK_HEARTBEAT: async (message) => {
         await sharedQueueTasks.taskHeartbeat(message.attemptFriendlyId);

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

@@ -125,7 +125,7 @@ export class DevQueueConsumer {
     logger.debug("Task run completed", { taskRunCompletion: completion, execution });
 
     const service = new CompleteAttemptService();
-    const result = await service.call(completion, execution, this.env);
+    const result = await service.call({ completion, execution, env: this.env });
 
     if (result === "COMPLETED") {
       this._inProgressRuns.delete(execution.run.id);

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

@@ -17,15 +17,27 @@ import { BaseService } from "./baseService.server";
 import { CancelAttemptService } from "./cancelAttempt.server";
 import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
 import { MAX_TASK_RUN_ATTEMPTS } from "~/consts";
+import { CreateCheckpointService } from "./createCheckpoint.server";
 
 type FoundAttempt = Awaited<ReturnType<typeof findAttempt>>;
 
+type CheckpointData = {
+  docker: boolean;
+  location: string;
+};
+
 export class CompleteAttemptService extends BaseService {
-  public async call(
-    completion: TaskRunExecutionResult,
-    execution: TaskRunExecution,
-    env?: AuthenticatedEnvironment
-  ): Promise<"COMPLETED" | "RETRIED"> {
+  public async call({
+    completion,
+    execution,
+    env,
+    checkpoint,
+  }: {
+    completion: TaskRunExecutionResult;
+    execution: TaskRunExecution;
+    env?: AuthenticatedEnvironment;
+    checkpoint?: CheckpointData;
+  }): Promise<"COMPLETED" | "RETRIED"> {
     const taskRunAttempt = await findAttempt(this._prisma, completion.id);
 
     if (!taskRunAttempt) {
@@ -47,7 +59,13 @@ export class CompleteAttemptService extends BaseService {
     if (completion.ok) {
       return await this.#completeAttemptSuccessfully(completion, taskRunAttempt, env);
     } else {
-      return await this.#completeAttemptFailed(completion, execution, taskRunAttempt, env);
+      return await this.#completeAttemptFailed(
+        completion,
+        execution,
+        taskRunAttempt,
+        env,
+        checkpoint
+      );
     }
   }
 
@@ -55,7 +73,7 @@ export class CompleteAttemptService extends BaseService {
     completion: TaskRunSuccessfulExecutionResult,
     taskRunAttempt: NonNullable<FoundAttempt>,
     env?: AuthenticatedEnvironment
-  ): Promise<"COMPLETED" | "RETRIED"> {
+  ): Promise<"COMPLETED"> {
     await this._prisma.taskRunAttempt.update({
       where: { friendlyId: completion.id },
       data: {
@@ -97,8 +115,9 @@ export class CompleteAttemptService extends BaseService {
     completion: TaskRunFailedExecutionResult,
     execution: TaskRunExecution,
     taskRunAttempt: NonNullable<FoundAttempt>,
-    env?: AuthenticatedEnvironment
-  ) {
+    env?: AuthenticatedEnvironment,
+    checkpoint?: CheckpointData
+  ): Promise<"COMPLETED" | "RETRIED"> {
     if (
       completion.error.type === "INTERNAL_ERROR" &&
       completion.error.code === "TASK_RUN_CANCELLED"
@@ -166,18 +185,32 @@ export class CompleteAttemptService extends BaseService {
       if (environment.type === "DEVELOPMENT") {
         // This is already an EXECUTE message so we can just NACK
         await marqs?.nackMessage(taskRunAttempt.taskRunId, completion.retry.timestamp);
-      } else {
-        // We have to replace a potential RESUME with EXECUTE to correctly retry the attempt
-        await marqs?.replaceMessage(
-          taskRunAttempt.taskRunId,
-          {
-            type: "EXECUTE",
-            taskIdentifier: taskRunAttempt.taskRun.taskIdentifier,
-          },
-          completion.retry.timestamp
-        );
+        return "RETRIED";
       }
 
+      if (checkpoint) {
+        const createCheckpoint = new CreateCheckpointService(this._prisma);
+        await createCheckpoint.call({
+          attemptId: execution.attempt.id,
+          docker: checkpoint.docker,
+          location: checkpoint.location,
+          reason: {
+            type: "RETRYING_AFTER_FAILURE",
+            attemptNumber: execution.attempt.number,
+          },
+        });
+      }
+
+      // We have to replace a potential RESUME with EXECUTE to correctly retry the attempt
+      await marqs?.replaceMessage(
+        taskRunAttempt.taskRunId,
+        {
+          type: "EXECUTE",
+          taskIdentifier: taskRunAttempt.taskRun.taskIdentifier,
+        },
+        completion.retry.timestamp
+      );
+
       return "RETRIED";
     } else {
       // No more retries, we need to fail the task run

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

@@ -14,14 +14,25 @@ export class CreateCheckpointService {
   }
 
   public async call(
-    params: InferSocketMessageSchema<typeof CoordinatorToPlatformMessages, "CHECKPOINT_CREATED">
+    params: Omit<
+      InferSocketMessageSchema<typeof CoordinatorToPlatformMessages, "CHECKPOINT_CREATED">,
+      "version"
+    >
   ): Promise<Checkpoint | undefined> {
     logger.debug(`Creating checkpoint`, params);
 
-    const attempt = await this.#prismaClient.taskRunAttempt.findUniqueOrThrow({
-      where: {
-        id: params.attemptId,
-      },
+    const isFriendlyId = (id: string) => {
+      return id.startsWith("attempt_");
+    };
+
+    const attempt = await this.#prismaClient.taskRunAttempt.findUnique({
+      where: isFriendlyId(params.attemptId)
+        ? {
+            friendlyId: params.attemptId,
+          }
+        : {
+            id: params.attemptId,
+          },
       include: {
         taskRun: true,
         backgroundWorker: {
@@ -36,6 +47,11 @@ export class CreateCheckpointService {
       },
     });
 
+    if (!attempt) {
+      logger.error("Attempt not found", { attemptId: params.attemptId });
+      return;
+    }
+
     const imageRef = attempt.backgroundWorker.deployment?.imageReference;
 
     if (!imageRef) {
@@ -63,7 +79,7 @@ export class CreateCheckpointService {
 
     await this.#prismaClient.taskRunAttempt.update({
       where: {
-        id: params.attemptId,
+        id: attempt.id,
       },
       data: {
         status: "PAUSED",

packages/core/src/v3/schemas/schemas.ts

@@ -180,6 +180,12 @@ export const CoordinatorToPlatformMessages = {
       version: z.literal("v1").default("v1"),
       execution: ProdTaskRunExecution,
       completion: TaskRunExecutionResult,
+      checkpoint: z
+        .object({
+          docker: z.boolean(),
+          location: z.string(),
+        })
+        .optional(),
     }),
   },
   TASK_HEARTBEAT: {