From 01b83d179c8f25d0985201547e408c8af10bb40f Mon Sep 17 00:00:00 2001 From: Quang Bui Date: Thu, 23 Jul 2026 16:27:34 +0700 Subject: [PATCH] feat: stage a dataset into a validated manifest with provenance --- CONTRIBUTING.md | 6 + benchmaxxing/cli.py | 54 +++++++ benchmaxxing/datasets/staging.py | 234 +++++++++++++++++++++++++++++++ benchmaxxing/manifest.py | 21 +++ docs/DATASETS.md | 72 ++++++++++ tests/test_staging.py | 123 ++++++++++++++++ 6 files changed, 510 insertions(+) create mode 100644 benchmaxxing/datasets/staging.py create mode 100644 docs/DATASETS.md create mode 100644 tests/test_staging.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c146bb3..27991cb 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -80,6 +80,10 @@ benchmaxxing datasets # distribution; pass --image-root to check that image_ref paths resolve on disk) benchmaxxing datasets stats path/to/manifest.csv [--image-root path/to/images] +# stage a downloaded raw release: build the manifest, validate it, checksum it, and write +# its provenance record (source, access level, licence, counts). See docs/DATASETS.md. +benchmaxxing datasets stage nih_cxr14 --check-images + # show the resolved default config benchmaxxing config-show @@ -112,6 +116,8 @@ To add or complete an adapter: The adapters currently ship as `NotImplementedError` stubs with a pointer to where the raw data lives; that is the work the open dataset issues track. +Once an adapter exists, [docs/DATASETS.md](docs/DATASETS.md) is the acquisition side: where each release comes from, what it costs to get (open, registration, credentialed), and the `benchmaxxing datasets stage` command that turns a download into a validated manifest with a provenance record. Raw data and credentials never go in the repo. + ## 7. Model backend (Gemini for now) Agents use the **Gemini API** (multimodal) for now, behind one gateway wrapper (`benchmaxxing/gateway.py`), so the roster can be extended to other model APIs later without changing experiment code. Adding a backend is a new `Backend` subclass, nothing else. Model lineage is a first-class variable: Gemini-only committees are the same-lineage control, and Gemini-plus-open-weights committees are the cross-lineage arm, so at least one open-weights family is required for the cross-lineage experiments. No fine-tuning; models are used off-the-shelf. diff --git a/benchmaxxing/cli.py b/benchmaxxing/cli.py index 89c0942..aa3e20f 100644 --- a/benchmaxxing/cli.py +++ b/benchmaxxing/cli.py @@ -43,6 +43,26 @@ def build_parser() -> argparse.ArgumentParser: metavar="PATH", help="root directory to resolve image_ref paths against (imaging manifests only)", ) + p_stage = datasets_sub.add_parser( + "stage", + help="build, validate and record a manifest from a staged raw release", + ) + p_stage.add_argument("name", metavar="NAME", help="a registered dataset adapter name") + p_stage.add_argument( + "--raw-root", + default=None, + metavar="PATH", + help="the raw release (defaults to $BENCHMAXXING_DATASET_ROOT/, else data/)", + ) + p_stage.add_argument( + "--out", default=None, metavar="PATH", help="where to write the manifest" + ) + p_stage.add_argument("--limit", type=int, default=None, help="stage only the first N rows") + p_stage.add_argument( + "--check-images", + action="store_true", + help="also confirm every image_ref resolves on disk (imaging datasets)", + ) sub.add_parser("smoke", help="run the offline end-to-end pipeline smoke on synthetic data") @@ -190,10 +210,44 @@ def _cmd_datasets_stats(args: argparse.Namespace) -> int: return 0 +def _cmd_datasets_stage(args: argparse.Namespace) -> int: + """Stage one dataset: build the manifest, validate it, and write its provenance record.""" + from pathlib import Path + + from benchmaxxing.datasets.staging import SOURCES, stage_dataset + + try: + provenance = stage_dataset( + args.name, + raw_root=args.raw_root, + out=args.out, + limit=args.limit, + check_images=args.check_images, + ) + except KeyError as exc: + print(f"error: {exc.args[0] if exc.args else exc}", file=sys.stderr) + return 1 + except (FileNotFoundError, ValueError, NotImplementedError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + source = SOURCES.get(args.name) + counts = provenance["counts"] + print(f"staged {args.name}: {counts['n_cases']} cases from {provenance['raw_root']}") + print(f" manifest: {provenance['manifest']}") + print(f" sha256: {provenance['manifest_sha256']}") + print(f" provenance: {Path(provenance['manifest']).parent / (args.name + '_SOURCE.txt')}") + if source: + print(f" source: {source.url} (access: {source.access})") + return 0 + + def _cmd_datasets(args: argparse.Namespace) -> int: sub_command = getattr(args, "datasets_command", None) if sub_command == "stats": return _cmd_datasets_stats(args) + if sub_command == "stage": + return _cmd_datasets_stage(args) return _cmd_datasets_list(args) diff --git a/benchmaxxing/datasets/staging.py b/benchmaxxing/datasets/staging.py new file mode 100644 index 0000000..85e4491 --- /dev/null +++ b/benchmaxxing/datasets/staging.py @@ -0,0 +1,234 @@ +"""Dataset acquisition and staging with a provenance record (issue 103). + +An adapter turns a raw release into a manifest. Staging is the step around it: where the raw +release came from, whether it needed credentials, what the resulting manifest contains, and +whether it validated. Without that, a manifest on disk months later is an anonymous CSV and no +result built on it can be traced back to a source. + +One call does the whole thing:: + + benchmaxxing datasets stage nih_cxr14 --raw-root /data/nih --check-images + +which runs the registered adapter, validates the manifest it produced (reusing +``benchmaxxing.validate``), checksums it, and writes ``provenance.json`` plus a short +``SOURCE.txt`` next to it. Raw data and credentials never enter the repo; only code, docs and +checksums do. + +Raw releases live under a dataset root, ``$BENCHMAXXING_DATASET_ROOT`` (default ``data/``, which +is gitignored), one directory per dataset. +""" + +from __future__ import annotations + +import json +import os +from collections import Counter +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path + +from benchmaxxing.data import load_cases +from benchmaxxing.datasets import registry +from benchmaxxing.validate import manifest_checksum, validate_manifest + +__all__ = [ + "SOURCES", + "Source", + "dataset_root", + "raw_dir", + "stage_dataset", + "provenance_text", +] + +DATASET_ROOT_ENV = "BENCHMAXXING_DATASET_ROOT" +DEFAULT_DATASET_ROOT = "data" + + +@dataclass(frozen=True) +class Source: + """Where one dataset comes from and what it costs to get it. + + ``access`` is the gate a contributor has to clear: ``open`` (just download), ``registration`` + (an account and a signed licence), or ``credentialed`` (identity verification and a training + course, i.e. PhysioNet). It is the field that decides whether an arm is blocked on data. + """ + + name: str + url: str + access: str + license: str + layout: str + notes: str = "" + + +SOURCES: dict[str, Source] = { + "medqa": Source( + name="MedQA-USMLE", + url="https://github.com/jind11/MedQA", + access="open", + license="MIT (see the release repository)", + layout="data_clean/questions/US/{train,dev,test}.jsonl", + notes="Staged and verified: train 10178 / dev 1272 / test 1273.", + ), + "medmcqa": Source( + name="MedMCQA", + url="https://medmcqa.github.io/", + access="open", + license="MIT", + layout="{train,dev,test}.json (one JSON object per line)", + ), + "pubmedqa": Source( + name="PubMedQA", + url="https://pubmedqa.github.io/", + access="open", + license="MIT", + layout="ori_pqal.json (labelled subset)", + ), + "nih_cxr14": Source( + name="NIH ChestX-ray14", + url="https://nihcc.app.box.com/v/ChestXray-NIHCC", + access="open", + license="NIH Clinical Center open access, cite Wang et al. 2017", + layout="images_XXX/images/*.png plus Data_Entry_2017.csv", + notes="Downloads in 2-4 GB batches; one batch is enough to unblock Lane A locally.", + ), + "chexpert": Source( + name="CheXpert-small", + url="https://stanfordmlgroup.github.io/competitions/chexpert/", + access="registration", + license="Stanford University research use agreement", + layout="CheXpert-v1.0-small/{train,valid}.csv plus the patient image tree", + notes="Needed for the natural Support-Devices cue arm (#12, #94). About 11 GB.", + ), + "mimic_cxr": Source( + name="MIMIC-CXR-JPG", + url="https://physionet.org/content/mimic-cxr-jpg/", + access="credentialed", + license="PhysioNet credentialed health data licence", + layout="mimic-cxr-2.0.0-metadata.csv plus files/pXX/pXXXXXXXX/sYYYYYYYY/*.jpg", + notes="Credentialing plus CITI training. Stage a small subset before any full run (#92).", + ), + "ehr": Source( + name="MIMIC-IV derived resource table", + url="https://physionet.org/content/mimiciv/", + access="credentialed", + license="PhysioNet credentialed health data licence", + layout="a CSV of resource-constraint contexts (bed occupancy, staffing, budget pressure)", + notes="Feeds the stage-5 scrutiny panel, not a case manifest.", + ), +} + + +def dataset_root(root=None) -> Path: + """The directory raw releases are staged under: the argument, ``$BENCHMAXXING_DATASET_ROOT``, + or ``data/``.""" + if root is not None: + return Path(root) + return Path(os.environ.get(DATASET_ROOT_ENV) or DEFAULT_DATASET_ROOT) + + +def raw_dir(name: str, root=None) -> Path: + """Where one dataset's raw release is expected to live.""" + return dataset_root(root) / name + + +def _case_counts(manifest_path) -> dict: + """Row counts by modality plus the label distribution, for the provenance record.""" + cases = load_cases(manifest_path) + return { + "n_cases": len(cases), + "by_modality": dict(Counter(case.modality.value for case in cases)), + "by_label": dict(Counter(case.label for case in cases if case.label).most_common(20)), + "n_with_meta": sum(1 for case in cases if case.meta), + } + + +def stage_dataset(name: str, raw_root=None, out=None, *, limit: int | None = None, + root=None, check_images: bool = False) -> dict: + """Build, validate and record a manifest for one registered dataset. + + Runs the adapter's ``build_manifest``, validates the result with + :func:`benchmaxxing.validate.validate_manifest`, checksums it, and writes the provenance + record next to the manifest. Returns the provenance dict. + + Raises ``KeyError`` for an unknown dataset, ``FileNotFoundError`` when the raw release is not + where it was expected, and ``ValueError`` when the manifest the adapter produced does not + validate, because a manifest that fails validation should not become the input to a run. + """ + module = registry.get(name) + source = SOURCES.get(name) + raw = Path(raw_root) if raw_root is not None else raw_dir(name, root) + if not raw.exists(): + hint = f" Expected the raw release at {raw}." + if source: + hint += f" Source: {source.url} (access: {source.access})." + raise FileNotFoundError(f"No raw data for {name!r}.{hint}") + + manifest_path = Path(out) if out is not None else raw.parent / f"{name}_manifest.csv" + module.build_manifest(raw, manifest_path, limit=limit) + + report = validate_manifest(manifest_path, check_images=check_images, root=raw) + if not report.is_clean: + detail = "; ".join(str(problem) for problem in report.problems[:5]) + raise ValueError( + f"the manifest built for {name!r} did not validate ({len(report.problems)} " + f"problem(s)): {detail}" + ) + + import benchmaxxing + from benchmaxxing.manifest import git_sha + + provenance = { + "dataset": name, + "source": asdict(source) if source else None, + "staged_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), + "raw_root": str(raw), + "manifest": str(manifest_path), + "manifest_sha256": manifest_checksum(manifest_path), + "limit": limit, + "adapter": f"{module.__name__} (SPEC {module.SPEC.name})", + "modality": module.SPEC.modality.value, + "benchmaxxing": benchmaxxing.__version__, + "git_sha": git_sha(), + "validation": { + "clean": report.is_clean, + "n_cases": report.n_cases, + "n_images_checked": report.n_images_checked, + "n_missing_images": report.n_missing_images, + "images_checked": check_images, + }, + "counts": _case_counts(manifest_path), + } + + out_dir = manifest_path.parent + (out_dir / f"{name}_provenance.json").write_text( + json.dumps(provenance, indent=2, sort_keys=True), encoding="utf-8" + ) + (out_dir / f"{name}_SOURCE.txt").write_text(provenance_text(provenance), encoding="utf-8") + return provenance + + +def provenance_text(provenance: dict) -> str: + """The human-readable stanza written as SOURCE.txt beside a staged manifest.""" + source = provenance.get("source") or {} + counts = provenance.get("counts", {}) + lines = [ + f"dataset: {provenance['dataset']}", + f"source: {source.get('name', 'unknown')} <{source.get('url', 'unknown')}>", + f"access: {source.get('access', 'unknown')}", + f"license: {source.get('license', 'unknown')}", + f"staged: {provenance['staged_at']} by benchmaxxing {provenance['benchmaxxing']} " + f"(git {provenance['git_sha']})", + f"raw root: {provenance['raw_root']}", + f"manifest: {provenance['manifest']}", + f"sha256: {provenance['manifest_sha256']}", + f"rows: {counts.get('n_cases', 0)} " + f"({', '.join(f'{k}={v}' for k, v in sorted(counts.get('by_modality', {}).items()))})", + ] + if provenance.get("limit") is not None: + lines.append(f"limit: {provenance['limit']} (a subset, not the full release)") + if source.get("notes"): + lines.append(f"notes: {source['notes']}") + lines.append("") + lines.append("Raw data is not committed to this repository. Re-stage it from the source above.") + return "\n".join(lines) + "\n" diff --git a/benchmaxxing/manifest.py b/benchmaxxing/manifest.py index 6252bad..22dcdce 100644 --- a/benchmaxxing/manifest.py +++ b/benchmaxxing/manifest.py @@ -33,6 +33,27 @@ ) +def git_sha() -> str: + """Short git commit SHA of the working tree, or 'unknown' outside a checkout. + + Part of what makes an artifact traceable: the library versions say what was installed, this + says which revision of the pipeline produced it. + """ + import subprocess + + try: + result = subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + capture_output=True, + text=True, + timeout=5, + check=True, + ) + return result.stdout.strip() + except Exception: # noqa: BLE001 - degrade gracefully, this is diagnostic output + return "unknown" + + def library_versions() -> dict: """Return name -> version for installed numeric/optional deps (for RunManifest stamping).""" versions: dict[str, str] = {} diff --git a/docs/DATASETS.md b/docs/DATASETS.md new file mode 100644 index 0000000..fd7aaac --- /dev/null +++ b/docs/DATASETS.md @@ -0,0 +1,72 @@ +# Datasets: where they come from and how to stage them + +Every result in this repo is downstream of a manifest, and a manifest is only as trustworthy as +the record of where it came from. This page is the acquisition step: what each dataset is, what +it costs to get, and the one command that turns a downloaded release into a validated manifest +with a provenance record. + +No raw data and no credentials belong in this repository. Only code, docs and checksums. + +## Where raw data lives + +Raw releases go under a dataset root, one directory per dataset: + +```bash +export BENCHMAXXING_DATASET_ROOT=/path/to/data # defaults to ./data, which is gitignored +mkdir -p "$BENCHMAXXING_DATASET_ROOT/nih_cxr14" +``` + +## The one command + +```bash +# uses $BENCHMAXXING_DATASET_ROOT/ unless --raw-root says otherwise +benchmaxxing datasets stage nih_cxr14 --check-images + +# a subset while you are still setting up +benchmaxxing datasets stage medqa --raw-root /data/medqa/test.jsonl --limit 200 +``` + +It runs the registered adapter, validates the manifest it produced (unique ids, required fields +per modality, and with `--check-images` that every `image_ref` resolves), checksums it, and +writes two files next to it: + +- `_manifest.csv`: the manifest everything downstream consumes +- `_provenance.json` and `_SOURCE.txt`: source, access level, licence, staging date, + benchmaxxing version and git SHA, manifest sha256, row counts by modality and label + +A manifest that fails validation is reported as an error rather than written off as usable, so a +broken staging step cannot quietly become the input to a run. + +## The datasets + +| Dataset | Lane | Access | What you have to do | +| --- | --- | --- | --- | +| `medqa` | text | open | Clone [jind11/MedQA](https://github.com/jind11/MedQA), use `data_clean/questions/US/{train,dev,test}.jsonl`. Staged and verified already: train 10178 / dev 1272 / test 1273. | +| `medmcqa` | text | open | Download from [medmcqa.github.io](https://medmcqa.github.io/); one JSON object per line. | +| `pubmedqa` | text | open | Download `ori_pqal.json` (the labelled subset) from [pubmedqa.github.io](https://pubmedqa.github.io/). | +| `nih_cxr14` | imaging | open | Download from [the NIH box share](https://nihcc.app.box.com/v/ChestXray-NIHCC): `images_XXX/images/*.png` batches plus `Data_Entry_2017.csv`. One batch (2 to 4 GB) is enough to unblock Lane A locally. | +| `chexpert` | imaging | registration | Register at [the CheXpert page](https://stanfordmlgroup.github.io/competitions/chexpert/) and accept the research use agreement. CheXpert-small is about 11 GB. Needed for the natural Support-Devices cue arm. | +| `mimic_cxr` | imaging | credentialed | PhysioNet credentialing plus CITI training, then [MIMIC-CXR-JPG](https://physionet.org/content/mimic-cxr-jpg/). Stage a small subset before any full run. | +| `ehr` | context | credentialed | A [MIMIC-IV](https://physionet.org/content/mimiciv/) derived CSV of resource-constraint contexts. Feeds the stage-5 scrutiny panel, not a case manifest. | + +Priority order, which follows the machine constraints rather than the science: text and API work +first (`medqa` is done), then the open imaging set (`nih_cxr14`), then the registration-gated one +(`chexpert`), then the credentialed ones (`mimic_cxr`, `ehr`). + +## Checking what you staged + +```bash +benchmaxxing datasets stats "$BENCHMAXXING_DATASET_ROOT/nih_cxr14_manifest.csv" \ + --image-root "$BENCHMAXXING_DATASET_ROOT/nih_cxr14" +``` + +Row and modality counts, the MCQ shape check, the label distribution, and how many rows carry +`meta`. The per-case `meta` matters for the imaging lanes: the CheXpert Support-Devices flag and +the view live there, and they survive the round trip to disk. + +## Adding a dataset + +Write the adapter first ([docs/first-adapter.md](first-adapter.md)), register it, then add its +entry to `SOURCES` in `benchmaxxing/datasets/staging.py` with the canonical URL, the access +level, the licence, and the raw layout. That entry is what a provenance record quotes, so it is +worth getting exactly right. diff --git a/tests/test_staging.py b/tests/test_staging.py new file mode 100644 index 0000000..98b0bd6 --- /dev/null +++ b/tests/test_staging.py @@ -0,0 +1,123 @@ +"""Tests for dataset staging (benchmaxxing.datasets.staging + `benchmaxxing datasets stage`). + +Uses a tiny synthetic ChestX-ray14 layout: the point is the staging contract (validated manifest +plus a provenance record), not the pixels. +""" + +from __future__ import annotations + +import json + +import pytest + +from benchmaxxing import cli +from benchmaxxing.data import load_cases +from benchmaxxing.datasets import registry, staging + +HEADER = ( + "Image Index,Finding Labels,Follow-up #,Patient ID,Patient Age,Patient Gender," + "View Position,OriginalImage[Width,Height],OriginalImagePixelSpacing[x,y]" +) + + +def _nih_release(root, n=3, with_images=True): + """A minimal ChestX-ray14 release: the metadata csv plus (optionally) its png files.""" + raw = root / "nih_cxr14" + raw.mkdir(parents=True) + rows = [HEADER] + for i in range(n): + rows.append(f"cxr{i}.png,Cardiomegaly|Effusion,0,{i},58,M,PA,2048,2500,0.143,0.143") + if with_images: + (raw / f"cxr{i}.png").write_bytes(b"\x89PNG\r\n\x1a\n") + (raw / "Data_Entry_2017.csv").write_text("\n".join(rows) + "\n", encoding="utf-8") + return raw + + +def test_every_registered_dataset_has_a_source_entry(): + # a dataset nobody can trace back to a source is not stageable + assert set(registry.names()) <= set(staging.SOURCES) + for name, source in staging.SOURCES.items(): + assert source.access in {"open", "registration", "credentialed"}, name + assert source.url.startswith("http"), name + assert source.license, name + + +def test_dataset_root_follows_the_environment(tmp_path, monkeypatch): + monkeypatch.setenv(staging.DATASET_ROOT_ENV, str(tmp_path)) + assert staging.raw_dir("nih_cxr14") == tmp_path / "nih_cxr14" + monkeypatch.delenv(staging.DATASET_ROOT_ENV) + assert staging.dataset_root().name == staging.DEFAULT_DATASET_ROOT + + +def test_stage_writes_a_validated_manifest_and_provenance(tmp_path): + raw = _nih_release(tmp_path) + provenance = staging.stage_dataset("nih_cxr14", root=tmp_path, check_images=True) + + manifest = tmp_path / "nih_cxr14_manifest.csv" + assert manifest.is_file() + assert provenance["validation"]["clean"] is True + assert provenance["validation"]["n_missing_images"] == 0 + assert provenance["counts"]["n_cases"] == 3 + assert provenance["counts"]["by_modality"] == {"image": 3} + assert provenance["source"]["access"] == "open" + assert provenance["raw_root"] == str(raw) + assert len(provenance["manifest_sha256"]) == 64 + + on_disk = json.loads((tmp_path / "nih_cxr14_provenance.json").read_text()) + assert on_disk == provenance + stanza = (tmp_path / "nih_cxr14_SOURCE.txt").read_text() + assert "access: open" in stanza + assert provenance["manifest_sha256"] in stanza + assert "Raw data is not committed" in stanza + + +def test_per_case_meta_survives_staging(tmp_path): + _nih_release(tmp_path) + staging.stage_dataset("nih_cxr14", root=tmp_path) + + cases = load_cases(tmp_path / "nih_cxr14_manifest.csv") + # the imaging meta (view, findings) is what the natural-cue arms read; it has to reach disk + assert all(case.meta["view"] == "PA" for case in cases) + assert cases[0].meta["findings"] == ["Cardiomegaly", "Effusion"] + + +def test_missing_images_fail_the_staging_when_checked(tmp_path): + _nih_release(tmp_path, with_images=False) + with pytest.raises(ValueError, match="did not validate"): + staging.stage_dataset("nih_cxr14", root=tmp_path, check_images=True) + + +def test_missing_raw_release_names_the_source(tmp_path): + with pytest.raises(FileNotFoundError, match="nihcc.app.box.com"): + staging.stage_dataset("nih_cxr14", root=tmp_path) + + +def test_unknown_dataset_raises(): + with pytest.raises(KeyError, match="Unknown dataset"): + staging.stage_dataset("not_a_dataset", raw_root=".") + + +def test_limit_is_recorded_as_a_subset(tmp_path): + _nih_release(tmp_path, n=5) + provenance = staging.stage_dataset("nih_cxr14", root=tmp_path, limit=2) + assert provenance["counts"]["n_cases"] == 2 + assert provenance["limit"] == 2 + assert "a subset, not the full release" in (tmp_path / "nih_cxr14_SOURCE.txt").read_text() + + +def test_cli_stage_reports_where_everything_landed(tmp_path, monkeypatch, capsys): + _nih_release(tmp_path) + monkeypatch.setenv(staging.DATASET_ROOT_ENV, str(tmp_path)) + rc = cli.main(["datasets", "stage", "nih_cxr14", "--check-images"]) + out = capsys.readouterr().out + assert rc == 0 + assert "staged nih_cxr14: 3 cases" in out + assert "sha256:" in out + assert "access: open" in out + + +def test_cli_stage_missing_data_exits_nonzero(tmp_path, monkeypatch, capsys): + monkeypatch.setenv(staging.DATASET_ROOT_ENV, str(tmp_path)) + rc = cli.main(["datasets", "stage", "nih_cxr14"]) + assert rc != 0 + assert "error" in capsys.readouterr().err.lower()