Skip to main content

MCP server for auditing ML experiments (W&B): confounded ablations, training-curve pathologies, and misleading sweep conclusions.

Project description

experiment-audit-mcp

Your agent can catch confounded ablations, training pathologies, and misleading sweep conclusions across dozens of runs that you'd otherwise have to notice by eye — this is leverage on researcher attention, not another dashboard.

experiment-audit-mcp is an MCP server that sits on top of your W&B project and gives an agent (or you, directly) eight tools split cleanly into two kinds: cheap deterministic retrieval, and heuristic judgment that always shows its work. It does not visualize anything and does not replace your dashboard — it exists for the one thing dashboards are bad at and LLMs are worse at: reliably comparing dozens of floats across dozens of runs without hand-waving.

Status: v1.0.0, W&B backend only. All ten roadmap milestones are complete and reviewed. Two verification steps remain genuinely blocked by this environment's sandbox constraints (no live W&B credentials, no live Anthropic API access) rather than skipped — see Known Gaps below before you rely on this in production.


Contents

Why this exists

Three specific, recurring failure modes motivated this project:

  1. Confounded ablations. You change use_memory: false to test an ablation, but batch_size also silently changed between runs. The metric delta you're about to write up isn't measuring what you think it's measuring.
  2. Training pathologies that are easy to miss by eye across dozens of runs and metrics — a NaN mid-curve that got silently dropped by a plotting library, a plateau that looks like convergence, a jagged oscillation.
  3. Misleading sweep conclusions — a "most important hyperparameter" claim from a 3-run sweep, or two hyperparameters that move together so an importance ranking attributes one's effect to the other.

experiment-audit-mcp gives an agent tools that refuse to be wrong quietly about any of these. See the design spec's non-negotiable design principles (docs/design-spec-v1.md §1) for the four rules every tool follows.

Install

Requires Python 3.11+.

Not yet published to PyPI (see Known Gaps — the package name is reserved but no release has been pushed there yet), so pip install experiment-audit-mcp will currently fail with "No matching distribution found." Until it's published, install from source instead:

git clone https://github.com/SreeDharshan-GJ/experiment-audit-mcp.git
cd experiment-audit-mcp
pip install -e .

Once a release is published to PyPI, pip install experiment-audit-mcp will also work.

Set your credentials (a read-only W&B API key is recommended — this server never writes to your project):

export WANDB_API_KEY="your-read-only-key"
# Optional: only needed if your key's default entity isn't the one you
# want to query against.
export WANDB_ENTITY="your-team-or-username"

Add it to your MCP client config. For Claude Desktop (claude_desktop_config.json):

{
  "mcpServers": {
    "experiment-audit": {
      "command": "experiment-audit-mcp",
      "env": {
        "WANDB_API_KEY": "your-read-only-key"
      }
    }
  }
}

For Claude Code:

claude mcp add -e WANDB_API_KEY=your-read-only-key experiment-audit -- experiment-audit-mcp

(Options like -e must come before the server name, not after — putting -e after the name has been a source of "Invalid environment variable format" errors in some Claude Code versions.)

Or run it directly (useful for testing with the MCP Inspector):

npx @modelcontextprotocol/inspector experiment-audit-mcp

Quick start

Once connected, ask your agent something like:

"Did I mess up my memory-ablation run? Compare mamfac-baseline and mamfac-no-memory in the mamfac project and check whether the only real difference is use_memory."

The agent will call audit_ablation, which returns a verdict (clean / confounded / uncertain), a confidence level, and the full config diff it based that verdict on — not just an assertion.

"Why did the reward on run xj29fk1a crash around step 40,000?"

The agent will call audit_training_curve, which fetches the metric's full history and returns scored signals (null values, sudden jumps, flat plateaus, oscillation) with the exact step range and evidence for each — never a bare label.

"Which hyperparameter actually mattered in my lr-sweep sweep?"

The agent will call audit_sweep, which refuses to rank anything below 10 usable runs, and flags any hyperparameter pairs that moved together so you don't mistake one's effect for the other's.

The eight tools

Tools are named so the trust level is visible in the name itself, not buried in the docs — a get_*/list_* result is a deterministic fact; a compare_* result is deterministic computation over multiple runs; an audit_* result is a heuristic judgment that always carries method, confidence, and evidence fields, enforced at the schema level (see docs/design-spec-v1.md §4.1).

Tool Kind What it does
test_connection retrieval Validates W&B credentials against the real API; call it explicitly (e.g. as your first tool call in a session) — it is not invoked automatically on server start. Only presence of WANDB_API_KEY is checked automatically at startup (fail-fast if missing), not that the key actually works.
list_runs retrieval Cheap, paginated run listing (id/name/tags/status only — no config or metrics).
get_run_summary retrieval Full config + summary metrics + data_completeness for one run.
get_metric_history retrieval Full point-by-point history for one metric on one run.
compare_runs diffing Config + metric diff across N runs (not just pairwise). No verdict — just the facts.
audit_training_curve judgment Scored signals over a metric history: null_values, sudden_jump, low_variance_plateau, high_frequency_oscillation.
audit_ablation judgment Verdict (clean/confounded/uncertain) on whether an ablation pair actually isolates claimed_variable.
audit_sweep judgment Hyperparameter importance ranking with a hard 10-run floor and co-varying-parameter warnings.

Full methodology — exact thresholds, formulas, and the reasoning behind each detector — lives in docs/audit-methods.md. Tool descriptions in the MCP schema itself stay short and point here, so they don't cost context budget on every conversational turn.

Architecture

experiment_audit/
├── models.py            # RunRef, Run, MetricPoint, MetricHistory, Sweep, Page[T]
├── errors.py             # ToolError + the frozen error_type taxonomy
├── auth.py               # env-var credential handling, fail-fast
├── server.py             # FastMCP entrypoint; registers all 8 tools
├── backends/
│   ├── base.py           # ExperimentBackend ABC, BackendCapability
│   ├── fake_backend.py   # in-memory test double (adversarial-state injectable)
│   └── wandb_backend.py  # real W&B implementation
└── analysis/
    ├── comparison.py      # compare_runs (pure diffing)
    ├── divergence.py       # audit_training_curve's 4 signal detectors
    ├── confound.py         # audit_ablation's allowlist + verdict logic
    └── sensitivity.py       # audit_sweep's correlation + significance testing

Two decisions worth understanding before you read the code:

  • Retrieval and judgment are structurally separate, all the way down. Every audit_* tool at the server.py layer does nothing but fetch data and translate backend errors into structured ToolError dicts; the actual heuristics live entirely in analysis/, which has no knowledge an MCP call is even involved. You can unit-test analysis/divergence.py's detectors against a hand-built curve with zero backend, zero MCP, zero network.
  • Backends are capability-declared, not blanket-abstract. list_sweeps has a default NotSupportedError implementation rather than being a required abstract method, so a future backend without a native sweep concept (MLflow, planned for v2) can declare capabilities = {ARTIFACTS} and get a clear refusal instead of a fake mapping. See Appendix A of the design spec for how this was validated against MLflow's actual shape before being frozen.

For the full frozen contract (every field, every tool signature, every adversarial case it's tested against) see docs/design-spec-v1.md. For how it was built, milestone by milestone, including every design flaw caught and fixed along the way, see docs/implementation-roadmap-v1.md.

API examples

Every audit tool call looks like this over MCP (shown here as the raw tool-call JSON an MCP client sends/receives — you won't normally write this by hand, your agent does):

Checking an ablation:

{
  "tool": "audit_ablation",
  "arguments": {
    "baseline": {"backend": "wandb", "entity": "your-team", "project": "mamfac", "run_id": "baseline-run-id"},
    "ablation": {"backend": "wandb", "entity": "your-team", "project": "mamfac", "run_id": "ablation-run-id"},
    "claimed_variable": "use_memory"
  }
}
{
  "verdict": "confounded",
  "confidence": "high",
  "differing_params": [
    {"param": "use_memory", "baseline_value": true, "ablation_value": false, "likely_intentional": true},
    {"param": "batch_size", "baseline_value": 64, "ablation_value": 32, "likely_intentional": false}
  ],
  "method": "full config diff against claimed_variable; params tagged intentional if name matches claimed_variable or is on the allowlist (seed, device, run name/id)",
  "evidence": { "...": "full compare_runs-style diff, config and metrics" }
}

Auditing a training curve:

{
  "tool": "audit_training_curve",
  "arguments": {
    "ref": {"backend": "wandb", "entity": "your-team", "project": "mamfac", "run_id": "xj29fk1a"},
    "metric": "reward"
  }
}
{
  "schema_version": 2,
  "metric_type_assumed": "reward",
  "signals": [
    {
      "signal": "sudden_jump",
      "score": 0.91,
      "step_range": [40120, 40140],
      "evidence": { "...": "the adjacent point pair and rate-of-change values" },
      "confidence": "high"
    }
  ],
  "method": "threshold-based, see docs/audit-methods.md#training-curve"
}

Every field in these responses is real output shape from the current implementation, not aspirational. See docs/design-spec-v1.md §4.2 for the complete, frozen schema of every tool.

Data handling

  • Data never leaves your machine except calls to your own W&B endpoint. This server is stateless and open source — read the code, there's nowhere for your data to go.
  • Credentials are read once from environment variables (WANDB_API_KEY, optionally WANDB_ENTITY), validated fail-fast on server start, and never logged or echoed back in any error message.
  • Use a read-only API key. This server has no write path to W&B — a read-only key is strictly sufficient and reduces what a misconfigured or compromised client could ever do.

Known gaps (honest status)

Two verification steps that this project's own completion criteria call for are written and ready to run, but have not actually been run, because this build environment has no live network path to the relevant services. Documented here rather than silently marked done:

  1. Fixture recording against a real W&B projecttests/ currently test WandbBackend against an in-memory fake client built from W&B's documented API shapes, not against fixtures recorded from a live project. scripts/record_wandb_fixtures.py exists and is ready to run against your own project (e.g. a MAMFAC/CARM++ project) to close this gap. See tests/fixtures/README.md.
  2. Tool-selection eval against a live MCP clientscripts/tool_selection_eval.py and its 15-prompt fixed set (scripts/tool_selection_prompts.py) exist and are exercised for correctness up to the actual API call, but have not been run against a live model (this environment has no ANTHROPIC_API_KEY / network path to api.anthropic.com). See docs/tool-selection-eval.md for the exact command to run this yourself.

Additionally, as of this release:

  • Package name availability: experiment-audit-mcp was confirmed unclaimed on both PyPI and npm as of 2026-07-10. This is a point-in-time check, not a lock — verify again immediately before publishing if time has passed.
  • MCP Registry / Glama / cursor.directory submission has not been performed from this environment (no outbound network access to those services from here). The package is publish-ready; registry submission is a step for whoever runs the actual release from a machine with that access.
  • This is a v1, W&B-only release. MLflow support is prototyped at the interface level (see Appendix A of the design spec) but not implemented.
  • audit_sweep's correlation-based ranking only detects linear relationships — a hyperparameter with a non-monotonic effect (e.g. an interior-optimum learning rate) can rank near the bottom despite mattering most. This is a documented method limitation, not a bug; see docs/audit-methods.md (sweep section).

None of these gaps are architectural. They are either genuine sandbox network limitations or explicitly deferred v2/v3 scope per the frozen roadmap — see Roadmap below.

Development

git clone https://github.com/SreeDharshan-GJ/experiment-audit-mcp.git
cd experiment-audit-mcp
pip install -e ".[dev]"
pytest                # 350 tests
ruff check .          # lint

Everything is developed against FakeBackend (backends/fake_backend.py), an in-memory test double that can inject every adversarial state named in the design spec (tiny sweeps, NaN mid-curve, correlated hyperparameters, partial data) on demand — no live W&B credentials or network access needed to run the full suite.

Contributing

Contributions are welcome. Please read CONTRIBUTING.md first — in short: the v1 design (docs/design-spec-v1.md) is frozen, so changes to existing tool schemas, model fields, or the backend interface need an explicit, logged design decision (a "Revision" entry in the spec, following the pattern of Revision 1 and Revision 2 already in the document), not a silent PR. New audit_* tools in future versions must implement the mandatory method / confidence / evidence schema from day one (spec §8).

Roadmap

See docs/implementation-roadmap-v1.md for the full v1 build history (10 milestones, 2 logged spec revisions, all approved). Looking ahead, per docs/design-spec-v1.md §10:

  • v2 — MLflow backend, versioned API compatibility matrix, first public case study from a real project.
  • v3 — RL-specific pathology signals (reward-hacking heuristics, proper multi-seed statistical tests), optional experimental claimed_variable inference for audit_ablation, Optuna/Ray Tune sweep support, open to external audit_* contributions.

License

MIT — see LICENSE.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

experiment_audit-1.0.0.tar.gz (146.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

experiment_audit-1.0.0-py3-none-any.whl (125.1 kB view details)

Uploaded Python 3

File details

Details for the file experiment_audit-1.0.0.tar.gz.

File metadata

  • Download URL: experiment_audit-1.0.0.tar.gz
  • Upload date:
  • Size: 146.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.9

File hashes

Hashes for experiment_audit-1.0.0.tar.gz
Algorithm Hash digest
SHA256 a26cf5fa6eed5c84ea3932936902d9ab5c2db988003c1d16ebd70c1aa5ce8a87
MD5 db41546a09eaff03dd9ee78191022610
BLAKE2b-256 0bc834b899266fbd283f7aa0c9c61458c190b142843a131526c85a1819d630fe

See more details on using hashes here.

File details

Details for the file experiment_audit-1.0.0-py3-none-any.whl.

File metadata

File hashes

Hashes for experiment_audit-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 63edf7cc2bbc3382204fe4835986c4e184b735d7a8dbcb4dac2e3c2d053577d7
MD5 c002f2c9f932ee0480a25971f8e93a4e
BLAKE2b-256 20f4ac061fe788758a329482aa4141487f6b1bcc462864f29f2e2258b925bedc

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page