Add CachedPredictor sharding — v5.6.0 (#133) (closes #128)

iskandr · web-flow · commit bf01e573d7ea · 2026-04-14T20:32:08.000-04:00
Final part of #128 — pluggable prediction sources reach feature-complete. - concat([caches]) + from_directory(path) merge multiple CachedPredictors through the core (name, version) invariant. - Overlap resolution: 'raise' (default) / 'last' / 'first' / callable(row_a, row_b) -> row. - v5.6.0 bump accumulates #131 (vaxrank polish), #132 (NetMHC loaders), and this PR (sharding). Closes #128.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,57 @@
 # Changelog
 
+## 5.6.0
+
+**Closes #128 — `CachedPredictor` reaches feature-complete.**
+
+**New loaders for the DTU NetMHC suite (#132):**
+
+- `CachedPredictor.from_netmhcpan_stdout(path, mode=…)` — auto-detects
+  NetMHCpan 2.8 / 3 / 4 / 4.1. `mode` selects `"binding_affinity"` or
+  `"elution_score"` for 4+.
+- `CachedPredictor.from_netmhc_stdout(path, version=…)` — classic
+  NetMHC 3 / 4 / 4.1.
+- `CachedPredictor.from_netmhcpan_cons_stdout(path)` — NetMHCcons.
+- `CachedPredictor.from_netmhciipan_stdout(path, version=…)` —
+  NetMHCIIpan legacy / 4 / 4.3.
+- `CachedPredictor.from_netmhcstabpan_stdout(path)` — NetMHCstabpan
+  pMHC-stability predictor.
+
+Each loader wraps an existing `mhctools.parsing.*_stdout` function
+(zero new parsing code) and parses the tool version out of the
+stdout preamble onto `predictor_version`. Parses stdout text, not
+the `-xlsfile` tab-delimited variant — flagged in `docs/cached.md`.
+
+**Sharding — `concat` + `from_directory`:**
+
+- `CachedPredictor.concat([caches], on_overlap=…)` — merge several
+  caches into one. All shards must share `(name, version)` per the
+  core invariant.
+- `CachedPredictor.from_directory(path, pattern="*", on_overlap=…)` —
+  glob a directory and concat every matching file through
+  `from_topiary_output`.
+- Overlap resolution policies (`on_overlap`): `"raise"` (default — fail
+  if any `(peptide, allele, peptide_length)` appears in more than one
+  shard), `"last"` (later shard wins), `"first"` (earlier wins), or a
+  user-supplied `callable(row_a, row_b) -> row` resolver.
+
+**Polish from vaxrank-consumer review on #130 (#131):**
+
+- `_fallback_resolve` filters fallback output to keys not already in
+  the index before merging, so a partial-allele cache (peptide P
+  present for allele A, missing for B) doesn't see its `(P, A)` row
+  silently overwritten by the fallback's all-alleles response.
+- Class docstring now flags silent peptide-length lock-in and
+  non-thread-safety.
+- `save()` raises on an empty never-queried cache with no identity,
+  so users don't write schema-only files that can't be round-tripped.
+
+**Tests:**
+
+- 59 tests in `tests/test_cached_predictor.py` (up from 41): 6 NetMHC
+  loader tests, 12 sharding tests. Full suite 1111 passed (up from
+  1093).
+
 ## 5.5.0
 
 **New feature — `CachedPredictor`:**
diff --git a/docs/api.md b/docs/api.md
@@ -41,6 +41,8 @@ pre-computed table. Pass as `models=cache` to `TopiaryPredictor`. See
 | `CachedPredictor.from_netmhcpan_cons_stdout(path)` | NetMHCcons stdout. |
 | `CachedPredictor.from_netmhciipan_stdout(path, version=...)` | NetMHCIIpan stdout (legacy / 4 / 4.3). |
 | `CachedPredictor.from_netmhcstabpan_stdout(path)` | NetMHCstabpan stdout (pMHC stability). |
+| `CachedPredictor.concat([caches], on_overlap=...)` | Merge shards (all must share name+version). `on_overlap`: `"raise"` / `"last"` / `"first"` / callable. |
+| `CachedPredictor.from_directory(path, pattern="*", on_overlap=...)` | Glob a dir and concat every matching file. |
 | `CachedPredictor.from_tsv(path, columns=..., prediction_method_name=..., predictor_version=...)` | Generic tab- or comma-delimited. |
 | `CachedPredictor.from_dataframe(df, ...)` | In-memory DataFrame. |
 | `CachedPredictor(fallback=live_predictor)` | Empty cache, lazy identity discovery — pure read-through over a live model. |
diff --git a/docs/cached.md b/docs/cached.md
@@ -116,6 +116,47 @@ when the file doesn't embed that identity.
 
 Pass `sep=","` for CSV files.
 
+### Sharding: merge multiple caches
+
+Predict per-allele or per-sample in parallel, persist each shard
+separately, then merge them:
+
+```python
+cache = CachedPredictor.concat([shard_a, shard_b, shard_c])
+
+# Or: load every matching file from a directory
+cache = CachedPredictor.from_directory(
+    "caches/",
+    pattern="*.parquet",
+)
+```
+
+Every shard must share the same
+`(prediction_method_name, predictor_version)` — the core invariant
+applies across shards the same way it applies inside one.
+
+**Overlap resolution** (`on_overlap=`):
+
+- `"raise"` (default) — fail if any `(peptide, allele, peptide_length)`
+  appears in more than one shard. A sample of conflicting keys is
+  included in the error. Use this if shards should be disjoint.
+- `"last"` — later shard in the input list wins. Useful when the
+  sort order represents "newer overwrites older."
+- `"first"` — earlier shard wins.
+- `callable(row_a, row_b) -> row` — custom resolver. Called pairwise
+  per duplicate group. Pattern for "keep stronger binder":
+
+  ```python
+  def keep_lower_affinity(a, b):
+      return a if a["affinity"] <= b["affinity"] else b
+
+  cache = CachedPredictor.concat(shards, on_overlap=keep_lower_affinity)
+  ```
+
+`from_directory` passes `on_overlap` through to `concat`; file order
+is sorted lexicographically, so `shard_a.tsv` is always earlier than
+`shard_b.tsv`.
+
 ### From an in-memory DataFrame
 
 ```python
diff --git a/tests/test_cached_predictor.py b/tests/test_cached_predictor.py
@@ -644,3 +644,130 @@ def test_from_netmhciipan_stdout_unsupported_version_rejects(self, tmp_path):
         path = self._write(tmp_path, "ii.out", _NETMHCPAN_41_STDOUT)
         with pytest.raises(ValueError, match="not supported"):
             CachedPredictor.from_netmhciipan_stdout(path, version="99")
+
+
+# ---------------------------------------------------------------------------
+# Sharding: concat + from_directory
+# ---------------------------------------------------------------------------
+
+
+class TestConcat:
+    def test_concat_merges_disjoint_shards(self):
+        a = CachedPredictor.from_dataframe(_df([_row(peptide="SIINFEKLA")]))
+        b = CachedPredictor.from_dataframe(_df([_row(peptide="GILGFVFTL")]))
+        merged = CachedPredictor.concat([a, b])
+        assert set(merged._df["peptide"]) == {"SIINFEKLA", "GILGFVFTL"}
+        assert merged.prediction_method_name == "random"
+        assert merged.predictor_version == "1.0"
+
+    def test_concat_empty_list_rejects(self):
+        with pytest.raises(ValueError, match="no caches given"):
+            CachedPredictor.concat([])
+
+    def test_concat_mixed_versions_rejects(self):
+        a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", predictor_version="1.0")]),
+        )
+        b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="GILGFVFTL", predictor_version="2.0")]),
+        )
+        with pytest.raises(ValueError, match="multiple"):
+            CachedPredictor.concat([a, b])
+
+    def test_concat_overlap_raises_by_default(self):
+        a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=100.0)]),
+        )
+        b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=999.0)]),
+        )
+        with pytest.raises(ValueError, match="overlapping"):
+            CachedPredictor.concat([a, b])
+
+    def test_concat_overlap_last_wins(self):
+        a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=100.0)]),
+        )
+        b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=999.0)]),
+        )
+        merged = CachedPredictor.concat([a, b], on_overlap="last")
+        assert merged._df.iloc[0]["affinity"] == 999.0
+
+    def test_concat_overlap_first_wins(self):
+        a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=100.0)]),
+        )
+        b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=999.0)]),
+        )
+        merged = CachedPredictor.concat([a, b], on_overlap="first")
+        assert merged._df.iloc[0]["affinity"] == 100.0
+
+    def test_concat_overlap_callable_resolver(self):
+        a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=100.0)]),
+        )
+        b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=999.0)]),
+        )
+        # Keep the lower affinity (stronger binder) — caller's choice.
+        def keep_lower(x, y):
+            return x if x["affinity"] <= y["affinity"] else y
+        merged = CachedPredictor.concat([a, b], on_overlap=keep_lower)
+        assert merged._df.iloc[0]["affinity"] == 100.0
+
+    def test_concat_overlap_invalid_policy_rejects(self):
+        a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=100.0)]),
+        )
+        b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=999.0)]),
+        )
+        with pytest.raises(ValueError, match="on_overlap"):
+            CachedPredictor.concat([a, b], on_overlap="random-policy")
+
+
+class TestFromDirectory:
+    def test_from_directory_loads_all_shards(self, tmp_path):
+        # Two shards with disjoint peptides
+        shard_a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA")]),
+        )
+        shard_b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="GILGFVFTL")]),
+        )
+        shard_a.save(tmp_path / "shard_a.tsv")
+        shard_b.save(tmp_path / "shard_b.tsv")
+
+        merged = CachedPredictor.from_directory(tmp_path, pattern="*.tsv")
+        assert set(merged._df["peptide"]) == {"SIINFEKLA", "GILGFVFTL"}
+
+    def test_from_directory_nonexistent_raises(self, tmp_path):
+        missing = tmp_path / "nope"
+        with pytest.raises(ValueError, match="not a directory"):
+            CachedPredictor.from_directory(missing)
+
+    def test_from_directory_no_matching_files_raises(self, tmp_path):
+        # Empty dir with a pattern that can't match
+        with pytest.raises(ValueError, match="no files matching"):
+            CachedPredictor.from_directory(tmp_path, pattern="*.parquet")
+
+    def test_from_directory_propagates_on_overlap(self, tmp_path):
+        shard_a = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=100.0)]),
+        )
+        shard_b = CachedPredictor.from_dataframe(
+            _df([_row(peptide="SIINFEKLA", affinity=999.0)]),
+        )
+        shard_a.save(tmp_path / "a.tsv")
+        shard_b.save(tmp_path / "b.tsv")
+        # Default policy raises
+        with pytest.raises(ValueError, match="overlapping"):
+            CachedPredictor.from_directory(tmp_path, pattern="*.tsv")
+        # last-wins policy succeeds
+        merged = CachedPredictor.from_directory(
+            tmp_path, pattern="*.tsv", on_overlap="last",
+        )
+        # sorted(["a.tsv", "b.tsv"]) → b wins (affinity=999.0)
+        assert merged._df.iloc[0]["affinity"] == 999.0
diff --git a/topiary/__init__.py b/topiary/__init__.py
@@ -49,7 +49,7 @@
 from .result import TopiaryResult, concat
 from .wide import detect_form, from_wide, to_wide
 
-__version__ = "5.5.0"
+__version__ = "5.6.0"
 
 __all__ = [
     "TopiaryPredictor",
diff --git a/topiary/cached.py b/topiary/cached.py
