diff --git a/apps/webapp/app/routes/api.v1.runs.$runParam.attempts.ts b/apps/webapp/app/routes/api.v1.runs.$runParam.attempts.ts new file mode 100644 index 000000000..ecc59815b --- /dev/null +++ b/apps/webapp/app/routes/api.v1.runs.$runParam.attempts.ts @@ -0,0 +1,45 @@ +import type { ActionFunctionArgs } from "@remix-run/server-runtime"; +import { json } from "@remix-run/server-runtime"; +import { z } from "zod"; +import { authenticateApiRequest } from "~/services/apiAuth.server"; +import { ServiceValidationError } from "~/v3/services/baseService.server"; +import { CreateTaskRunAttemptService } from "~/v3/services/createTaskRunAttempt.server"; + +const ParamsSchema = z.object({ + /* This is the run friendly ID */ + runParam: z.string(), +}); + +export async function action({ request, params }: ActionFunctionArgs) { + // Authenticate the request + const authenticationResult = await authenticateApiRequest(request); + + if (!authenticationResult) { + return json({ error: "Invalid or Missing API Key" }, { status: 401 }); + } + + const parsed = ParamsSchema.safeParse(params); + + if (!parsed.success) { + return json({ error: "Invalid or missing run ID" }, { status: 400 }); + } + + const { runParam } = parsed.data; + + const service = new CreateTaskRunAttemptService(); + + try { + const execution = await service.call(runParam, authenticationResult.environment); + + return json(execution, { status: 200 }); + } catch (error) { + if (error instanceof ServiceValidationError) { + return json({ error: error.message }, { status: error.status ?? 422 }); + } + + return json( + { error: error instanceof Error ? error.message : "Internal Server Error" }, + { status: 500 } + ); + } +} diff --git a/apps/webapp/app/v3/authenticatedSocketConnection.server.ts b/apps/webapp/app/v3/authenticatedSocketConnection.server.ts index 209954f34..33b903819 100644 --- a/apps/webapp/app/v3/authenticatedSocketConnection.server.ts +++ b/apps/webapp/app/v3/authenticatedSocketConnection.server.ts @@ -54,7 +54,10 @@ export class AuthenticatedSocketConnection { schema: clientWebsocketMessages, messages: { READY_FOR_TASKS: async (payload) => { - await this._consumer.registerBackgroundWorker(payload.backgroundWorkerId); + await this._consumer.registerBackgroundWorker( + payload.backgroundWorkerId, + payload.inProgressRuns ?? [] + ); }, BACKGROUND_WORKER_DEPRECATED: async (payload) => { await this._consumer.deprecateBackgroundWorker(payload.backgroundWorkerId); @@ -73,6 +76,10 @@ export class AuthenticatedSocketConnection { await this._consumer.taskHeartbeat(payload.backgroundWorkerId, payload.data.id); break; } + case "TASK_RUN_HEARTBEAT": { + await this._consumer.taskRunHeartbeat(payload.backgroundWorkerId, payload.data.id); + break; + } } }, }, diff --git a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts index d13683851..98b6a298c 100644 --- a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts @@ -1,6 +1,7 @@ import { Context, ROOT_CONTEXT, Span, SpanKind, context, trace } from "@opentelemetry/api"; import { TaskRunExecution, + TaskRunExecutionLazyAttemptPayload, TaskRunExecutionPayload, TaskRunExecutionResult, serverWebsocketMessages, @@ -14,10 +15,9 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { marqs } from "~/v3/marqs/index.server"; import { EnvironmentVariablesRepository } from "../environmentVariables/environmentVariablesRepository.server"; -import { generateFriendlyId } from "../friendlyIdentifiers"; -import { CancelAttemptService } from "../services/cancelAttempt.server"; import { CancelTaskRunService } from "../services/cancelTaskRun.server"; import { CompleteAttemptService } from "../services/completeAttempt.server"; +import { CreateTaskRunAttemptService } from "../services/createTaskRunAttempt.server"; import { SEMINTATTRS_FORCE_RECORDING, attributesFromAuthenticatedEnv, @@ -54,7 +54,6 @@ export class DevQueueConsumer { private _taskSuccesses: number = 0; private _currentSpan: Span | undefined; private _endSpanInNextIteration = false; - private _inProgressAttempts: Map = new Map(); // Keys are task attempt friendly IDs, values are TaskRun ids/queue message ids private _inProgressRuns: Map = new Map(); // Keys are task run friendly IDs, values are TaskRun internal ids/queue message ids constructor( @@ -78,7 +77,7 @@ export class DevQueueConsumer { this._backgroundWorkers.delete(id); } - public async registerBackgroundWorker(id: string) { + public async registerBackgroundWorker(id: string, inProgressRuns: string[] = []) { const backgroundWorker = await prisma.backgroundWorker.findUnique({ where: { friendlyId: id, runtimeEnvironmentId: this.env.id }, include: { @@ -92,7 +91,10 @@ export class DevQueueConsumer { this._backgroundWorkers.set(backgroundWorker.id, backgroundWorker); - logger.debug("Registered background worker", { backgroundWorker: backgroundWorker.id }); + logger.debug("Registered background worker", { + backgroundWorker: backgroundWorker.id, + inProgressRuns, + }); const subscriber = await devPubSub.subscribe(`backgroundWorker:${backgroundWorker.id}:*`); @@ -109,6 +111,10 @@ export class DevQueueConsumer { this._backgroundWorkerSubscriber.set(backgroundWorker.id, subscriber); + for (const runId of inProgressRuns) { + this._inProgressRuns.set(runId, runId); + } + // Start reading from the queue if we haven't already await this.#enable(); } @@ -118,15 +124,16 @@ export class DevQueueConsumer { completion: TaskRunExecutionResult, execution: TaskRunExecution ) { - this._inProgressAttempts.delete(execution.attempt.id); - if (completion.ok) { this._taskSuccesses++; } else { this._taskFailures++; } - logger.debug("Task run completed", { taskRunCompletion: completion, execution }); + logger.debug("[DevQueueConsumer] taskAttemptCompleted()", { + taskRunCompletion: completion, + execution, + }); const service = new CompleteAttemptService(); const result = await service.call({ completion, execution, env: this.env }); @@ -136,7 +143,12 @@ export class DevQueueConsumer { } } + /** + * @deprecated Use `taskRunHeartbeat` instead + */ public async taskHeartbeat(workerId: string, id: string, seconds: number = 60) { + logger.debug("[DevQueueConsumer] taskHeartbeat()", { id, seconds }); + const taskRunAttempt = await prisma.taskRunAttempt.findUnique({ where: { friendlyId: id }, }); @@ -148,6 +160,12 @@ export class DevQueueConsumer { await marqs?.heartbeatMessage(taskRunAttempt.taskRunId, seconds); } + public async taskRunHeartbeat(workerId: string, id: string, seconds: number = 60) { + logger.debug("[DevQueueConsumer] taskRunHeartbeat()", { id, seconds }); + + await marqs?.heartbeatMessage(id, seconds); + } + public async stop(reason: string = "CLI disconnected") { if (!this._enabled) { return; @@ -180,66 +198,23 @@ export class DevQueueConsumer { } 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(); - 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, - cancelAttemptService, - cancelledAt, - reason - ); - } - - for (const runId of inProgressRunsWithNoInProgressAttempts) { + for (const [_, runId] of inProgressRuns) { await this.#cancelInProgressRun(runId, cancelTaskRunService, cancelledAt, reason); } } - async #cancelInProgressAttempt( - attemptId: string, - messageId: string, - cancelAttemptService: CancelAttemptService, - cancelledAt: Date, - reason: string - ) { - logger.debug("Cancelling in progress attempt", { attemptId, messageId }); - - try { - await cancelAttemptService.call(attemptId, messageId, cancelledAt, reason, this.env); - } catch (e) { - logger.error("Failed to cancel in progress attempt", { - attemptId, - messageId, - error: e, - }); - } - } - async #cancelInProgressRun( runId: string, service: CancelTaskRunService, @@ -248,16 +223,20 @@ export class DevQueueConsumer { ) { logger.debug("Cancelling in progress run", { runId }); - const taskRun = await prisma.taskRun.findUnique({ - where: { id: runId }, - }); + const taskRun = runId.startsWith("run_") + ? await prisma.taskRun.findUnique({ + where: { friendlyId: runId }, + }) + : await prisma.taskRun.findUnique({ + where: { id: runId }, + }); if (!taskRun) { return; } try { - await service.call(taskRun, { reason, cancelAttempts: false, cancelledAt }); + await service.call(taskRun, { reason, cancelAttempts: true, cancelledAt }); } catch (e) { logger.error("Failed to cancel in progress run", { runId, @@ -446,154 +425,132 @@ export class DevQueueConsumer { return; } - const queue = await prisma.taskQueue.findUnique({ - where: { - runtimeEnvironmentId_name: { runtimeEnvironmentId: this.env.id, name: lockedTaskRun.queue }, - }, - }); - - if (!queue) { - await marqs?.nackMessage(message.messageId); - setTimeout(() => this.#doWork(), 1000); - return; - } - if (!this._enabled) { + logger.debug("Dev queue consumer is disabled", { env: this.env, queueMessage: message }); + await marqs?.nackMessage(message.messageId); return; } - const taskRunAttempt = await prisma.taskRunAttempt.create({ - data: { - number: lockedTaskRun.attempts[0] ? lockedTaskRun.attempts[0].number + 1 : 1, - friendlyId: generateFriendlyId("attempt"), - taskRunId: lockedTaskRun.id, - startedAt: new Date(), - backgroundWorkerId: backgroundTask.workerId, - backgroundWorkerTaskId: backgroundTask.id, - status: "EXECUTING" as const, - queueId: queue.id, - runtimeEnvironmentId: this.env.id, - }, - }); - - const execution: TaskRunExecution = { - task: { - id: backgroundTask.slug, - filePath: backgroundTask.filePath, - exportName: backgroundTask.exportName, - }, - attempt: { - id: taskRunAttempt.friendlyId, - number: taskRunAttempt.number, - startedAt: taskRunAttempt.startedAt ?? taskRunAttempt.createdAt, - backgroundWorkerId: backgroundWorker.id, - backgroundWorkerTaskId: backgroundTask.id, - status: "EXECUTING" as const, - }, - run: { - id: lockedTaskRun.friendlyId, - payload: lockedTaskRun.payload, - payloadType: lockedTaskRun.payloadType, - context: lockedTaskRun.context, - createdAt: lockedTaskRun.createdAt, - tags: lockedTaskRun.tags.map((tag) => tag.name), - isTest: lockedTaskRun.isTest, - idempotencyKey: lockedTaskRun.idempotencyKey ?? undefined, - }, - queue: { - id: queue.friendlyId, - name: queue.name, - }, - environment: { - id: this.env.id, - slug: this.env.slug, - type: this.env.type, - }, - organization: { - id: this.env.organization.id, - slug: this.env.organization.slug, - name: this.env.organization.title, - }, - project: { - id: this.env.project.id, - ref: this.env.project.externalRef, - slug: this.env.project.slug, - name: this.env.project.name, - }, - batch: - lockedTaskRun.batchItems[0] && lockedTaskRun.batchItems[0].batchTaskRun - ? { id: lockedTaskRun.batchItems[0].batchTaskRun.friendlyId } - : undefined, - }; - const environmentRepository = new EnvironmentVariablesRepository(); const variables = await environmentRepository.getEnvironmentVariables( this.env.project.id, this.env.id ); - const payload: TaskRunExecutionPayload = { - execution, - traceContext: lockedTaskRun.traceContext as Record, - environment: variables.reduce((acc: Record, curr) => { - acc[curr.key] = curr.value; - return acc; - }, {}), - }; + if (backgroundWorker.supportsLazyAttempts) { + const payload: TaskRunExecutionLazyAttemptPayload = { + traceContext: lockedTaskRun.traceContext as Record, + environment: variables.reduce((acc: Record, curr) => { + acc[curr.key] = curr.value; + return acc; + }, {}), + runId: lockedTaskRun.friendlyId, + messageId: lockedTaskRun.id, + isTest: lockedTaskRun.isTest, + }; - try { - // TODO: send trace context down to the CLI - await this._sender.send("BACKGROUND_WORKER_MESSAGE", { - backgroundWorkerId: backgroundWorker.friendlyId, - data: { - type: "EXECUTE_RUNS", - payloads: [payload], - }, - }); - - logger.debug("Saving the in progress attempt", { - taskRunAttempt: taskRunAttempt.id, - messageId: message.messageId, - }); - - this._inProgressAttempts.set(taskRunAttempt.friendlyId, message.messageId); - this._inProgressRuns.set(lockedTaskRun.friendlyId, message.messageId); - } catch (e) { - if (e instanceof Error) { - this._currentSpan?.recordException(e); - } else { - this._currentSpan?.recordException(new Error(String(e))); - } - - this._endSpanInNextIteration = true; - - // We now need to unlock the task run and delete the task run attempt - await prisma.$transaction([ - prisma.taskRun.update({ - where: { - id: lockedTaskRun.id, - }, + try { + await this._sender.send("BACKGROUND_WORKER_MESSAGE", { + backgroundWorkerId: backgroundWorker.friendlyId, data: { - lockedAt: null, - lockedById: null, - status: "PENDING", + type: "EXECUTE_RUN_LAZY_ATTEMPT", + payload, }, - }), - prisma.taskRunAttempt.delete({ - where: { - id: taskRunAttempt.id, + }); + + logger.debug("Executing the run", { + messageId: message.messageId, + }); + + this._inProgressRuns.set(lockedTaskRun.friendlyId, message.messageId); + } catch (e) { + if (e instanceof Error) { + this._currentSpan?.recordException(e); + } else { + this._currentSpan?.recordException(new Error(String(e))); + } + + this._endSpanInNextIteration = true; + + // We now need to unlock the task run and delete the task run attempt + await prisma.$transaction([ + prisma.taskRun.update({ + where: { + id: lockedTaskRun.id, + }, + data: { + lockedAt: null, + lockedById: null, + status: "PENDING", + }, + }), + ]); + + this._inProgressRuns.delete(lockedTaskRun.friendlyId); + + // Finally we need to nack the message so it can be retried + await marqs?.nackMessage(message.messageId); + } finally { + setTimeout(() => this.#doWork(), 100); + } + } else { + const service = new CreateTaskRunAttemptService(); + const execution = await service.call(lockedTaskRun.friendlyId, this.env); + + const payload: TaskRunExecutionPayload = { + traceContext: lockedTaskRun.traceContext as Record, + environment: variables.reduce((acc: Record, curr) => { + acc[curr.key] = curr.value; + return acc; + }, {}), + execution, + }; + + try { + await this._sender.send("BACKGROUND_WORKER_MESSAGE", { + backgroundWorkerId: backgroundWorker.friendlyId, + data: { + type: "EXECUTE_RUNS", + payloads: [payload], }, - }), - ]); + }); - this._inProgressAttempts.delete(taskRunAttempt.friendlyId); - this._inProgressRuns.delete(lockedTaskRun.friendlyId); + logger.debug("Executing the run", { + messageId: message.messageId, + }); - // Finally we need to nack the message so it can be retried - await marqs?.nackMessage(message.messageId); - } finally { - setTimeout(() => this.#doWork(), 100); + this._inProgressRuns.set(lockedTaskRun.friendlyId, message.messageId); + } catch (e) { + if (e instanceof Error) { + this._currentSpan?.recordException(e); + } else { + this._currentSpan?.recordException(new Error(String(e))); + } + + this._endSpanInNextIteration = true; + + // We now need to unlock the task run and delete the task run attempt + await prisma.$transaction([ + prisma.taskRun.update({ + where: { + id: lockedTaskRun.id, + }, + data: { + lockedAt: null, + lockedById: null, + status: "PENDING", + }, + }), + ]); + + this._inProgressRuns.delete(lockedTaskRun.friendlyId); + + // Finally we need to nack the message so it can be retried + await marqs?.nackMessage(message.messageId); + } finally { + setTimeout(() => this.#doWork(), 100); + } } } diff --git a/apps/webapp/app/v3/services/baseService.server.ts b/apps/webapp/app/v3/services/baseService.server.ts index 6892118c6..e6b9d0252 100644 --- a/apps/webapp/app/v3/services/baseService.server.ts +++ b/apps/webapp/app/v3/services/baseService.server.ts @@ -34,7 +34,7 @@ export abstract class BaseService { } export class ServiceValidationError extends Error { - constructor(message: string) { + constructor(message: string, public status?: number) { super(message); this.name = "ServiceValidationError"; } diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts index 4137d7a16..79a5e2878 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts @@ -63,6 +63,7 @@ export class CreateBackgroundWorkerService extends BaseService { contentHash: body.metadata.contentHash, cliVersion: body.metadata.cliPackageVersion, sdkVersion: body.metadata.packageVersion, + supportsLazyAttempts: body.supportsLazyAttempts, }, }); diff --git a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts new file mode 100644 index 000000000..b72ecfe06 --- /dev/null +++ b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts @@ -0,0 +1,131 @@ +import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { BaseService, ServiceValidationError } from "./baseService.server"; +import { TaskRunExecution } from "@trigger.dev/core/v3"; +import { prisma } from "~/db.server"; +import { generateFriendlyId } from "../friendlyIdentifiers"; +import { logger } from "~/services/logger.server"; + +export class CreateTaskRunAttemptService extends BaseService { + public async call( + runFriendlyId: string, + environment: AuthenticatedEnvironment + ): Promise { + return await this.traceWithEnv("call()", environment, async (span) => { + span.setAttribute("taskRunId", runFriendlyId); + + const taskRun = await this._prisma.taskRun.findUnique({ + where: { + friendlyId: runFriendlyId, + runtimeEnvironmentId: environment.id, + }, + include: { + tags: true, + attempts: true, + lockedBy: { + include: { + worker: true, + }, + }, + batchItems: { + include: { + batchTaskRun: true, + }, + }, + }, + }); + + logger.debug("Creating a task run attempt", { taskRun }); + + if (!taskRun) { + throw new ServiceValidationError("Task run not found", 404); + } + + if (taskRun.status === "CANCELED") { + throw new ServiceValidationError("Task run is cancelled", 400); + } + + if (!taskRun.lockedBy) { + throw new ServiceValidationError("Task run is not locked", 400); + } + + const queue = await this._prisma.taskQueue.findUnique({ + where: { + runtimeEnvironmentId_name: { + runtimeEnvironmentId: environment.id, + name: taskRun.queue, + }, + }, + }); + + if (!queue) { + throw new ServiceValidationError("Queue not found", 404); + } + + const taskRunAttempt = await prisma.taskRunAttempt.create({ + data: { + number: taskRun.attempts[0] ? taskRun.attempts[0].number + 1 : 1, + friendlyId: generateFriendlyId("attempt"), + taskRunId: taskRun.id, + startedAt: new Date(), + backgroundWorkerId: taskRun.lockedBy.worker.id, + backgroundWorkerTaskId: taskRun.lockedBy.id, + status: "EXECUTING" as const, + queueId: queue.id, + runtimeEnvironmentId: environment.id, + }, + }); + + const execution: TaskRunExecution = { + task: { + id: taskRun.lockedBy.slug, + filePath: taskRun.lockedBy.filePath, + exportName: taskRun.lockedBy.exportName, + }, + attempt: { + id: taskRunAttempt.friendlyId, + number: taskRunAttempt.number, + startedAt: taskRunAttempt.startedAt ?? taskRunAttempt.createdAt, + backgroundWorkerId: taskRun.lockedBy.worker.id, + backgroundWorkerTaskId: taskRun.lockedBy.id, + status: "EXECUTING" as const, + }, + run: { + id: taskRun.friendlyId, + payload: taskRun.payload, + payloadType: taskRun.payloadType, + context: taskRun.context, + createdAt: taskRun.createdAt, + tags: taskRun.tags.map((tag) => tag.name), + isTest: taskRun.isTest, + idempotencyKey: taskRun.idempotencyKey ?? undefined, + }, + queue: { + id: queue.friendlyId, + name: queue.name, + }, + environment: { + id: environment.id, + slug: environment.slug, + type: environment.type, + }, + organization: { + id: environment.organization.id, + slug: environment.organization.slug, + name: environment.organization.title, + }, + project: { + id: environment.project.id, + ref: environment.project.externalRef, + slug: environment.project.slug, + name: environment.project.name, + }, + batch: + taskRun.batchItems[0] && taskRun.batchItems[0].batchTaskRun + ? { id: taskRun.batchItems[0].batchTaskRun.friendlyId } + : undefined, + }; + + return execution; + }); + } +} diff --git a/packages/cli-v3/src/apiClient.ts b/packages/cli-v3/src/apiClient.ts index 7c7751813..347ca1d82 100644 --- a/packages/cli-v3/src/apiClient.ts +++ b/packages/cli-v3/src/apiClient.ts @@ -14,6 +14,7 @@ import { GetDeploymentResponseBody, GetProjectsResponseBody, GetProjectResponseBody, + TaskRunExecution, } from "@trigger.dev/core/v3"; export class CliApiClient { @@ -103,6 +104,20 @@ export class CliApiClient { ); } + async createTaskRunAttempt(runFriendlyId: string) { + if (!this.accessToken) { + throw new Error("creatTaskRunAttempt: No access token"); + } + + return zodfetch(TaskRunExecution, `${this.apiURL}/api/v1/runs/${runFriendlyId}/attempts`, { + method: "POST", + headers: { + Authorization: `Bearer ${this.accessToken}`, + "Content-Type": "application/json", + }, + }); + } + async getProjectEnv({ projectRef, env, @@ -198,11 +213,11 @@ type ApiResult = error: string; }; -async function zodfetch( - schema: z.Schema, +async function zodfetch( + schema: T, url: string, requestInit?: RequestInit -): Promise> { +): Promise>> { try { const response = await fetch(url, requestInit); diff --git a/packages/cli-v3/src/commands/dev.tsx b/packages/cli-v3/src/commands/dev.tsx index c9e684847..adc61a3e3 100644 --- a/packages/cli-v3/src/commands/dev.tsx +++ b/packages/cli-v3/src/commands/dev.tsx @@ -278,6 +278,7 @@ function useDev({ websocket.addEventListener("close", (event) => {}); websocket.addEventListener("error", (event) => {}); + // This is the deprecated task heart beat that uses the friendly attempt ID backgroundWorkerCoordinator.onWorkerTaskHeartbeat.attach( async ({ worker, backgroundWorkerId, id }) => { await sender.send("BACKGROUND_WORKER_MESSAGE", { @@ -290,6 +291,19 @@ function useDev({ } ); + // "Task Run Heartbeat" id is the actual run ID that corresponds to the MarQS message ID + backgroundWorkerCoordinator.onWorkerTaskRunHeartbeat.attach( + async ({ worker, backgroundWorkerId, id }) => { + await sender.send("BACKGROUND_WORKER_MESSAGE", { + backgroundWorkerId, + data: { + type: "TASK_RUN_HEARTBEAT", + id, + }, + }); + } + ); + backgroundWorkerCoordinator.onTaskCompleted.attach( async ({ backgroundWorkerId, completion, execution }) => { await sender.send("BACKGROUND_WORKER_MESSAGE", { @@ -327,6 +341,7 @@ function useDev({ for (const worker of backgroundWorkerCoordinator.currentWorkers) { await sender.send("READY_FOR_TASKS", { backgroundWorkerId: worker.id, + inProgressRuns: worker.worker.inProgressRuns, }); } }, @@ -495,20 +510,24 @@ function useDev({ const processEnv = await gatherProcessEnv(); - const backgroundWorker = new BackgroundWorker(fullPath, { - projectConfig: config, - dependencies, - env: { - ...processEnv, - TRIGGER_API_URL: apiUrl, - TRIGGER_SECRET_KEY: apiKey, - ...(environmentVariablesResponse.success - ? environmentVariablesResponse.data.variables - : {}), + const backgroundWorker = new BackgroundWorker( + fullPath, + { + projectConfig: config, + dependencies, + env: { + ...processEnv, + TRIGGER_API_URL: apiUrl, + TRIGGER_SECRET_KEY: apiKey, + ...(environmentVariablesResponse.success + ? environmentVariablesResponse.data.variables + : {}), + }, + debuggerOn, + debugOtel, }, - debuggerOn, - debugOtel, - }); + environmentClient + ); try { await backgroundWorker.initialize(); @@ -565,6 +584,7 @@ function useDev({ tasks: taskResources, contentHash: contentHash, }, + supportsLazyAttempts: true, }; const backgroundWorkerRecord = await environmentClient.createBackgroundWorker( @@ -816,18 +836,9 @@ function createDuplicateTaskIdOutputErrorMessage( async function gatherProcessEnv() { const env = { + ...process.env, NODE_ENV: process.env.NODE_ENV ?? "development", - PATH: process.env.PATH, - USER: process.env.USER, - SHELL: process.env.SHELL, - NVM_INC: process.env.NVM_INC, - NVM_DIR: process.env.NVM_DIR, - NVM_BIN: process.env.NVM_BIN, - LANG: process.env.LANG, - TERM: process.env.TERM, NODE_PATH: await amendNodePathWithPnpmNodeModules(process.env.NODE_PATH), - HOME: process.env.HOME, - BUN_INSTALL: process.env.BUN_INSTALL, }; // Filter out undefined values diff --git a/packages/cli-v3/src/workers/dev/backgroundWorker.ts b/packages/cli-v3/src/workers/dev/backgroundWorker.ts index 0b2b2d715..c0b75fdba 100644 --- a/packages/cli-v3/src/workers/dev/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/dev/backgroundWorker.ts @@ -9,6 +9,7 @@ import { TaskRunError, TaskRunErrorCodes, TaskRunExecution, + TaskRunExecutionLazyAttemptPayload, TaskRunExecutionPayload, TaskRunExecutionResult, childToWorkerMessages, @@ -37,6 +38,7 @@ import { safeDeleteFileSync } from "../../utilities/fileSystem.js"; import { installPackages } from "../../utilities/installPackages.js"; import { logger } from "../../utilities/logger.js"; import { TaskMetadataParseError, UncaughtExceptionError } from "../common/errors.js"; +import { CliApiClient } from "../../apiClient.js"; export type CurrentWorkers = BackgroundWorkerCoordinator["currentWorkers"]; export class BackgroundWorkerCoordinator { @@ -51,11 +53,20 @@ export class BackgroundWorkerCoordinator { id: string; record: CreateBackgroundWorkerResponse; }> = new Evt(); + + /** + * @deprecated use onWorkerTaskRunHeartbeat instead + */ public onWorkerTaskHeartbeat: Evt<{ id: string; backgroundWorkerId: string; worker: BackgroundWorker; }> = new Evt(); + public onWorkerTaskRunHeartbeat: Evt<{ + id: string; + backgroundWorkerId: string; + worker: BackgroundWorker; + }> = new Evt(); public onWorkerDeprecated: Evt<{ worker: BackgroundWorker; id: string }> = new Evt(); private _backgroundWorkers: Map = new Map(); private _records: Map = new Map(); @@ -106,6 +117,10 @@ export class BackgroundWorkerCoordinator { worker.onTaskHeartbeat.attach((id) => { this.onWorkerTaskHeartbeat.post({ id, backgroundWorkerId: record.id, worker }); }); + + worker.onTaskRunHeartbeat.attach((id) => { + this.onWorkerTaskRunHeartbeat.post({ id, backgroundWorkerId: record.id, worker }); + }); } close() { @@ -135,10 +150,39 @@ export class BackgroundWorkerCoordinator { } await worker.cancelRun(message.taskRunId); + break; + } + case "EXECUTE_RUN_LAZY_ATTEMPT": { + await this.#executeTaskRunLazyAttempt(id, message.payload); } } } + async #executeTaskRunLazyAttempt(id: string, payload: TaskRunExecutionLazyAttemptPayload) { + const worker = this._backgroundWorkers.get(id); + + if (!worker) { + logger.error(`Could not find worker ${id}`); + return; + } + + const record = this._records.get(id); + + if (!record) { + logger.error(`Could not find worker record ${id}`); + return; + } + + const { completion, execution } = await worker.executeTaskRunLazyAttempt(payload, this.baseURL); + + this.onTaskCompleted.post({ + completion, + execution, + worker, + backgroundWorkerId: id, + }); + } + async #executeTaskRun(id: string, payload: TaskRunExecutionPayload) { const worker = this._backgroundWorkers.get(id); @@ -154,82 +198,14 @@ export class BackgroundWorkerCoordinator { return; } - const { execution } = payload; + const completion = await worker.executeTaskRun(payload, this.baseURL); - // ○ Mar 27 09:17:25.653 -> View logs | 20240326.20 | create-avatar | run_slufhjdfiv8ejnrkw9dsj.1 - - const logsUrl = `${this.baseURL}/runs/${execution.run.id}`; - - const pipe = chalkGrey("|"); - const bullet = chalkGrey("○"); - const link = chalkLink(terminalLink("View logs", logsUrl)); - let timestampPrefix = chalkGrey(prettyPrintDate(payload.execution.attempt.startedAt)); - const workerPrefix = chalkWorker(record.version); - const taskPrefix = chalkTask(execution.task.id); - const runId = chalkRun(`${execution.run.id}.${execution.attempt.number}`); - - logger.log( - `${bullet} ${timestampPrefix} ${chalkGrey( - "->" - )} ${link} ${pipe} ${workerPrefix} ${pipe} ${taskPrefix} ${pipe} ${runId}` - ); - - const now = performance.now(); - - const completion = await worker.executeTaskRun(payload); - - const elapsed = performance.now() - now; - - const retryingText = chalkGrey( - !completion.ok && completion.skippedRetrying - ? " (retrying skipped)" - : !completion.ok && completion.retry !== undefined - ? ` (retrying in ${completion.retry.delay}ms)` - : "" - ); - - const resultText = !completion.ok - ? completion.error.type === "INTERNAL_ERROR" && - (completion.error.code === TaskRunErrorCodes.TASK_EXECUTION_ABORTED || - completion.error.code === TaskRunErrorCodes.TASK_RUN_CANCELLED) - ? chalkWarning("Cancelled") - : `${chalkError("Error")}${retryingText}` - : chalkSuccess("Success"); - - const errorText = !completion.ok - ? this.#formatErrorLog(completion.error) - : "retry" in completion - ? `retry in ${completion.retry}ms` - : ""; - - const elapsedText = chalkGrey(`(${formatDurationMilliseconds(elapsed, { style: "short" })})`); - - timestampPrefix = chalkGrey(prettyPrintDate()); - - logger.log( - `${bullet} ${timestampPrefix} ${chalkGrey( - "->" - )} ${link} ${pipe} ${workerPrefix} ${pipe} ${taskPrefix} ${pipe} ${runId} ${pipe} ${resultText} ${elapsedText}${errorText}` - ); - - this.onTaskCompleted.post({ completion, execution, worker, backgroundWorkerId: id }); - } - - #formatErrorLog(error: TaskRunError) { - switch (error.type) { - case "INTERNAL_ERROR": { - return ""; - } - case "STRING_ERROR": { - return `\n\n${chalkError("X Error:")} ${error.raw}\n`; - } - case "CUSTOM_ERROR": { - return `\n\n${chalkError("X Error:")} ${error.raw}\n`; - } - case "BUILT_IN_ERROR": { - return `\n\n${error.stackTrace.replace(/^Error: /, chalkError("X Error: "))}\n`; - } - } + this.onTaskCompleted.post({ + completion, + execution: payload.execution, + worker, + backgroundWorkerId: id, + }); } } @@ -264,13 +240,18 @@ export type BackgroundWorkerParams = { debuggerOn: boolean; debugOtel?: boolean; }; + export class BackgroundWorker { private _initialized: boolean = false; private _handler = new ZodMessageHandler({ schema: childToWorkerMessages, }); + /** + * @deprecated use onTaskRunHeartbeat instead + */ public onTaskHeartbeat: Evt = new Evt(); + public onTaskRunHeartbeat: Evt = new Evt(); private _onClose: Evt = new Evt(); public tasks: Array = []; @@ -282,7 +263,8 @@ export class BackgroundWorker { constructor( public path: string, - private params: BackgroundWorkerParams + private params: BackgroundWorkerParams, + private apiClient: CliApiClient ) {} close() { @@ -293,6 +275,7 @@ export class BackgroundWorker { this._closed = true; this.onTaskHeartbeat.detach(); + this.onTaskRunHeartbeat.detach(); // We need to close all the task run processes for (const taskRunProcess of this._taskRunProcesses.values()) { @@ -306,6 +289,10 @@ export class BackgroundWorker { safeDeleteFileSync(`${this.path}.map`); } + get inProgressRuns(): Array { + return Array.from(this._taskRunProcesses.keys()); + } + async initialize() { if (this._initialized) { throw new Error("Worker already initialized"); @@ -393,14 +380,18 @@ export class BackgroundWorker { } } - async #initializeTaskRunProcess(payload: TaskRunExecutionPayload): Promise { + async #initializeTaskRunProcess( + payload: TaskRunExecutionPayload, + messageId?: string + ): Promise { if (!this.metadata) { throw new Error("Worker not registered"); } if (!this._taskRunProcesses.has(payload.execution.run.id)) { const taskRunProcess = new TaskRunProcess( - payload.execution, + payload.execution.run.id, + payload.execution.run.isTest, this.path, { ...this.params.env, @@ -408,7 +399,8 @@ export class BackgroundWorker { ...this.#readEnvVars(), }, this.metadata, - this.params + this.params, + messageId ); taskRunProcess.onExit.attach(() => { @@ -419,6 +411,10 @@ export class BackgroundWorker { this.onTaskHeartbeat.post(id); }); + taskRunProcess.onTaskRunHeartbeat.attach((id) => { + this.onTaskRunHeartbeat.post(id); + }); + await taskRunProcess.initialize(); this._taskRunProcesses.set(payload.execution.run.id, taskRunProcess); @@ -437,10 +433,104 @@ export class BackgroundWorker { await taskRunProcess.cancel(); } + async executeTaskRunLazyAttempt(payload: TaskRunExecutionLazyAttemptPayload, baseURL: string) { + const attemptResponse = await this.apiClient.createTaskRunAttempt(payload.runId); + + if (!attemptResponse.success) { + throw new Error(`Failed to create task run attempt: ${attemptResponse.error}`); + } + + const execution = attemptResponse.data; + + const completion = await this.executeTaskRun( + { execution, traceContext: payload.traceContext, environment: payload.environment }, + baseURL, + payload.messageId + ); + + return { execution, completion }; + } + // We need to fork the process before we can execute any tasks - async executeTaskRun(payload: TaskRunExecutionPayload): Promise { + async executeTaskRun( + payload: TaskRunExecutionPayload, + baseURL: string, + messageId?: string + ): Promise { + if (this._closed) { + throw new Error("Worker is closed"); + } + + if (!this.metadata) { + throw new Error("Worker not registered"); + } + + const { execution } = payload; + // ○ Mar 27 09:17:25.653 -> View logs | 20240326.20 | create-avatar | run_slufhjdfiv8ejnrkw9dsj.1 + + const logsUrl = `${baseURL}/runs/${execution.run.id}`; + + const pipe = chalkGrey("|"); + const bullet = chalkGrey("○"); + const link = chalkLink(terminalLink("View logs", logsUrl)); + let timestampPrefix = chalkGrey(prettyPrintDate(payload.execution.attempt.startedAt)); + const workerPrefix = chalkWorker(this.metadata.version); + const taskPrefix = chalkTask(execution.task.id); + const runId = chalkRun(`${execution.run.id}.${execution.attempt.number}`); + + logger.log( + `${bullet} ${timestampPrefix} ${chalkGrey( + "->" + )} ${link} ${pipe} ${workerPrefix} ${pipe} ${taskPrefix} ${pipe} ${runId}` + ); + + const now = performance.now(); + + const completion = await this.#doExecuteTaskRun(payload, messageId); + + const elapsed = performance.now() - now; + + const retryingText = chalkGrey( + !completion.ok && completion.skippedRetrying + ? " (retrying skipped)" + : !completion.ok && completion.retry !== undefined + ? ` (retrying in ${completion.retry.delay}ms)` + : "" + ); + + const resultText = !completion.ok + ? completion.error.type === "INTERNAL_ERROR" && + (completion.error.code === TaskRunErrorCodes.TASK_EXECUTION_ABORTED || + completion.error.code === TaskRunErrorCodes.TASK_RUN_CANCELLED) + ? chalkWarning("Cancelled") + : `${chalkError("Error")}${retryingText}` + : chalkSuccess("Success"); + + const errorText = !completion.ok + ? formatErrorLog(completion.error) + : "retry" in completion + ? `retry in ${completion.retry}ms` + : ""; + + const elapsedText = chalkGrey(`(${formatDurationMilliseconds(elapsed, { style: "short" })})`); + + timestampPrefix = chalkGrey(prettyPrintDate()); + + logger.log( + `${bullet} ${timestampPrefix} ${chalkGrey( + "->" + )} ${link} ${pipe} ${workerPrefix} ${pipe} ${taskPrefix} ${pipe} ${runId} ${pipe} ${resultText} ${elapsedText}${errorText}` + ); + + return completion; + } + + async #doExecuteTaskRun( + payload: TaskRunExecutionPayload, + messageId?: string + ): Promise { try { - const taskRunProcess = await this.#initializeTaskRunProcess(payload); + const taskRunProcess = await this.#initializeTaskRunProcess(payload, messageId); const result = await taskRunProcess.executeTaskRun(payload); // Kill the worker if the task was successful or if it's not going to be retried); @@ -553,15 +643,21 @@ class TaskRunProcess { private _currentExecution: TaskRunExecution | undefined; 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(); constructor( - private execution: TaskRunExecution, + private runId: string, + private isTest: boolean, private path: string, private env: NodeJS.ProcessEnv, private metadata: BackgroundWorkerProperties, - private worker: BackgroundWorkerParams + private worker: BackgroundWorkerParams, + private messageId?: string ) { this._sender = new ZodMessageSender({ schema: workerToChildMessages, @@ -581,7 +677,7 @@ class TaskRunProcess { async initialize() { const fullEnv = { - ...(this.execution.run.isTest ? { TRIGGER_LOG_LEVEL: "debug" } : {}), + ...(this.isTest ? { TRIGGER_LOG_LEVEL: "debug" } : {}), ...this.env, OTEL_RESOURCE_ATTRIBUTES: JSON.stringify({ [SemanticInternalAttributes.PROJECT_DIR]: this.worker.projectConfig.projectDir, @@ -592,7 +688,7 @@ class TaskRunProcess { const cwd = dirname(this.path); - logger.debug(`[${this.execution.run.id}] initializing task run process`, { + logger.debug(`[${this.runId}] initializing task run process`, { env: fullEnv, path: this.path, cwd, @@ -618,7 +714,7 @@ class TaskRunProcess { return; } - logger.debug(`[${this.execution.run.id}] cleaning up task run process`, { kill }); + logger.debug(`[${this.runId}] cleaning up task run process`, { kill }); await this._sender.send("CLEANUP", { flush: true, @@ -630,7 +726,7 @@ class TaskRunProcess { // Set a timeout to kill the child process if it hasn't been killed within 5 seconds setTimeout(() => { if (this._child && !this._child.killed) { - logger.debug(`[${this.execution.run.id}] killing task run process after timeout`); + logger.debug(`[${this.runId}] killing task run process after timeout`); this._child.kill(); } @@ -673,12 +769,12 @@ class TaskRunProcess { return; } - if (execution.run.id === this.execution.run.id) { + if (execution.run.id === this.runId) { // We don't need to notify the task run process if it's the same as the one we're running return; } - logger.debug(`[${this.execution.run.id}] task run completed notification`, { + logger.debug(`[${this.runId}] task run completed notification`, { completion, execution, }); @@ -717,14 +813,18 @@ class TaskRunProcess { break; } case "READY_TO_DISPOSE": { - logger.debug(`[${this.execution.run.id}] task run process is ready to dispose`); + logger.debug(`[${this.runId}] task run process is ready to dispose`); this.#kill(); break; } case "TASK_HEARTBEAT": { - this.onTaskHeartbeat.post(message.payload.id); + if (this.messageId) { + this.onTaskRunHeartbeat.post(this.messageId); + } else { + this.onTaskHeartbeat.post(message.payload.id); + } break; } @@ -735,7 +835,7 @@ class TaskRunProcess { } async #handleExit(code: number) { - logger.debug(`[${this.execution.run.id}] task run process exiting`, { code }); + logger.debug(`[${this.runId}] task run process exiting`, { code }); // Go through all the attempts currently pending and reject them for (const [id, status] of this._attemptStatuses.entries()) { @@ -801,9 +901,26 @@ class TaskRunProcess { #kill() { if (this._child && !this._child.killed) { - logger.debug(`[${this.execution.run.id}] killing task run process`); + logger.debug(`[${this.runId}] killing task run process`); this._child?.kill(); } } } + +function formatErrorLog(error: TaskRunError) { + switch (error.type) { + case "INTERNAL_ERROR": { + return ""; + } + case "STRING_ERROR": { + return `\n\n${chalkError("X Error:")} ${error.raw}\n`; + } + case "CUSTOM_ERROR": { + return `\n\n${chalkError("X Error:")} ${error.raw}\n`; + } + case "BUILT_IN_ERROR": { + return `\n\n${error.stackTrace.replace(/^Error: /, chalkError("X Error: "))}\n`; + } + } +} diff --git a/packages/core/src/v3/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index 78013422c..30101309c 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -41,6 +41,7 @@ export type GetProjectEnvResponse = z.infer; export const CreateBackgroundWorkerRequestBody = z.object({ localOnly: z.boolean(), metadata: BackgroundWorkerMetadata, + supportsLazyAttempts: z.boolean().optional(), }); export type CreateBackgroundWorkerRequestBody = z.infer; diff --git a/packages/core/src/v3/schemas/common.ts b/packages/core/src/v3/schemas/common.ts index c190d1631..58ae633c7 100644 --- a/packages/core/src/v3/schemas/common.ts +++ b/packages/core/src/v3/schemas/common.ts @@ -50,7 +50,7 @@ export const TaskRunInternalError = z.object({ "TASK_RUN_CANCELLED", "TASK_OUTPUT_ERROR", "HANDLE_ERROR_ERROR", - "GRACEFUL_EXIT_TIMEOUT" + "GRACEFUL_EXIT_TIMEOUT", ]), message: z.string().optional(), }); diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index 32c66c808..becdc9fda 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -1,11 +1,13 @@ import { z } from "zod"; import { TaskRunExecution, TaskRunExecutionResult } from "./common"; + import { EnvironmentType, Machine, ProdTaskRunExecution, ProdTaskRunExecutionPayload, TaskMetadataWithFilePath, + TaskRunExecutionLazyAttemptPayload, TaskRunExecutionPayload, WaitReason, } from "./schemas"; @@ -34,6 +36,10 @@ export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [ projectId: z.string(), runId: z.string(), }), + z.object({ + type: z.literal("EXECUTE_RUN_LAZY_ATTEMPT"), + payload: TaskRunExecutionLazyAttemptPayload, + }), ]); export type BackgroundWorkerServerMessages = z.infer; @@ -62,6 +68,11 @@ export const BackgroundWorkerClientMessages = z.discriminatedUnion("type", [ type: z.literal("TASK_HEARTBEAT"), id: z.string(), }), + z.object({ + version: z.literal("v1").default("v1"), + type: z.literal("TASK_RUN_HEARTBEAT"), + id: z.string(), + }), ]); export type BackgroundWorkerClientMessages = z.infer; @@ -78,6 +89,7 @@ export const clientWebsocketMessages = { READY_FOR_TASKS: z.object({ version: z.literal("v1").default("v1"), backgroundWorkerId: z.string(), + inProgressRuns: z.string().array().optional(), }), BACKGROUND_WORKER_DEPRECATED: z.object({ version: z.literal("v1").default("v1"), diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index fe82a1be3..5a609a21d 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -223,3 +223,13 @@ export type ResolvedConfig = RequireKeys< export const WaitReason = z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]); export type WaitReason = z.infer; + +export const TaskRunExecutionLazyAttemptPayload = z.object({ + runId: z.string(), + messageId: z.string(), + isTest: z.boolean(), + traceContext: z.record(z.unknown()), + environment: z.record(z.string()).optional(), +}); + +export type TaskRunExecutionLazyAttemptPayload = z.infer; diff --git a/packages/database/prisma/migrations/20240430101936_/migration.sql b/packages/database/prisma/migrations/20240430101936_/migration.sql new file mode 100644 index 000000000..28b3b3909 --- /dev/null +++ b/packages/database/prisma/migrations/20240430101936_/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "BackgroundWorker" ADD COLUMN "supportsLazyAttempts" BOOLEAN NOT NULL DEFAULT false; diff --git a/packages/database/prisma/schema.prisma b/packages/database/prisma/schema.prisma index d36c3061e..da9207652 100644 --- a/packages/database/prisma/schema.prisma +++ b/packages/database/prisma/schema.prisma @@ -1539,6 +1539,8 @@ model BackgroundWorker { deployment WorkerDeployment? + supportsLazyAttempts Boolean @default(false) + @@unique([projectId, runtimeEnvironmentId, version]) } diff --git a/references/v3-catalog/src/trigger/longRunning.ts b/references/v3-catalog/src/trigger/longRunning.ts index 09462a1f1..17b761a85 100644 --- a/references/v3-catalog/src/trigger/longRunning.ts +++ b/references/v3-catalog/src/trigger/longRunning.ts @@ -19,10 +19,11 @@ export const longRunningParent = task({ run: async (payload: { message: string }) => { logger.info("Long running parent", { payload }); - await longRunning.triggerAndWait({ message: "child" }); + const result = await longRunning.triggerAndWait({ message: "child" }); return { finished: new Date().toISOString(), + result, }; }, });