How to version ML models with MLflow on Databricks: step by step

This tutorial shows how to register, version and load an ML model using MLflow on Databricks, making it easier to promote models from development to production. The ability to version models in Databricks with MLflow helps track experiments, compare metrics and manage deployments in a controlled way — essential when there are multiple iterations of the same model (for example, 5–20 runs per experiment) and when you want to reproduce results in production.

Prerequisites

• Databricks account with a Workspace configured and permissions to create models/experiments. Ideally, a user with Developer or Admin permissions.
• Running Databricks cluster with a runtime that includes MLflow (e.g.: Databricks Runtime ML). A cluster with 4 cores and 16 GB of RAM is sufficient for this example.
• Basic knowledge of Python and scikit-learn or similar. Basic command-line and Databricks UI knowledge helps with navigation.

Step 1: Create an MLflow experiment

An experiment groups runs. Creating an experiment allows you to organize and compare models by objective (e.g.: "classificacao_iris_v1"). In real environments, create one experiment per project or per pipeline (for example, "fraude_cartoes_production").

import mlflow
mlflow.set_experiment("exemplo_classificacao_iris")

Why: with an experiment you can easily filter runs by date, author or hyperparameters. E.g.: if you produce 10 runs with n_estimators = [10, 50, 100], you can quickly compare aggregated metrics (mean/std) between groups.

Step 2: Train and register a model with a run

Inside a run, we log metrics, parameters and produce the model artifact. Here we use scikit-learn for a simple example with the Iris dataset — useful to demonstrate the full flow before applying it to real data.

from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import mlflow.sklearn

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

with mlflow.start_run(run_name="rf_iris_v1") as run:
    params = {"n_estimators": 50, "random_state": 42}
    mlflow.log_params(params)

    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    preds = model.predict(X_test)
    acc = accuracy_score(y_test, preds)
    mlflow.log_metric("accuracy", acc)

    # Registar o modelo no MLflow Model Registry (name define o modelo versionado)
    mlflow.sklearn.log_model(model, artifact_path="model", registered_model_name="IrisRF")

    print(f"Run ID: {run.info.run_id}, accuracy: {acc}")

Practical details: the artifact is stored in the workspace storage (DBFS or OneLake depending on configuration). The registered_model_name field creates a unique entry in the Model Registry called "IrisRF" and automatically generates a new version (e.g.: v1, v2) for each log_model with the same name.

Step 3: View versions in the Model Registry

After registering, the model appears in the Databricks Model Registry. The UI shows versions, stages (Staging, Production, Archived) and metadata such as creation date, author and notes. Using the API allows you to automate checks and generate reports (for example, list all versions created in the last 30 days).

from mlflow.tracking import MlflowClient
client = MlflowClient()

# List versions of the registered model
for v in client.get_latest_versions("IrisRF"):
    print(v.version, v.current_stage, v.status)

# Get information about a specific version
ver = client.get_model_version("IrisRF", "1")
print(ver.version, ver.current_stage, ver.source)

Example expected output: "1 Staging READY" or "1 None READY". If you see states like PENDING_REGISTRATION or FAILED, check the run logs for errors in the artifact upload.

Step 4: Promote a version to Staging/Production

Promoting versions enables validation and deployment workflows. Common practice is to promote first to Staging to run A/B tests or batch validation, and only move to Production once validated. Use manual approval policies or CI/CD pipelines to control who can promote.

# Promote version 1 to Staging
client.transition_model_version_stage(name="IrisRF", version="1", stage="Staging")

# After tests, promote to Production
# client.transition_model_version_stage(name="IrisRF", version="1", stage="Production")

Best practices: include a tag or description with the code commit and the dataset hash used (e.g.: git_sha: abc123, data_snapshot: 2026-10-01) for full traceability.

Step 5: Load a specific version for prediction

To perform inference, load the desired version from the Model Registry. You can use the URI 'models:/<name>/<stage>' to always point to the version in Production or Staging, or 'models:/<name>/<version>' to bind to an immutable version (e.g.: v1).

import mlflow
# Load the version in Staging
model = mlflow.pyfunc.load_model("models:/IrisRF/Staging")

# Make an example prediction
import numpy as np
sample = X_test[0:2]
pred = model.predict(sample)
print(pred)

Note that loading by stage is useful in automated deployment (inference service points to "models:/IrisRF/Production"). If you need full reproducibility, load by version.

Verify the result

Confirm in the Databricks UI: in the Models section you should see 'IrisRF' with versions and stages. In the notebook console you will see the Run ID, metrics (accuracy) and the example prediction. Also check the artifact size (for example 2–10 MB for simple models) and metadata. If the prediction returns valid classes and the version appears in the Registry, the process worked.

Conclusion

You now know how to register, version, promote and load models with MLflow on Databricks, which facilitates traceability and deployment. Next practical steps: automate evaluations via CI/CD, add unit tests for predictions and integrate A/B testing between versions. Final tip: keep clear naming conventions (e.g.: projeto_modelo_versao) and always record parameters, metrics and data hashes to avoid regressions over time.