ML Experimentation Pipeline on Databricks

Build a versioned ML pipeline with Delta Lake, MLflow, and Optuna on Databricks.

Introduction

30 Second Summary

Every ML practitioner knows the frustration of getting different results after retraining a model. Was it the code, the hyperparameters, or the data that changed?

In this project, you will build an ML experimentation pipeline on Databricks that connects data versioning to experiment tracking. You will load wine quality data into a versioned table, tune a classifier, and retrain on earlier data versions to see exactly how data changes affect model performance.

What You'll Build

An end-to-end ML experimentation pipeline on Databricks Free Edition that uses Delta Lake for data versioning, scikit-learn with Optuna for model training and hyperparameter tuning, and MLflow for experiment tracking.

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

  • A Delta Lake table with multiple versions stored in Unity Catalog.
  • An MLflow experiment with 30+ tracked runs from hyperparameter tuning across multiple training sessions.
  • A reproducible workflow that retrains models on earlier data versions using Delta Lake time travel and compares performance across data states.
  • Secret Mission: Register your best model in Unity Catalog and run batch inference on new data.

Are there any prerequisites?

You should be familiar with Python and ML model training (classification). No Databricks account is needed, as this project includes full signup instructions.

Set Up Databricks Free Edition

Before you can build an ML pipeline, you need a workspace that ties together data storage, experiment tracking, and compute. That is exactly what Databricks provides out of the box.

In this step, you will sign up for Databricks Free Edition, create a Python notebook, and confirm that Unity Catalog is ready to go.

In this step, get ready to:
  • Sign up for Databricks Free Edition and land in your workspace.
  • Create a Python notebook connected to serverless compute.
  • Verify your Unity Catalog catalog and schema.
Sign up for Databricks Free Edition

Databricks Free Edition is a forever-free tier designed for learning and training. No credit card is required.

  • Open the Databricks Free Edition signup page in your browser:
https://login.databricks.com/signup?provider=DB_FREE_TIER
  • Choose a signup method. You can sign up with email OTP, Sign in with Google, or Sign in with Microsoft.
  • Complete the signup flow. Databricks will create a serverless workspace for you automatically.

Once signup finishes, you land in the Databricks workspace. You should see a sidebar on the left with options including New, Workspace, and Catalog.

What is Databricks Free Edition?

Free Edition gives you a fully managed Databricks workspace with serverless compute. It has daily usage limits but is plenty for completing this project.

Free Edition replaces the older Community Edition, which was retired on January 1, 2026.

Having trouble signing up?

Make sure you are using a personal email address. Databricks Free Edition is designed for personal learning, not company accounts.

If you chose email OTP and did not receive a code, check your spam folder.

help me troubleshoot signing up for Databricks Free Edition

Create a notebook

Databricks notebooks are where you write and run code. Each notebook connects to a compute resource that executes your cells.

  • In the sidebar, click New.
  • Click Notebook.

A blank Python notebook opens, already connected to serverless compute.

  • At the top of the notebook, confirm the compute selector shows Serverless.
  • Click the notebook title at the top left of the page.
  • Rename the notebook to Wine Quality ML Pipeline.

What is serverless compute?

Serverless compute means Databricks manages the infrastructure for you. You do not need to configure clusters or wait for them to start. Free Edition only supports serverless compute, so the selector should already show the correct option.

Don't see the Notebook option?

Make sure you are clicking New in the left sidebar, not a menu inside another page. The Notebook option should appear in the dropdown.

help me create a new notebook in Databricks Free Edition

Verify your catalog and schema

Unity Catalog is the governance layer that organizes your tables, models, and data. Free Edition workspaces have Unity Catalog enabled by default with a workspace catalog and the default schema ready to use.

You need to know your catalog name because you will reference it every time you create or read a table.

  • Confirm Unity Catalog is working by pasting this code into the first cell of your notebook and pressing Shift+Enter to run it:
# Print the current catalog and schema names
catalog_name = spark.catalog.currentCatalog()
schema_name = spark.catalog.currentDatabase()

print(f"Catalog: {catalog_name}")
print(f"Schema: {schema_name}")

What does this code do?

spark.catalog.currentCatalog() returns the name of the catalog your notebook is connected to. The workspace catalog shares a name with your workspace.

spark.catalog.currentDatabase() returns the current schema. For Free Edition users, this defaults to default.

You should see output similar to this:

Catalog: your_workspace_name
Schema: default
  • Record the catalog name printed in the output. It shares a name with your workspace. You will use this catalog and the default schema throughout the project.

Find your catalog name in the sidebar

You can also find your catalog name by clicking Catalog in the sidebar. The workspace catalog is listed in the catalog panel and shares its name with your workspace.

Seeing an error when running the cell?

Make sure the compute selector at the top of the notebook shows Serverless. If it shows something else, click the selector and choose Serverless.

Serverless compute can take a few seconds to start on the first run. Wait for the cell to finish before moving on.

help me troubleshoot running my first cell in a Databricks Free Edition notebook

Your Databricks workspace is ready, your notebook is connected to serverless compute, and you know your catalog name. Next up, you will load wine quality data into a Delta Lake table and start building your ML pipeline.

Load Data into Delta Lake

Your Databricks workspace is set up and your notebook is connected to serverless compute. Now it's time to bring in the data you'll work with throughout this project.

You'll load the UCI Wine Quality dataset from Databricks' built-in sample data. After cleaning and combining the red and white wine files, you'll store the result as a Delta Lake table in Unity Catalog.

In this step, get ready to:
  • Read and combine the red and white wine quality datasets.
  • Save the data as a Delta Lake table in Unity Catalog.
  • Verify the table's version history with DESCRIBE HISTORY.
Read and combine the wine quality datasets

The wine quality dataset is pre-loaded in every Databricks workspace under /databricks-datasets. It comes as two separate CSV files for red and white wine.

  • In your notebook, click into a new empty cell and paste the following code to load both wine datasets, clean column names, and combine them:
from pyspark.sql.functions import lit

# Read the red wine CSV from databricks-datasets
red_wine = (spark.read
    .option("header", True)
    .option("sep", ";")
    .option("inferSchema", True)
    .csv("/databricks-datasets/wine-quality/winequality-red.csv"))

# Read the white wine CSV
white_wine = (spark.read
    .option("header", True)
    .option("sep", ";")
    .option("inferSchema", True)
    .csv("/databricks-datasets/wine-quality/winequality-white.csv"))

# Clean column names by replacing spaces with underscores
for c in red_wine.columns:
    red_wine = red_wine.withColumnRenamed(c, c.replace(" ", "_"))
for c in white_wine.columns:
    white_wine = white_wine.withColumnRenamed(c, c.replace(" ", "_"))

# Add is_red column to distinguish wine types
red_wine = red_wine.withColumn("is_red", lit(True))
white_wine = white_wine.withColumn("is_red", lit(False))

# Combine both datasets into a single DataFrame
wine_df = red_wine.union(white_wine)

print(f"Combined dataset: {wine_df.count()} rows")
display(wine_df.limit(5))

What does this code do?

The wine quality CSVs use semicolons as separators instead of commas. The sep=";" option tells Spark to parse them correctly. The inferSchema option automatically detects column data types (like double for decimal values) instead of reading everything as strings.

Column names in the original CSV contain spaces (e.g. "fixed acidity"). The withColumnRenamed loop replaces spaces with underscores so you can reference columns cleanly in code later.

The is_red column tags each row with its wine type before the two DataFrames are combined with union.

  • Run the cell by pressing Shift+Enter.
  • You should see a row count printed (around 6,497 rows) and a preview table showing 5 rows with underscored column names and an is_red column on the right.

Seeing only one column in the output?

If all the data appears crammed into a single column, the CSV separator is not being parsed correctly. Make sure your .option("sep", ";") line uses a semicolon inside the quotes, not a comma.

Why is my wine quality data loading into a single column?

Save the combined DataFrame as a Delta Lake table

With your data cleaned and combined, you'll persist it as a Delta Lake table. Delta Lake is the default storage format on Databricks, and every write is automatically versioned.

  • In a new cell, paste the following code to save the DataFrame as a table and verify it:
# Save the combined DataFrame as a Delta Lake table in Unity Catalog
wine_df.write.saveAsTable("wine_quality")

# Read the table back to verify it was created
wine_table = spark.table("wine_quality")

# Display row count, schema, and sample rows
print(f"Row count: {wine_table.count()}")
wine_table.printSchema()
display(wine_table.limit(5))

What does this code do?

The saveAsTable("wine_quality") call writes your DataFrame to your workspace catalog's default schema as a Delta table. Because Delta Lake is the default format on Databricks, you don't need to specify a format.

Reading the table back with spark.table("wine_quality") confirms the data was persisted. The printSchema() method prints a tree of all column names and their data types so you can verify the schema looks correct.

  • Run the cell.
  • You should see the row count printed, a schema tree listing each column with its type (e.g. fixed_acidity: double), and a 5-row preview table.

Getting a table already exists error?

This happens if you ran the cell more than once. The default save mode raises an error when the table already exists. Drop the table first by running this in a new cell:

spark.sql("DROP TABLE IF EXISTS wine_quality")

Then re-run the saveAsTable cell.

Why am I getting a table already exists error?

View the table's version history

Every write to a Delta table is recorded as a new version. You can inspect this history at any time to see what operations were performed and when.

  • In a new cell, paste the following code to view the version history of your table:
# View the version history of the Delta table
display(spark.sql("DESCRIBE HISTORY wine_quality"))

What is DESCRIBE HISTORY?

DESCRIBE HISTORY is a SQL command that returns the full version log of a Delta table. Each row represents a write operation, showing the version number, timestamp, operation type, and other metadata.

Right now you should see a single row: version 0 with a WRITE operation. This is your baseline data state. In the next step, you'll create new versions and use time travel to query any prior state.

  • Run the cell.
  • You should see a table with one row showing version 0, the operation WRITE, and a timestamp.

Not seeing any history output?

Double-check that the table name wine_quality in your DESCRIBE HISTORY command matches exactly what you used in saveAsTable. Table names are case-sensitive.

Why isn't DESCRIBE HISTORY returning any results?

✔️ Awesome, I've got everything!

Great. Double check you've run all three cells and your notebook is auto-saved.

ⓧ I'd like to double check the full code

Here are all three cells from this step combined. Compare yours against this reference:

# --- Cell 1: Read and combine wine quality datasets ---

from pyspark.sql.functions import lit

# Read the red wine CSV from databricks-datasets
red_wine = (spark.read
    .option("header", True)
    .option("sep", ";")
    .option("inferSchema", True)
    .csv("/databricks-datasets/wine-quality/winequality-red.csv"))

# Read the white wine CSV
white_wine = (spark.read
    .option("header", True)
    .option("sep", ";")
    .option("inferSchema", True)
    .csv("/databricks-datasets/wine-quality/winequality-white.csv"))

# Clean column names by replacing spaces with underscores
for c in red_wine.columns:
    red_wine = red_wine.withColumnRenamed(c, c.replace(" ", "_"))
for c in white_wine.columns:
    white_wine = white_wine.withColumnRenamed(c, c.replace(" ", "_"))

# Add is_red column to distinguish wine types
red_wine = red_wine.withColumn("is_red", lit(True))
white_wine = white_wine.withColumn("is_red", lit(False))

# Combine both datasets into a single DataFrame
wine_df = red_wine.union(white_wine)

print(f"Combined dataset: {wine_df.count()} rows")
display(wine_df.limit(5))
# --- Cell 2: Save as a Delta Lake table and verify ---

# Save the combined DataFrame as a Delta Lake table in Unity Catalog
wine_df.write.saveAsTable("wine_quality")

# Read the table back to verify it was created
wine_table = spark.table("wine_quality")

# Display row count, schema, and sample rows
print(f"Row count: {wine_table.count()}")
wine_table.printSchema()
display(wine_table.limit(5))
# --- Cell 3: View version history ---

# View the version history of the Delta table
display(spark.sql("DESCRIBE HISTORY wine_quality"))

Your baseline data is now versioned in Delta Lake. In the next step, you'll engineer features and create new table versions to see time travel in action.

Engineer Features and Version Data

Your Delta Lake table holds the raw wine data at version 0. Raw columns like fixed_acidity and quality are useful on their own, but a classifier needs more signal to work with.

In this step, you will engineer new features, convert the quality score into a binary label, and then deliberately mutate the data to create multiple table versions. Delta Lake tracks every overwrite automatically, so you can query any past state of the table at any time using time travel.

In this step, get ready to:
  • Engineer features and create a binary classification target.
  • Simulate a data update to create a second table version.
  • Query historical versions using Delta Lake time travel.
Create the classification target and engineer features

The wine quality dataset scores each wine from 1 to 10. For a binary classifier, you need a True/False label instead of a numeric score. You will also combine existing columns into new features that capture chemical relationships a model can learn from.

  • In your notebook, hover below the last cell until a + icon appears, then click it to add a new code cell.
  • Add the feature engineering code by pasting the following into the new cell:
from pyspark.sql.functions import col, when

# Read the current wine quality table
df = spark.table("wine_quality")

# Create binary classification target
# High quality = True when quality score is 7 or above
df = df.withColumn("high_quality", when(col("quality") >= 7, True).otherwise(False))

# Engineer derived features
df = df.withColumn("total_acidity", col("fixed_acidity") + col("volatile_acidity"))
df = df.withColumn("alcohol_to_sugar_ratio", col("alcohol") / col("residual_sugar"))
df = df.withColumn("sulfur_ratio", col("free_sulfur_dioxide") / col("total_sulfur_dioxide"))

# Overwrite the Delta table to save as version 1
df.write.mode("overwrite").saveAsTable("wine_quality")

# Verify the new columns exist
print("Version 1 created with engineered features.")
display(spark.table("wine_quality").select(
    "quality", "high_quality", "total_acidity",
    "alcohol_to_sugar_ratio", "sulfur_ratio"
).limit(5))

What does this code do?

  • high_quality converts the 1-10 quality score into a binary label. Wines scoring 7 or above are True. Everything else is False.
  • total_acidity sums fixed and volatile acidity into a single measure of overall acid content.
  • alcohol_to_sugar_ratio captures the balance between alcohol and residual sugar, which differs sharply between wine types.
  • sulfur_ratio compares free sulfur dioxide to total sulfur dioxide. This ratio signals how much of the preservative is in its active form.
  • The .write.mode("overwrite").saveAsTable() call replaces all the data in the table. Delta Lake automatically records this as a new version.

Why overwrite instead of creating separate tables?

Delta Lake records every write to a table as a new version. Overwriting the same table gives you a complete history of how the data evolved, without cluttering your catalog with wine_quality_v1, wine_quality_v2 copies.

You can always go back to any prior state with a time travel query. One table, full history.

  • Run the cell by pressing Shift+Enter.
  • You should see a table with five sample rows showing the new columns: high_quality, total_acidity, alcohol_to_sugar_ratio, and sulfur_ratio.

Seeing a AnalysisException or column not found error?

  • Make sure the cells from Step 2 that create the wine_quality table have been run in this session. If you restarted the notebook kernel, re-run all previous cells first.
  • Check that the column names in your table use underscores (e.g. fixed_acidity), not spaces. The column-cleaning cell from Step 2 handles this.
  • Still stuck? Help me debug this column not found error in my Databricks notebook.
Simulate a data update to create version 2

Low-alcohol wines (below 9.0) tend to add noise to a quality classifier. Filtering them out creates a more focused dataset for training. This overwrite simulates a real-world scenario where a data team refines a dataset after an initial analysis.

  • Add a new cell below the previous one.
  • Filter out low-alcohol wines and overwrite the table by pasting the following:
# Read the table (now version 1 with engineered features)
df = spark.table("wine_quality")

# Filter out low-alcohol wines (below 9.0) to reduce noise
df_filtered = df.filter(col("alcohol") >= 9.0)

# Overwrite the Delta table to create version 2
df_filtered.write.mode("overwrite").saveAsTable("wine_quality")

# Verify the row count after filtering
print(f"Version 2 created. Row count: {spark.table('wine_quality').count()}")

What does this code do?

  • The .filter(col("alcohol") >= 9.0) keeps only wines with alcohol content of 9.0 or higher. This removes a chunk of low-alcohol entries that could confuse the classifier.
  • The overwrite creates version 2 of the table. Delta Lake still has versions 0 and 1 intact underneath.
  • Run the cell by pressing Shift+Enter.
  • The output should print a row count that is lower than what you saw in Step 2. The filter removed a portion of the wines.

Row count looks the same as before?

  • Double check the filter condition. It should be >= 9.0 (greater than or equal). A typo like >= 0.9 would keep nearly everything.
  • Make sure the feature engineering cell (version 1) ran successfully before this one.
  • Need a hand? Help me figure out why my alcohol filter didn't reduce the row count.
Confirm versions and query with time travel

Delta Lake time travel lets you read any past version of a table using a simple SQL clause. You will confirm the table's version history, then query two versions side-by-side to see the impact of the alcohol filter.

  • Add a new cell and view the table's full version history by pasting the following:
# View the version history of the Delta table
spark.sql("DESCRIBE HISTORY wine_quality").select("version", "timestamp", "operation").show(truncate=False)
  • Run the cell. You should see three rows in the output: versions 0, 1, and 2, each with a WRITE operation.

Now compare the two most recent versions to see the data difference in action.

  • Add one more cell and query both versions using time travel by pasting this code:
# Query specific versions using Delta Lake time travel
v1_count = spark.sql("SELECT COUNT(*) FROM wine_quality VERSION AS OF 1").collect()[0][0]
v2_count = spark.sql("SELECT COUNT(*) FROM wine_quality VERSION AS OF 2").collect()[0][0]

# Display the comparison
print(f"Version 1 (all features, no filter): {v1_count} rows")
print(f"Version 2 (alcohol >= 9.0 filter):   {v2_count} rows")
print(f"Difference: {v1_count - v2_count} rows removed")

What does this code do?

  • VERSION AS OF 1 tells Delta Lake to read the table as it existed at version 1, before the alcohol filter was applied.
  • VERSION AS OF 2 reads the current version with the filter applied.
  • The printed difference tells you exactly how many wines the alcohol filter removed. Later, you will retrain a model on version 1 data to see how this filtering decision affects model performance.
  • Run the cell. You should see version 1 with a higher row count and version 2 with a lower count, along with the difference.

Getting a version not found error?

  • Make sure both the feature engineering cell (version 1) and the alcohol filter cell (version 2) ran successfully before this one.
  • Run the DESCRIBE HISTORY cell again to confirm three versions exist. If you only see one or two, re-run the missing cells above.
  • Need help? Help me troubleshoot a Delta Lake time travel VERSION AS OF query that is failing.

✔️ Awesome, I've got everything!

Make sure all four cells in this step ran without errors and that your time travel query shows two different row counts.

ⓧ I'd like to double check the full code

Here are all four cells from this step. Compare each one with the matching cell in your notebook.

# --- Cell 1: Feature Engineering ---
from pyspark.sql.functions import col, when

# Read the current wine quality table
df = spark.table("wine_quality")

# Create binary classification target
# High quality = True when quality score is 7 or above
df = df.withColumn("high_quality", when(col("quality") >= 7, True).otherwise(False))

# Engineer derived features
df = df.withColumn("total_acidity", col("fixed_acidity") + col("volatile_acidity"))
df = df.withColumn("alcohol_to_sugar_ratio", col("alcohol") / col("residual_sugar"))
df = df.withColumn("sulfur_ratio", col("free_sulfur_dioxide") / col("total_sulfur_dioxide"))

# Overwrite the Delta table to save as version 1
df.write.mode("overwrite").saveAsTable("wine_quality")

# Verify the new columns exist
print("Version 1 created with engineered features.")
display(spark.table("wine_quality").select(
    "quality", "high_quality", "total_acidity",
    "alcohol_to_sugar_ratio", "sulfur_ratio"
).limit(5))

# --- Cell 2: Data Filter ---
# Read the table (now version 1 with engineered features)
df = spark.table("wine_quality")

# Filter out low-alcohol wines (below 9.0) to reduce noise
df_filtered = df.filter(col("alcohol") >= 9.0)

# Overwrite the Delta table to create version 2
df_filtered.write.mode("overwrite").saveAsTable("wine_quality")

# Verify the row count after filtering
print(f"Version 2 created. Row count: {spark.table('wine_quality').count()}")

# --- Cell 3: Version History ---
# View the version history of the Delta table
spark.sql("DESCRIBE HISTORY wine_quality").select("version", "timestamp", "operation").show(truncate=False)

# --- Cell 4: Time Travel Comparison ---
# Query specific versions using Delta Lake time travel
v1_count = spark.sql("SELECT COUNT(*) FROM wine_quality VERSION AS OF 1").collect()[0][0]
v2_count = spark.sql("SELECT COUNT(*) FROM wine_quality VERSION AS OF 2").collect()[0][0]

# Display the comparison
print(f"Version 1 (all features, no filter): {v1_count} rows")
print(f"Version 2 (alcohol >= 9.0 filter):   {v2_count} rows")
print(f"Difference: {v1_count - v2_count} rows removed")

Your data is versioned and queryable at any point in time. Next up, you will train and tune a classification model on this data, with every experiment tracked in MLflow.

Train and Track Models with MLflow

Your Delta Lake table has three clean, versioned snapshots of wine data. The next question is: which model performs best on this data, and how do you prove it?

In this step, you will train a baseline classifier, then use Optuna to search for better hyperparameters. Every run is automatically tracked by MLflow, so you can compare models and reproduce results later.

In this step, get ready to:
  • Train a baseline GradientBoostingClassifier and track it with MLflow.
  • Run 20 hyperparameter tuning trials with Optuna, logged as nested MLflow runs.
  • Find and evaluate the best model using the MLflow API.
Set up the ML workflow and train a baseline model

Before tuning anything, you need a baseline. A baseline model uses default settings so you have a reference point to measure improvements against.

  • In your notebook, create a new cell and load the latest table version as a Pandas DataFrame by pasting the following code:
import numpy as np
import pandas as pd
import sklearn.datasets
import sklearn.metrics
import sklearn.model_selection
import sklearn.ensemble
import optuna
from mlflow.optuna.storage import MlflowStorage
from mlflow.pyspark.optuna.study import MlflowSparkStudy

# Load version 2 of the wine quality table as a Pandas DataFrame
df = spark.read.table(f"{CATALOG_NAME}.default.wine_quality").toPandas()

# Separate features and target label
X = df.drop(columns=["quality", "high_quality"])
y = df["high_quality"]

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

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

What does this code do?

This cell pulls all the imports you will need for this step. It then reads version 2 of the wine_quality table (the filtered, feature-engineered data) and converts it from a Spark DataFrame into a Pandas DataFrame that scikit-learn can work with.

The quality and high_quality columns are separated out because high_quality is the label you are predicting and quality is the raw score it was derived from. Keeping either in the feature set would leak the answer to the model.

  • Run the cell.
  • You should see the training and test set sizes printed. The exact numbers depend on version 2 of your table.

Seeing a NameError for CATALOG_NAME?

The variable CATALOG_NAME was defined in an earlier step. Make sure you have run all the cells above this one in order. In Databricks, you can click **Run all above** from the cell menu to re-execute everything.

Help me fix my NameError for CATALOG_NAME in my Databricks notebook.

  • Create a new cell and enable MLflow autologging, then train the baseline model by pasting the following code:
import mlflow

# Enable MLflow autologging for scikit-learn
mlflow.sklearn.autolog()

# Train a baseline GradientBoostingClassifier with default parameters
with mlflow.start_run(run_name="baseline_gbc") as run:
    model = sklearn.ensemble.GradientBoostingClassifier(random_state=0)
    # Models, parameters, and training metrics are tracked automatically
    model.fit(X_train, y_train)

    predicted_probs = model.predict_proba(X_test)
    # Calculate the AUC on the test set
    roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:, 1])

    # Log test AUC manually (autolog does not capture test metrics)
    mlflow.log_metric("test_auc", roc_auc)

    print(f"Baseline test AUC: {roc_auc:.4f}")

What does this code do?

The call to mlflow.sklearn.autolog() tells MLflow to automatically capture parameters, metrics, and the model artifact every time you train a scikit-learn model. You do not need to manually log each parameter.

The mlflow.start_run() context manager opens a tracked experiment run. Inside it, you train a GradientBoostingClassifier with default settings. Autologging captures the training details, but does not automatically compute test metrics. That is why you manually log test_auc using mlflow.log_metric().

  • Run the cell.
  • You should see the baseline test AUC printed. This is your starting point for improvement.
  • Click the **Experiment** icon (flask icon) at the upper right of the notebook to open the experiment sidebar.
  • You should see your baseline_gbc run listed with its parameters and metrics.

What is MLflow autologging?

Databricks provides a hosted MLflow tracking server that stores experiment data in your workspace with no setup required. When you call mlflow.sklearn.autolog(), MLflow hooks into scikit-learn's fit() method and logs every parameter, training metric, and model artifact automatically.

If no active experiment is set, runs are logged to the notebook experiment automatically. The notebook experiment shares the same name and ID as your notebook.

Don't see the experiment sidebar?

Look for the flask-shaped **Experiment** icon at the upper right of the notebook. If the sidebar is empty, click the refresh icon to fetch the latest runs. Make sure the cell finished executing without errors before checking.

Help me find my MLflow experiment run in Databricks.

Run hyperparameter tuning with Optuna

A baseline with default parameters is a starting point, not the finish line. Hyperparameter tuning systematically searches for the combination of settings that produces the best model.

Optuna is a hyperparameter optimization framework that Databricks recommends for serverless compute. Each trial trains a model with different settings and logs the result as a nested MLflow run.

  • Create a new cell and define the Optuna objective function and tuning loop by pasting the following code:
def objective(trial):
    # Enable autologging on each worker
    mlflow.sklearn.autolog()

    with mlflow.start_run(nested=True):
        # Define the hyperparameter search space
        params = {
            'n_estimators': trial.suggest_int('n_estimators', 20, 1000),
            'learning_rate': trial.suggest_float('learning_rate', 0.05, 1.0, log=True),
            'max_depth': trial.suggest_int('max_depth', 2, 5),
        }

        model_hp = sklearn.ensemble.GradientBoostingClassifier(
            random_state=0,
            **params
        )
        model_hp.fit(X_train, y_train)

        predicted_probs = model_hp.predict_proba(X_test)
        # Tune based on the test AUC
        # In production, you could use a separate validation set instead
        roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:, 1])
        mlflow.log_metric('test_auc', roc_auc)

        # Negate the AUC because Optuna minimizes the objective by default
        return -roc_auc

with mlflow.start_run(run_name='gb_optuna') as run:
    # Use the MLflow Tracking Server as the Optuna storage backend
    experiment_id = mlflow.active_run().info.experiment_id
    mlflow_storage = MlflowStorage(experiment_id=experiment_id)

    # MlflowSparkStudy distributes the tuning using Spark workers
    mlflow_study = MlflowSparkStudy(
        study_name="gb-optuna-tuning",
        storage=mlflow_storage,
    )

    # Run 20 tuning trials across 4 parallel workers
    mlflow_study.optimize(objective, n_trials=20, n_jobs=4)

What does this code do?

The objective function is called once per trial. Each trial receives a trial object that suggests values for n_estimators, learning_rate, and max_depth from the defined ranges. The function returns the negated AUC because Optuna minimizes by default.

Each trial opens a nested MLflow run under the parent gb_optuna run. MlflowStorage uses the MLflow Tracking Server as Optuna's storage backend, so trial results persist alongside your other experiment data. MlflowSparkStudy distributes the tuning across Spark workers for faster execution.

Why Optuna instead of Hyperopt?

Databricks recommends Optuna for hyperparameter tuning on serverless compute. Hyperopt is the legacy approach for Databricks Runtime 16.4 LTS ML and below. Optuna integrates directly with MLflow through MlflowStorage and MlflowSparkStudy, making it the natural fit for this project.

  • Run the cell. This will take a few minutes as 20 trials train and log sequentially across workers.
  • When the cell finishes, open the experiment sidebar by clicking the **Experiment** icon. You should see the parent gb_optuna run with 20 nested child runs underneath it.

Cell running for a very long time?

20 trials with gradient boosting can take several minutes on serverless compute. If a single trial takes more than 5 minutes, check whether Optuna suggested a very high n_estimators value. This is normal. Wait for the cell to complete before proceeding.

Help me understand why my Optuna tuning is slow on Databricks serverless.

Review the experiment results and load the best model

With 20+ runs logged, the MLflow experiment page gives you a bird's-eye view of what worked. You can sort, compare, and drill into individual runs from one place.

  • In the experiment sidebar, click the experiment page icon (the external-link icon next to the experiment name) to open the full MLflow experiment page.
  • On the experiment page, click the test_auc column header to sort runs by AUC. The top rows show your best-performing hyperparameter combinations.
  • Compare how different values of n_estimators, learning_rate, and max_depth affected performance. This is one of the core benefits of experiment tracking.
  • Switch back to your notebook. Create a new cell and find the best run programmatically by pasting the following code:
# Sort runs by test AUC (descending), break ties by most recent
best_run = mlflow.search_runs(
    order_by=['metrics.test_auc DESC', 'start_time DESC'],
    max_results=10,
).iloc[0]

print(f"Best Run")
print(f"AUC: {best_run['metrics.test_auc']:.4f}")
print(f"Num Estimators: {best_run['params.n_estimators']}")
print(f"Max Depth: {best_run['params.max_depth']}")
print(f"Learning Rate: {best_run['params.learning_rate']}")

What does this code do?

The mlflow.search_runs() API returns a Pandas DataFrame of all runs in the current experiment. Sorting by metrics.test_auc DESC puts the highest-performing run first. The .iloc[0] grabs that top row so you can inspect the winning hyperparameters.

  • Run the cell. You should see the best AUC and its corresponding hyperparameters printed.
  • Create a new cell and load the best model for evaluation by pasting the following code:
# Load the best model from the MLflow tracking server
best_model = mlflow.pyfunc.load_model(
    f"runs:/{best_run.run_id}/model"
)

# Generate predictions using the loaded model
y_pred = best_model.predict(X_test)

# Calculate the final AUC
final_auc = sklearn.metrics.roc_auc_score(
    y_test, best_model.predict(X_test)
)
print(f"Final AUC: {final_auc:.4f}")

# Print the full classification report
print("\nClassification Report:")
print(sklearn.metrics.classification_report(y_test, y_pred))

What does this code do?

The mlflow.pyfunc.load_model() call downloads the serialized model artifact from the MLflow tracking server and makes it available for predictions. This proves the model is fully reproducible. Anyone with access to the experiment can load the same model and get identical results.

The classification report shows precision, recall, and F1-score for each class (regular quality vs. high quality). Combined with the AUC, this gives you a complete picture of model performance.

  • Run the cell. You should see the final AUC score and a classification report with precision, recall, and F1-score for each class printed below.

Seeing a KeyError or empty results?

If mlflow.search_runs() returns an empty DataFrame, the Optuna cell may not have completed successfully. Go back and re-run the tuning cell. Make sure you see the 20 nested runs in the experiment sidebar before running the search cell.

Help me debug why mlflow.search_runs returns no results.

✔️ Awesome, I've got everything!

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

ⓧ I'd like to double check the full code

Here is all the code you should have added in this step across your notebook cells. Compare cell by cell to make sure nothing is missing.

Cell 1: Imports and data loading.

import numpy as np
import pandas as pd
import sklearn.datasets
import sklearn.metrics
import sklearn.model_selection
import sklearn.ensemble
import optuna
from mlflow.optuna.storage import MlflowStorage
from mlflow.pyspark.optuna.study import MlflowSparkStudy

# Load version 2 of the wine quality table as a Pandas DataFrame
df = spark.read.table(f"{CATALOG_NAME}.default.wine_quality").toPandas()

# Separate features and target label
X = df.drop(columns=["quality", "high_quality"])
y = df["high_quality"]

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

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

Cell 2: Baseline model training.

import mlflow

# Enable MLflow autologging for scikit-learn
mlflow.sklearn.autolog()

# Train a baseline GradientBoostingClassifier with default parameters
with mlflow.start_run(run_name="baseline_gbc") as run:
    model = sklearn.ensemble.GradientBoostingClassifier(random_state=0)
    # Models, parameters, and training metrics are tracked automatically
    model.fit(X_train, y_train)

    predicted_probs = model.predict_proba(X_test)
    # Calculate the AUC on the test set
    roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:, 1])

    # Log test AUC manually (autolog does not capture test metrics)
    mlflow.log_metric("test_auc", roc_auc)

    print(f"Baseline test AUC: {roc_auc:.4f}")

Cell 3: Optuna hyperparameter tuning.

def objective(trial):
    # Enable autologging on each worker
    mlflow.sklearn.autolog()

    with mlflow.start_run(nested=True):
        # Define the hyperparameter search space
        params = {
            'n_estimators': trial.suggest_int('n_estimators', 20, 1000),
            'learning_rate': trial.suggest_float('learning_rate', 0.05, 1.0, log=True),
            'max_depth': trial.suggest_int('max_depth', 2, 5),
        }

        model_hp = sklearn.ensemble.GradientBoostingClassifier(
            random_state=0,
            **params
        )
        model_hp.fit(X_train, y_train)

        predicted_probs = model_hp.predict_proba(X_test)
        # Tune based on the test AUC
        # In production, you could use a separate validation set instead
        roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:, 1])
        mlflow.log_metric('test_auc', roc_auc)

        # Negate the AUC because Optuna minimizes the objective by default
        return -roc_auc

with mlflow.start_run(run_name='gb_optuna') as run:
    # Use the MLflow Tracking Server as the Optuna storage backend
    experiment_id = mlflow.active_run().info.experiment_id
    mlflow_storage = MlflowStorage(experiment_id=experiment_id)

    # MlflowSparkStudy distributes the tuning using Spark workers
    mlflow_study = MlflowSparkStudy(
        study_name="gb-optuna-tuning",
        storage=mlflow_storage,
    )

    # Run 20 tuning trials across 4 parallel workers
    mlflow_study.optimize(objective, n_trials=20, n_jobs=4)

Cell 4: Find the best run.

# Sort runs by test AUC (descending), break ties by most recent
best_run = mlflow.search_runs(
    order_by=['metrics.test_auc DESC', 'start_time DESC'],
    max_results=10,
).iloc[0]

print(f"Best Run")
print(f"AUC: {best_run['metrics.test_auc']:.4f}")
print(f"Num Estimators: {best_run['params.n_estimators']}")
print(f"Max Depth: {best_run['params.max_depth']}")
print(f"Learning Rate: {best_run['params.learning_rate']}")

Cell 5: Load best model and evaluate.

# Load the best model from the MLflow tracking server
best_model = mlflow.pyfunc.load_model(
    f"runs:/{best_run.run_id}/model"
)

# Generate predictions using the loaded model
y_pred = best_model.predict(X_test)

# Calculate the final AUC
final_auc = sklearn.metrics.roc_auc_score(
    y_test, best_model.predict(X_test)
)
print(f"Final AUC: {final_auc:.4f}")

# Print the full classification report
print("\nClassification Report:")
print(sklearn.metrics.classification_report(y_test, y_pred))

You now have an MLflow experiment packed with 20+ tracked runs, a clear winner identified by AUC, and a fully reproducible model you can reload anytime. Next up, you will put Delta Lake time travel to work by retraining on an earlier data version and comparing how data changes affect model performance.

Retrain on a Different Data Version

In the previous step, you trained and tuned a GradientBoostingClassifier on version 2 of your Delta Lake table. You found the best hyperparameters using Optuna and tracked every run in MLflow. But here is the real question: how much did your data changes affect the model?

Version 2 removed low-alcohol wines to reduce noise. In this step, you will retrain on version 1 (the full dataset before filtering) using Delta Lake time travel. By comparing results across data versions, you will see exactly how data changes ripple through to model performance.

In this step, get ready to:
  • Load version 1 of your wine quality table using Delta Lake time travel.
  • Retrain the best model on version 1 data and track the run with MLflow.
  • Build a summary comparison showing model performance across data versions.
Load version 1 data using time travel

Version 1 contains the full engineered dataset before the alcohol filter was applied. Loading it with the same train/test split ensures any performance difference comes from the data, not from randomization.

  • Create a new cell in your notebook and load version 1 of the wine quality table by pasting this code:
# Load version 1 of the wine quality table using Delta Lake time travel
df_v1 = spark.sql(f"SELECT * FROM {catalog_name}.default.wine_quality VERSION AS OF 1")

# Convert to a Pandas DataFrame for scikit-learn
pdf_v1 = df_v1.toPandas()

# Use the same feature columns and target as version 2 training
feature_cols = [c for c in pdf_v1.columns if c not in ["quality", "high_quality"]]
X_v1 = pdf_v1[feature_cols]
y_v1 = pdf_v1["high_quality"].astype(int)

# Split with the same random seed and ratio as version 2
from sklearn.model_selection import train_test_split
X_train_v1, X_test_v1, y_train_v1, y_test_v1 = train_test_split(
    X_v1, y_v1, test_size=0.2, random_state=42
)

print(f"Version 1 dataset: {len(pdf_v1)} total rows")
print(f"Training: {len(X_train_v1)} rows, Test: {len(X_test_v1)} rows")

What does this code do?

The SQL query uses VERSION AS OF 1 to read the table exactly as it was after the feature engineering step, before low-alcohol wines were removed.

Using the same random_state=42 and test_size=0.2 as the version 2 split is critical. It ensures any difference in AUC is caused by the data content, not by a different random split.

  • Run the cell by pressing Shift+Enter.
  • You should see the total row count for version 1 and the train/test split sizes. Version 1 has more rows than version 2 because the alcohol filter has not been applied.

Seeing a table not found error?

Make sure catalog_name is still defined from your earlier cells. If your notebook session restarted, re-run the cell from Step 1 that captures your catalog name.

help me fix a table not found error when querying my Delta Lake table with VERSION AS OF

Retrain with the best hyperparameters on version 1

You already know which hyperparameters performed best on version 2 from the Optuna tuning in Step 4. Now you will use those same hyperparameters on version 1 data to isolate the effect of the data change.

  • Create a new cell and retrain the model on version 1 data by pasting this code:
import mlflow
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score, classification_report

# Extract the best hyperparameters from the version 2 tuning
best_n_estimators = int(best_run["params.n_estimators"])
best_learning_rate = float(best_run["params.learning_rate"])
best_max_depth = int(best_run["params.max_depth"])

# Turn off autologging to control logging manually
mlflow.sklearn.autolog(disable=True)

# Train on version 1 data with the best hyperparameters
with mlflow.start_run(run_name="GBT_retrain_data_v1") as run:
    # Tag this run with the data version for filtering later
    mlflow.set_tag("data_version", "1")
    mlflow.log_param("n_estimators", best_n_estimators)
    mlflow.log_param("learning_rate", best_learning_rate)
    mlflow.log_param("max_depth", best_max_depth)
    mlflow.log_param("data_rows", len(pdf_v1))

    # Build the classifier with the same winning configuration
    model_v1 = GradientBoostingClassifier(
        n_estimators=best_n_estimators,
        learning_rate=best_learning_rate,
        max_depth=best_max_depth,
        random_state=42
    )
    model_v1.fit(X_train_v1, y_train_v1)

    # Evaluate on the test set
    y_prob_v1 = model_v1.predict_proba(X_test_v1)[:, 1]
    auc_v1 = roc_auc_score(y_test_v1, y_prob_v1)
    mlflow.log_metric("test_auc", auc_v1)
    mlflow.sklearn.log_model(model_v1, "model")

    # Save the run ID for later comparison
    v1_run_id = run.info.run_id

print(f"Version 1 AUC: {auc_v1:.4f}")
print(f"Version 2 Best AUC: {best_auc:.4f}")
print(f"Difference: {auc_v1 - best_auc:.4f}")
print("\nClassification Report (Version 1 Data):")
print(classification_report(y_test_v1, model_v1.predict(X_test_v1)))

What does this code do?

mlflow.set_tag("data_version", "1") attaches metadata to this run so you can filter and group runs by data version later. Tags are free-form key-value strings stored alongside every run.

The code explicitly logs each hyperparameter with mlflow.log_param() because autologging is disabled. This gives you full control over what appears in the experiment. The final comparison prints the AUC from both data versions side by side.

  • Run the cell by pressing Shift+Enter.
  • You should see the version 1 AUC, the version 2 best AUC, and the difference between them. You will also see a full classification report showing precision, recall, and F1 scores for version 1.

Seeing a NameError for best_run or best_auc?

These variables were created in Step 4 when you searched for the best Optuna trial. If your notebook session restarted, re-run the cell from Step 4 that calls mlflow.search_runs() to find the best run.

help me fix a NameError for best_run in my Databricks notebook after retraining on a different data version

Create a summary comparison across data versions

The individual AUC comparison is useful, but a structured summary table makes the relationship between data state and model performance unmistakable. You will tag the version 2 best run, then use mlflow.search_runs() to pull both runs into a single comparison view.

  • Create a new cell and build the comparison table by pasting this code:
from mlflow import MlflowClient
import pandas as pd

# Tag the version 2 best run so it appears in the comparison
client = MlflowClient()
client.set_tag(best_run["run_id"], "data_version", "2")

# Get row counts from each data version using time travel
v1_rows = spark.sql(
    f"SELECT COUNT(*) as cnt FROM {catalog_name}.default.wine_quality VERSION AS OF 1"
).collect()[0]["cnt"]
v2_rows = spark.sql(
    f"SELECT COUNT(*) as cnt FROM {catalog_name}.default.wine_quality VERSION AS OF 2"
).collect()[0]["cnt"]

# Search for all runs tagged with a data version
comparison_runs = mlflow.search_runs(
    filter_string="tags.data_version IS NOT NULL",
    order_by=["metrics.test_auc DESC"]
)

# Build a summary table with the key metrics per data version
summary_data = []
for version, rows in [("1", v1_rows), ("2", v2_rows)]:
    version_runs = comparison_runs[comparison_runs["tags.data_version"] == version]
    if not version_runs.empty:
        best = version_runs.iloc[0]
        summary_data.append({
            "Data Version": version,
            "Row Count": rows,
            "Best AUC": round(best["metrics.test_auc"], 4),
            "n_estimators": best["params.n_estimators"],
            "learning_rate": best["params.learning_rate"],
            "max_depth": best["params.max_depth"]
        })

summary = pd.DataFrame(summary_data)
print("Model Performance vs. Data Version")
print("=" * 55)
display(summary)

What does this code do?

client.set_tag(best_run["run_id"], "data_version", "2") retroactively tags the version 2 best run. This lets you filter all comparable runs with a single search_runs query.

The filter_string="tags.data_version IS NOT NULL" pulls only runs you explicitly tagged. The loop then groups results by data version and picks the best AUC for each, creating a clean comparison of data version, row count, best AUC, and winning hyperparameters.

  • Run the cell by pressing Shift+Enter.
  • You should see a table with two rows. One row for data version 1 and one for data version 2. Each row shows the row count, best AUC, and the hyperparameters that produced it.

Why does this matter?

This is the payoff of connecting Delta Lake to MLflow. You can see exactly how a data change (filtering low-alcohol wines) affected model performance. Both the data state and the model artifacts are fully reproducible. You can reload any version of the data and retrain at any time.

In a production ML workflow, this pattern lets you answer questions like: did last week's data refresh improve the model, or make it worse?

Comparison table only shows one row?

Make sure the client.set_tag() call completed without errors. If best_run["run_id"] is empty or invalid, the version 2 run will not be tagged and the search will not find it.

help me debug why my MLflow search_runs comparison table only shows one data version instead of two

✔️ Awesome, I've got everything!

Great. Double check you have run all three cells and that your summary table shows both data versions.

ⓧ I'd like to double check the full code

Here is the full code added in this step across three cells.

Cell 1: Load version 1 data

# Load version 1 of the wine quality table using Delta Lake time travel
df_v1 = spark.sql(f"SELECT * FROM {catalog_name}.default.wine_quality VERSION AS OF 1")

# Convert to a Pandas DataFrame for scikit-learn
pdf_v1 = df_v1.toPandas()

# Use the same feature columns and target as version 2 training
feature_cols = [c for c in pdf_v1.columns if c not in ["quality", "high_quality"]]
X_v1 = pdf_v1[feature_cols]
y_v1 = pdf_v1["high_quality"].astype(int)

# Split with the same random seed and ratio as version 2
from sklearn.model_selection import train_test_split
X_train_v1, X_test_v1, y_train_v1, y_test_v1 = train_test_split(
    X_v1, y_v1, test_size=0.2, random_state=42
)

print(f"Version 1 dataset: {len(pdf_v1)} total rows")
print(f"Training: {len(X_train_v1)} rows, Test: {len(X_test_v1)} rows")

Cell 2: Retrain with best hyperparameters

import mlflow
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score, classification_report

# Extract the best hyperparameters from the version 2 tuning
best_n_estimators = int(best_run["params.n_estimators"])
best_learning_rate = float(best_run["params.learning_rate"])
best_max_depth = int(best_run["params.max_depth"])

# Turn off autologging to control logging manually
mlflow.sklearn.autolog(disable=True)

# Train on version 1 data with the best hyperparameters
with mlflow.start_run(run_name="GBT_retrain_data_v1") as run:
    # Tag this run with the data version for filtering later
    mlflow.set_tag("data_version", "1")
    mlflow.log_param("n_estimators", best_n_estimators)
    mlflow.log_param("learning_rate", best_learning_rate)
    mlflow.log_param("max_depth", best_max_depth)
    mlflow.log_param("data_rows", len(pdf_v1))

    # Build the classifier with the same winning configuration
    model_v1 = GradientBoostingClassifier(
        n_estimators=best_n_estimators,
        learning_rate=best_learning_rate,
        max_depth=best_max_depth,
        random_state=42
    )
    model_v1.fit(X_train_v1, y_train_v1)

    # Evaluate on the test set
    y_prob_v1 = model_v1.predict_proba(X_test_v1)[:, 1]
    auc_v1 = roc_auc_score(y_test_v1, y_prob_v1)
    mlflow.log_metric("test_auc", auc_v1)
    mlflow.sklearn.log_model(model_v1, "model")

    # Save the run ID for later comparison
    v1_run_id = run.info.run_id

print(f"Version 1 AUC: {auc_v1:.4f}")
print(f"Version 2 Best AUC: {best_auc:.4f}")
print(f"Difference: {auc_v1 - best_auc:.4f}")
print("\nClassification Report (Version 1 Data):")
print(classification_report(y_test_v1, model_v1.predict(X_test_v1)))

Cell 3: Summary comparison

from mlflow import MlflowClient
import pandas as pd

# Tag the version 2 best run so it appears in the comparison
client = MlflowClient()
client.set_tag(best_run["run_id"], "data_version", "2")

# Get row counts from each data version using time travel
v1_rows = spark.sql(
    f"SELECT COUNT(*) as cnt FROM {catalog_name}.default.wine_quality VERSION AS OF 1"
).collect()[0]["cnt"]
v2_rows = spark.sql(
    f"SELECT COUNT(*) as cnt FROM {catalog_name}.default.wine_quality VERSION AS OF 2"
).collect()[0]["cnt"]

# Search for all runs tagged with a data version
comparison_runs = mlflow.search_runs(
    filter_string="tags.data_version IS NOT NULL",
    order_by=["metrics.test_auc DESC"]
)

# Build a summary table with the key metrics per data version
summary_data = []
for version, rows in [("1", v1_rows), ("2", v2_rows)]:
    version_runs = comparison_runs[comparison_runs["tags.data_version"] == version]
    if not version_runs.empty:
        best = version_runs.iloc[0]
        summary_data.append({
            "Data Version": version,
            "Row Count": rows,
            "Best AUC": round(best["metrics.test_auc"], 4),
            "n_estimators": best["params.n_estimators"],
            "learning_rate": best["params.learning_rate"],
            "max_depth": best["params.max_depth"]
        })

summary = pd.DataFrame(summary_data)
print("Model Performance vs. Data Version")
print("=" * 55)
display(summary)

Secret mission

Register and Serve Your Best Model

Your experiment pipeline tracks every run and every data version. But in a real ML workflow, the best model needs to move beyond your notebook. Can you register it in Unity Catalog and prove it works on new data?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely on Databricks Free Edition, so there are no ongoing costs. Decide whether to keep your resources in the workspace, or delete them entirely.

Resources you used:

  • Notebook: Wine Quality ML Pipeline in your Databricks workspace.
  • Delta Lake table: wine_quality in Unity Catalog (3 versions).
  • MLflow experiment with 20+ tracked runs (tied to the notebook).
  • Registered model: wine_quality_model in Unity Catalog (if you completed the Secret Mission).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with the notebook or want to reference your MLflow runs later.

  • Your Databricks Free Edition account has no ongoing charges. Daily usage limits reset automatically.
  • Your Delta Lake table, MLflow experiment, and notebook remain in the workspace until you delete them or your account becomes inactive.
  • Databricks may delete Free Edition accounts that are inactive for a prolonged period.

Pause - I'll come back to this later

Since Databricks Free Edition uses serverless compute, there are no running clusters to stop. Your workspace resources stay saved at no cost.

  • Simply close the browser tab. Your notebook, table, and experiment data persist in the workspace.
  • When you return, open your workspace and navigate back to the Wine Quality ML Pipeline notebook to pick up where you left off.

Delete - I don't want to use this again

Remove all project resources from your workspace. Follow these steps in order.

  • If you completed the Secret Mission, delete the registered model first. Click Catalog in the sidebar, navigate to your catalog, then the default schema, and click Models. Click on wine_quality_model.
  • On the model page, click the kebab menu (three dots) in the upper-right corner and select Delete. Click Delete to confirm.
  • Drop the Delta Lake table by adding a new cell in your notebook and running this SQL command:
DROP TABLE IF EXISTS wine_quality;
  • Delete the notebook. Click Workspace in the sidebar, find your Wine Quality ML Pipeline notebook, right-click it, and select Move to Trash. This also deletes the associated MLflow experiment and all its runs.

What about my Databricks account?

Databricks Free Edition has no ongoing costs, so keeping your account open costs nothing. There is no self-serve option to delete a Free Edition account.

If you want to remove the account entirely, email help@databricks.com and request account deletion. Alternatively, inactive Free Edition accounts are automatically deactivated and deleted after a prolonged period.

That's a Wrap!

That's a Wrap!

Nice work! You just built an end-to-end ML experimentation pipeline on Databricks Free Edition that connects data versioning to experiment tracking for fully reproducible ML workflows.

You've learned how to:

  • Load and transform data with PySpark into a Delta Lake table, create multiple versions through data mutations, and query any historical state using time travel.
  • Train and hyperparameter-tune a classification model using scikit-learn and Optuna, with every run automatically tracked in MLflow.
  • Retrain on an earlier data version and compare model performance across data states, proving that both the data and model artifacts are fully reproducible.
  • Secret Mission: Registered the best model in Unity Catalog and performed batch inference on new data.

Ready to quiz yourself?