Building Deep Learning Models on Databricks

Complete course on building, training, deploying, and managing the lifecycle of deep learning models on Databricks.

Table of Contents

1. Module 1 — TensorBoard Integration in Databricks
   • 1.1 Introduction to Large-Scale Training
   • 1.2 Single-node GPU vs Distributed Clusters
   • 1.3 TorchDistributor and Data Parallelism
   • 1.4 Large-Scale Model Evaluation
   • 1.5 Experiment Tracking with MLflow
   • 1.6 Model Registration and Deployment
   • 1.7 Introduction to PyTorch and Dynamic Computation Graphs
2. Module 2 — Building a Neural Network
   • 2.1 GPU-Accelerated Training with Databricks Runtime for ML
   • 2.2 Tracking with MLflow Autologging
   • 2.3 Distributed Training
   • 2.4 Model Evaluation and Performance Visualization
   • 2.5 Registering and Managing PyTorch Models
3. Module 3 — Data, Hyperparameter, and Resource Optimization
   • 3.1 Data Loading Optimization with Delta Lake
   • 3.2 Hyperparameter Tuning with Optuna
   • 3.3 GPU and Cluster Usage Monitoring
4. Module 4 — Serving, Versioning, and Model Lifecycle
   • 4.1 Serving Deep Learning Models
   • 4.2 Integrating Served Models into Applications
   • 4.3 Versioning and Lifecycle Management with Unity Catalog

Module 1 — TensorBoard Integration in Databricks

1.1 Introduction to Large-Scale Training

Training and evaluating models at scale starts with choosing the right compute environment. For moderate workloads or experimentation phases, a GPU on a single node is often sufficient and allows for rapid iteration. However, as data and model complexity grows, distributed clusters allow for efficient scaling and reduced training time.

It is also important to use scalable training techniques such as:

• Data parallelism
• Validation monitoring
• Early stopping

These techniques ensure that, while increasing performance, model quality is maintained and overfitting is avoided.

1.2 Single-node GPU vs Distributed Clusters

flowchart LR
    A[Workload] --> B{Data size\nand complexity?}
    B -- Moderate / Experimentation --> C[Single-Node GPU]
    B -- Large scale / Production --> D[Distributed Cluster]

    C --> C1[✔ Simple to configure]
    C --> C2[✔ Low overhead]
    C --> C3[✔ Fast iteration]
    C --> C4[✔ Easy debugging]

    D --> D1[✔ Multi-node parallelization]
    D --> D2[✔ Reduced training time]
    D --> D3[✔ Production-grade]
    D --> D4[⚠ More complex setup]

1.3 TorchDistributor and Data Parallelism

TorchDistributor facilitates scaling PyTorch training across multiple GPUs or nodes in Databricks. It uses data parallelism where each worker processes a subset of the data and computes gradients locally. These gradients are then synchronized to update a single global model.

How Data Parallelism Works with TorchDistributor

flowchart TD
    DS[(Full Dataset)] --> S[Split into mini-batches]
    S --> W1[Worker 1\nModel copy]
    S --> W2[Worker 2\nModel copy]
    S --> W3[Worker N\nModel copy]

    W1 --> G1[Local gradients 1]
    W2 --> G2[Local gradients 2]
    W3 --> G3[Local gradients N]

    G1 --> SYNC[DDP Synchronization\nDistributed Data Parallel]
    G2 --> SYNC
    G3 --> SYNC

    SYNC --> UPDATE[Global model update]
    UPDATE --> W1
    UPDATE --> W2
    UPDATE --> W3

Detailed Steps

1. Model replication: Each worker starts with an identical copy of the model (same initial weights).
2. Data distribution: The dataset is split into mini-batches distributed across workers.
3. DDP synchronization: After each training step, gradients from all workers are synchronized via Distributed Data Parallel (DDP).
4. Consistent update: Each worker updates its parameters in the same way — all replicas stay aligned.

Key Advantages of TorchDistributor on Databricks

1.4 Large-Scale Model Evaluation

At scale, it is easy to assume a model is performing well simply because training is progressing. However, scaling can mask problems like overfitting or underfitting. Without adequate monitoring, different workers can learn inconsistently.

Validation Metrics for Classification Problems

┌─────────────────────────────────────────────────────────────────┐
│                  CLASSIFICATION METRICS                         │
├──────────────┬──────────────────────────────────────────────────┤
│ Accuracy     │ Overall rate of correct predictions.             │
│              │ ⚠ Misleading with imbalanced data               │
├──────────────┼──────────────────────────────────────────────────┤
│ Precision    │ Reliability of positive predictions              │
│              │ TP / (TP + FP)                                   │
├──────────────┼──────────────────────────────────────────────────┤
│ Recall       │ Ability to capture true positives                │
│              │ TP / (TP + FN)                                   │
├──────────────┼──────────────────────────────────────────────────┤
│ F1 Score     │ Precision / Recall balance                       │
│              │ 2 × (P × R) / (P + R)                           │
├──────────────┼──────────────────────────────────────────────────┤
│ AUC-ROC      │ Class separation across thresholds               │
└──────────────┴──────────────────────────────────────────────────┘

Metrics for Regression Problems

Early Stopping

Early stopping is a simple but powerful technique to prevent overfitting. Instead of training for a fixed number of epochs, validation performance is monitored and training stops when it no longer improves.

flowchart LR
    E[Epoch N] --> CHECK{Validation\nimprovement?}
    CHECK -- Yes --> CONTINUE[Continue training\nSave best model]
    CHECK -- No --> PATIENCE{Patience\nreached?}
    PATIENCE -- No, wait --> E
    PATIENCE -- Yes --> STOP[Stop training\nRestore best model]
    CONTINUE --> E

Early stopping parameters:

1.5 Experiment Tracking with MLflow

Experiment tracking is the practice of recording everything that happens during a training run (parameters, metrics, outputs) to compare experiments, reproduce results, and improve models over time.

MLflow Autologging

MLflow autologging removes the need to manually track experiments by automatically capturing key details during training.

# Enable autologging in a single line
import mlflow
mlflow.pytorch.autolog()

# From this point, everything is logged automatically:
# - Parameters: learning rate, epochs, batch size
# - Metrics: loss, accuracy, validation scores
# - Artifacts: model weights, TensorBoard event logs, library versions

What Is Automatically Logged

mindmap
  root((MLflow Autologging))
    Parameters
      Learning rate
      Number of epochs
      Batch size
      Model architecture
    Metrics
      Training loss
      Validation loss
      Accuracy
      RMSE / MAE
    Artifacts
      Model weights
      TensorBoard event logs
      Python environment
      Dependencies
    Metadata
      Run ID
      Timestamps
      User tags

MLflow UI — Run Comparison

In Databricks, MLflow is natively integrated. Here is what you can do in the interface:

• Compare metrics: loss, accuracy, RMSE across different runs
• Visualize curves: observe metric evolution over time (convergence, fluctuations, divergence)
• Compare hyperparameters: identify which configuration produced the best result
• Inspect artifacts: examine model weights, TensorBoard graphs
• Tag and filter: organize runs by version, dataset, or configuration

1.6 Model Registration and Deployment

Demo: Model Registration and Deployment (Module 1)

Step 1 — Import Libraries

import mlflow
import mlflow.pytorch
import torch
import torch.nn as nn
import torch.optim as optim
import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

Step 2 — Load and Preprocess the Dataset

# Load dataset
X, y = load_diabetes(return_X_y=True)

# Train/test split + normalization
X_train, X_test, y_train, y_test = train_test_split(X, y)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)

# Convert to PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).view(-1, 1)

Step 3 — Define the Neural Network

# Simple neural network: 10 features → hidden layer → 1 output (regression)
model = nn.Sequential(
    nn.Linear(10, 8),   # reduced hidden layer
    nn.ReLU(),
    nn.Linear(8, 1)
)

Step 4 — Training and Logging in MLflow

# Train with Adam and MSE Loss
optimizer = optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

with mlflow.start_run() as run:
    for _ in range(20):
        preds = model(X_train)
        loss = loss_fn(preds, y_train)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Log model with input example (for signature inference)
    mlflow.pytorch.log_model(model, name="model", input_example=X[:1].astype("float32"))
    run_id = run.info.run_id

print("Run ID:", run_id)

Step 5 — Register in the Model Registry (Unity Catalog)

model_uri = f"runs:/{run_id}/model"
model_name = "workspace.ml_models.diabetes_predictor"

mlflow.register_model(model_uri, model_name)

The model is now versioned and governed in Unity Catalog. Each registration automatically creates a new version.

1.7 Introduction to PyTorch and Dynamic Computation Graphs

PyTorch is a popular open-source deep learning framework that allows building and training neural networks with flexibility. In Databricks, it runs on scalable clusters, integrates with distributed data pipelines, and supports collaborative development.

Dynamic Computation Graph

flowchart LR
    subgraph Static["Static Framework (e.g., TF v1)"]
        S1[Define the full graph\nin advance] --> S2[Compile] --> S3[Execute]
    end

    subgraph Dynamic["PyTorch - Dynamic Graph"]
        D1[Execute Python code\nnormally] --> D2[Graph is built\non the fly] --> D3[Debug with print()\nand standard breakpoints]
    end

Advantages of PyTorch’s dynamic computation graph:

• No prior definition: the graph is built step by step as data flows through
• Native Python constructs: loops and conditions can be used directly in the model
• On-the-fly modification: architecture can be changed at any time without rebuilding the entire graph
• Standard debugging: inspect variables, print outputs, debug like ordinary Python

Three Fundamental PyTorch Abstractions

classDiagram
    class Tensor {
        +dtype: float32/int64/...
        +device: cpu/cuda
        +grad: Tensor
        Multi-dimensional data structure
        Stores inputs, parameters, outputs
        Can run on CPU or GPU
    }

    class Autograd {
        +requires_grad: bool
        +backward()
        +grad_fn: Function
        Automatic differentiation
        Dynamically traces operations
        Automatically computes gradients
    }

    class nn_Module {
        +parameters()
        +forward()
        +state_dict()
        Base class for all models
        Encapsulates layers and parameters
        Defines the forward pass
    }

    Tensor --> Autograd : tracks operations
    Autograd --> nn_Module : provides gradients
    nn_Module --> Tensor : produces tensors

Module 2 — Building a Neural Network

2.1 GPU-Accelerated Training with Databricks Runtime for ML

Databricks Runtime for ML provides a fully pre-configured environment with:

• Frameworks: PyTorch, TensorFlow, scikit-learn
• Built-in GPU acceleration with optimized CUDA drivers
• Native integration with notebooks, jobs, and MLflow
• Transparent scalability from single node to distributed training

flowchart TD
    subgraph DBML["Databricks Runtime for ML"]
        PY[PyTorch / TensorFlow / scikit-learn]
        CUDA[Optimized CUDA Drivers]
        MLF[Integrated MLflow]
        SPARK[Apache Spark]
    end

    subgraph ENV["Development Environment"]
        NB[Interactive notebooks]
        JOBS[Scheduled jobs]
        COLLAB[Team collaboration]
    end

    DBML --> ENV
    CUDA --> GPU[GPU Cluster]
    SPARK --> DIST[Distributed training]

Trade-off: Batch Size and Latency

When grouping requests into batches, hardware utilization and overall throughput improve, but individual requests wait longer before being executed.

2.2 Tracking with MLflow Autologging

Enabling Autologging

import mlflow
import mlflow.pytorch

# Enable autologging BEFORE the training loop
mlflow.pytorch.autolog()

# From this call, MLflow monitors everything during training:
# - Training parameters
# - Metrics epoch by epoch
# - Intermediate checkpoints
# - Final model

Complete Workflow with Autologging

import mlflow
import mlflow.pytorch
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# Data preparation
X, y = load_diabetes(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).view(-1, 1)

# Model
model = nn.Sequential(
    nn.Linear(10, 8),
    nn.ReLU(),
    nn.Linear(8, 1)
)

# Enable autologging
mlflow.pytorch.autolog()

# Training — everything is logged automatically
optimizer = optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

with mlflow.start_run() as run:
    for epoch in range(20):
        preds = model(X_train)
        loss = loss_fn(preds, y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    run_id = run.info.run_id

What MLflow Autologging Records for PyTorch

┌─────────────────────────────────────────────────────────────────────┐
│                    MLflow Autologging — PyTorch                     │
├───────────────────────┬─────────────────────────────────────────────┤
│ PARAMETERS            │ learning_rate, optimizer, loss_fn, epochs   │
│                       │ batch_size, architecture                    │
├───────────────────────┼─────────────────────────────────────────────┤
│ METRICS               │ train_loss (per epoch)                      │
│                       │ val_loss, val_accuracy (if defined)         │
├───────────────────────┼─────────────────────────────────────────────┤
│ ARTIFACTS             │ model.pth — model weights                   │
│                       │ TensorBoard event logs                      │
│                       │ requirements.txt — dependencies             │
│                       │ conda.yaml — reproducible environment       │
├───────────────────────┼─────────────────────────────────────────────┤
│ METADATA              │ run_id, start_time, end_time, status        │
└───────────────────────┴─────────────────────────────────────────────┘

2.3 Distributed Training

Databricks offers several tools to scale PyTorch training beyond a single machine.

Comparison of Distributed Training Tools

quadrantChart
    title Distributed Training Tools — Complexity vs Capacity
    x-axis Low complexity --> High complexity
    y-axis Standard capacity --> Maximum capacity
    quadrant-1 Powerful and complex
    quadrant-2 Simple and high-performing
    quadrant-3 Limited and simple
    quadrant-4 Complex but limited
    TorchDistributor: [0.2, 0.5]
    DeepSpeed: [0.7, 0.9]
    Ray: [0.8, 0.75]

TorchDistributor — Architecture

flowchart TD
    DRIVER[Driver Node\nTorchDistributor.run\(\)] --> |Launch workers| W1[Worker 1\nGPU 0]
    DRIVER --> |Launch workers| W2[Worker 2\nGPU 1]
    DRIVER --> |Launch workers| W3[Worker N\nGPU N]

    W1 <-->|DDP Gradient Sync| W2
    W2 <-->|DDP Gradient Sync| W3
    W1 <-->|DDP Gradient Sync| W3

    W1 --> M[Unified global model]
    W2 --> M
    W3 --> M

Considerations for Distributed Training

Moving from a single GPU to distributed training requires efficient data transfer:

• Inter-GPU communication: bandwidth (NVLink, InfiniBand) is critical
• Data consistency: ensure each worker sees a different and coherent subset
• Error handling: a failing worker must not corrupt the overall training
• Linear scalability: ideally, doubling GPUs should halve training time

2.4 Model Evaluation and Performance Visualization

Key Metrics to Monitor

flowchart LR
    EVAL[Model evaluation] --> ACC[Accuracy\nQuick overview]
    EVAL --> LOSS[Loss Curves\nConvergence / Overfitting]
    EVAL --> CONF[Confusion Matrix\nClassification error detail]

    LOSS --> OV[Overfitting:\ntrain_loss ↓↓\nval_loss ↑]
    LOSS --> UN[Underfitting:\ntrain_loss stays high\nval_loss stays high]
    LOSS --> OK[Convergence:\ntrain_loss ≈ val_loss ↓]

Diagnosis via Learning Curves

Diagnosing training/validation loss curves:

  Overfitting                Underfitting              Good convergence

  loss                       loss                      loss
   │ train ╲                  │                          │ train╲
   │        ╲                 │ train╲                   │       ╲___
   │         ╲___             │       ╲___               │ val    ╲__
   │ val   /‾‾‾‾‾             │ val    ╲___              │
   └──────────── epochs       └──────────── epochs       └──────────── epochs

   train_loss ↓↓             train_loss stays high       train_loss ≈ val_loss
   val_loss ↑                val_loss stays high         both converge

Confusion Matrix

A confusion matrix provides a detailed breakdown of prediction results across all classes:

                    Prediction
                    Positive  Negative
Actual  Positive  │    TP    │    FN    │  ← Recall = TP/(TP+FN)
        Negative  │    FP    │    TN    │
                  └──────────┴──────────┘
                       ↑
                  Precision = TP/(TP+FP)

2.5 Registering and Managing PyTorch Models

Complete Demo: Model Lifecycle (Module 2)

This demo covers the full MLOps workflow: training → logging → registration → serving → batch inference.

Steps 1-5: Same as Module 1 (see section 1.6)

Step 6 — Batch Inference with Spark

from pyspark.sql import SparkSession
import pandas as pd
import numpy as np

spark = SparkSession.builder.getOrCreate()

# Prepare test data (same preprocessing as training!)
X_test_scaled = scaler.transform(X_test).astype(np.float32)

# Convert to Pandas then to Spark DataFrame
pdf = pd.DataFrame(X_test_scaled)
spark_df = spark.createDataFrame(pdf)

Step 7 — Load Model from Registry as Spark UDF

import mlflow.pyfunc

# Reference the versioned model from Unity Catalog
model_uri = f"models:/{model_name}/1"

predict_udf = mlflow.pyfunc.spark_udf(
    spark,
    model_uri=model_uri,
    result_type="double"
)

By loading the model as a Spark UDF, we guarantee use of a governed, versioned, and reproducible artifact. Spark automatically distributes inference across all cluster nodes.

Step 8 — Execute Batch Inference

from pyspark.sql.functions import struct

# Apply model to all rows of the distributed DataFrame
predictions_df = spark_df.withColumn(
    "prediction",
    predict_udf(struct(*spark_df.columns))
)

display(predictions_df.limit(10))

Complete MLOps Workflow Architecture

flowchart LR
    A[Raw data] --> B[Preprocessing\nStandardScaler]
    B --> C[Model definition\nnn.Sequential]
    C --> D[Training\nAdam + MSELoss]
    D --> E[MLflow logging\nmlflow.pytorch.log_model]
    E --> F[Model Registry\nUnity Catalog]
    F --> G1[Real-time Serving\nMosaic AI REST Endpoint]
    F --> G2[Batch Inference\nSpark UDF on cluster]

    style A fill:#e8f5e9
    style F fill:#fff3e0
    style G1 fill:#e3f2fd
    style G2 fill:#e3f2fd

Module 3 — Data, Hyperparameter, and Resource Optimization

3.1 Data Loading Optimization with Delta Lake

Delta Lake — Transactional Layer on Parquet

Delta Lake adds a transactional layer to your data, enabling versioned access to datasets. It allows you to:

• Track changes and reproduce experiments
• Perform time travel to previous versions
• Optimize reading by processing only incremental changes

Delta Lake Transaction Log

flowchart LR
    subgraph Storage["Storage"]
        P1[Parquet File 1]
        P2[Parquet File 2]
        P3[Parquet File 3]
        TL[Transaction Log\n_delta_log/]
    end

    subgraph Ops["Recorded Operations"]
        I[INSERT]
        U[UPDATE]
        D[DELETE]
    end

    Ops --> TL
    TL --> |Determines which files to read| P1
    TL --> |Determines which files to read| P2

    subgraph Benefits["Advantages"]
        B1[✔ Consistent multi-worker reads]
        B2[✔ Time travel — access previous versions]
        B3[✔ Selective scan — only necessary files]
        B4[✔ ML experiment reproducibility]
    end

Comparison: Parquet vs Delta Lake

Mosaic Streaming for Large Datasets

Mosaic Streaming (MosaicML) complements Delta Lake for deep learning training by enabling efficient streaming of data directly from object storage (S3, ADLS) to GPUs, without loading everything into memory.

Optimized data pipeline:

  Delta Lake       Mosaic Streaming        Training
  ┌──────────┐     ┌──────────────────┐    ┌──────────────┐
  │ Parquet  │────▶│ Parallel shard   │───▶│  GPU Worker  │
  │ + TX Log │     │ streaming        │    │  Batch N     │
  └──────────┘     └──────────────────┘    └──────────────┘
                         │
                    Prefetching
                    + Local cache

3.2 Hyperparameter Tuning with Optuna

What is Hyperparameter Tuning?

Hyperparameter tuning is the process of finding the best configuration for your model to achieve optimal performance. Instead of manually guessing values like learning rate, batch size, or number of layers, an search space is defined and automated methods explore it efficiently.

Optuna + Databricks

flowchart TD
    SEARCH[Search space\nlearning_rate, batch_size, layers, ...] --> OPTUNA[Optuna\nIntelligent search TPE / CMA-ES]
    OPTUNA --> T1[Trial 1\nlr=0.001, bs=32]
    OPTUNA --> T2[Trial 2\nlr=0.01, bs=64]
    OPTUNA --> T3[Trial N\nlr=0.1, bs=128]

    T1 --> R1[Result 1\nval_loss=0.45]
    T2 --> R2[Result 2\nval_loss=0.31]
    T3 --> R3[Result N\nval_loss=0.52]

    R1 --> OPTUNA
    R2 --> OPTUNA
    R3 --> OPTUNA

    OPTUNA --> BEST[Best hyperparameters\nval_loss=0.31, lr=0.01, bs=64]

    style BEST fill:#c8e6c9

Code Example — Optuna in Databricks

import optuna
import torch
import torch.nn as nn
import torch.optim as optim
import mlflow

def objective(trial):
    """Objective function to minimize for each Optuna trial."""
    # Define the search space
    lr = trial.suggest_float("lr", 1e-4, 1e-1, log=True)
    batch_size = trial.suggest_categorical("batch_size", [32, 64, 128, 256])
    hidden_size = trial.suggest_int("hidden_size", 16, 128)
    n_epochs = trial.suggest_int("n_epochs", 10, 50)

    # Build model with trial hyperparameters
    model = nn.Sequential(
        nn.Linear(10, hidden_size),
        nn.ReLU(),
        nn.Linear(hidden_size, 1)
    )

    optimizer = optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

    # Training
    with mlflow.start_run(nested=True):
        mlflow.log_params(trial.params)
        for epoch in range(n_epochs):
            # ... training loop ...
            pass

        # Return metric to minimize
        val_loss = evaluate(model)
        mlflow.log_metric("val_loss", val_loss)

    return val_loss

# Launch Optuna study
study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=50)

print("Best hyperparameters:", study.best_params)
print("Best val_loss:", study.best_value)

Optuna Search Strategies

Advantages of Combining Optuna + Databricks

• Parallel trials: Databricks can run multiple trials simultaneously on different nodes
• Integrated MLflow tracking: each trial is automatically logged as an MLflow child run
• Pruning: Optuna can stop unpromising trials early (Median Pruner, Hyperband)
• Scalability: Databricks clusters dynamically adjust based on load

3.3 GPU and Cluster Usage Monitoring

Demo: Real-time GPU/CPU Monitoring (Module 3)

Step 1 — Setup and GPU Detection

import torch
import time
import psutil

try:
    import pynvml
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    gpu_available = True
except Exception:
    gpu_available = False
    handle = None
    print("No GPU detected — running on CPU")

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Device:", device)

Step 2 — Synthetic Data Generation

# Synthetic dataset to isolate system performance observation
data = torch.randn(10000, 100)
labels = torch.randint(0, 2, (10000,))

Step 3 — Simple Model Definition

# Simple architecture to focus on system metrics
model = torch.nn.Linear(100, 2).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()

Step 4 — Metrics Collection Function

def get_metrics():
    """Collects real-time system metrics."""
    metrics = {"cpu": psutil.cpu_percent()}

    if gpu_available:
        try:
            gpu = pynvml.nvmlDeviceGetUtilizationRates(handle)
            mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
            metrics["gpu_util"] = gpu.gpu
            metrics["gpu_mem_MB"] = mem.used / 1024**2
        except Exception:
            metrics["gpu_util"] = "N/A"
            metrics["gpu_mem_MB"] = "N/A"
    else:
        metrics["gpu_util"] = "N/A"
        metrics["gpu_mem_MB"] = "N/A"

    return metrics

Step 5 — Training Loop with Monitoring

BATCH_SIZE = 256

for i in range(0, len(data), BATCH_SIZE):
    start = time.time()

    x = data[i:i+BATCH_SIZE].to(device)
    y = labels[i:i+BATCH_SIZE].to(device)

    optimizer.zero_grad()
    out = model(x)
    loss = loss_fn(out, y)
    loss.backward()
    optimizer.step()

    step_time = time.time() - start
    throughput = len(x) / step_time   # samples/sec

    m = get_metrics()

    print(f"""
Step {i//BATCH_SIZE}
Loss: {loss.item():.3f}
Throughput: {throughput:.0f} samples/sec
GPU: {m['gpu_util']}% | GPU Mem: {m['gpu_mem_MB']} MB | CPU: {m['cpu']}%
""")

Signals to Monitor During Training

flowchart LR
    subgraph Signals["Monitoring Metrics"]
        GPU_UTIL[GPU Utilization\ngpu_util %]
        GPU_MEM[GPU Memory\ngpu_mem_MB]
        CPU[CPU Utilization\ncpu %]
        THROUGHPUT[Throughput\nsamples/sec]
    end

    subgraph Diagnosis
        GPU_UTIL --> |GPU < 50%| BOTTLENECK_DATA[Data bottleneck\nDataLoader too slow]
        GPU_UTIL --> |GPU > 90%| COMPUTE_BOUND[Compute-bound\nGood sign!]
        GPU_MEM --> |OOM Error| REDUCE_BATCH[Reduce batch_size\nor gradient checkpointing]
        CPU --> |CPU > 95%| CPU_BOUND[CPU-bound\nData preprocessing]
        THROUGHPUT --> |Variable throughput| UNEVEN[Uneven batches\nor irregular I/O]
    end

Common Diagnostics and Solutions

Module 4 — Serving, Versioning, and Model Lifecycle

4.1 Serving Deep Learning Models

Once a model is trained, the next step is deployment so it can make predictions on new data — moving from experimentation to production where the model continuously delivers business value.

Two Main Serving Modes

flowchart TD
    DEPLOY[Deployed model] --> RT[Real-time Inference]
    DEPLOY --> BATCH[Batch Inference]

    RT --> RT1[Instant predictions\nvia REST API]
    RT --> RT2[Use cases: fraud detection,\nrecommendations, chatbots]
    RT --> RT3[Mosaic AI Model Serving\nHTTP endpoint]

    BATCH --> B1[Processing large volumes\nat scheduled intervals]
    BATCH --> B2[Use cases: scoring\nmillions of records]
    BATCH --> B3[Apache Spark + Spark UDF]

    style RT fill:#e3f2fd
    style BATCH fill:#f3e5f5

Important Serving Terminology

Decision: Real-time vs Batch Inference

flowchart LR
    Q1{Immediate response\nrequired?} -- Yes --> Q2{Volume < a few\nthousand/sec?}
    Q1 -- No --> BATCH[Batch Inference\nSpark]

    Q2 -- Yes --> RT[Real-time Serving\nMosaic AI Endpoint]
    Q2 -- No --> Q3{GPU budget\navailable?}
    Q3 -- Yes --> RT
    Q3 -- No --> BATCH

4.2 Integrating Served Models into Applications

Calling a Mosaic AI Endpoint from an External Application

Once deployed, any external system can call the model via HTTP:

import requests
import json

# Endpoint configuration
ENDPOINT_URL = "https://<databricks-workspace>.cloud.databricks.com/serving-endpoints/<endpoint-name>/invocations"
TOKEN = "<databricks-token>"

# Prepare input data in JSON format
input_data = {
    "inputs": [[0.5, -0.3, 1.2, 0.8, -0.1, 0.4, -0.7, 0.2, 0.9, -0.5]]
}

# Send POST request
response = requests.post(
    ENDPOINT_URL,
    headers={
        "Authorization": f"Bearer {TOKEN}",
        "Content-Type": "application/json"
    },
    data=json.dumps(input_data)
)

# Retrieve prediction
prediction = response.json()
print("Prediction:", prediction)

Recommended Integration Patterns

flowchart TD
    APP[Application / Service] --> FS[Feature Store\nCentralized]
    FS --> |Consistent features| MODEL[Deployed model\nMosaic AI Endpoint]
    MODEL --> |Predictions| IT[Inference Table\nPrediction storage]
    IT --> MONITOR[Monitoring\nData drift / Model drift]
    IT --> DOWNSTREAM[Downstream systems\nDashboards, alerts, actions]

    subgraph Advantages
        A1[✔ No training-serving skew]
        A2[✔ Reusable features]
        A3[✔ Full traceability]
        A4[✔ Early drift detection]
    end

4.3 Versioning and Lifecycle Management with Unity Catalog

Unity Catalog — Centralized Model Governance

Unity Catalog serves as a centralized governance layer, allowing teams to register, discover, and manage models across the organization in a consistent and controlled way.

flowchart LR
    subgraph UC["Unity Catalog — Model Registry"]
        M1[Version 1\nDev]
        M2[Version 2\nStaging]
        M3[Version 3\nProduction]
        M4[Version 4\nArchived]

        M1 -->|Promotion after validation| M2
        M2 -->|Integration tests OK| M3
        M3 -->|New version available| M4
    end

    subgraph Alias
        PROD_ALIAS[Alias: 'production'\n→ points to Version 3]
        STAGING_ALIAS[Alias: 'staging'\n→ points to Version 2]
        LATEST_ALIAS[Alias: 'latest'\n→ points to Version 4]
    end

    subgraph RBAC
        ADMIN[Admin\nRead / Write / Deploy]
        DS[Data Scientist\nRead / Write]
        ANALYST[Analyst\nRead only]
    end

Key Unity Catalog Features for Models

Model Promotion Workflow

stateDiagram-v2
    [*] --> Development : New model trained
    Development --> Staging : Metric validation
    Staging --> Production : Integration tests + approval
    Production --> Archived : Replaced by a newer version

    Production --> Staging : Rollback (degradation detected)
    Staging --> Development : Rollback (failed tests)

    note right of Production
        Continuous monitoring:
        - Accuracy
        - Data drift
        - Latency
    end note

Rollback Strategy

Rollback means switching your production model back to a previously stable version. It is your safety net when a new deployment introduces unexpected problems.

Signals triggering a rollback:

• Dropping accuracy on production data
• Detection of data drift (changing data distribution)
• Unusual predictions or anomalies detected
• Business metric degradation (e.g., missed fraud rate)

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

# Find the previous stable version
model_name = "workspace.ml_models.diabetes_predictor"

# List all versions
versions = client.search_model_versions(f"name='{model_name}'")
for v in versions:
    print(f"Version {v.version} — Status: {v.status} — Tags: {v.tags}")

# Rollback: assign the 'production' alias to a previous stable version
client.set_registered_model_alias(
    name=model_name,
    alias="production",
    version="2"  # roll back to version 2
)

print("Rollback complete: the 'production' alias now points to version 2")

Best Practices — Versioning and Lifecycle

┌─────────────────────────────────────────────────────────────────────┐
│              BEST PRACTICES — MODEL LIFECYCLE                       │
├─────────────────────────────────────────────────────────────────────┤
│ ✔ Always store model artifacts, dependencies, and configs           │
│   → Guaranteed reproducibility for any rollback                     │
├─────────────────────────────────────────────────────────────────────┤
│ ✔ Use aliases rather than version numbers                           │
│   → Application code does not change during updates                 │
├─────────────────────────────────────────────────────────────────────┤
│ ✔ Define clear stages and require validations                       │
│   → No promotion without tests + approval                           │
├─────────────────────────────────────────────────────────────────────┤
│ ✔ Continuously monitor production models                            │
│   → Detect data drift, accuracy degradation, anomalies              │
├─────────────────────────────────────────────────────────────────────┤
│ ✔ Maintain a clean version history in Unity Catalog                 │
│   → Faster and more reliable rollback decisions                     │
├─────────────────────────────────────────────────────────────────────┤
│ ✔ Test rollbacks regularly in staging                               │
│   → Ensure the process works before you need it                     │
└─────────────────────────────────────────────────────────────────────┘

General Summary

flowchart TD
    subgraph M1["Module 1 — Foundations"]
        M1A[PyTorch + Dynamic graphs]
        M1B[TorchDistributor + DDP]
        M1C[MLflow Experiment Tracking]
        M1D[Early Stopping + Validation Metrics]
    end

    subgraph M2["Module 2 — Model Building"]
        M2A[Databricks Runtime for ML\nGPU Acceleration]
        M2B[MLflow Autologging]
        M2C[Distributed Training\nTorchDistributor, DeepSpeed, Ray]
        M2D[Evaluation: Loss curves, Confusion Matrix]
        M2E[Batch Inference with Spark UDF]
    end

    subgraph M3["Module 3 — Optimization"]
        M3A[Delta Lake + Mosaic Streaming]
        M3B[Hyperparameter Tuning with Optuna]
        M3C[GPU/CPU Monitoring with pynvml + psutil]
    end

    subgraph M4["Module 4 — Deployment & Governance"]
        M4A[Real-time Serving via REST API]
        M4B[Large-scale Batch Inference]
        M4C[Unity Catalog — Versioning + RBAC]
        M4D[Rollback and Lifecycle Management]
    end

    M1 --> M2 --> M3 --> M4

    style M1 fill:#e8f5e9
    style M2 fill:#e3f2fd
    style M3 fill:#fff3e0
    style M4 fill:#fce4ec

Technology Stack Overview

Course: Building Deep Learning Models on Databricks — 4 modules — ~84 minutes

Search Terms

deep · models · databricks · azure · spark · data · engineering · analytics · model · mlflow · autologging · optuna · batch · catalog · delta · distributed · gpu · lake · lifecycle · metrics · pytorch · serving · torchdistributor · unity