fix(v7-c02): cache bug + backward-compat single-file + opened-set assertions

Closes AC gaps in V7-C02 audit:

- Bug fix: load_sharded_for_rank used 'out.get("__loaded__", set())'
  where out is dict[str, mx.array] — re-opened the same shard once
  per parameter (O(num_params) instead of O(num_owned_shards)).
  Replaced with explicit seen_shards set keyed on shard filename.
- AC#3: new load_with_backward_compat accepts legacy single-file
  safetensors checkpoints AND sharded index.json AND a directory
  containing model.index.json.
- AC#2: both loaders accept a test-only _opened list sink so tests
  can assert exactly which shard files were opened (and not opened).
- _shard_index_from_name helper deduplicates the index-parsing logic.

5 new pytest cover: owned-shards-only opened, no duplicate opens for
load_sharded, legacy single-file roundtrip, sharded via index, dir
discovery, missing-dir error.

bd: cppmega-mlx-9vv (AC#2/AC#3 closure)
11/11 shard tests + 4/4 topology regression all green.

Author: apstenku123

cppmega_v4/runtime/checkpoint_shard.py

@@ -63,48 +63,103 @@ def save_sharded(state: dict[str, mx.array],
     return index
 
 
-def load_sharded(index_path: str | pathlib.Path) -> dict[str, mx.array]:
+def _shard_index_from_name(shard_name: str) -> int | None:
+    """Parse shard index from '<prefix>.shard-NNNNN-of-MMMMM.safetensors'."""
+    try:
+        return int(shard_name.split("-")[-3]) - 1
+    except (ValueError, IndexError):
+        return None
+
+
+def load_sharded(
+    index_path: str | pathlib.Path,
+    *, _opened: list[str] | None = None,
+) -> dict[str, mx.array]:
+    """Read every shard listed in the index manifest into a single dict.
+
+    ``_opened`` is a test-only sink — when supplied, each shard's absolute
+    path is appended exactly once. Lets tests assert which files were
+    actually opened on disk (V7-C02 AC#2)."""
     index_path = pathlib.Path(index_path)
     index = json.loads(index_path.read_text())
     base = index_path.parent
     out: dict[str, mx.array] = {}
     seen_shards: set[str] = set()
-    for k, shard_name in index["weight_map"].items():
+    for shard_name in index["weight_map"].values():
         if shard_name in seen_shards:
             continue
         seen_shards.add(shard_name)
-        out.update(st.load_file(str(base / shard_name)))
+        path = str(base / shard_name)
+        if _opened is not None:
+            _opened.append(path)
+        out.update(st.load_file(path))
     return out
 
 
-def load_sharded_for_rank(index_path: str | pathlib.Path,
-                           *, rank: int, world_size: int
-                           ) -> dict[str, mx.array]:
+def load_sharded_for_rank(
+    index_path: str | pathlib.Path,
+    *, rank: int, world_size: int,
+    _opened: list[str] | None = None,
+) -> dict[str, mx.array]:
     """Round-robin shard ownership: rank r reads shards i where
-    i % world_size == r."""
+    i % world_size == r. Each owned shard is opened **exactly once**.
+
+    ``_opened`` is the same test-only sink as ``load_sharded``."""
     if world_size <= 0 or not (0 <= rank < world_size):
-        raise ValueError("invalid rank/world_size")
+        raise ValueError(
+            f"invalid rank/world_size: rank={rank}, world_size={world_size}")
     index_path = pathlib.Path(index_path)
     index = json.loads(index_path.read_text())
     base = index_path.parent
     num_shards = int(index["metadata"]["num_shards"])
     owned = {i for i in range(num_shards) if i % world_size == rank}
     out: dict[str, mx.array] = {}
-    for k, shard_name in index["weight_map"].items():
-        # Reverse-derive shard index from filename "*-NNNNN-of-MMMMM"
-        try:
-            idx = int(shard_name.split("-")[-3]) - 1
-        except (ValueError, IndexError):
+    # Deduplicate shard filenames before opening — the previous impl
+    # iterated weight_map keys and re-opened the same file once per
+    # tensor, which scaled O(num_params) instead of O(num_owned_shards).
+    seen_shards: set[str] = set()
+    for shard_name in index["weight_map"].values():
+        if shard_name in seen_shards:
             continue
-        if idx in owned:
-            if shard_name not in {n for n in
-                                  (index["weight_map"][kk]
-                                   for kk in index["weight_map"])}:
-                continue
-            # Load each shard once.
-            if shard_name not in out.get("__loaded__", set()):
-                out.update(st.load_file(str(base / shard_name)))
+        idx = _shard_index_from_name(shard_name)
+        if idx is None or idx not in owned:
+            continue
+        seen_shards.add(shard_name)
+        path = str(base / shard_name)
+        if _opened is not None:
+            _opened.append(path)
+        out.update(st.load_file(path))
     return out
 
 
-__all__ = ["save_sharded", "load_sharded", "load_sharded_for_rank"]
+def load_with_backward_compat(
+    path: str | pathlib.Path,
+    *, _opened: list[str] | None = None,
+) -> dict[str, mx.array]:
+    """V7-C02 AC#3: accept either a sharded index.json OR a legacy
+    single-file safetensors checkpoint.
+
+    If ``path`` ends in ``.index.json`` (or names an index.json file
+    that exists), delegate to :func:`load_sharded`. Otherwise treat the
+    path as a flat safetensors file and read it directly.
+    """
+    p = pathlib.Path(path)
+    if p.is_dir():
+        # Convenience: caller passed a checkpoint directory — pick the
+        # canonical model.index.json.
+        cand = p / "model.index.json"
+        if cand.is_file():
+            return load_sharded(cand, _opened=_opened)
+        raise FileNotFoundError(f"no model.index.json under {p}")
+    if p.name.endswith(".index.json"):
+        return load_sharded(p, _opened=_opened)
+    # Legacy single-file path.
+    if _opened is not None:
+        _opened.append(str(p))
+    return st.load_file(str(p))
+
+
+__all__ = [
+    "save_sharded", "load_sharded", "load_sharded_for_rank",
+    "load_with_backward_compat",
+]

tests/v4/test_checkpoint_shard.py

@@ -7,8 +7,12 @@
 import mlx.core as mx
 import pytest
 
+import mlx.core as mx
+import safetensors.mlx as st
+
 from cppmega_v4.runtime.checkpoint_shard import (
-    load_sharded, load_sharded_for_rank, save_sharded,
+    load_sharded, load_sharded_for_rank,
+    load_with_backward_compat, save_sharded,
 )
 
 
@@ -67,3 +71,72 @@ def test_v7_c02_load_sharded_for_rank_round_robin(tmp_path):
 def test_v7_c02_num_shards_validation(tmp_path):
     with pytest.raises(ValueError):
         save_sharded({"k": mx.zeros((2,))}, tmp_path, num_shards=0)
+
+
+def test_v7_c02_load_sharded_for_rank_opens_owned_only_once(tmp_path):
+    """AC#2: rank R must open ONLY its own shards, each exactly once.
+
+    Previously the impl re-opened the same shard once per parameter
+    (buggy `out.get('__loaded__')` check) — opened count scaled with
+    num_params instead of num_owned_shards."""
+    save_sharded(_state(), tmp_path, num_shards=4)
+    opened: list[str] = []
+    out = load_sharded_for_rank(
+        tmp_path / "model.index.json",
+        rank=0, world_size=2, _opened=opened)
+    # rank 0 / world 2 owns shards 0 and 2 (i % 2 == 0).
+    assert len(opened) == 2, opened
+    assert all("shard-00001-of-00004" in p or "shard-00003-of-00004" in p
+               for p in opened)
+    # Other shards (1, 3) must NOT have been opened.
+    assert not any("shard-00002-of-00004" in p for p in opened)
+    assert not any("shard-00004-of-00004" in p for p in opened)
+    # Coverage still right for rank's owned subset.
+    assert len(out) > 0
+
+
+def test_v7_c02_load_sharded_opens_each_shard_exactly_once(tmp_path):
+    save_sharded(_state(), tmp_path, num_shards=3)
+    opened: list[str] = []
+    load_sharded(tmp_path / "model.index.json", _opened=opened)
+    assert len(opened) == 3
+    # No duplicates.
+    assert len(set(opened)) == 3
+
+
+def test_v7_c02_backward_compat_loads_legacy_single_file(tmp_path):
+    """AC#3: single-file safetensors checkpoints saved before the
+    sharded format still load via load_with_backward_compat."""
+    state = _state()
+    legacy = tmp_path / "legacy.safetensors"
+    st.save_file(state, str(legacy))
+    opened: list[str] = []
+    out = load_with_backward_compat(legacy, _opened=opened)
+    assert set(out.keys()) == set(state.keys())
+    for k in state:
+        assert mx.allclose(out[k], state[k], atol=0.0)
+    assert opened == [str(legacy)]
+
+
+def test_v7_c02_backward_compat_loads_sharded_via_index(tmp_path):
+    save_sharded(_state(), tmp_path, num_shards=2)
+    opened: list[str] = []
+    out = load_with_backward_compat(
+        tmp_path / "model.index.json", _opened=opened)
+    assert set(out.keys()) == set(_state().keys())
+    # 2 shards opened, no duplicates.
+    assert len(opened) == 2
+    assert len(set(opened)) == 2
+
+
+def test_v7_c02_backward_compat_directory_picks_index(tmp_path):
+    save_sharded(_state(), tmp_path, num_shards=2)
+    out = load_with_backward_compat(tmp_path)
+    assert set(out.keys()) == set(_state().keys())
+
+
+def test_v7_c02_backward_compat_missing_directory_raises(tmp_path):
+    empty = tmp_path / "empty"
+    empty.mkdir()
+    with pytest.raises(FileNotFoundError):
+        load_with_backward_compat(empty)