Merge branch 'main' into v3/self-hosting-latest

Author: nicktrn

apps/coordinator/src/index.ts

@@ -162,6 +162,49 @@ class TaskCoordinator {
 
           taskSocket.emit("RESUME_AFTER_DEPENDENCY", message);
         },
+        RESUME_AFTER_DEPENDENCY_WITH_ACK: async (message) => {
+          const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
+
+          if (!taskSocket) {
+            logger.log("Socket for attempt not found", {
+              attemptFriendlyId: message.attemptFriendlyId,
+            });
+            return {
+              success: false,
+              error: {
+                name: "SocketNotFoundError",
+                message: "Socket for attempt not found",
+              },
+            };
+          }
+
+          //if this is set, we want to kill the process because it will be resumed with the checkpoint from the queue
+          if (taskSocket.data.requiresCheckpointResumeWithMessage) {
+            logger.log("RESUME_AFTER_DEPENDENCY_WITH_ACK: Checkpoint is set so going to nack", {
+              socketData: taskSocket.data,
+            });
+
+            return {
+              success: false,
+              error: {
+                name: "CheckpointMessagePresentError",
+                message:
+                  "Checkpoint message is present, so we need to kill the process and resume from the queue.",
+              },
+            };
+          }
+
+          await chaosMonkey.call();
+
+          // In case the task resumed faster than we could checkpoint
+          this.#cancelCheckpoint(message.runId);
+
+          taskSocket.emit("RESUME_AFTER_DEPENDENCY", message);
+
+          return {
+            success: true,
+          };
+        },
         RESUME_AFTER_DURATION: async (message) => {
           const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
 
@@ -792,6 +835,18 @@ class TaskCoordinator {
             return;
           }
 
+          logger.log("WAIT_FOR_TASK checkpoint created", {
+            checkpoint,
+            socketData: socket.data,
+          });
+
+          //setting this means we can only resume from a checkpoint
+          socket.data.requiresCheckpointResumeWithMessage = `location:${checkpoint.location}-docker:${checkpoint.docker}`;
+          logger.log("WAIT_FOR_TASK set requiresCheckpointResumeWithMessage", {
+            checkpoint,
+            socketData: socket.data,
+          });
+
           const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", {
             version: "v1",
             attemptFriendlyId: message.attemptFriendlyId,
@@ -804,6 +859,7 @@ class TaskCoordinator {
           });
 
           if (ack?.keepRunAlive) {
+            socket.data.requiresCheckpointResumeWithMessage = undefined;
             logger.log("keeping run alive after task checkpoint", { runId: socket.data.runId });
             return;
           }
@@ -862,6 +918,18 @@ class TaskCoordinator {
             return;
           }
 
+          logger.log("WAIT_FOR_BATCH checkpoint created", {
+            checkpoint,
+            socketData: socket.data,
+          });
+
+          //setting this means we can only resume from a checkpoint
+          socket.data.requiresCheckpointResumeWithMessage = `location:${checkpoint.location}-docker:${checkpoint.docker}`;
+          logger.log("WAIT_FOR_BATCH set checkpoint", {
+            checkpoint,
+            socketData: socket.data,
+          });
+
           const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", {
             version: "v1",
             attemptFriendlyId: message.attemptFriendlyId,
@@ -875,6 +943,7 @@ class TaskCoordinator {
           });
 
           if (ack?.keepRunAlive) {
+            socket.data.requiresCheckpointResumeWithMessage = undefined;
             logger.log("keeping run alive after batch checkpoint", { runId: socket.data.runId });
             return;
           }

apps/webapp/app/hooks/useSearchParam.ts

@@ -18,13 +18,13 @@ export function useSearchParams() {
         }
 
         if (typeof value === "string") {
-          search.set(param, encodeURIComponent(value));
+          search.set(param, value);
           continue;
         }
 
         search.delete(param);
         for (const v of value) {
-          search.append(param, encodeURIComponent(v));
+          search.append(param, v);
         }
       }
     },

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

@@ -725,20 +725,47 @@ export class SharedQueueConsumer {
         }
 
         try {
-          logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY", {
-            runId: resumableAttempt.taskRunId,
-            attemptId: resumableAttempt.id,
-          });
-
-          // The attempt should still be running so we can broadcast to all coordinators to resume immediately
-          socketIo.coordinatorNamespace.emit("RESUME_AFTER_DEPENDENCY", {
-            version: "v1",
+          const resumeMessage = {
+            version: "v1" as const,
             runId: resumableAttempt.taskRunId,
             attemptId: resumableAttempt.id,
             attemptFriendlyId: resumableAttempt.friendlyId,
             completions,
             executions,
+          };
+
+          logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY_WITH_ACK", { resumeMessage, message });
+
+          // The attempt should still be running so we can broadcast to all coordinators to resume immediately
+          const responses = await socketIo.coordinatorNamespace
+            .timeout(10_000)
+            .emitWithAck("RESUME_AFTER_DEPENDENCY_WITH_ACK", resumeMessage);
+
+          logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK received", {
+            resumeMessage,
+            responses,
+            message,
           });
+
+          if (responses.length === 0) {
+            logger.error("RESUME_AFTER_DEPENDENCY_WITH_ACK no response", {
+              resumeMessage,
+              message,
+            });
+            await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval, 5_000);
+            return;
+          }
+
+          const hasSuccess = responses.some((response) => response.success);
+          if (!hasSuccess) {
+            logger.warn("RESUME_AFTER_DEPENDENCY_WITH_ACK failed", {
+              resumeMessage,
+              responses,
+              message,
+            });
+            await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval, 5_000);
+            return;
+          }
         } catch (e) {
           if (e instanceof Error) {
             this._currentSpan?.recordException(e);
@@ -748,7 +775,12 @@ export class SharedQueueConsumer {
 
           this._endSpanInNextIteration = true;
 
-          await this.#nackAndDoMoreWork(message.messageId);
+          logger.error("RESUME_AFTER_DEPENDENCY_WITH_ACK threw, nacking with delay", {
+            message,
+            error: e,
+          });
+
+          await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval, 5_000);
           return;
         }
 

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

@@ -122,7 +122,9 @@ export class ResumeBatchRunService extends BaseService {
         // When the checkpoint is created, it will continue the run
         logger.error("ResumeBatchRunService: attempt is paused but there's no checkpoint event", {
           batchRunId: batchRun.id,
-          dependentTaskAttemptId: batchRun.dependentTaskAttempt.id,
+          dependentTaskAttempt: batchRun.dependentTaskAttempt,
+          checkpointEventId: batchRun.checkpointEventId,
+          hasCheckpointEvent: !!batchRun.checkpointEventId,
         });
         return;
       }
@@ -132,7 +134,9 @@ export class ResumeBatchRunService extends BaseService {
       if (wasUpdated) {
         logger.debug("ResumeBatchRunService: Resuming dependent run without checkpoint", {
           batchRunId: batchRun.id,
-          dependentTaskAttemptId: batchRun.dependentTaskAttempt.id,
+          dependentTaskAttempt: batchRun.dependentTaskAttempt,
+          checkpointEventId: batchRun.checkpointEventId,
+          hasCheckpointEvent: !!batchRun.checkpointEventId,
         });
         await marqs?.replaceMessage(dependentRun.id, {
           type: "RESUME",

docs/examples/dall-e3-generate-image.mdx

@@ -29,7 +29,6 @@ export const generateContent = task({
     maxAttempts: 3, // Retry up to 3 times
   },
   run: async ({ theme, description }: Payload) => {
-
     // Generate text
     const textResult = await openai.chat.completions.create({
       model: "gpt-4o",
@@ -64,4 +63,15 @@ function generateTextPrompt(theme: string, description: string): any {
 function generateImagePrompt(theme: string, description: string): any {
   return `Theme: ${theme}\n\nDescription: ${description}`;
 }
-```
+```
+
+## Testing your task
+
+To test this task in the dashboard, you can use the following payload:
+
+```json
+{
+  "theme": "A beautiful sunset",
+  "description": "A sunset over the ocean with a tiny yacht in the distance."
+}
+```