diff --git a/README.md b/README.md index fd64403..44fee18 100644 --- a/README.md +++ b/README.md @@ -198,6 +198,25 @@ against the known truth). Export the JSONL with `pixi run frugalmind-export`; th drop-in FrugalMind suite lives in [`integrations/frugalmind/`](integrations/frugalmind). +### Scaling the sweep (Fargate / AWS Batch) + +`codameter-bench` scores a grid of processing configs against every golden case +and records recovery RMS. The work is deterministic and embarrassingly parallel, +so it shards cleanly across an array of tasks: + +```bash +codameter-bench plan --grid multiverse --shard 0/64 # size the job (no compute) +codameter-bench sweep --grid multiverse --shard 0/64 --jobs 4 --out s3://bucket/run/ +codameter-bench aggregate --src s3://bucket/run/ --out s3://bucket/run/agg/ +``` + +Each task writes one `shard--of-.jsonl`, so retries are idempotent and no +coordination is needed. On AWS Batch the shard is read from +`AWS_BATCH_JOB_ARRAY_INDEX` + `CODAMETER_SHARDS`, so one image runs every array +task. Container image and a Batch job definition are in +[`docker/`](docker); `s3://` output needs the `aws` extra +(`pip install "codameter[aws]"`). + --- ## Models, hyperparameters, and physical bounds diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..bac45b7 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,24 @@ +# Container for the codameter config-sweep benchmark (codameter-bench). +# Built to run as one task of an array job (AWS Batch on Fargate, or any array +# runner): each task computes one --shard of the case x config grid and writes a +# shard--of-.jsonl to the --out path (local or s3://). +FROM python:3.11-slim + +# Non-interactive, no .pyc, unbuffered logs (so CloudWatch sees progress live). +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 + +WORKDIR /app + +# Install the package (with the aws extra for s3:// output). Only the bits the +# sweep needs are imported at runtime; wheels cover the scientific stack. +COPY pyproject.toml README.md ./ +COPY src ./src +COPY tests/data/golden/manifest.json ./tests/data/golden/manifest.json +RUN pip install ".[aws]" + +# The array runner sets the shard: pass --shard k/N explicitly, or set +# CODAMETER_SHARDS=N and let the CLI read AWS_BATCH_JOB_ARRAY_INDEX. +ENTRYPOINT ["codameter-bench"] +CMD ["--help"] diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 0000000..9148bc1 --- /dev/null +++ b/docker/README.md @@ -0,0 +1,75 @@ +# Scaling the codameter sweep on Fargate / AWS Batch + +`codameter-bench` scores a grid of processing configs against every golden case +and records recovery RMS. The work is deterministic and embarrassingly parallel, +so it fans out cleanly across an array of Fargate tasks: each task runs one +`--shard k/N` and writes one `shard--of-.jsonl` to S3. + +## Size the job first + +```bash +codameter-bench plan --grid multiverse --shard 0/64 +# grid=multiverse cases=10 configs/case=648 +# total cells=6480 shards=64 cells/shard: min=101 max=102 +``` + +`plan` needs no compute; use it to pick the array size and vCPUs per task. Grids: +`compact` (smoke), `multiverse` (default), `wide` (overnight). + +## Run locally (one machine, all vCPUs) + +```bash +codameter-bench sweep --grid multiverse --out ./out --jobs 8 --shard 0/1 +codameter-bench aggregate --src ./out --out ./agg +``` + +`--jobs` spreads one shard's cells across worker processes. Splitting into shards +*and* using `--jobs` composes: shards across tasks, jobs across a task's vCPUs. + +## Run on AWS Batch (Fargate) + +1. **Build and push** the image: + + ```bash + docker build -f docker/Dockerfile -t codameter-bench . + aws ecr get-login-password --region | docker login --username AWS \ + --password-stdin .dkr.ecr..amazonaws.com + docker tag codameter-bench .dkr.ecr..amazonaws.com/codameter-bench:latest + docker push .dkr.ecr..amazonaws.com/codameter-bench:latest + ``` + +2. **Register** the job definition (edit placeholders + `CODAMETER_SHARDS` first): + + ```bash + aws batch register-job-definition --cli-input-json file://docker/aws-batch-job-def.json + ``` + +3. **Submit** an array of N tasks (N must equal `CODAMETER_SHARDS`): + + ```bash + aws batch submit-job --job-name codameter-sweep --job-queue \ + --job-definition codameter-sweep --array-properties size=64 + ``` + + Batch sets `AWS_BATCH_JOB_ARRAY_INDEX` on each task; the CLI reads it together + with `CODAMETER_SHARDS` to select the shard, so every task runs the same image + and command. The `jobRoleArn` needs `s3:PutObject` on the output prefix. + +4. **Aggregate** when the array finishes (locally or as a final one-off task): + + ```bash + codameter-bench aggregate --src s3:///codameter-sweep/run-001/ \ + --out s3:///codameter-sweep/run-001/agg/ + ``` + +## Notes + +- **Idempotent retries.** Each shard writes exactly its own file, named by shard, + so a Batch retry overwrites only that shard. No coordination, no partial merges. +- **Even load.** Sharding is round-robin (`items[k::N]`), so the slow configs + (moving / inversion references, which rerun the estimator per epoch) spread + evenly instead of piling into one task. +- **s3://** output needs the `aws` extra (`pip install "codameter[aws]"`, already + in the image). Local paths need nothing. +- **Cost knob.** vCPUs per task x array size sets throughput; `plan` tells you the + cell count so you can trade array size against per-task `--jobs`. diff --git a/docker/aws-batch-job-def.json b/docker/aws-batch-job-def.json new file mode 100644 index 0000000..29acd60 --- /dev/null +++ b/docker/aws-batch-job-def.json @@ -0,0 +1,28 @@ +{ + "_comment": "Example AWS Batch job definition for the codameter sweep on Fargate. Replace , , , and the role ARNs. Register with: aws batch register-job-definition --cli-input-json file://aws-batch-job-def.json. Submit an array of N tasks with: aws batch submit-job --job-name codameter-sweep --job-queue --job-definition codameter-sweep --array-properties size=64 . Batch sets AWS_BATCH_JOB_ARRAY_INDEX on each task; CODAMETER_SHARDS below must equal the array size.", + "jobDefinitionName": "codameter-sweep", + "type": "container", + "platformCapabilities": ["FARGATE"], + "containerProperties": { + "image": ".dkr.ecr..amazonaws.com/codameter-bench:latest", + "command": [ + "sweep", + "--grid", "multiverse", + "--out", "s3:///codameter-sweep/run-001/", + "--jobs", "4" + ], + "environment": [ + { "name": "CODAMETER_SHARDS", "value": "64" } + ], + "resourceRequirements": [ + { "type": "VCPU", "value": "4" }, + { "type": "MEMORY", "value": "8192" } + ], + "executionRoleArn": "arn:aws:iam:::role/ecsTaskExecutionRole", + "jobRoleArn": "arn:aws:iam:::role/codameter-sweep-task-role", + "fargatePlatformConfiguration": { "platformVersion": "LATEST" }, + "networkConfiguration": { "assignPublicIp": "ENABLED" } + }, + "retryStrategy": { "attempts": 2 }, + "timeout": { "attemptDurationSeconds": 3600 } +} diff --git a/pixi.toml b/pixi.toml index 4cac691..1aa6a2e 100644 --- a/pixi.toml +++ b/pixi.toml @@ -71,6 +71,9 @@ golden = { cmd = "python -m codameter.golden", description = "Regenerate tests/d # Export the golden dv/v datasets as FrugalMind-compatible JSONL + manifest frugalmind-export = { cmd = "python -m codameter.frugalmind --out datasets", description = "Export dv/v benchmark rows as JSONL (FrugalMind format)" } +# Shardable config-sweep benchmark (plan / sweep / aggregate); scales to Fargate +bench = { cmd = "python -m codameter.bench", description = "Config-sweep benchmark over case x config grids" } + # Run only the Parkfield notebook-friendly example parkfield = { cmd = "python examples/01_parkfield_synthetic.py", description = "Synthetic Parkfield demo" } diff --git a/pyproject.toml b/pyproject.toml index 785a8f5..eb10414 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,10 +62,12 @@ dev = [ "mypy>=1.6", "pre-commit>=3.5", ] +aws = ["boto3>=1.28"] # s3:// output for codameter-bench (Fargate/Batch) all = ["codameter[kernels,mcmc]"] [project.scripts] codameter = "codameter.cli:main" +codameter-bench = "codameter.bench:main" [project.urls] Homepage = "https://github.com/Denolle-Lab/codameter" diff --git a/src/codameter/bench.py b/src/codameter/bench.py new file mode 100644 index 0000000..db463db --- /dev/null +++ b/src/codameter/bench.py @@ -0,0 +1,347 @@ +r"""Shardable config-sweep benchmark: recovery RMS over case x config grids. + +The scientific payload is embarrassingly parallel and deterministic: for every +golden case (:mod:`codameter.golden`) and every processing config in a grid +built around that case's recommended choice (:mod:`codameter.use_cases`), run the +pipeline (:func:`codameter.deviations.run_pipeline`) on the known synthetic and +record how well the recovered dv/v tracks the truth. The result is one row per +``(case, config)`` cell: a sensitivity map of which deviations cost how much. + +This module is built to fan out across a cluster (AWS Batch array on Fargate, or +any array runner): + +- ``--shard k/N`` deterministically partitions the ordered work list, so array + task ``k`` of ``N`` computes its slice and nothing else. Round-robin, so load + is even even when some configs (moving/inversion references) are far slower. +- ``--jobs`` runs the shard's cells across local vCPUs. +- ``--out`` takes a local dir or an ``s3://`` prefix; each shard writes one + ``shard--of-.jsonl`` so shards never collide and Batch retries are + idempotent. ``aggregate`` merges the shard files into one table. +- ``plan`` prints the work-item and per-shard counts so you can size the array. + +Run ``codameter-bench --help``. +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +from functools import lru_cache +from itertools import product +from pathlib import Path +from typing import Any, Iterator + +import numpy as np + +from . import golden +from . import use_cases as uc +from .deviations import metrics, run_pipeline + +# --------------------------------------------------------------------------- +# Config grids, built relative to each case's recommended config. +# Each grid names, per axis, how to vary it: "rec" = keep the recommended value, +# "variants" = the generator below, or an explicit list of values. band/window +# variants scale around the recommended value so the grid stays physical for the +# use case (a landslide keeps a high-frequency band; a volcano a low one). +# --------------------------------------------------------------------------- +ESTIMATORS_ALL = ["stretching (TS)", "MWCS", "WCS", "DTW"] + +GRIDS: dict[str, dict[str, Any]] = { + # Small, fast: for `plan`, tests, and smoke runs. + "compact": { + "estimator": ["stretching (TS)", "MWCS"], + "band": "rec", "window": "rec", "stack": "rec", + "reference": ["fixed", "moving"], "gate": [True], + }, + # The default massive sweep: hundreds of cells per case. + "multiverse": { + "estimator": ESTIMATORS_ALL, + "band": "variants", "window": "variants", "stack": [None, 1, 45], + "reference": ["fixed", "moving", "inversion"], "gate": [True, False], + }, + # Everything, for an overnight run. + "wide": { + "estimator": ESTIMATORS_ALL, + "band": "variants", "window": "variants", "stack": [None, 1, 5, 20, 60], + "reference": ["fixed", "moving", "inversion"], "gate": [True, False], + }, +} + + +def _band_variants(band) -> list[tuple[float, float]]: + lo, hi = float(band[0]), float(band[1]) + out = [(lo, hi), (lo * 0.5, hi * 0.5), (lo * 1.5, hi * 1.5)] + # Keep physical: strictly positive and ordered. + return [(max(a, 1e-3), b) for a, b in out if b > max(a, 1e-3)] + + +def _window_variants(window) -> list[tuple[float, float]]: + a, b = float(window[0]), float(window[1]) + span = b - a + return [(a, b), (a, a + 0.5 * span), (a, a + 1.8 * span)] + + +def _axis_values(rec_value, mode, generator) -> list: + if mode == "rec": + return [rec_value] + if mode == "variants": + return generator(rec_value) + return list(mode) # explicit list + + +def build_grid(use_case: str, grid: str = "multiverse") -> list[dict]: + """Return the list of configs to score for one use case under ``grid``.""" + if grid not in GRIDS: + raise ValueError(f"unknown grid {grid!r}; known: {sorted(GRIDS)}") + spec = GRIDS[grid] + rec = uc.recommend(use_case) + bands = _axis_values(rec["band"], spec["band"], _band_variants) + windows = _axis_values(rec["window"], spec["window"], _window_variants) + stacks = [rec["stack"] if s is None else s for s in + _axis_values(rec["stack"], spec["stack"], None)] + configs = [] + for est, band, window, stack, ref, gate in product( + spec["estimator"], bands, windows, stacks, spec["reference"], spec["gate"] + ): + configs.append({"estimator": est, "band": tuple(band), + "window": tuple(window), "stack": int(stack), + "reference": ref, "gate": bool(gate)}) + return configs + + +# --------------------------------------------------------------------------- +# Work list + sharding +# --------------------------------------------------------------------------- +def work_items(case_ids: list[str], grid: str) -> list[tuple[str, int, dict]]: + """Every ``(case_id, config_index, config)`` cell, in a stable order.""" + items = [] + for cid in case_ids: + use_case = golden.CASES_BY_ID[cid]["use_case"] + for i, cfg in enumerate(build_grid(use_case, grid)): + items.append((cid, i, cfg)) + return items + + +def shard(items: list, k: int, n: int) -> list: + """Round-robin slice ``k`` of ``n`` (``0 <= k < n``).""" + if not (0 <= k < n): + raise ValueError(f"shard index {k} out of range for {n} shards") + return items[k::n] + + +def parse_shard(spec: str | None) -> tuple[int, int]: + """Parse ``--shard k/N``; fall back to AWS Batch array env; default 0/1. + + AWS Batch array jobs set ``AWS_BATCH_JOB_ARRAY_INDEX``; pass ``--shards N`` + (or ``CODAMETER_SHARDS``) as the array size and the index is read from the + environment, so the same image runs every array task. + """ + if spec: + k, n = spec.split("/") + return int(k), int(n) + idx = os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX") or os.environ.get("CODAMETER_SHARD_INDEX") + cnt = os.environ.get("CODAMETER_SHARDS") + if idx is not None and cnt is not None: + return int(idx), int(cnt) + return 0, 1 + + +# --------------------------------------------------------------------------- +# Scoring one cell +# --------------------------------------------------------------------------- +@lru_cache(maxsize=None) +def _case(case_id: str) -> dict: + """Load a case's arrays once per process (deterministic, disk-cached).""" + return golden.generate(case_id) + + +def score_cell(case_id: str, config_index: int, cfg: dict) -> dict: + """Score one ``(case, config)`` cell into a JSON-serializable row.""" + use_case = golden.CASES_BY_ID[case_id]["use_case"] + row = { + "case_id": case_id, "use_case": use_case, "config_index": config_index, + "estimator": cfg["estimator"], "band": list(cfg["band"]), + "window": list(cfg["window"]), "stack": cfg["stack"], + "reference": cfg["reference"], "gate": cfg["gate"], + "eps_max": uc.eps_max(use_case), + } + try: + d = _case(case_id) + dvv, valid = run_pipeline(d["ccfs"], d["t"], d["fs"], cfg, + eps_max=row["eps_max"]) + m = metrics(dvv, d["truth"], d["days"], valid) + row.update(rms=golden._rms(dvv, d["truth"], d["days"], valid), + drop_err=m["drop_err"], n_valid=int(np.sum(valid)), ok=True, + error=None) + except Exception as exc: # a bad cell must not kill the shard + row.update(rms=None, drop_err=None, n_valid=0, ok=False, error=str(exc)) + return row + + +def _score_star(args): + return score_cell(*args) + + +# --------------------------------------------------------------------------- +# Sweep driver +# --------------------------------------------------------------------------- +def run_sweep(*, case_ids: list[str], grid: str, k: int, n: int, + jobs: int = 1, progress_every: int = 200) -> list[dict]: + """Score this shard's cells and return the rows (unwritten).""" + items = shard(work_items(case_ids, grid), k, n) + # Pre-generate the shard's unique cases to disk once, so parallel workers + # read the cache rather than racing to write it. + for cid in dict.fromkeys(c for c, _, _ in items): + golden.generate(cid) + + rows: list[dict] = [] + if jobs and jobs > 1: + from concurrent.futures import ProcessPoolExecutor + with ProcessPoolExecutor(max_workers=jobs) as ex: + for i, row in enumerate(ex.map(_score_star, items, chunksize=4), 1): + rows.append(row) + if i % progress_every == 0: + print(f" {i}/{len(items)} cells", file=sys.stderr, flush=True) + else: + for i, (cid, ci, cfg) in enumerate(items, 1): + rows.append(score_cell(cid, ci, cfg)) + if i % progress_every == 0: + print(f" {i}/{len(items)} cells", file=sys.stderr, flush=True) + return rows + + +# --------------------------------------------------------------------------- +# Output (local or s3://) +# --------------------------------------------------------------------------- +def _write_jsonl(rows: list[dict], out: str, name: str) -> str: + """Write ``rows`` to ``/``; ``out`` may be a local dir or s3://.""" + body = "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows) + if out.startswith("s3://"): + import tempfile + + import boto3 # optional; only needed for s3 output + + bucket, _, prefix = out[len("s3://"):].partition("/") + key = f"{prefix.rstrip('/')}/{name}" if prefix else name + with tempfile.NamedTemporaryFile("w", suffix=".jsonl", delete=False) as fh: + fh.write(body) + tmp = fh.name + boto3.client("s3").upload_file(tmp, bucket, key) + os.unlink(tmp) + return f"s3://{bucket}/{key}" + dest = Path(out) + dest.mkdir(parents=True, exist_ok=True) + path = dest / name + path.write_text(body) + return str(path) + + +def _read_jsonl_dir(src: str) -> Iterator[dict]: + if src.startswith("s3://"): + import boto3 + + bucket, _, prefix = src[len("s3://"):].partition("/") + s3 = boto3.client("s3") + for obj in s3.get_paginator("list_objects_v2").paginate( + Bucket=bucket, Prefix=prefix): + for it in obj.get("Contents", []): + if it["Key"].endswith(".jsonl"): + body = s3.get_object(Bucket=bucket, Key=it["Key"])["Body"].read() + for line in body.decode().splitlines(): + if line.strip(): + yield json.loads(line) + else: + for path in sorted(Path(src).glob("shard-*.jsonl")): + for line in path.read_text().splitlines(): + if line.strip(): + yield json.loads(line) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +def _all_case_ids(selector: str | None) -> list[str]: + if not selector or selector == "all": + return [c["id"] for c in golden.CASES] + return [s.strip() for s in selector.split(",") if s.strip()] + + +def _cmd_plan(args) -> int: + case_ids = _all_case_ids(args.cases) + items = work_items(case_ids, args.grid) + k, n = parse_shard(args.shard) + per = [len(shard(items, i, n)) for i in range(n)] + print(f"grid={args.grid} cases={len(case_ids)} configs/case=" + f"{len(build_grid(golden.CASES_BY_ID[case_ids[0]]['use_case'], args.grid))}") + print(f"total cells={len(items)} shards={n} " + f"cells/shard: min={min(per)} max={max(per)}") + return 0 + + +def _cmd_sweep(args) -> int: + case_ids = _all_case_ids(args.cases) + k, n = parse_shard(args.shard) + print(f"sweep grid={args.grid} shard={k}/{n} cases={len(case_ids)} " + f"jobs={args.jobs}", file=sys.stderr) + rows = run_sweep(case_ids=case_ids, grid=args.grid, k=k, n=n, jobs=args.jobs) + name = f"shard-{k:05d}-of-{n:05d}.jsonl" + where = _write_jsonl(rows, args.out, name) + ok = sum(r["ok"] for r in rows) + print(f"wrote {len(rows)} rows ({ok} ok) -> {where}", file=sys.stderr) + return 0 + + +def _cmd_aggregate(args) -> int: + rows = list(_read_jsonl_dir(args.src)) + if not rows: + print("no shard rows found", file=sys.stderr) + return 1 + where = _write_jsonl(rows, args.out, "sweep.jsonl") + # Compact per-case summary: best config by RMS. + best: dict[str, dict] = {} + for r in rows: + if not r.get("ok") or r.get("rms") is None: + continue + cur = best.get(r["case_id"]) + if cur is None or r["rms"] < cur["rms"]: + best[r["case_id"]] = r + print(f"aggregated {len(rows)} rows -> {where}") + for cid in sorted(best): + b = best[cid] + print(f" {cid:<24} best rms={b['rms']*100:.4f}% " + f"{b['estimator']} band={b['band']} ref={b['reference']}") + return 0 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(prog="codameter-bench", description=__doc__.split("\n")[0]) + sub = p.add_subparsers(dest="command", required=True) + + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--grid", default="multiverse", choices=sorted(GRIDS)) + common.add_argument("--cases", default="all", + help="'all' or a comma-separated list of case ids") + common.add_argument("--shard", default=None, + help="k/N; else read AWS_BATCH_JOB_ARRAY_INDEX + CODAMETER_SHARDS") + + sp = sub.add_parser("plan", parents=[common], help="count work items / shards") + sp.set_defaults(func=_cmd_plan) + + ss = sub.add_parser("sweep", parents=[common], help="score this shard's cells") + ss.add_argument("--out", required=True, help="local dir or s3:// prefix") + ss.add_argument("--jobs", type=int, default=int(os.environ.get("CODAMETER_JOBS", "1")), + help="parallel worker processes for this shard") + ss.set_defaults(func=_cmd_sweep) + + sa = sub.add_parser("aggregate", help="merge shard-*.jsonl into one table") + sa.add_argument("--src", required=True, help="dir or s3:// prefix of shard files") + sa.add_argument("--out", required=True, help="local dir or s3:// prefix") + sa.set_defaults(func=_cmd_aggregate) + + args = p.parse_args(argv) + return args.func(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_bench.py b/tests/test_bench.py new file mode 100644 index 0000000..99e97a2 --- /dev/null +++ b/tests/test_bench.py @@ -0,0 +1,83 @@ +"""Tests for the shardable config-sweep benchmark (codameter.bench).""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +from codameter import bench +from codameter import golden + + +def test_build_grid_shapes_and_counts(): + cfgs = bench.build_grid("volcano", "compact") + assert len(cfgs) == 4 # 2 estimators x 2 references, rec band/window/stack, gate=[True] + for c in cfgs: + assert set(c) == {"estimator", "band", "window", "stack", "reference", "gate"} + assert isinstance(c["band"], tuple) and isinstance(c["window"], tuple) + # The default massive grid is hundreds of cells per case. + assert len(bench.build_grid("volcano", "multiverse")) == 648 + + +def test_band_variants_stay_physical(): + for v in bench._band_variants((0.4, 1.0)): + assert 0 < v[0] < v[1] + + +def test_shard_partition_is_disjoint_and_covers_all(): + items = bench.work_items([c["id"] for c in golden.CASES], "compact") + n = 7 + seen = [] + for k in range(n): + seen.extend(bench.shard(items, k, n)) + # Every item appears exactly once across the shards. + assert len(seen) == len(items) + assert {id(x) for x in seen} == {id(x) for x in items} + + +def test_parse_shard_explicit_env_and_default(monkeypatch): + assert bench.parse_shard("3/8") == (3, 8) + monkeypatch.delenv("AWS_BATCH_JOB_ARRAY_INDEX", raising=False) + monkeypatch.delenv("CODAMETER_SHARDS", raising=False) + monkeypatch.delenv("CODAMETER_SHARD_INDEX", raising=False) + assert bench.parse_shard(None) == (0, 1) + monkeypatch.setenv("AWS_BATCH_JOB_ARRAY_INDEX", "5") + monkeypatch.setenv("CODAMETER_SHARDS", "16") + assert bench.parse_shard(None) == (5, 16) + + +def test_score_cell_recovers_and_handles_errors(): + from codameter import use_cases as uc + good = uc.recommend("volcano") + row = bench.score_cell("volcano_mainstream", 0, good) + assert row["ok"] and row["rms"] is not None and np.isfinite(row["rms"]) + assert row["rms"] < 1e-3 + + bad = dict(good, estimator="NOT_AN_ESTIMATOR") + brow = bench.score_cell("volcano_mainstream", 1, bad) + assert brow["ok"] is False and brow["rms"] is None and brow["error"] + + +def test_run_sweep_and_roundtrip(tmp_path): + rows = bench.run_sweep(case_ids=["volcano_mainstream"], grid="compact", + k=0, n=1, jobs=1) + assert len(rows) == 4 + assert all(r["case_id"] == "volcano_mainstream" for r in rows) + # Write + read back as JSONL (the shard artifact). + dest = bench._write_jsonl(rows, str(tmp_path), "shard-00000-of-00001.jsonl") + assert dest.endswith(".jsonl") + back = list(bench._read_jsonl_dir(str(tmp_path))) + assert len(back) == 4 + json.dumps(back) # serializable + + +def test_best_config_is_the_recommended_one(): + # Across the compact grid, the recommended (stretching, fixed) should win. + rows = bench.run_sweep(case_ids=["volcano_mainstream"], grid="compact", + k=0, n=1, jobs=1) + ok = [r for r in rows if r["ok"]] + best = min(ok, key=lambda r: r["rms"]) + assert best["estimator"] == "stretching (TS)" + assert best["reference"] == "fixed"