-
-
The exact features implicated: refund_count and days_since_last_refund, both derived from data that postdates the outcome.
-
Sentinel's write-back: tags, rewritten description, and the real AUC gap it's flagging — 0.85 with the leaky features vs. 0.78 without.
-
Both findings written into the model's docs, each linked to a full incident memo — real DataHub state, not a mockup.
Inspiration
Production ML models fail silently in two specific, expensive ways that no monitor catches: target leakage (a feature secretly derived from data that only exists after the outcome you're predicting) and silent upstream drift (a schema change nobody repointed the pipeline to). Both look completely fine on a dashboard. Both are only visible if you actually follow the data's lineage — which is exactly what DataHub is built to expose. That felt like the real "why DataHub" story for the Production ML Agents track, rather than another read-only chat-with-your-catalog agent.
What it does
Sentinel audits every ML model in a DataHub catalog. For each one, it:
- Reads real dataset descriptions and asks an LLM to judge, from the prose alone, whether the data is causally downstream of a business outcome — discovering leakage risk itself rather than reading a flag planted for it.
- Proves the finding deterministically by walking DataHub's lineage graph via MCP.
- Explains it in plain English.
- Writes the finding back into DataHub: tags, an updated model description, and a linked incident memo.
- Generates the fix as a real diff against the dbt project that built the feature, and can open a genuine, mergeable pull request.
- Does this for every model sharing the risky data — proven with two models sharing one feature store, both caught independently from a single triage pass.
- Investigates further: once a finding is proven, hands the LLM the read-only DataHub toolset and lets it choose its own next call -- step by step, based on what the previous call returned -- to gather context beyond the finding itself, before concluding.
How I built it
- DataHub OSS (
docker quickstart) as the metadata graph, ingested via the Python SDK: full ML lineage chain (MLFeatureTable -> MLFeature -> MLModel -> MLModelGroup -> MLModelDeployment), including real column-level (fine-grained) lineage, not just table-level. - All reads and writes go through
mcp-server-datahubover stdio — the official MCP server, the same way Claude Code or Cursor would talk to it. No direct REST/GraphQL calls from detector, triage, or action code. - Detection is split deliberately: an LLM (Ollama locally, or Gemini) only discovers governance risk from prose and explains proven findings — it never invents a finding. The actual leakage/drift detectors are pure, deterministic graph traversals with no LLM in the loop, so the same metadata always produces the same finding.
- A synthetic e-commerce dbt + DuckDB pipeline with a scikit-learn churn model, with two real (not hardcoded-in-the-detector) bugs baked into the data itself.
- GitHub Actions CI runs the entire pipeline — real DataHub, ingestion, agent — on every push, and fails the build if Sentinel doesn't actually find both bugs. Not a smoke test; a standing proof.
- An agentic investigation layer (investigate.py): the one place the LLM actually drives, choosing its own MCP tool calls around an already-proven finding. If it runs out of steps without concluding on its own, a final forced-conclusion call guarantees a grounded answer instead of trailing off.
- A GitHub Codespaces environment (.devcontainer/) so judges can test without installing Docker locally -- click the badge, same make commands work.
- Two detectors (detect_missing_ownership, detect_undocumented_upstream) that fire on any real DataHub catalog, not just this repo's fixtures -- validated against DataHub's actual showcase-ecommerce sample pack rather than only this project's own synthetic data.
Challenges I ran into
- Making the leakage detection discovery, not staged: an earlier version pre-tagged the leaky tables at ingest time, which is exactly the kind of thing a skeptical reviewer should distrust. I rebuilt it so an LLM reads honest, human-written descriptions and decides on its own — then validated it three ways, including against real third-party DataHub sample data it had never seen (13/13, zero false positives).
- A real DataHub search-index race condition: the OpenSearch index updates asynchronously after a write/ingest is acknowledged, so a CI run that ingests and immediately searches can miss freshly written data. I found this via CI (not locally), and fixed it with a poll-until-consistent step.
- I investigated using DataHub's real column-level (fine-grained) lineage to make the leakage detector name the exact tainted column instead of matching on feature-name substrings. I confirmed the lineage is stored correctly, but no current MCP tool actually surfaces it back — rather than break this project's own MCP-only architecture to work around that gap, I filed it as a documented finding upstream instead.
- A targeted, skeptical review of the less-tested code paths (real PR
opening, Slack notifications,
--watchmode) turned up four real reliability bugs — including one that would crash the agent on a second run — all fixed and covered by regression tests.
What I learned
- The credibility of "the agent found it" only survives if the finding is provably not something you told it to find. It's worth the extra build time to prove that with independent data the project never touched.
- Reliability bugs hide in the paths you don't demo every day (retry logic, opt-in features, long-running modes) — a deliberately skeptical re-read found more real issues than another pass over the happy path.
- Real production metadata graphs have real gaps in the tooling around them. Two of those gaps (search read-after-write consistency, and no MCP tool exposing column-level lineage) were worth documenting upstream rather than working around silently.
What's next for Sentinel
- Wire detection into
--watchmode against DataHub's real-time change events, so it audits new models continuously rather than on a manual run. - Extend the detector set: staleness beyond training-data age, and ownership-gap detection (models with no on-call owner in DataHub).
- Once
mcp-server-datahubexposes column-level lineage (see the filed PR), upgrade the leakage detector to name the exact tainted column instead of the table.
Built With
- datahub
- dbt
- docker
- duckdb
- gemini
- github-actions
- github-codespaces
- langgraph
- mcp
- model-context-protocol
- ollama
- pytest
- python
- scikit-learn
- sql
Log in or sign up for Devpost to join the conversation.