feat(supervisor): release a run whose container or pod failed before it started

The supervisor can now give back a run whose container or pod could not be created, or whose pod failed before the run started, so the run is requeued at once instead of waiting out the pending execution timeout. Both parts are off by default. Set `RELEASE_UNSTARTED_RUNS_ENABLED=true` on the supervisor for creates that fail, and `RELEASE_FAILED_POD_RUNS_ENABLED=true` for pods that fail.

Mono-RevId: 7b425fa0ec7bf68fe11de4d1f3a44319e933bcc1

Author: myftija

.server-changes/release-failed-pod-runs.md

@@ -0,0 +1,6 @@
+---
+area: supervisor
+type: improvement
+---
+
+Runs whose container or pod could not be created, or whose pod failed before the run started, can now be requeued straight away instead of waiting for the pending execution timeout. Both are off by default: set `RELEASE_UNSTARTED_RUNS_ENABLED=true` on a supervisor for creates that fail, and `RELEASE_FAILED_POD_RUNS_ENABLED=true` for pods that fail.

apps/supervisor/src/env.test.ts

@@ -185,6 +185,31 @@ describe("Env superRefine - backpressure source awareness", () => {
   });
 });
 
+describe("releasing runs that never started", () => {
+  it("is off by default for both failed creates and failed pods", () => {
+    const env = Env.parse(base);
+    expect(env.RELEASE_UNSTARTED_RUNS_ENABLED).toBe(false);
+    expect(env.RELEASE_FAILED_POD_RUNS_ENABLED).toBe(false);
+  });
+
+  it("limits releases to five different runs a minute unless told otherwise, never fewer than none", () => {
+    expect(Env.parse(base).RUN_RELEASES_PER_MINUTE).toBe(5);
+    expect(Env.parse({ ...base, RUN_RELEASES_PER_MINUTE: "0" }).RUN_RELEASES_PER_MINUTE).toBe(0);
+    expect(() => Env.parse({ ...base, RUN_RELEASES_PER_MINUTE: "-1" })).toThrow();
+    expect(Env.parse({ ...base, RUN_RELEASES_PER_MINUTE: " " }).RUN_RELEASES_PER_MINUTE).toBe(5);
+  });
+
+  it("turns on each path on its own", () => {
+    const creates = Env.parse({ ...base, RELEASE_UNSTARTED_RUNS_ENABLED: "true" });
+    expect(creates.RELEASE_UNSTARTED_RUNS_ENABLED).toBe(true);
+    expect(creates.RELEASE_FAILED_POD_RUNS_ENABLED).toBe(false);
+
+    const pods = Env.parse({ ...base, RELEASE_FAILED_POD_RUNS_ENABLED: "1" });
+    expect(pods.RELEASE_UNSTARTED_RUNS_ENABLED).toBe(false);
+    expect(pods.RELEASE_FAILED_POD_RUNS_ENABLED).toBe(true);
+  });
+});
+
 describe("worker cell name", () => {
   it("treats an empty value as unset", () => {
     expect(

apps/supervisor/src/env.ts

@@ -363,6 +363,17 @@ export const Env = z
     // Report a run whose runner the runtime killed, so it fails as out of memory only
     // when the runtime says so. Without it the run waits out the heartbeat timeout.
     RUNNER_KILL_REPORTS_ENABLED: BoolEnv.default(true),
+    // Release a run whose container or pod could not be created, so it is requeued at once instead
+    // of waiting out the pending execution timeout.
+    RELEASE_UNSTARTED_RUNS_ENABLED: BoolEnv.default(false),
+    // Release a run whose pod failed before the run started, likewise. Independent of the flag above.
+    RELEASE_FAILED_POD_RUNS_ENABLED: BoolEnv.default(false),
+    // Different runs released within a minute, by either flag, past which every release waits out
+    // the timeout until a slot frees. 0 withholds every release; blank means the default.
+    RUN_RELEASES_PER_MINUTE: z.preprocess(
+      (value) => (typeof value === "string" && value.trim() === "" ? undefined : value),
+      z.coerce.number().int().min(0).default(5)
+    ),
     FAILED_POD_HANDLER_RECONNECT_INTERVAL_MS: z.coerce.number().int().default(1000),
     FAILED_POD_HANDLER_WATCH_TIMEOUT_SECONDS: z.coerce.number().int().positive().default(300),
 

apps/supervisor/src/index.ts

@@ -42,6 +42,8 @@ import { register } from "./metrics.js";
 import { PodCleaner } from "./services/podCleaner.js";
 import { FailedPodHandler } from "./services/failedPodHandler.js";
 import { RunnerKillReporter } from "./services/runnerKills.js";
+import { dequeuedRun, RunReleaser } from "./services/runReleases.js";
+import { WorkloadNotCreatedError } from "./workloadManager/createErrors.js";
 import { DequeueDrain } from "./services/dequeueDrain.js";
 import { getWorkerToken } from "./workerToken.js";
 import { mintDeploymentToken } from "./workloadToken.js";
@@ -124,6 +126,7 @@ class ManagedSupervisor {
   private readonly restoreWatcher?: RestoreWatcher;
   private readonly runnerKillReporter?: RunnerKillReporter;
   private readonly runnerKillInformer?: RunnerKillInformer;
+  private readonly runReleaser?: RunReleaser;
   private readonly logger = new SimpleStructuredLogger("managed-supervisor");
   private readonly resourceMonitor: ResourceMonitor;
   private readonly checkpointClient?: CheckpointClient;
@@ -180,8 +183,18 @@ class ManagedSupervisor {
       additionalEnvVars: env.RUNNER_ADDITIONAL_ENV_VARS,
       dockerAutoremove: env.DOCKER_AUTOREMOVE_EXITED_CONTAINERS,
       checkpointsEnabled: !!env.TRIGGER_CHECKPOINT_URL,
+      releaseUnstartedRuns: env.RELEASE_UNSTARTED_RUNS_ENABLED,
     } satisfies WorkloadManagerOptions;
 
+    this.runReleaser =
+      env.RELEASE_UNSTARTED_RUNS_ENABLED || env.RELEASE_FAILED_POD_RUNS_ENABLED
+        ? new RunReleaser({
+            register,
+            client: () => this.workerSession.httpClient,
+            runsReleasedAMinute: env.RUN_RELEASES_PER_MINUTE,
+          })
+        : undefined;
+
     this.resourceMonitor = env.RESOURCE_MONITOR_ENABLED
       ? this.isKubernetes
         ? new KubernetesResourceMonitor(createK8sApi(), env.TRIGGER_WORKER_INSTANCE_NAME)
@@ -287,6 +300,7 @@ class ManagedSupervisor {
           namespace: env.KUBERNETES_NAMESPACE,
           reconnectIntervalMs: env.FAILED_POD_HANDLER_RECONNECT_INTERVAL_MS,
         });
+        const runReleaser = env.RELEASE_FAILED_POD_RUNS_ENABLED ? this.runReleaser : undefined;
         this.failedPodHandler = new FailedPodHandler({
           register,
           namespace: env.KUBERNETES_NAMESPACE,
@@ -295,6 +309,7 @@ class ManagedSupervisor {
           onRunnerKilled: this.runnerKillReporter
             ? (runnerId, kill) => this.runnerKillReporter!.onRunnerKilled(runnerId, kill)
             : undefined,
+          onRunPodFailed: runReleaser ? (run) => runReleaser.release(run, "PodFailed") : undefined,
         });
       } else {
         this.logger.warn("Failed pod handler disabled");
@@ -303,6 +318,12 @@ class ManagedSupervisor {
       this.logger.log("🧹 Pod cleaner and failed pod handler disabled, the operator owns run pods");
     }
 
+    if (env.RELEASE_FAILED_POD_RUNS_ENABLED && !this.failedPodHandler) {
+      this.logger.warn(
+        "RELEASE_FAILED_POD_RUNS_ENABLED has no effect without the failed pod handler"
+      );
+    }
+
     if (env.TRIGGER_DEQUEUE_INTERVAL_MS > env.TRIGGER_DEQUEUE_IDLE_INTERVAL_MS) {
       this.logger.warn(
         `⚠️  TRIGGER_DEQUEUE_INTERVAL_MS (${env.TRIGGER_DEQUEUE_INTERVAL_MS}) is greater than TRIGGER_DEQUEUE_IDLE_INTERVAL_MS (${env.TRIGGER_DEQUEUE_IDLE_INTERVAL_MS}) - did you mix them up?`
@@ -968,6 +989,10 @@ class ManagedSupervisor {
         runId: message.run.friendlyId,
         error,
       });
+
+      if (error instanceof WorkloadNotCreatedError) {
+        await this.runReleaser?.release(dequeuedRun(message), "CreateFailed");
+      }
     }
   }
 

apps/supervisor/src/services/failedPodHandler.ts

@@ -7,6 +7,7 @@ import { createK8sApi } from "../clients/kubernetes.js";
 import { ReconnectingInformer } from "../clients/reconnectingInformer.js";
 import { register } from "../metrics.js";
 import { podKill, type RunnerKill } from "./runnerKills.js";
+import { runOfPod, type DequeuedRun } from "./runReleases.js";
 
 // Operator-owned pods carry app=task-run too, since the network policy selects
 // on it, so they have to be excluded by their own label instead. Unconditional:
@@ -25,6 +26,8 @@ export type FailedPodHandlerOptions = {
   register?: Registry;
   /** Called with the pod's name, its runner id, when the runtime killed its run-controller. */
   onRunnerKilled?: (runnerId: string, kill: RunnerKill) => Promise<unknown>;
+  /** Called with the run a failed pod was dequeued for, unless the pod shut down gracefully. */
+  onRunPodFailed?: (run: DequeuedRun) => Promise<unknown>;
 };
 
 export class FailedPodHandler {
@@ -33,6 +36,7 @@ export class FailedPodHandler {
   private readonly k8s: K8sApi;
   private readonly namespace: string;
   private readonly onRunnerKilled?: FailedPodHandlerOptions["onRunnerKilled"];
+  private readonly onRunPodFailed?: FailedPodHandlerOptions["onRunPodFailed"];
 
   private isRunning = false;
 
@@ -58,6 +62,7 @@ export class FailedPodHandler {
 
     this.namespace = opts.namespace;
     this.onRunnerKilled = opts.onRunnerKilled;
+    this.onRunPodFailed = opts.onRunPodFailed;
 
     this.watch = new ReconnectingInformer({
       name: "failed-pod-informer",
@@ -217,16 +222,21 @@ export class FailedPodHandler {
   private async onPodFailed(pod: V1Pod) {
     this.logger.debug("pod-failed", this.podSummary(pod));
 
+    let run: DequeuedRun | undefined;
     try {
-      await this.processFailedPod(pod);
+      run = await this.processFailedPod(pod);
     } catch (error) {
       this.logger.error("pod-failed: error processing pod", this.podSummary(pod), { error });
-    } finally {
-      await this.deletePod(pod);
+    }
+    const deleted = await this.deletePod(pod);
+
+    // The run's next pod takes this pod's name, so it is released only once this one is gone.
+    if (deleted && run && this.onRunPodFailed) {
+      await this.onRunPodFailed(run);
     }
   }
 
-  private async processFailedPod(pod: V1Pod) {
+  private async processFailedPod(pod: V1Pod): Promise<DequeuedRun | undefined> {
     this.logger.verbose("pod-failed: processing pod", this.podSummary(pod));
 
     const mainContainer = pod.status?.containerStatuses?.find((c) => c.name === "run-controller");
@@ -252,25 +262,41 @@ export class FailedPodHandler {
     if (kill && this.onRunnerKilled) {
       await this.onRunnerKilled(pod.metadata!.name!, kill);
     }
+
+    return runOfPod(pod);
   }
 
-  private async deletePod(pod: V1Pod) {
+  private async deletePod(pod: V1Pod): Promise<boolean> {
     this.logger.verbose("pod-failed: deleting pod", this.podSummary(pod));
+    // With a release the run's next pod takes this name within seconds, so the delete must spare it.
+    const releasing = this.onRunPodFailed !== undefined;
     try {
       await this.k8s.core.deleteNamespacedPod({
         name: pod.metadata!.name!,
         namespace: this.namespace,
+        ...(releasing ? { body: { preconditions: { uid: pod.metadata!.uid! } } } : {}),
       });
       this.deletedPodsTotal.inc({
         namespace: this.namespace,
         status: this.podStatus(pod),
       });
+      return true;
     } catch (error) {
+      const code = (error as { code?: unknown } | null)?.code;
+      if (releasing && code === 404) {
+        this.logger.info("pod-failed: pod already gone", this.podSummary(pod));
+        return true;
+      }
+      if (releasing && code === 409) {
+        this.logger.info("pod-failed: pod already replaced, leaving it", this.podSummary(pod));
+        return false;
+      }
       this.logger.error("pod-failed: error deleting pod", this.podSummary(pod), { error });
       this.deletionErrorsTotal.inc({
         namespace: this.namespace,
         error_type: error instanceof Error ? error.name : "unknown",
       });
+      return false;
     }
   }