Inspiration

DataHub's own team has pointed out that most ML production failures don't come from model drift — they come from upstream data problems that never get traced back to the models depending on them. At the same time, feature stores have a structural blind spot: a feature's declared source column is written once, at registration time, and almost never checked again against the live warehouse. Schemas evolve constantly — columns get renamed in refactors, dropped when tables are restructured, retyped when upstream ETL logic changes — and nothing closes the loop back to the feature store.

We wanted to build the thing that closes that loop: an agent that treats feature-store metadata as a claim, and checks it against reality.

What it does

Ghost Features is an ML observability agent that detects orphaned feature- store entries — features whose upstream warehouse column has been renamed, dropped, or changed type without the feature store being told. It:

  1. Reads a feature table's registered features and their declared source columns from DataHub.
  2. Reads the current live schema of each source table.
  3. Compares declared vs. actual — flagging GHOSTED (column gone) or TYPE_CHANGED (column exists, but its type no longer matches what was recorded).
  4. Cross-references DataHub's ML model catalog to identify which production models depend on the broken feature.
  5. Writes a remediation Document back into DataHub for each finding — contributing to the graph, not just reading from it.

In testing, it correctly caught three independent, realistic failures in a single run, all traced back to the same production model (fraud_risk_model_v1):

Feature Failure Detail
customer_email_domain Renamed column cust_emailemail_address
avg_order_value Type change unit_price FLOAT → VARCHAR
customer_lifetime_orders Dropped column customer_id removed entirely

How we built it

  • DataHub (open-source, v1.5.0.6, local via docker quickstart) is the system of record for both warehouse schema metadata and ML feature/model metadata.
  • DataHub's MCP Server (acryldata/mcp-server-datahub) is the primary interface the agent uses at runtime — get_entities, list_schema_fields, and save_document, via a small async Python wrapper (scripts/datahub_mcp_client.py) around the official mcp client SDK.
  • Impact tracing reads the same relationship metadata DataHub's ML lineage view is built from — mlFeatureProperties.sources and mlModelProperties.mlFeatures — to walk from a broken feature to every production model consuming it.
  • Two narrow, explicitly documented calls fall back to DataHub's GMS REST API, because we verified live that the MCP server's GraphQL layer doesn't yet expose mlFeatureProperties via get_entities, and rejects ML_MODEL as a search entity type on this GMS version. We filed this upstream (see "open-source contributions" below) rather than silently working around it.
  • Seeding and simulation scripts (seed_ml_features.py, simulate_schema_change.py, simulate_type_change.py, simulate_dropped_column.py) use the DataHub SDK directly — this is one-time environment setup to build a realistic demo scenario, not the agent's runtime path.

What we learned

  • DataHub's ML metadata model (MLFeatureTableMLFeatureMLModel.mlFeatures) is expressive enough to model real feature-store dependency graphs, but the MCP server's current tool surface doesn't fully cover ML entity types yet — a gap worth fixing upstream, so we did.
  • Running a full DataHub quickstart stack (GMS, Kafka, OpenSearch, MySQL) on an 8GB laptop is workable but genuinely tight — we hit a real OpenSearch field-mapping limit and Kafka consumer lag under memory pressure, both documented below.
  • Verifying against a live system beats assuming: partway through, an agent fell back to a mock server when the real one was down and reported clean results — a false success that would have shipped broken behavior baked in if we hadn't caught it and insisted on a real end-to-end re-run.

Challenges we ran into

  • OpenSearch field-mapping limit. DataHub's top_queries analytics index hit its default 1000-field cap (Limit of total fields [1000] has been exceeded), which surfaced as a generic "Something went wrong" error on the DataHub homepage and empty search results — with no obvious link to the real cause. Fixed by raising index.mapping.total_fields.limit.
  • Kafka consumer lag under memory pressure. The generic-mae-consumer- job-client consumer group (which populates the search index from the metadata change log) intermittently showed "no active members" and thousands of messages of lag, meaning newly-written entities were correct in the primary datastore (confirmed via direct datahub get calls) but invisible in search and the UI. We designed the agent to read by known URN through MCP tools rather than relying on search, which sidesteps this entirely and is arguably more realistic for how a production agent would operate anyway.
  • MCP server ML entity coverage gap. get_entities doesn't return mlFeatureProperties for ML feature entities, and the search tool's GraphQL layer raises a ValidationError for ML_MODEL as an entity type. Verified with a live introspection of the MCP server's real tool schemas (not assumed), documented in docs/mcp-server-ml-entity-gap.md, and drafted as an upstream contribution.
  • CLI/server version mismatch. Running a newer datahub CLI (1.6.0.16) against an older GMS (1.5.0.6) allowed writes to silently partially succeed — entities existed but were missing aspects needed for proper indexing — with only a small warning easy to miss in verbose output. Pinning CLI and server to matching versions fixed it.

Built With

Share this project:

Updates