Skip to content

Commit 4bc64d1

Browse files
committed
feat(e-audit-03): checkpoint history UI + compress/strict toggles + mid-run trigger
V7-Q03 (E-AUDIT-03, cppmega-mlx-9avx) — Lane 6 audit gap closure. Q03.1 ckpt.list_history RPC: scans dir recursively for *.safetensors, reads metadata via read_ckpt_metadata, sorts by mtime desc, caps at max_entries. Returns path/mtime/size/arch_hash/opt_kind/global_step/ has_opt_sidecar. Registered in dispatcher. Q03.2 CheckpointHistoryDropdown component: lazy-fetch on first open, click row -> onSelect(path) fills train-checkpoint-load-path. Wired into TopBar (renders only when rpc prop provided). Q03.3 Compress + strict toggles: train-opt-compress dropdown (none/ weights-int8/opt-fp16/both), train-opt-ckpt-strict checkbox, train-opt-opt-state-strict checkbox. Threaded through onRunPipeline to App.tsx then stage_train via stage_options.train. Q03.4 V7-H05 mid-run checkpoint trigger: new stages.request_checkpoint + consume_checkpoint_trigger; train loop drains the queue between steps and saves weights + opt-state sidecar via the existing _save_compressed path. Tests (all green): - tests/v4/test_ckpt_list_history.py: 7 cases (empty/metadata/ sidecar/sort/cap/missing/recursive) - tests/v4/test_ckpt_trigger_mid_run.py: 3 cases (queue roundtrip, unknown-token, full end-to-end save during train) - vbgui/tests/CheckpointHistoryDropdown.test.tsx: 5 cases Regression: backend pytest checkpoint+train 191/191 PASS (81s); vbgui vitest 428->440 (+12) PASS. Closes Lane 6 P1 audit gap from docs/UI-TO-TRAIN-AUDIT-2026-05-23.md.
1 parent af1f62b commit 4bc64d1

9 files changed

Lines changed: 747 additions & 1 deletion

File tree

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
"""V7-Q03.1: ckpt.list_history RPC.
2+
3+
Scans a directory recursively for ``*.safetensors`` files, reads
4+
self-describing metadata (header-only, no full tensor load) and
5+
returns a sorted-by-mtime descending list capped at 100. Pairs with
6+
``ckpt.inspect`` for the single-file inspector view and feeds the
7+
``CheckpointHistoryDropdown`` UI component that lets the operator pick
8+
a past checkpoint for resume without copy-pasting paths.
9+
10+
Closes Lane 6 audit gap from docs/UI-TO-TRAIN-AUDIT-2026-05-23.md.
11+
"""
12+
13+
from __future__ import annotations
14+
15+
from pathlib import Path
16+
from typing import Any
17+
18+
from pydantic import BaseModel, ConfigDict, Field
19+
20+
from cppmega_v4.jsonrpc.cache import LRUCache
21+
22+
23+
class CkptListHistoryParams(BaseModel):
24+
model_config = ConfigDict(extra="forbid")
25+
26+
directory: str = "."
27+
# Cap to prevent runaway scans on large workspaces.
28+
max_entries: int = 100
29+
30+
31+
class CkptHistoryEntry(BaseModel):
32+
model_config = ConfigDict(extra="allow")
33+
34+
path: str
35+
mtime: float
36+
size_bytes: int
37+
arch_hash: str | None = None
38+
opt_kind: str | None = None
39+
global_step: int | None = None
40+
has_opt_sidecar: bool = False
41+
42+
43+
class CkptListHistoryResult(BaseModel):
44+
model_config = ConfigDict(extra="allow")
45+
46+
directory: str
47+
scanned: int = 0
48+
entries: list[CkptHistoryEntry] = Field(default_factory=list)
49+
error: str | None = None
50+
51+
52+
def ckpt_list_history(
53+
params: CkptListHistoryParams, *, cache: LRUCache | None = None,
54+
) -> CkptListHistoryResult:
55+
root = Path(params.directory).expanduser()
56+
try:
57+
root = root.resolve()
58+
except Exception as exc: # noqa: BLE001 -- surface to caller
59+
return CkptListHistoryResult(
60+
directory=str(params.directory), error=f"resolve failed: {exc}",
61+
)
62+
if not root.exists():
63+
return CkptListHistoryResult(
64+
directory=str(root), error="directory does not exist",
65+
)
66+
if not root.is_dir():
67+
return CkptListHistoryResult(
68+
directory=str(root), error="path is not a directory",
69+
)
70+
71+
from cppmega_v4.runner.stages import read_ckpt_metadata
72+
73+
raw_entries: list[CkptHistoryEntry] = []
74+
scanned = 0
75+
for path in root.rglob("*.safetensors"):
76+
# Skip opt-state sidecars — they're surfaced via has_opt_sidecar.
77+
name = path.name
78+
if name.endswith(".opt.safetensors") or name.endswith(".opt"):
79+
continue
80+
scanned += 1
81+
try:
82+
stat = path.stat()
83+
meta: dict[str, Any] | None = None
84+
try:
85+
meta = read_ckpt_metadata(str(path))
86+
except Exception:
87+
meta = None
88+
arch = (meta or {}).get("arch")
89+
train = (meta or {}).get("train")
90+
opt = (meta or {}).get("opt")
91+
arch_hash = (
92+
arch.get("config_hash") if isinstance(arch, dict) else None
93+
)
94+
opt_kind = (
95+
opt.get("kind") if isinstance(opt, dict) else None
96+
)
97+
global_step = (
98+
train.get("global_step") if isinstance(train, dict) else None
99+
)
100+
sidecar = path.with_suffix(path.suffix + ".opt")
101+
raw_entries.append(CkptHistoryEntry(
102+
path=str(path),
103+
mtime=stat.st_mtime,
104+
size_bytes=stat.st_size,
105+
arch_hash=arch_hash,
106+
opt_kind=opt_kind,
107+
global_step=global_step,
108+
has_opt_sidecar=sidecar.is_file(),
109+
))
110+
except Exception:
111+
# Unreadable file — skip silently.
112+
continue
113+
114+
raw_entries.sort(key=lambda e: e.mtime, reverse=True)
115+
capped = raw_entries[: max(1, int(params.max_entries))]
116+
return CkptListHistoryResult(
117+
directory=str(root), scanned=scanned, entries=capped,
118+
)
119+
120+
121+
__all__ = [
122+
"CkptListHistoryParams",
123+
"CkptHistoryEntry",
124+
"CkptListHistoryResult",
125+
"ckpt_list_history",
126+
]

cppmega_v4/jsonrpc/dispatcher.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,9 @@
4646
from cppmega_v4.jsonrpc.ckpt_inspect_method import (
4747
CkptInspectParams, ckpt_inspect,
4848
)
49+
from cppmega_v4.jsonrpc.ckpt_history_method import (
50+
CkptListHistoryParams, ckpt_list_history,
51+
)
4952
from cppmega_v4.jsonrpc.dtype_cost_method import (
5053
DtypeCostParams, dtype_cost_estimate,
5154
)
@@ -176,6 +179,10 @@
176179
CkptInspectParams,
177180
lambda p, c: ckpt_inspect(p, cache=c),
178181
),
182+
"ckpt.list_history": (
183+
CkptListHistoryParams,
184+
lambda p, c: ckpt_list_history(p, cache=c),
185+
),
179186
"dtype.cost_estimate": (
180187
DtypeCostParams,
181188
lambda p, c: dtype_cost_estimate(p, cache=c),

cppmega_v4/runner/stages.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1895,6 +1895,46 @@ def _compiled_step_value_and_grad(_emb, _targets):
18951895
val_losses.append(vL)
18961896
except Exception:
18971897
pass
1898+
# V7-Q03.4: drain mid-run checkpoint trigger queue. WS
1899+
# handler can request a save without aborting; we save via
1900+
# the same _save_compressed path as end-of-run + sidecar.
1901+
if abort_token is not None:
1902+
trig_path = consume_checkpoint_trigger(str(abort_token))
1903+
if trig_path:
1904+
try:
1905+
from cppmega_v4.runtime.checkpoint_quantize import (
1906+
save_state_compressed as _save_mid,
1907+
)
1908+
mid_flat = dict(nn.utils.tree_flatten(
1909+
all_modules.parameters()))
1910+
mid_meta = _build_ckpt_metadata(
1911+
ctx=ctx, optimizer_kind=optimizer_kind,
1912+
n_steps=step + 1, lr=lr,
1913+
)
1914+
_save_mid(mid_flat, trig_path,
1915+
compress="none",
1916+
metadata=mid_meta,
1917+
role="weights")
1918+
# Sidecar opt-state for full resumable bundle.
1919+
try:
1920+
import safetensors.mlx as _stmlx_mid
1921+
mid_opt_flat = dict(nn.utils.tree_flatten(
1922+
opt.state))
1923+
mid_opt_arrays = {
1924+
k: v for k, v in mid_opt_flat.items()
1925+
if hasattr(v, "shape")
1926+
}
1927+
if hasattr(rng_key, "shape"):
1928+
mid_opt_arrays["_rng_key"] = rng_key
1929+
_stmlx_mid.save_file(
1930+
mid_opt_arrays, trig_path + ".opt",
1931+
metadata=mid_meta,
1932+
)
1933+
except Exception:
1934+
pass
1935+
except Exception:
1936+
# Best-effort; don't fail the run on save error.
1937+
pass
18981938
if abort_token is not None and abort_token in _ABORT_TOKENS:
18991939
partial_delta = 0.0
19001940
if probe_key is not None and probe_before is not None:
@@ -2516,6 +2556,29 @@ def clear_abort(token: str) -> None:
25162556
_ABORT_TOKENS.discard(token)
25172557

25182558

2559+
# V7-Q03.4: in-process trigger-checkpoint queue. Caller (WS handler)
2560+
# inserts `{abort_token: save_path}` and the train loop drains it on
2561+
# the next step boundary, saving weights + opt-state sidecar via the
2562+
# existing _save_compressed path. Live mid-run checkpointing without
2563+
# affecting the post-train save_path.
2564+
_TRIGGER_CHECKPOINT_QUEUE: dict[str, str] = {}
2565+
2566+
2567+
def request_checkpoint(token: str, save_path: str) -> None:
2568+
"""V7-Q03.4: queue a mid-run checkpoint for the named run.
2569+
2570+
The train loop polls _TRIGGER_CHECKPOINT_QUEUE between steps; when
2571+
it finds an entry keyed by the active abort_token, it saves to
2572+
save_path (and the sidecar `<save_path>.opt`) then pops the entry.
2573+
"""
2574+
_TRIGGER_CHECKPOINT_QUEUE[token] = save_path
2575+
2576+
2577+
def consume_checkpoint_trigger(token: str) -> str | None:
2578+
"""Pop and return any pending checkpoint save_path for token."""
2579+
return _TRIGGER_CHECKPOINT_QUEUE.pop(token, None)
2580+
2581+
25192582
_REWRITER_FACTORIES: dict[str, Callable[..., Any]] = {}
25202583

25212584

tests/v4/test_ckpt_list_history.py

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
"""V7-Q03.1: ckpt.list_history RPC regression.
2+
3+
Pins: scan returns sorted-by-mtime descending list of safetensors files
4+
with metadata extracted via read_ckpt_metadata. Skips opt-state sidecars.
5+
Surfaces has_opt_sidecar flag when `<path>.opt` exists alongside.
6+
"""
7+
8+
from __future__ import annotations
9+
10+
import os
11+
import tempfile
12+
import time
13+
14+
import mlx.core as mx
15+
import safetensors.mlx as stmlx
16+
17+
from cppmega_v4.jsonrpc.ckpt_history_method import (
18+
CkptListHistoryParams, ckpt_list_history,
19+
)
20+
21+
22+
def _make_ckpt(path: str, *, with_sidecar: bool = False,
23+
arch_hash: str = "abc123def456", step: int = 7) -> None:
24+
weights = {"layers.0.weight": mx.zeros((4, 4))}
25+
meta = {
26+
"cppmega_version": "test",
27+
"arch": '{"config_hash": "' + arch_hash + '"}',
28+
"train": '{"global_step": ' + str(step) + '}',
29+
"opt": '{"kind": "adamw", "lr": 0.001}',
30+
}
31+
stmlx.save_file(weights, path, metadata=meta)
32+
if with_sidecar:
33+
opt_state = {"opt.0.m": mx.zeros((4, 4))}
34+
stmlx.save_file(opt_state, path + ".opt", metadata=meta)
35+
36+
37+
def test_list_history_empty_dir() -> None:
38+
with tempfile.TemporaryDirectory() as td:
39+
res = ckpt_list_history(CkptListHistoryParams(directory=td))
40+
assert res.error is None
41+
assert res.scanned == 0
42+
assert res.entries == []
43+
44+
45+
def test_list_history_returns_metadata() -> None:
46+
with tempfile.TemporaryDirectory() as td:
47+
p = os.path.join(td, "a.safetensors")
48+
_make_ckpt(p, arch_hash="aaa111", step=4)
49+
res = ckpt_list_history(CkptListHistoryParams(directory=td))
50+
assert res.error is None
51+
assert res.scanned == 1
52+
assert len(res.entries) == 1
53+
e = res.entries[0]
54+
assert e.path.endswith("a.safetensors")
55+
assert e.arch_hash == "aaa111"
56+
assert e.opt_kind == "adamw"
57+
assert e.global_step == 4
58+
assert e.size_bytes > 0
59+
assert e.has_opt_sidecar is False
60+
61+
62+
def test_list_history_detects_sidecar() -> None:
63+
with tempfile.TemporaryDirectory() as td:
64+
p = os.path.join(td, "withopt.safetensors")
65+
_make_ckpt(p, with_sidecar=True)
66+
res = ckpt_list_history(CkptListHistoryParams(directory=td))
67+
# Sidecar file is NOT a top-level entry but IS detected on the
68+
# parent's has_opt_sidecar flag.
69+
assert res.scanned == 1
70+
assert len(res.entries) == 1
71+
assert res.entries[0].has_opt_sidecar is True
72+
73+
74+
def test_list_history_sorted_by_mtime_desc() -> None:
75+
with tempfile.TemporaryDirectory() as td:
76+
a = os.path.join(td, "older.safetensors")
77+
b = os.path.join(td, "newer.safetensors")
78+
_make_ckpt(a)
79+
time.sleep(0.05) # mtime granularity
80+
_make_ckpt(b)
81+
res = ckpt_list_history(CkptListHistoryParams(directory=td))
82+
assert [e.path.split("/")[-1] for e in res.entries] == [
83+
"newer.safetensors", "older.safetensors",
84+
]
85+
86+
87+
def test_list_history_caps_at_max_entries() -> None:
88+
with tempfile.TemporaryDirectory() as td:
89+
for i in range(5):
90+
_make_ckpt(os.path.join(td, f"c{i}.safetensors"))
91+
res = ckpt_list_history(
92+
CkptListHistoryParams(directory=td, max_entries=3))
93+
assert res.scanned == 5
94+
assert len(res.entries) == 3
95+
96+
97+
def test_list_history_missing_dir() -> None:
98+
res = ckpt_list_history(
99+
CkptListHistoryParams(directory="/nonexistent/dir/xyzzy"))
100+
assert res.error is not None
101+
assert "does not exist" in res.error
102+
103+
104+
def test_list_history_recursive() -> None:
105+
with tempfile.TemporaryDirectory() as td:
106+
sub = os.path.join(td, "sub")
107+
os.makedirs(sub)
108+
_make_ckpt(os.path.join(td, "top.safetensors"))
109+
_make_ckpt(os.path.join(sub, "deep.safetensors"))
110+
res = ckpt_list_history(CkptListHistoryParams(directory=td))
111+
assert res.scanned == 2
112+
assert {e.path.split("/")[-1] for e in res.entries} == {
113+
"top.safetensors", "deep.safetensors",
114+
}

0 commit comments

Comments
 (0)