feat(v7-d06,v7-h06): dtype.cost_estimate RPC + pause/resume integration

Two paired UI honest-closure gaps closed:

V7-D06: master_dtype dropdown previously showed only dtype names.
The bench_dtype_cast_cost.py script produced static CSV/HTML offline
but the live UI surfaced no real cost. After D06 wiring, the dropdown
fetches dtype.cost_estimate on first open and renders ms/token next
to each option.

V7-H06: pause / resume mid-train. pipeline.pause / pipeline.resume
RPC handlers were declared but their handler functions did not exist.
This commit completes the dispatcher side and adds the UI button.

- jsonrpc/dtype_cost_method.py: tiny inline probe (cast+fwd+bwd
  warm-loop, ms/token derived) returns one row per dtype.
- jsonrpc/dispatcher.py + schema.py: registers dtype.cost_estimate,
  pipeline.pause, pipeline.resume in METHOD_REGISTRY.
- runner/stages.py: stage_train polls job_control.wait_while_paused
  at the top of each step (cooperative with abort_token).
- TopBar: dtype-cost-summary indicator, ms/tok suffix on each option;
  Pause button next to Cancel, toggles label to Resume.
- App.tsx: handlePauseTrain / handleResumeTrain callbacks; trainPaused
  state cleared whenever a new run starts.
- tests/v4/test_dtype_cost_method.py: 3-row response + param validation.
- tests/v4/test_pipeline_pause_resume_rpc.py: pause marks token,
  resume clears it, both methods appear in /schema/methods.
- e2e 84/85: visual assertions for dropdown labels and button toggle.

Author: apstenku123

cppmega_v4/jsonrpc/dispatcher.py

@@ -45,6 +45,9 @@
 from cppmega_v4.jsonrpc.ckpt_inspect_method import (
     CkptInspectParams, ckpt_inspect,
 )
+from cppmega_v4.jsonrpc.dtype_cost_method import (
+    DtypeCostParams, dtype_cost_estimate,
+)
 from cppmega_v4.jsonrpc.schema import (
     BuildPresetSpecsParams,
     CatalogExplainParams,
@@ -99,6 +102,14 @@
         PipelineAbortParams,
         lambda p, c: _pipeline_abort(p),
     ),
+    "pipeline.pause": (
+        PipelineAbortParams,
+        lambda p, c: _pipeline_pause(p),
+    ),
+    "pipeline.resume": (
+        PipelineAbortParams,
+        lambda p, c: _pipeline_resume(p),
+    ),
     "tokenizer.encode_visualize": (
         EncodeVisualizeParams,
         lambda p, c: encode_visualize(p, cache=c),
@@ -131,6 +142,10 @@
         CkptInspectParams,
         lambda p, c: ckpt_inspect(p, cache=c),
     ),
+    "dtype.cost_estimate": (
+        DtypeCostParams,
+        lambda p, c: dtype_cost_estimate(p, cache=c),
+    ),
 }
 
 
@@ -153,6 +168,20 @@ def _pipeline_abort(params: PipelineAbortParams) -> PipelineAbortResult:
     return PipelineAbortResult(run_id=params.run_id)
 
 
+def _pipeline_pause(params: PipelineAbortParams) -> PipelineAbortResult:
+    """V7-H06: mark a train run as paused; the loop waits between steps."""
+    from cppmega_v4.runtime.job_control import pause
+    pause(params.run_id)
+    return PipelineAbortResult(run_id=params.run_id)
+
+
+def _pipeline_resume(params: PipelineAbortParams) -> PipelineAbortResult:
+    """V7-H06: clear the paused flag so the train loop proceeds."""
+    from cppmega_v4.runtime.job_control import resume
+    resume(params.run_id)
+    return PipelineAbortResult(run_id=params.run_id)
+
+
 def dispatch(
     payload: Mapping[str, Any] | JsonRpcRequest,
     *,

cppmega_v4/jsonrpc/dtype_cost_method.py

@@ -0,0 +1,139 @@
+"""V7-D06: dtype.cost_estimate RPC.
+
+Honest-closure: scripts/bench_dtype_cast_cost.py produced a static
+CSV/HTML report of cast_overhead_ms + fwd_ms + fwdbwd_ms per dtype,
+but the UI dtype dropdown showed no numbers and the user could not
+see the real cost of fp32 vs bf16 vs fp16 before clicking Train.
+
+This method runs an inline tiny-model probe (n_iter=5 by default,
+B=1 S=8 H=128) so the UI can fetch fresh numbers on mount and render
+ms/token alongside each option in the dropdown.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from cppmega_v4.jsonrpc.cache import LRUCache
+
+DtypeName = Literal["fp32", "bf16", "fp16"]
+
+
+class DtypeCostParams(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    n_iter: int = Field(5, ge=1, le=200)
+    hidden: int = Field(128, ge=16, le=4096)
+    batch: int = Field(1, ge=1, le=32)
+    seq: int = Field(8, ge=1, le=2048)
+
+
+class DtypeCostRow(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    dtype: DtypeName
+    supported: bool
+    cast_overhead_ms: float | None = None
+    fwd_ms: float | None = None
+    fwdbwd_ms: float | None = None
+    fwd_ms_per_token: float | None = None
+    fwdbwd_ms_per_token: float | None = None
+    error: str | None = None
+
+
+class DtypeCostResult(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    rows: list[DtypeCostRow] = Field(default_factory=list)
+    n_iter: int
+    shape: list[int]
+
+
+def _build_tiny(hidden: int) -> nn.Module:
+    class _Block(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.q = nn.Linear(hidden, hidden, bias=False)
+            self.up = nn.Linear(hidden, 4 * hidden, bias=False)
+            self.down = nn.Linear(4 * hidden, hidden, bias=False)
+
+        def __call__(self, x: mx.array) -> mx.array:
+            return self.down(nn.silu(self.up(self.q(x))))
+
+    return _Block()
+
+
+def _bench_one(dtype: DtypeName, *, n_iter: int, hidden: int,
+               batch: int, seq: int) -> DtypeCostRow:
+    dtype_map = {"fp32": mx.float32, "bf16": mx.bfloat16,
+                 "fp16": mx.float16}
+    if dtype not in dtype_map:
+        return DtypeCostRow(dtype=dtype, supported=False,
+                             error=f"unknown dtype {dtype}")
+    target = dtype_map[dtype]
+    try:
+        model = _build_tiny(hidden)
+        t0 = time.perf_counter()
+        model.set_dtype(target)
+        mx.eval(model.parameters())
+        cast_ms = (time.perf_counter() - t0) * 1000.0
+
+        x = mx.random.normal(shape=(batch, seq, hidden),
+                              key=mx.random.key(0)).astype(target)
+        # Warm.
+        warm = model(x)
+        mx.eval(warm)
+        t0 = time.perf_counter()
+        for _ in range(n_iter):
+            out = model(x)
+            mx.eval(out)
+        fwd_ms = (time.perf_counter() - t0) * 1000.0 / n_iter
+
+        def _loss(m, _x):
+            return mx.mean(m(_x).astype(mx.float32) ** 2)
+
+        lvg = nn.value_and_grad(model, _loss)
+        warm_lvg = lvg(model, x)
+        mx.eval(warm_lvg)
+        t0 = time.perf_counter()
+        for _ in range(n_iter):
+            loss, grads = lvg(model, x)
+            mx.eval(loss, grads)
+        fwdbwd_ms = (time.perf_counter() - t0) * 1000.0 / n_iter
+
+        tokens = batch * seq
+        return DtypeCostRow(
+            dtype=dtype, supported=True,
+            cast_overhead_ms=round(cast_ms, 4),
+            fwd_ms=round(fwd_ms, 4),
+            fwdbwd_ms=round(fwdbwd_ms, 4),
+            fwd_ms_per_token=round(fwd_ms / tokens, 6),
+            fwdbwd_ms_per_token=round(fwdbwd_ms / tokens, 6),
+        )
+    except Exception as exc:
+        return DtypeCostRow(dtype=dtype, supported=False, error=str(exc))
+
+
+def dtype_cost_estimate(
+    params: DtypeCostParams, *, cache: LRUCache | None = None,
+) -> DtypeCostResult:
+    dtype_names: tuple[DtypeName, ...] = ("fp32", "bf16", "fp16")
+    rows = [
+        _bench_one(dt, n_iter=params.n_iter, hidden=params.hidden,
+                   batch=params.batch, seq=params.seq)
+        for dt in dtype_names
+    ]
+    return DtypeCostResult(
+        rows=rows, n_iter=params.n_iter,
+        shape=[params.batch, params.seq, params.hidden],
+    )
+
+
+__all__ = ["DtypeCostParams", "DtypeCostResult", "DtypeCostRow",
+            "dtype_cost_estimate"]

cppmega_v4/jsonrpc/schema.py

@@ -750,6 +750,9 @@ class PipelineRunResult(BaseModel):
     "data.roundtrip_check",
     "ablation.run",
     "ckpt.inspect",
+    "dtype.cost_estimate",
+    "pipeline.pause",
+    "pipeline.resume",
 })
 
 

cppmega_v4/runner/stages.py

@@ -1421,7 +1421,10 @@ def _base(k: str) -> str | None:
         val_losses: list[float] = []
         # G09: check abort flag set via opts.abort or _ABORT_TOKENS set
         abort_token = opts.get("abort_token")
+        # V7-H06: block while job is paused via job_control.pause(token).
+        from cppmega_v4.runtime.job_control import wait_while_paused
         for step in range(n_steps):
+            wait_while_paused(abort_token, poll_s=0.05, max_wait_s=600.0)
             if abort_token is not None and abort_token in _ABORT_TOKENS:
                 # Stop early; return partial extras with cancellation flag.
                 return _cancelled_train_result(

tests/v4/test_dtype_cost_method.py

@@ -0,0 +1,47 @@
+"""V7-D06: dtype.cost_estimate RPC round-trip."""
+
+from __future__ import annotations
+
+import pytest
+from fastapi.testclient import TestClient
+
+from cppmega_v4.jsonrpc import create_app
+
+
+@pytest.fixture
+def client():
+    return TestClient(create_app(cache_capacity=2))
+
+
+def test_dtype_cost_estimate_returns_three_rows(client):
+    payload = {
+        "jsonrpc": "2.0", "id": "d1", "method": "dtype.cost_estimate",
+        "params": {"n_iter": 2, "hidden": 32, "batch": 1, "seq": 4},
+    }
+    r = client.post("/rpc", json=payload)
+    assert r.status_code == 200
+    body = r.json()
+    assert "error" not in body, body
+    res = body["result"]
+    assert {row["dtype"] for row in res["rows"]} == {"fp32", "bf16", "fp16"}
+    for row in res["rows"]:
+        if row["supported"]:
+            assert row["fwd_ms"] is not None and row["fwd_ms"] >= 0
+            assert row["fwd_ms_per_token"] is not None
+            assert row["fwd_ms_per_token"] >= 0
+            assert row["fwdbwd_ms"] is not None
+        else:
+            # Even when unsupported, the row still appears and carries
+            # an error string the UI can show.
+            assert row["error"] is not None
+
+
+def test_dtype_cost_estimate_rejects_huge_params(client):
+    payload = {
+        "jsonrpc": "2.0", "id": "d2", "method": "dtype.cost_estimate",
+        "params": {"n_iter": 10000},
+    }
+    r = client.post("/rpc", json=payload)
+    body = r.json()
+    assert "error" in body
+    assert body["error"]["code"] == -32602

tests/v4/test_pipeline_pause_resume_rpc.py

@@ -0,0 +1,53 @@
+"""V7-H06: pipeline.pause / pipeline.resume RPC dispatcher round-trips.
+
+These are the dispatcher-side handlers. The end-to-end pause-during-
+train flow is covered by the e2e Playwright test (84_pause_resume).
+"""
+
+from __future__ import annotations
+
+import pytest
+from fastapi.testclient import TestClient
+
+from cppmega_v4.jsonrpc import create_app
+from cppmega_v4.runtime import job_control
+
+
+@pytest.fixture(autouse=True)
+def _clean_state():
+    job_control.reset()
+    yield
+    job_control.reset()
+
+
+@pytest.fixture
+def client():
+    return TestClient(create_app(cache_capacity=2))
+
+
+def test_pipeline_pause_marks_token_paused(client):
+    payload = {"jsonrpc": "2.0", "id": "p1", "method": "pipeline.pause",
+               "params": {"run_id": "run-abc"}}
+    r = client.post("/rpc", json=payload)
+    assert r.status_code == 200
+    body = r.json()
+    assert body["result"]["run_id"] == "run-abc"
+    assert job_control.is_paused("run-abc") is True
+
+
+def test_pipeline_resume_clears_pause(client):
+    job_control.pause("run-xyz")
+    assert job_control.is_paused("run-xyz") is True
+    payload = {"jsonrpc": "2.0", "id": "r1", "method": "pipeline.resume",
+               "params": {"run_id": "run-xyz"}}
+    r = client.post("/rpc", json=payload)
+    assert r.status_code == 200
+    assert r.json()["result"]["run_id"] == "run-xyz"
+    assert job_control.is_paused("run-xyz") is False
+
+
+def test_pause_resume_methods_appear_in_registry(client):
+    r = client.get("/schema/methods")
+    methods = set(r.json()["methods"])
+    assert "pipeline.pause" in methods
+    assert "pipeline.resume" in methods

vbgui/e2e/scenarios/84_dtype_cost_display.spec.ts

@@ -0,0 +1,29 @@
+// V7-D06: master_dtype dropdown renders real ms/token cost from
+// dtype.cost_estimate RPC. Honest-closure: previously the bench script
+// produced a static CSV/HTML report but the live UI showed only dtype
+// names — now each option carries " · X.XXXX ms/tok" next to its label.
+
+import { test, expect } from "@playwright/test";
+import { gotoApp, selectPreset } from "../fixtures";
+
+test("V7-D06: master_dtype options show ms/tok cost after dropdown opens",
+  async ({ page }) => {
+    test.setTimeout(60_000);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+
+    await page.getByTestId("run-pipeline-toggle").click();
+    // Probe fires once when the menu first opens — wait for either
+    // measured indicator OR the loading indicator to disappear.
+    await expect(page.getByTestId("dtype-cost-summary"))
+      .toBeVisible({ timeout: 30_000 });
+
+    const select = page.getByTestId("top-bar-precision-mode");
+    const fp32Text = await select.locator("option[value='fp32']").textContent();
+    const bf16Text = await select.locator("option[value='bf16']").textContent();
+    const fp16Text = await select.locator("option[value='fp16']").textContent();
+    // Each supported dtype carries the ms/tok suffix.
+    expect(fp32Text).toMatch(/ms\/tok/);
+    expect(bf16Text).toMatch(/ms\/tok/);
+    expect(fp16Text).toMatch(/ms\/tok/);
+  });

vbgui/e2e/scenarios/85_pause_resume.spec.ts

@@ -0,0 +1,49 @@
+// V7-H06: pause / resume mid-train. UI button toggles between Pause
+// and Resume, calls pipeline.pause / pipeline.resume RPCs, and the
+// backend job_control module blocks the train loop until resumed.
+
+import { test, expect } from "@playwright/test";
+import { gotoApp, selectPreset, closeModal } from "../fixtures";
+
+test("V7-H06: pause button toggles label and disables when no run is active",
+  async ({ page }) => {
+    test.setTimeout(30_000);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+
+    // No train running yet → pause button visible but disabled.
+    const pauseBtn = page.getByTestId("run-pipeline-pause");
+    await expect(pauseBtn).toBeVisible();
+    await expect(pauseBtn).toBeDisabled();
+    await expect(pauseBtn).toHaveText(/Pause/i);
+  });
+
+test("V7-H06: long-running train can be paused then resumed",
+  async ({ page }) => {
+    test.setTimeout(120_000);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+
+    await page.getByTestId("run-pipeline-toggle").click();
+    // Use enough steps so we have time to hit pause mid-run.
+    await page.getByTestId("train-num-steps").fill("8");
+    await page.getByTestId("run-pipeline-train").click();
+
+    const pauseBtn = page.getByTestId("run-pipeline-pause");
+    // Wait until trainInFlight engages — pause button becomes enabled.
+    await expect(pauseBtn).toBeEnabled({ timeout: 10_000 });
+
+    await pauseBtn.click();
+    // Button label flips to Resume.
+    await expect(pauseBtn).toHaveText(/Resume/i, { timeout: 5_000 });
+
+    // After short hold, click Resume.
+    await page.waitForTimeout(500);
+    await pauseBtn.click();
+    await expect(pauseBtn).toHaveText(/Pause/i, { timeout: 5_000 });
+
+    // Eventually the modal appears with status ok.
+    const modal = page.getByTestId("run-result-modal");
+    await modal.waitFor({ timeout: 120_000 });
+    await closeModal(page);
+  });