Make sourceTaskAttemptId optional on resumeBatchRun

Author: matt-aitken

apps/webapp/app/services/worker.server.ts

@@ -131,7 +131,7 @@ const workerCatalog = {
   }),
   "v3.resumeBatchRun": z.object({
     batchRunId: z.string(),
-    sourceTaskAttemptId: z.string(),
+    sourceTaskAttemptId: z.string().optional(),
   }),
   "v3.resumeTaskDependency": z.object({
     dependencyId: z.string(),

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

@@ -617,15 +617,15 @@ export class SharedQueueConsumer {
           return;
         }
 
-        if (messageBody.data.completedAttemptIds.length < 1) {
-          logger.error("No attempt IDs provided", {
-            queueMessage: message.data,
-            messageId: message.messageId,
-          });
-
-          await this.#ackAndDoMoreWork(message.messageId);
-          return;
-        }
+        // if (messageBody.data.completedAttemptIds.length < 1) {
+        //   logger.error("No attempt IDs provided", {
+        //     queueMessage: message.data,
+        //     messageId: message.messageId,
+        //   });
+
+        //   await this.#ackAndDoMoreWork(message.messageId);
+        //   return;
+        // }
 
         const resumableRun = await prisma.taskRun.findUnique({
           where: {

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

@@ -7,6 +7,8 @@ import { marqs } from "~/v3/marqs/index.server";
 import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server";
 import { BaseService } from "./baseService.server";
 import { isFinalRunStatus, isFreezableAttemptStatus, isFreezableRunStatus } from "../taskStatus";
+import { ResumeBatchRunService } from "./resumeBatchRun.server";
+import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
 
 export class CreateCheckpointService extends BaseService {
   public async call(
@@ -89,6 +91,9 @@ export class CreateCheckpointService extends BaseService {
       };
     }
 
+    //sleep for 20 seconds
+    // await new Promise((resolve) => setTimeout(resolve, 10_000));
+
     const checkpoint = await this._prisma.checkpoint.create({
       data: {
         friendlyId: generateFriendlyId("checkpoint"),
@@ -154,10 +159,33 @@ export class CreateCheckpointService extends BaseService {
           batchDependencyFriendlyId: reason.batchFriendlyId,
         });
 
-        keepRunAlive = await this.#isBatchCompleted(reason.batchFriendlyId);
+        if (checkpointEvent) {
+          const batchRun = await this._prisma.batchTaskRun.findFirst({
+            select: {
+              id: true,
+            },
+            where: {
+              friendlyId: reason.batchFriendlyId,
+            },
+          });
 
-        if (!keepRunAlive) {
-          await marqs?.acknowledgeMessage(attempt.taskRunId);
+          if (!batchRun) {
+            logger.error("Batch not found", { friendlyId: reason.batchFriendlyId });
+            await marqs?.acknowledgeMessage(attempt.taskRunId);
+
+            return {
+              success: false,
+            };
+          }
+
+          await ResumeBatchRunService.enqueue(batchRun.id, undefined, this._prisma);
+
+          return {
+            success: true,
+            checkpoint,
+            event: checkpointEvent,
+            keepRunAlive: false,
+          };
         }
 
         break;

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

@@ -5,7 +5,7 @@ import { BaseService } from "./baseService.server";
 import { logger } from "~/services/logger.server";
 
 export class ResumeBatchRunService extends BaseService {
-  public async call(batchRunId: string, sourceTaskAttemptId: string) {
+  public async call(batchRunId: string, sourceTaskAttemptId?: string) {
     const batchRun = await this._prisma.batchTaskRun.findFirst({
       where: {
         id: batchRunId,
@@ -41,34 +41,43 @@ export class ResumeBatchRunService extends BaseService {
       return;
     }
 
-    // We need to update the batchRun status so we don't resume it again
-    await this._prisma.batchTaskRun.update({
-      where: {
-        id: batchRun.id,
-      },
-      data: {
-        status: "COMPLETED",
-      },
-    });
-
     // This batch has a dependent attempt and just finalized, we should resume that attempt
     const environment = batchRun.dependentTaskAttempt.runtimeEnvironment;
 
     // If we are in development, we don't need to resume the dependent task (that will happen automatically)
     if (environment.type === "DEVELOPMENT") {
+      // We need to update the batchRun status so we don't resume it again
+      await this._prisma.batchTaskRun.update({
+        where: {
+          id: batchRun.id,
+        },
+        data: {
+          status: "COMPLETED",
+        },
+      });
       return;
     }
 
     const dependentRun = batchRun.dependentTaskAttempt.taskRun;
 
     if (batchRun.dependentTaskAttempt.status === "PAUSED" && batchRun.checkpointEventId) {
+      // We need to update the batchRun status so we don't resume it again
+      await this._prisma.batchTaskRun.update({
+        where: {
+          id: batchRun.id,
+        },
+        data: {
+          status: "COMPLETED",
+        },
+      });
+
       await marqs?.enqueueMessage(
         environment,
         dependentRun.queue,
         dependentRun.id,
         {
           type: "RESUME",
-          completedAttemptIds: [sourceTaskAttemptId],
+          completedAttemptIds: sourceTaskAttemptId ? [sourceTaskAttemptId] : [],
           resumableAttemptId: batchRun.dependentTaskAttempt.id,
           checkpointEventId: batchRun.checkpointEventId,
           projectId: batchRun.dependentTaskAttempt.runtimeEnvironment.projectId,
@@ -86,18 +95,12 @@ export class ResumeBatchRunService extends BaseService {
           dependentTaskAttemptId: batchRun.dependentTaskAttempt.id,
         });
       }
-
-      await marqs?.replaceMessage(dependentRun.id, {
-        type: "RESUME",
-        completedAttemptIds: batchRun.items.map((item) => item.taskRunAttemptId).filter(Boolean),
-        resumableAttemptId: batchRun.dependentTaskAttempt.id,
-      });
     }
   }
 
   static async enqueue(
     batchRunId: string,
-    sourceTaskAttemptId: string,
+    sourceTaskAttemptId: string | undefined,
     tx: PrismaClientOrTransaction,
     runAt?: Date
   ) {