feat(supervisor): carry the deployment's template key to the Runner

Carry a deployment's prebuilt template key on the dequeued run message, so the supervisor can hand it to the runner it creates for both new runs and restores. On a worker group that starts runs from templates, a run that cannot start from one goes back to the region's container worker group instead.

Mono-RevId: 866804b09fe4724fc98af488aded1b85364f0098

Author: 1stvamp

apps/supervisor/src/index.ts

@@ -22,7 +22,9 @@ import {
   RunCrdWorkloadManager,
   RunnerKillInformer,
   RunnerRestoreInformer,
+  reportTemplateKeyRefusal,
   settleRestoreFailure,
+  type RunnerStartFailure,
 } from "./workloadManager/runCrd.js";
 import { workerLaneFor, type WorkloadManagerBackend } from "./workerLane.js";
 import { RestoreWatcher, type RestoreFailure } from "./services/restoreWatcher.js";
@@ -89,6 +91,20 @@ const outboundRequestDuration = new Histogram({
   registers: [register],
 });
 
+const startReportsTotal = new Counter({
+  name: "supervisor_start_reports_total",
+  help: "Cold starts the operator failed before they began, moved to the fallback lane through the platform, by result: ok, conflict when the run had already moved on, or error.",
+  labelNames: ["result"],
+  registers: [register],
+});
+
+const templateKeyRefusalsTotal = new Counter({
+  name: "supervisor_template_key_refusals_total",
+  help: "Dequeued runs this lane created no Runner for, as they had no valid template key, by outcome (fallback or fail) and result: ok, conflict when the run had already moved on, or error.",
+  labelNames: ["outcome", "result"],
+  registers: [register],
+});
+
 const restoreReportsTotal = new Counter({
   name: "supervisor_restore_reports_total",
   help: "Failed restores reported to the platform, by outcome (requeue or fail) and result: ok, conflict when the run had already moved on, or error.",
@@ -183,6 +199,9 @@ class ManagedSupervisor {
         this.runnerKillInformer = new RunnerKillInformer({
           namespace: env.KUBERNETES_NAMESPACE,
           onRunnerKilled: (runnerId, kill) => void killReporter.onRunnerKilled(runnerId, kill),
+          ...(env.KUBERNETES_RUN_CRD_ENABLED
+            ? { onRunnerStartFailed: (failure) => void this.reportStartFailure(failure) }
+            : {}),
         });
       }
     }
@@ -197,10 +216,16 @@ class ManagedSupervisor {
         this.restoreInformer = new RunnerRestoreInformer({
           namespace: env.KUBERNETES_NAMESPACE,
           onUnwatchedFailure: (failure) => void this.restoreWatcher?.report(failure),
-          onUnwatchedRestore: ({ runnerId, uid, runFriendlyId, snapshotFriendlyId }) => {
+          onUnwatchedRestore: ({
+            runnerId,
+            uid,
+            runFriendlyId,
+            snapshotFriendlyId,
+            environmentId,
+          }) => {
             if (runFriendlyId && snapshotFriendlyId) {
               void this.restoreWatcher?.watch(
-                { runFriendlyId, snapshotFriendlyId },
+                { runFriendlyId, snapshotFriendlyId, environmentId },
                 { runnerId, uid }
               );
             }
@@ -505,6 +530,11 @@ class ManagedSupervisor {
 
             const { checkpoint, ...rest } = message;
 
+            if (this.runCrdManager && (await this.refuseWithoutTemplateKey(message))) {
+              setExtra(fromContext(), "path_taken", "no_template_key");
+              return;
+            }
+
             if (checkpoint) {
               setExtra(fromContext(), "path_taken", "restore");
               this.logger.debug("Restoring run", { runId: message.run.id });
@@ -655,6 +685,7 @@ class ManagedSupervisor {
       envId: message.environment.id,
       envType: message.environment.type,
       image: message.image,
+      templateKey: message.templateKey,
       machine: message.run.machine,
       orgId: message.organization.id,
       projectId: message.project.id,
@@ -691,6 +722,7 @@ class ManagedSupervisor {
         {
           runFriendlyId: message.run.friendlyId,
           snapshotFriendlyId: message.snapshot.friendlyId,
+          environmentId: message.environment.id,
         },
         restore
       );
@@ -734,6 +766,7 @@ class ManagedSupervisor {
   private async reportRestoreFailure({
     runFriendlyId,
     snapshotFriendlyId,
+    environmentId,
     runnerId,
     outcome,
   }: RestoreFailure) {
@@ -760,7 +793,8 @@ class ManagedSupervisor {
             runFriendlyId,
             snapshotFriendlyId,
             body,
-            runnerId
+            runnerId,
+            environmentId
           ),
       }
     );
@@ -806,6 +840,105 @@ class ManagedSupervisor {
     }
   }
 
+  /** Reports a run this lane needs a template key for and has none, rather than starting it. */
+  private async refuseWithoutTemplateKey(message: DequeuedMessage): Promise<boolean> {
+    const refused = await reportTemplateKeyRefusal(
+      {
+        runtime: env.KUBERNETES_RUNNER_RUNTIME,
+        templateKey: message.templateKey,
+        checkpoint: message.checkpoint,
+      },
+      (body) =>
+        this.workerSession.httpClient.reportRestoreOutcome(
+          message.run.friendlyId,
+          message.snapshot.friendlyId,
+          body,
+          undefined,
+          message.environment.id
+        )
+    );
+    if (!refused) {
+      return false;
+    }
+    templateKeyRefusalsTotal.inc({ outcome: refused.outcome, result: refused.result });
+    const fields = {
+      runId: message.run.friendlyId,
+      snapshotId: message.snapshot.friendlyId,
+      outcome: refused.outcome,
+      reason: refused.reason,
+      result: refused.result,
+    };
+    if (refused.result === "error") {
+      this.logger.error("No Runner created without a template key, and the report failed", {
+        ...fields,
+        error: refused.error,
+      });
+    } else {
+      this.logger.warn("No Runner created without a template key", fields);
+    }
+    return true;
+  }
+
+  /**
+   * Moves a cold start the operator failed before its attempt began to the platform's fallback
+   * lane, rather than leaving it to the stall timeout. Without one it is requeued here, and the
+   * queue's redelivery budget bounds how often a run comes back this way.
+   */
+  private async reportStartFailure(failure: RunnerStartFailure) {
+    const { runFriendlyId, snapshotFriendlyId, runnerId, environmentId, outcome } = failure;
+    const fields = { runFriendlyId, snapshotFriendlyId, runnerId, reason: outcome.reason };
+    this.logger.warn("Runner start failed (run-crd)", { ...fields, error: outcome.error });
+    const manager = this.runCrdManager;
+    if (!manager) {
+      this.logger.warn(
+        "Runner start failure left to the stall timeout, no run-crd manager",
+        fields
+      );
+      return;
+    }
+    await this.workerSession.httpClient.sendDebugLog(runFriendlyId, {
+      time: new Date(),
+      message: "runner failed to start on the node",
+      properties: { runnerId, error: outcome.error },
+    });
+    const settled = await settleRestoreFailure(
+      { runnerId, outcome, action: { outcome: "fallback", deleteRunner: true } },
+      {
+        deleteRunner: (name, uid) => manager.deleteFailedRunner(name, uid),
+        report: (body) =>
+          this.workerSession.httpClient.reportRestoreOutcome(
+            runFriendlyId,
+            snapshotFriendlyId,
+            body,
+            runnerId,
+            environmentId
+          ),
+      }
+    );
+    switch (settled.action) {
+      case "none":
+        return;
+      case "kept":
+        this.logger.warn("Failed start Runner not deleted, so the run was not moved", {
+          ...fields,
+          error: settled.error,
+        });
+        return;
+      case "reported":
+        startReportsTotal.inc({ result: settled.result });
+        if (settled.result === "error") {
+          this.logger.warn("Start failure report failed", { ...fields, error: settled.error });
+        } else {
+          this.logger.info(
+            settled.result === "conflict"
+              ? "Start failure not applied, the run has moved on"
+              : "Start failure reported, the run was moved",
+            fields
+          );
+        }
+    }
+  }
+
   private async createWorkload(message: DequeuedMessage, timings: WarmStartTimings) {
     const createStart = performance.now();
     try {

apps/supervisor/src/services/restoreWatcher.test.ts

@@ -50,6 +50,28 @@ describe("RestoreWatcher", () => {
     });
   });
 
+  it("reports a failed restore with the environment it was dequeued for", async () => {
+    const reported: RestoreFailure[] = [];
+    const { watcher, outcomes } = setup({ reportFailure: async (f) => void reported.push(f) });
+
+    const watched = watcher.watch(
+      { ...target, environmentId: "env_a" },
+      { runnerId: "runner-a", uid: "uid-a" }
+    );
+    outcomes[0]!.resolve(failed);
+    await watched;
+
+    expect(reported).toEqual([
+      {
+        runFriendlyId: "run_a",
+        snapshotFriendlyId: "snapshot_1",
+        environmentId: "env_a",
+        runnerId: "runner-a",
+        outcome: failed,
+      },
+    ]);
+  });
+
   it("names the snapshot dequeued last when the restore is redelivered", async () => {
     const { watcher, outcomes, awaitRestore, reportFailure } = setup();
 

apps/supervisor/src/services/restoreWatcher.ts

@@ -11,12 +11,14 @@ import {
 type RestoreWatchTarget = {
   runFriendlyId: string;
   snapshotFriendlyId: string;
+  environmentId?: string;
 };
 
 /** A failed restore, with what the watch or the Runner knew of its run. */
 export type RestoreFailure = {
   runFriendlyId?: string;
   snapshotFriendlyId?: string;
+  environmentId?: string;
   runnerId: string;
   outcome: RestoreWatchResult & { ok: false };
 };
@@ -136,6 +138,7 @@ export class RestoreWatcher {
     await this.report({
       runFriendlyId,
       snapshotFriendlyId: entry.snapshotFriendlyId,
+      environmentId: entry.environmentId,
       runnerId,
       outcome,
     });

apps/supervisor/src/services/runnerKills.test.ts

@@ -11,7 +11,13 @@ import type { AddressInfo } from "node:net";
 import { Registry } from "prom-client";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import type { K8sApi } from "../clients/kubernetes.js";
-import { RunnerKillInformer, runnerKill, SUSPEND_ANNOTATION } from "../workloadManager/runCrd.js";
+import {
+  RESTORE_LABEL,
+  RunnerKillInformer,
+  runnerKill,
+  SUSPEND_ANNOTATION,
+  type RunnerStartFailure,
+} from "../workloadManager/runCrd.js";
 import { FailedPodHandler } from "./failedPodHandler.js";
 import { podKill, RunnerKillReporter, type RunnerKill } from "./runnerKills.js";
 
@@ -497,3 +503,66 @@ describe("RunnerKillInformer", () => {
     ]);
   });
 });
+
+describe("RunnerKillInformer start failures", () => {
+  function templateMissing(name: string, labels?: Record<string, string>) {
+    const runner = failedRunner(
+      name,
+      "TemplateNotFound",
+      "template abc is not available on node node-a: not in the store",
+      "run_a"
+    );
+    return {
+      ...runner,
+      metadata: { ...runner.metadata, ...(labels ? { labels } : {}) },
+      spec: { ...runner.spec, bootstrap: { runFriendlyID: "run_a", snapshotFriendlyID: "snap_a" } },
+    } as KubernetesObject;
+  }
+
+  it("passes on a cold start whose template is missing once, and not as a kill", async () => {
+    const cold = templateMissing("runner-cold");
+    const { k8s, watchCallbacks } = fakeK8s([
+      cold,
+      templateMissing("runner-resume", { [RESTORE_LABEL]: "true" }),
+    ]);
+    const killed: string[] = [];
+    const failed: RunnerStartFailure[] = [];
+    const informer = new RunnerKillInformer({
+      namespace: "runs",
+      k8s,
+      onRunnerKilled: (runnerId) => killed.push(runnerId),
+      onRunnerStartFailed: (failure) => failed.push(failure),
+    });
+
+    await informer.start();
+    await vi.waitFor(() => expect(failed).toHaveLength(1));
+    await vi.waitFor(() => expect(watchCallbacks).toHaveLength(1));
+    watchCallbacks[0]!("MODIFIED", {
+      ...cold,
+      metadata: { ...cold.metadata, resourceVersion: "2" },
+    });
+    await informer.stop();
+
+    expect(killed).toEqual([]);
+    expect(
+      failed.map((f) => [f.runnerId, f.runFriendlyId, f.snapshotFriendlyId, f.outcome.uid])
+    ).toEqual([["runner-cold", "run_a", "snap_a", "runner-cold-uid"]]);
+  });
+
+  it("passes on nothing for it without a start failure handler", async () => {
+    const { k8s, watchCallbacks } = fakeK8s([templateMissing("runner-cold")]);
+    const killed: string[] = [];
+    const informer = new RunnerKillInformer({
+      namespace: "runs",
+      k8s,
+      onRunnerKilled: (runnerId) => killed.push(runnerId),
+    });
+
+    await informer.start();
+    // The watch starts once the list is handled.
+    await vi.waitFor(() => expect(watchCallbacks).toHaveLength(1));
+    await informer.stop();
+
+    expect(killed).toEqual([]);
+  });
+});