v3: checkpoint and restore with sim for local dev (#933)

* zod ipc connection with acks

* pass sender through to ipc handlers

* bits and bobs

* add host networking back in

* disable verbose logs

* restore after tasks and batches.. almost

* restore and resume all the things

* fix for systems without checkpoint support

* Get deployment indexing errors to work with the new ZodIpc stuff

---------

Co-authored-by: Eric Allam <eallam@icloud.com>
This commit is contained in:
nicktrn
2024-03-11 15:20:33 +00:00
committed by GitHub
parent 883feab257
commit 20aa1cbe9b
24 changed files with 1421 additions and 364 deletions
+156 -42
View File
@@ -30,25 +30,44 @@ const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "coordinator-secret";
const logger = new SimpleLogger(`[${NODE_NAME}]`);
type CheckpointerInitializeReturn = {
canCheckpoint: boolean;
willSimulate: boolean;
};
class Checkpointer {
#initialized = false;
#canCheckpoint = false;
#dockerMode = true;
#dockerMode = !process.env.KUBERNETES_PORT;
#logger = new SimpleLogger("[checkptr]");
async initialize() {
constructor(private opts = { forceSimulate: false }) {}
async initialize(): Promise<CheckpointerInitializeReturn> {
if (this.#initialized) {
return;
return this.#getInitializeReturn();
}
this.#logger.log(`${this.#dockerMode ? "Docker" : "Kubernetes"} mode`);
if (this.opts.forceSimulate) {
this.#logger.log(
"Forced simulation enabled. Will simulate regardless of checkpoint support."
);
}
try {
await $`criu --version`;
} catch (error) {
this.#logger.error("No checkpoint support: Missing CRIU binary");
if (this.#dockerMode) {
this.#logger.error("Will simulate instead");
}
this.#canCheckpoint = false;
this.#initialized = true;
return;
return this.#getInitializeReturn();
}
if (this.#dockerMode) {
@@ -58,29 +77,41 @@ class Checkpointer {
this.#logger.error(
"No checkpoint support: Docker needs to have experimental features enabled"
);
this.#logger.error("Will simulate instead");
this.#canCheckpoint = false;
this.#initialized = true;
return;
return this.#getInitializeReturn();
}
}
this.#logger.log(
`Full checkpoint support with docker ${this.#dockerMode ? "enabled" : "disabled"}`
`Full checkpoint support in ${this.#dockerMode ? "docker" : "kubernetes"} mode`
);
this.#initialized = true;
this.#canCheckpoint = true;
return this.#getInitializeReturn();
}
async checkpointAndPush(podName: string) {
#getInitializeReturn(): CheckpointerInitializeReturn {
return {
canCheckpoint: this.#canCheckpoint,
willSimulate: this.#dockerMode && (!this.#canCheckpoint || this.opts.forceSimulate),
};
}
async checkpointAndPush(podName: string, leaveRunning = false) {
await this.initialize();
if (!this.#canCheckpoint) {
if (!this.#dockerMode && !this.#canCheckpoint) {
this.#logger.error("No checkpoint support. Simulation requires docker.");
return;
}
try {
const { path } = await this.#checkpointContainer(podName);
const { path } = await this.#checkpointContainer(podName, leaveRunning);
const { tag } = await this.#buildImage(path, podName);
const { destination } = await this.#pushImage(tag);
@@ -102,20 +133,27 @@ class Checkpointer {
}
}
async #checkpointContainer(podName: string) {
async #checkpointContainer(podName: string, leaveRunning = false) {
await this.initialize();
if (!this.#canCheckpoint) {
throw new Error("No checkpoint support");
}
if (this.#dockerMode) {
this.#logger.log("Checkpointing:", podName);
const path = randomUUID();
try {
this.#logger.debug(await $`docker checkpoint create --leave-running ${podName} ${path}`);
if (this.opts.forceSimulate || !this.#canCheckpoint) {
this.#logger.log("Simulating checkpoint");
this.#logger.debug(await $`docker pause ${podName}`);
} else {
if (leaveRunning) {
this.#logger.debug(
await $`docker checkpoint create --leave-running ${podName} ${path}`
);
} else {
this.#logger.debug(await $`docker checkpoint create ${podName} ${path}`);
}
}
} catch (error: any) {
this.#logger.error(error.stderr);
}
@@ -123,6 +161,10 @@ class Checkpointer {
return { path };
}
if (!this.#canCheckpoint) {
throw new Error("No checkpoint support. Simulation requires docker.");
}
const containerId = this.#logger.debug(
// @ts-expect-error
await $`crictl ps`
@@ -146,15 +188,15 @@ class Checkpointer {
async #buildImage(checkpointPath: string, tag: string) {
await this.initialize();
if (!this.#canCheckpoint) {
throw new Error("No checkpoint support");
}
if (this.#dockerMode) {
// Nothing to do here
return { tag };
}
if (!this.#canCheckpoint) {
throw new Error("No checkpoint support. Simulation requires docker.");
}
const container = this.#logger.debug(await $`buildah from scratch`);
this.#logger.debug(await $`buildah add ${container} ${checkpointPath} /`);
this.#logger.debug(
@@ -171,15 +213,15 @@ class Checkpointer {
async #pushImage(tag: string) {
await this.initialize();
if (!this.#canCheckpoint) {
throw new Error("No checkpoint support");
}
if (this.#dockerMode) {
// Nothing to do here
return { destination: "" };
}
if (!this.#canCheckpoint) {
throw new Error("No checkpoint support. Simulation requires docker.");
}
const destination = `${REGISTRY_FQDN}/${REPO_NAME}:${tag}`;
this.#logger.debug(await $`buildah push --tls-verify=${REGISTRY_TLS_VERIFY} ${destination}`);
@@ -191,7 +233,7 @@ class Checkpointer {
class TaskCoordinator {
#httpServer: ReturnType<typeof createServer>;
#checkpointer = new Checkpointer();
#checkpointer = new Checkpointer({ forceSimulate: true });
#prodWorkerNamespace: ZodNamespace<
typeof ProdWorkerToCoordinatorMessages,
@@ -249,6 +291,16 @@ class TaskCoordinator {
taskSocket.emit("RESUME", message);
},
RESUME_AFTER_DURATION: async (message) => {
const taskSocket = await this.#getAttemptSocket(message.attemptId);
if (!taskSocket) {
logger.log("Socket for attempt not found", { attemptId: message.attemptId });
return;
}
taskSocket.emit("RESUME_AFTER_DURATION", message);
},
},
});
@@ -341,34 +393,49 @@ class TaskCoordinator {
return;
}
// FIXME: shouldn't wait for completion here
const completionAck = await socket.emitWithAck("EXECUTE_TASK_RUN", {
socket.emit("EXECUTE_TASK_RUN", {
version: "v1",
executionPayload: executionAck.payload,
});
});
if (!completionAck.success) {
logger.error("completion unsuccessful", { attemptId: socket.data.attemptId });
return;
}
socket.on("READY_FOR_RESUME", async (message) => {
logger.log("[READY_FOR_RESUME]", message);
this.#platformSocket?.send("READY_FOR_RESUME", message);
});
logger.log("completed task", { completionId: completionAck.completion.id });
socket.on("TASK_RUN_COMPLETED", async (message, callback) => {
logger.log("completed task", { completionId: message.completion.id });
this.#platformSocket?.send("TASK_RUN_COMPLETED", {
version: "v1",
execution: executionAck.payload.execution,
completion: completionAck.completion,
execution: message.execution,
completion: message.completion,
});
callback();
});
socket.on("WAIT_FOR_DURATION", async (message, callback) => {
logger.log("[WAIT_FOR_DURATION]", message);
const { canCheckpoint, willSimulate } = await this.#checkpointer.initialize();
callback({ willCheckpointAndRestore: canCheckpoint || willSimulate });
if (!canCheckpoint) {
return;
}
// Wait for attempt to reach checkpointable state
// TODO: The worker should let us know when to checkpoint so we don't have to guess
await new Promise((resolve) => setTimeout(resolve, 2_000));
const checkpoint = await this.#checkpointer.checkpointAndPush(socket.data.podName);
if (!checkpoint) {
logger.error("Failed to checkpoint", { podName: socket.data.podName });
callback({ success: false });
// TODO: We have to let the worker know about failures so it can use its own timer
return;
}
@@ -377,10 +444,63 @@ class TaskCoordinator {
attemptId: socket.data.attemptId,
docker: checkpoint.docker,
location: checkpoint.destination,
reason: "WAIT_FOR_DURATION",
reason: {
type: "WAIT_FOR_DURATION",
ms: message.ms,
},
});
});
callback({ success: true });
socket.on("WAIT_FOR_TASK", async (message, callback) => {
logger.log("[WAIT_FOR_TASK]", message);
const { canCheckpoint, willSimulate } = await this.#checkpointer.initialize();
callback({ willCheckpointAndRestore: canCheckpoint || willSimulate });
const checkpoint = await this.#checkpointer.checkpointAndPush(socket.data.podName);
if (!checkpoint) {
logger.error("Failed to checkpoint", { podName: socket.data.podName });
return;
}
this.#platformSocket?.send("CHECKPOINT_CREATED", {
version: "v1",
attemptId: socket.data.attemptId,
docker: checkpoint.docker,
location: checkpoint.destination,
reason: {
type: "WAIT_FOR_TASK",
id: message.id,
},
});
});
socket.on("WAIT_FOR_BATCH", async (message, callback) => {
logger.log("[WAIT_FOR_BATCH]", message);
const { canCheckpoint, willSimulate } = await this.#checkpointer.initialize();
callback({ willCheckpointAndRestore: canCheckpoint || willSimulate });
const checkpoint = await this.#checkpointer.checkpointAndPush(socket.data.podName);
if (!checkpoint) {
logger.error("Failed to checkpoint", { podName: socket.data.podName });
return;
}
this.#platformSocket?.send("CHECKPOINT_CREATED", {
version: "v1",
attemptId: socket.data.attemptId,
docker: checkpoint.docker,
location: checkpoint.destination,
reason: {
type: "WAIT_FOR_BATCH",
id: message.id,
},
});
});
socket.on("INDEX_TASKS", async (message, callback) => {
@@ -428,12 +548,6 @@ class TaskCoordinator {
TASK_HEARTBEAT: async (message) => {
this.#platformSocket?.send("TASK_HEARTBEAT", message);
},
WAIT_FOR_BATCH: async (message) => {
// this.#checkpointer.checkpointAndPush(socket.data.podName);
},
WAIT_FOR_TASK: async (message) => {
// this.#checkpointer.checkpointAndPush(socket.data.podName);
},
},
});
+20 -12
View File
@@ -11,6 +11,8 @@ const OTEL_EXPORTER_OTLP_ENDPOINT =
const logger = new SimpleLogger(`[${MACHINE_NAME}]`);
class DockerTaskOperations implements TaskOperations {
constructor(private opts = { forceSimulate: false }) {}
async index(opts: {
contentHash: string;
imageTag: string;
@@ -26,7 +28,7 @@ class DockerTaskOperations implements TaskOperations {
});
const { exitCode } = logger.debug(
await $`docker run --rm -e TRIGGER_SECRET_KEY=${opts.apiKey} -e TRIGGER_API_URL=${opts.apiUrl} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e INDEX_TASKS=true --name=${containerName} ${opts.imageTag}`
await $`docker run --network=host --rm -e TRIGGER_SECRET_KEY=${opts.apiKey} -e TRIGGER_API_URL=${opts.apiUrl} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e INDEX_TASKS=true --name=${containerName} ${opts.imageTag}`
);
if (exitCode !== 0) {
@@ -38,7 +40,7 @@ class DockerTaskOperations implements TaskOperations {
const containerName = this.#getRunContainerName(opts.attemptId);
const { exitCode } = logger.debug(
await $`docker run -d -e OTEL_EXPORTER_OTLP_ENDPOINT=${OTEL_EXPORTER_OTLP_ENDPOINT} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e TRIGGER_ATTEMPT_ID=${opts.attemptId} --name=${containerName} ${opts.image}`
await $`docker run --network=host -d -e OTEL_EXPORTER_OTLP_ENDPOINT=${OTEL_EXPORTER_OTLP_ENDPOINT} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e TRIGGER_ATTEMPT_ID=${opts.attemptId} --name=${containerName} ${opts.image}`
);
if (exitCode !== 0) {
@@ -46,18 +48,23 @@ class DockerTaskOperations implements TaskOperations {
}
}
async restore(opts: {
attemptId: string;
runId: string;
image: string;
name: string;
checkpointId: string;
machine: Machine;
}) {
async restore(opts: { attemptId: string; checkpointRef: string; machine: Machine }) {
const containerName = this.#getRunContainerName(opts.attemptId);
if (this.opts.forceSimulate) {
logger.log("Simulating restore");
const { exitCode } = logger.debug(await $`docker unpause ${containerName}`);
if (exitCode !== 0) {
throw new Error("docker unpause command failed");
}
return;
}
const { exitCode } = logger.debug(
await $`docker start --checkpoint=${opts.checkpointId} ${containerName}`
await $`docker start --checkpoint=${opts.checkpointRef} ${containerName}`
);
if (exitCode !== 0) {
@@ -83,7 +90,8 @@ class DockerTaskOperations implements TaskOperations {
}
const provider = new ProviderShell({
tasks: new DockerTaskOperations(),
tasks: new DockerTaskOperations({ forceSimulate: true }),
type: "docker",
});
provider.listen();
+1
View File
@@ -389,6 +389,7 @@ class KubernetesTaskOperations implements TaskOperations {
const provider = new ProviderShell({
tasks: new KubernetesTaskOperations(),
type: "kubernetes",
});
provider.listen();
+5 -1
View File
@@ -14,10 +14,10 @@ import { SharedSocketConnection } from "./sharedSocketConnection";
import { CreateCheckpointService } from "./services/createCheckpoint.server";
import { sharedQueueTasks } from "./marqs/sharedQueueConsumer.server";
import { CompleteAttemptService } from "./services/completeAttempt.server";
import { CreateBackgroundWorkerService } from "./services/createBackgroundWorker.server";
import { logger } from "~/services/logger.server";
import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
import { CreateDeployedBackgroundWorkerService } from "./services/createDeployedBackgroundWorker.server";
import { ResumeAttemptService } from "./services/resumeAttempt.server";
import { DeploymentIndexFailed } from "./services/deploymentIndexFailed.server";
export const socketIo = singleton("socketIo", initalizeIoServer);
@@ -58,6 +58,10 @@ function createCoordinatorNamespace(io: Server) {
return { success: true, payload };
}
},
READY_FOR_RESUME: async (message) => {
const resumeAttempt = new ResumeAttemptService();
await resumeAttempt.call(message);
},
TASK_RUN_COMPLETED: async (message) => {
const completeAttempt = new CompleteAttemptService();
await completeAttempt.call(message.completion, message.execution);
@@ -32,6 +32,9 @@ const MessageBody = z.discriminatedUnion("type", [
type: z.literal("RESUME"),
completedAttemptIds: z.string().array(),
}),
z.object({
type: z.literal("RESUME_AFTER_DURATION"),
}),
]);
type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTask[] };
@@ -500,7 +503,7 @@ export class SharedQueueConsumer {
const resumableAttempt = resumableRun?.attempts[0];
if (!resumableAttempt) {
logger.error("Task run attempt to resume not found", {
logger.error("Resumable attempt not found", {
queueMessage: message.data,
messageId: message.messageId,
});
@@ -509,67 +512,6 @@ export class SharedQueueConsumer {
return;
}
const deployment = await prisma.workerDeployment.findFirst({
where: {
environmentId: resumableRun.runtimeEnvironmentId,
projectId: resumableRun.projectId,
status: "DEPLOYED",
imageReference: {
not: null,
},
},
orderBy: {
createdAt: "desc",
},
include: {
worker: {
include: {
tasks: true,
},
},
},
});
if (!deployment || !deployment.worker) {
logger.error("No matching deployment found for task run", {
queueMessage: message.data,
messageId: message.messageId,
});
await marqs?.acknowledgeMessage(message.messageId);
setTimeout(() => this.#doWork(), 100);
return;
}
if (!deployment.imageReference) {
logger.error("Deployment is missing an image reference", {
queueMessage: message.data,
messageId: message.messageId,
deployment: deployment.id,
});
await marqs?.acknowledgeMessage(message.messageId);
setTimeout(() => this.#doWork(), 100);
return;
}
const backgroundTask = deployment.worker.tasks.find(
(task) => task.slug === resumableRun.taskIdentifier
);
if (!backgroundTask) {
logger.warn("No matching background task found for task run", {
taskRun: resumableRun.id,
taskIdentifier: resumableRun.taskIdentifier,
deployment: deployment.id,
backgroundWorker: deployment.worker.id,
taskSlugs: deployment.worker.tasks.map((task) => task.slug),
});
await marqs?.acknowledgeMessage(message.messageId);
setTimeout(() => this.#doWork(), 100);
return;
}
const queue = await prisma.taskQueue.findUnique({
where: {
runtimeEnvironmentId_name: {
@@ -590,6 +532,43 @@ 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 (!latestCheckpoint) {
logger.error("No checkpoint found", {
queueMessage: message.data,
messageId: message.messageId,
resumableAttemptId: resumableAttempt.id,
});
await marqs?.acknowledgeMessage(message.messageId);
setTimeout(() => this.#doWork(), 100);
return;
}
await prisma.taskRunAttempt.update({
where: {
id: resumableAttempt.id,
},
data: {
status: "EXECUTING",
},
});
socketIo.providerNamespace.emit("RESTORE", {
version: "v1",
id: latestCheckpoint.id,
attemptId: latestCheckpoint.attemptId,
type: latestCheckpoint.type,
location: latestCheckpoint.location,
reason: latestCheckpoint.reason ?? undefined,
});
setTimeout(() => this.#doWork(), 100);
return;
}
const completions: TaskRunExecutionResult[] = [];
const executions: TaskRunExecution[] = [];
@@ -643,30 +622,108 @@ export class SharedQueueConsumer {
}
try {
// The attempt should still be running so we can broadcast to all coordinators to resume immediately
socketIo.coordinatorNamespace.emit("RESUME", {
version: "v1",
attemptId: resumableAttempt.id,
completions,
executions,
});
} 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);
} finally {
setTimeout(() => this.#doWork(), 100);
}
break;
}
// Resume after duration-based wait
case "RESUME_AFTER_DURATION": {
const resumableRun = await prisma.taskRun.findFirst({
where: {
id: message.messageId,
},
include: {
attempts: {
orderBy: {
createdAt: "desc",
},
take: 1,
include: {
checkpoints: {
take: 1,
orderBy: {
createdAt: "desc",
},
},
},
},
},
});
const resumableAttempt = resumableRun?.attempts[0];
if (!resumableAttempt) {
logger.error("Resumable attempt not found", {
queueMessage: message.data,
messageId: message.messageId,
});
await marqs?.acknowledgeMessage(message.messageId);
setTimeout(() => this.#doWork(), 100);
return;
}
if (resumableAttempt.status !== "PAUSED") {
logger.error("Attempt not paused", {
queueMessage: message.data,
messageId: message.messageId,
});
await marqs?.acknowledgeMessage(message.messageId);
setTimeout(() => this.#doWork(), 100);
return;
}
try {
// We need to restore the attempt from the latest checkpoint before we can resume
const latestCheckpoint = resumableAttempt.checkpoints[0];
if (!latestCheckpoint) {
// No checkpoint means the task should still be running
// We can broadcast to all coordinators to resume immediately
socketIo.coordinatorNamespace.emit("RESUME", {
version: "v1",
attemptId: resumableAttempt.id,
image: deployment.imageReference,
completions,
executions,
});
} else {
// There's a checkpoint we need to restore first
// TODO: Send RESUME message once the restored task has checked in
socketIo.providerNamespace.emit("RESTORE", {
version: "v1",
id: latestCheckpoint.id,
attemptId: latestCheckpoint.attemptId,
type: latestCheckpoint.type,
location: latestCheckpoint.location,
reason: latestCheckpoint.reason ?? undefined,
logger.error("No checkpoint found", {
queueMessage: message.data,
messageId: message.messageId,
resumableAttemptId: resumableAttempt.id,
});
await marqs?.acknowledgeMessage(message.messageId);
setTimeout(() => this.#doWork(), 100);
return;
}
await prisma.taskRunAttempt.update({
where: {
id: resumableAttempt.id,
},
data: {
status: "EXECUTING",
},
});
// The attempt will resume automatically after restore
socketIo.providerNamespace.emit("RESTORE", {
version: "v1",
id: latestCheckpoint.id,
attemptId: latestCheckpoint.attemptId,
type: latestCheckpoint.type,
location: latestCheckpoint.location,
reason: latestCheckpoint.reason ?? undefined,
});
} catch (e) {
if (e instanceof Error) {
this._currentSpan?.recordException(e);
@@ -182,7 +182,11 @@ export class CompleteAttemptService extends BaseService {
},
},
include: {
dependentTaskAttempt: true,
dependentTaskAttempt: {
include: {
taskRun: true,
},
},
items: {
include: {
taskRun: {
@@ -220,12 +224,27 @@ export class CompleteAttemptService extends BaseService {
return "ACKNOWLEDGED";
}
await marqs?.replaceMessage(finalizedBatchRun.dependentTaskAttempt.taskRunId, {
type: "RESUME",
completedAttemptIds: finalizedBatchRun.items.map(
(item) => item.taskRun.attempts[0]?.id
),
});
const dependentRun = finalizedBatchRun.dependentTaskAttempt.taskRun;
if (finalizedBatchRun.dependentTaskAttempt.status === "PAUSED") {
await marqs?.enqueueMessage(
environment,
dependentRun.queue,
dependentRun.id,
{
type: "RESUME",
completedAttemptIds: [taskRunAttempt.id],
},
dependentRun.concurrencyKey ?? undefined
);
} else {
await marqs?.replaceMessage(dependentRun.id, {
type: "RESUME",
completedAttemptIds: finalizedBatchRun.items.map(
(item) => item.taskRun.attempts[0]?.id
),
});
}
}
}
@@ -263,10 +282,23 @@ export class CompleteAttemptService extends BaseService {
return "FAILED";
}
await marqs?.replaceMessage(dependentRun.id, {
type: "RESUME",
completedAttemptIds: [taskRunAttempt.id],
});
if (dependency.dependentAttempt.status === "PAUSED") {
await marqs?.enqueueMessage(
environment,
dependentRun.queue,
dependentRun.id,
{
type: "RESUME",
completedAttemptIds: [taskRunAttempt.id],
},
dependentRun.concurrencyKey ?? undefined
);
} else {
await marqs?.replaceMessage(dependentRun.id, {
type: "RESUME",
completedAttemptIds: [taskRunAttempt.id],
});
}
}
}
@@ -34,12 +34,37 @@ export class CreateCheckpointService {
attemptId: attempt.id,
location: params.location,
type: params.docker ? "DOCKER" : "KUBERNETES",
reason: params.reason,
reason: params.reason.type,
},
});
// TODO: Can't heartbeat when checkpointed, so we ACK to prevent automatic requeue
// await marqs?.acknowledgeMessage(attempt.taskRunId);
await this.#prismaClient.taskRunAttempt.update({
where: {
id: params.attemptId,
},
data: {
status: "PAUSED",
},
});
switch (params.reason.type) {
case "WAIT_FOR_DURATION": {
await marqs?.replaceMessage(
attempt.taskRunId,
{ type: "RESUME_AFTER_DURATION" },
Date.now() + params.reason.ms
);
break;
}
case "WAIT_FOR_TASK":
case "WAIT_FOR_BATCH": {
await marqs?.acknowledgeMessage(attempt.taskRunId);
break;
}
default: {
break;
}
}
return checkpoint;
}
@@ -0,0 +1,178 @@
import {
CoordinatorToPlatformMessages,
InferSocketMessageSchema,
TaskRunExecution,
TaskRunExecutionResult,
} from "@trigger.dev/core/v3";
import { PrismaClient, prisma } from "~/db.server";
import { logger } from "~/services/logger.server";
import { marqs } from "../marqs.server";
import { socketIo } from "../handleSocketIo.server";
import { sharedQueueTasks } from "../marqs/sharedQueueConsumer.server";
export class ResumeAttemptService {
#prismaClient: PrismaClient;
constructor(prismaClient: PrismaClient = prisma) {
this.#prismaClient = prismaClient;
}
public async call(
params: InferSocketMessageSchema<typeof CoordinatorToPlatformMessages, "READY_FOR_RESUME">
): Promise<void> {
logger.debug(`ResumeAttemptService.call()`, params);
const attempt = await this.#prismaClient.taskRunAttempt.findUnique({
where: {
id: params.attemptId,
},
include: {
taskRun: true,
taskRunDependency: {
include: {
taskRun: {
include: {
attempts: {
orderBy: {
number: "desc",
},
take: 1,
select: {
id: true,
},
},
},
},
},
},
batchTaskRunDependency: {
include: {
items: {
include: {
taskRun: {
include: {
attempts: {
orderBy: {
number: "desc",
},
take: 1,
select: {
id: true,
},
},
},
},
},
},
},
},
},
});
if (!attempt) {
logger.error("Could not find attempt", { attemptId: params.attemptId });
return;
}
switch (params.type) {
case "WAIT_FOR_DURATION": {
// Nothing to do, but thanks for checking in!
break;
}
case "WAIT_FOR_TASK":
case "WAIT_FOR_BATCH": {
let completedAttemptIds: string[] = [];
if (attempt.taskRunDependency) {
const dependentAttempt = attempt.taskRunDependency.taskRun.attempts[0];
if (!dependentAttempt) {
logger.error("No dependent attempt", { attemptId: params.attemptId });
return;
}
completedAttemptIds = [dependentAttempt.id];
} else if (attempt.batchTaskRunDependency) {
const dependentBatchItems = attempt.batchTaskRunDependency.items;
if (!dependentBatchItems) {
logger.error("No dependent batch items", { attemptId: params.attemptId });
return;
}
completedAttemptIds = dependentBatchItems.map((item) => item.taskRun.attempts[0]?.id);
} else {
logger.error("No dependencies", { attemptId: params.attemptId });
return;
}
if (completedAttemptIds.length === 0) {
logger.error("No completed attempt IDs", { attemptId: params.attemptId });
return;
}
const completions: TaskRunExecutionResult[] = [];
const executions: TaskRunExecution[] = [];
for (const completedAttemptId of completedAttemptIds) {
const completedAttempt = await prisma.taskRunAttempt.findUnique({
where: {
id: completedAttemptId,
taskRun: {
lockedAt: {
not: null,
},
lockedById: {
not: null,
},
},
},
});
if (!completedAttempt) {
logger.error("Completed attempt not found", {
attemptId: params.attemptId,
completedAttemptId,
});
await marqs?.acknowledgeMessage(attempt.taskRunId);
return;
}
const completion = await sharedQueueTasks.getCompletionPayloadFromAttempt(
completedAttempt.id
);
if (!completion) {
await marqs?.acknowledgeMessage(attempt.taskRunId);
return;
}
completions.push(completion);
const executionPayload = await sharedQueueTasks.getExecutionPayloadFromAttempt(
completedAttempt.id,
false
);
if (!executionPayload) {
await marqs?.acknowledgeMessage(attempt.taskRunId);
return;
}
executions.push(executionPayload.execution);
}
socketIo.coordinatorNamespace.emit("RESUME", {
version: "v1",
attemptId: params.attemptId,
completions,
executions,
});
break;
}
default: {
break;
}
}
}
}
+2 -2
View File
@@ -587,12 +587,12 @@ async function compileProject(config: ResolvedConfig, options: DeployCommandOpti
);
const workerSetupPath = new URL(
importResolve("./workers/common/worker-setup.js", import.meta.url)
importResolve("./workers/prod/worker-setup.js", import.meta.url)
).href.replace("file://", "");
const workerContents = workerFacade
.replace("__TASKS__", createTaskFileImports(taskFiles))
.replace("__WORKER_SETUP__", `import { tracingSDK, sender } from "${workerSetupPath}";`);
.replace("__WORKER_SETUP__", `import { tracingSDK } from "${workerSetupPath}";`);
const result = await build({
stdin: {
+1 -1
View File
@@ -305,7 +305,7 @@ function useDev({
);
const workerSetupPath = new URL(
importResolve("./workers/common/worker-setup.js", import.meta.url)
importResolve("./workers/dev/worker-setup.js", import.meta.url)
).href.replace("file://", "");
const entryPointContents = workerFacade
@@ -1,7 +1,9 @@
import {
BackgroundWorkerProperties,
CreateBackgroundWorkerResponse,
ProdChildToWorkerMessages,
ProdTaskRunExecutionPayload,
ProdWorkerToChildMessages,
SemanticInternalAttributes,
TaskMetadataWithFilePath,
TaskRunBuiltInError,
@@ -9,11 +11,8 @@ import {
TaskRunExecution,
TaskRunExecutionPayload,
TaskRunExecutionResult,
ZodMessageHandler,
ZodMessageSender,
childToWorkerMessages,
ZodIpcConnection,
correctErrorStackTrace,
workerToChildMessages,
} from "@trigger.dev/core/v3";
import { Evt } from "evt";
import { ChildProcess, fork } from "node:child_process";
@@ -45,15 +44,14 @@ type BackgroundWorkerParams = {
export class ProdBackgroundWorker {
private _initialized: boolean = false;
private _handler = new ZodMessageHandler({
schema: childToWorkerMessages,
});
public onTaskHeartbeat: Evt<string> = new Evt();
public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt();
public onWaitForDuration: Evt<{ version?: "v1"; ms: number }> = new Evt();
public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt();
public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt();
public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
private _onClose: Evt<void> = new Evt();
@@ -116,20 +114,28 @@ export class ProdBackgroundWorker {
reject(new Error("Worker timed out"));
}, 10_000);
child.on("message", async (msg: any) => {
const message = this._handler.parseMessage(msg);
if (message.type === "TASKS_READY" && !resolved) {
clearTimeout(timeout);
resolved = true;
resolve(message.payload.tasks);
child.kill();
} else if (message.type === "UNCAUGHT_EXCEPTION") {
clearTimeout(timeout);
resolved = true;
reject(new UncaughtExceptionError(message.payload.error, message.payload.origin));
child.kill();
}
new ZodIpcConnection({
listenSchema: ProdChildToWorkerMessages,
emitSchema: ProdWorkerToChildMessages,
process: child,
handlers: {
TASKS_READY: async (message) => {
if (!resolved) {
clearTimeout(timeout);
resolved = true;
resolve(message.tasks);
child.kill();
}
},
UNCAUGHT_EXCEPTION: async (message) => {
if (!resolved) {
clearTimeout(timeout);
resolved = true;
reject(new UncaughtExceptionError(message.error, message.origin));
child.kill();
}
},
},
});
child.stdout?.on("data", (data) => {
@@ -170,6 +176,11 @@ export class ProdBackgroundWorker {
taskRunProcess.taskRunCompletedNotification(completion, execution);
}
}
async waitCompletedNotification() {
for (const taskRunProcess of this._taskRunProcesses.values()) {
taskRunProcess.waitCompletedNotification();
}
}
async #initializeTaskRunProcess(payload: ProdTaskRunExecutionPayload): Promise<TaskRunProcess> {
const metadata = this.getMetadata(
@@ -208,6 +219,10 @@ export class ProdBackgroundWorker {
this.onWaitForTask.post(message);
});
this.preCheckpointNotification.attach((message) => {
taskRunProcess.preCheckpointNotification.post(message);
});
await taskRunProcess.initialize();
this._taskRunProcesses.set(payload.execution.run.id, taskRunProcess);
@@ -291,11 +306,12 @@ export class ProdBackgroundWorker {
}
class TaskRunProcess {
private _handler = new ZodMessageHandler({
schema: childToWorkerMessages,
});
private _sender: ZodMessageSender<typeof workerToChildMessages>;
private _child: ChildProcess | undefined;
private _ipc?: ZodIpcConnection<
typeof ProdChildToWorkerMessages,
typeof ProdWorkerToChildMessages
>;
private _child?: ChildProcess;
private _attemptPromises: Map<
string,
{ resolver: (value: TaskRunExecutionResult) => void; rejecter: (err?: any) => void }
@@ -311,21 +327,14 @@ class TaskRunProcess {
public onWaitForDuration: Evt<{ version?: "v1"; ms: number }> = new Evt();
public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt();
public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
constructor(
private path: string,
private env: NodeJS.ProcessEnv,
private metadata: BackgroundWorkerProperties,
private worker: BackgroundWorkerParams
) {
this._sender = new ZodMessageSender({
schema: workerToChildMessages,
sender: async (message) => {
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
this._child?.send?.(message);
}
},
});
}
) {}
async initialize() {
this._child = fork(this.path, {
@@ -339,7 +348,57 @@ class TaskRunProcess {
},
});
this._child.on("message", this.#handleMessage.bind(this));
this._ipc = new ZodIpcConnection({
listenSchema: ProdChildToWorkerMessages,
emitSchema: ProdWorkerToChildMessages,
process: this._child,
handlers: {
TASK_RUN_COMPLETED: async (message) => {
const { result, execution } = message;
const promiseStatus = this._attemptStatuses.get(execution.attempt.id);
if (promiseStatus !== "PENDING") {
return;
}
this._attemptStatuses.set(execution.attempt.id, "RESOLVED");
const attemptPromise = this._attemptPromises.get(execution.attempt.id);
if (!attemptPromise) {
return;
}
const { resolver } = attemptPromise;
resolver(result);
},
READY_TO_DISPOSE: async (message) => {
this.#kill();
},
TASK_HEARTBEAT: async (message) => {
this.onTaskHeartbeat.post(message.id);
},
TASKS_READY: async (message) => {},
WAIT_FOR_BATCH: async (message) => {
this.onWaitForBatch.post(message);
},
WAIT_FOR_DURATION: async (message) => {
this.onWaitForDuration.post(message);
// The coordinator will let us know if a checkpoint is about to happen
// We then pass this back down to the runtime in the child process
const { willCheckpointAndRestore } = await this.preCheckpointNotification.waitFor();
return { willCheckpointAndRestore };
},
WAIT_FOR_TASK: async (message) => {
this.onWaitForTask.post(message);
},
},
});
this._child.on("exit", this.#handleExit.bind(this));
this._child.stdout?.on("data", this.#handleLog.bind(this));
this._child.stderr?.on("data", this.#handleStdErr.bind(this));
@@ -350,7 +409,7 @@ class TaskRunProcess {
return;
}
await this._sender.send("CLEANUP", {
await this._ipc?.send("CLEANUP", {
flush: true,
kill,
});
@@ -376,11 +435,13 @@ class TaskRunProcess {
this._currentExecution = execution;
await this._sender.send("EXECUTE_TASK_RUN", {
execution,
traceContext,
metadata: this.metadata,
});
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
await this._ipc?.send("EXECUTE_TASK_RUN", {
execution,
traceContext,
metadata: this.metadata,
});
}
const result = await promise;
@@ -394,67 +455,17 @@ class TaskRunProcess {
return;
}
this._sender.send("TASK_RUN_COMPLETED_NOTIFICATION", {
completion,
execution,
});
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
this._ipc?.send("TASK_RUN_COMPLETED_NOTIFICATION", {
completion,
execution,
});
}
}
async #handleMessage(msg: any) {
const message = this._handler.parseMessage(msg);
switch (message.type) {
case "TASK_RUN_COMPLETED": {
const { result, execution } = message.payload;
const promiseStatus = this._attemptStatuses.get(execution.attempt.id);
if (promiseStatus !== "PENDING") {
return;
}
this._attemptStatuses.set(execution.attempt.id, "RESOLVED");
const attemptPromise = this._attemptPromises.get(execution.attempt.id);
if (!attemptPromise) {
return;
}
const { resolver } = attemptPromise;
resolver(result);
break;
}
case "READY_TO_DISPOSE": {
this.#kill();
break;
}
case "TASK_HEARTBEAT": {
this.onTaskHeartbeat.post(message.payload.id);
break;
}
case "TASKS_READY": {
break;
}
case "WAIT_FOR_BATCH": {
this.onWaitForBatch.post(message.payload);
break;
}
case "WAIT_FOR_DURATION": {
this.onWaitForDuration.post(message.payload);
break;
}
case "WAIT_FOR_TASK": {
this.onWaitForTask.post(message.payload);
break;
}
waitCompletedNotification() {
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
this._ipc?.send("WAIT_COMPLETED_NOTIFICATION", {});
}
}
+111 -26
View File
@@ -30,6 +30,9 @@ class ProdWorker {
private executing = false;
private completed = false;
private paused = false;
private nextResumeAfter: "WAIT_FOR_DURATION" | "WAIT_FOR_TASK" | "WAIT_FOR_BATCH" | undefined;
#httpPort: number;
#backgroundWorker: ProdBackgroundWorker;
@@ -55,18 +58,76 @@ class ProdWorker {
},
contentHash: this.contentHash,
});
this.#backgroundWorker.onTaskHeartbeat.attach((attemptFriendlyId) => {
this.#coordinatorSocket.send("TASK_HEARTBEAT", { attemptFriendlyId });
});
this.#backgroundWorker.onWaitForBatch.attach((message) => {
this.#coordinatorSocket.send("WAIT_FOR_BATCH", message);
// TODO: Switch to .send() once coordinator uses zod handler for all messages
this.#coordinatorSocket.socket.emit("TASK_HEARTBEAT", { version: "v1", attemptFriendlyId });
});
this.#backgroundWorker.onWaitForDuration.attach(async (message) => {
const { success } = await this.#coordinatorSocket.sendWithAck("WAIT_FOR_DURATION", message);
logger.log("WAIT_FOR_DURATION", { success });
// 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 }
);
logger.log("WAIT_FOR_DURATION", { willCheckpointAndRestore });
this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore });
setTimeout(() => {
if (willCheckpointAndRestore) {
this.paused = true;
this.nextResumeAfter = "WAIT_FOR_DURATION";
}
// Forcing a reconnect will ensure the connection handler runs to trigger automatic resume
this.#coordinatorSocket.close();
this.#coordinatorSocket.connect();
}, 3_000);
});
this.#backgroundWorker.onWaitForTask.attach((message) => {
this.#coordinatorSocket.send("WAIT_FOR_TASK", message);
this.#backgroundWorker.onWaitForTask.attach(async (message) => {
// 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 }
);
logger.log("WAIT_FOR_TASK", { willCheckpointAndRestore });
this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore });
setTimeout(() => {
if (willCheckpointAndRestore) {
this.paused = true;
this.nextResumeAfter = "WAIT_FOR_TASK";
}
// Forcing a reconnect will ensure the connection handler runs to trigger automatic resume
this.#coordinatorSocket.close();
this.#coordinatorSocket.connect();
}, 3_000);
});
this.#backgroundWorker.onWaitForBatch.attach(async (message) => {
// 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 }
);
logger.log("WAIT_FOR_BATCH", { willCheckpointAndRestore });
this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore });
setTimeout(() => {
if (willCheckpointAndRestore) {
this.paused = true;
this.nextResumeAfter = "WAIT_FOR_BATCH";
}
// Forcing a reconnect will ensure the connection handler runs to trigger automatic resume
this.#coordinatorSocket.close();
this.#coordinatorSocket.connect();
}, 3_000);
});
this.#httpPort = port;
@@ -112,11 +173,13 @@ class ProdWorker {
this.#backgroundWorker.taskRunCompletedNotification(completion, execution);
}
},
RESUME_AFTER_DURATION: async (message) => {
this.#backgroundWorker.waitCompletedNotification();
},
EXECUTE_TASK_RUN: async (message) => {
if (this.executing || this.completed) {
return {
success: false,
};
logger.error("dropping execute request, already executing or completed");
return;
}
this.executing = true;
@@ -127,15 +190,13 @@ class ProdWorker {
this.completed = true;
this.executing = false;
setTimeout(() => {
process.exit(0);
}, 2000);
// TODO: replace ack with emit
return {
success: true,
await this.#coordinatorSocket.socket.emitWithAck("TASK_RUN_COMPLETED", {
version: "v1",
execution: message.executionPayload.execution,
completion,
};
});
process.exit(0);
},
},
onConnection: async (socket, handler, sender, logger) => {
@@ -209,12 +270,33 @@ class ProdWorker {
process.exit(1);
}, 200);
}
} else {
socket.emit("READY_FOR_EXECUTION", {
}
if (this.paused) {
if (!this.nextResumeAfter) {
return;
}
socket.emit("READY_FOR_RESUME", {
version: "v1",
attemptId: this.attemptId,
type: this.nextResumeAfter,
});
this.#backgroundWorker.waitCompletedNotification();
this.paused = false;
return;
}
if (this.executing) {
return;
}
socket.emit("READY_FOR_EXECUTION", {
version: "v1",
attemptId: this.attemptId,
});
},
});
@@ -247,11 +329,14 @@ class ProdWorker {
return reply.text(this.contentHash);
case "/wait":
const { success } = await this.#coordinatorSocket.sendWithAck("WAIT_FOR_DURATION", {
version: "v1",
ms: 60_000,
});
logger.log("WAIT_FOR_DURATION", { success });
const { willCheckpointAndRestore } = await this.#coordinatorSocket.sendWithAck(
"WAIT_FOR_DURATION",
{
version: "v1",
ms: 60_000,
}
);
logger.log("WAIT_FOR_DURATION", { willCheckpointAndRestore });
// this is required when C/Ring established connections
this.#coordinatorSocket.close();
return reply.text("sent WAIT");
@@ -1,4 +1,9 @@
import { type TracingSDK } from "@trigger.dev/core/v3";
import {
ProdChildToWorkerMessages,
ProdWorkerToChildMessages,
ZodIpcConnection,
type TracingSDK,
} from "@trigger.dev/core/v3";
import "source-map-support/register.js";
__WORKER_SETUP__;
@@ -20,16 +25,12 @@ import {
TaskRunExecution,
TaskRunExecutionRetry,
TriggerTracer,
ZodMessageHandler,
ZodMessageSender,
accessoryAttributes,
calculateNextRetryDelay,
childToWorkerMessages,
logger,
parseError,
runtime,
taskContextManager,
workerToChildMessages,
type BackgroundWorkerProperties,
} from "@trigger.dev/core/v3";
import * as packageJson from "../../../package.json";
@@ -40,12 +41,6 @@ import { TaskMetadataWithFunctions } from "../../types";
const tracer = new TriggerTracer({ tracer: otelTracer, logger: otelLogger });
const consoleInterceptor = new ConsoleInterceptor(otelLogger);
declare const sender: ZodMessageSender<typeof childToWorkerMessages>;
const prodRuntimeManager = new ProdRuntimeManager(sender);
runtime.setGlobalRuntimeManager(prodRuntimeManager);
const otelTaskLogger = new OtelTaskLogger({
logger: otelLogger,
tracer: tracer,
@@ -233,10 +228,12 @@ for (const task of tasks) {
let _execution: TaskRunExecution | undefined;
let _isRunning = false;
const handler = new ZodMessageHandler({
schema: workerToChildMessages,
messages: {
EXECUTE_TASK_RUN: async ({ execution, traceContext, metadata }) => {
const zodIpc = new ZodIpcConnection({
listenSchema: ProdWorkerToChildMessages,
emitSchema: ProdChildToWorkerMessages,
process,
handlers: {
EXECUTE_TASK_RUN: async ({ execution, traceContext, metadata }, sender) => {
if (_isRunning) {
console.error("Worker is already running a task");
@@ -308,7 +305,10 @@ const handler = new ZodMessageHandler({
TASK_RUN_COMPLETED_NOTIFICATION: async ({ completion, execution }) => {
prodRuntimeManager.resumeTask(completion, execution);
},
CLEANUP: async ({ flush, kill }) => {
WAIT_COMPLETED_NOTIFICATION: async () => {
prodRuntimeManager.resumeAfterRestore();
},
CLEANUP: async ({ flush, kill }, sender) => {
if (kill) {
await tracingSDK.flush();
// Now we need to exit the process
@@ -322,11 +322,11 @@ const handler = new ZodMessageHandler({
},
});
process.on("message", async (msg: any) => {
await handler.handleMessage(msg);
});
const prodRuntimeManager = new ProdRuntimeManager(zodIpc);
sender.send("TASKS_READY", { tasks: getTaskMetadata() }).catch((err) => {
runtime.setGlobalRuntimeManager(prodRuntimeManager);
zodIpc.send("TASKS_READY", { tasks: getTaskMetadata() }).catch((err) => {
console.error("Failed to send TASKS_READY message", err);
});
@@ -337,7 +337,7 @@ async function asyncHeartbeat(initialDelayInSeconds: number = 30, intervalInSeco
while (true) {
if (_isRunning && _execution) {
try {
await sender.send("TASK_HEARTBEAT", { id: _execution.attempt.id });
await zodIpc.send("TASK_HEARTBEAT", { id: _execution.attempt.id });
} catch (err) {
console.error("Failed to send HEARTBEAT message", err);
}
@@ -0,0 +1,36 @@
import { Resource } from "@opentelemetry/resources";
import { OpenAIInstrumentation } from "@traceloop/instrumentation-openai";
import {
SemanticInternalAttributes,
TracingDiagnosticLogLevel,
TracingSDK,
ZodMessageSender,
childToWorkerMessages,
} from "@trigger.dev/core/v3";
export const tracingSDK = new TracingSDK({
url: process.env.OTEL_EXPORTER_OTLP_ENDPOINT ?? "http://0.0.0.0:4318",
resource: new Resource({
[SemanticInternalAttributes.CLI_VERSION]: "3.0.0",
}),
instrumentations: [new OpenAIInstrumentation()],
diagLogLevel: (process.env.OTEL_LOG_LEVEL as TracingDiagnosticLogLevel) ?? "none",
});
process.on("uncaughtException", (error, origin) => {
process.send?.({
type: "EVENT",
message: {
type: "UNCAUGHT_EXCEPTION",
payload: {
error: {
name: error.name,
message: error.message,
stack: error.stack,
},
origin,
},
version: "v1",
},
});
});
+2 -1
View File
@@ -6,8 +6,9 @@ export default defineConfig({
tsconfig: "tsconfig.json",
splitting: false,
entry: [
"src/workers/dev/worker-setup.ts",
"src/workers/dev/worker-facade.ts",
"src/workers/common/worker-setup.ts",
"src/workers/prod/worker-setup.ts",
"src/workers/prod/worker-facade.ts",
"src/workers/prod/entry-point.ts",
],
+22 -3
View File
@@ -30,6 +30,7 @@ export interface TaskOperations {
type ProviderShellOptions = {
tasks: TaskOperations;
type: "docker" | "kubernetes";
host?: string;
port?: number;
};
@@ -111,7 +112,7 @@ export class ProviderShell implements Provider {
serverMessages: PlatformToProviderMessages,
authToken: PLATFORM_SECRET,
extraHeaders: {
"x-trigger-provider-type": "docker",
"x-trigger-provider-type": this.options.type,
},
handlers: {
DELETE: async (message) => {
@@ -139,10 +140,28 @@ export class ProviderShell implements Provider {
apiUrl: message.apiUrl,
});
} catch (error) {
logger.error("task index failed", error);
logger.error("index failed", error);
}
},
RESTORE: async (message) => {
if (message.type.toLowerCase() !== this.options.type.toLowerCase()) {
logger.error(
`restore failed: ${this.options.type} provider can't restore ${message.type} checkpoints`
);
return;
}
try {
await this.tasks.restore({
attemptId: message.attemptId,
checkpointRef: message.location,
// TODO
// machine: message.machine,
});
} catch (error) {
logger.error("restore failed", error);
}
},
RESTORE: async (message) => {},
},
});
+1
View File
@@ -5,6 +5,7 @@ export * from "./apiClient";
export * from "./zodMessageHandler";
export * from "./zodNamespace";
export * from "./zodSocket";
export * from "./zodIpc";
export * from "./errors";
export * from "./runtime-api";
export * from "./logger-api";
@@ -1,12 +1,13 @@
import {
BatchTaskRunExecutionResult,
ProdChildToWorkerMessages,
ProdWorkerToChildMessages,
TaskMetadataWithFilePath,
TaskRunContext,
TaskRunExecution,
TaskRunExecutionResult,
childToWorkerMessages,
} from "../schemas";
import { ZodMessageSender } from "../zodMessageHandler";
import { ZodIpcConnection } from "../zodIpc";
import { RuntimeManager } from "./manager";
export class ProdRuntimeManager implements RuntimeManager {
@@ -20,9 +21,16 @@ export class ProdRuntimeManager implements RuntimeManager {
{ resolve: (value: BatchTaskRunExecutionResult) => void; reject: (err?: any) => void }
> = new Map();
_waitForRestore: { resolve: (value?: any) => void; reject: (err?: any) => void } | undefined;
_tasks: Map<string, TaskMetadataWithFilePath> = new Map();
constructor(private sender: ZodMessageSender<typeof childToWorkerMessages>) {}
constructor(
private ipc: ZodIpcConnection<
typeof ProdWorkerToChildMessages,
typeof ProdChildToWorkerMessages
>
) {}
disable(): void {
// do nothing
@@ -39,15 +47,42 @@ export class ProdRuntimeManager implements RuntimeManager {
}
async waitForDuration(ms: number): Promise<void> {
if (ms > 30_000) {
// TODO: sender with ack support
await this.sender.send("WAIT_FOR_DURATION", { ms });
// TODO: resolve after resume signal instead
let timeout: NodeJS.Timeout | undefined;
const resolveAfterDuration = new Promise((resolve) => {
timeout = setTimeout(resolve, ms);
});
if (ms < 10_000) {
await resolveAfterDuration;
return;
}
return new Promise((resolve) => {
setTimeout(resolve, ms);
const waitForRestore = new Promise<TaskRunExecutionResult>((resolve, reject) => {
this._waitForRestore = { resolve, reject };
});
// There is a slight delay before actually checkpointing, so this has a chance to return
const { willCheckpointAndRestore } = await this.ipc.sendWithAck("WAIT_FOR_DURATION", { ms });
if (!willCheckpointAndRestore) {
await resolveAfterDuration;
return;
}
// Checkpointing should happen after this line
await waitForRestore;
clearTimeout(timeout);
}
resumeAfterRestore(): void {
if (!this._waitForRestore) {
return;
}
this._waitForRestore.resolve();
this._waitForRestore = undefined;
}
async waitUntil(date: Date): Promise<void> {
@@ -59,7 +94,7 @@ export class ProdRuntimeManager implements RuntimeManager {
this._taskWaits.set(params.id, { resolve, reject });
});
await this.sender.send("WAIT_FOR_TASK", {
await this.ipc.send("WAIT_FOR_TASK", {
id: params.id,
});
@@ -83,7 +118,7 @@ export class ProdRuntimeManager implements RuntimeManager {
})
);
await this.sender.send("WAIT_FOR_BATCH", {
await this.ipc.send("WAIT_FOR_BATCH", {
id: params.id,
runs: params.runs,
});
+91 -9
View File
@@ -187,6 +187,16 @@ export const TaskMetadataWithFilePath = TaskMetadata.extend({
export type TaskMetadataWithFilePath = z.infer<typeof TaskMetadataWithFilePath>;
export const UncaughtExceptionMessage = z.object({
version: z.literal("v1").default("v1"),
error: z.object({
name: z.string(),
message: z.string(),
stack: z.string().optional(),
}),
origin: z.enum(["uncaughtException", "unhandledRejection"]),
});
export const childToWorkerMessages = {
TASK_RUN_COMPLETED: z.object({
version: z.literal("v1").default("v1"),
@@ -215,13 +225,85 @@ export const childToWorkerMessages = {
id: z.string(),
runs: z.string().array(),
}),
UNCAUGHT_EXCEPTION: z.object({
version: z.literal("v1").default("v1"),
error: z.object({
name: z.string(),
message: z.string(),
stack: z.string().optional(),
}),
origin: z.enum(["uncaughtException", "unhandledRejection"]),
}),
UNCAUGHT_EXCEPTION: UncaughtExceptionMessage,
};
export const ProdChildToWorkerMessages = {
TASK_RUN_COMPLETED: {
message: z.object({
version: z.literal("v1").default("v1"),
execution: TaskRunExecution,
result: TaskRunExecutionResult,
}),
},
TASKS_READY: {
message: z.object({
version: z.literal("v1").default("v1"),
tasks: TaskMetadataWithFilePath.array(),
}),
},
TASK_HEARTBEAT: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
}),
},
READY_TO_DISPOSE: {
message: z.undefined(),
},
WAIT_FOR_DURATION: {
message: z.object({
version: z.literal("v1").default("v1"),
ms: z.number(),
}),
callback: z.object({
willCheckpointAndRestore: z.boolean(),
}),
},
WAIT_FOR_TASK: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
}),
},
WAIT_FOR_BATCH: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
runs: z.string().array(),
}),
},
UNCAUGHT_EXCEPTION: {
message: UncaughtExceptionMessage,
},
};
export const ProdWorkerToChildMessages = {
EXECUTE_TASK_RUN: {
message: z.object({
version: z.literal("v1").default("v1"),
execution: TaskRunExecution,
traceContext: z.record(z.unknown()),
metadata: BackgroundWorkerProperties,
}),
},
TASK_RUN_COMPLETED_NOTIFICATION: {
message: z.object({
version: z.literal("v1").default("v1"),
completion: TaskRunExecutionResult,
execution: TaskRunExecution,
}),
},
CLEANUP: {
message: z.object({
version: z.literal("v1").default("v1"),
flush: z.boolean().default(false),
kill: z.boolean().default(true),
}),
},
WAIT_COMPLETED_NOTIFICATION: {
message: z.object({
version: z.literal("v1").default("v1"),
}),
},
};
+68 -31
View File
@@ -145,6 +145,13 @@ export const CoordinatorToPlatformMessages = {
}),
]),
},
READY_FOR_RESUME: {
message: z.object({
version: z.literal("v1").default("v1"),
attemptId: z.string(),
type: z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]),
}),
},
TASK_RUN_COMPLETED: {
message: z.object({
version: z.literal("v1").default("v1"),
@@ -164,7 +171,20 @@ export const CoordinatorToPlatformMessages = {
attemptId: z.string(),
docker: z.boolean(),
location: z.string(),
reason: z.string().optional(),
reason: z.discriminatedUnion("type", [
z.object({
type: z.literal("WAIT_FOR_DURATION"),
ms: z.number(),
}),
z.object({
type: z.literal("WAIT_FOR_BATCH"),
id: z.string(),
}),
z.object({
type: z.literal("WAIT_FOR_TASK"),
id: z.string(),
}),
]),
}),
},
INDEXING_FAILED: {
@@ -185,11 +205,16 @@ export const PlatformToCoordinatorMessages = {
message: z.object({
version: z.literal("v1").default("v1"),
attemptId: z.string(),
image: z.string(),
completions: TaskRunExecutionResult.array(),
executions: TaskRunExecution.array(),
}),
},
RESUME_AFTER_DURATION: {
message: z.object({
version: z.literal("v1").default("v1"),
attemptId: z.string(),
}),
},
};
export const ClientToSharedQueueMessages = {
@@ -260,37 +285,53 @@ export const ProdWorkerToCoordinatorMessages = {
attemptId: z.string(),
}),
},
READY_FOR_RESUME: {
message: z.object({
version: z.literal("v1").default("v1"),
attemptId: z.string(),
type: z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]),
}),
},
TASK_HEARTBEAT: {
message: z.object({
version: z.literal("v1").default("v1"),
attemptFriendlyId: z.string(),
}),
},
TASK_RUN_COMPLETED: {
message: z.object({
version: z.literal("v1").default("v1"),
execution: ProdTaskRunExecution,
completion: TaskRunExecutionResult,
}),
callback: z.void(),
},
WAIT_FOR_DURATION: {
message: z.object({
version: z.literal("v1").default("v1"),
ms: z.number(),
}),
callback: z.object({
willCheckpointAndRestore: z.boolean(),
}),
},
WAIT_FOR_TASK: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
}),
callback: z.object({
willCheckpointAndRestore: z.boolean(),
}),
},
WAIT_FOR_BATCH: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
runs: z.string().array(),
}),
},
WAIT_FOR_DURATION: {
message: z.object({
version: z.literal("v1").default("v1"),
ms: z.number(),
}),
callback: z.discriminatedUnion("success", [
z.object({
success: z.literal(false),
}),
z.object({
success: z.literal(true),
}),
]),
},
WAIT_FOR_TASK: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
callback: z.object({
willCheckpointAndRestore: z.boolean(),
}),
},
INDEXING_FAILED: {
@@ -311,25 +352,21 @@ export const CoordinatorToProdWorkerMessages = {
message: z.object({
version: z.literal("v1").default("v1"),
attemptId: z.string(),
image: z.string(),
completions: TaskRunExecutionResult.array(),
executions: TaskRunExecution.array(),
}),
},
RESUME_AFTER_DURATION: {
message: z.object({
version: z.literal("v1").default("v1"),
attemptId: z.string(),
}),
},
EXECUTE_TASK_RUN: {
message: z.object({
version: z.literal("v1").default("v1"),
executionPayload: ProdTaskRunExecutionPayload,
}),
callback: z.discriminatedUnion("success", [
z.object({
success: z.literal(false),
}),
z.object({
success: z.literal(true),
completion: TaskRunExecutionResult,
}),
]),
},
};
+328
View File
@@ -0,0 +1,328 @@
import { randomUUID } from "crypto";
import {
GetSocketCallbackSchema,
GetSocketMessageSchema,
GetSocketMessagesWithCallback,
GetSocketMessagesWithoutCallback,
MessagesFromSocketCatalog,
SocketMessageHasCallback,
ZodSocketMessageCatalogSchema,
} from "./zodSocket";
import { z } from "zod";
interface ZodIpcMessageSender<TEmitCatalog extends ZodSocketMessageCatalogSchema> {
send<K extends GetSocketMessagesWithoutCallback<TEmitCatalog>>(
type: K,
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>
): Promise<void>;
sendWithAck<K extends GetSocketMessagesWithCallback<TEmitCatalog>>(
type: K,
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>
): Promise<z.infer<GetSocketCallbackSchema<TEmitCatalog, K>>>;
}
type ZodIpcMessageHandlers<
TListenCatalog extends ZodSocketMessageCatalogSchema,
TEmitCatalog extends ZodSocketMessageCatalogSchema,
> = Partial<{
[K in keyof TListenCatalog]: (
payload: z.infer<GetSocketMessageSchema<TListenCatalog, K>>,
sender: ZodIpcMessageSender<TEmitCatalog>
) => Promise<
SocketMessageHasCallback<TListenCatalog, K> extends true
? z.input<GetSocketCallbackSchema<TListenCatalog, K>>
: void
>;
}>;
const messageSchema = z.object({
version: z.literal("v1").default("v1"),
type: z.string(),
payload: z.unknown(),
});
type ZodIpcMessageHandlerOptions<
TListenCatalog extends ZodSocketMessageCatalogSchema,
TEmitCatalog extends ZodSocketMessageCatalogSchema,
> = {
schema: TListenCatalog;
handlers?: ZodIpcMessageHandlers<TListenCatalog, TEmitCatalog>;
sender: ZodIpcMessageSender<TEmitCatalog>;
};
class ZodIpcMessageHandler<
TListenCatalog extends ZodSocketMessageCatalogSchema,
TEmitCatalog extends ZodSocketMessageCatalogSchema,
> {
#schema: TListenCatalog;
#handlers: ZodIpcMessageHandlers<TListenCatalog, TEmitCatalog> | undefined;
#sender: ZodIpcMessageSender<TEmitCatalog>;
constructor(options: ZodIpcMessageHandlerOptions<TListenCatalog, TEmitCatalog>) {
this.#schema = options.schema;
this.#handlers = options.handlers;
this.#sender = options.sender;
}
public async handleMessage(message: unknown) {
const parsedMessage = this.parseMessage(message);
if (!this.#handlers) {
throw new Error("No handlers provided");
}
const handler = this.#handlers[parsedMessage.type];
if (!handler) {
// console.error(`No handler for message type: ${String(parsedMessage.type)}`);
return;
}
const ack = await handler(parsedMessage.payload, this.#sender);
return ack;
}
public parseMessage(message: unknown): MessagesFromSocketCatalog<TListenCatalog> {
const parsedMessage = messageSchema.safeParse(message);
if (!parsedMessage.success) {
throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`);
}
const schema = this.#schema[parsedMessage.data.type]["message"];
if (!schema) {
throw new Error(`Unknown message type: ${parsedMessage.data.type}`);
}
const parsedPayload = schema.safeParse(parsedMessage.data.payload);
if (!parsedPayload.success) {
throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
}
return {
type: parsedMessage.data.type,
payload: parsedPayload.data,
};
}
}
const Packet = z.discriminatedUnion("type", [
z.object({
type: z.literal("CONNECT"),
sessionId: z.string().optional(),
}),
z.object({
type: z.literal("ACK"),
message: z.any(),
id: z.number(),
}),
z.object({
type: z.literal("EVENT"),
message: z.any(),
id: z.number().optional(),
}),
]);
type Packet = z.infer<typeof Packet>;
interface ZodIpcConnectionOptions<
TListenCatalog extends ZodSocketMessageCatalogSchema,
TEmitCatalog extends ZodSocketMessageCatalogSchema,
> {
listenSchema: TListenCatalog;
emitSchema: TEmitCatalog;
process: {
send?: (message: any) => any;
on?: (event: "message", listener: (message: any) => void) => void;
};
handlers?: ZodIpcMessageHandlers<TListenCatalog, TEmitCatalog>;
}
export class ZodIpcConnection<
TListenCatalog extends ZodSocketMessageCatalogSchema,
TEmitCatalog extends ZodSocketMessageCatalogSchema,
> {
#sessionId?: string;
#messageCounter: number = 0;
#handler: ZodIpcMessageHandler<TListenCatalog, TEmitCatalog>;
#acks: Map<
number,
{
resolve: (value: unknown) => void;
reject: (reason?: any) => void;
timeout: NodeJS.Timeout;
}
> = new Map();
constructor(private opts: ZodIpcConnectionOptions<TListenCatalog, TEmitCatalog>) {
this.#handler = new ZodIpcMessageHandler({
schema: opts.listenSchema,
handlers: opts.handlers,
sender: {
send: this.send.bind(this),
sendWithAck: this.sendWithAck.bind(this),
},
});
this.#registerHandlers();
// this.connect();
}
async #registerHandlers() {
if (!this.opts.process.on) {
return;
}
this.opts.process.on("message", async (message) => {
this.#handlePacket(message);
});
}
async connect() {
this.#sendPacket({ type: "CONNECT" });
}
async #handlePacket(packet: Packet): Promise<void> {
const parsedPacket = Packet.safeParse(packet);
if (!parsedPacket.success) {
// console.error("dropping invalid packet", packet);
return;
}
// console.log("<-", packet);
switch (parsedPacket.data.type) {
case "ACK": {
// Check our list of ACKs and resolve with the message
const ack = this.#acks.get(parsedPacket.data.id);
if (!ack) {
return;
}
clearTimeout(ack.timeout);
ack.resolve(parsedPacket.data.message);
break;
}
case "CONNECT": {
if (!parsedPacket.data.sessionId) {
// This is a client trying to connect, so we generate and send back a session ID
const id = randomUUID();
await this.#sendPacket({ type: "CONNECT", sessionId: id });
return;
}
// This is a server replying to our connect message
if (this.#sessionId) {
// We're already connected
return;
}
this.#sessionId = parsedPacket.data.sessionId;
break;
}
case "EVENT": {
const result = await this.#handler.handleMessage(parsedPacket.data.message);
if (typeof parsedPacket.data.id === "undefined") {
return;
}
// There's an ID so we should ACK
await this.#sendPacket({
type: "ACK",
id: parsedPacket.data.id,
message: result,
});
break;
}
default: {
break;
}
}
}
async #sendPacket(packet: Packet) {
// console.log("->", packet);
await this.opts.process.send?.(packet);
}
async send<K extends GetSocketMessagesWithoutCallback<TEmitCatalog>>(
type: K,
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>
): Promise<void> {
const schema = this.opts.emitSchema[type]["message"];
if (!schema) {
throw new Error(`Unknown message type: ${type as string}`);
}
const parsedPayload = schema.safeParse(payload);
if (!parsedPayload.success) {
throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
}
await this.#sendPacket({
type: "EVENT",
message: {
type,
payload,
version: "v1",
},
});
}
public async sendWithAck<K extends GetSocketMessagesWithCallback<TEmitCatalog>>(
type: K,
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>,
timeoutInMs?: number
): Promise<z.infer<GetSocketCallbackSchema<TEmitCatalog, K>>> {
const currentId = this.#messageCounter++;
return new Promise(async (resolve, reject) => {
// Timeout if the ACK takes too long to get back to us
const timeout = setTimeout(() => {
reject("timeout");
}, timeoutInMs ?? 2000);
this.#acks.set(currentId, { resolve, reject, timeout });
const schema = this.opts.emitSchema[type]["message"];
if (!schema) {
clearTimeout(timeout);
return reject(`Unknown message type: ${type as string}`);
}
const parsedPayload = schema.safeParse(payload);
if (!parsedPayload.success) {
clearTimeout(timeout);
return reject(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
}
await this.#sendPacket({
type: "EVENT",
message: {
type,
payload,
version: "v1",
},
id: currentId,
});
});
}
}
+1
View File
@@ -96,6 +96,7 @@ export class ZodNamespace<
this.namespace = this.io.of(opts.name);
// FIXME: There's a bug here, this sender should not accept Socket schemas with callbacks
this.sender = new ZodMessageSender({
schema: opts.serverMessages,
sender: async (message) => {
+8 -6
View File
@@ -22,7 +22,7 @@ export type ZodMessageCatalogToSocketIoEvents<TCatalog extends ZodSocketMessageC
: (message: z.infer<GetSocketMessageSchema<TCatalog, K>>) => void;
};
type GetSocketMessageSchema<
export type GetSocketMessageSchema<
TRPCCatalog extends ZodSocketMessageCatalogSchema,
TMessageType extends keyof TRPCCatalog,
> = TRPCCatalog[TMessageType]["message"];
@@ -32,7 +32,7 @@ export type InferSocketMessageSchema<
TMessageType extends keyof TRPCCatalog,
> = z.infer<GetSocketMessageSchema<TRPCCatalog, TMessageType>>;
type GetSocketCallbackSchema<
export type GetSocketCallbackSchema<
TRPCCatalog extends ZodSocketMessageCatalogSchema,
TMessageType extends keyof TRPCCatalog,
> = TRPCCatalog[TMessageType] extends { callback: any }
@@ -44,7 +44,7 @@ export type InferSocketCallbackSchema<
TMessageType extends keyof TRPCCatalog,
> = z.infer<GetSocketCallbackSchema<TRPCCatalog, TMessageType>>;
type SocketMessageHasCallback<
export type SocketMessageHasCallback<
TRPCCatalog extends ZodSocketMessageCatalogSchema,
TMessageType extends keyof TRPCCatalog,
> = GetSocketCallbackSchema<TRPCCatalog, TMessageType> extends never ? false : true;
@@ -74,7 +74,7 @@ type MessageFromSocketSchema<
payload: z.input<GetSocketMessageSchema<TMessageCatalog, K>>;
};
type MessagesFromSocketCatalog<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
export type MessagesFromSocketCatalog<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
[K in keyof TMessageCatalog]: MessageFromSocketSchema<K, TMessageCatalog>;
}[keyof TMessageCatalog];
@@ -173,13 +173,15 @@ export type ZodSocketMessageSenderOptions<TMessageCatalog extends ZodSocketMessa
socket: ZodSocket<any, TMessageCatalog>;
};
type GetSocketMessagesWithCallback<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
export type GetSocketMessagesWithCallback<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
[K in keyof TMessageCatalog]: SocketMessageHasCallback<TMessageCatalog, K> extends true
? K
: never;
}[keyof TMessageCatalog];
type GetSocketMessagesWithoutCallback<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
export type GetSocketMessagesWithoutCallback<
TMessageCatalog extends ZodSocketMessageCatalogSchema,
> = {
[K in keyof TMessageCatalog]: SocketMessageHasCallback<TMessageCatalog, K> extends true
? never
: K;