E2E MLOps Part 1: How to build and govern models with AutoML, MLflow, and Unity Catalog

[GabFernandes]

In enterprise environments, the bottleneck of Machine Learning is rarely the algorithms themselves—it is the engineering around them. Data scientists often spend days configuring environments, managing infrastructure, and manually tracking model iterations.

To bridge this gap, Databricks offers a powerful tandem: AutoML to rapidly accelerate model exploration, and MLflow integrated with Unity Catalog to enforce robust corporate governance, lineage, and lifecycle tracking.

This 3-part series will guide you through building a production-ready, End-to-End MLOps pipeline on Databricks.

• Part 1: Rapid baseline generation with AutoML, MLflow tracking, and model registration in Unity Catalog.

• Part 2: Diving deeper—Feature Stores, dataset lineage, and customizing AutoML-generated trial notebooks.

• Part 3: Model Serving—Deploying real-time inference endpoints and setting up Lakehouse monitoring for data drift.

Let's dive into Part 1.

The End-to-End MLOps Architecture

Before writing code, let’s understand the data flow we are building today:

[ Delta Lake Table (Unity Catalog) ] 
               │
               ▼
   [ Databricks AutoML Run ]  ──(Automatic Tracking)──► [ MLflow Experiments ]
               │                                                 │
               ▼                                                 ▼
[ Selected Champion Model ] ──────────────────────────► [ Unity Catalog Model Registry ]

By keeping all assets under the Unity Catalog umbrella, we secure full lineage from the raw data features up to the production-registered model binary.

Step-by-Step Implementation

Step 1: Initiating AutoML via Python API

While Databricks AutoML provides a great UI, executing it programmatically via the Python API allows you to integrate training into your automated CI/CD pipelines or nightly data orchestration jobs.

For this guide, we are assuming a classic Customer Churn dataset registered in Unity Catalog.

Python

from databricks import automl

# Define the source table in Unity Catalog
dataset_path = "main.gold_analytics.customer_churn_features"

# Execute AutoML for a Binary Classification task
summary = automl.classify(
    dataset=spark.table(dataset_path),
    target_col="churn_flag",
    primary_metric="f1",
    timeout_minutes=15  # Constraining execution time for demo purposes
)

# Retrieve the Champion Run ID
best_run_id = summary.best_trial.mlflow_run_id
print(f"Champion Run ID: {best_run_id}")

Step 2: Unlocking the MLflow Autologging Metadata

Databricks AutoML automatically configures MLflow Autologging. This means that during the 15-minute search, every single hyperparameter, loss curve, ROC curve, and confusion matrix was captured.

We can query the MLflow Client API to programmatically inspect the winning model's parameters and performance metrics:

Python

import mlflow

# Fetch run metadata from MLflow
client = mlflow.tracking.MlflowClient()
run = client.get_run(best_run_id)

metrics = run.data.metrics
params = run.data.params

print(f"Winning Algorithm: {params.get('classifier')}")
print(f"Validation F1-Score: {metrics.get('val_f1_score'):.4f}")
print(f"Validation Accuracy: {metrics.get('val_accuracy_score'):.4f}")

Step 3: Registering the Champion Model in Unity Catalog

Now that we have identified our champion, we need to register it. In modern Databricks architectures, we avoid the legacy workspace model registry and use Unity Catalog Model Registry for three-tier namespace support (catalog.schema.model).

This ensures our model inherits the same robust security, access controls, and tags as any standard Delta table.

Python

# Path to the model artifact within MLflow
model_uri = f"runs:/{best_run_id}/model"

# Secure three-level namespace destination
registered_model_name = "main.gold_analytics.customer_churn_predictor"

# Register the model to Unity Catalog
model_details = mlflow.register_model(
    model_uri=model_uri, 
    name=registered_model_name
)

print(f"Model successfully registered in Unity Catalog!")
print(f"Active Version: {model_details.version}")

Key Takeaways for Enterprise Architectures

1. Zero Black Boxes: Unlike traditional black-box AutoML tools, Databricks AutoML generates fully documented source-code notebooks for every single trial. If your team wants to tune the champion model further, you can open the notebook, edit the code, and log it manually.

2. Unified Governance: Registering models directly to Unity Catalog bridges the gap between Data Engineering and Data Science. Your ML models now respect the same catalog boundaries and governance policies as your production databases.

Next Up in Part 2

In the next article, we will look at how to construct a Databricks Feature Store to feed this pipeline, prevent training-serving skew, and deep-dive into customizing the PySpark code generated by the AutoML trial notebooks.

What are your thoughts? Do you trigger your AutoML runs programmatically or do you prefer the UI approach for quick exploration? Let's discuss in the comments below!