ML Model Guardian

Blocks schema changes that break ML models. It asks DataHub which feature dies with the column, how much the model actually used it, and turns that answer into a failing CI check.

Comment

Inspiration

Every data engineer has shipped a DROP COLUMN that passed review, passed the linter, and took down a model three days later. The diff looked harmless. The information that would have made it look dangerous existed — it just lived in the catalog, not the pull request.

I wanted to close that gap literally: put the catalog in front of the merge button.

What it does

Model Guardian receives a schema diff from CI and answers a question a linter cannot: does any ML model actually depend on this column, and how much?

It walks DataHub's lineage graph downstream, breadth-first. A table is never wired straight to a model — the real path runs customers.email → feature table → training job → model. Then it reads DataHub's fine-grained lineage to name the individual feature that dies with the column, along with its transform:

customers.email → email_domain, computed as SPLIT(email, '@')[1]

Then it does the thing that makes it usable rather than annoying. It reads the model's measured permutation importance back out of DataHub and separates two questions that most tooling conflates:

email_domain measures -0.0018 — accuracy would not suffer. But the training job reads the column directly and dies with MissingSourceColumn. Remove the feature from the pipeline in the same change and this becomes safe.

The verdict becomes a failing commit status on the real pull request, making the merge button unavailable. It also lands back in DataHub as a tag and a schema Assertion with a FAILURE run event, so the block is visible to anyone who opens the table. A human can override; the same status flips to green carrying the approver's name and reason, so a passing check never misrepresents an accepted risk.

The same guard is also an MCP server, so an agent can ask before the pull request exists.

How I built it

The agent is FastAPI over the DataHub Python SDK. Lineage traversal uses scroll_lineage; semantics come from get_aspect on EditableSchemaMetadata and GlossaryTerms; write-back goes through DatahubRestEmitter.

The reasoning is deterministic, not an LLM. Schema comparison and graph traversal are exact operations — a model guessing at them adds risk without adding information. The graph is the only source of truth, so when DataHub is unreachable the agent returns 503 and refuses to rule rather than emitting a false "safe".

The ML layer is real. ml/train.py trains a HistGradientBoostingRegressor on 24,699 customers to predict spend in the four months after the feature window, split temporally so features never see the label period, and publishes the metrics it measured (MAE 2054, R² 0.38) plus per-feature permutation importance to DataHub. Nothing on the model card is typed by hand.

ml/features.py is the single source of truth for the feature contract. The lineage published to DataHub is generated from the same spec the pipeline runs, so the catalog cannot describe features nobody computes — the way dbt and Spark listeners do it, where metadata is a by-product of the job rather than a description someone has to remember to update.

Challenges I ran into

DataHubGraph.get_lineage does not exist in acryl-datahub 1.6. My first implementation called it, swallowed the AttributeError, and fell back to a hardcoded model — so the lineage traversal never ran even once while appearing to work. The real API is scroll_lineage.

upstreamLineage is rejected on an mlModel entity — GMS answers 422 Unknown aspect. Dataset-to-model lineage has to run through a DataJob, which is why a training job sits in the middle of our graph.

schemaField URNs are not indexed by the lineage endpoint. Column-level traversal returns zero relationships, so fine-grained edges have to be read off each downstream dataset's upstreamLineage aspect instead.

A malformed assertion URN was accepted silently. We built it by embedding the dataset URN, whose :, ( and , are reserved in an assertion id. GMS returned success and the entity was never addressable — the "failed Assertion" we were claiming simply did not exist until we hashed the id.

globalTags is replace-style. Writing MLGuardBlocked on its own deleted the pii classification the guard exists to protect. Read-modify-write was required.

GitHub refuses check-run creation from user tokens (403 You must authenticate via a GitHub App), so the default gate is the commit status API — which blocks merges just as hard once required, and needs no App installation to reproduce.

Accomplishments that I am proud of

It blocks a real pull request. PR #1 is red right now and the merge button is gone.

It goes to column level, not table level. "A model downstream is affected" is a warning; "the feature email_domain stops being computable" is something you can act on.

It measures instead of asserting. The importance in the verdict was produced by an actual training run, and python -m ml.predict --drop-column email demonstrates the failure rather than describing it.

It refuses to guess. No downstream model means allowed. An unreachable graph means 503, never "safe". 54 tests cover exactly the paths where an earlier version quietly invented an answer.

What I learned

The hard part of metadata-driven agents isn't the agent — it's that a catalog will happily accept and return things that aren't true. Three of our worst bugs were writes DataHub accepted while producing nothing usable. We ended up treating "no evidence" as a first-class outcome distinct from "no problem", which changed the design more than any feature did.

I also learned that hop distance is a bad proxy for severity. Once a feature table sits between the table and the training job, the model is three hops away while its input genuinely disappears. Evidence beat distance.

What's next for ML Model Guardian

Authentication on the override endpoint — a governance tool where anyone who can reach the port can lift a block is ironic, and it's the first thing Id fix. Beyond that: a GitHub App installation so the full report renders on the Checks tab, semantic invariants past drop-and-retype (unit changes, nullability, cardinality shifts), and reading impact from serving lineage rather than training lineage so the blast radius includes live inference.

Built With

Share this project:

Updates