fix(v7-i03): _RUN_CACHE locked helpers + race-safe warm-start

Closes plan item #44 / bd cppmega-mlx-lak7.

stages.py _RUN_CACHE was a bare dict accessed unsynchronised from
both the train loop and the API thread. On Python 3.13+ where dict
ops are no longer atomic over compound get/set/eviction sequences a
concurrent reader could observe a torn opt.state or hit KeyError
mid-eviction.

- Add module-level threading.Lock guarding _RUN_CACHE.
- New helpers _run_cache_get / _run_cache_set / _run_cache_contains
  hold the lock only across the dict op itself — the train step
  outside the helper is never serialised.
- _run_cache_set takes the LRU cap (default 32, callers pass 8 to
  preserve the prior behaviour).
- Replace the two raw _RUN_CACHE access sites (line ~958 warm-start
  load, line ~2009 post-train cache) with the helpers.

3/3 pytest cover: 1000 mixed reads+writes from 8 threads / 2 threads
running warm-start train concurrently against the same cached
run_id / cache eviction returns None not KeyError for stale keys.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -955,12 +955,16 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
         opt_state_carried = False
         run_id = str(opts.get("run_id") or id(opt))
         continue_from = opts.get("continue_from_run_id")
-        if continue_from and continue_from in _RUN_CACHE:
-            try:
-                opt.state = _RUN_CACHE[continue_from]
-                opt_state_carried = True
-            except Exception:
-                pass
+        if continue_from:
+            # V7-I03: single locked read, then assign outside the lock so
+            # the train step never serialises on the cache.
+            cached_state = _run_cache_get(continue_from)
+            if cached_state is not None:
+                try:
+                    opt.state = cached_state
+                    opt_state_carried = True
+                except Exception:
+                    pass
         # V4-9: when hybrid, count params routed to each bucket so e2e can
         # prove the split predicate actually saw 2D vs 1D/3D parameters.
         muon_group_size: int | None = None
@@ -2000,11 +2004,11 @@ def _compiled_step_value_and_grad(_emb, _targets):
             )
             opt.state = nn.utils.tree_unflatten(sorted_state_pairs)
 
-        # G10: cache opt.state for future warm-start lookups (capped LRU)
+        # G10 / V7-I03: cache opt.state for future warm-start lookups
+        # via the locked helper so a concurrent reader can't observe a
+        # partial dict mid-eviction.
         try:
-            _RUN_CACHE[run_id] = opt.state
-            if len(_RUN_CACHE) > 8:
-                _RUN_CACHE.pop(next(iter(_RUN_CACHE)))
+            _run_cache_set(run_id, opt.state, cap=8)
         except Exception:
             pass
 
@@ -2465,7 +2469,37 @@ def _validate_parquet_stream_tokenizer_fingerprints(
 
 # G10: in-process LRU cache of opt.state by run_id. Used for warm-start
 # across sequential Train clicks in the same backend session.
+# V7-I03: _RUN_CACHE is read+written from the train loop AND from
+# concurrent UI clicks (two pipeline.run calls landing simultaneously
+# in different worker threads). dict mutation is not atomic across
+# the get/set pair in stage_train, so race readers could observe a
+# partially-populated opt.state. The lock is held only during the
+# get/set/eviction operations — never across the train step itself
+# (which would serialise actual training).
+import threading as _threading
 _RUN_CACHE: dict[str, Any] = {}
+_RUN_CACHE_LOCK = _threading.Lock()
+
+
+def _run_cache_get(key: str) -> Any | None:
+    """V7-I03: thread-safe read of _RUN_CACHE."""
+    with _RUN_CACHE_LOCK:
+        return _RUN_CACHE.get(key)
+
+
+def _run_cache_set(key: str, value: Any, *, cap: int = 32) -> None:
+    """V7-I03: thread-safe write with LRU-style eviction."""
+    with _RUN_CACHE_LOCK:
+        _RUN_CACHE[key] = value
+        # Drop oldest entries past cap (insertion order in py3.7+).
+        while len(_RUN_CACHE) > cap:
+            oldest = next(iter(_RUN_CACHE))
+            _RUN_CACHE.pop(oldest, None)
+
+
+def _run_cache_contains(key: str) -> bool:
+    with _RUN_CACHE_LOCK:
+        return key in _RUN_CACHE
 
 # G09: in-process set of abort tokens. Caller sets opts.abort_token to
 # some unique string; another caller (e.g. WS handler) inserts the same

tests/v4/test_concurrent_train_no_race.py

@@ -0,0 +1,144 @@
+"""V7-I03 (cppmega-mlx-lak7): _RUN_CACHE warm-start is race-safe.
+
+Two threads launch pipeline.run with continue_from=<same run_id>
+that's already cached. Neither should observe a partially-evicted
+or torn opt.state (i.e. both stage_train calls succeed end-to-end
+without exception).
+
+Pre-fix the cache was a bare dict mutated from the train loop +
+the API thread without synchronisation; under load the eviction
+path could pop an entry mid-read on Python 3.13+ where dict
+operations are no longer GIL-atomic for compound operations.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+
+import pytest
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+from cppmega_v4.runner.stages import (
+    _RUN_CACHE, _RUN_CACHE_LOCK,
+    _run_cache_contains, _run_cache_get, _run_cache_set,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_cache():
+    with _RUN_CACHE_LOCK:
+        _RUN_CACHE.clear()
+    yield
+    with _RUN_CACHE_LOCK:
+        _RUN_CACHE.clear()
+
+
+def _spec() -> VerifyParams:
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention",
+                 "params": {"num_heads": 4, "head_dim": 64}},
+                {"id": "mlp", "kind": "mlp", "params": {}},
+            ],
+            "edges": [{"src": "attn", "dst": "mlp"}],
+        },
+        "dim_env": {"B": 1, "S": 8, "H": 128,
+                    "nh": 2, "nkv": 1, "head_dim": 64},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["mlp"]},
+        "optim": {"kind": "adamw",
+                  "groups": [{"matcher": "all", "lr": 1e-3,
+                              "weight_decay": 0.01,
+                              "betas": [0.9, 0.95]}]},
+    })
+
+
+def _train(opts: dict) -> dict:
+    rep = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": opts},
+    }))
+    tr = next(s for s in rep.stages if s.name == "train")
+    return {"status": tr.status, "error": tr.error,
+            "extras": tr.extras}
+
+
+def test_v7_i03_run_cache_helpers_are_thread_safe_under_load():
+    """1000 mixed reads + writes from 8 threads — no exception, no
+    corrupted entries (every value retrieved equals what was set)."""
+    errors: list[BaseException] = []
+
+    def _hammer(thread_id: int) -> None:
+        try:
+            for i in range(125):
+                key = f"k-{thread_id}-{i}"
+                payload = {"thread": thread_id, "i": i, "data": [i] * 4}
+                _run_cache_set(key, payload, cap=64)
+                if _run_cache_contains(key):
+                    got = _run_cache_get(key)
+                    # cap=64 → many entries evicted by other threads;
+                    # if the key is still present its value must be
+                    # exactly what we wrote, never a torn merge.
+                    if got is not None and got != payload:
+                        raise AssertionError(
+                            f"torn read: set={payload} got={got}")
+        except BaseException as exc:
+            errors.append(exc)
+
+    threads = [threading.Thread(target=_hammer, args=(t,))
+                for t in range(8)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=10.0)
+    assert not errors, errors
+
+
+def test_v7_i03_concurrent_train_warm_start_completes_both(tmp_path):
+    """First train populates _RUN_CACHE['shared']. Then two threads
+    each launch a fresh train with continue_from='shared' — both
+    must finish status='ok' (no exception from a torn read)."""
+    base = _train({"num_steps": 2, "run_id": "shared"})
+    assert base["status"] == "ok"
+    assert _run_cache_contains("shared")
+
+    results: list[dict] = []
+    errors: list[BaseException] = []
+
+    def _worker(idx: int) -> None:
+        try:
+            r = _train({
+                "num_steps": 2, "run_id": f"warm-{idx}",
+                "continue_from_run_id": "shared",
+            })
+            results.append(r)
+        except BaseException as exc:
+            errors.append(exc)
+
+    t1 = threading.Thread(target=_worker, args=(0,))
+    t2 = threading.Thread(target=_worker, args=(1,))
+    t1.start(); t2.start()
+    t1.join(timeout=30.0); t2.join(timeout=30.0)
+    assert not errors, errors
+    assert len(results) == 2
+    for r in results:
+        assert r["status"] == "ok", r["error"]
+        # Both runs report opt-state actually carried from the shared
+        # warm-start entry.
+        assert r["extras"].get("opt_state_carried") is True
+
+
+def test_v7_i03_cache_eviction_does_not_strand_concurrent_readers():
+    """Pre-fix path: reader could call cache[k] right after writer
+    popped k via eviction → KeyError. Post-fix: _run_cache_get returns
+    None safely. Pin that contract."""
+    for i in range(100):
+        _run_cache_set(f"k{i}", {"i": i}, cap=4)
+    # cap=4 enforced → only ~last 4 keys remain.
+    surviving = sum(1 for i in range(100)
+                     if _run_cache_contains(f"k{i}"))
+    assert surviving <= 4
+    # Stale key returns None, not KeyError.
+    assert _run_cache_get("k0") is None