ML Incident Copilot

Agent that walks DataHub lineage graphs to catch schema drift & PII exposure before they break production ML models.

Comment

Inspiration

ML models fail silently. A source table gets a column renamed, a new PII-tagged field gets added upstream, a data source quietly goes stale — and nothing breaks loudly. Model quality just slowly degrades until someone notices weeks later, usually the hard way. The information to catch this earlier already exists: DataHub's lineage graph knows exactly what feeds into every model. The problem is that a human has to remember to go look. We wanted an agent that looks automatically.

What it does

ML Incident Copilot walks a production ML model's real upstream lineage graph in DataHub — training data → features → model — using the DataHub MCP Server. At each node it checks for:

  • Schema drift — columns added, removed, or renamed since the last check
  • PII exposure — sensitive/PII-tagged data feeding a model without documented justification
  • Stale data — sources that haven't refreshed within their SLA
  • Deprecated sources — deprecated/archived datasets still live in a production lineage path

When it finds something, it doesn't just alert — it writes the incident directly back onto the affected entity in DataHub (a tag, a structured property, and a saved incident document), so the next engineer or agent who opens that table in DataHub inherits full context immediately. It also opens a GitHub issue and posts to Slack, so a human gets pulled in outside the catalog too.

How we built it

The core is a Python agent (agent/main.py) that connects to mcp-server-datahub over stdio using the official MCP client library, walks lineage via the get_lineage tool, pulls schema via list_schema_fields, runs a set of pluggable risk checks (agent/risk_checks.py) against each node, and hands off anything it finds to a write-back layer (agent/writeback.py) that talks to DataHub, GitHub, and Slack.

We seeded a realistic lineage graph ourselves: raw.customers (PII-tagged) and raw.transactions feeding two real MLFeature entities, which feed an MLModel entity (churn_predictor_v1) via its mlFeatures property — deliberately using real DataHub entity types and a real, graph-indexed lineage path rather than a shortcut.

Challenges we ran into

Almost everything here was discovered the hard way against a live DataHub instance rather than assumed from docs:

  • Our first attempt at model→feature lineage used the mlModelTrainingData aspect — it emits successfully but isn't indexed by DataHub's lineage graph at all. We only found this by checking the UI and seeing an empty Lineage tab, then found the underlying GitHub issue confirming it. Switching to real MLFeature entities with an indexed sources field fixed it for good.
  • The MCP server's tool response shapes are not uniform. search wraps results in searchResults, get_lineage nests them inside upstreams/downstreams, list_schema_fields returns a flat fields array, and get_entities wraps everything in result. We had to inspect each one individually via the MCP Inspector against live data — an "empty" test call with total: 0 looks identical to a working call with the wrong field name, so we learned to always verify with a populated result before trusting a shape.
  • Tags came back in a nested shape we didn't expect — {"tags": [{"tag": {"urn": "urn:li:tag:PII"}}]} — with the tag name only recoverable by parsing the urn suffix, not a clean name field.
  • A subtle Python bug: caching a set() to JSON via json.dumps(..., default=str) silently corrupts it into a string instead of erroring, which made schema-drift detection appear to work at first but quietly never actually compare anything.
  • Windows-specific friction throughout — python3 vs python, set vs export, Python 3.14 failing to build pydantic-core (fixed by pinning to 3.11), and Docker Desktop needing more memory than its default to run the full DataHub quickstart stack.

Accomplishments that we're proud of

  • A fully working, end-to-end agent against a real DataHub instance — not a mocked demo. Every tool call, every response shape, every write-back was verified against a live server, not assumed.
  • Getting model→feature→dataset lineage genuinely indexed and visible in DataHub's own Lineage UI, after an earlier approach silently failed to show up there at all.
  • Catching and fixing a real, non-obvious bug (the set-to-JSON serialization issue) that would have made schema-drift detection quietly never work, without ever throwing an error to flag it.
  • A complete round trip: an actual schema change in DataHub triggers detection, a tag + structured property + saved document written back onto the entity, a GitHub issue, and a Slack message — all from one command.
  • Documentation good enough that someone else (or a judge) could actually get this running, including exact steps for getting a DataHub token, a GitHub token, and a Slack webhook — the parts that are usually left as an exercise for the reader.

What we learned

Trust the live response over the docstring. Almost every real bug here came from an assumption about a tool's output shape that turned out to be subtly wrong — and every one of them was resolved fastest by adding a raw debug print and looking at the literal dictionary, rather than guessing again. We also learned that DataHub's various aspects aren't interchangeable for lineage purposes — which aspect you use to link entities determines whether the graph is actually queryable, not just whether the emit call succeeds.

What's next for ML Incident Copilot

The README includes a full "Path to Production" section, but in short: move the local JSON cache to shared storage so multiple runs don't clobber each other, move from manual/cron runs to DataHub's own Actions framework so this reacts to changes in real time instead of on a schedule, add incident state tracking (open/acknowledged/resolved) so it stops re-notifying on already-known issues, and paginate lineage walks for catalogs with real scale instead of a flat 30-result call.

Built With

Share this project:

Updates