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

checkpoint readiness and cancel signals

Browse Source
This commit is contained in:
nicktrn
2024-03-22 08:41:31 +00:00
parent 6eae7270c2
commit b017af4e6d
6 changed files with 165 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/coordinator/src/index.ts
+110 -27
View File
@@ -52,6 +52,7 @@ class Checkpointer {
#dockerMode = !process.env.KUBERNETES_PORT;
#logger = new SimpleLogger("[checkptr]");
#abortControllers = new Map<string, AbortController>();
constructor(private opts = { forceSimulate: false }) {}
@@ -138,7 +139,6 @@ class Checkpointer {
const result = await this.#checkpointAndPush(opts);
const end = performance.now();
logger.log(`checkpointAndPush() end`, {
start,
end,
@@ -150,6 +150,22 @@ class Checkpointer {
return result;
}
isCheckpointing(runId: string) {
return this.#abortControllers.has(runId);
}
cancelCheckpoint(runId: string) {
const controller = this.#abortControllers.get(runId);
if (!controller) {
logger.debug("Nothing to cancel", { runId });
return;
}
controller.abort("cancelCheckpointing()");
this.#abortControllers.delete(runId);
}
async #checkpointAndPush(opts: CheckpointAndPushOptions): Promise<CheckpointData | undefined> {
await this.initialize();
@@ -158,15 +174,25 @@ class Checkpointer {
return;
}
const shortCode = nanoid(8);
const imageRef = this.#getImageRef(opts.projectRef, opts.deploymentVersion, shortCode);
const exportLocation = this.#getExportLocation(
opts.projectRef,
opts.deploymentVersion,
shortCode
);
if (this.#abortControllers.has(opts.runId)) {
logger.error("Checkpoint procedure already in progress", { opts });
return;
}
const controller = new AbortController();
this.#abortControllers.set(opts.runId, controller);
const $$ = $({ signal: controller.signal });
try {
const shortCode = nanoid(8);
const imageRef = this.#getImageRef(opts.projectRef, opts.deploymentVersion, shortCode);
const exportLocation = this.#getExportLocation(
opts.projectRef,
opts.deploymentVersion,
shortCode
);
this.#logger.log("Checkpointing:", { opts });
const containterName = this.#getRunContainerName(opts.runId);
@@ -176,15 +202,15 @@ class Checkpointer {
try {
if (this.opts.forceSimulate || !this.#canCheckpoint) {
this.#logger.log("Simulating checkpoint");
this.#logger.debug(await $`docker pause ${containterName}`);
this.#logger.debug(await $$`docker pause ${containterName}`);
} else {
if (opts.leaveRunning) {
this.#logger.debug(
await $`docker checkpoint create --leave-running ${containterName} ${exportLocation}`
await $$`docker checkpoint create --leave-running ${containterName} ${exportLocation}`
);
} else {
this.#logger.debug(
await $`docker checkpoint create ${containterName} ${exportLocation}`
await $$`docker checkpoint create ${containterName} ${exportLocation}`
);
}
}
@@ -211,36 +237,40 @@ class Checkpointer {
const containerId = this.#logger.debug(
// @ts-expect-error
await $`crictl ps`
.pipeStdout($({ stdin: "pipe" })`grep ${containterName}`)
.pipeStdout($({ stdin: "pipe" })`cut -f1 ${"-d "}`)
await $$`crictl ps`
.pipeStdout($$({ stdin: "pipe" })`grep ${containterName}`)
.pipeStdout($$({ stdin: "pipe" })`cut -f1 ${"-d "}`)
);
if (!containerId.stdout) {
throw new Error("could not find container id");
}
this.#logger.debug(await $`crictl checkpoint --export=${exportLocation} ${containerId}`);
this.#logger.debug(await $$`crictl checkpoint --export=${exportLocation} ${containerId}`);
// Create image from checkpoint
const container = this.#logger.debug(await $`buildah from scratch`);
this.#logger.debug(await $`buildah add ${container} ${exportLocation} /`);
const container = this.#logger.debug(await $$`buildah from scratch`);
this.#logger.debug(await $$`buildah add ${container} ${exportLocation} /`);
this.#logger.debug(
await $`buildah config --annotation=io.kubernetes.cri-o.annotations.checkpoint.name=counter ${container}`
await $$`buildah config --annotation=io.kubernetes.cri-o.annotations.checkpoint.name=counter ${container}`
);
this.#logger.debug(await $`buildah commit ${container} ${imageRef}`);
this.#logger.debug(await $`buildah rm ${container}`);
this.#logger.debug(await $$`buildah commit ${container} ${imageRef}`);
this.#logger.debug(await $$`buildah rm ${container}`);
// Push checkpoint image
this.#logger.debug(await $`buildah push --tls-verify=${REGISTRY_TLS_VERIFY} ${imageRef}`);
this.#logger.debug(await $$`buildah push --tls-verify=${REGISTRY_TLS_VERIFY} ${imageRef}`);
this.#logger.log("Checkpointed and pushed image to:", { location: imageRef });
try {
await $`rm ${exportLocation}`;
await $$`rm ${exportLocation}`;
this.#logger.log("Deleted checkpoint archive", { exportLocation });
// Disabled for now as this will increase restore time by having to pull the image again
// await $`buildah rmi ${imageRef}`;
// this.#logger.log("Deleted checkpoint image", { imageRef });
} catch (error) {
this.#logger.error("Failed to delete checkpoint archive", { exportLocation });
this.#logger.error("Failed during checkpoint cleanup", { exportLocation });
this.#logger.debug(error);
}
@@ -251,6 +281,8 @@ class Checkpointer {
} catch (error) {
this.#logger.error("checkpoint failed", { options: opts, error });
return;
} finally {
this.#abortControllers.delete(opts.runId);
}
}
@@ -273,6 +305,11 @@ class TaskCoordinator {
typeof PlatformToCoordinatorMessages
>;
#checkpointableTasks = new Map<
string,
{ resolve: (value: void) => void; reject: (err?: any) => void }
>();
constructor(
private port: number,
private host = "0.0.0.0"
@@ -529,9 +566,42 @@ class TaskCoordinator {
}
});
socket.on("READY_FOR_CHECKPOINT", async (message) => {
logger.log("[READY_FOR_CHECKPOINT]", message);
const checkpointable = this.#checkpointableTasks.get(socket.data.runId);
if (!checkpointable) {
logger.error("No checkpoint scheduled", { runId: socket.data.runId });
return;
}
checkpointable.resolve();
});
socket.on("CANCEL_CHECKPOINT", async (message) => {
logger.log("[CANCEL_CHECKPOINT]", message);
const checkpointWait = this.#checkpointableTasks.get(socket.data.runId);
if (checkpointWait) {
// Stop waiting for task to reach checkpointable state
checkpointWait.reject("Checkpoint cancelled");
}
// Cancel checkpointing procedure
this.#checkpointer.cancelCheckpoint(socket.data.runId);
});
socket.on("WAIT_FOR_DURATION", async (message, callback) => {
logger.log("[WAIT_FOR_DURATION]", message);
if (this.#checkpointableTasks.has(socket.data.runId)) {
logger.error("Checkpoint already in progress", { runId: socket.data.runId });
callback({ willCheckpointAndRestore: false });
return;
}
const { canCheckpoint, willSimulate } = await this.#checkpointer.initialize();
const willCheckpointAndRestore = canCheckpoint || willSimulate;
@@ -542,9 +612,22 @@ class TaskCoordinator {
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 isCheckpointable = new Promise((resolve, reject) => {
// We set a reasonable timeout to prevent waiting forever
setTimeout(reject, 10_000);
this.#checkpointableTasks.set(socket.data.runId, { resolve, reject });
});
try {
await isCheckpointable;
} catch (error) {
logger.error("Error while waiting for checkpointable state", { error });
// TODO: We may want to cancel the task as it's unlikely to recover
return;
} finally {
this.#checkpointableTasks.delete(socket.data.runId);
}
const checkpoint = await this.#checkpointer.checkpointAndPush({
runId: socket.data.runId,
@@ -553,8 +636,8 @@ class TaskCoordinator {
});
if (!checkpoint) {
// The task container will keep running until the wait duration has elapsed
logger.error("Failed to checkpoint", { runId: socket.data.runId });
// TODO: We have to let the worker know about failures so it can use its own timer
return;
}
packages/cli-v3/src/workers/prod/backgroundWorker.ts
+19
View File
@@ -60,6 +60,8 @@ export class ProdBackgroundWorker {
public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt();
public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>();
public onCancelCheckpoint = Evt.create<{ version?: "v1" }>();
private _onClose: Evt<void> = new Evt();
@@ -220,6 +222,15 @@ export class ProdBackgroundWorker {
this.onWaitForTask.post(message);
});
taskRunProcess.onReadyForCheckpoint.attach((message) => {
this.onReadyForCheckpoint.post(message);
});
taskRunProcess.onCancelCheckpoint.attach((message) => {
this.onCancelCheckpoint.post(message);
});
// Notify down the chain
this.preCheckpointNotification.attach((message) => {
taskRunProcess.preCheckpointNotification.post(message);
});
@@ -346,6 +357,8 @@ class TaskRunProcess {
public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt();
public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>();
public onCancelCheckpoint = Evt.create<{ version?: "v1" }>();
constructor(
private path: string,
@@ -414,6 +427,12 @@ class TaskRunProcess {
WAIT_FOR_TASK: async (message) => {
this.onWaitForTask.post(message);
},
READY_FOR_CHECKPOINT: async (message) => {
this.onReadyForCheckpoint.post(message);
},
CANCEL_CHECKPOINT: async (message) => {
this.onCancelCheckpoint.post(message);
},
},
});
packages/cli-v3/src/workers/prod/entry-point.ts
+9
View File
@@ -68,6 +68,14 @@ class ProdWorker {
this.#coordinatorSocket.socket.emit("TASK_HEARTBEAT", { version: "v1", attemptFriendlyId });
});
this.#backgroundWorker.onReadyForCheckpoint.attach(async (message) => {
this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" });
});
this.#backgroundWorker.onCancelCheckpoint.attach(async (message) => {
this.#coordinatorSocket.socket.emit("CANCEL_CHECKPOINT", { version: "v1" });
});
this.#backgroundWorker.onWaitForDuration.attach(async (message) => {
// TODO: Switch to .send() once coordinator uses zod handler for all messages
const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck(
@@ -376,6 +384,7 @@ class ProdWorker {
{
version: "v1",
ms: 60_000,
now: Date.now(),
}
);
logger.log("WAIT_FOR_DURATION", { willCheckpointAndRestore });
packages/core/src/v3/runtime/prodRuntimeManager.ts
+7 -3
View File
@@ -64,7 +64,6 @@ export class ProdRuntimeManager implements RuntimeManager {
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,
now,
@@ -75,9 +74,14 @@ export class ProdRuntimeManager implements RuntimeManager {
return;
}
// Checkpointing should happen after this line
this.ipc.send("READY_FOR_CHECKPOINT", {});
// Don't wait for checkpoint beyond the requested wait duration
await Promise.race([waitForRestore, resolveAfterDuration]);
// The coordinator can then cancel any in-progress checkpoints
this.ipc.send("CANCEL_CHECKPOINT", {});
await waitForRestore;
clearTimeout(timeout);
}
packages/core/src/v3/schemas/messages.ts
+10
View File
@@ -258,6 +258,16 @@ export const ProdChildToWorkerMessages = {
READY_TO_DISPOSE: {
message: z.undefined(),
},
READY_FOR_CHECKPOINT: {
message: z.object({
version: z.literal("v1").default("v1"),
}),
},
CANCEL_CHECKPOINT: {
message: z.object({
version: z.literal("v1").default("v1"),
}),
},
WAIT_FOR_DURATION: {
message: z.object({
version: z.literal("v1").default("v1"),
packages/core/src/v3/schemas/schemas.ts
+10
View File
@@ -333,6 +333,16 @@ export const ProdWorkerToCoordinatorMessages = {
type: z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]),
}),
},
READY_FOR_CHECKPOINT: {
message: z.object({
version: z.literal("v1").default("v1"),
}),
},
CANCEL_CHECKPOINT: {
message: z.object({
version: z.literal("v1").default("v1"),
}),
},
TASK_HEARTBEAT: {
message: z.object({
version: z.literal("v1").default("v1"),