mlgate

DataHub knows your owners, PII tags and freshness. It has no opinion. mlgate is that opinion: it gates ML releases on lineage and writes the verdict back.

Comment

Inspiration

DataHub already knows everything you need to answer one question: is this model safe to ship? It stores owners, PII tags, freshness assertions, lineage, schema history. What it doesn't do is form an opinion about any of it.

So before a model goes to production, the check means opening the catalog and walking a dozen tables by hand — which means in practice nobody does it. Models reach production trained on tables nobody owns, containing PII nobody flagged, from sources that quietly stopped updating weeks ago.

mlgate is that missing opinion, and it puts the answer back in the catalog.

What it does

Point it at a model. It walks the lineage backwards through DataHub — deployment → model → features → source datasets — evaluates a rule pack against every asset, and returns PASS / BLOCK with named blockers. Exit code 1, so it drops straight into CI.

Then it writes the verdict back into DataHub through five paths: a Governance:ReleaseBlocked tag on the model, per-finding governance tags on each offending asset, readiness properties, the release dossier rendered into DataHub's Documentation tab, and each run recorded as a DataProcessInstance so the history is queryable.

Beyond a single check: scan gates the whole fleet concurrently, fix remediates what it can honestly, watch re-gates on a schedule and reports only what changed, and --override is an attributed break-glass path that never erases the block.

How we built it

Python on DataHub Core v1.7.0 (docker quickstart), built on the showcase-ecommerce datapack — which ships no ML entities at all, so the seeder emits the full dataset → MLFeatureTable → MLModel → MLModelDeployment chain itself. The gate therefore traverses genuine catalog data, not only fixtures.

It reads five real DataHub surfaces:

  • the entity API for aspects (never search — see below)
  • the relationship graph (OUTGOING/DownstreamOf) for transitive upstream lineage
  • GraphQL for assertions with their run outcomes
  • the Timeline API for real schema changes since the model trained
  • DataJobInputOutput.inputDatasetFields for column-level lineage

Rule logic is Python; rule policy is YAML, so a team adds a pack without touching the engine. It also ships as a DataHub Skill so an agent can drive it conversationally.

Challenges we ran into

The search index lags the entity store. Right after a write, search returns zero hits while the entity is already durably stored. A gate that walked lineage via search would find nothing, report no findings, and print PASS — approving a model because an index was three seconds behind. So: entity API only, and an empty walk raises an error (exit 2) rather than passing.

Two of our own rules were meaningless. Freshness originally read a custom property our seeder wrote, and drift compared two integers we invented — both proved nothing on a real catalog. They now read DataHub's actual Assertion model and Timeline API.

no_owner checked whether an Ownership aspect existed, not whether it had owners. A dataset with zero owners passed. Testing found it — and that turned out to be a class, not an instance: three more paths where absent metadata produced no finding and read as "fine."

DataHub rejects a schemaField URN in MLFeatureProperties.sources (HTTP 422). Column-level truth belongs on the job, not the feature.

EditableMLModelProperties shadowed MLModelProperties in our aspect lookup, blanking the model's features. It surfaced as exit 2, not a false PASS — the safety design catching a bug we introduced hours later.

Accomplishments that we're proud of

Every safety property is tested by its negative case. The gate is proven to pass a clean model. The PII rule is proven not to fire on a documented column. A waiver is proven to stop working when it expires. And the acceptance suite itself was mutation-tested — we broke the verdict logic and disabled a rule to confirm it can actually fail.

make verify fails the build if the broken model ever passes. A gate nothing can fail is not measuring anything.

Three states, not two: compliant, violating, and unverifiable. "We could not check" is never reported as "it's fine."

31 unit tests, 45 end-to-end checks, 16 committed sample artifacts.

What we learned

That the hard part of a governance tool isn't the rules — it's refusing to report false comfort. Almost every bug we found had the same shape: absent data rendering as clean data. Catching one instance was worth less than recognising the class and sweeping for the rest.

Also that existence is not governance. A freshness assertion that has never run proves nothing; one that is failing means the data is already stale.

What's next for mlgate

  • Subscribe instead of poll — watch re-gates on an interval; DataHub emits metadata change events, which would make it instant.
  • Signed dossiers — the current hash proves integrity, not provenance.
  • Fleet-scale exemption policy — glob matching won't scale to hundreds of models.

Known limits are documented in the README rather than hidden. The dossier records data-provenance readiness — a precondition for NIST AI RMF and EU AI Act technical documentation. It is not a compliance certification and asserts conformity with no regulation.

Built With

  • airflow
  • apache-datahub
  • ci-cd
  • data-lineage
  • data-quality
  • datahub
  • dbt
  • docker
  • github-actions
  • graphql
  • metadata
  • ml-governance
  • mlflow
  • mlops
  • model-governance
  • pytest
  • python
  • rest-api
  • snowflake
  • yaml
Share this project:

Updates