github-featured/triggerdotdev--trigger.dev
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

lazy attempt creation for prod workers

Browse Source
This commit is contained in:
nicktrn
2024-05-02 15:58:02 +01:00
parent 14992f43ae
commit c75bbfdf2c
10 changed files with 536 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/coordinator/src/index.ts
+76
View File
@@ -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);
},
},
});
apps/webapp/app/v3/eventRepository.server.ts
+31
View File
@@ -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.
apps/webapp/app/v3/failedTaskRun.server.ts
+2 -36
View File
@@ -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,
};
}
}
}
apps/webapp/app/v3/handleSocketIo.server.ts
+48
View File
@@ -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();
apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts
+62 -22
View File
@@ -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<TaskRunExecutionLazyAttemptPayload | undefined> {
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<string, unknown>,
environment: variables.reduce((acc: Record<string, string>, 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 });
apps/webapp/app/v3/services/createTaskRunAttempt.server.ts
+62 -22
View File
@@ -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;
}
packages/cli-v3/src/workers/prod/backgroundWorker.ts
+78 -6
View File
@@ -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<string> = new Evt();
public onTaskRunHeartbeat: Evt<string> = new Evt();
public onWaitForBatch: Evt<
InferSocketMessageSchema<typeof ProdChildToWorkerMessages, "WAIT_FOR_BATCH">
@@ -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<void> = new Evt();
public tasks: Array<TaskMetadataWithFilePath> = [];
@@ -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<TaskRunProcess> {
async #initializeTaskRunProcess(
payload: ProdTaskRunExecutionPayload,
messageId?: string
): Promise<TaskRunProcess> {
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<TaskRunExecutionResult> {
async executeTaskRun(
payload: ProdTaskRunExecutionPayload,
messageId?: string
): Promise<TaskRunExecutionResult> {
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<string> = new Evt();
public onTaskRunHeartbeat: Evt<string> = new Evt();
public onExit: Evt<number> = 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) => {
packages/cli-v3/src/workers/prod/entry-point.ts
+94 -2
View File
@@ -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,
packages/core-apps/src/provider.ts
-2
View File
@@ -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);
packages/core/src/v3/schemas/messages.ts
+83 -1
View File
@@ -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"),