3131from mcp_server .core import ablation as _ablation # noqa: E402
3232from mcp_server .core .ablation import Mechanism # noqa: E402
3333from benchmarks .lib .db_snapshot import ( # noqa: E402
34- fingerprint , restore_snapshot ,
34+ fingerprint ,
35+ restore_snapshot ,
3536)
3637from benchmarks .lib import db_setup # noqa: E402
3738
@@ -73,16 +74,18 @@ def _patches_for(mech: Mechanism) -> list[PatchSpec]:
7374 no_decay : PatchSpec = (
7475 "mcp_server.core.thermodynamics" ,
7576 "compute_decay" ,
76- lambda : ( lambda current_heat , * a , ** k : current_heat ) ,
77+ lambda : lambda current_heat , * a , ** k : current_heat ,
7778 )
7879 table : dict [Mechanism , list [PatchSpec ]] = {
7980 Mechanism .OSCILLATORY_CLOCK : [no_decay ],
8081 Mechanism .ADAPTIVE_DECAY : [no_decay ],
81- Mechanism .HOMEOSTATIC_PLASTICITY : [(
82- "mcp_server.core.ablation" ,
83- "neutral_scaling_factor" ,
84- lambda : _ablation .neutral_scaling_factor ,
85- )],
82+ Mechanism .HOMEOSTATIC_PLASTICITY : [
83+ (
84+ "mcp_server.core.ablation" ,
85+ "neutral_scaling_factor" ,
86+ lambda : _ablation .neutral_scaling_factor ,
87+ )
88+ ],
8689 }
8790 return table .get (mech , [])
8891
@@ -126,11 +129,11 @@ def _parse_metrics(stdout: str) -> BenchMetrics:
126129 """
127130 r10 = mrr = 0.0
128131 n = 0
129- if ( m := _RE_R10 .search (stdout ) ):
132+ if m := _RE_R10 .search (stdout ):
130133 r10 = float (m .group (1 )) / 100.0
131- if ( m := _RE_MRR_LINE .search (stdout ) ):
134+ if m := _RE_MRR_LINE .search (stdout ):
132135 mrr = float (m .group (1 ))
133- if ( m := _RE_OVERALL .search (stdout ) ):
136+ if m := _RE_OVERALL .search (stdout ):
134137 mrr = float (m .group (1 ))
135138 r10 = float (m .group (3 )) / 100.0
136139 n = int (m .group (4 ))
@@ -146,16 +149,19 @@ def _resolve_benchmark(bench_id: str, quick: bool) -> tuple[Callable[[], None],
146149 """Return (callable that runs the benchmark, declared n_queries-or-0)."""
147150 if bench_id == "longmemeval-s" :
148151 from benchmarks .longmemeval .run_benchmark import run_benchmark as r
152+
149153 data_path = _ROOT / "benchmarks" / "longmemeval" / "longmemeval_s.json"
150154 limit = QUICK_LIMIT if quick else 0
151155 return (lambda : r (str (data_path ), limit = limit , verbose = False ), limit )
152156 if bench_id == "locomo" :
153157 from benchmarks .locomo .run_benchmark import run_benchmark as r
158+
154159 data_path = _ROOT / "benchmarks" / "locomo" / "locomo10.json"
155160 limit = max (1 , QUICK_LIMIT // 10 ) if quick else None
156161 return (lambda : r (str (data_path ), limit = limit , verbose = False ), limit or 0 )
157162 if bench_id == "beam-100K" :
158163 from benchmarks .beam .run_benchmark import run_benchmark as r
164+
159165 limit = max (1 , QUICK_LIMIT // 10 ) if quick else None
160166 return (lambda : r (split = "100K" , limit = limit , verbose = False ), limit or 0 )
161167 raise ValueError (f"unknown benchmark: { bench_id } " )
@@ -310,19 +316,23 @@ def _select_mechanisms(args: argparse.Namespace) -> list[Mechanism]:
310316
311317
312318def _maybe_restore (
313- snapshot_path : Path | None , target_db_url : str | None ,
314- * , run_id : str = "ablation" ,
319+ snapshot_path : Path | None ,
320+ target_db_url : str | None ,
321+ * ,
322+ run_id : str = "ablation" ,
315323) -> dict | str :
316324 """Restore snapshot (if any), apply deterministic DB+session config,
317325 ANALYZE, return db_seed payload (playbook §4.3, §4.10, §4.15)."""
318326 if snapshot_path is None or target_db_url is None :
319327 return "not-snapshotted"
320328 sha = fingerprint (snapshot_path )
321- print (f" [snapshot] restoring { snapshot_path .name } (sha8={ sha [:8 ]} ) -> "
322- f"{ target_db_url } ..." )
329+ print (
330+ f" [snapshot] restoring { snapshot_path .name } (sha8={ sha [:8 ]} ) -> "
331+ f"{ target_db_url } ..."
332+ )
323333 report = restore_snapshot (target_db_url , snapshot_path )
324334 if not report .success :
325- msg = '; ' .join (report .mismatch + report .version_drift )
335+ msg = "; " .join (report .mismatch + report .version_drift )
326336 raise RuntimeError (f"snapshot restore failed: { msg } " )
327337 os .environ ["DATABASE_URL" ] = target_db_url
328338 # Post-restore deterministic setup (playbook §8 SRP split).
@@ -348,8 +358,11 @@ def _maybe_restore(
348358
349359
350360def _run_baseline_if_needed (
351- benchmark : str , quick : bool , force : bool ,
352- snapshot : Path | None , target_db : str | None ,
361+ benchmark : str ,
362+ quick : bool ,
363+ force : bool ,
364+ snapshot : Path | None ,
365+ target_db : str | None ,
353366) -> BenchMetrics :
354367 cached = None if force else _load_baseline (benchmark )
355368 if cached is not None :
@@ -358,8 +371,9 @@ def _run_baseline_if_needed(
358371 print (f" [baseline] running ({ benchmark } , quick={ quick } ) ..." )
359372 run_id = f"{ benchmark } _BASELINE"
360373 db_seed = _maybe_restore (snapshot , target_db , run_id = run_id )
361- metrics , wall , rss , _ = run_one (benchmark , None , quick = quick ,
362- run_id = run_id if snapshot else None )
374+ metrics , wall , rss , _ = run_one (
375+ benchmark , None , quick = quick , run_id = run_id if snapshot else None
376+ )
363377 _save (benchmark , None , metrics , None , wall , rss , db_seed = db_seed )
364378 print (
365379 f" [baseline] r@10={ metrics .r_at_10 :.3f} mrr={ metrics .mrr :.3f} "
@@ -376,26 +390,43 @@ def main() -> int:
376390 p .add_argument ("--baseline-only" , action = "store_true" )
377391 p .add_argument ("--quick" , action = "store_true" )
378392 p .add_argument ("--force" , action = "store_true" )
379- p .add_argument ("--from-snapshot" , dest = "from_snapshot" , default = None ,
380- help = "path to a custom-format pg_dump; restored before each trial" )
381- p .add_argument ("--target-db" , dest = "target_db" , default = None ,
382- help = "benchmark-only DB URL/name to restore the snapshot into" )
393+ p .add_argument (
394+ "--from-snapshot" ,
395+ dest = "from_snapshot" ,
396+ default = None ,
397+ help = "path to a custom-format pg_dump; restored before each trial" ,
398+ )
399+ p .add_argument (
400+ "--target-db" ,
401+ dest = "target_db" ,
402+ default = None ,
403+ help = "benchmark-only DB URL/name to restore the snapshot into" ,
404+ )
383405 args = p .parse_args ()
384406
385407 snapshot = Path (args .from_snapshot ) if args .from_snapshot else None
386408 target_db = None
387409 if snapshot is not None :
388410 if not args .target_db :
389411 raise SystemExit ("--from-snapshot requires --target-db" )
390- target_db = (args .target_db if "://" in args .target_db
391- else f"postgresql://localhost:5432/{ args .target_db } " )
412+ target_db = (
413+ args .target_db
414+ if "://" in args .target_db
415+ else f"postgresql://localhost:5432/{ args .target_db } "
416+ )
392417 else :
393- print ("WARNING: --from-snapshot not set; results lack DB-state determinism. "
394- "HNSW non-determinism + ingest order will leak into the deltas." )
418+ print (
419+ "WARNING: --from-snapshot not set; results lack DB-state determinism. "
420+ "HNSW non-determinism + ingest order will leak into the deltas."
421+ )
395422
396423 print (f"== ablation runner :: benchmark={ args .benchmark } quick={ args .quick } ==" )
397424 baseline = _run_baseline_if_needed (
398- args .benchmark , args .quick , args .force , snapshot , target_db ,
425+ args .benchmark ,
426+ args .quick ,
427+ args .force ,
428+ snapshot ,
429+ target_db ,
399430 )
400431 if args .baseline_only :
401432 return 0
@@ -413,9 +444,12 @@ def main() -> int:
413444 try :
414445 run_id = f"{ args .benchmark } _{ mech .name } "
415446 db_seed = _maybe_restore (snapshot , target_db , run_id = run_id )
416- metrics , wall , rss , _ = run_one (args .benchmark , mech ,
417- quick = args .quick ,
418- run_id = run_id if snapshot else None )
447+ metrics , wall , rss , _ = run_one (
448+ args .benchmark ,
449+ mech ,
450+ quick = args .quick ,
451+ run_id = run_id if snapshot else None ,
452+ )
419453 except Exception as exc : # source: ablation must be fail-soft per mech
420454 print (f" [{ mech .name } ] FAILED: { exc !r} " )
421455 continue
0 commit comments