WIP worker TaskRunAttempt creation

This commit is contained in:
Eric Allam
2024-04-30 13:27:03 +01:00
parent 62700245a3
commit 5ed700dad8
16 changed files with 622 additions and 310 deletions
@@ -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 }
);
}
}
@@ -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;
}
}
},
},
@@ -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<string, string> = new Map(); // Keys are task attempt friendly IDs, values are TaskRun ids/queue message ids
private _inProgressRuns: Map<string, string> = 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<string, unknown>,
environment: variables.reduce((acc: Record<string, string>, curr) => {
acc[curr.key] = curr.value;
return acc;
}, {}),
};
if (backgroundWorker.supportsLazyAttempts) {
const payload: TaskRunExecutionLazyAttemptPayload = {
traceContext: lockedTaskRun.traceContext as Record<string, unknown>,
environment: variables.reduce((acc: Record<string, string>, 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<string, unknown>,
environment: variables.reduce((acc: Record<string, string>, 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);
}
}
}
@@ -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";
}
@@ -63,6 +63,7 @@ export class CreateBackgroundWorkerService extends BaseService {
contentHash: body.metadata.contentHash,
cliVersion: body.metadata.cliPackageVersion,
sdkVersion: body.metadata.packageVersion,
supportsLazyAttempts: body.supportsLazyAttempts,
},
});
@@ -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<TaskRunExecution> {
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;
});
}
}
+18 -3
View File
@@ -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<TSuccessResult> =
error: string;
};
async function zodfetch<TResponseBody extends any>(
schema: z.Schema<TResponseBody>,
async function zodfetch<T extends z.ZodTypeAny>(
schema: T,
url: string,
requestInit?: RequestInit
): Promise<ApiResult<TResponseBody>> {
): Promise<ApiResult<z.infer<T>>> {
try {
const response = await fetch(url, requestInit);
+34 -23
View File
@@ -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
@@ -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<string, BackgroundWorker> = new Map();
private _records: Map<string, CreateBackgroundWorkerResponse> = 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<string> = new Evt();
public onTaskRunHeartbeat: Evt<string> = new Evt();
private _onClose: Evt<void> = new Evt();
public tasks: Array<TaskMetadataWithFilePath> = [];
@@ -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<string> {
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<TaskRunProcess> {
async #initializeTaskRunProcess(
payload: TaskRunExecutionPayload,
messageId?: string
): Promise<TaskRunProcess> {
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<TaskRunExecutionResult> {
async executeTaskRun(
payload: TaskRunExecutionPayload,
baseURL: string,
messageId?: string
): Promise<TaskRunExecutionResult> {
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<TaskRunExecutionResult> {
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<string> = new Evt();
public onTaskRunHeartbeat: Evt<string> = new Evt();
public onExit: Evt<number> = 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`;
}
}
}
+1
View File
@@ -41,6 +41,7 @@ export type GetProjectEnvResponse = z.infer<typeof GetProjectEnvResponse>;
export const CreateBackgroundWorkerRequestBody = z.object({
localOnly: z.boolean(),
metadata: BackgroundWorkerMetadata,
supportsLazyAttempts: z.boolean().optional(),
});
export type CreateBackgroundWorkerRequestBody = z.infer<typeof CreateBackgroundWorkerRequestBody>;
+1 -1
View File
@@ -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(),
});
+12
View File
@@ -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<typeof BackgroundWorkerServerMessages>;
@@ -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<typeof BackgroundWorkerClientMessages>;
@@ -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"),
+10
View File
@@ -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<typeof WaitReason>;
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<typeof TaskRunExecutionLazyAttemptPayload>;
@@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "BackgroundWorker" ADD COLUMN "supportsLazyAttempts" BOOLEAN NOT NULL DEFAULT false;
+2
View File
@@ -1539,6 +1539,8 @@ model BackgroundWorker {
deployment WorkerDeployment?
supportsLazyAttempts Boolean @default(false)
@@unique([projectId, runtimeEnvironmentId, version])
}
@@ -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,
};
},
});