diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index 45d72023c..d05a02a9f 100644 --- a/apps/coordinator/src/index.ts +++ b/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) => { diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts index d4601cb68..02c2c63dc 100644 --- a/apps/webapp/app/v3/handleSocketIo.server.ts +++ b/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); diff --git a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts index 637ae2c5c..46bad40f7 100644 --- a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts +++ b/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); diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index 9c76a3436..b098906ff 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/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>; +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, 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, - 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 diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index c0f34f18b..9a91e0120 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -14,14 +14,25 @@ export class CreateCheckpointService { } public async call( - params: InferSocketMessageSchema + params: Omit< + InferSocketMessageSchema, + "version" + > ): Promise { 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", diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index 90c517b7b..85c6da359 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/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: {