rework suspend restore

Author: nicktrn

apps/supervisor/src/env.ts

@@ -45,6 +45,12 @@ const Env = z.object({
   // Used by the resource monitor
   OVERRIDE_CPU_TOTAL: z.coerce.number().optional(),
   OVERRIDE_MEMORY_TOTAL_GB: z.coerce.number().optional(),
+
+  // Kubernetes specific settings
+  KUBERNETES_FORCE_ENABLED: BoolEnv.default(false),
+  KUBERNETES_NAMESPACE: z.string().default("default"),
+  EPHEMERAL_STORAGE_SIZE_LIMIT: z.string().default("10Gi"),
+  EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"),
 });
 
 export const env = Env.parse(stdEnv);

apps/supervisor/src/index.ts

@@ -13,7 +13,11 @@ import {
 } from "./resourceMonitor.js";
 import { KubernetesWorkloadManager } from "./workloadManager/kubernetes.js";
 import { DockerWorkloadManager } from "./workloadManager/docker.js";
-import { HttpServer, CheckpointClient } from "@trigger.dev/core/v3/serverOnly";
+import {
+  HttpServer,
+  CheckpointClient,
+  isKubernetesEnvironment,
+} from "@trigger.dev/core/v3/serverOnly";
 import { createK8sApi, RUNTIME_ENV } from "./clients/kubernetes.js";
 
 class ManagedSupervisor {
@@ -25,7 +29,7 @@ class ManagedSupervisor {
   private readonly resourceMonitor: ResourceMonitor;
   private readonly checkpointClient?: CheckpointClient;
 
-  private readonly isKubernetes = RUNTIME_ENV === "kubernetes";
+  private readonly isKubernetes = isKubernetesEnvironment(env.KUBERNETES_FORCE_ENABLED);
   private readonly warmStartUrl = env.TRIGGER_WARM_START_URL;
 
   constructor() {
@@ -94,6 +98,7 @@ class ManagedSupervisor {
       this.checkpointClient = new CheckpointClient({
         apiUrl: new URL(env.TRIGGER_CHECKPOINT_URL),
         workerClient: this.workerSession.httpClient,
+        orchestrator: this.isKubernetes ? "KUBERNETES" : "DOCKER",
       });
     }
 
@@ -127,7 +132,9 @@ class ManagedSupervisor {
         return;
       }
 
-      if (message.checkpoint) {
+      const { checkpoint, ...rest } = message;
+
+      if (checkpoint) {
         this.logger.log("[ManagedWorker] Restoring run", { runId: message.run.id });
 
         if (!this.checkpointClient) {
@@ -139,7 +146,10 @@ class ManagedSupervisor {
           const didRestore = await this.checkpointClient.restoreRun({
             runFriendlyId: message.run.friendlyId,
             snapshotFriendlyId: message.snapshot.friendlyId,
-            checkpoint: message.checkpoint,
+            body: {
+              ...rest,
+              checkpoint,
+            },
           });
 
           if (didRestore) {

apps/supervisor/src/workloadManager/kubernetes.ts

@@ -9,17 +9,14 @@ import type { EnvironmentType, MachinePreset } from "@trigger.dev/core/v3";
 import { env } from "../env.js";
 import { type K8sApi, createK8sApi, type k8s } from "../clients/kubernetes.js";
 
-const POD_EPHEMERAL_STORAGE_SIZE_LIMIT = process.env.POD_EPHEMERAL_STORAGE_SIZE_LIMIT || "10Gi";
-const POD_EPHEMERAL_STORAGE_SIZE_REQUEST = process.env.POD_EPHEMERAL_STORAGE_SIZE_REQUEST || "2Gi";
-
 type ResourceQuantities = {
   [K in "cpu" | "memory" | "ephemeral-storage"]?: string;
 };
 
 export class KubernetesWorkloadManager implements WorkloadManager {
   private readonly logger = new SimpleStructuredLogger("kubernetes-workload-provider");
   private k8s: K8sApi;
-  private namespace = "default";
+  private namespace = env.KUBERNETES_NAMESPACE;
 
   constructor(private opts: WorkloadManagerOptions) {
     this.k8s = createK8sApi();
@@ -205,13 +202,13 @@ export class KubernetesWorkloadManager implements WorkloadManager {
 
   get #defaultResourceRequests(): ResourceQuantities {
     return {
-      "ephemeral-storage": POD_EPHEMERAL_STORAGE_SIZE_REQUEST,
+      "ephemeral-storage": env.EPHEMERAL_STORAGE_SIZE_REQUEST,
     };
   }
 
   get #defaultResourceLimits(): ResourceQuantities {
     return {
-      "ephemeral-storage": POD_EPHEMERAL_STORAGE_SIZE_LIMIT,
+      "ephemeral-storage": env.EPHEMERAL_STORAGE_SIZE_LIMIT,
     };
   }
 

apps/supervisor/src/workloadServer/index.ts

@@ -94,8 +94,8 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
     this.websocketServer = this.createWebsocketServer();
   }
 
-  private runnerIdFromRequest(req: IncomingMessage): string | undefined {
-    const value = req.headers[WORKLOAD_HEADERS.RUNNER_ID];
+  private headerValueFromRequest(req: IncomingMessage, headerName: string): string | undefined {
+    const value = req.headers[headerName];
 
     if (Array.isArray(value)) {
       return value[0];
@@ -104,6 +104,22 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
     return value;
   }
 
+  private runnerIdFromRequest(req: IncomingMessage): string | undefined {
+    return this.headerValueFromRequest(req, WORKLOAD_HEADERS.RUNNER_ID);
+  }
+
+  private deploymentIdFromRequest(req: IncomingMessage): string | undefined {
+    return this.headerValueFromRequest(req, WORKLOAD_HEADERS.DEPLOYMENT_ID);
+  }
+
+  private deploymentVersionFromRequest(req: IncomingMessage): string | undefined {
+    return this.headerValueFromRequest(req, WORKLOAD_HEADERS.DEPLOYMENT_VERSION);
+  }
+
+  private projectRefFromRequest(req: IncomingMessage): string | undefined {
+    return this.headerValueFromRequest(req, WORKLOAD_HEADERS.PROJECT_REF);
+  }
+
   private createHttpServer({ host, port }: { host: string; port: number }) {
     return new HttpServer({ port, host })
       .route(
@@ -213,8 +229,10 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
             }
 
             const runnerId = this.runnerIdFromRequest(req);
+            const deploymentVersion = this.deploymentVersionFromRequest(req);
+            const projectRef = this.projectRefFromRequest(req);
 
-            if (!runnerId) {
+            if (!runnerId || !deploymentVersion || !projectRef) {
               console.error("Invalid headers for suspend request", {
                 ...params,
                 headers: req.headers,
@@ -241,8 +259,13 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
             const suspendResult = await this.checkpointClient.suspendRun({
               runFriendlyId: params.runFriendlyId,
               snapshotFriendlyId: params.snapshotFriendlyId,
-              containerId: runnerId,
-              runnerId,
+              body: {
+                runnerId,
+                runId: params.runFriendlyId,
+                snapshotId: params.snapshotFriendlyId,
+                projectRef,
+                deploymentVersion,
+              },
             });
 
             if (!suspendResult) {

packages/cli-v3/src/entryPoints/managed-run-controller.ts

@@ -149,8 +149,10 @@ class ManagedRunController {
 
     this.httpClient = new WorkloadHttpClient({
       workerApiUrl: this.workerApiUrl,
-      deploymentId: env.TRIGGER_DEPLOYMENT_ID,
       runnerId: env.TRIGGER_RUNNER_ID,
+      deploymentId: env.TRIGGER_DEPLOYMENT_ID,
+      deploymentVersion: env.TRIGGER_DEPLOYMENT_VERSION,
+      projectRef: env.TRIGGER_PROJECT_REF,
     });
 
     if (env.TRIGGER_WARM_START_URL) {

packages/core/package.json

@@ -191,6 +191,7 @@
     "nanoid": "^3.3.4",
     "socket.io": "4.7.4",
     "socket.io-client": "4.7.5",
+    "std-env": "^3.8.1",
     "superjson": "^2.2.1",
     "tinyexec": "^0.3.2",
     "zod": "3.23.8",

packages/core/src/utils.ts

@@ -2,11 +2,11 @@ export function assertExhaustive(x: never): never {
   throw new Error("Unexpected object: " + x);
 }
 
-export async function tryCatch<T, E = Error>(promise: Promise<T>): Promise<[T, null] | [null, E]> {
+export async function tryCatch<T, E = Error>(promise: Promise<T>): Promise<[null, T] | [E, null]> {
   try {
     const data = await promise;
-    return [data, null];
+    return [null, data];
   } catch (error) {
-    return [null, error as E];
+    return [error as E, null];
   }
 }

packages/core/src/v3/runEngineWorker/consts.ts

@@ -8,4 +8,6 @@ export const WORKER_HEADERS = {
 export const WORKLOAD_HEADERS = {
   DEPLOYMENT_ID: "x-trigger-workload-deployment-id",
   RUNNER_ID: "x-trigger-workload-runner-id",
+  DEPLOYMENT_VERSION: "x-trigger-workload-deployment-version",
+  PROJECT_REF: "x-trigger-workload-project-ref",
 };

packages/core/src/v3/runEngineWorker/workload/types.ts

@@ -1,5 +1,7 @@
 export type WorkloadClientCommonOptions = {
   workerApiUrl: string;
-  deploymentId: string;
   runnerId: string;
+  deploymentId: string;
+  deploymentVersion: string;
+  projectRef: string;
 };

packages/core/src/v3/runEngineWorker/workload/util.ts

@@ -8,5 +8,7 @@ export function getDefaultWorkloadHeaders(
   return createHeaders({
     [WORKLOAD_HEADERS.DEPLOYMENT_ID]: options.deploymentId,
     [WORKLOAD_HEADERS.RUNNER_ID]: options.runnerId,
+    [WORKLOAD_HEADERS.DEPLOYMENT_VERSION]: options.deploymentVersion,
+    [WORKLOAD_HEADERS.PROJECT_REF]: options.projectRef,
   });
 }

packages/core/src/v3/schemas/checkpoints.ts

@@ -1,4 +1,4 @@
-import { CheckpointType } from "./runEngine.js";
+import { CheckpointType, DequeuedMessage } from "./runEngine.js";
 import z from "zod";
 
 const CallbackUrl = z
@@ -8,20 +8,12 @@ const CallbackUrl = z
 
 export const CheckpointServiceSuspendRequestBody = z.object({
   type: CheckpointType,
-  containerId: z.string(),
-  simulate: z.boolean().optional(),
-  leaveRunning: z.boolean().optional(),
+  runId: z.string(),
+  snapshotId: z.string(),
+  runnerId: z.string(),
+  projectRef: z.string(),
+  deploymentVersion: z.string(),
   reason: z.string().optional(),
-  callbacks: z
-    .object({
-      /** These headers will sent to all callbacks */
-      headers: z.record(z.string()).optional(),
-      /** This will be hit before suspending the container. Suspension will proceed unless we receive an error response. */
-      preSuspend: CallbackUrl.optional(),
-      /** This will be hit after suspending or failure to suspend the container */
-      completion: CallbackUrl.optional(),
-    })
-    .optional(),
 });
 
 export type CheckpointServiceSuspendRequestBody = z.infer<
@@ -39,16 +31,7 @@ export type CheckpointServiceSuspendResponseBody = z.infer<
   typeof CheckpointServiceSuspendResponseBody
 >;
 
-export const CheckpointServiceRestoreRequestBody = z.discriminatedUnion("type", [
-  z.object({
-    type: z.literal(CheckpointType.Enum.DOCKER),
-    containerId: z.string(),
-  }),
-  z.object({
-    type: z.literal(CheckpointType.Enum.KUBERNETES),
-    containerId: z.string(),
-  }),
-]);
+export const CheckpointServiceRestoreRequestBody = DequeuedMessage.required({ checkpoint: true });
 
 export type CheckpointServiceRestoreRequestBody = z.infer<
   typeof CheckpointServiceRestoreRequestBody

packages/core/src/v3/schemas/runEngine.ts

@@ -123,52 +123,6 @@ export const ExecutionResult = z.object({
 
 export type ExecutionResult = z.infer<typeof ExecutionResult>;
 
-/** This is sent to a Worker when a run is dequeued (a new run or continuing run) */
-export const DequeuedMessage = z.object({
-  version: z.literal("1"),
-  snapshot: ExecutionSnapshot,
-  dequeuedAt: z.coerce.date(),
-  image: z.string().optional(),
-  checkpoint: z
-    .object({
-      id: z.string(),
-      type: z.string(),
-      location: z.string(),
-      reason: z.string().nullish(),
-    })
-    .optional(),
-  completedWaitpoints: z.array(CompletedWaitpoint),
-  backgroundWorker: z.object({
-    id: z.string(),
-    friendlyId: z.string(),
-    version: z.string(),
-  }),
-  deployment: z.object({
-    id: z.string().optional(),
-    friendlyId: z.string().optional(),
-  }),
-  run: z.object({
-    id: z.string(),
-    friendlyId: z.string(),
-    isTest: z.boolean(),
-    machine: MachinePreset,
-    attemptNumber: z.number(),
-    masterQueue: z.string(),
-    traceContext: z.record(z.unknown()),
-  }),
-  environment: z.object({
-    id: z.string(),
-    type: EnvironmentType,
-  }),
-  organization: z.object({
-    id: z.string(),
-  }),
-  project: z.object({
-    id: z.string(),
-  }),
-});
-export type DequeuedMessage = z.infer<typeof DequeuedMessage>;
-
 /** The response to the Worker when starting an attempt */
 export const StartRunAttemptResult = ExecutionResult.and(
   z.object({
@@ -256,3 +210,51 @@ export const MachineResources = z.object({
   memory: z.number(),
 });
 export type MachineResources = z.infer<typeof MachineResources>;
+
+export const DequeueMessageCheckpoint = z.object({
+  id: z.string(),
+  type: CheckpointType,
+  location: z.string(),
+  imageRef: z.string(),
+  reason: z.string().nullish(),
+});
+export type DequeueMessageCheckpoint = z.infer<typeof DequeueMessageCheckpoint>;
+
+/** This is sent to a Worker when a run is dequeued (a new run or continuing run) */
+export const DequeuedMessage = z.object({
+  version: z.literal("1"),
+  snapshot: ExecutionSnapshot,
+  dequeuedAt: z.coerce.date(),
+  image: z.string().optional(),
+  checkpoint: DequeueMessageCheckpoint.optional(),
+  completedWaitpoints: z.array(CompletedWaitpoint),
+  backgroundWorker: z.object({
+    id: z.string(),
+    friendlyId: z.string(),
+    version: z.string(),
+  }),
+  deployment: z.object({
+    id: z.string().optional(),
+    friendlyId: z.string().optional(),
+  }),
+  run: z.object({
+    id: z.string(),
+    friendlyId: z.string(),
+    isTest: z.boolean(),
+    machine: MachinePreset,
+    attemptNumber: z.number(),
+    masterQueue: z.string(),
+    traceContext: z.record(z.unknown()),
+  }),
+  environment: z.object({
+    id: z.string(),
+    type: EnvironmentType,
+  }),
+  organization: z.object({
+    id: z.string(),
+  }),
+  project: z.object({
+    id: z.string(),
+  }),
+});
+export type DequeuedMessage = z.infer<typeof DequeuedMessage>;