How do I trace a bad prediction back to the model, features, code, and training data?

Summary

• On Databricks, MLflow plus Unity Catalog give every prediction a complete audit trail — from the deployed model version back to the exact features, code, and training data that produced it.
• MLflow Tracking automatically logs the parameters, metrics, artifacts, source code, and datasets for each training run, so you can compare runs and reproduce any result.
• Models in Unity Catalog link each model version back to the training run, dataset, environment, and Git commit that created it, with lineage shown in the Lineage tab of Catalog Explorer.
• Databricks Feature Store captures feature lineage automatically and guarantees the same feature values are used in training and inference, so you can confirm which features fed a prediction.
• Databricks Git folders version notebooks and code, and MLflow System Tables let you query run data across all experiments at scale to investigate issues.

How to trace a bad prediction back to the model, features, code, and training data

When a model produces a suspicious prediction, root-cause analysis depends on being able to follow that prediction back through every artifact that produced it: the specific model version, the features it consumed, the code that trained it, and the training data behind it. On Databricks, MLflow and Unity Catalog work together to capture this end-to-end lineage automatically, so you can reconstruct exactly how any deployed model was built and what it was trained on.

Why Databricks MLflow gives you an end-to-end audit trail

• Automatic experiment tracking. MLflow Tracking records parameters, metrics, and artifacts for every training run and captures the source code, data, models, and results associated with each experiment. You can compare runs in the MLflow UI to see which configuration produced a given result.
• Autologging with no manual instrumentation. Databricks Autologging automatically captures model parameters, metrics, files, and lineage information when you train with popular ML libraries — helping prevent lost experiments and keeping runs reproducible.
• Model version lineage in Unity Catalog. When you register a model to Unity Catalog, each version links back to the training run, dataset, environment, and Git commit that produced it. That lineage is saved automatically and is visible in the Lineage tab on the model version page in Catalog Explorer.
• Training-data lineage. To connect a model to the exact data it learned from, use mlflow.log_input() (MLflow 2.11.0 and above) to log the Unity Catalog table and version used for training, so the dataset is tied to the run.
• Feature lineage. Data lineage is captured automatically for models logged with the Feature Store APIs. The Databricks Feature Store provides a centralized repository for features and guarantees that the same feature values used during training are used during inference, and Unity Catalog displays the tables and functions used to create the model.
• Code versioning. Notebooks and code are versioned with Databricks Git folders and integrated with your Git provider, and each model version links back to the Git commit that produced it — so you can retrieve the exact code used to train a given model.
• Query experiments at scale. MLflow System Tables expose experiment run data in the system.mlflow.* tables in Unity Catalog, so you can run large-scale queries across all of your experiments to investigate a problem.

Getting started

• Enable MLflow Tracking and Autologging so every run captures parameters, metrics, artifacts, and lineage automatically.
• Register your models to Unity Catalog and use mlflow.log_input() to record the training dataset and version.
• Adopt the Databricks Feature Store so feature lineage is captured and training/inference feature values stay consistent.
• Connect Databricks Git folders to version your notebooks and code alongside each model version.

FAQs

How do I find which model version produced a prediction?

Register your models to Unity Catalog. Each model version links back to the training run, dataset, environment, and Git commit that created it, and the lineage is shown in the Lineage tab on the model version page in Catalog Explorer.

How do I tie a model back to its training data?

Use mlflow.log_input() (MLflow 2.11.0 and above) to log the Unity Catalog table and version used for training. Models logged with the Feature Store APIs capture feature lineage automatically.

How do I recover the exact code used to train a model?

Version notebooks and code with Databricks Git folders. Because each model version links back to the Git commit that produced it, you can retrieve the precise code used for that training run.

Can I analyze runs across many experiments at once?

Yes. MLflow System Tables expose run tracking data in the system.mlflow.* tables in Unity Catalog, so you can query experiment data at scale across your workspaces.