TrainGate

TrainGate: Detect. Heal. Verify. Promote - Continuous training, gated by live DataHub.

Comment

Inspiration

Everyone is shipping AI agents and continuous ML training. The failure mode that actually wakes up on-call is quieter: an upstream pipeline changes units, renames a column, or lags labels. The retrain job still "succeeds," and production predictions go wrong before anyone notices.

Lineage scanners are great at reporting drift after the damage. We wanted something that sits on the promote path: refuse to ship a model unless live DataHub context says the data supply chain is safe. When it isn't, heal what we can, prove it again, then promote.

DataHub's MCP Server, ML lineage, and Skills made that feel buildable for a hackathon and useful for a real platform team.

What it does

TrainGate is a continuous-training control plane gated by live DataHub:

  1. Connect to DataHub GMS
  2. Inspect feature-table lineage, schema, and custom properties (e.g. distance_unit)
  3. Diagnose with a deterministic Guard (lineage, schema, unit contract, nulls, label freshness)
  4. Raise a first-class DataHub Incident when checks fail
  5. Remediate known failures (unit scale, schema rename, null spike, label lag)
  6. Re-verify against live DataHub
  7. Resolve the incident with heal notes
  8. Train + promote only if clean, registering mlModel versions, tags, descriptions, and deployment lineage back into DataHub

The UI shows a live supply-chain graph (from DataHub lineage), customer impact (e.g. ETA ~74 min to ~37 min after a unit poison heal), and a step-by-step agent timeline. An MCP server lets Cursor/Claude drive the same gate while the dashboard updates live. Sample models: delivery ETA and order cancel-risk.

How we built it

  • Backend: FastAPI agent loop (detect to incident to heal to verify to promote), scikit-learn training, parquet feature store, chaos injectors for silent failures
  • DataHub: official Python SDK for entities/lineage/tags/models; GraphQL for raise/resolve Incidents
  • Frontend: React + Vite dashboard (lineage graph, animated agent timeline, run history)
  • MCP: zero-dependency stdio MCP server proxying the TrainGate API
  • OSS: DataHub Skill at skills/traingate-promote-gate/ documenting the promote-gate procedure
  • Local stack: DataHub quickstart (GMS + UI) plus API plus frontend

Challenges we ran into

  • Silent failures aren't always visible. Tree models don't extrapolate wildly, so "poisoned train / clean serve" didn't look dramatic. We modeled reality: the broken upstream feeds serving too, so customers feel unit poison immediately.
  • DataHub Incidents API details. Resolve mutations needed the correct GraphQL input type for our quickstart version. Freshly raised incidents can lag in search indexes, so we resolve by the URN we just created.

Accomplishments that we're proud of

  • A full write-back loop to DataHub, not read-only chat over metadata
  • Real Incidents raised and resolved with heal notes, visible in the DataHub UI
  • An MCP server so any agent can run the same promote gate
  • End-to-end tests (Guard, pipeline, MCP) plus a live demo that judges can reproduce from the README

What we learned

  • Metadata platforms win when agents act and write back, not only retrieve context
  • Deterministic gates beat LLM-only decisions to promote safety (LLMs can narrate; they shouldn't decide ALLOW/BLOCK alone)

What's next for TrainGate

  • More remediation playbooks and multi-table contracts
  • Webhooks / Slack when an incident is raised, or a promote is blocked
  • Propose the promote-gate Skill upstream to DataHub Skills
  • Hosted MVP once DataHub + API can run affordably in the cloud

Tech: DataHub, FastAPI, scikit-learn, React, MCP, Apache 2.0
Data: synthetic delivery features (orders, couriers, weather to eta_features)

Built With

Share this project:

Updates

Submission history