diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index 6fd842efa..7277d8c13 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -624,6 +624,44 @@ class TaskCoordinator { } }); + socket.on("READY_FOR_LAZY_ATTEMPT", async (message) => { + logger.log("[READY_FOR_LAZY_ATTEMPT]", message); + + try { + const lazyAttempt = await this.#platformSocket?.sendWithAck("READY_FOR_LAZY_ATTEMPT", { + ...message, + envId: socket.data.envId, + }); + + if (!lazyAttempt) { + logger.error("no lazy attempt ack", { runId: socket.data.runId }); + + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + + return; + } + + if (!lazyAttempt.success) { + logger.error("failed to get lazy attempt payload", { runId: socket.data.runId }); + + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + + return; + } + + socket.emit("EXECUTE_TASK_RUN_LAZY_ATTEMPT", { + version: "v1", + lazyPayload: lazyAttempt.lazyPayload, + }); + } catch (error) { + logger.error("Error", { error }); + } + }); + socket.on("READY_FOR_RESUME", async (message) => { logger.log("[READY_FOR_RESUME]", message); @@ -714,6 +752,19 @@ class TaskCoordinator { } }); + socket.on("TASK_RUN_FAILED_TO_RUN", async ({ completion }) => { + logger.log("completed task", { completionId: completion.id }); + + this.#platformSocket?.send("TASK_RUN_FAILED_TO_RUN", { + version: "v1", + completion, + }); + + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + }); + socket.on("READY_FOR_CHECKPOINT", async (message) => { logger.log("[READY_FOR_CHECKPOINT]", message); @@ -918,6 +969,28 @@ class TaskCoordinator { error: message.error, }); }); + + socket.on("CREATE_TASK_RUN_ATTEMPT", async (message, callback) => { + logger.log("[CREATE_TASK_RUN_ATTEMPT]", message); + + const createAttempt = await this.#platformSocket?.sendWithAck("CREATE_TASK_RUN_ATTEMPT", { + runId: message.runId, + envId: socket.data.envId, + }); + + if (!createAttempt?.success) { + logger.debug("no ack while creating attempt", message); + callback({ success: false }); + return; + } + + socket.data.attemptFriendlyId = createAttempt.executionPayload.execution.attempt.id; + + callback({ + success: true, + executionPayload: createAttempt.executionPayload, + }); + }); }, onDisconnect: async (socket, handler, sender, logger) => { this.#platformSocket?.send("LOG", { @@ -929,6 +1002,9 @@ class TaskCoordinator { TASK_HEARTBEAT: async (message) => { this.#platformSocket?.send("TASK_HEARTBEAT", message); }, + TASK_RUN_HEARTBEAT: async (message) => { + this.#platformSocket?.send("TASK_RUN_HEARTBEAT", message); + }, }, }); diff --git a/apps/webapp/app/v3/eventRepository.server.ts b/apps/webapp/app/v3/eventRepository.server.ts index bfac9b700..8c94f44e2 100644 --- a/apps/webapp/app/v3/eventRepository.server.ts +++ b/apps/webapp/app/v3/eventRepository.server.ts @@ -10,6 +10,7 @@ import { SpanEvents, SpanMessagingEvent, TaskEventStyle, + TaskRunError, correctErrorStackTrace, createPacketAttributesAsJson, flattenAttributes, @@ -864,6 +865,36 @@ export function stripAttributePrefix(attributes: Attributes, prefix: string) { return result; } +export function createExceptionPropertiesFromError(error: TaskRunError): ExceptionEventProperties { + switch (error.type) { + case "BUILT_IN_ERROR": { + return { + type: error.name, + message: error.message, + stacktrace: error.stackTrace, + }; + } + case "CUSTOM_ERROR": { + return { + type: "Error", + message: error.raw, + }; + } + case "INTERNAL_ERROR": { + return { + type: "Internal error", + message: [error.code, error.message].filter(Boolean).join(": "), + }; + } + case "STRING_ERROR": { + return { + type: "Error", + message: error.raw, + }; + } + } +} + /** * Filters out partial events from a batch of creatable events, excluding those that have a corresponding full event. * @param batch - The batch of creatable events to filter. diff --git a/apps/webapp/app/v3/failedTaskRun.server.ts b/apps/webapp/app/v3/failedTaskRun.server.ts index b9cffeaf6..79594e73c 100644 --- a/apps/webapp/app/v3/failedTaskRun.server.ts +++ b/apps/webapp/app/v3/failedTaskRun.server.ts @@ -1,13 +1,9 @@ -import { - ExceptionEventProperties, - TaskRunError, - TaskRunFailedExecutionResult, -} from "@trigger.dev/core/v3"; +import { TaskRunFailedExecutionResult } from "@trigger.dev/core/v3"; import { logger } from "~/services/logger.server"; import { marqs } from "~/v3/marqs/index.server"; import { TaskRunStatus } from "@trigger.dev/database"; -import { eventRepository } from "./eventRepository.server"; +import { createExceptionPropertiesFromError, eventRepository } from "./eventRepository.server"; import { BaseService } from "./services/baseService.server"; const FAILABLE_TASK_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "PENDING", "WAITING_FOR_DEPLOY"]; @@ -68,33 +64,3 @@ export class FailedTaskRunService extends BaseService { }); } } - -function createExceptionPropertiesFromError(error: TaskRunError): ExceptionEventProperties { - switch (error.type) { - case "BUILT_IN_ERROR": { - return { - type: error.name, - message: error.message, - stacktrace: error.stackTrace, - }; - } - case "CUSTOM_ERROR": { - return { - type: "Error", - message: error.raw, - }; - } - case "INTERNAL_ERROR": { - return { - type: "Internal error", - message: [error.code, error.message].filter(Boolean).join(": "), - }; - } - case "STRING_ERROR": { - return { - type: "Error", - message: error.raw, - }; - } - } -} diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts index 81e8c4683..2b917e45a 100644 --- a/apps/webapp/app/v3/handleSocketIo.server.ts +++ b/apps/webapp/app/v3/handleSocketIo.server.ts @@ -22,6 +22,7 @@ import { DeploymentIndexFailed } from "./services/deploymentIndexFailed.server"; import { Redis } from "ioredis"; import { createAdapter } from "@socket.io/redis-adapter"; import { CrashTaskRunService } from "./services/crashTaskRun.server"; +import { CreateTaskRunAttemptService } from "./services/createTaskRunAttempt.server"; export const socketIo = singleton("socketIo", initalizeIoServer); @@ -91,6 +92,23 @@ function createCoordinatorNamespace(io: Server) { return { success: true, payload }; } }, + READY_FOR_LAZY_ATTEMPT: async (message) => { + try { + const payload = await sharedQueueTasks.getLazyAttemptPayload( + message.envId, + message.runId + ); + + if (!payload) { + logger.error("Failed to retrieve lazy attempt payload", message); + return { success: false, reason: "Failed to retrieve payload" }; + } + + return { success: true, lazyPayload: payload }; + } catch (error) { + return { success: false }; + } + }, READY_FOR_RESUME: async (message) => { const resumeAttempt = new ResumeAttemptService(); await resumeAttempt.call(message); @@ -103,6 +121,9 @@ function createCoordinatorNamespace(io: Server) { checkpoint: message.checkpoint, }); }, + TASK_RUN_FAILED_TO_RUN: async (message) => { + await sharedQueueTasks.taskRunFailed(message.completion); + }, TASK_HEARTBEAT: async (message) => { await sharedQueueTasks.taskHeartbeat(message.attemptFriendlyId); }, @@ -135,6 +156,33 @@ function createCoordinatorNamespace(io: Server) { return { success: false }; } }, + CREATE_TASK_RUN_ATTEMPT: async (message) => { + try { + const environment = await findEnvironmentById(message.envId); + + if (!environment) { + logger.error("Environment not found", { id: message.envId }); + return { success: false, reason: "Environment not found" }; + } + + const service = new CreateTaskRunAttemptService(); + const { attempt } = await service.call(message.runId, environment, false); + + const payload = await sharedQueueTasks.getExecutionPayloadFromAttempt(attempt.id, true); + + if (!payload) { + logger.error("Failed to retrieve payload after attempt creation", { + id: message.envId, + }); + return { success: false, reason: "Failed to retrieve payload" }; + } + + return { success: true, executionPayload: payload }; + } catch (error) { + logger.error("Error while creating attempt", { error }); + return { success: false }; + } + }, INDEXING_FAILED: async (message) => { try { const service = new DeploymentIndexFailed(); diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index f88e7a3bc..4e75c2047 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -1,9 +1,11 @@ import { Context, ROOT_CONTEXT, Span, SpanKind, context, trace } from "@opentelemetry/api"; import { + Machine, ProdTaskRunExecution, ProdTaskRunExecutionPayload, TaskRunError, TaskRunExecution, + TaskRunExecutionLazyAttemptPayload, TaskRunExecutionResult, TaskRunFailedExecutionResult, TaskRunSuccessfulExecutionResult, @@ -30,6 +32,7 @@ import { tracer } from "../tracer.server"; import { CrashTaskRunService } from "../services/crashTaskRun.server"; import { FailedTaskRunService } from "../failedTaskRun.server"; import { CreateTaskRunAttemptService } from "../services/createTaskRunAttempt.server"; +import { findEnvironmentById } from "~/models/runtimeEnvironment.server"; const WithTraceContext = z.object({ traceparent: z.string().optional(), @@ -393,6 +396,7 @@ export class SharedQueueConsumer { createdAt: "desc", }, }, + lockedBy: true, }, }); @@ -442,34 +446,29 @@ export class SharedQueueConsumer { await this.#ackAndDoMoreWork(message.messageId); return; } - } else if (isRetry) { + + break; + } + + if (!deployment.worker.supportsLazyAttempts) { + const service = new CreateTaskRunAttemptService(); + await service.call(lockedTaskRun.friendlyId, undefined, false); + } + + if (isRetry) { socketIo.coordinatorNamespace.emit("READY_FOR_RETRY", { version: "v1", runId: lockedTaskRun.id, }); } else { - const environment = await prisma.runtimeEnvironment.findUniqueOrThrow({ - where: { - id: lockedTaskRun.runtimeEnvironmentId, - }, - include: { - project: true, - organization: true, - }, - }); + const machineConfig = lockedTaskRun.lockedBy?.machineConfig; + const machine = Machine.safeParse(machineConfig ?? {}); - const service = new CreateTaskRunAttemptService(); - const { attempt, machine } = await service.call( - lockedTaskRun.friendlyId, - environment, - false - ); - - if (!machine) { - logger.error("Missing machine config", { + if (!machine.success) { + logger.error("Failed to parse machine config", { queueMessage: message.data, messageId: message.messageId, - attemptId: attempt.id, + machineConfig, }); await this.#ackAndDoMoreWork(message.messageId); @@ -482,9 +481,9 @@ export class SharedQueueConsumer { type: "SCHEDULE_ATTEMPT", image: deployment.imageReference, version: deployment.version, - machine: machine, + machine: machine.data, // identifiers - id: attempt.id, + id: "placeholder", // TODO: Remove this completely in a future release envId: lockedTaskRun.runtimeEnvironment.id, envType: lockedTaskRun.runtimeEnvironment.type, orgId: lockedTaskRun.runtimeEnvironment.organizationId, @@ -1063,6 +1062,47 @@ class SharedQueueTasks { return this.getExecutionPayloadFromAttempt(latestAttempt.id, setToExecuting, isRetrying); } + async getLazyAttemptPayload( + envId: string, + runId: string + ): Promise { + const environment = await findEnvironmentById(envId); + + if (!environment) { + logger.error("Environment not found", { id: envId }); + return; + } + + const run = await prisma.taskRun.findUnique({ + where: { + id: runId, + runtimeEnvironmentId: environment.id, + }, + }); + + if (!run) { + logger.error("Run not found", { id: runId, envId }); + return; + } + + const environmentRepository = new EnvironmentVariablesRepository(); + const variables = await environmentRepository.getEnvironmentVariables( + environment.projectId, + environment.id + ); + + return { + traceContext: run.traceContext as Record, + environment: variables.reduce((acc: Record, curr) => { + acc[curr.key] = curr.value; + return acc; + }, {}), + runId: run.friendlyId, + messageId: run.id, + isTest: run.isTest, + } satisfies TaskRunExecutionLazyAttemptPayload; + } + async taskHeartbeat(attemptFriendlyId: string, seconds: number = 60) { logger.debug("[SharedQueueConsumer] taskHeartbeat()", { id: attemptFriendlyId, seconds }); diff --git a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts index cdb2b8a9c..90ad49273 100644 --- a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts +++ b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts @@ -1,5 +1,5 @@ -import { Machine, TaskRunExecution } from "@trigger.dev/core/v3"; -import { $transaction } from "~/db.server"; +import { TaskRunExecution } from "@trigger.dev/core/v3"; +import { $transaction, PrismaClientOrTransaction, prisma } from "~/db.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; @@ -8,26 +8,47 @@ import { TaskRun, TaskRunAttempt } from "@trigger.dev/database"; export class CreateTaskRunAttemptService extends BaseService { public async call( - runFriendlyId: string, - environment: AuthenticatedEnvironment, + runId: string, + env?: AuthenticatedEnvironment, setToExecuting = true ): Promise<{ execution: TaskRunExecution; run: TaskRun; attempt: TaskRunAttempt; - machine?: Machine; }> { + let environment: AuthenticatedEnvironment | undefined = env; + + if (!environment) { + environment = await getAuthenticatedEnvironmentFromRun(runId, this._prisma); + + if (!environment) { + throw new ServiceValidationError("Environment not found", 404); + } + } + + const isFriendlyId = runId.startsWith("run_"); + return await this.traceWithEnv("call()", environment, async (span) => { - span.setAttribute("taskRunId", runFriendlyId); + if (isFriendlyId) { + span.setAttribute("taskRunFriendlyId", runId); + } else { + span.setAttribute("taskRunId", runId); + } const taskRun = await this._prisma.taskRun.findUnique({ where: { - friendlyId: runFriendlyId, + id: !isFriendlyId ? runId : undefined, + friendlyId: isFriendlyId ? runId : undefined, runtimeEnvironmentId: environment.id, }, include: { tags: true, - attempts: true, + attempts: { + take: 1, + orderBy: { + number: "desc", + }, + }, lockedBy: { include: { worker: true, @@ -47,6 +68,9 @@ export class CreateTaskRunAttemptService extends BaseService { throw new ServiceValidationError("Task run not found", 404); } + span.setAttribute("taskRunId", taskRun.id); + span.setAttribute("taskRunFriendlyId", taskRun.friendlyId); + if (taskRun.status === "CANCELED") { throw new ServiceValidationError("Task run is cancelled", 400); } @@ -68,10 +92,12 @@ export class CreateTaskRunAttemptService extends BaseService { throw new ServiceValidationError("Queue not found", 404); } + const nextAttemptNumber = taskRun.attempts[0] ? taskRun.attempts[0].number + 1 : 1; + const taskRunAttempt = await $transaction(this._prisma, async (tx) => { const taskRunAttempt = await tx.taskRunAttempt.create({ data: { - number: taskRun.attempts[0] ? taskRun.attempts[0].number + 1 : 1, + number: nextAttemptNumber, friendlyId: generateFriendlyId("attempt"), taskRunId: taskRun.id, startedAt: new Date(), @@ -82,6 +108,7 @@ export class CreateTaskRunAttemptService extends BaseService { runtimeEnvironmentId: environment.id, }, include: { + backgroundWorker: true, backgroundWorkerTask: true, }, }); @@ -101,6 +128,7 @@ export class CreateTaskRunAttemptService extends BaseService { }); if (!taskRunAttempt) { + logger.error("Failed to create task run attempt", { runId: taskRun.id, nextAttemptNumber }); throw new ServiceValidationError("Failed to create task run attempt", 500); } @@ -154,24 +182,36 @@ export class CreateTaskRunAttemptService extends BaseService { : undefined, }; - const { machineConfig } = taskRunAttempt.backgroundWorkerTask; - const machine = Machine.safeParse(machineConfig ?? {}); - - if (!machine.success) { - logger.error("Failed to parse machine config", { - run: taskRun.id, - attempt: taskRunAttempt.id, - backgroundWorkerTask: taskRunAttempt.backgroundWorkerTask.id, - machineConfig, - }); - } - return { execution, run: taskRun, attempt: taskRunAttempt, - machine: machine.success ? machine.data : undefined, }; }); } } + +async function getAuthenticatedEnvironmentFromRun( + friendlyId: string, + prismaClient?: PrismaClientOrTransaction +) { + const taskRun = await (prismaClient ?? prisma).taskRun.findUnique({ + where: { + friendlyId, + }, + include: { + runtimeEnvironment: { + include: { + organization: true, + project: true, + }, + }, + }, + }); + + if (!taskRun) { + return; + } + + return taskRun?.runtimeEnvironment; +} diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index 8f2bbe897..7766b9801 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -11,6 +11,7 @@ import { TaskRunBuiltInError, TaskRunErrorCodes, TaskRunExecution, + TaskRunExecutionLazyAttemptPayload, TaskRunExecutionPayload, TaskRunExecutionResult, WaitReason, @@ -56,7 +57,11 @@ type BackgroundWorkerParams = { export class ProdBackgroundWorker { private _initialized: boolean = false; + /** + * @deprecated use onTaskRunHeartbeat instead + */ public onTaskHeartbeat: Evt = new Evt(); + public onTaskRunHeartbeat: Evt = new Evt(); public onWaitForBatch: Evt< InferSocketMessageSchema @@ -74,6 +79,18 @@ export class ProdBackgroundWorker { public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>(); public onCancelCheckpoint = Evt.create<{ version?: "v1" | "v2"; reason?: WaitReason }>(); + public onCreateTaskRunAttempt = Evt.create<{ version?: "v1"; runId: string }>(); + public attemptCreatedNotification = Evt.create< + | { + success: false; + reason?: string; + } + | { + success: true; + execution: ProdTaskRunExecution; + } + >(); + private _onClose: Evt = new Evt(); public tasks: Array = []; @@ -95,6 +112,7 @@ export class ProdBackgroundWorker { this._closed = true; this.onTaskHeartbeat.detach(); + this.onTaskRunHeartbeat.detach(); // We need to close the task run process await this._taskRunProcess?.cleanup(true); @@ -204,7 +222,10 @@ export class ProdBackgroundWorker { this._taskRunProcess?.waitCompletedNotification(); } - async #initializeTaskRunProcess(payload: ProdTaskRunExecutionPayload): Promise { + async #initializeTaskRunProcess( + payload: ProdTaskRunExecutionPayload, + messageId?: string + ): Promise { const metadata = this.getMetadata( payload.execution.worker.id, payload.execution.worker.version @@ -219,7 +240,8 @@ export class ProdBackgroundWorker { ...(payload.environment ?? {}), }, metadata, - this.params + this.params, + messageId ); taskRunProcess.onExit.attach(() => { @@ -230,6 +252,10 @@ export class ProdBackgroundWorker { this.onTaskHeartbeat.post(id); }); + taskRunProcess.onTaskRunHeartbeat.attach((id) => { + this.onTaskRunHeartbeat.post(id); + }); + taskRunProcess.onWaitForBatch.attach((message) => { this.onWaitForBatch.post(message); }); @@ -267,9 +293,12 @@ export class ProdBackgroundWorker { } // We need to fork the process before we can execute any tasks - async executeTaskRun(payload: ProdTaskRunExecutionPayload): Promise { + async executeTaskRun( + payload: ProdTaskRunExecutionPayload, + messageId?: string + ): Promise { try { - const taskRunProcess = await this.#initializeTaskRunProcess(payload); + const taskRunProcess = await this.#initializeTaskRunProcess(payload, messageId); const result = await taskRunProcess.executeTaskRun(payload); @@ -342,6 +371,40 @@ export class ProdBackgroundWorker { await this._taskRunProcess?.cancel(); } + async executeTaskRunLazyAttempt(payload: TaskRunExecutionLazyAttemptPayload) { + // Post to coordinator + this.onCreateTaskRunAttempt.post({ runId: payload.runId }); + + let execution: ProdTaskRunExecution; + + try { + // ..and wait for response + const attemptCreated = await this.attemptCreatedNotification.waitFor(30_000); + + if (!attemptCreated.success) { + throw new Error( + `Failed to create attempt${attemptCreated.reason ? `: ${attemptCreated.reason}` : ""}` + ); + } + + execution = attemptCreated.execution; + } catch (error) { + console.error("Error while creating attempt", error); + throw new Error(`Failed to create task run attempt: ${error}`); + } + + const completion = await this.executeTaskRun( + { + execution, + traceContext: payload.traceContext, + environment: payload.environment, + }, + payload.messageId + ); + + return { execution, completion }; + } + async #correctError( error: TaskRunBuiltInError, execution: TaskRunExecution @@ -369,7 +432,11 @@ class TaskRunProcess { private _isBeingKilled: boolean = false; private _isBeingCancelled: boolean = false; + /** + * @deprecated use onTaskRunHeartbeat instead + */ public onTaskHeartbeat: Evt = new Evt(); + public onTaskRunHeartbeat: Evt = new Evt(); public onExit: Evt = new Evt(); public onWaitForBatch: Evt< @@ -393,7 +460,8 @@ class TaskRunProcess { private path: string, private env: NodeJS.ProcessEnv, private metadata: BackgroundWorkerProperties, - private worker: BackgroundWorkerParams + private worker: BackgroundWorkerParams, + private messageId?: string ) {} async initialize() { @@ -439,7 +507,11 @@ class TaskRunProcess { process.exit(0); }, TASK_HEARTBEAT: async (message) => { - this.onTaskHeartbeat.post(message.id); + if (this.messageId) { + this.onTaskRunHeartbeat.post(this.messageId); + } else { + this.onTaskHeartbeat.post(message.id); + } }, TASKS_READY: async (message) => {}, WAIT_FOR_TASK: async (message) => { diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index 641414c72..f6fe5a0f3 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -5,6 +5,7 @@ import { PreStopCauses, ProdWorkerToCoordinatorMessages, TaskResource, + TaskRunFailedExecutionResult, WaitReason, } from "@trigger.dev/core/v3"; import { ZodSocketConnection } from "@trigger.dev/core/v3/zodSocket"; @@ -78,6 +79,10 @@ class ProdWorker { this.#coordinatorSocket.socket.emit("TASK_HEARTBEAT", { version: "v1", attemptFriendlyId }); }); + this.#backgroundWorker.onTaskRunHeartbeat.attach((runId) => { + this.#coordinatorSocket.socket.emit("TASK_RUN_HEARTBEAT", { version: "v1", runId }); + }); + this.#backgroundWorker.onReadyForCheckpoint.attach(async (message) => { // Flush before checkpointing so we don't flush the same spans again after restore await this.#backgroundWorker.flushTelemetry(); @@ -108,6 +113,40 @@ class ProdWorker { this.#backgroundWorker.checkpointCanceledNotification.post({ checkpointCanceled }); }); + this.#backgroundWorker.onCreateTaskRunAttempt.attach(async (message) => { + logger.log("onCreateTaskRunAttempt()", { message }); + + const createAttempt = await this.#coordinatorSocket.socket.emitWithAck( + "CREATE_TASK_RUN_ATTEMPT", + { + version: "v1", + runId: message.runId, + } + ); + + if (!createAttempt.success) { + this.#backgroundWorker.attemptCreatedNotification.post({ + success: false, + reason: createAttempt.reason, + }); + return; + } + + this.#backgroundWorker.attemptCreatedNotification.post({ + success: true, + execution: createAttempt.executionPayload.execution, + }); + }); + + this.#backgroundWorker.attemptCreatedNotification.attach((message) => { + if (!message.success) { + return; + } + + // Workers with lazy attempt support set their friendly ID here + this.attemptFriendlyId = message.execution.attempt.id; + }); + this.#backgroundWorker.onWaitForDuration.attach(async (message) => { if (!this.attemptFriendlyId) { logger.error("Failed to send wait message, attempt friendly ID not set", { message }); @@ -420,6 +459,59 @@ class ProdWorker { this.#prepareForRetry(willCheckpointAndRestore, shouldExit); }, + EXECUTE_TASK_RUN_LAZY_ATTEMPT: async (message) => { + if (this.executing) { + logger.error("dropping execute request, already executing"); + return; + } + + this.executing = true; + + try { + const { completion, execution } = + await this.#backgroundWorker.executeTaskRunLazyAttempt(message.lazyPayload); + + logger.log("completed", completion); + + this.completed.add(execution.attempt.id); + + const { willCheckpointAndRestore, shouldExit } = + await this.#coordinatorSocket.socket.emitWithAck("TASK_RUN_COMPLETED", { + version: "v1", + execution, + completion, + }); + + logger.log("completion acknowledged", { willCheckpointAndRestore, shouldExit }); + + this.#prepareForRetry(willCheckpointAndRestore, shouldExit); + } catch (error) { + const completion: TaskRunFailedExecutionResult = { + ok: false, + id: message.lazyPayload.runId, + retry: undefined, + error: + error instanceof Error + ? { + type: "BUILT_IN_ERROR", + name: error.name, + message: error.message, + stackTrace: error.stack ?? "", + } + : { + type: "BUILT_IN_ERROR", + name: "UnknownError", + message: String(error), + stackTrace: "", + }, + }; + + this.#coordinatorSocket.socket.emit("TASK_RUN_FAILED_TO_RUN", { + version: "v1", + completion, + }); + } + }, REQUEST_ATTEMPT_CANCELLATION: async (message) => { if (!this.executing) { return; @@ -436,7 +528,7 @@ class ProdWorker { return; } - this.#coordinatorSocket.socket.emit("READY_FOR_EXECUTION", { + this.#coordinatorSocket.socket.emit("READY_FOR_LAZY_ATTEMPT", { version: "v1", runId: this.runId, totalCompletions: this.completed.size, @@ -564,7 +656,7 @@ class ProdWorker { return; } - socket.emit("READY_FOR_EXECUTION", { + socket.emit("READY_FOR_LAZY_ATTEMPT", { version: "v1", runId: this.runId, totalCompletions: this.completed.size, diff --git a/packages/core-apps/src/provider.ts b/packages/core-apps/src/provider.ts index cb1b0c4d5..159c81dbc 100644 --- a/packages/core-apps/src/provider.ts +++ b/packages/core-apps/src/provider.ts @@ -46,7 +46,6 @@ export interface TaskOperationsCreateOptions { orgId: string; projectId: string; runId: string; - attemptId: string; } export interface TaskOperationsRestoreOptions { @@ -129,7 +128,6 @@ export class ProviderShell implements Provider { orgId: message.data.orgId, projectId: message.data.projectId, runId: message.data.runId, - attemptId: message.data.id, }); } catch (error) { logger.error("create failed", error); diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index b81425487..2c31e047f 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -29,7 +29,7 @@ export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [ version: z.string(), machine: Machine, // identifiers - id: z.string(), // attempt + id: z.string().optional(), // TODO: Remove this completely in a future release envId: z.string(), envType: EnvironmentType, orgId: z.string(), @@ -451,6 +451,23 @@ export const CoordinatorToPlatformMessages = { }), ]), }, + CREATE_TASK_RUN_ATTEMPT: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + envId: z.string(), + }), + callback: z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + reason: z.string().optional(), + }), + z.object({ + success: z.literal(true), + executionPayload: ProdTaskRunExecutionPayload, + }), + ]), + }, READY_FOR_EXECUTION: { message: z.object({ version: z.literal("v1").default("v1"), @@ -467,6 +484,24 @@ export const CoordinatorToPlatformMessages = { }), ]), }, + READY_FOR_LAZY_ATTEMPT: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + envId: z.string(), + totalCompletions: z.number(), + }), + callback: z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + reason: z.string().optional(), + }), + z.object({ + success: z.literal(true), + lazyPayload: TaskRunExecutionLazyAttemptPayload, + }), + ]), + }, READY_FOR_RESUME: { message: z.object({ version: z.literal("v1").default("v1"), @@ -487,6 +522,12 @@ export const CoordinatorToPlatformMessages = { .optional(), }), }, + TASK_RUN_FAILED_TO_RUN: { + message: z.object({ + version: z.literal("v1").default("v1"), + completion: TaskRunFailedExecutionResult, + }), + }, TASK_HEARTBEAT: { message: z.object({ version: z.literal("v1").default("v1"), @@ -652,6 +693,13 @@ export const ProdWorkerToCoordinatorMessages = { totalCompletions: z.number(), }), }, + READY_FOR_LAZY_ATTEMPT: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + totalCompletions: z.number(), + }), + }, READY_FOR_RESUME: { message: z.object({ version: z.literal("v1").default("v1"), @@ -688,6 +736,12 @@ export const ProdWorkerToCoordinatorMessages = { attemptFriendlyId: z.string(), }), }, + TASK_RUN_HEARTBEAT: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + }), + }, TASK_RUN_COMPLETED: { message: z.object({ version: z.literal("v1").default("v1"), @@ -699,6 +753,12 @@ export const ProdWorkerToCoordinatorMessages = { shouldExit: z.boolean(), }), }, + TASK_RUN_FAILED_TO_RUN: { + message: z.object({ + version: z.literal("v1").default("v1"), + completion: TaskRunFailedExecutionResult, + }), + }, WAIT_FOR_DURATION: { message: z.object({ version: z.literal("v1").default("v1"), @@ -744,6 +804,22 @@ export const ProdWorkerToCoordinatorMessages = { }), }), }, + CREATE_TASK_RUN_ATTEMPT: { + message: z.object({ + version: z.literal("v1").default("v1"), + runId: z.string(), + }), + callback: z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + reason: z.string().optional(), + }), + z.object({ + success: z.literal(true), + executionPayload: ProdTaskRunExecutionPayload, + }), + ]), + }, }; export const CoordinatorToProdWorkerMessages = { @@ -767,6 +843,12 @@ export const CoordinatorToProdWorkerMessages = { executionPayload: ProdTaskRunExecutionPayload, }), }, + EXECUTE_TASK_RUN_LAZY_ATTEMPT: { + message: z.object({ + version: z.literal("v1").default("v1"), + lazyPayload: TaskRunExecutionLazyAttemptPayload, + }), + }, REQUEST_ATTEMPT_CANCELLATION: { message: z.object({ version: z.literal("v1").default("v1"),