diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index 2bcb86163..23217c757 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -52,6 +52,7 @@ class Checkpointer { #dockerMode = !process.env.KUBERNETES_PORT; #logger = new SimpleLogger("[checkptr]"); + #abortControllers = new Map(); 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 { 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; } diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index 4593432a6..caad5638a 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -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 = 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); + }, }, }); diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index f518e6c26..cda670ec0 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -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 }); diff --git a/packages/core/src/v3/runtime/prodRuntimeManager.ts b/packages/core/src/v3/runtime/prodRuntimeManager.ts index fe8a1450d..c6bfb5b5f 100644 --- a/packages/core/src/v3/runtime/prodRuntimeManager.ts +++ b/packages/core/src/v3/runtime/prodRuntimeManager.ts @@ -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); } diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index 3023ed91b..b6ae8d3bc 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -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"), diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index 7bb72ba05..85685200d 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -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"),