Cancel in progress runs when disconnecting the dev CLI (in between attempts)

Author: ericallam

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

@@ -611,14 +611,14 @@ export class MarQS {
       String(this.options.defaultConcurrency ?? 10)
     );
 
-    logger.debug("Dequeue message result", {
-      result,
-    });
-
     if (!result) {
       return;
     }
 
+    logger.debug("Dequeue message result", {
+      result,
+    });
+
     if (result.length !== 2) {
       return;
     }

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

@@ -18,6 +18,7 @@ import { CancelAttemptService } from "../services/cancelAttempt.server";
 import { CompleteAttemptService } from "../services/completeAttempt.server";
 import { attributesFromAuthenticatedEnv } from "../tracer.server";
 import { DevSubscriber, devPubSub } from "./devPubSub.server";
+import { CancelTaskRunService } from "../services/cancelTaskRun.server";
 
 const tracer = trace.getTracer("devQueueConsumer");
 
@@ -50,6 +51,7 @@ export class DevQueueConsumer {
   private _currentSpan: Span | undefined;
   private _endSpanInNextIteration = false;
   private _inProgressAttempts: Map<string, string> = new Map(); // Keys are task attempt friendly IDs, values are TaskRun ids/queue message ids
+  private _inProgressRuns: Map<string, string> = new Map(); // Keys are task run friendly IDs, values are TaskRun internal ids/queue message ids
 
   constructor(
     public env: AuthenticatedEnvironment,
@@ -123,7 +125,11 @@ export class DevQueueConsumer {
     logger.debug("Task run completed", { taskRunCompletion: completion, execution });
 
     const service = new CompleteAttemptService();
-    await service.call(completion, execution, this.env);
+    const result = await service.call(completion, execution, this.env);
+
+    if (result === "COMPLETED") {
+      this._inProgressRuns.delete(execution.run.id);
+    }
   }
 
   public async taskHeartbeat(workerId: string, id: string, seconds: number = 60) {
@@ -148,7 +154,7 @@ export class DevQueueConsumer {
     this._enabled = false;
 
     // We need to cancel all the in progress task run attempts and ack the messages so they will stop processing
-    await this.#cancelInProgressAttempts(reason);
+    await this.#cancelInProgressRunsAndAttempts(reason);
 
     // We need to unsubscribe from the background worker channels
     for (const [id, subscriber] of this._backgroundWorkerSubscriber) {
@@ -161,21 +167,44 @@ export class DevQueueConsumer {
     }
   }
 
-  async #cancelInProgressAttempts(reason: string) {
-    const service = new CancelAttemptService();
+  async #cancelInProgressRunsAndAttempts(reason: string) {
+    const cancelAttemptService = new CancelAttemptService();
+    const cancelTaskRunService = new CancelTaskRunService();
 
     const cancelledAt = new Date();
 
     const inProgressAttempts = new Map(this._inProgressAttempts);
+    const inProgressRuns = new Map(this._inProgressRuns);
 
     this._inProgressAttempts.clear();
+    this._inProgressRuns.clear();
 
-    logger.debug("Cancelling in progress attempts", {
+    const inProgressRunsWithNoInProgressAttempts: string[] = [];
+    const inProgressAttemptRunIds = new Set(inProgressAttempts.values());
+
+    for (const [runId, messageId] of inProgressRuns) {
+      if (!inProgressAttemptRunIds.has(messageId)) {
+        inProgressRunsWithNoInProgressAttempts.push(messageId);
+      }
+    }
+
+    logger.debug("Cancelling in progress runs and attempts", {
       attempts: Array.from(inProgressAttempts.keys()),
+      runs: Array.from(inProgressRuns.keys()),
     });
 
     for (const [attemptId, messageId] of inProgressAttempts) {
-      await this.#cancelInProgressAttempt(attemptId, messageId, service, cancelledAt, reason);
+      await this.#cancelInProgressAttempt(
+        attemptId,
+        messageId,
+        cancelAttemptService,
+        cancelledAt,
+        reason
+      );
+    }
+
+    for (const runId of inProgressRunsWithNoInProgressAttempts) {
+      await this.#cancelInProgressRun(runId, cancelTaskRunService, cancelledAt, reason);
     }
   }
 
@@ -199,6 +228,32 @@ export class DevQueueConsumer {
     }
   }
 
+  async #cancelInProgressRun(
+    runId: string,
+    service: CancelTaskRunService,
+    cancelledAt: Date,
+    reason: string
+  ) {
+    logger.debug("Cancelling in progress run", { runId });
+
+    const taskRun = await prisma.taskRun.findUnique({
+      where: { id: runId },
+    });
+
+    if (!taskRun) {
+      return;
+    }
+
+    try {
+      await service.call(taskRun, { reason, cancelAttempts: false, cancelledAt });
+    } catch (e) {
+      logger.error("Failed to cancel in progress run", {
+        runId,
+        error: e,
+      });
+    }
+  }
+
   #enable() {
     if (this._enabled) {
       return;
@@ -481,6 +536,7 @@ export class DevQueueConsumer {
       });
 
       this._inProgressAttempts.set(taskRunAttempt.friendlyId, message.messageId);
+      this._inProgressRuns.set(lockedTaskRun.friendlyId, message.messageId);
     } catch (e) {
       if (e instanceof Error) {
         this._currentSpan?.recordException(e);
@@ -499,6 +555,7 @@ export class DevQueueConsumer {
           data: {
             lockedAt: null,
             lockedById: null,
+            status: "PENDING",
           },
         }),
         prisma.taskRunAttempt.delete({
@@ -508,6 +565,9 @@ export class DevQueueConsumer {
         }),
       ]);
 
+      this._inProgressAttempts.delete(taskRunAttempt.friendlyId);
+      this._inProgressRuns.delete(lockedTaskRun.friendlyId);
+
       // Finally we need to nack the message so it can be retried
       await marqs?.nackMessage(message.messageId);
     } finally {

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

@@ -23,8 +23,21 @@ const CANCELLABLE_ATTEMPT_STATUSES: Array<TaskRunAttemptStatus> = [
   "PENDING",
 ];
 
+export type CancelTaskRunServiceOptions = {
+  reason?: string;
+  cancelAttempts?: boolean;
+  cancelledAt?: Date;
+};
+
 export class CancelTaskRunService extends BaseService {
-  public async call(taskRun: TaskRun) {
+  public async call(taskRun: TaskRun, options?: CancelTaskRunServiceOptions) {
+    const opts = {
+      reason: "Task run was cancelled by user",
+      cancelAttempts: true,
+      cancelledAt: new Date(),
+      ...options,
+    };
+
     // Make sure the task run is in a cancellable state
     if (!CANCELLABLE_STATUSES.includes(taskRun.status)) {
       return;
@@ -68,59 +81,61 @@ export class CancelTaskRunService extends BaseService {
 
     await Promise.all(
       inProgressEvents.map((event) => {
-        return eventRepository.cancelEvent(event, new Date(), "Task run was cancelled by user");
+        return eventRepository.cancelEvent(event, opts.cancelledAt, opts.reason);
       })
     );
 
     // Cancel any in progress attempts
-    for (const attempt of cancelledTaskRun.attempts) {
-      if (attempt.runtimeEnvironment.type === "DEVELOPMENT") {
-        // Signal the task run attempt to stop
-        await devPubSub.publish(
-          `backgroundWorker:${attempt.backgroundWorkerId}:${attempt.id}`,
-          "CANCEL_ATTEMPT",
-          {
-            attemptId: attempt.friendlyId,
-            backgroundWorkerId: attempt.backgroundWorker.friendlyId,
-            taskRunId: cancelledTaskRun.friendlyId,
-          }
-        );
-      } else {
-        switch (attempt.status) {
-          case "EXECUTING": {
-            // We need to send a cancel message to the coordinator
-            socketIo.coordinatorNamespace.emit("REQUEST_ATTEMPT_CANCELLATION", {
-              version: "v1",
-              attemptId: attempt.id,
-            });
+    if (opts.cancelAttempts) {
+      for (const attempt of cancelledTaskRun.attempts) {
+        if (attempt.runtimeEnvironment.type === "DEVELOPMENT") {
+          // Signal the task run attempt to stop
+          await devPubSub.publish(
+            `backgroundWorker:${attempt.backgroundWorkerId}:${attempt.id}`,
+            "CANCEL_ATTEMPT",
+            {
+              attemptId: attempt.friendlyId,
+              backgroundWorkerId: attempt.backgroundWorker.friendlyId,
+              taskRunId: cancelledTaskRun.friendlyId,
+            }
+          );
+        } else {
+          switch (attempt.status) {
+            case "EXECUTING": {
+              // We need to send a cancel message to the coordinator
+              socketIo.coordinatorNamespace.emit("REQUEST_ATTEMPT_CANCELLATION", {
+                version: "v1",
+                attemptId: attempt.id,
+              });
 
-            break;
-          }
-          case "PENDING":
-          case "PAUSED": {
-            logger.debug("Cancelling pending or paused attempt", {
-              attempt,
-            });
+              break;
+            }
+            case "PENDING":
+            case "PAUSED": {
+              logger.debug("Cancelling pending or paused attempt", {
+                attempt,
+              });
 
-            const service = new CancelAttemptService();
+              const service = new CancelAttemptService();
 
-            await service.call(
-              attempt.friendlyId,
-              taskRun.id,
-              new Date(),
-              "Task run was cancelled by user"
-            );
+              await service.call(
+                attempt.friendlyId,
+                taskRun.id,
+                new Date(),
+                "Task run was cancelled by user"
+              );
 
-            break;
-          }
-          case "CANCELED":
-          case "COMPLETED":
-          case "FAILED": {
-            // Do nothing
-            break;
-          }
-          default: {
-            assertUnreachable(attempt.status);
+              break;
+            }
+            case "CANCELED":
+            case "COMPLETED":
+            case "FAILED": {
+              // Do nothing
+              break;
+            }
+            default: {
+              assertUnreachable(attempt.status);
+            }
           }
         }
       }

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

@@ -25,7 +25,7 @@ export class CompleteAttemptService extends BaseService {
     completion: TaskRunExecutionResult,
     execution: TaskRunExecution,
     env?: AuthenticatedEnvironment
-  ) {
+  ): Promise<"COMPLETED" | "RETRIED"> {
     const taskRunAttempt = await findAttempt(this._prisma, completion.id);
 
     if (!taskRunAttempt) {
@@ -41,7 +41,7 @@ export class CompleteAttemptService extends BaseService {
         },
       });
 
-      return "FAILED";
+      return "COMPLETED";
     }
 
     if (completion.ok) {
@@ -55,7 +55,7 @@ export class CompleteAttemptService extends BaseService {
     completion: TaskRunSuccessfulExecutionResult,
     taskRunAttempt: NonNullable<FoundAttempt>,
     env?: AuthenticatedEnvironment
-  ) {
+  ): Promise<"COMPLETED" | "RETRIED"> {
     await this._prisma.taskRunAttempt.update({
       where: { friendlyId: completion.id },
       data: {
@@ -90,7 +90,7 @@ export class CompleteAttemptService extends BaseService {
       await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma);
     }
 
-    return "ACKNOWLEDGED";
+    return "COMPLETED";
   }
 
   async #completeAttemptFailed(
@@ -106,13 +106,15 @@ export class CompleteAttemptService extends BaseService {
       // We need to cancel the task run instead of fail it
       const cancelService = new CancelAttemptService();
 
-      return await cancelService.call(
+      await cancelService.call(
         taskRunAttempt.friendlyId,
         taskRunAttempt.taskRunId,
         new Date(),
         "Cancelled by user",
         env
       );
+
+      return "COMPLETED";
     }
 
     await this._prisma.taskRunAttempt.update({
@@ -204,7 +206,7 @@ export class CompleteAttemptService extends BaseService {
         await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma);
       }
 
-      return "ACKNOWLEDGED";
+      return "COMPLETED";
     }
   }
 

references/v3-catalog/src/trigger/openai.ts

@@ -25,7 +25,7 @@ export const openaiTask = task({
 
       return {
         error: new Error("Custom OpenAI API error"),
-        retryDelayInMs: 1000,
+        retryDelayInMs: 10000,
       };
     }
   },