Inspiration

Some of the most expensive ML failures do not crash a job or break a schema. The model keeps training and serving, but its quality quietly gets worse. When a team finally notices, the upstream change may be buried under weeks of unrelated work.

DataHub already contains the evidence needed to explain such failures: model versions, training runs, metrics, input datasets, column-level lineage, and historical column statistics. Cold Case joins those facts across time to answer one question:

Why did this model quietly get worse?

What it does

Cold Case starts from an ML model in DataHub and autonomously:

  1. Discovers the model using the DataHub Agent Context Kit.
  2. Reconstructs its model-version and training-run history.
  3. Finds the exact version where a quality metric regressed.
  4. Compares the last good and first bad runs.
  5. Rules out changes to model settings, features, input tables, label prevalence, and training volume.
  6. Isolates the new data window seen only by the bad version.
  7. Walks column-level lineage upstream and examines historical column profiles.
  8. Detects transforms such as COALESCE that can hide broken source data.
  9. Produces an evidence-backed diagnosis with confidence and recommended fixes.
  10. Writes a post-mortem, tags, and links back into DataHub.

Cold Case then reads its writes back through 11 verification checks. When run again, it finds its previous post-mortem and stops instead of repeating the investigation. The next person—or agent—inherits the answer.

Demo incident

The repository creates a reproducible synthetic payments warehouse and trains three real fraud-detection models using scikit-learn and MLflow.

Version 3 drops from 0.8803 to 0.8507 ROC AUC, even though the model settings remain unchanged and no job fails.

Cold Case discovers that payments.raw.devices.device_reputation changed from 0% to 40% missing values on June 8. A downstream COALESCE replaced those missing values with zero. Therefore, the feature table still showed 0% missing values, every downstream null check passed, and the model treated the substituted value as a real measurement.

Only version 3 trained on the affected days.

Although the incident is synthetic, the models, MLflow registration, metric regression, DataHub metadata, lineage traversal, historical profiles, and write-back are real.

A three-seed treated-versus-control audit confirms causality: the injected corruption causes an average 0.0334 ROC AUC loss, while unaffected training windows remain exactly equal.

How we built it

The project has two main parts:

  • seed/ generates 90 days of payments data, injects the silent failure, builds warehouse tables, trains and registers three MLflow model versions, ingests their metadata, and emits schemas, column-level lineage, and 634 daily profiles into DataHub.
  • coldcase/ implements the investigation workflow: memory, model timeline, run comparison, historical lineage analysis, evidence ranking, and verified write-back.

Cold Case uses:

  • DataHub Agent Context Kit for initial model discovery
  • DataHub Python SDK for metadata aspects
  • GraphQL for relationships and graph traversal
  • Timeseries aspects for historical dataset profiles
  • Timeline API for schema history
  • OpenAPI and REST metadata change proposals
  • MLflow for experiment tracking and model registration
  • scikit-learn for real model training

Challenges we faced

DataHub’s ML entities form a second graph layered on the dataset graph. Cold Case must cross model groups, model versions, training runs, feature references, and warehouse datasets without confusing their relationship directions.

Historical profiles are timeseries metadata, so the investigation must examine the correct training window rather than only the current state. MLflow can also create platform-local dataset references, requiring traversal to continue upstream to the real warehouse table.

Search indexing is eventually consistent, so Cold Case verifies write-back through authoritative aspect reads and bounded polling instead of assuming an HTTP success means the metadata is visible.

Accomplishments

  • A complete end-to-end investigation verified on DataHub 1.7
  • A 95% confidence diagnosis backed by a 15-item evidence chain
  • Eleven read-back checks proving the document, relationships and tags were saved
  • A credential-free self-test against a completely different model and failure
  • A causal A/B audit proving the injected corruption caused the measured loss
  • Real transcripts, JSON reports and post-mortems in the repository
  • A reusable DataHub ML Lineage Skill contributed upstream

What we learned

The table closest to a model is often the worst place to detect a data failure. Imputation can make a feature table look complete while silently destroying information the model depends on.

The answer only becomes visible when model timing, column-level lineage, transformation semantics and historical profiles are examined together.

We also learned that trustworthy diagnostic agents should separate fact collection from judgment. Cold Case never lets a language model decide what the graph says, and no API key is required.

Open-source contribution

We contributed the reusable investigation workflow as a new DataHub Agent Skill:

datahub-ml-lineage pull request

The contribution is currently under review.

What’s next

  • Detect gradual drift as well as sudden changes
  • Trigger investigations automatically when model metrics regress
  • Support additional ML platforms and deployment patterns
  • Turn remediation into approval-gated DataHub assertions
  • Let related models reuse knowledge from previous incidents

Built With

Share this project:

Updates