fix(supervisor): watch a restore Runner and report a failed restore

Supervisor: with the Runner CRD backend, a restore that fails on its
node is detected and reported with the reason, instead of only surfacing
when the platform's heartbeat timeout fires.
Mono-RevId: c95a9ebc98d5b336717fba5cae03510a46c5b19c

Author: 1stvamp

apps/supervisor/src/index.ts

@@ -98,6 +98,7 @@ class ManagedSupervisor {
   private readonly failedPodHandler?: FailedPodHandler;
   private readonly tracing?: OtlpTraceService;
   private readonly backpressureMonitors: BackpressureMonitor[] = [];
+  private readonly watchedRestores = new Set<string>();
   private readonly backpressureRedis?: Redis;
 
   private readonly isKubernetes = isKubernetesEnvironment(env.KUBERNETES_FORCE_ENABLED);
@@ -709,10 +710,12 @@ class ManagedSupervisor {
   ) {
     const restoreStart = performance.now();
     try {
-      await manager.restore(await this.createOptionsFor(message), checkpoint);
+      const runnerId = 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 });
+      // Not awaited: the resume can take minutes, and the dequeue is done.
+      void this.watchRestore(manager, message.run.friendlyId, runnerId);
     } catch (error) {
       recordPhaseSince(
         "restore",
@@ -724,6 +727,36 @@ class ManagedSupervisor {
     }
   }
 
+  /** A resume that fails on the node is otherwise silent until the run's heartbeat stalls. */
+  private async watchRestore(
+    manager: RunCrdWorkloadManager,
+    runFriendlyId: string,
+    runnerId: string
+  ) {
+    // A redelivered restore finds the same Runner, which needs only one watch.
+    if (this.watchedRestores.has(runnerId)) {
+      return;
+    }
+    this.watchedRestores.add(runnerId);
+    const outcome = await manager
+      .awaitRestore(runnerId)
+      .finally(() => this.watchedRestores.delete(runnerId));
+    if (outcome.ok) {
+      this.logger.debug("Runner restore started", { runFriendlyId, runnerId });
+      return;
+    }
+    this.logger.error("Runner restore failed (run-crd)", {
+      runFriendlyId,
+      runnerId,
+      error: outcome.error,
+    });
+    await this.workerSession.httpClient.sendDebugLog(runFriendlyId, {
+      time: new Date(),
+      message: "restore failed on the node",
+      properties: { runnerId, error: outcome.error },
+    });
+  }
+
   private async createWorkload(message: DequeuedMessage, timings: WarmStartTimings) {
     const createStart = performance.now();
     try {

apps/supervisor/src/workloadManager/runCrd.test.ts

@@ -5,9 +5,11 @@ import {
   SUSPEND_ANNOTATION,
   SUSPEND_RUN_ANNOTATION,
   SUSPEND_SUBMITTED_ANNOTATION,
+  awaitRestoreOf,
   checkpointLocation,
   parseCheckpointLocation,
   publishedSuspend,
+  restoreOutcome,
   runnerBodyFor,
   runnerTokenSecretName,
   suspendOutcome,
@@ -517,7 +519,9 @@ describe("RunCrdWorkloadManager.restore", () => {
   }
 
   it("creates a Runner named from the checkpoint that restores its location", async () => {
-    await manager().restore(createOptions(), checkpoint);
+    await expect(manager().restore(createOptions(), checkpoint)).resolves.toBe(
+      getRestoreRunnerId("run_abc123", "checkpoint_abc")
+    );
 
     const { body } = createRunner.mock.calls[0]![0];
     expect(body.metadata.name).toBe(getRestoreRunnerId("run_abc123", "checkpoint_abc"));
@@ -538,7 +542,9 @@ describe("RunCrdWorkloadManager.restore", () => {
       createRunner.mockRejectedValue({ code: 409 });
       getRunner.mockResolvedValue(existing(phase));
 
-      await expect(manager().restore(createOptions(), checkpoint)).resolves.toBeUndefined();
+      await expect(manager().restore(createOptions(), checkpoint)).resolves.toBe(
+        getRestoreRunnerId("run_abc123", "checkpoint_abc")
+      );
       expect(getRunner).toHaveBeenCalledWith(
         expect.objectContaining({ name: getRestoreRunnerId("run_abc123", "checkpoint_abc") })
       );
@@ -589,6 +595,127 @@ describe("RunCrdWorkloadManager.restore", () => {
   });
 });
 
+describe("restoreOutcome", () => {
+  it.each([undefined, "Pending", "Admitted", "Scheduling", "Restoring"])(
+    "is still waiting while the Runner is %s",
+    (phase) => {
+      expect(restoreOutcome({ status: phase ? { phase } : undefined })).toBeUndefined();
+    }
+  );
+
+  it.each(["Running", "Suspending", "Succeeded"])("is a success once the Runner is %s", (phase) => {
+    expect(restoreOutcome({ status: { phase } })).toEqual({ ok: true });
+  });
+
+  it("is a failure carrying the operator's reason", () => {
+    const runner = {
+      status: {
+        phase: "Failed",
+        conditions: [
+          {
+            type: "Failed",
+            status: "True",
+            reason: "StartError",
+            message: "pulling the image: 401",
+          },
+        ],
+      },
+    };
+    expect(restoreOutcome(runner)).toEqual({
+      ok: false,
+      error: "StartError: pulling the image: 401",
+    });
+  });
+
+  it("is a failure even when the operator recorded no reason", () => {
+    expect(restoreOutcome({ status: { phase: "Failed" } })).toEqual({
+      ok: false,
+      error: "the Runner failed with no reason recorded",
+    });
+  });
+});
+
+describe("awaitRestoreOf", () => {
+  /** Answers each read with the next response, repeating the last. `{ throw: err }` fails that read. */
+  function reads(...responses: unknown[]) {
+    const log: unknown[] = [];
+    let i = 0;
+    const readRunner = async () => {
+      const response = responses[Math.min(i++, responses.length - 1)];
+      log.push(response);
+      if (response && typeof response === "object" && "throw" in response) {
+        throw response.throw;
+      }
+      return response;
+    };
+    return { readRunner, log };
+  }
+
+  function awaitWith(readRunner: () => Promise<unknown>, timeoutMs = 1_000) {
+    return awaitRestoreOf(readRunner, { pollMs: 1, timeoutMs, onReadError: () => {} });
+  }
+
+  it("waits through Restoring until the Runner is Running", async () => {
+    const { readRunner, log } = reads(
+      { status: { phase: "Pending" } },
+      { status: { phase: "Restoring" } },
+      { status: { phase: "Running" } }
+    );
+
+    await expect(awaitWith(readRunner)).resolves.toEqual({ ok: true });
+    expect(log).toHaveLength(3);
+  });
+
+  it("reports a failed restore with the operator's reason", async () => {
+    const { readRunner } = reads(
+      { status: { phase: "Restoring" } },
+      {
+        status: {
+          phase: "Failed",
+          conditions: [{ type: "Failed", reason: "SnapshotNodeGone", message: "node a is gone" }],
+        },
+      }
+    );
+
+    await expect(awaitWith(readRunner)).resolves.toEqual({
+      ok: false,
+      error: "SnapshotNodeGone: node a is gone",
+    });
+  });
+
+  it("keeps polling through a read that may succeed next time", async () => {
+    const readErrors: unknown[] = [];
+    const { readRunner } = reads({ throw: { code: 500 } }, { status: { phase: "Running" } });
+
+    await expect(
+      awaitRestoreOf(readRunner, {
+        pollMs: 1,
+        timeoutMs: 1_000,
+        onReadError: (err) => readErrors.push(err),
+      })
+    ).resolves.toEqual({ ok: true });
+    expect(readErrors).toEqual([{ code: 500 }]);
+  });
+
+  it("fails when the Runner is gone", async () => {
+    const { readRunner } = reads({ throw: { code: 404 } });
+
+    await expect(awaitWith(readRunner)).resolves.toEqual({
+      ok: false,
+      error: "the Runner no longer exists",
+    });
+  });
+
+  it("gives up after its timeout", async () => {
+    const { readRunner } = reads({ status: { phase: "Restoring" } });
+
+    await expect(awaitWith(readRunner, 20)).resolves.toEqual({
+      ok: false,
+      error: "the Runner did not start within 20ms",
+    });
+  });
+});
+
 describe("RunCrdWorkloadManager.suspend", () => {
   function manager(snapshots = { enabled: true, delayMs: 5_000, dispatchLimit: 10 }) {
     return new RunCrdWorkloadManager({

apps/supervisor/src/workloadManager/runCrd.ts

@@ -60,8 +60,15 @@ export type RunCrdWorkloadManagerOptions = WorkloadManagerOptions & {
    * than its own snapshot budget, so its failure arrives first and carries the reason.
    */
   suspendTimeoutMs?: number;
+  /**
+   * How long a resume has to reach Running. Longer than the operator's pod start
+   * deadline, so its failure arrives first and carries the reason.
+   */
+  restoreTimeoutMs?: number;
 };
 
+export type RunnerRestoreResult = { ok: true } | { ok: false; error: string };
+
 /**
  * Creates a Runner and stops; the operator builds the pod, so uid, node selection
  * and labels are decided in one place. A Runner serves every later run warm start
@@ -75,6 +82,7 @@ export class RunCrdWorkloadManager implements WorkloadManager, RunnerSnapshotter
   private readonly snapshots?: RunCrdWorkloadManagerOptions["snapshots"];
   private readonly suspendPollMs: number;
   private readonly suspendTimeoutMs: number;
+  private readonly restoreTimeoutMs: number;
 
   constructor(opts: RunCrdWorkloadManagerOptions) {
     this.k8s = createK8sApi();
@@ -83,6 +91,7 @@ export class RunCrdWorkloadManager implements WorkloadManager, RunnerSnapshotter
     this.snapshots = opts.snapshots;
     this.suspendPollMs = opts.suspendPollMs ?? 1_000;
     this.suspendTimeoutMs = opts.suspendTimeoutMs ?? 6 * 60_000;
+    this.restoreTimeoutMs = opts.restoreTimeoutMs ?? 16 * 60_000;
   }
 
   get snapshotsEnabled(): boolean {
@@ -229,6 +238,22 @@ export class RunCrdWorkloadManager implements WorkloadManager, RunnerSnapshotter
     };
   }
 
+  /**
+   * Waits for a resume to start or to fail. Nothing else watches one: a resume
+   * that fails never has a runner to connect and say so.
+   */
+  async awaitRestore(runnerId: string): Promise<RunnerRestoreResult> {
+    return awaitRestoreOf(() => this.getRunner(runnerId), {
+      pollMs: this.suspendPollMs,
+      timeoutMs: this.restoreTimeoutMs,
+      onReadError: (err) =>
+        this.logger.warn("[RunCrdWorkloadManager] Runner read failed during restore", {
+          runnerId,
+          rawError: err,
+        }),
+    });
+  }
+
   private getRunner(name: string): Promise<unknown> {
     return this.k8s.custom.getNamespacedCustomObject({
       group: GROUP,
@@ -255,15 +280,20 @@ export class RunCrdWorkloadManager implements WorkloadManager, RunnerSnapshotter
    * Creates a resume: a Runner the operator restores from the checkpoint's
    * snapshot, on the node holding it, instead of cold-starting. Named from the
    * checkpoint, so a redelivered restore finds the first Runner and leaves it
-   * rather than restoring twice.
+   * rather than restoring twice. Returns the Runner's name, also for a resume
+   * found in the way, which a restarted supervisor has no other watch on.
    */
-  async restore(opts: WorkloadManagerCreateOptions, checkpoint: { id: string; location: string }) {
+  async restore(
+    opts: WorkloadManagerCreateOptions,
+    checkpoint: { id: string; location: string }
+  ): Promise<string> {
     const runnerId = getRestoreRunnerId(opts.runFriendlyId, checkpoint.id);
     const restore = parseCheckpointLocation(checkpoint.location);
     const created = await this.createRunner(opts, runnerId, restore);
     if (!created) {
       await this.checkExistingRestore(runnerId, restore);
     }
+    return runnerId;
   }
 
   /**
@@ -528,6 +558,61 @@ export function suspendOutcome(runner: unknown, request: string): RunnerSuspendR
   return undefined;
 }
 
+/** Polls a resume's Runner, read by `readRunner`, until it starts, fails or times out. */
+export async function awaitRestoreOf(
+  readRunner: () => Promise<unknown>,
+  opts: { pollMs: number; timeoutMs: number; onReadError: (err: unknown) => void }
+): Promise<RunnerRestoreResult> {
+  const deadline = Date.now() + opts.timeoutMs;
+  while (Date.now() < deadline) {
+    await sleep(opts.pollMs);
+    let runner: unknown;
+    try {
+      runner = await readRunner();
+    } catch (err: unknown) {
+      const code = statusCodeOf(err);
+      if (code === 404) {
+        return { ok: false, error: "the Runner no longer exists" };
+      }
+      if (code !== undefined && TERMINAL_READ_CODES.has(code)) {
+        return { ok: false, error: `Runner read failed: ${messageOf(err)}` };
+      }
+      opts.onReadError(err);
+      continue;
+    }
+    const outcome = restoreOutcome(runner);
+    if (outcome) {
+      return outcome;
+    }
+  }
+  return { ok: false, error: `the Runner did not start within ${opts.timeoutMs}ms` };
+}
+
+/**
+ * Whether a resume started, or undefined while the operator is still restoring
+ * it. A Runner that ended is a success only if it ran to completion.
+ */
+export function restoreOutcome(runner: unknown): RunnerRestoreResult | undefined {
+  const status = (runner as RunnerSuspendStatus | null)?.status;
+  switch (status?.phase) {
+    case "Running":
+    case "Suspending":
+    case "Succeeded":
+      return { ok: true };
+    case "Failed": {
+      const condition = status.conditions?.find((c) => c.type === "Failed");
+      return {
+        ok: false,
+        error: condition
+          ? `${condition.reason}: ${condition.message}`
+          : "the Runner failed with no reason recorded",
+      };
+    }
+    default:
+      return undefined;
+  }
+}
+
 /**
  * The Runner's latest suspend, when it has an outcome the platform has not
  * been marked as accepting. A request made without the run annotation, by an