How do I version models, code, features, and training data together?

Summary

• MLflow's LoggedModel is a unified version hub that links your trained model, application code, configuration, and training metadata under a single immutable identifier, enabling reproducibility and comparison across versions.
• Version code by Git commit hash, stored as metadata in the LoggedModel, so you can trace any deployed model back to its exact source code and configuration state.
• Use mlflow.set_active_model() for simple versioning that automatically creates a LoggedModel and links all subsequent traces, or mlflow.create_external_model() for full control over parameters and tags.
• Training data provenance is recorded in MLflow Tracking, which captures datasets, parameters, and hyperparameters alongside metrics and artifacts for each experiment run.
• In production, use environment variables (MLFLOW_ACTIVE_MODEL_ID) to link live traces to your deployed application version without code changes, enabling comparison and debugging across versions.

How do I version models, code, features, and training data together?
Version your machine learning artifacts together using MLflow's unified metadata system: the LoggedModel captures a specific version of your application, storing references to code (as a Git commit hash), configuration parameters, feature definitions, and links to training runs that produced the model. This centralized versioning approach eliminates the manual tracking of separate code, model, feature, and data versions by tying them all to a single, immutable identifier that you can reference, compare, and redeploy.

What is a LoggedModel?

A LoggedModel represents a snapshot of your application at a specific point in time, functioning as a metadata hub that connects all components of your ML system. It stores not the application code itself (which lives in Git) but rather a pointer to where that code is managed, typically as a Git commit hash. Each LoggedModel has a unique ID (starting with m-) that serves as a stable reference for deployment, evaluation, and debugging across your ML lifecycle.

Key components to version together

Component	How it's versioned	Why it matters
Application code	Git commit hash stored in LoggedModel tags	Enables reproducibility and lets you trace any deployed model back to its exact source
Model artifact	MLflow Model Registry or Tracking	The serialized model file (PyTorch, Sklearn, or GenAI agent code) with transformers and preprocessing
Features	Stored as metadata/tags in LoggedModel parameters	Feature definitions, transformations, and selection logic tied to your model version
Training data	Linked via MLflow Tracking experiment runs	Dataset paths, versioning info, and data split definitions recorded in experiment metadata
Hyperparameters	Logged in MLflow Tracking during training	All configuration (learning rate, batch size, model size) tied to the experiment that produced this model
Dependencies	Recorded in model environment (conda, Python versions)	Ensures your model runs with the same library versions used during training

How Databricks approaches version control with MLflow

Use MLflow to create unified versions in two ways:

Simple versioning with mlflow.set_active_model()

This method creates a LoggedModel automatically and links all subsequent traces:

import mlflow

# Creates a new LoggedModel if needed
mlflow.set_active_model(name="customer-churn-predictor-v1.2")

# All traces (logs, metrics, etc.) within this context are linked to this version
with mlflow.set_active_model(name="customer-churn-predictor-v1.3") as model:
    print(f"Model ID: {model.model_id}")

Full control with mlflow.create_external_model()

For comprehensive versioning, use this method to attach Git commit hashes, parameters, and tags:

import mlflow
import subprocess

# Get current Git commit for code versioning
git_commit = subprocess.check_output(
    ["git", "rev-parse", "HEAD"]
).decode().strip()[:8]

model = mlflow.create_external_model(
    name="recommendation-engine-v2.1",
    model_type="genai-agent",
    params={
        "llm_model": "gpt-4",
        "temperature": "0.7",
        "retrieval_k": "5"
    },
    tags={
        "git_commit": git_commit,
        "training_dataset": "user-behavior-q3-2026",
        "feature_version": "v2.1",
        "hyperparameters_source": "hyperopt-run-42"
    }
)

End-to-end versioning workflow

1. Train and log: Use MLflow Tracking to log your training run, which captures training data source, hyperparameters, metrics, and the trained model artifact.
2. Create a LoggedModel: Link your trained model to your Git commit and feature definitions by creating a LoggedModel with the Git hash and feature metadata.
3. Tag metadata: Attach training data version, feature engineering version, and hyperparameter source as tags in the LoggedModel.
4. Deploy: Register the LoggedModel version in Model Serving and reference it by its unique ID.
5. Compare and debug: Use MLflow's evaluation tools to compare metrics, costs, and latency across versions. Reference the Git commit to inspect code differences.

Production deployment

For production environments, set the MLFLOW_ACTIVE_MODEL_ID environment variable instead of embedding version tracking in code:

export MLFLOW_ACTIVE_MODEL_ID="m-9dcd3a2b8e764e5bb4e8a89b4e4c6f21"

This approach allows you to change which model version is active by updating the environment variable without redeploying code, and all inference traces are automatically linked to the correct application version.

Use cases

Data teams use unified versioning when A/B testing different feature sets or model architectures: by versioning features, hyperparameters, and code together, you can objectively measure which combination performs best and easily roll back to a prior version. ML engineering teams use it for compliance and audit trails, tying deployed models to their exact training data, code version, and configuration. GenAI application teams use unified versioning to compare prompt versions, retrieval strategies, and LLM model choices side-by-side before promoting a new version to production.

FAQs

Do I need to store my code in the LoggedModel itself?

No. The LoggedModel stores a reference to your code (typically a Git commit hash), not the code itself. This keeps the metadata lean and ensures the single source of truth remains your Git repository.

How do I link training data to my LoggedModel?

During training, log the dataset source and version as a tag or parameter in your MLflow Tracking experiment run. When you create the LoggedModel, reference that experiment run ID or include the training data identifier as a tag in the LoggedModel.

Can I change the active version in production without redeploying?

Yes. Instead of calling mlflow.set_active_model() in code, use the MLFLOW_ACTIVE_MODEL_ID environment variable. You can change which model version is active by updating the environment variable, and all traces are automatically linked to the new version.

What if I have multiple models or feature sets to version?

Create a separate LoggedModel for each combination. Tag each one with its feature set version, training data version, and code version. MLflow's evaluation tools let you compare all combinations to find the best one.

How does versioning help with compliance and debugging?

Each LoggedModel is immutable and links to specific code, training data, hyperparameters, and results. If a model behaves unexpectedly in production, you can trace it back to its exact training configuration, feature definitions, and code version for root-cause analysis.