ML Pipeline on Databricks with MLflow

Train, tune, and register a wine quality model using MLflow and Unity Catalog.

Introduction

30 Second Summary

Studying for a certification by reading docs is like learning to cook by memorizing recipes. The concepts only stick when you actually build something with them.

In this project, you will build a complete wine quality prediction pipeline on Databricks that covers the core MLOps workflow tested in the Databricks ML Associate certification. You will train models, tune hyperparameters, track experiments, and register a production-ready model.

What You'll Build

When you finish, you will have a registered model in Unity Catalog with a "Champion" alias that you can load with a single line of code to run predictions on new data.

By the end of this project, you'll have:

  • A full MLflow experiment dashboard where you can compare training runs side-by-side and see exactly which hyperparameters produced the best metrics.
  • A Hyperopt-tuned model that visibly outperforms your baseline, proving that Bayesian search finds better parameters in fewer iterations than manual guessing.
  • A registered model in Unity Catalog with Champion and Challenger aliases, showing full lineage from raw data to production-ready artifact.
  • Secret Mission: Integrate the Feature Store to create a reusable feature table and train a model using FeatureLookup for production-grade feature management.

Are there any prerequisites?

You should be comfortable with basic Python and foundational ML concepts like train/test splits and evaluation metrics. No Databricks account is needed beforehand. The project walks you through creating a free account in Step 1.

Before We Start

Before we dive into Databricks, let's lock in what you're about to build and why it matters. This step is your chance to commit to the project goals before any hands-on work begins.

Set Up Your Databricks Workspace

The Databricks ML Associate certification tests your ability to use specific APIs and tools hands-on. Studying documentation alone won't build the muscle memory you need for the exam's 48 questions in 90 minutes.

In this step, you'll create a free Databricks workspace, configure a notebook with pre-installed ML libraries, and upload the dataset you'll use throughout the project.

In this step, get ready to:
  • Create a Databricks Free Edition account.
  • Configure a notebook with Serverless compute and the ML base environment.
  • Upload the Wine Quality dataset to a Unity Catalog volume.
Create your Databricks Free Edition account

Databricks Free Edition is a permanent, no-cost workspace designed for students and hobbyists. It uses serverless compute with quota limits, but provides the same Unity Catalog and MLflow APIs that the certification tests.

  • Open the Databricks Free Edition signup page in your browser.
  • Sign up with your personal email address.
  • Choose your preferred signup method (email OTP, Google, or Microsoft).
  • Wait for Databricks to create your workspace automatically.

Once your workspace loads, you should see the welcome page with navigation options in the left sidebar, including links to create a new notebook or open the SQL editor.

What is Databricks Free Edition?

Free Edition is a serverless-only, quota-limited environment with no time limit. Unlike a 14-day free trial, it never expires. It includes Unity Catalog enabled by default, giving you the same model registry and data governance tools that production teams use.

Workspace not loading?

Workspace provisioning can take 1-2 minutes. If you see a loading spinner, wait for it to complete.

If you receive an error about your email domain, make sure you're using a personal email address (Gmail, Outlook, etc.) rather than a corporate or university email that may already be linked to an existing Databricks account.

Help me troubleshoot my Databricks Free Edition signup

Configure your ML notebook environment

Your workspace needs a notebook connected to compute that has ML libraries pre-installed. The ML base environment includes scikit-learn, MLflow, pandas, and matplotlib out of the box.

  • In your workspace, click New in the left sidebar.
  • Select Notebook.
  • In the notebook, click the compute dropdown at the top of the page.
  • Select Serverless as your compute type.
  • Open the Environment side pane by clicking the gear icon.
  • Under the Base environment dropdown, select ML.
  • Click Apply to confirm the environment selection.

You should see ML displayed as the active base environment in the side pane. Your notebook is now ready with all the ML libraries pre-installed.

What does the ML base environment include?

The ML environment (serverless environment version 5) comes with scikit-learn 1.6.1, mlflow-skinny 3.8.1, pandas 2.2.3, matplotlib 3.10.0, and databricks-feature-engineering 0.13.0.1 pre-installed.

This means you won't need to install most packages manually. The one exception is Hyperopt, which you'll install later with %pip install.

Don't see the ML option?

Make sure you selected Serverless compute first. The base environment options (Standard, ML, AI) only appear when Serverless is selected as the compute type.

Help me configure the ML environment in my Databricks notebook

Upload the Wine Quality dataset

Free Edition restricts outbound internet access from your workspace. This means you cannot download files directly using Python code in your notebook. Instead, you'll download the CSV to your local machine first and then upload it through the workspace UI.

  • On your local machine, open this link in a new browser tab: Wine Quality CSV.
  • Save the file to your computer (your browser may download it automatically, or you can right-click and choose Save As).
  • Back in your Databricks workspace, click New in the left sidebar.
  • Select Add or upload data.
  • Click Upload files to a volume.
  • Select the workspace catalog and default schema.
  • Upload the winequality-red.csv file you downloaded.
  • Note the volume path displayed after the upload completes. It follows the format /Volumes/workspace/default/<volume_name>/winequality-red.csv.

After the upload finishes, you should see winequality-red.csv listed in the volume's files. Copy or write down the full volume path. You'll use it in the next step to load the data into a Spark DataFrame.

Why upload manually instead of using code?

Free Edition restricts outbound internet access from your workspace notebooks. Commands like urllib.request.urlretrieve() or wget would fail with a connection error. The workspace UI upload is the supported path for Free Edition users.

Upload failing or volume not appearing?

Make sure you selected the workspace catalog and default schema. Free Edition users have access to these by default without additional configuration.

If the file fails to upload, check that it's the original CSV (about 100 KB). Some browsers may have renamed it or added a .txt extension during download.

Help me upload my Wine Quality CSV to a Databricks volume

Your Databricks workspace is ready with a configured ML notebook and the Wine Quality dataset loaded into Unity Catalog. Next up, you'll load this data into a Spark DataFrame and explore it to understand what you're working with.

Explore Data with Spark DataFrames

Your Databricks workspace is ready and the Wine Quality dataset is uploaded. Before training any model, you need to understand what the data looks like.

Exploring data first reveals patterns, distributions, and quality issues that directly impact model performance. The Databricks ML Associate certification tests your ability to compute summary statistics, handle outliers, and visualize features using Spark. This step builds that muscle memory.

In this step, get ready to:
  • Load the Wine Quality CSV into a Spark DataFrame and compute summary statistics.
  • Identify and remove outliers using the IQR method.
  • Save the cleaned data as a Delta table in Unity Catalog for downstream steps.
Load data and compute summary statistics

The Wine Quality CSV uses semicolons as separators, not commas. You need to tell Spark about this when loading the file.

  • In your notebook, add a new cell and load the CSV into a Spark DataFrame by pasting this code:
# Load the Wine Quality CSV from your Unity Catalog volume
df = spark.read.csv(
    "/Volumes/workspace/default/[[VOLUME_NAME=\"your-volume-name\"]]/winequality-red.csv",
    header=True,
    inferSchema=True,
    sep=";"
)

# Compute summary statistics for all columns
display(df.summary())

What does this code do?

The sep=";" parameter tells Spark that fields are separated by semicolons instead of commas. The Wine Quality dataset from UCI uses this format.

The .summary() method computes count, mean, standard deviation, min, and max for every column. This gives you a quick snapshot of the entire dataset.

  • Run the cell by pressing Shift+Enter.

You should see a table with rows for count, mean, stddev, min, and max across all 12 columns. The quality column is your target variable, with scores ranging from 3 to 8.

Seeing a single column with garbled data?

If all data appears in one column, the separator was not applied. Double-check that you included sep=";" in the spark.read.csv() call. Without it, Spark defaults to comma separation and treats each row as a single string.

Also confirm the volume path matches the one you noted after uploading. Help me fix my CSV loading issue in Databricks.

Now visualize how wine quality scores are distributed across the dataset.

  • Add a new cell and display the quality distribution by pasting this code:
# Group by quality score and count occurrences
quality_counts = df.groupBy("quality").count().orderBy("quality")

# Display the result as a table
display(quality_counts)
  • Run the cell with Shift+Enter.
  • Click the + button above the results table, then select Visualization.
  • In the Visualization Type dropdown, select Bar.
  • Set the X column to quality and the Y column to count.
  • Click Save.

You should see a bar chart showing that most wines cluster around quality scores 5 and 6, with very few at the extremes (3 or 8).

Why does the quality distribution matter?

An imbalanced target variable means your model will see far more examples of "average" wines than excellent or poor ones. This tells you upfront that the model may struggle to predict extreme scores accurately.

Handle data quality issues

Before training, you need to check for missing values and remove extreme outliers that would distort the model.

  • Add a new cell and check for null values across all columns by pasting this code:
from pyspark.sql.functions import col

# Check for missing values in each column
for column in df.columns:
    null_count = df.where(col(column).isNull()).count()
    print(f"{column}: {null_count} nulls")
  • Run the cell with Shift+Enter.

You should see 0 nulls for every column. The Wine Quality dataset is complete, so no imputation is needed.

Why check for nulls anyway?

In production pipelines, you always validate assumptions about data completeness before training. The certification expects you to know how to check, even when the dataset happens to be clean.

Next, remove extreme outliers using the Interquartile Range (IQR) method. Outliers in volatile_acidity and residual_sugar can skew model predictions.

  • Add a new cell and remove outliers beyond 3x IQR by pasting this code:
# Count rows before cleaning
rows_before = df.count()
print(f"Rows before cleaning: {rows_before}")

# Calculate IQR for volatile_acidity
quantiles_va = df.approxQuantile("volatile_acidity", [0.25, 0.75], 0.01)
q1_va, q3_va = quantiles_va[0], quantiles_va[1]
iqr_va = q3_va - q1_va

# Calculate IQR for residual_sugar
quantiles_rs = df.approxQuantile("residual_sugar", [0.25, 0.75], 0.01)
q1_rs, q3_rs = quantiles_rs[0], quantiles_rs[1]
iqr_rs = q3_rs - q1_rs

# Remove outliers beyond 3x IQR for both columns
df_cleaned = df.filter(
    (col("volatile_acidity") >= q1_va - 3 * iqr_va) &
    (col("volatile_acidity") <= q3_va + 3 * iqr_va) &
    (col("residual_sugar") >= q1_rs - 3 * iqr_rs) &
    (col("residual_sugar") <= q3_rs + 3 * iqr_rs)
)

# Count rows after cleaning
rows_after = df_cleaned.count()
print(f"Rows after cleaning: {rows_after}")
print(f"Rows removed: {rows_before - rows_after}")

What does this code do?

The IQR is the range between the 25th and 75th percentiles. Any value more than 3 times the IQR above Q3 or below Q1 is considered an extreme outlier.

The approxQuantile method computes approximate percentiles efficiently on distributed data. The third argument (0.01) controls precision.

Using 3x IQR (instead of the standard 1.5x) means you only remove the most extreme values, preserving as much data as possible for training.

  • Run the cell with Shift+Enter.

You should see the original row count (1599), the cleaned row count, and how many rows were removed. Expect a small number of removals (typically fewer than 30 rows).

Seeing zero rows removed?

If no rows were removed, the 3x IQR threshold may not catch any outliers in this dataset (which is possible). Verify by printing q1_va, q3_va, and iqr_va to confirm the quantiles computed correctly. If the values look reasonable, your data simply has no extreme outliers at this threshold.

Help me verify my IQR outlier removal is working correctly.

Save cleaned data as a Delta table

Saving the cleaned data as a Delta table in Unity Catalog means every subsequent step can read from a single, versioned source of truth instead of re-running the cleaning logic.

  • Add a new cell and save the cleaned DataFrame by pasting this code:
# Save cleaned DataFrame as a Delta table in Unity Catalog
table_name = "workspace.default.wine_quality_cleaned"
df_cleaned.write.format("delta").mode("overwrite").saveAsTable(table_name)

print(f"Cleaned data saved to {table_name}")

What does this code do?

The three-level name workspace.default.wine_quality_cleaned follows the Unity Catalog naming convention: catalog, schema, then table name.

Delta format gives you ACID transactions, time travel, and schema enforcement. These are features you will rely on in production ML pipelines.

  • Run the cell with Shift+Enter.

You should see the confirmation message: Cleaned data saved to workspace.default.wine_quality_cleaned.

Before you move on, predict what you will find: do you think the table appears in the Catalog Explorer?

  • In the left sidebar, click Catalog.
  • Expand workspace > default.
  • Confirm that wine_quality_cleaned appears in the list of tables.

You should see your new table listed under the default schema. This confirms the cleaned data is persisted and ready for model training.

Table not appearing in Catalog Explorer?

Make sure you expanded the correct path: workspace catalog, then default schema. If the table still does not appear, click the refresh icon at the top of the Catalog pane.

If you see an error about permissions or catalog not found, confirm your notebook is using Serverless compute. Free Edition workspaces only support the workspace catalog on Serverless.

Help me troubleshoot my Delta table not showing in Unity Catalog.

✔️ Awesome, I've got everything!

Great. Make sure your notebook is saved before moving on.

ⓧ I'd like to double check the full code

Your notebook should contain five cells in this order. Scroll up and confirm each is present:

  • Cell 1: Load CSV with spark.read.csv() using sep=";", then display(df.summary()).
  • Cell 2: groupBy("quality").count().orderBy("quality") with display() and a bar chart visualization saved.
  • Cell 3: Null check loop using df.where(col(column).isNull()).count().
  • Cell 4: IQR outlier removal on volatile_acidity and residual_sugar with row count logging.
  • Cell 5: saveAsTable("workspace.default.wine_quality_cleaned") with Delta format.

Your data is explored, cleaned, and persisted. Next up, you will train a baseline model with MLflow tracking to see how a deliberately simple model performs on this data.

Train a Baseline Model with MLflow Tracking

Your cleaned wine data is ready in Unity Catalog. Now it is time to train a model and see how well it predicts wine quality.

A baseline model gives you a reference point to measure all future improvements against. In this step, you will train a deliberately weak model and manually log every parameter and metric to MLflow. This teaches you the exact API calls that the Databricks ML Associate certification tests.

In this step, get ready to:
  • Load cleaned data from the Delta table and split it into training and test sets.
  • Train a baseline RandomForestRegressor with MLflow tracking.
  • Create a visualization artifact and log it to your MLflow run.
Load cleaned data and prepare train/test split

Before training, you need to load the cleaned wine data from your Delta table and split it into two sets: one for training the model, one for evaluating it.

  • Create a new cell in your notebook by hovering below your last cell and clicking + Code.
  • Paste the following code to load and split the data:
# Load cleaned wine data from Unity Catalog Delta table
df = spark.table("workspace.default.wine_quality_cleaned").toPandas()

# Separate features (X) from target variable (y)
X = df.drop("quality", axis=1)
y = df["quality"]

# Split into training and test sets (80/20)
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

print(f"Training set: {X_train.shape[0]} rows")
print(f"Test set: {X_test.shape[0]} rows")

What does this code do?

spark.table() reads your cleaned Delta table directly from Unity Catalog using the three-level name format: catalog.schema.table.

The .toPandas() call converts the Spark DataFrame to a pandas DataFrame so scikit-learn can work with it. The 80/20 split keeps 80% of rows for training and reserves 20% for honest evaluation.

  • Run the cell by pressing Shift+Enter.
  • You should see output like Training set: ~1050 rows and Test set: ~260 rows (exact numbers depend on how many outliers were removed).

Seeing a table not found error?

Make sure the table name matches exactly what you saved in the previous step. Confirm it exists by running spark.sql("SHOW TABLES IN workspace.default").show() in a new cell.

help me fix a table not found error when loading my Delta table

Train a baseline model inside an MLflow run

The certification tests your ability to manually log parameters, metrics, and models using the MLflow API. You will now train a RandomForestRegressor with intentionally weak parameters to establish a baseline.

Before you run this next cell, think about what R-squared score you expect from a random forest with only 10 trees and a maximum depth of 3. Would you expect it to perform well on a regression task with 11 features?

  • Create a new cell and paste the following training code:
import mlflow
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import numpy as np

# Train baseline model with deliberately weak parameters
with mlflow.start_run(run_name="baseline_rf") as run:
    n_estimators = 10
    max_depth = 3

    # Log parameters manually
    mlflow.log_param("n_estimators", n_estimators)
    mlflow.log_param("max_depth", max_depth)

    # Train the model
    rf = RandomForestRegressor(
        n_estimators=n_estimators,
        max_depth=max_depth,
        random_state=42
    )
    rf.fit(X_train, y_train)

    # Compute metrics on test set
    predictions = rf.predict(X_test)
    rmse = np.sqrt(mean_squared_error(y_test, predictions))
    r2 = r2_score(y_test, predictions)

    # Log metrics and model
    mlflow.log_metric("rmse", rmse)
    mlflow.log_metric("r2_score", r2)
    mlflow.sklearn.log_model(rf, "model")

What does this code do?

mlflow.start_run() creates a new tracked run inside the notebook experiment. Everything inside the with block is recorded as part of that run.

  • mlflow.log_param() records the hyperparameters you chose so you can compare them later.
  • mlflow.log_metric() records the model's performance scores (RMSE and R-squared).
  • mlflow.sklearn.log_model() saves the trained model as an artifact you can reload or deploy later.

The parameters are deliberately weak. Only 10 trees with a max depth of 3 cannot capture the complexity of 11 wine features.

  • Run the cell by pressing Shift+Enter.

Check the output. You should see an R-squared around 0.35 to 0.45. That means the model explains less than half the variance in wine quality. This is the mediocre baseline you will dramatically improve with hyperparameter tuning in the next step.

Not seeing any output?

Add print(f"RMSE: {rmse:.4f}") and print(f"R-squared: {r2:.4f}") at the end of the with block (inside it, after the log_model line) to display the values explicitly.

help me debug my MLflow training run that isn't showing output

Explore the MLflow UI to verify your run

Databricks automatically creates an MLflow experiment for your notebook. Every run you log is captured there with full details.

  • In the left sidebar of your workspace, click Experiments.
  • Find your experiment in the list. It has the same name as your notebook.
  • Click the experiment name to open it.
  • Click the run named baseline_rf to open its details.

Inside the run details, confirm you can see:

  • Parameters: n_estimators = 10, max_depth = 3.
  • Metrics: rmse and r2_score (R-squared should be around 0.35 to 0.45).
  • Artifacts: a model folder containing the serialized scikit-learn model.

Why is the R-squared so low?

An R-squared of 0.35 to 0.45 means the model only explains about a third of the variance in wine quality. With just 10 shallow trees, the model underfits the data. It cannot learn the relationships between the 11 chemical features and the quality score.

This gap is intentional. In the next step, you will use Hyperopt to find much better parameters and close this performance gap.

Create a visualization artifact

A scatter plot of actual vs predicted values is one of the most useful ML artifacts. It shows at a glance where your model struggles.

  • Create a new cell and paste the following code to generate and log the plot:
import matplotlib.pyplot as plt

# Create actual vs predicted scatter plot
fig, ax = plt.subplots(figsize=(8, 6))
ax.scatter(y_test, predictions, alpha=0.5, edgecolors="k", linewidth=0.5)
ax.plot(
    [y_test.min(), y_test.max()],
    [y_test.min(), y_test.max()],
    "r--", lw=2
)
ax.set_xlabel("Actual Quality")
ax.set_ylabel("Predicted Quality")
ax.set_title("Baseline Model: Actual vs Predicted")

# Save plot as PNG file
fig.savefig("/tmp/actual_vs_predicted.png", dpi=100, bbox_inches="tight")
plt.show()

# Log the plot to the existing MLflow run
with mlflow.start_run(run_id=run.info.run_id):
    mlflow.log_artifact("/tmp/actual_vs_predicted.png")

print("Plot logged to MLflow artifacts.")

What does this code do?

The scatter plot puts actual quality scores on the x-axis and predicted scores on the y-axis. Points close to the red dashed line are accurate predictions. Points far from the line are errors.

mlflow.start_run(run_id=run.info.run_id) reopens your existing baseline run so the artifact is logged to the same place as your parameters and metrics. The plot file is saved to /tmp/ first, then uploaded to MLflow as an artifact.

  • Run the cell. You should see the scatter plot rendered inline in your notebook.
  • Notice how the predictions cluster in a narrow band. The model cannot distinguish between wines rated 5, 6, or 7 very well.

Now verify the artifact was logged to your MLflow run.

  • Return to the Experiments page in the left sidebar.
  • Click your experiment name, then click the baseline_rf run.
  • Click the Artifacts tab. You should see actual_vs_predicted.png listed alongside the model folder.

Plot not appearing in artifacts?

Make sure run.info.run_id refers to the variable from your training cell. If you restarted the notebook kernel between cells, the run variable may be lost. Re-run the training cell first, then re-run the plot cell.

help me log an artifact to an existing MLflow run

✔️ Awesome, I've got everything!

Great. Confirm all three cells from this step are saved in your notebook.

ⓧ I'd like to double check the full code

Scroll through the cells you added in this step and confirm they appear in this order:

  • Cell 1: Load data from Delta table, split into X/y, train_test_split, print row counts.
  • Cell 2: Import mlflow and sklearn, mlflow.start_run with log_param, train RandomForestRegressor, compute RMSE and R-squared, log_metric, log_model.
  • Cell 3: Import matplotlib, create scatter plot, save to /tmp, log_artifact to the existing run.

Your baseline is established. An R-squared below 0.45 means there is plenty of room to improve. Next up, you will use Hyperopt's Bayesian optimization to find dramatically better hyperparameters in just 20 iterations.

Tune Hyperparameters with Hyperopt

Your baseline model from the previous step logged an R-squared around 0.35 to 0.45. That score is mediocre because you deliberately used weak hyperparameters (n_estimators=10, max_depth=3).

In this step, you will use Hyperopt's fmin function with Bayesian optimization (TPE) to automatically search for better hyperparameters. This is exactly the API the Databricks ML Associate certification tests, and it finds strong parameters in far fewer iterations than a grid search.

In this step, get ready to:
  • Define a Hyperopt search space and objective function with cross-validation.
  • Run fmin with TPE and retrieve the best parameters.
  • Compare the tuned model against the baseline in the MLflow UI.
Define the Hyperopt search space and objective function

Hyperopt is not included in the serverless ML environment by default. You need to install it before importing.

  • In a new cell at the bottom of your notebook, install Hyperopt by running:
%pip install hyperopt

You should see output confirming the package was installed successfully, followed by a message to restart the Python environment.

Seeing a network error?

Databricks Free Edition restricts outbound internet access by default. If %pip install fails with a connection error, you need to verify your identity with LinkedIn first. Go to your workspace settings, find the LinkedIn verification option, and complete it. This unlocks outbound internet access instantly.

help me fix a network error when running %pip install on Databricks Free Edition

Now you can import Hyperopt's core components and define the search space that tells Hyperopt which parameter ranges to explore.

  • Add a new cell with the imports and search space definition by pasting this code:
from hyperopt import fmin, tpe, hp, Trials, STATUS_OK, space_eval
from sklearn.model_selection import cross_val_score
from sklearn.ensemble import RandomForestRegressor
import numpy as np
import mlflow

# Define ranges for Hyperopt to search
search_space = {
    'n_estimators': hp.choice('n_estimators', [50, 100, 200, 300]),
    'max_depth': hp.uniform('max_depth', 5, 30),
    'min_samples_split': hp.uniform('min_samples_split', 2, 20)
}

What does this code do?

  • hp.choice('n_estimators', [50, 100, 200, 300]) picks from a list of discrete values. Hyperopt returns the index of the chosen option, not the value itself.
  • hp.uniform('max_depth', 5, 30) samples any float between 5 and 30. You will cast it to an integer inside the objective function.
  • hp.uniform('min_samples_split', 2, 20) explores the minimum number of samples required to split a node.
  • Run the cell to confirm imports succeed with no errors.

Next, you need an objective function that Hyperopt calls on each trial. This function trains a model, evaluates it with cross-validation, logs to MLflow, and returns the loss.

  • Add a new cell with the objective function by pasting this code:
def objective(params):
    """Train and evaluate a model for one Hyperopt trial."""
    with mlflow.start_run(nested=True):
        # Extract and cast parameters
        n_estimators = params['n_estimators']
        max_depth = int(params['max_depth'])
        min_samples_split = int(params['min_samples_split'])

        rf = RandomForestRegressor(
            n_estimators=n_estimators,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            random_state=42
        )

        # 3-fold cross-validation using negative MSE
        scores = cross_val_score(
            rf, X_train, y_train,
            cv=3, scoring='neg_mean_squared_error'
        )
        rmse = np.sqrt(-scores.mean())

        # Log to MLflow for each trial
        mlflow.log_param("n_estimators", n_estimators)
        mlflow.log_param("max_depth", max_depth)
        mlflow.log_param("min_samples_split", min_samples_split)
        mlflow.log_metric("cv_rmse", rmse)

        return {'loss': rmse, 'status': STATUS_OK}

What does this code do?

mlflow.start_run(nested=True) creates a child run under the parent tuning run, so all 20 trials are grouped together in the MLflow UI.

cross_val_score with cv=3 splits the training data three ways and averages the scores. This gives a more reliable estimate than a single train/test split.

The function returns a dictionary with loss (the RMSE to minimize) and status (tells Hyperopt the trial succeeded).

  • Run the cell. You should see no output, which confirms the function was defined without errors.
Run fmin with TPE and train the final model

With the search space and objective ready, you can now launch the Hyperopt search. The TPE (Tree of Parzen Estimators) algorithm uses past trial results to choose the next set of parameters intelligently, converging faster than random search.

  • Add a new cell and run the fmin search by pasting this code:
# Create a Trials object to record each evaluation
trials = Trials()

# Wrap in a parent run so all trials are grouped
with mlflow.start_run(run_name="hyperopt_tuning"):
    best_params = fmin(
        fn=objective,
        space=search_space,
        algo=tpe.suggest,
        max_evals=20,
        trials=trials
    )

print("Best parameters (raw):", best_params)

What does this code do?

  • Trials() stores the result of every evaluation so you can inspect them later.
  • fmin runs 20 evaluations (max_evals=20), each time picking parameters guided by tpe.suggest.
  • The parent mlflow.start_run groups all 20 nested child runs under one experiment entry named hyperopt_tuning.

You should see a progress bar counting through 20 trials, then a printed dictionary of the best parameters found. This cell may take 1 to 2 minutes to complete.

Seeing a NameError for X_train?

If you see NameError: name 'X_train' is not defined, your notebook session lost the variables from an earlier cell. Re-run the cell that loads the cleaned data from the Delta table and performs the train/test split, then try this cell again.

help me fix NameError for X_train in my Hyperopt objective function

Notice that best_params shows an index for n_estimators (like 2) instead of the actual value (like 200). This is because hp.choice returns the index of the chosen option. You need to map it back to the real value using space_eval.

  • Add a new cell to retrieve the actual best parameter values:
# Map indices back to actual values
best_actual = space_eval(search_space, best_params)

best_n_estimators = best_actual['n_estimators']
best_max_depth = int(best_actual['max_depth'])
best_min_samples_split = int(best_actual['min_samples_split'])

print(f"Best n_estimators: {best_n_estimators}")
print(f"Best max_depth: {best_max_depth}")
print(f"Best min_samples_split: {best_min_samples_split}")

Why use space_eval?

When you use hp.choice(), Hyperopt returns the index of the choice list, not the actual value. space_eval() maps those indices back to real parameter values. This is a Databricks-documented best practice for working with Hyperopt results.

  • Run the cell. You should see the actual best values printed (for example, Best n_estimators: 200).

Now train a final model with these best parameters and log it as a standalone MLflow run so you can easily compare it to the baseline.

  • Add a new cell to train and log the final tuned model:
from sklearn.metrics import mean_squared_error, r2_score

with mlflow.start_run(run_name="best_tuned_model"):
    # Train with the best hyperparameters
    best_rf = RandomForestRegressor(
        n_estimators=best_n_estimators,
        max_depth=best_max_depth,
        min_samples_split=best_min_samples_split,
        random_state=42
    )
    best_rf.fit(X_train, y_train)

    # Evaluate on the held-out test set
    predictions = best_rf.predict(X_test)
    rmse = np.sqrt(mean_squared_error(y_test, predictions))
    r2 = r2_score(y_test, predictions)

    # Log everything to MLflow
    mlflow.log_param("n_estimators", best_n_estimators)
    mlflow.log_param("max_depth", best_max_depth)
    mlflow.log_param("min_samples_split", best_min_samples_split)
    mlflow.log_metric("rmse", rmse)
    mlflow.log_metric("r2_score", r2)
    mlflow.sklearn.log_model(best_rf, "model")

    print(f"Tuned model RMSE: {rmse:.4f}")
    print(f"Tuned model R-squared: {r2:.4f}")

What does this code do?

This trains a fresh RandomForestRegressor with the best parameters Hyperopt found, evaluates it on the test set (which was not used during tuning), and logs the full result to MLflow as a top-level run named best_tuned_model.

The R-squared metric measures how well the model explains variance in wine quality. A score closer to 1.0 is better.

You should see the tuned model's RMSE and R-squared printed. The R-squared should be noticeably higher than the baseline's 0.35 to 0.45 range.

R-squared not much higher than baseline?

If your tuned R-squared is similar to the baseline, check that you used the best_actual values from space_eval and not the raw indices from best_params. Using raw indices (like 2 instead of 200 for n_estimators) would train a very weak model.

help me figure out why my tuned model is not outperforming my baseline

Compare runs in the MLflow UI

Before you check, predict: how much do you think the tuned model's R-squared improved over the baseline's score of roughly 0.35 to 0.45?

  • Click the Experiments icon in the left sidebar of your Databricks workspace.
  • Click on your experiment (it has the same name as your notebook).
  • Select the checkbox next to your baseline run (the one with n_estimators=10) and the best_tuned_model run.
  • Click Compare in the toolbar above the runs table.

The Compare view shows both runs' parameters and metrics side-by-side. You should see the tuned model's R-squared is at least 0.15 points higher than the baseline. The parameter differences explain why: more trees, deeper splits, and a better minimum-samples threshold all contribute to the improvement.

✔️ Awesome, I've got everything!

Great. Make sure you have saved your notebook with all the cells from this step.

ⓧ I'd like to double check the full code

Scroll up through your notebook and confirm these cells are present in order:

  • Cell 1: %pip install hyperopt.
  • Cell 2: Imports (fmin, tpe, hp, Trials, STATUS_OK, space_eval) and search_space dictionary with hp.choice and hp.uniform.
  • Cell 3: objective function with nested=True, 3-fold cross_val_score, and MLflow logging.
  • Cell 4: fmin call with tpe.suggest, max_evals=20, and Trials() wrapped in mlflow.start_run(run_name="hyperopt_tuning").
  • Cell 5: space_eval to retrieve actual best values and print them.
  • Cell 6: Final model trained with best parameters, logged as mlflow.start_run(run_name="best_tuned_model") with RMSE and R-squared printed.

Hyperopt just found significantly better hyperparameters in 20 automated trials. Next up, you will register the best model in Unity Catalog and promote it with Champion/Challenger aliases.

Register and Promote Your Model in Unity Catalog

Your tuned model outperforms the baseline, but a good model sitting in an experiment is not a production model. Production teams need a central registry where they can track which model version is live, who approved it, and what it does.

In this step, you will register both models in Unity Catalog and manage their lifecycle with aliases and tags. This is exactly how the Databricks ML Associate certification tests model deployment (12% of the exam).

In this step, get ready to:
  • Register both models in Unity Catalog as versioned artifacts.
  • Assign Champion and Challenger aliases to control which version serves traffic.
  • Load the Champion model and run batch predictions to prove it works.
Configure Unity Catalog registry and register both models

Unity Catalog uses a three-level namespace (catalog.schema.model_name) to organise registered models. You need to tell MLflow to use Unity Catalog as the registry backend before registering anything.

  • Create a new cell in your notebook and register both models by running this code:
import mlflow

# Point MLflow at Unity Catalog as the model registry
mlflow.set_registry_uri("databricks-uc")

# Three-level name: catalog.schema.model_name
model_name = "workspace.default.wine_quality_model"

# Search for all runs in the current experiment
runs = mlflow.search_runs(order_by=["start_time ASC"])

# Find the baseline run (weak parameters: n_estimators=10)
baseline_run_id = runs[
    runs["params.n_estimators"] == "10"
].iloc[0].run_id

# Register baseline as Version 1
baseline_uri = f"runs:/{baseline_run_id}/model"
result_v1 = mlflow.register_model(baseline_uri, model_name)
print(f"Baseline registered as Version: {result_v1.version}")

What does this code do?

mlflow.set_registry_uri("databricks-uc") switches the registry backend from the legacy Workspace Model Registry to Unity Catalog. Every model you register after this call lands in Unity Catalog's three-level namespace.

mlflow.search_runs() retrieves all runs from the active experiment as a pandas DataFrame. You filter for the baseline run using its distinctive parameter (n_estimators=10).

mlflow.register_model() takes the run's model artifact URI and publishes it to the registry under the three-level name workspace.default.wine_quality_model.

  • Run the cell. You should see output confirming Baseline registered as Version: 1.
  • Now register the tuned model as Version 2 by running this code in a new cell:
# Filter for top-level runs only (exclude nested Hyperopt runs)
top_level_runs = runs[
    runs["tags.mlflow.parentRunId"].isna()
]

# The last top-level run is the final tuned model
tuned_run_id = top_level_runs.iloc[-1].run_id

# Register tuned model as Version 2
tuned_uri = f"runs:/{tuned_run_id}/model"
result_v2 = mlflow.register_model(tuned_uri, model_name)
print(f"Tuned model registered as Version: {result_v2.version}")

Why filter out nested runs?

Hyperopt created 20 nested child runs during tuning. Filtering with tags.mlflow.parentRunId being null isolates only the top-level runs you explicitly created. The last one chronologically is your final best-tuned model.

  • Run the cell. You should see Tuned model registered as Version: 2 in the output.

Seeing a permissions error?

Confirm you are using the workspace catalog and default schema that Free Edition provides. If you changed the model_name variable, make sure it follows the exact three-level format catalog.schema.model_name.

help me fix a permissions error when registering a model in Unity Catalog

Set aliases and tags on your registered model

Aliases are mutable, named references that point to a specific model version. Production teams use them to control which version serves traffic without changing deployment code.

  • Assign aliases and tags to your registered model by running this code in a new cell:
from mlflow import MlflowClient

# Create a client for model management operations
client = MlflowClient()

# Assign aliases: Challenger for baseline, Champion for tuned
client.set_registered_model_alias(model_name, "Challenger", 1)
client.set_registered_model_alias(model_name, "Champion", 2)

# Tag the registered model with its task type
client.set_registered_model_tag(model_name, "task", "regression")

# Tag the Champion version as approved for production
client.set_model_version_tag(model_name, "2", "validation_status", "approved")

print("Aliases and tags applied successfully")

What do aliases replace?

Unity Catalog replaced the old four fixed stages (None, Staging, Production, Archived) with up to 10 custom aliases. This gives teams more flexibility to define their own promotion workflow.

set_registered_model_alias() assigns a name like "Champion" to a specific version number. Later, you can load that model by alias instead of by version number. If you retrain and get a better model, you reassign the alias to the new version without touching any deployment code.

  • Run the cell. You should see Aliases and tags applied successfully in the output.

Getting a model not found error?

Make sure the previous cell completed successfully and the model_name variable matches exactly what you used during registration. Re-run the registration cell if needed.

help me fix a model not found error when setting aliases in Unity Catalog

Load the Champion model and run batch predictions

The real test of a registered model is whether you can load it by alias and run inference. This proves the full loop works, from registration to serving predictions.

  • Load and test your Champion model by running this code in a new cell:
import pandas as pd

# Load the Champion model using its alias
champion_model = mlflow.pyfunc.load_model(
    f"models:/{model_name}@Champion"
)

# Run batch predictions on a sample of test data
sample_data = X_test[:5]
predictions = champion_model.predict(sample_data)

# Display predictions alongside actual values
results = pd.DataFrame({
    "Actual": y_test[:5].values,
    "Predicted": predictions.round(2)
})
print(results)

What does the model URI format mean?

The URI models:/workspace.default.wine_quality_model@Champion tells MLflow to look up the model named wine_quality_model in the workspace catalog and default schema, then resolve the "Champion" alias to a version number and load that version.

This means your deployment code never hard-codes a version number. When you retrain and promote a new Champion, existing code automatically picks up the new version.

  • Run the cell. You should see a table with 5 rows showing actual wine quality scores next to the model's predicted scores.

Predictions look wrong or all identical?

If every prediction is the same value, double-check that Version 2 (the tuned model) has the "Champion" alias, not Version 1. You can verify by running client.get_model_version_by_alias(model_name, "Champion") and checking the returned version number is 2.

help me verify which model version my Champion alias points to

Before you check the Catalog Explorer, predict what you expect to see there. How many model versions should exist, and what alias should each have?

  • Click Catalog in the left sidebar to open Catalog Explorer.
  • In the Catalog Explorer tree on the left, expand the workspace catalog, then expand the default schema.
  • Under Models, click wine_quality_model.

You should see the model details page showing Version 1 with alias "Challenger" and Version 2 with alias "Champion". The registered model tag task = regression is visible on the model page.

Why does this matter for the certification?

The Databricks ML Associate exam dedicates 12% of questions to model deployment. You just used the exact API calls tested: register_model(), set_registered_model_alias(), set_registered_model_tag(), set_model_version_tag(), and load_model() with an alias URI.

✔️ Awesome, I've got everything!

Great. Confirm you have saved your notebook with all the new cells.

ⓧ I'd like to double check the full code

Scroll up through your notebook and confirm these cells are present, in order:

  • Cell 1: mlflow.set_registry_uri("databricks-uc"), model name definition, baseline run search, and mlflow.register_model() for Version 1.
  • Cell 2: Top-level run filter and mlflow.register_model() for the tuned model as Version 2.
  • Cell 3: MlflowClient(), set_registered_model_alias() calls for Challenger and Champion, set_registered_model_tag(), and set_model_version_tag().
  • Cell 4: mlflow.pyfunc.load_model() with the Champion alias, sample predictions, and results DataFrame.

Secret mission

Integrate the Feature Store

The Feature Store is heavily tested in Section 1 of the Databricks ML Associate exam. In this challenge, you'll create a reusable feature table in Unity Catalog, train a model with FeatureLookup, and perform batch scoring with automatic feature lookup. This decouples your features from your training code, exactly how production ML teams operate.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project uses Databricks Free Edition, so there are no ongoing compute costs. Your data and tables remain stored in your workspace at no charge.

Resources you used:

  • Databricks notebook (your ML pipeline code).
  • Unity Catalog Delta table (wine_quality_cleaned).
  • Unity Catalog Feature table (wine_features).
  • Unity Catalog Registered Model (wine_quality_model with 3 versions).
  • MLflow experiment (with baseline, Hyperopt, and Feature Store runs).
  • Uploaded CSV file in the Unity Catalog volume.

Keep everything running

No action needed. Choose this if you plan to keep practicing with your ML pipeline or want to revisit the project for certification study.

  • Your notebook, tables, registered model, and experiment all remain in your workspace.
  • Databricks Free Edition has no ongoing costs. Your workspace stays active indefinitely as long as you log in occasionally.
  • Databricks may delete Free Edition accounts that are inactive for a prolonged period, so sign in periodically if you want to keep your work.

Pause - I'll come back to this later

Since Databricks Free Edition uses serverless compute that only runs when you execute cells, there is nothing to shut down. Your workspace is effectively paused whenever you close the browser.

  • Close your Databricks browser tab. No compute continues running after you stop executing cells.
  • Your notebook, tables, model, and experiment persist in the workspace until you return.
  • Sign in again whenever you want to resume. Your data and settings will still be there.

Delete - I don't want to use this again

Remove all project resources from your workspace. Follow each action below in order.

  • Drop the Delta table and Feature table. In your notebook, add a new cell and run the SQL commands below.
DROP TABLE IF EXISTS wine_quality_cleaned;
DROP TABLE IF EXISTS wine_features;

This removes both tables and their underlying data from Unity Catalog.

  • Delete the registered model. Click the Catalog icon in the left sidebar to open Catalog Explorer. Navigate to your model (wine_quality_model). On the model page, click the kebab menu (three dots) in the upper-right corner, then select Delete. Confirm the deletion in the dialog.
  • Delete the MLflow experiment. Click Experiments in the left sidebar. Find your experiment in the list. Click the kebab menu in the rightmost column of the experiment row, then click Delete.
  • Delete the notebook. In the workspace file browser, right-click your notebook and select Move to Trash.
  • Delete the uploaded CSV. Click the Catalog icon in the left sidebar. Navigate to your volume in the default schema. Select the winequality-red.csv file, click the kebab menu, and select Delete.

What about the Databricks account itself?

There is no self-service account deletion for Free Edition. Databricks automatically deletes accounts that remain inactive for a prolonged period.

If you want immediate deletion, you can file a support ticket at help.databricks.com.

Nice Work!

Nice Work!

You just built a complete wine quality prediction pipeline on Databricks, from raw data exploration all the way to a production-registered model with Champion and Challenger aliases.

You've learned how to:

  • Build an end-to-end ML pipeline on Databricks using Spark DataFrames for exploration and MLflow for experiment tracking, logging parameters, metrics, artifacts, and models.
  • Systematically improve model performance using Hyperopt's fmin with Bayesian optimization (TPE), comparing tuned runs against a baseline in the MLflow UI.
  • Register models in Unity Catalog and manage their lifecycle with aliases (Champion/Challenger) and tags to control which version is promoted to production.
  • Secret Mission: Integrated the Feature Store with FeatureLookup and create_training_set for production-ready feature management and automatic batch scoring.

Ready to quiz yourself?