Databricks Tutorials – Part 3 – AI/ML end-to-end lab using Serverless notebook compute

Uncategorized
Posted on | by rajeshkumar

Absolutely. Here is a notebook-first, fresh-workspace, self-contained Databricks AI/ML tutorial built around Serverless notebook compute. I based it on the current Databricks docs as of late March 2026, including serverless notebooks, the official ML quickstart, Unity Catalog model lifecycle, Model Serving, AI Playground, and the current retrieval-agent tutorial. The key fit for your setup is that serverless notebooks are the right place for Python, MLflow, training, and experiments, while sample data is already available in Databricks through samples and /databricks-datasets. (Databricks Documentation)

What this tutorial will show

This lab lets a student do the full flow in one workspace: create a notebook, load built-in sample data, write Python, train a model, track experiments with MLflow, register the model in Unity Catalog, deploy it with Mosaic AI Model Serving, then move into the GenAI side with AI Playground and a Databricks-provided retrieval-agent notebook that is explicitly described as standalone and ready to run with no setup or data required. (Databricks Documentation)

Important reality check

For this tutorial, Serverless notebook compute is enough for the Python/ML portion. Databricks says that if serverless interactive compute is enabled, all users in the workspace can attach notebooks to serverless compute with no extra permissions beyond the workspace capability itself. For model registration in Unity Catalog, the student still needs the target catalog/schema permissions: USE CATALOG, USE SCHEMA, CREATE TABLE, and CREATE MODEL. (Databricks Documentation)

Databricks AI/ML end-to-end lab using Serverless notebook compute

1. Learning outcome

By the end of this lab, the student will have created and run a real notebook-driven workflow that covers:

  • Python development in a Databricks notebook
  • MLflow experiment tracking
  • model training and comparison
  • Unity Catalog model registration
  • custom model deployment with Mosaic AI Model Serving
  • no-code agent prototyping in AI Playground
  • a code-first agent capstone using Databricks’ official retrieval-agent tutorial notebook (Databricks Documentation)

2. High-level flow

Use this mental model throughout the class:

flowchart LR
    A[Serverless Notebook] --> B[Create schema in Unity Catalog]
    B --> C[Load built-in wine sample data]
    C --> D[Train model with scikit-learn]
    D --> E[Track runs in MLflow]
    E --> F[Register best model in Unity Catalog]
    F --> G[Deploy with Mosaic AI Model Serving]
    G --> H[Test endpoint]
    A --> I[AI Playground]
    I --> J[Export agent notebook]
    J --> K[Run official retrieval-agent notebook]

This mirrors the current Databricks platform flow: notebooks for development, MLflow for tracking, Unity Catalog for governed model lifecycle, Model Serving for deployment, and AI Playground / Agent Framework for agent prototyping and deployment workflows. (Databricks Documentation)

3. Prerequisites

Before starting, verify these in your workspace:

  1. Unity Catalog is enabled.
  2. Serverless notebook compute is available.
  3. Serving appears in the left menu.
  4. Playground appears in the left menu.
  5. You can create objects in at least one catalog and schema.

Databricks documents Unity Catalog as required for serverless notebooks and for the current recommended model lifecycle in Unity Catalog. AI Playground also requires workspace access to foundation models plus Unity Catalog and Agent Framework availability. (Databricks Documentation)

Instructor note
If main is not writable in your workspace, keep the tutorial exactly the same and only change CATALOG_NAME to a catalog where the student has create permissions.

4. Create the notebook

Create a new notebook in Workspace called:

01_databricks_ai_ml_end_to_end

Attach it to Serverless from the compute drop-down. Databricks says that in serverless-enabled workspaces, new notebooks default to serverless on execution if nothing else is selected. (Databricks Documentation)

Expected output

The notebook shows Serverless as the active compute.

5. Optional notebook environment setup

Databricks now provides an Environment side panel for serverless notebooks where you can manage dependencies, memory, usage policy, and base environment. If you later hit package or memory issues, use that panel instead of guessing. The same panel can switch the notebook to higher memory if needed. (Databricks Documentation)

Instructor note
For this lab, standard serverless memory is usually enough. Only move to high memory if the notebook throws out-of-memory errors.

6. Install Python libraries

Cell 1 — Python

%pip install -U mlflow scikit-learn pandas matplotlib hyperopt
dbutils.library.restartPython()

Expected output

Package installation logs, then Python restarts.

This package set matches the official Databricks ML getting-started flow: scikit-learn for the model, MLflow for experiment tracking, and Hyperopt for automated hyperparameter tuning. (Databricks Documentation)

7. Define your catalog, schema, and model names

Cell 2 — Python

import mlflow
import pandas as pd
import sklearn.metrics
import sklearn.model_selection
import sklearn.ensemble
import matplotlib.pyplot as plt

from hyperopt import fmin, tpe, hp, SparkTrials, STATUS_OK
from hyperopt.pyll import scope

# Change only if main is not writable in your workspace
CATALOG_NAME = "main"
SCHEMA_NAME = "ai_ml_starter"
MODEL_NAME = f"{CATALOG_NAME}.{SCHEMA_NAME}.wine_quality_model"

# Explicitly use Unity Catalog for model registry
mlflow.set_registry_uri("databricks-uc")

spark.sql(f"CREATE SCHEMA IF NOT EXISTS {CATALOG_NAME}.{SCHEMA_NAME}")

print("Catalog:", CATALOG_NAME)
print("Schema:", SCHEMA_NAME)
print("Model:", MODEL_NAME)
print("Registry URI set to databricks-uc")

Expected output

Printed values for catalog, schema, and full model name.

Databricks recommends Models in Unity Catalog for governing and deploying models, and the current docs note that in MLflow 3 the default registry URI is databricks-uc, which is the Unity Catalog-backed registry. Setting it explicitly here keeps the tutorial unambiguous for students. (Databricks Documentation)

8. Load built-in sample data

Databricks provides sample datasets in two convenient places: the samples catalog for table-style data and /databricks-datasets for file-based data. The official ML tutorial uses the wine-quality dataset from /databricks-datasets, so we’ll use that exact path to keep the workflow fresh-workspace friendly and self-contained. (Databricks Documentation)

Cell 3 — Python

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

red_wine = spark.read.csv(
    "/databricks-datasets/wine-quality/winequality-red.csv",
    sep=";",
    header=True,
    inferSchema=True
)

for c in white_wine.columns:
    white_wine = white_wine.withColumnRenamed(c, c.replace(" ", "_"))

for c in red_wine.columns:
    red_wine = red_wine.withColumnRenamed(c, c.replace(" ", "_"))

display(white_wine.limit(5))
display(red_wine.limit(5))

Expected output

Two preview tables showing red and white wine records.

9. Save the sample data as your own governed tables

Cell 4 — Python

white_wine.write.mode("overwrite").saveAsTable(f"{CATALOG_NAME}.{SCHEMA_NAME}.white_wine")
red_wine.write.mode("overwrite").saveAsTable(f"{CATALOG_NAME}.{SCHEMA_NAME}.red_wine")

print(f"Created {CATALOG_NAME}.{SCHEMA_NAME}.white_wine")
print(f"Created {CATALOG_NAME}.{SCHEMA_NAME}.red_wine")

Expected output

Printed confirmation that both tables were created.

This is an important teaching step because it moves the student from “reading built-in sample files” to “working with their own managed Unity Catalog tables,” which is the right foundation for lineage, governance, and downstream model lifecycle. (Databricks Documentation)

10. Validate the data with SQL inside the same notebook

Cell 5 — SQL

SELECT quality, COUNT(*) AS cnt
FROM main.ai_ml_starter.red_wine
GROUP BY quality
ORDER BY quality;

Expected output

A result grid showing counts for each wine quality score.

Instructor note
Pause here and explain the first platform lesson: one serverless notebook can mix Python and SQL while operating on the same Unity Catalog tables.

11. Prepare the machine learning dataset

The official Databricks tutorial frames this as a classification problem: predict whether a wine is “high quality” from its chemical properties. It also combines the red and white datasets and adds an is_red feature. We’ll follow that same pattern so the student is aligned with the official learning path. (Databricks Documentation)

Cell 6 — Python

white_pdf = spark.read.table(f"{CATALOG_NAME}.{SCHEMA_NAME}.white_wine").toPandas()
red_pdf = spark.read.table(f"{CATALOG_NAME}.{SCHEMA_NAME}.red_wine").toPandas()

white_pdf["is_red"] = 0.0
red_pdf["is_red"] = 1.0

data_df = pd.concat([white_pdf, red_pdf], axis=0)

# Label: high-quality wine means quality >= 7
labels = data_df["quality"].astype("int") >= 7
features = data_df.drop(["quality"], axis=1)

X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(
    features,
    labels,
    test_size=0.2,
    random_state=1
)

print("Train shape:", X_train.shape)
print("Test shape:", X_test.shape)
print("Positive rate:", float(labels.mean()))

Expected output

Training shape, test shape, and the proportion of “high-quality” wines.

12. Train a baseline model and log it with MLflow

Databricks’ ML tutorial uses MLflow to track the development process and recommends comparing runs rather than treating the notebook as a dead-end script. That is exactly the behavior we want students to see. (Databricks Documentation)

Cell 7 — Python

mlflow.autolog()

with mlflow.start_run(run_name="gradient_boost_baseline"):
    model = sklearn.ensemble.GradientBoostingClassifier(random_state=0)
    model.fit(X_train, y_train)

    predicted_probs = model.predict_proba(X_test)
    roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:, 1])

    roc_display = sklearn.metrics.RocCurveDisplay.from_estimator(model, X_test, y_test)
    roc_display.figure_.savefig("/tmp/roc_curve.png")

    mlflow.log_metric("test_auc", roc_auc)
    mlflow.log_artifact("/tmp/roc_curve.png")

    print("Baseline test AUC:", roc_auc)

Expected output

AUC printed in the notebook, plus an MLflow run link in the output.

Instructor note
Ask the student to click the run and inspect metrics, parameters, artifacts, and the model. This is the first moment they really “see” experiments rather than just code execution.

13. Run hyperparameter tuning to create multiple experiments

The official tutorial uses Hyperopt to automate hyperparameter tuning and compare model runs. We’ll keep the sweep small so it is classroom-friendly, but still enough to visualize what experimentation looks like. (Databricks Documentation)

Cell 8 — Python

search_space = {
    "n_estimators": scope.int(hp.quniform("n_estimators", 20, 300, 1)),
    "learning_rate": hp.loguniform("learning_rate", -3, 0),
    "max_depth": scope.int(hp.quniform("max_depth", 2, 6, 1)),
}

def train_model(params):
    mlflow.autolog()
    with mlflow.start_run(nested=True):
        model_hp = sklearn.ensemble.GradientBoostingClassifier(
            random_state=0,
            **params
        )
        model_hp.fit(X_train, y_train)

        predicted_probs = model_hp.predict_proba(X_test)
        auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:, 1])

        mlflow.log_metric("test_auc", auc)
        return {"loss": -auc, "status": STATUS_OK}

spark_trials = SparkTrials(parallelism=4)

with mlflow.start_run(run_name="gb_hyperopt"):
    best_params = fmin(
        fn=train_model,
        space=search_space,
        algo=tpe.suggest,
        max_evals=8,
        trials=spark_trials
    )

print("Best parameters:", best_params)

Expected output

A printed best-parameter dictionary and several nested MLflow runs.

Cell 9 — Python

best_run = mlflow.search_runs(
    order_by=["metrics.test_auc DESC", "start_time DESC"],
    max_results=1
).iloc[0]

print("Best run_id:", best_run.run_id)
print("Best test_auc:", best_run["metrics.test_auc"])

Expected output

The best run_id and best AUC score.

Instructor note
Open the MLflow experiment UI and sort by test_auc. This is the clearest way to teach “write code → experiment code → compare results.”

14. Register the best model in Unity Catalog

Databricks recommends the Unity Catalog model registry for the full ML model lifecycle. Models there inherit centralized access control, lineage, auditing, and model discovery across workspaces. (Databricks Documentation)

Cell 10 — Python

model_uri = f"runs:/{best_run.run_id}/model"
registered_model = mlflow.register_model(model_uri, MODEL_NAME)

print("Registered model:", registered_model.name)
print("Version:", registered_model.version)

Expected output

A printed model name like main.ai_ml_starter.wine_quality_model and a version number such as 1.

What to do in the UI

Open Catalog and navigate to your catalog and schema. The model should now appear as a governed model object alongside your tables. Databricks explicitly supports models as Unity Catalog assets managed through the same governance plane. (Databricks Documentation)

15. Serve the model

Mosaic AI Model Serving is Databricks’ managed deployment layer for AI and ML models. Databricks describes it as a unified interface for deploying, governing, and querying models for real-time and batch inference, exposed as REST APIs and backed by serverless compute. (Databricks Documentation)

Step 15.1 — Create the endpoint in the UI

  1. Click Serving in the sidebar.
  2. Click Create serving endpoint.
  3. Name it:

wine-quality-endpoint

  1. In the served model section, choose your Unity Catalog model:
    main.ai_ml_starter.wine_quality_model
  2. Select the latest version.
  3. Leave the defaults unless you need special scaling.
  4. Create the endpoint.

Databricks supports creating custom model endpoints through the Serving UI, REST API, or MLflow Deployments SDK; for a student lab, the Serving UI is the simplest path. (Databricks Documentation)

Expected output

The endpoint status moves from provisioning to Ready.

Step 15.2 — Test the endpoint in the UI

Databricks says the easiest way to query a served custom model is from the Query endpoint panel in the Serving UI, using an accepted input format such as dataframe_records. (Databricks Documentation)

Use this request body:

{
  "dataframe_records": [
    {
      "fixed_acidity": 7.4,
      "volatile_acidity": 0.70,
      "citric_acid": 0.00,
      "residual_sugar": 1.9,
      "chlorides": 0.076,
      "free_sulfur_dioxide": 11.0,
      "total_sulfur_dioxide": 34.0,
      "density": 0.9978,
      "pH": 3.51,
      "sulphates": 0.56,
      "alcohol": 9.4,
      "is_red": 1.0
    }
  ]
}

Expected output

A prediction response showing the model’s classification output.

Instructor note
This is the “release and serve” moment. Make the student say out loud: “My notebook code became a model, and my model became an endpoint.”

16. Prototype an agent in AI Playground

Databricks’ current no-code GenAI tutorial uses AI Playground to compare LLMs, prototype tool-calling agents, and export the result to code. It also states that a tools-enabled model can call the built-in Unity Catalog function system.ai.python_exec to execute Python in a sandboxed environment. (Databricks Documentation)

Step 16.1 — Compare LLMs

  1. Click Playground.
  2. Ask:

What is a machine learning classifier?

  1. Add another model with the + button.
  2. Ask the same question with sync enabled.

Expected output

Two model responses side by side.

Step 16.2 — Prototype a tool-calling agent

  1. Choose a model labeled Tools enabled.
  2. Open Tools.
  3. Add the built-in Unity Catalog function:

system.ai.python_exec

  1. Ask:

Use Python to calculate the average of 9.4, 9.8, 10.0, and 10.2.

Expected output

A response showing that the model invoked the tool and returned the computed result.

Step 16.3 — Export the agent to code

In AI Playground, click:

Get code → Create agent notebook

Databricks says this action generates a Python notebook that defines the agent and deploys it to a model serving endpoint, though the docs also note that this exported notebook currently follows a legacy Model Serving–based authoring workflow and that Databricks now recommends Databricks Apps for long-term agent authoring. For a student lab, this export is still a very strong bridge from no-code to code. (Databricks Documentation)

Expected output

A new notebook or notebook folder is created in Workspace containing the generated agent code.

17. Run the official Databricks retrieval-agent notebook

For the capstone, do not invent your own RAG corpus. Databricks already provides an official Build, evaluate, and deploy a retrieval agent tutorial notebook. The docs say the example notebook contains all of the code used in the tutorial, uses a sample document corpus, and is ready to run with no setup or data required. (Databricks Documentation)

Step 17.1 — Import the notebook

Open the Databricks tutorial page for Build, evaluate, and deploy a retrieval agent and use the Open notebook in new tab or Copy link for import action from that page. (Databricks Documentation)

Step 17.2 — Run it from top to bottom

That notebook demonstrates a fuller agent workflow, including tool use, MLflow tracing, evaluation, and deployment. The same page shows examples of using Unity Catalog functions as tools, enabling mlflow.langchain.autolog() for traces, evaluating the agent with mlflow.evaluate(..., model_type="databricks-agent"), and deploying the agent to serving after registering it in Unity Catalog. (Databricks Documentation)

Expected output

The student sees an end-to-end agent path beyond the simpler AI Playground prototype: tool registration, traces, evaluation, and deployment.

Instructor note
This is the best capstone because it is official, current, and requires no custom data prep from the student.

18. What the student has visualized end to end

At the end of this lab, the student has seen the platform in a complete and realistic order:

  1. Notebook on Serverless
  2. built-in sample data
  3. own Unity Catalog tables
  4. Python model code
  5. MLflow experiment runs
  6. best-model selection
  7. Unity Catalog model registration
  8. serverless deployment with Model Serving
  9. tool-calling agent prototype in AI Playground
  10. official code-first retrieval-agent workflow (Databricks Documentation)

That is a strong, honest visualization of Databricks AI/ML capabilities in a fresh workspace.

19. Troubleshooting notes

If package management becomes messy, use the notebook’s Environment side panel instead of repeatedly reinstalling things. Databricks documents that panel as the central place for dependencies, base environment, and memory settings for serverless notebooks. (Databricks Documentation)

If a student cannot create the schema or register the model, the most likely issue is missing Unity Catalog privileges. The required permissions for the catalog and schema are explicitly listed in the Databricks ML quickstart. (Databricks Documentation)

If Model Serving deployment fails, use the endpoint’s logs and remember that the endpoint runs with the creator’s associated identity for accessing Unity Catalog resources. Databricks documents both the creation flow and the creator-identity behavior for serving endpoints. (Databricks Documentation)

If Playground is missing or the tools-enabled model list is empty, the workspace may not have access to foundation models, Unity Catalog-backed AI features, or Agent Framework in that region. Databricks lists those as the prerequisites for the no-code AI Playground / agent workflow. (Databricks Documentation)

20. Cleanup

When the class is over, delete the serving endpoint first, then remove the model and tables if you do not need them. Model Serving uses serverless infrastructure, so deleting the endpoint is the most important cost-control step. (Databricks Documentation)

Optional cleanup cell — SQL

DROP MODEL IF EXISTS main.ai_ml_starter.wine_quality_model;
DROP TABLE IF EXISTS main.ai_ml_starter.red_wine;
DROP TABLE IF EXISTS main.ai_ml_starter.white_wine;
DROP SCHEMA IF EXISTS main.ai_ml_starter;

21. One-line teaching summary

This notebook-first lab shows how Databricks lets a student start from built-in sample data, write Python in a serverless notebook, track and compare ML experiments with MLflow, govern models in Unity Catalog, deploy them with Model Serving, and then move into agent development through AI Playground and the official retrieval-agent workflow. (Databricks Documentation)

Reference

Leave a Reply