diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index 23217c757..93f839da9 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -455,36 +455,39 @@ class TaskCoordinator { socket.on("READY_FOR_EXECUTION", async (message) => { logger.log("[READY_FOR_EXECUTION]", message); - const executionAck = await this.#platformSocket?.sendWithAck("READY_FOR_EXECUTION", { - version: "v1", - attemptId: message.attemptId, - runId: message.runId, - }); + try { + const executionAck = await this.#platformSocket?.sendWithAck( + "READY_FOR_EXECUTION", + message + ); - if (!executionAck) { - logger.error("no execution ack", { attemptId: socket.data.attemptId }); + if (!executionAck) { + logger.error("no execution ack", { attemptId: socket.data.attemptId }); - socket.emit("REQUEST_EXIT", { + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + + return; + } + + if (!executionAck.success) { + logger.error("failed to get execution payload", { attemptId: socket.data.attemptId }); + + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + + return; + } + + socket.emit("EXECUTE_TASK_RUN", { version: "v1", + executionPayload: executionAck.payload, }); - - return; + } catch (error) { + logger.error("Error", { error }); } - - if (!executionAck.success) { - logger.error("failed to get execution payload", { attemptId: socket.data.attemptId }); - - socket.emit("REQUEST_EXIT", { - version: "v1", - }); - - return; - } - - socket.emit("EXECUTE_TASK_RUN", { - version: "v1", - executionPayload: executionAck.payload, - }); }); socket.on("READY_FOR_RESUME", async (message) => { @@ -649,7 +652,7 @@ class TaskCoordinator { this.#platformSocket?.send("CHECKPOINT_CREATED", { version: "v1", - attemptId: socket.data.attemptId, + attemptFriendlyId: message.attemptFriendlyId, docker: checkpoint.docker, location: checkpoint.location, reason: { @@ -692,7 +695,7 @@ class TaskCoordinator { this.#platformSocket?.send("CHECKPOINT_CREATED", { version: "v1", - attemptId: socket.data.attemptId, + attemptFriendlyId: message.attemptFriendlyId, docker: checkpoint.docker, location: checkpoint.location, reason: { @@ -734,7 +737,7 @@ class TaskCoordinator { this.#platformSocket?.send("CHECKPOINT_CREATED", { version: "v1", - attemptId: socket.data.attemptId, + attemptFriendlyId: message.attemptFriendlyId, docker: checkpoint.docker, location: checkpoint.location, reason: { diff --git a/apps/docker-provider/src/index.ts b/apps/docker-provider/src/index.ts index 4419614f3..e13e0bec2 100644 --- a/apps/docker-provider/src/index.ts +++ b/apps/docker-provider/src/index.ts @@ -118,8 +118,6 @@ class DockerTaskOperations implements TaskOperations { stdout: error.stdout, stderr: error.stderr, }); - - throw new Error(`Index failed with: ${error.stderr || error.stdout}`); } } @@ -157,8 +155,6 @@ class DockerTaskOperations implements TaskOperations { stdout: error.stdout, stderr: error.stderr, }); - - throw new Error(`Create failed with: ${error.stderr || error.stdout}`); } } diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts index 02c2c63dc..7bdb7ee6e 100644 --- a/apps/webapp/app/v3/handleSocketIo.server.ts +++ b/apps/webapp/app/v3/handleSocketIo.server.ts @@ -52,7 +52,8 @@ function createCoordinatorNamespace(io: Server) { READY_FOR_EXECUTION: async (message) => { const payload = await sharedQueueTasks.getLatestExecutionPayloadFromRun( message.runId, - true + true, + !!message.totalCompletions ); if (!payload) { diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index e60838faa..01b643de5 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -10,7 +10,12 @@ import { ZodMessageSender, serverWebsocketMessages, } from "@trigger.dev/core/v3"; -import { BackgroundWorker, BackgroundWorkerTask } from "@trigger.dev/database"; +import { + BackgroundWorker, + BackgroundWorkerTask, + TaskRunAttemptStatus, + TaskRunStatus, +} from "@trigger.dev/database"; import { z } from "zod"; import { prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; @@ -29,15 +34,18 @@ const MessageBody = z.discriminatedUnion("type", [ z.object({ type: z.literal("EXECUTE"), taskIdentifier: z.string(), + checkpointEventId: z.string().optional(), }), z.object({ type: z.literal("RESUME"), completedAttemptIds: z.string().array(), resumableAttemptId: z.string(), + checkpointEventId: z.string().optional(), }), z.object({ type: z.literal("RESUME_AFTER_DURATION"), resumableAttemptId: z.string(), + checkpointEventId: z.string(), }), ]); @@ -180,7 +188,7 @@ export class SharedQueueConsumer { this._taskFailures = 0; this._taskSuccesses = 0; - this.#doWork().finally(() => { }); + this.#doWork().finally(() => {}); } async #doWork() { @@ -302,11 +310,13 @@ export class SharedQueueConsumer { return; } + const retryingFromCheckpoint = !!messageBody.data.checkpointEventId; + if ( - existingTaskRun.status !== "PENDING" && - existingTaskRun.status !== "RETRYING_AFTER_FAILURE" + (retryingFromCheckpoint && existingTaskRun.status !== "WAITING_TO_RESUME") || + (!retryingFromCheckpoint && existingTaskRun.status !== "PENDING") ) { - logger.debug("Task run is not pending, aborting", { + logger.debug("Task run has invalid status for execution", { queueMessage: message.data, messageId: message.messageId, taskRun: existingTaskRun.id, @@ -432,37 +442,38 @@ export class SharedQueueConsumer { }); try { - const latestCheckpoint = lockedTaskRun.checkpoints[0]; + if (messageBody.data.checkpointEventId) { + const restoreService = new RestoreCheckpointService(); - if (lockedTaskRun.status === "RETRYING_AFTER_FAILURE" && latestCheckpoint) { - if (latestCheckpoint.reason !== "RETRYING_AFTER_FAILURE") { - logger.error("Latest checkpoint is invalid", { + const checkpoint = await restoreService.call({ + eventId: messageBody.data.checkpointEventId, + isRetry: taskRunAttempt.number > 1, + }); + + if (!checkpoint) { + logger.error("Failed to restore checkpoint", { queueMessage: message.data, messageId: message.messageId, - resumableAttemptId: taskRunAttempt.id, - latestCheckpointId: latestCheckpoint.id, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); return; } - const restoreService = new RestoreCheckpointService(); - await restoreService.call({ checkpointId: latestCheckpoint.id }); - } else { - await this._sender.send("BACKGROUND_WORKER_MESSAGE", { - backgroundWorkerId: deployment.worker.friendlyId, - data: { - type: "SCHEDULE_ATTEMPT", - id: taskRunAttempt.id, - image: deployment.imageReference, - envId: environment.id, - runId: taskRunAttempt.taskRunId, - version: deployment.version, - }, - }); + return; } + await this._sender.send("BACKGROUND_WORKER_MESSAGE", { + backgroundWorkerId: deployment.worker.friendlyId, + data: { + type: "SCHEDULE_ATTEMPT", + id: taskRunAttempt.id, + image: deployment.imageReference, + envId: environment.id, + runId: taskRunAttempt.taskRunId, + version: deployment.version, + }, + }); + this._inProgressAttempts.set(taskRunAttempt.friendlyId, message.messageId); } catch (e) { if (e instanceof Error) { @@ -500,6 +511,39 @@ export class SharedQueueConsumer { } // Resume after dependency completed with no remaining retries case "RESUME": { + if (messageBody.data.checkpointEventId) { + try { + const restoreService = new RestoreCheckpointService(); + + const checkpoint = await restoreService.call({ + eventId: messageBody.data.checkpointEventId, + }); + + if (!checkpoint) { + logger.error("Failed to restore checkpoint", { + queueMessage: message.data, + messageId: message.messageId, + }); + await marqs?.acknowledgeMessage(message.messageId); + return; + } + } catch (e) { + if (e instanceof Error) { + this._currentSpan?.recordException(e); + } else { + this._currentSpan?.recordException(new Error(String(e))); + } + + this._endSpanInNextIteration = true; + + // Finally we need to nack the message so it can be retried + await marqs?.nackMessage(message.messageId); + return; + } finally { + setTimeout(() => this.#doWork(), this._options.interval); + } + } + if (messageBody.data.completedAttemptIds.length < 1) { logger.error("No attempt IDs provided", { queueMessage: message.data, @@ -570,26 +614,37 @@ export class SharedQueueConsumer { return; } - if (resumableAttempt.status === "PAUSED") { - // We need to restore the attempt from the latest checkpoint before we can resume - const latestCheckpoint = resumableAttempt.checkpoints[0]; + if (messageBody.data.checkpointEventId) { + try { + const restoreService = new RestoreCheckpointService(); - if (!latestCheckpoint) { - logger.error("No checkpoint found", { - queueMessage: message.data, - messageId: message.messageId, - resumableAttemptId: resumableAttempt.id, + const checkpoint = await restoreService.call({ + eventId: messageBody.data.checkpointEventId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + if (!checkpoint) { + logger.error("Failed to restore checkpoint", { + queueMessage: message.data, + messageId: message.messageId, + }); + await marqs?.acknowledgeMessage(message.messageId); + return; + } + } catch (e) { + if (e instanceof Error) { + this._currentSpan?.recordException(e); + } else { + this._currentSpan?.recordException(new Error(String(e))); + } + + this._endSpanInNextIteration = true; + + // Finally we need to nack the message so it can be retried + await marqs?.nackMessage(message.messageId); return; + } finally { + setTimeout(() => this.#doWork(), this._options.interval); } - - const restoreService = new RestoreCheckpointService(); - await restoreService.call({ checkpointId: latestCheckpoint.id }); - - setTimeout(() => this.#doWork(), this._options.interval); - return; } const completions: TaskRunExecutionResult[] = []; @@ -669,59 +724,21 @@ export class SharedQueueConsumer { } // Resume after duration-based wait case "RESUME_AFTER_DURATION": { - const resumableAttempt = await prisma.taskRunAttempt.findUnique({ - where: { - id: messageBody.data.resumableAttemptId, - }, - include: { - checkpoints: { - take: 1, - orderBy: { - createdAt: "desc", - }, - }, - taskRun: true, - }, - }); - - if (!resumableAttempt) { - logger.error("Resumable attempt not found", { - queueMessage: message.data, - messageId: message.messageId, - }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); - return; - } - - if (resumableAttempt.status !== "PAUSED") { - logger.error("Attempt not paused", { - queueMessage: message.data, - messageId: message.messageId, - }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); - return; - } - try { - // We need to restore the attempt from the latest checkpoint before we can resume - const latestCheckpoint = resumableAttempt.checkpoints[0]; + const restoreService = new RestoreCheckpointService(); - if (!latestCheckpoint) { - logger.error("No checkpoint found", { + const checkpoint = await restoreService.call({ + eventId: messageBody.data.checkpointEventId, + }); + + if (!checkpoint) { + logger.error("Failed to restore checkpoint", { queueMessage: message.data, messageId: message.messageId, - resumableAttemptId: resumableAttempt.id, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); return; } - - // The attempt will resume automatically after restore - const restoreService = new RestoreCheckpointService(); - await restoreService.call({ checkpointId: latestCheckpoint.id }); } catch (e) { if (e instanceof Error) { this._currentSpan?.recordException(e); @@ -800,7 +817,8 @@ class SharedQueueTasks { async getExecutionPayloadFromAttempt( id: string, - setToExecuting?: boolean + setToExecuting?: boolean, + isRetrying?: boolean ): Promise { const attempt = await prisma.taskRunAttempt.findUnique({ where: { @@ -859,26 +877,42 @@ class SharedQueueTasks { } if (setToExecuting) { + const FINAL_RUN_STATUSES: TaskRunStatus[] = [ + "CANCELED", + "COMPLETED_SUCCESSFULLY", + "COMPLETED_WITH_ERRORS", + "INTERRUPTED", + "SYSTEM_FAILURE", + ]; + const FINAL_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["CANCELED", "COMPLETED", "FAILED"]; + + if ( + FINAL_ATTEMPT_STATUSES.includes(attempt.status) || + FINAL_RUN_STATUSES.includes(attempt.taskRun.status) + ) { + logger.error("Status already in final state", { + attempt: { + id: attempt.id, + status: attempt.status, + }, + run: { + id: attempt.taskRunId, + status: attempt.taskRun.status, + }, + }); + return; + } + await prisma.taskRunAttempt.update({ where: { id, - taskRun: { - status: { - notIn: [ - "CANCELED", - "COMPLETED_SUCCESSFULLY", - "COMPLETED_WITH_ERRORS", - "SYSTEM_FAILURE", - ], - }, - }, }, data: { status: "EXECUTING", taskRun: { update: { data: { - status: "EXECUTING", + status: isRetrying ? "RETRYING_AFTER_FAILURE" : "EXECUTING", }, }, }, @@ -961,7 +995,8 @@ class SharedQueueTasks { async getLatestExecutionPayloadFromRun( id: string, - setToExecuting?: boolean + setToExecuting?: boolean, + isRetrying?: boolean ): Promise { const run = await prisma.taskRun.findUnique({ where: { @@ -984,7 +1019,7 @@ class SharedQueueTasks { return; } - return this.getExecutionPayloadFromAttempt(latestAttempt.id, setToExecuting); + return this.getExecutionPayloadFromAttempt(latestAttempt.id, setToExecuting, isRetrying); } async taskHeartbeat(attemptFriendlyId: string, seconds: number = 60) { diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index b098906ff..8a5e729d8 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/apps/webapp/app/v3/services/completeAttempt.server.ts @@ -18,6 +18,7 @@ import { CancelAttemptService } from "./cancelAttempt.server"; import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server"; import { MAX_TASK_RUN_ATTEMPTS } from "~/consts"; import { CreateCheckpointService } from "./createCheckpoint.server"; +import { TaskRun } from "@trigger.dev/database"; type FoundAttempt = Awaited>; @@ -188,27 +189,42 @@ export class CompleteAttemptService extends BaseService { 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, - }, - }); + if (!checkpoint) { + await this.#enqueueRetry(taskRunAttempt.taskRun, completion.retry.timestamp); + return "RETRIED"; } - // 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, + const createCheckpoint = new CreateCheckpointService(this._prisma); + const checkpointCreateResult = await createCheckpoint.call({ + attemptFriendlyId: execution.attempt.id, + docker: checkpoint.docker, + location: checkpoint.location, + reason: { + type: "RETRYING_AFTER_FAILURE", + attemptNumber: execution.attempt.number, }, - completion.retry.timestamp + }); + + if (!checkpointCreateResult) { + logger.error("Failed to create checkpoint", { checkpoint, execution: execution.run.id }); + + // Update the task run to be failed + await this._prisma.taskRun.update({ + where: { + friendlyId: execution.run.id, + }, + data: { + status: "SYSTEM_FAILURE", + }, + }); + + return "COMPLETED"; + } + + await this.#enqueueRetry( + taskRunAttempt.taskRun, + completion.retry.timestamp, + checkpointCreateResult.event.id ); return "RETRIED"; @@ -243,6 +259,19 @@ export class CompleteAttemptService extends BaseService { } } + async #enqueueRetry(run: TaskRun, retryTimestamp: number, checkpointEventId?: string) { + // We have to replace a potential RESUME with EXECUTE to correctly retry the attempt + return await marqs?.replaceMessage( + run.id, + { + type: "EXECUTE", + taskIdentifier: run.taskIdentifier, + checkpointEventId: checkpointEventId, + }, + retryTimestamp + ); + } + #generateMetadataAttributesForNextAttempt(execution: TaskRunExecution) { const context = TaskRunContext.parse(execution); diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index b62bb8e09..84ed45ff7 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -1,42 +1,32 @@ import { CoordinatorToPlatformMessages, InferSocketMessageSchema } from "@trigger.dev/core/v3"; -import type { Checkpoint } from "@trigger.dev/database"; -import { PrismaClient, prisma } from "~/db.server"; +import type { TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database"; import { logger } from "~/services/logger.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; import { marqs } from "../marqs.server"; import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server"; +import { BaseService } from "./baseService.server"; -export class CreateCheckpointService { - #prismaClient: PrismaClient; - - constructor(prismaClient: PrismaClient = prisma) { - this.#prismaClient = prismaClient; - } +const FREEZABLE_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "RETRYING_AFTER_FAILURE"]; +const FREEZABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["EXECUTING", "FAILED"]; +export class CreateCheckpointService extends BaseService { public async call( params: Omit< InferSocketMessageSchema, "version" > - ): Promise { + ) { logger.debug(`Creating checkpoint`, params); - 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, - }, + const attempt = await this._prisma.taskRunAttempt.findUnique({ + where: { + friendlyId: params.attemptFriendlyId, + }, include: { taskRun: true, backgroundWorker: { select: { + id: true, deployment: { select: { imageReference: true, @@ -48,18 +38,38 @@ export class CreateCheckpointService { }); if (!attempt) { - logger.error("Attempt not found", { attemptId: params.attemptId }); + logger.error("Attempt not found", { attemptFriendlyId: params.attemptFriendlyId }); + return; + } + + if ( + !FREEZABLE_ATTEMPT_STATUSES.includes(attempt.status) || + !FREEZABLE_RUN_STATUSES.includes(attempt.taskRun.status) + ) { + logger.error("Unfreezable state", { + attempt: { + id: attempt.id, + status: attempt.status, + }, + run: { + id: attempt.taskRunId, + status: attempt.taskRun.status, + }, + }); return; } const imageRef = attempt.backgroundWorker.deployment?.imageReference; if (!imageRef) { - logger.error("No image ref", { attemptId: params.attemptId }); + logger.error("Missing deployment or image ref", { + attemptId: attempt.id, + workerId: attempt.backgroundWorker.id, + }); return; } - const checkpoint = await this.#prismaClient.checkpoint.create({ + const checkpoint = await this._prisma.checkpoint.create({ data: { friendlyId: generateFriendlyId("checkpoint"), runtimeEnvironmentId: attempt.taskRun.runtimeEnvironmentId, @@ -74,21 +84,29 @@ export class CreateCheckpointService { }, }); - const eventService = new CreateCheckpointRestoreEventService(this.#prismaClient); - await eventService.call({ checkpointId: checkpoint.id, type: "CHECKPOINT" }); + const eventService = new CreateCheckpointRestoreEventService(this._prisma); + const checkpointEvent = await eventService.call({ + checkpointId: checkpoint.id, + type: "CHECKPOINT", + }); - await this.#prismaClient.taskRunAttempt.update({ + if (!checkpointEvent) { + logger.error("No checkpoint event", { + attemptId: attempt.id, + checkpointId: checkpoint.id, + }); + return; + } + + await this._prisma.taskRunAttempt.update({ where: { id: attempt.id, }, data: { - status: "PAUSED", + status: params.reason.type === "RETRYING_AFTER_FAILURE" ? undefined : "PAUSED", taskRun: { update: { - status: - params.reason.type === "RETRYING_AFTER_FAILURE" - ? "RETRYING_AFTER_FAILURE" - : "WAITING_TO_RESUME", + status: "WAITING_TO_RESUME", }, }, }, @@ -98,11 +116,16 @@ export class CreateCheckpointService { case "WAIT_FOR_DURATION": { await marqs?.replaceMessage( attempt.taskRunId, - { type: "RESUME_AFTER_DURATION", resumableAttemptId: attempt.id }, + { + type: "RESUME_AFTER_DURATION", + resumableAttemptId: attempt.id, + checkpointEventId: checkpointEvent.id, + }, params.reason.now + params.reason.ms ); break; } + // TODO: Attach the checkpoint event ID to in-progress dependencies case "WAIT_FOR_TASK": case "WAIT_FOR_BATCH": { await marqs?.acknowledgeMessage(attempt.taskRunId); @@ -117,6 +140,9 @@ export class CreateCheckpointService { } } - return checkpoint; + return { + checkpoint, + event: checkpointEvent, + }; } } diff --git a/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts b/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts index f4acf5adc..f2ea1a50d 100644 --- a/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts +++ b/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts @@ -3,17 +3,21 @@ import { logger } from "~/services/logger.server"; import { BaseService } from "./baseService.server"; export class CreateCheckpointRestoreEventService extends BaseService { - public async call(params: { checkpointId: string; type: CheckpointRestoreEventType; }): Promise { - const checkpoint = await this._prisma.checkpoint.findUniqueOrThrow({ + const checkpoint = await this._prisma.checkpoint.findUnique({ where: { id: params.checkpointId, }, }); + if (!checkpoint) { + logger.error("Checkpoint not found", { id: params.checkpointId }); + return; + } + logger.debug(`Creating checkpoint/restore event`, params); const checkpointEvent = await this._prisma.checkpointRestoreEvent.create({ diff --git a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts index cd8676b45..ea046bce1 100644 --- a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts @@ -1,26 +1,69 @@ -import { type Checkpoint } from "@trigger.dev/database"; -import { PrismaClient, prisma } from "~/db.server"; +import { TaskRunStatus, type Checkpoint, TaskRunAttemptStatus } from "@trigger.dev/database"; import { logger } from "~/services/logger.server"; import { socketIo } from "../handleSocketIo.server"; import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server"; +import { BaseService } from "./baseService.server"; -export class RestoreCheckpointService { - #prismaClient: PrismaClient; +const RESTORABLE_RUN_STATUSES: TaskRunStatus[] = ["WAITING_TO_RESUME"]; +const RESTORABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["PAUSED"]; - constructor(prismaClient: PrismaClient = prisma) { - this.#prismaClient = prismaClient; - } - - public async call(params: { checkpointId: string }): Promise { +export class RestoreCheckpointService extends BaseService { + public async call(params: { + eventId: string; + isRetry?: boolean; + }): Promise { logger.debug(`Restoring checkpoint`, params); - const checkpoint = await this.#prismaClient.checkpoint.findUniqueOrThrow({ + const checkpointEvent = await this._prisma.checkpointRestoreEvent.findUnique({ where: { - id: params.checkpointId, + id: params.eventId, + type: "CHECKPOINT", + }, + include: { + checkpoint: { + include: { + run: { + select: { + status: true, + }, + }, + attempt: { + select: { + status: true, + }, + }, + }, + }, }, }); - const eventService = new CreateCheckpointRestoreEventService(this.#prismaClient); + if (!checkpointEvent) { + logger.error("Checkpoint event not found", params); + return; + } + + const checkpoint = checkpointEvent.checkpoint; + + const runIsRestorable = RESTORABLE_RUN_STATUSES.includes(checkpoint.run.status); + const attemptIsRestorable = RESTORABLE_ATTEMPT_STATUSES.includes(checkpoint.attempt.status); + + if (!runIsRestorable) { + logger.error("Run is unrestorable", { + id: checkpoint.runId, + status: checkpoint.run.status, + }); + return; + } + + if (!attemptIsRestorable && !params.isRetry) { + logger.error("Attempt is unrestorable", { + id: checkpoint.attemptId, + status: checkpoint.attempt.status, + }); + return; + } + + const eventService = new CreateCheckpointRestoreEventService(this._prisma); await eventService.call({ checkpointId: checkpoint.id, type: "RESTORE" }); socketIo.providerNamespace.emit("RESTORE", { diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index caad5638a..26ed569f0 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -2,6 +2,7 @@ import { BackgroundWorkerProperties, Config, CreateBackgroundWorkerResponse, + InferSocketMessageSchema, ProdChildToWorkerMessages, ProdTaskRunExecutionPayload, ProdWorkerToChildMessages, @@ -55,9 +56,15 @@ export class ProdBackgroundWorker { public onTaskHeartbeat: Evt = new Evt(); - public onWaitForDuration: Evt<{ version?: "v1"; ms: number; now: number }> = new Evt(); - public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt(); - public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt(); + public onWaitForBatch: Evt< + InferSocketMessageSchema + > = new Evt(); + public onWaitForDuration: Evt< + InferSocketMessageSchema + > = new Evt(); + public onWaitForTask: Evt< + InferSocketMessageSchema + > = new Evt(); public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>(); public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>(); @@ -352,9 +359,15 @@ class TaskRunProcess { public onTaskHeartbeat: Evt = new Evt(); public onExit: Evt = new Evt(); - public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt(); - public onWaitForDuration: Evt<{ version?: "v1"; ms: number; now: number }> = new Evt(); - public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt(); + public onWaitForBatch: Evt< + InferSocketMessageSchema + > = new Evt(); + public onWaitForDuration: Evt< + InferSocketMessageSchema + > = new Evt(); + public onWaitForTask: Evt< + InferSocketMessageSchema + > = new Evt(); public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>(); public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>(); diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index cda670ec0..b5db3bb77 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -35,6 +35,7 @@ class ProdWorker { private executing = false; private completed = new Set(); private paused = false; + private attemptFriendlyId?: string; private nextResumeAfter: "WAIT_FOR_DURATION" | "WAIT_FOR_TASK" | "WAIT_FOR_BATCH" | undefined; @@ -77,10 +78,18 @@ class ProdWorker { }); this.#backgroundWorker.onWaitForDuration.attach(async (message) => { + if (!this.attemptFriendlyId) { + logger.error("Failed to send wait message, attempt friendly ID not set", { message }); + return; + } + // TODO: Switch to .send() once coordinator uses zod handler for all messages const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck( "WAIT_FOR_DURATION", - { version: "v1", ...message } + { + ...message, + attemptFriendlyId: this.attemptFriendlyId, + } ); logger.log("WAIT_FOR_DURATION", { willCheckpointAndRestore }); @@ -99,10 +108,18 @@ class ProdWorker { }); this.#backgroundWorker.onWaitForTask.attach(async (message) => { + if (!this.attemptFriendlyId) { + logger.error("Failed to send wait message, attempt friendly ID not set", { message }); + return; + } + // TODO: Switch to .send() once coordinator uses zod handler for all messages const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck( "WAIT_FOR_TASK", - { version: "v1", ...message } + { + ...message, + attemptFriendlyId: this.attemptFriendlyId, + } ); logger.log("WAIT_FOR_TASK", { willCheckpointAndRestore }); @@ -121,10 +138,18 @@ class ProdWorker { }); this.#backgroundWorker.onWaitForBatch.attach(async (message) => { + if (!this.attemptFriendlyId) { + logger.error("Failed to send wait message, attempt friendly ID not set", { message }); + return; + } + // TODO: Switch to .send() once coordinator uses zod handler for all messages const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck( "WAIT_FOR_BATCH", - { version: "v1", ...message } + { + ...message, + attemptFriendlyId: this.attemptFriendlyId, + } ); logger.log("WAIT_FOR_BATCH", { willCheckpointAndRestore }); @@ -208,12 +233,14 @@ class ProdWorker { } this.executing = true; + this.attemptFriendlyId = executionPayload.execution.attempt.id; const completion = await this.#backgroundWorker.executeTaskRun(executionPayload); logger.log("completed", completion); this.completed.add(executionPayload.execution.attempt.id); this.executing = false; + this.attemptFriendlyId = undefined; await this.#backgroundWorker.flushTelemetry(); @@ -378,20 +405,6 @@ class ProdWorker { case "/whoami": return reply.text(this.contentHash); - case "/wait": - const { willCheckpointAndRestore } = await this.#coordinatorSocket.sendWithAck( - "WAIT_FOR_DURATION", - { - version: "v1", - ms: 60_000, - now: Date.now(), - } - ); - logger.log("WAIT_FOR_DURATION", { willCheckpointAndRestore }); - // this is required when C/Ring established connections - this.#coordinatorSocket.close(); - return reply.text("sent WAIT"); - case "/connect": this.#coordinatorSocket.connect(); return reply.empty(); @@ -430,6 +443,7 @@ class ProdWorker { version: "v1", attemptId: this.attemptId, runId: this.runId, + totalCompletions: this.completed.size, }); return reply.empty(); diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index 85685200d..9ae08afbe 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -157,6 +157,7 @@ export const CoordinatorToPlatformMessages = { version: z.literal("v1").default("v1"), attemptId: z.string(), runId: z.string(), + totalCompletions: z.number(), }), callback: z.discriminatedUnion("success", [ z.object({ @@ -197,7 +198,7 @@ export const CoordinatorToPlatformMessages = { CHECKPOINT_CREATED: { message: z.object({ version: z.literal("v1").default("v1"), - attemptId: z.string(), + attemptFriendlyId: z.string(), docker: z.boolean(), location: z.string(), reason: z.discriminatedUnion("type", [ @@ -324,6 +325,7 @@ export const ProdWorkerToCoordinatorMessages = { version: z.literal("v1").default("v1"), attemptId: z.string(), runId: z.string(), + totalCompletions: z.number(), }), }, READY_FOR_RESUME: { @@ -365,6 +367,7 @@ export const ProdWorkerToCoordinatorMessages = { version: z.literal("v1").default("v1"), ms: z.number(), now: z.number(), + attemptFriendlyId: z.string(), }), callback: z.object({ willCheckpointAndRestore: z.boolean(), @@ -374,6 +377,7 @@ export const ProdWorkerToCoordinatorMessages = { message: z.object({ version: z.literal("v1").default("v1"), id: z.string(), + attemptFriendlyId: z.string(), }), callback: z.object({ willCheckpointAndRestore: z.boolean(), @@ -384,6 +388,7 @@ export const ProdWorkerToCoordinatorMessages = { version: z.literal("v1").default("v1"), id: z.string(), runs: z.string().array(), + attemptFriendlyId: z.string(), }), callback: z.object({ willCheckpointAndRestore: z.boolean(),