DriftGuard — ML Lineage Guardian for DataHub

Autonomous agents that catch silent data changes breaking production ML models — trace DataHub lineage, measure the AUC drop, open a fix PR, and write the warning back via the MCP Server.

Comment

Inspiration

A data engineer renames one column, or switches its unit from cents to dollars. Nothing throws an error. CI stays green. But three hops downstream, that column feeds a feature that feeds churn_predictor v3, live in production — and the model now sees values off by 100×. Its AUC silently drops from 0.85 to 0.76. Nobody notices, until it costs money.

DataHub already knows what connects to what — but its lineage map is passive. We wanted to make it active: an agent that watches the graph, reasons about which changes are dangerous, fixes them, and writes what happened back so the next person or agent inherits the knowledge.

What it does

DriftGuard runs a closed read → act → write-back loop, automatically:

  • Detect — traces the live DataHub lineage from the changed column through the dbt transforms to the production models it endangers (column-precise, no hardcoding).
  • Analyze — an Investigator agent drives DataHub tools, classifies the change as a silent break, and quantifies the damage as a real AUC delta.
  • Remediate — a Remediator agent reads the real schema from DataHub and opens a GitHub pull request with a fail-loud guard.
  • Write back — through the DataHub MCP Server (add_tags, add_owners, save_document) plus the SDK for a native deprecation banner and incident.

A watcher fires the whole loop the moment a watched table's schema changes — the dashboard is a pure live monitor, no button.

How we built it

  • DataHub OSS as the context graph; two lineage layers — a real Postgres + dbt warehouse (raw_orders → stg_orders → customer_orders / customer_360) ingested with DataHub's Postgres + dbt connectors, connected to emitted ML lineage for two production models.
  • Two agents on GMI Cloud (google/gemini-3.5-flash-lite, OpenAI-compatible) via function-calling.
  • Real MCP integration — the write-back drives the official mcp-server-datahub mutation tools (schemas pinned from the live list_tools()).
  • A genuine gradient-boosted churn model (scikit-learn) so the silent break is a measurable AUC drop, proven by tests.
  • Next.js + FastAPI dashboard, PyGithub for PRs, GitHub Actions to run the Guardian on data/model PRs. One make compose-up brings the whole stack up.

Challenges we ran into

  • Gemini's multi-turn tool-calling 400'd on re-sent assistant tool calls (a "thought_signature" quirk) → we split the agent loop into two phases (choose tools → synthesize a JSON verdict).
  • MCP mutation tools validate that referenced entities exist — add_tags/add_owners reject unknown tag/user URNs — so we seed the tag + agent-user vocabulary via the SDK first.
  • The stdio MCP client raises a benign ExceptionGroup at teardown → populate the result inside the session and swallow the noise.
  • A resolved DataHub incident stayed "active" in the search index → switch to a hard delete on reset so re-runs start truly clean.
  • A dependency conflict (acryl-datahub[postgres] pins sqlalchemy<2) would have broken a clean make install — caught and fixed.

What we learned

The model is trained on correct-unit features, so a serving-time unit change is training-serving skew: it degrades AUC even though no exception is ever raised. Tree models are not rank-invariant to a fixed-threshold unit shift — our test suite proves the drop. And DataHub is far more than a catalog to read: with the MCP Server you can genuinely act on and contribute back to the graph, turning passive lineage into an early-warning system.

Accomplishments we're proud of

  • Real, verified MCP write-back (read back from the graph), not just a configured server.
  • Precise blast radius — a unit change on purchase_amount_cents flags churn_predictor and correctly leaves the unrelated ltv_predictor alone.
  • Two open-source DataHub Skills contributed upstream.

What's next

  • Route the Investigator's read tools through the MCP server too (write already does).
  • More detectors: stale training data, type/nullability drift.
  • Finish the GitHub Actions gate that comments the blast radius and fails on a CRITICAL finding.

Built With

Share this project:

Updates