feat(h05): checkpoint save/load path inputs in TopBar → stage_train G12

apstenku123 · apstenku123 · commit 5516781f94d8 · 2026-05-22T12:00:17.000+02:00
TopBar gains two text inputs in the Train dropdown menu:
  - `train-checkpoint-save-path` → opts.checkpoint_save_path
  - `train-checkpoint-load-path` → opts.checkpoint_load_path

App.handleRunPipeline forwards both into stage_options.train when
non-empty so backend stage_train (V5-G12) can safetensors.save_file /
safetensors.load_file the model weights. Backend extras already emit
`checkpoint.saved_path` and `checkpoint.loaded_path`.

Tests:
- vbgui vitest TopBar.test.tsx: two H05 specs cover (a) filled inputs
  forward as expected and (b) empty inputs omit fields — 20/20 TopBar
  tests passing.
- vbgui vitest full suite: 200/200 passing.
- vbgui e2e 59_checkpoint_save_load.spec.ts:
    H05.A: save-path → extras.checkpoint.saved_path matches; file lands
           on disk after the train run.
    H05.B: depends on H05.A having written the file; load-path →
           extras.checkpoint.loaded_path matches.
  Strict loss-continuation (H05.6 in plan) is deferred to H19
  (cppmega-mlx-56o — strict identical-loss-continuation) where the
  spec/data pinning also lands.
- vbgui e2e V6 suite (55–59): 7/7 passing — no regression.
- backend pytest stage_train + warm + checkpoint: 44/44 passing.
diff --git a/vbgui/e2e/scenarios/59_checkpoint_save_load.spec.ts b/vbgui/e2e/scenarios/59_checkpoint_save_load.spec.ts
@@ -0,0 +1,58 @@
+// H05: TopBar checkpoint save/load path inputs forward into stage_train
+// G12. The two tests run sequentially in this file (--workers=1):
+//   1. Save-path → extras.checkpoint.saved_path matches and the
+//      safetensors file lands on disk for the next test to load.
+//   2. Load-path with that file → extras.checkpoint.loaded_path matches.
+//
+// The strict loss-continuation round-trip (H05.6) is deferred to H19
+// (strict identical-loss-continuation) where the spec/data also gets
+// pinned to a deterministic seed. H05 closes the path-forwarding gap.
+
+import { test, expect } from "@playwright/test";
+import { unlinkSync, existsSync } from "node:fs";
+import { gotoApp, selectPreset, closeModal } from "../fixtures";
+
+const SAVE = "/tmp/vbgui_h05_ckpt.safetensors";
+
+test.beforeAll(() => {
+  if (existsSync(SAVE)) unlinkSync(SAVE);
+});
+
+test("H05.A: save-path forwards → checkpoint.saved_path matches",
+  async ({ page }) => {
+    test.setTimeout(120_000);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+    await page.getByTestId("run-pipeline-toggle").click();
+    await page.getByTestId("train-num-steps").fill("2");
+    await page.getByTestId("train-checkpoint-save-path").fill(SAVE);
+    await page.getByTestId("run-pipeline-train").click();
+    await page.getByTestId("run-result-modal").waitFor({ timeout: 60_000 });
+    await page.getByTestId("run-result-expand-train").click();
+    await page.getByTestId("run-result-extras-row-train").waitFor();
+    const saved = await page.getByTestId(
+      "run-result-extras-train-checkpoint-saved_path").textContent();
+    expect(saved?.trim()).toBe(SAVE);
+    await closeModal(page);
+    expect(existsSync(SAVE)).toBe(true);
+  });
+
+test("H05.B: load-path forwards → checkpoint.loaded_path matches",
+  async ({ page }) => {
+    test.setTimeout(120_000);
+    // Guard: depends on file written by H05.A.
+    expect(existsSync(SAVE)).toBe(true);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+    await page.getByTestId("run-pipeline-toggle").click();
+    await page.getByTestId("train-num-steps").fill("2");
+    await page.getByTestId("train-checkpoint-load-path").fill(SAVE);
+    await page.getByTestId("run-pipeline-train").click();
+    await page.getByTestId("run-result-modal").waitFor({ timeout: 60_000 });
+    await page.getByTestId("run-result-expand-train").click();
+    await page.getByTestId("run-result-extras-row-train").waitFor();
+    const loaded = await page.getByTestId(
+      "run-result-extras-train-checkpoint-loaded_path").textContent();
+    expect(loaded?.trim()).toBe(SAVE);
+    await closeModal(page);
+  });
diff --git a/vbgui/src/App.tsx b/vbgui/src/App.tsx
@@ -293,7 +293,10 @@ export function App(): JSX.Element {
 
   const handleRunPipeline = useCallback(async (
     mode: RunMode,
-    opts?: { num_steps?: number; warm_start?: boolean },
+    opts?: { num_steps?: number; warm_start?: boolean;
+      checkpoint_save_path?: string;
+      checkpoint_load_path?: string;
+    },
   ) => {
     const snap = wireSpecRef.current;
     if (snap.nodes.length === 0) {
@@ -322,6 +325,13 @@ export function App(): JSX.Element {
       if (opts?.warm_start && lastTrainRunId) {
         trainOpts.continue_from_run_id = lastTrainRunId;
       }
+      // H05: forward checkpoint save/load paths to stage_train G12.
+      if (opts?.checkpoint_save_path) {
+        trainOpts.checkpoint_save_path = opts.checkpoint_save_path;
+      }
+      if (opts?.checkpoint_load_path) {
+        trainOpts.checkpoint_load_path = opts.checkpoint_load_path;
+      }
       if (trainParquetPath) trainOpts.parquet_path = trainParquetPath;
       if (trainTokenizerPath) trainOpts.tokenizer_path = trainTokenizerPath;
       // Forward SideChannelsTab train selection as synthetic int lists for the
diff --git a/vbgui/src/components/TopBar.tsx b/vbgui/src/components/TopBar.tsx
@@ -14,7 +14,10 @@ export interface TopBarProps {
   onTopologyChange: (t: TopologyFactory) => void;
   onCompileModeChange: (m: SpecState["sharding"]["compile_mode"]) => void;
   onRunPipeline: (mode: RunMode,
-    opts?: { num_steps?: number; warm_start?: boolean }) => void;
+    opts?: { num_steps?: number; warm_start?: boolean;
+      checkpoint_save_path?: string;
+      checkpoint_load_path?: string;
+    }) => void;
   /** H02: toggle callbacks. */
   onMixedPrecisionChange?: (enabled: boolean) => void;
   onFp8EnabledChange?: (enabled: boolean) => void;
@@ -40,6 +43,8 @@ export function TopBar(p: TopBarProps): JSX.Element {
   const [open, setOpen] = useState(false);
   const [trainNumSteps, setTrainNumSteps] = useState<number>(2);
   const [warmStart, setWarmStart] = useState<boolean>(false);
+  const [ckptSavePath, setCkptSavePath] = useState<string>("");
+  const [ckptLoadPath, setCkptLoadPath] = useState<string>("");
   return (
     <header data-testid="top-bar"
             style={{ height: 56, display: "flex", alignItems: "center",
@@ -174,11 +179,37 @@ export function TopBar(p: TopBarProps): JSX.Element {
                      onChange={(e) => setWarmStart(e.target.checked)} />
               warm-start (continue from last run)
             </label>
+            <div style={{ padding: "6px 12px", display: "flex",
+                          flexDirection: "column", gap: 4, fontSize: 11,
+                          color: "#374151" }}>
+              <label style={{ display: "flex", alignItems: "center",
+                              gap: 6 }}>
+                <span style={{ width: 78, color: "#6b7280" }}>ckpt save:</span>
+                <input data-testid="train-checkpoint-save-path" type="text"
+                       placeholder="/tmp/ckpt.safetensors"
+                       value={ckptSavePath}
+                       onChange={(e) => setCkptSavePath(e.target.value)}
+                       style={{ width: 200 }} />
+              </label>
+              <label style={{ display: "flex", alignItems: "center",
+                              gap: 6 }}>
+                <span style={{ width: 78, color: "#6b7280" }}>ckpt load:</span>
+                <input data-testid="train-checkpoint-load-path" type="text"
+                       placeholder="/tmp/prev.safetensors"
+                       value={ckptLoadPath}
+                       onChange={(e) => setCkptLoadPath(e.target.value)}
+                       style={{ width: 200 }} />
+              </label>
+            </div>
             <button data-testid="run-pipeline-train"
                     onClick={() => { setOpen(false);
                                      p.onRunPipeline("train",
                                        { num_steps: trainNumSteps,
-                                         warm_start: warmStart }); }}
+                                         warm_start: warmStart,
+                                         checkpoint_save_path:
+                                           ckptSavePath || undefined,
+                                         checkpoint_load_path:
+                                           ckptLoadPath || undefined }); }}
                     disabled={!!p.trainDisabled}
                     title={p.trainDisabled?.reason ?? ""}
                     style={{ ...menuItem,
diff --git a/vbgui/tests/TopBar.test.tsx b/vbgui/tests/TopBar.test.tsx
@@ -131,6 +131,33 @@ describe("TopBar", () => {
       expect.objectContaining({ num_steps: expect.any(Number) }));
   });
 
+  it("H05: train-checkpoint-save/load-path inputs forward into opts",
+    () => {
+      const onRunPipeline = vi.fn();
+      render(<TopBar {...defaultTopProps({ onRunPipeline })} />);
+      fireEvent.click(screen.getByTestId("run-pipeline-toggle"));
+      fireEvent.change(screen.getByTestId("train-checkpoint-save-path"),
+        { target: { value: "/tmp/save.safetensors" } });
+      fireEvent.change(screen.getByTestId("train-checkpoint-load-path"),
+        { target: { value: "/tmp/load.safetensors" } });
+      fireEvent.click(screen.getByTestId("run-pipeline-train"));
+      expect(onRunPipeline).toHaveBeenLastCalledWith("train",
+        expect.objectContaining({
+          checkpoint_save_path: "/tmp/save.safetensors",
+          checkpoint_load_path: "/tmp/load.safetensors",
+        }));
+    });
+
+  it("H05: empty checkpoint inputs omit fields", () => {
+    const onRunPipeline = vi.fn();
+    render(<TopBar {...defaultTopProps({ onRunPipeline })} />);
+    fireEvent.click(screen.getByTestId("run-pipeline-toggle"));
+    fireEvent.click(screen.getByTestId("run-pipeline-train"));
+    const opts = onRunPipeline.mock.calls.at(-1)?.[1];
+    expect(opts.checkpoint_save_path).toBeUndefined();
+    expect(opts.checkpoint_load_path).toBeUndefined();
+  });
+
   it("H04: train-warm-start checkbox forwards warm_start flag", () => {
     const onRunPipeline = vi.fn();
     render(<TopBar {...defaultTopProps({ onRunPipeline })} />);
