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

restore from checkpoint events and fix statuses

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