ZacharyZhang-NY
29163632af
The per-arm runner recomputed BM25 and the dense matmul for every arm, which put the full ablation at roughly 16 hours and left no way to survive a kill. JointRunner computes each signal once per query, attributes latency to the components each arm uses, and checkpoints every query to runs/baseline/rankings.jsonl. - bench/systems.py: JointRunner + ARM_COMPONENTS signal attribution. - bench/run.py: joint pass, rankings cache with a run fingerprint, per-query resume, latency summary, resident-memory progress lines. - bench/tests/test_joint.py: joint rankings equal the per-arm rankings for all five arms; the sparse corpus is scored once even with three arms that need it; cache fingerprint mismatch and malformed rows fail fast. - bench/reports/EXTERNAL_APPENDIX.md: what a real Mem0/Zep run would cost, and which published numbers are comparable at all. Verification: pytest bench/tests -q -> 57 passed; ruff check bench/ -> all checks passed; mypy (fresh cache) -> no issues in 25 source files.
100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
"""JointRunner must rank exactly like ArmRunner, and the rankings cache must resume.
|
|
|
|
The ablation switched to a single pass that computes each signal once. If that pass
|
|
drifted from the per-arm path, the baseline would silently measure a different system.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from bench import systems
|
|
from bench.run import Query, fingerprint, latency, load_rankings
|
|
from bench.systems import ARMS, ArmRunner, Atom, JointRunner
|
|
from bench.tests.test_systems import ATOMS
|
|
|
|
MODEL_ARMS = ["bm25_only", "graph_only", "dense_only", "no_rerank", "fusion_rerank"]
|
|
|
|
|
|
def _queries() -> list[str]:
|
|
return ["alpha report", "beta invoice", "gamma", ""]
|
|
|
|
|
|
@pytest.mark.parametrize("arm", MODEL_ARMS)
|
|
def test_joint_ranking_matches_the_per_arm_runner(arm: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
|
|
queries = _queries()
|
|
joint = JointRunner([arm], ATOMS, systems.MODELS)
|
|
per_arm = ArmRunner(arm, ATOMS, systems.MODELS)
|
|
for query in queries:
|
|
expected = per_arm.rank(query)
|
|
ranked, seconds = joint.rank_all(query)[arm]
|
|
assert ranked == expected, f"{arm} drifted from the per-arm path on {query!r}"
|
|
assert seconds >= 0.0
|
|
|
|
|
|
def test_joint_runner_computes_the_sparse_signal_once_per_query(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
|
|
joint = JointRunner(["bm25_only", "no_rerank", "fusion_rerank"], ATOMS, systems.MODELS)
|
|
calls = 0
|
|
assert joint._sparse is not None
|
|
original = joint._sparse.scores
|
|
|
|
def counting_scores(query: str):
|
|
nonlocal calls
|
|
calls += 1
|
|
return original(query)
|
|
|
|
monkeypatch.setattr(joint._sparse, "scores", counting_scores)
|
|
joint.rank_all("alpha report")
|
|
assert calls == 1, "three arms that all need BM25 must still score the corpus once"
|
|
|
|
|
|
def test_joint_runner_rejects_unknown_or_empty_arm_lists() -> None:
|
|
with pytest.raises(SystemExit, match="unknown arm"):
|
|
JointRunner(["vibes"], ATOMS, systems.MODELS)
|
|
with pytest.raises(SystemExit, match="at least one arm"):
|
|
JointRunner([], ATOMS, systems.MODELS)
|
|
|
|
|
|
def test_rankings_cache_resumes_and_refuses_a_mismatched_run(tmp_path: Path) -> None:
|
|
atoms = [Atom(atom_id=f"a{index}", source_id=f"s{index}", text=f"text {index}") for index in range(3)]
|
|
queries = [
|
|
Query(
|
|
query_id=f"q{index}",
|
|
text="t",
|
|
modality="pdf",
|
|
gold=frozenset({"s0"}),
|
|
holder="h",
|
|
answerable=True,
|
|
as_of=None,
|
|
)
|
|
for index in range(2)
|
|
]
|
|
mark = fingerprint(["bm25_only"], queries, atoms)
|
|
path = tmp_path / "rankings.jsonl"
|
|
assert load_rankings(path, mark) == {}
|
|
rows = [json.dumps({"fingerprint": mark}), json.dumps({"query_id": "q0", "arms": {"bm25_only": {"ranked": ["s0"]}}})]
|
|
path.write_text("\n".join(rows) + "\n", encoding="utf-8")
|
|
assert load_rankings(path, mark) == {"q0": {"bm25_only": {"ranked": ["s0"]}}}
|
|
with pytest.raises(SystemExit, match="different run"):
|
|
load_rankings(path, "other")
|
|
path.write_text("\n".join([rows[0], "{not json"]) + "\n", encoding="utf-8")
|
|
with pytest.raises(SystemExit, match="not JSON"):
|
|
load_rankings(path, mark)
|
|
|
|
|
|
def test_latency_needs_samples() -> None:
|
|
assert latency([0.5, 1.0])["p50"] == pytest.approx(0.75)
|
|
with pytest.raises(SystemExit, match="no latency samples"):
|
|
latency([])
|
|
|
|
|
|
def test_every_arm_declares_the_signals_it_consumes() -> None:
|
|
assert set(systems.ARM_COMPONENTS) == set(ARMS)
|