Replace the Go compute runtime with Rust, with run-CRD suspend and resume

Supervisor: with the Runner CRD backend on the microvm runtime, suspends
are requested on the Runner and completed from its status, a completed
suspend is still delivered after a supervisor restart, and a restored
run is placed on the node that holds its snapshot. A restore whose
Runner has already ended is reported as failed.
Mono-RevId: 2707a91994f6ab5fb2f7f9d25e575eb21ff78a41

Author: 1stvamp

apps/supervisor/src/env.test.ts

@@ -62,3 +62,16 @@ describe("Env superRefine - backpressure source awareness", () => {
     ).toThrow();
   });
 });
+
+describe("Env superRefine - compute snapshots", () => {
+  it("needs the metadata URL and workload API domain only for the gateway", () => {
+    expect(() => Env.parse({ ...base, COMPUTE_SNAPSHOTS_ENABLED: "true" })).not.toThrow();
+    expect(() =>
+      Env.parse({
+        ...base,
+        COMPUTE_SNAPSHOTS_ENABLED: "true",
+        COMPUTE_GATEWAY_URL: "http://gateway:8080",
+      })
+    ).toThrow(/TRIGGER_METADATA_URL[\s\S]*TRIGGER_WORKLOAD_API_DOMAIN/);
+  });
+});

apps/supervisor/src/env.ts

@@ -179,17 +179,14 @@ export const Env = z
 
     // Kubernetes settings
     KUBERNETES_FORCE_ENABLED: BoolEnv.default(false),
-    // Create a Runner and let the operator build the pod, instead of building
-    // the pod here. The two are mutually exclusive by construction: whichever
-    // one runs is the only thing that creates a workload for a cold start.
+    // Create a Runner for the operator to build the pod from, instead of building it
+    // here. Mutually exclusive with the pod backend, so one thing creates each cold start.
     KUBERNETES_RUN_CRD_ENABLED: BoolEnv.default(false),
-    // Which isolation lane a Runner asks for. Cell-wide rather than per run,
-    // because a cell's node pools decide what it can serve and nothing on a
-    // dequeued message can express the choice. Ignored unless the run-crd
-    // backend is the one running: the pod backends have no guest lane.
-    //
-    // Not to be confused with the task runtime (node-24, bun), which rides on
-    // the same Runner as taskRuntime and says which interpreter the task needs.
+    // Isolation lane for Runners, cell-wide because a cell's node pools decide what it
+    // can serve. Only read by the run-crd backend; the pod backends have no guest lane.
+    // Under microvm, checkpoints resume as restoring Runners, since only that lane's node
+    // runtime can restore, and COMPUTE_SNAPSHOTS_ENABLED sends suspends to the operator.
+    // Unrelated to the task runtime (node-24, bun), which is the Runner's taskRuntime.
     KUBERNETES_RUNNER_RUNTIME: z.enum(["container", "microvm"]).default("container"),
     KUBERNETES_NAMESPACE: z.string().default("default"),
     KUBERNETES_WORKER_NODETYPE_LABEL: NodeLabelValue.default("v4-worker"),
@@ -236,13 +233,13 @@ export const Env = z
     KUBERNETES_RUNNER_SECURITY_CONTEXT: z.enum(["off", "baseline", "restricted"]).default("off"),
     KUBERNETES_RUNNER_RUN_AS_USER: z.coerce.number().int().min(1).default(1000),
 
-    // Pod DNS config — override the cluster default ndots to `KUBERNETES_POD_DNS_NDOTS`.
+    // Pod DNS config: override the cluster default ndots to `KUBERNETES_POD_DNS_NDOTS`.
     // Default k8s ndots is 5: any name with fewer than 5 dots (e.g. `api.example.com`, 2 dots) is first walked
     // through every entry in the cluster search list (`<ns>.svc.cluster.local`, `svc.cluster.local`, `cluster.local`)
     // before being tried as-is, turning one resolution into 4+ CoreDNS queries (×2 with A+AAAA).
     // Overriding the default can be useful to cut CoreDNS query amplification for external domains.
     // Note: before enabling, make sure no code path relies on search-list expansion for names with dots ≥ the value
-    // set here — those names will now hit their as-is form first and could resolve externally before falling back.
+    // set here, since those names will now hit their as-is form first and could resolve externally before falling back.
     KUBERNETES_POD_DNS_NDOTS_OVERRIDE_ENABLED: BoolEnv.default(false),
     KUBERNETES_POD_DNS_NDOTS: z.coerce.number().int().min(1).max(15).default(2),
     // Large machine affinity settings - large-* presets prefer a dedicated pool
@@ -357,14 +354,20 @@ export const Env = z
         }
       }
     }
-    if (data.COMPUTE_SNAPSHOTS_ENABLED && !data.TRIGGER_METADATA_URL) {
+    // Only the gateway reads these for a snapshot: it builds its callback URL
+    // from the domain and hands the metadata URL to the instance.
+    if (data.COMPUTE_SNAPSHOTS_ENABLED && data.COMPUTE_GATEWAY_URL && !data.TRIGGER_METADATA_URL) {
       ctx.addIssue({
         code: z.ZodIssueCode.custom,
         message: "TRIGGER_METADATA_URL is required when COMPUTE_SNAPSHOTS_ENABLED is true",
         path: ["TRIGGER_METADATA_URL"],
       });
     }
-    if (data.COMPUTE_SNAPSHOTS_ENABLED && !data.TRIGGER_WORKLOAD_API_DOMAIN) {
+    if (
+      data.COMPUTE_SNAPSHOTS_ENABLED &&
+      data.COMPUTE_GATEWAY_URL &&
+      !data.TRIGGER_WORKLOAD_API_DOMAIN
+    ) {
       ctx.addIssue({
         code: z.ZodIssueCode.custom,
         message: "TRIGGER_WORKLOAD_API_DOMAIN is required when COMPUTE_SNAPSHOTS_ENABLED is true",

apps/supervisor/src/index.ts

@@ -3,7 +3,11 @@ import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLog
 import { formatLogLine, startTelnetLogServer } from "@trigger.dev/core/v3/telnetLogServer";
 import { env } from "./env.js";
 import { WorkloadServer } from "./workloadServer/index.js";
-import type { WorkloadManagerOptions, WorkloadManager } from "./workloadManager/types.js";
+import type {
+  WorkloadManagerCreateOptions,
+  WorkloadManagerOptions,
+  WorkloadManager,
+} from "./workloadManager/types.js";
 import Docker from "dockerode";
 import { z } from "zod";
 import { type DequeuedMessage } from "@trigger.dev/core/v3";
@@ -84,6 +88,7 @@ class ManagedSupervisor {
   private readonly workloadManager: WorkloadManager;
   private readonly workloadManagerBackend: "compute" | "kubernetes" | "run-crd" | "docker";
   private readonly computeManager?: ComputeWorkloadManager;
+  private readonly runCrdManager?: RunCrdWorkloadManager;
   private readonly logger = new SimpleStructuredLogger("managed-supervisor");
   private readonly resourceMonitor: ResourceMonitor;
   private readonly checkpointClient?: CheckpointClient;
@@ -192,11 +197,18 @@ class ManagedSupervisor {
       this.workloadManager = computeManager;
       this.workloadManagerBackend = "compute";
     } else if (this.isKubernetes && env.KUBERNETES_RUN_CRD_ENABLED) {
-      this.workloadManager = new RunCrdWorkloadManager({
+      const runCrdManager = new RunCrdWorkloadManager({
         ...workloadManagerOptions,
         namespace: env.KUBERNETES_NAMESPACE,
         runtime: env.KUBERNETES_RUNNER_RUNTIME,
+        snapshots: {
+          enabled: env.COMPUTE_SNAPSHOTS_ENABLED,
+          delayMs: env.COMPUTE_SNAPSHOT_DELAY_MS,
+          dispatchLimit: env.COMPUTE_SNAPSHOT_DISPATCH_LIMIT,
+        },
       });
+      this.runCrdManager = runCrdManager;
+      this.workloadManager = runCrdManager;
       this.workloadManagerBackend = "run-crd";
     } else if (this.isKubernetes) {
       this.workloadManager = new KubernetesWorkloadManager(workloadManagerOptions);
@@ -521,6 +533,11 @@ class ManagedSupervisor {
                 return;
               }
 
+              if (this.runCrdManager?.restores(checkpoint)) {
+                await this.restoreRunner(this.runCrdManager, message, checkpoint);
+                return;
+              }
+
               if (!this.checkpointClient) {
                 this.logger.error("No checkpoint client", { runId: message.run.id });
                 return;
@@ -608,6 +625,7 @@ class ManagedSupervisor {
       workerClient: this.workerSession.httpClient,
       checkpointClient: this.checkpointClient,
       computeManager: this.computeManager,
+      runnerSnapshotter: this.runCrdManager,
       tracing: this.tracing,
       snapshotCallbackSecret: workerToken,
       wideEventOpts: this.wideEventOpts,
@@ -630,57 +648,86 @@ class ManagedSupervisor {
     this.workerSession.unsubscribeFromRunNotifications([run.friendlyId]);
   }
 
-  private async createWorkload(message: DequeuedMessage, timings: WarmStartTimings) {
-    const createStart = performance.now();
-    try {
-      if (!message.deployment.friendlyId) {
-        // mostly a type guard, deployments always exists for deployed environments
-        // a proper fix would be to use a discriminated union schema to differentiate between dequeued runs in dev and in deployed environments.
-        throw new Error("Deployment is missing");
-      }
+  private async createOptionsFor(
+    message: DequeuedMessage,
+    timings?: WarmStartTimings
+  ): Promise<WorkloadManagerCreateOptions> {
+    if (!message.deployment.friendlyId) {
+      // mostly a type guard, deployments always exists for deployed environments
+      // a proper fix would be to use a discriminated union schema to differentiate between dequeued runs in dev and in deployed environments.
+      throw new Error("Deployment is missing");
+    }
 
-      if (!message.image) {
-        // same type-guard situation as deployment above
-        throw new Error("Image is missing");
-      }
+    if (!message.image) {
+      // same type-guard situation as deployment above
+      throw new Error("Image is missing");
+    }
 
-      const deploymentToken = await mintDeploymentToken({
-        deployment: message.deployment.friendlyId,
-        deployment_version: message.backgroundWorker.version,
-        environment_id: message.environment.id,
-        environment_type: message.environment.type,
-        org_id: message.organization.id,
-        project_id: message.project.id,
-      });
+    const deploymentToken = await mintDeploymentToken({
+      deployment: message.deployment.friendlyId,
+      deployment_version: message.backgroundWorker.version,
+      environment_id: message.environment.id,
+      environment_type: message.environment.type,
+      org_id: message.organization.id,
+      project_id: message.project.id,
+    });
 
-      await this.workloadManager.create({
-        dequeuedAt: message.dequeuedAt,
-        dequeueResponseMs: timings.dequeueResponseMs,
-        pollingIntervalMs: timings.pollingIntervalMs,
-        warmStartCheckMs: timings.warmStartCheckMs,
-        envId: message.environment.id,
-        envType: message.environment.type,
-        image: message.image,
-        machine: message.run.machine,
-        orgId: message.organization.id,
-        projectId: message.project.id,
-        deploymentFriendlyId: message.deployment.friendlyId,
-        deploymentVersion: message.backgroundWorker.version,
-        runtime: message.backgroundWorker.runtime,
-        deploymentToken,
-        runId: message.run.id,
-        runFriendlyId: message.run.friendlyId,
-        version: message.version,
-        nextAttemptNumber: message.run.attemptNumber,
-        snapshotId: message.snapshot.id,
-        snapshotFriendlyId: message.snapshot.friendlyId,
-        // Carry the run's storage route to the cold-start pod so its start request echoes it back.
-        snapshotRoute: message.snapshotRoute,
-        placementTags: message.placementTags,
-        traceContext: message.run.traceContext,
-        annotations: message.run.annotations,
-        hasPrivateLink: message.organization.hasPrivateLink,
-      });
+    return {
+      dequeuedAt: message.dequeuedAt,
+      dequeueResponseMs: timings?.dequeueResponseMs,
+      pollingIntervalMs: timings?.pollingIntervalMs,
+      warmStartCheckMs: timings?.warmStartCheckMs,
+      envId: message.environment.id,
+      envType: message.environment.type,
+      image: message.image,
+      machine: message.run.machine,
+      orgId: message.organization.id,
+      projectId: message.project.id,
+      deploymentFriendlyId: message.deployment.friendlyId,
+      deploymentVersion: message.backgroundWorker.version,
+      runtime: message.backgroundWorker.runtime,
+      deploymentToken,
+      runId: message.run.id,
+      runFriendlyId: message.run.friendlyId,
+      version: message.version,
+      nextAttemptNumber: message.run.attemptNumber,
+      snapshotId: message.snapshot.id,
+      snapshotFriendlyId: message.snapshot.friendlyId,
+      // Carry the run's storage route to the runner pod so its start request echoes it back.
+      snapshotRoute: message.snapshotRoute,
+      placementTags: message.placementTags,
+      traceContext: message.run.traceContext,
+      annotations: message.run.annotations,
+      hasPrivateLink: message.organization.hasPrivateLink,
+    };
+  }
+
+  private async restoreRunner(
+    manager: RunCrdWorkloadManager,
+    message: DequeuedMessage,
+    checkpoint: { id: string; location: string }
+  ) {
+    const restoreStart = performance.now();
+    try {
+      await manager.restore(await this.createOptionsFor(message), checkpoint);
+      recordPhaseSince("restore", restoreStart, undefined);
+      setExtra(fromContext(), "did_restore", true);
+      this.logger.debug("Runner restore created", { runId: message.run.id });
+    } catch (error) {
+      recordPhaseSince(
+        "restore",
+        restoreStart,
+        error instanceof Error ? error : new Error(String(error))
+      );
+      setExtra(fromContext(), "did_restore", false);
+      this.logger.error("Failed to restore run (run-crd)", { runId: message.run.id, error });
+    }
+  }
+
+  private async createWorkload(message: DequeuedMessage, timings: WarmStartTimings) {
+    const createStart = performance.now();
+    try {
+      await this.workloadManager.create(await this.createOptionsFor(message, timings));
       recordPhaseSince("workload_create", createStart, undefined);
       workloadCreateDuration.observe(
         { backend: this.workloadManagerBackend, outcome: "success" },
@@ -722,10 +769,8 @@ class ManagedSupervisor {
     const headers: Record<string, string> = {
       "Content-Type": "application/json",
     };
-    // Propagate the inbound W3C traceparent so the upstream warm-start
-    // receiver continues the same trace instead of minting a new one. Gated
-    // by the same kill switch as the wide-event emission so the whole PR is
-    // a no-op on the wire when disabled.
+    // Lets the warm-start receiver continue this trace instead of minting a new one.
+    // Gated by the wide-events kill switch, so the wire is unchanged when disabled.
     if (this.wideEventOpts.enabled && traceparent) {
       headers.traceparent = traceparent;
     }