diff --git a/apps/supervisor/src/workloadServer/index.ts b/apps/supervisor/src/workloadServer/index.ts index e14060e609..72d3b0e12b 100644 --- a/apps/supervisor/src/workloadServer/index.ts +++ b/apps/supervisor/src/workloadServer/index.ts @@ -60,6 +60,13 @@ const checkpointDeleteRequests = new Counter({ registers: [register], }); +const checkpointCancelRequests = new Counter({ + name: "checkpoint_cancel_requests_total", + help: "Checkpoint cancel requests attempted when a run continues, by outcome", + labelNames: ["result"], + registers: [register], +}); + const WorkloadActionParams = z.object({ runFriendlyId: z.string(), snapshotFriendlyId: z.string(), @@ -283,6 +290,30 @@ export class WorkloadServer extends EventEmitter { checkpointDeleteRequests.inc({ result: "sent" }); } + private async cancelCheckpointsAfterReply(runFriendlyId: string): Promise { + if (!this.checkpointClient) { + checkpointCancelRequests.inc({ result: "no_client" }); + return; + } + + if (this.snapshotService) { + checkpointCancelRequests.inc({ result: "not_applicable" }); + return; + } + + const [error, accepted] = await tryCatch( + this.checkpointClient.cancelCheckpoints({ runFriendlyId }) + ); + + if (error || !accepted) { + checkpointCancelRequests.inc({ result: "http_error" }); + this.logger.error("Failed to request checkpoint cancel", { runFriendlyId, error }); + return; + } + + checkpointCancelRequests.inc({ result: "sent" }); + } + /** * Sets common route meta on the wide-event state from URL params. */ @@ -643,6 +674,8 @@ export class WorkloadServer extends EventEmitter { } reply.json(continuationResult.data as WorkloadContinueRunExecutionResponseBody); + + await this.cancelCheckpointsAfterReply(params.runFriendlyId); } ), } diff --git a/packages/core/src/v3/serverOnly/checkpointClient.ts b/packages/core/src/v3/serverOnly/checkpointClient.ts index 936557c0ac..f89c1d1f43 100644 --- a/packages/core/src/v3/serverOnly/checkpointClient.ts +++ b/packages/core/src/v3/serverOnly/checkpointClient.ts @@ -13,6 +13,8 @@ export type CheckpointClientOptions = { orchestrator: CheckpointType; }; +const CANCEL_TIMEOUT_MS = 5_000; + export class CheckpointClient { private readonly logger = new SimpleStructuredLogger("checkpoint-client"); @@ -157,4 +159,21 @@ export class CheckpointClient { return true; } + + async cancelCheckpoints({ runFriendlyId }: { runFriendlyId: string }): Promise { + const res = await fetch( + new URL(`/api/v1/runs/${runFriendlyId}/checkpoints/cancel`, this.opts.apiUrl), + { method: "POST", signal: AbortSignal.timeout(CANCEL_TIMEOUT_MS) } + ); + + if (!res.ok) { + this.logger.error("[CheckpointClient] Cancel checkpoints request failed", { + runFriendlyId, + status: res.status, + }); + return false; + } + + return true; + } }