Advanced

Machine Learning with Azure Databricks

The Databricks ML lifecycle: MLflow tracking, tuning with Ray, the model registry, serving and AutoML.

Level: Intermediate / Advanced | Prerequisites: Databricks Premium Workspace

Table of Contents

  1. Why Databricks for Machine Learning?
  2. Databricks ML Runtime
  3. MLflow — ML Lifecycle Management Platform
  4. Exploratory Data Analysis (EDA)
  5. Data Preprocessing and Preparation
  6. Training Models with scikit-learn
  7. MLflow Tracking — Experiment Tracking
  8. MLflow Autologging
  9. Hyperparameter Tuning with Ray Tune
  10. Model Registry with Unity Catalog
  11. Model Serving — Deploying as REST API
  12. Predictions from a Deployed Model
  13. Automated Model Retraining
  14. Databricks AutoML
  15. ML Orchestration with Azure Data Factory
  16. Databricks Feature Store
  17. Distributed Machine Learning with Spark MLlib
  18. MLOps Best Practices on Databricks
  19. Summary and Tool Comparison
  20. Glossary

1. Why Databricks for Machine Learning?

1.1 Platform Overview

Azure Databricks offers a unified platform for the entire machine learning lifecycle:

graph TB
    subgraph "Data"
        DL[Delta Lake\nReliable Storage]
        UC[Unity Catalog\nGovernance]
        FS[Feature Store\nReusable Features]
    end
    
    subgraph "ML Development"
        EDA[EDA / Exploration]
        Prep[Data Preparation\nscikit-learn, PySpark]
        Train[Training\nscikit-learn, TF, PyTorch, XGBoost]
        Eval[Evaluation\ncross-validation, metrics]
    end
    
    subgraph "Lifecycle Management"
        MLflow[MLflow Tracking\nExperiments + Runs]
        Reg[Model Registry\nVersioning + Staging]
        HPT[Hyperparameter Tuning\nRay Tune / Hyperopt]
        AutoML2[AutoML\nFull Automation]
    end
    
    subgraph "Deployment"
        Serving[Model Serving\nReal-time REST API]
        Batch[Batch Inference\nDatabricks Jobs]
        Workflow[Automated Workflows]
    end
    
    DL & UC & FS --> EDA & Prep
    Prep --> Train
    Train --> Eval
    Eval --> MLflow
    MLflow --> Reg
    Reg --> Serving & Batch
    HPT --> Train
    AutoML2 --> Reg
    Workflow --> Train

1.2 Databricks Advantages for ML

CapabilityDescriptionBenefit
Distributed ScaleApache Spark + Distributed MLTrain on millions of rows
Built-in Librariesscikit-learn, TF, PyTorch, XGBoost, MLlibNo complex installation
Native MLflowTracking, Registry, Serving integratedFull lifecycle on one platform
AutoMLNo-code/low-code with automatic HPTRapid prototyping
Feature StoreCentralized feature repositoryReusability and consistency
Delta LakeData versioningModel reproducibility
GPU SupportCUDA, cuDNN via ML Runtime GPUHigh-performance Deep Learning
Unity CatalogModel governanceCompliance and auditability

1.3 MLflow vs Competitors

ToolTypeStrengthsLimitations
MLflowOpen SourceNative Databricks integration, multi-frameworkLess rich UI
Amazon SageMakerCloud (AWS)Managed, full AWS integrationAWS lock-in
Google Vertex AICloud (GCP)Advanced AutoML, BigQuery integrationGCP lock-in
Azure Machine LearningCloud (Azure)Native Azure integrationSeparate from Databricks
Weights & BiasesSaaSRich visualizations, collaborationCostly for large teams
KubeflowOpen Source (K8s)Portable, Kubernetes nativeOperational complexity
ClearMLOpen SourceFull MLOps, orchestrationLess mature

2. Databricks ML Runtime

2.1 What is the Databricks ML Runtime?

The Databricks ML Runtime is a pre-configured and optimized environment for machine learning:

graph LR
    subgraph "Databricks ML Runtime"
        subgraph "DL Frameworks"
            TF[TensorFlow\n2.x]
            PT[PyTorch\n2.x]
        end
        
        subgraph "ML Libraries"
            SK[scikit-learn]
            XGB[XGBoost]
            LGB[LightGBM]
            SP[Spark MLlib]
        end
        
        subgraph "ML Infrastructure"
            MLF[MLflow\nautomatic]
            RAY[Ray Tune\nDistributed HPT]
            DASK[Dask\nDistributed pandas]
        end
        
        subgraph "Optimizations"
            CUDA[CUDA/cuDNN\nif GPU cluster]
            RAPIDS[RAPIDS\nif GPU cluster]
            Arrow[Apache Arrow\nPandas ↔ Spark]
        end
    end

2.2 Creating an ML Runtime Cluster

// ML Runtime cluster configuration
{
  "cluster_name": "data-ml-cluster",
  "spark_version": "16.4.x-cpu-ml-scala2.12",
  "node_type_id": "Standard_DS4_v2",
  "driver_node_type_id": "Standard_DS4_v2",
  "num_workers": 0,  // Single node for dev
  "autotermination_minutes": 30,
  "spark_conf": {
    "spark.databricks.delta.preview.enabled": "true"
  },
  "custom_tags": {
    "team": "data-science",
    "project": "sales-ml"
  }
}

Version Naming Convention:

  • 16.4.x-cpu-ml-scala2.12: DBR 16.4, CPU, ML Runtime
  • 16.4.x-gpu-ml-scala2.12: DBR 16.4, GPU, ML Runtime
  • 16.4.x-scala2.12: DBR 16.4 standard (without ML libraries)

3. MLflow — ML Lifecycle Management Platform

3.1 MLflow Architecture

graph TB
    subgraph "MLflow Components"
        Track[Tracking Server\n• Experiments\n• Runs\n• Params/Metrics\n• Artifacts]
        Registry[Model Registry\n• Versions\n• Stages\n• Annotations\n• Tags]
        Models[Models\n• Standard Format\n• Signatures\n• Environments]
        Serve[Serving\n• REST Endpoints\n• Batch Inference\n• A/B Testing]
    end
    
    subgraph "ML Workflow"
        Code[Notebook Code] -->|mlflow.start_run| Track
        Code -->|mlflow.log_model| Registry
        Registry -->|Deploy| Serve
    end
    
    subgraph "Storage"
        MetadataDB["(Metadata DB\nExperiments, Runs)"]
        ArtifactStore["(Artifact Store\nModels, Files)"]
    end
    
    Track --> MetadataDB
    Track --> ArtifactStore
    Registry --> ArtifactStore

3.2 Core MLflow Concepts

import mlflow
import mlflow.sklearn
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
import numpy as np

# ── Experiment ────────────────────────────────────────────────
# An experiment groups runs related to the same ML problem
mlflow.set_experiment("/Users/user@company.com/sales_price_prediction")

# ── Run ───────────────────────────────────────────────────────
# A run = one training execution with its parameters/metrics
with mlflow.start_run(run_name="linear_regression_v1") as run:
    
    # Train the model
    model = LinearRegression()
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    
    # Metrics
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    r2 = r2_score(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)
    
    # ── Log parameters ────────────────────────────────────────
    mlflow.log_param("model_type", "LinearRegression")
    mlflow.log_param("training_rows", len(X_train))
    mlflow.log_param("n_features", X_train.shape[1])
    
    # ── Log metrics ───────────────────────────────────────────
    mlflow.log_metric("rmse", rmse)
    mlflow.log_metric("r2_score", r2)
    mlflow.log_metric("mae", mae)
    
    # ── Log artifacts ─────────────────────────────────────────
    # Save a residuals plot
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.scatter(y_test, y_pred - y_test, alpha=0.5)
    ax.axhline(y=0, color='r', linestyle='--')
    ax.set_xlabel("Actual Values")
    ax.set_ylabel("Residuals")
    ax.set_title("Residual Analysis - Linear Regression")
    mlflow.log_figure(fig, "residuals_plot.png")
    plt.close()
    
    # ── Infer and log model signature ─────────────────────────
    from mlflow.models import infer_signature
    signature = infer_signature(X_train, model.predict(X_train))
    
    # ── Log the model ─────────────────────────────────────────
    result = mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        signature=signature,
        registered_model_name="sales_regression_model"  # Registers directly
    )
    
    model_uri = result.model_uri
    print(f"Model URI: {model_uri}")
    print(f"Run ID: {run.info.run_id}")
    print(f"R² Score: {r2:.4f}")
    print(f"RMSE: {rmse:.2f}")

4. Exploratory Data Analysis (EDA)

4.1 Load and Explore Data

from pyspark.sql import SparkSession, functions as F
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

spark = SparkSession.builder.getOrCreate()

# Load data from a Delta Table
sales_df = spark.table("analytics_db.default.sales_data")

print(f"Dimensions: {sales_df.count()} rows × {len(sales_df.columns)} columns")
print(f"\nColumns: {sales_df.columns}")
print(f"\nSchema:")
sales_df.printSchema()

# Display a preview
display(sales_df)

# Complete descriptive statistics
sales_df.describe().show()

4.2 Null Values and Outlier Detection

# Count null values
null_counts = sales_df.select([
    F.count(F.when(F.col(c).isNull(), c)).alias(c)
    for c in sales_df.columns
])
print("Null values by column:")
null_counts.show()

# Identify numeric columns
numeric_cols = [
    field.name for field in sales_df.schema.fields
    if field.dataType.__class__.__name__ in 
    ['IntegerType', 'LongType', 'FloatType', 'DoubleType']
]
print(f"Numeric columns: {numeric_cols}")

# Visualize price distribution
sales_pd = sales_df.select("revenue", "quantity", "category", "region").toPandas()

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Revenue histogram
axes[0].hist(sales_pd["revenue"], bins=50, edgecolor='black', color='steelblue')
axes[0].set_title("Revenue Distribution")
axes[0].set_xlabel("Revenue ($)")
axes[0].set_ylabel("Frequency")

# Box plot by category
sales_pd.boxplot(column="revenue", by="category", ax=axes[1], rot=45)
axes[1].set_title("Revenue by Category")
axes[1].set_xlabel("Category")
axes[1].set_ylabel("Revenue ($)")

plt.tight_layout()
plt.show()
mlflow.log_figure(fig, "eda_distributions.png")

5. Data Preprocessing and Preparation

5.1 Removing Null Values and Outliers

from pyspark.sql import functions as F
from pyspark.sql import DataFrame

def remove_outliers_iqr(df: DataFrame, column: str, factor: float = 1.5) -> DataFrame:
    """
    Remove outliers from a numeric column using the IQR method.
    
    Args:
        df: Spark DataFrame
        column: Name of the column to process
        factor: IQR multiplier factor (default: 1.5)
    
    Returns:
        DataFrame without outliers
    """
    # Calculate quantiles
    quantiles = df.approxQuantile(column, [0.25, 0.75], 0.05)
    q1 = quantiles[0]
    q3 = quantiles[1]
    iqr = q3 - q1
    
    lower_bound = q1 - factor * iqr
    upper_bound = q3 + factor * iqr
    
    count_before = df.count()
    df_filtered = df.filter(
        (F.col(column) >= lower_bound) & 
        (F.col(column) <= upper_bound)
    )
    count_after = df_filtered.count()
    
    print(f"  {column}: {count_before - count_after} outliers removed "
          f"(bounds: [{lower_bound:.2f}, {upper_bound:.2f}])")
    
    return df_filtered

# Complete cleaning pipeline
def clean_sales_data(df: DataFrame) -> DataFrame:
    """Sales data cleaning pipeline."""
    
    print("=== Data Cleaning ===")
    initial_count = df.count()
    print(f"Initial rows: {initial_count:,}")
    
    # 1. Remove null values
    df = df.dropna()
    print(f"After removing nulls: {df.count():,}")
    
    # 2. Identify numeric columns
    numeric_cols = [
        field.name for field in df.schema.fields
        if field.dataType.__class__.__name__ in 
        ['IntegerType', 'LongType', 'FloatType', 'DoubleType']
    ]
    
    # 3. Remove outliers for each numeric column
    print("Removing outliers:")
    for col in numeric_cols:
        df = remove_outliers_iqr(df, col)
    
    final_count = df.count()
    print(f"\nFinal rows: {final_count:,}")
    print(f"Data removed: {((initial_count - final_count) / initial_count * 100):.1f}%")
    
    return df

# Apply cleaning
sales_clean = clean_sales_data(sales_df)

5.2 Encoding and Normalization (transition to pandas)

import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, LabelEncoder, OneHotEncoder
from sklearn.model_selection import train_test_split
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline

# Convert to pandas for scikit-learn (reasonable size data)
sales_pd = sales_clean.toPandas()

# Identify column types
categorical_cols = sales_pd.select_dtypes(include=['object']).columns.tolist()
numeric_cols_pd = sales_pd.select_dtypes(include=['int64', 'float64']).columns.tolist()

# Remove target (revenue) from features
target = 'revenue'
feature_cols = [c for c in sales_pd.columns if c != target]
numeric_features = [c for c in numeric_cols_pd if c != target]

print(f"Categorical columns: {categorical_cols}")
print(f"Numeric columns: {numeric_features}")

# Convert numerics to float64 (best practice for MLflow)
sales_pd[numeric_features] = sales_pd[numeric_features].astype(np.float64)
sales_pd[target] = sales_pd[target].astype(np.float64)

# Features and target
X = sales_pd[feature_cols]
y = sales_pd[target]

# Train/test split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, 
    test_size=0.2, 
    random_state=42
)

print(f"\nTrain size: {X_train.shape}")
print(f"Test size:  {X_test.shape}")

# Preprocessing pipeline with ColumnTransformer
preprocessor = ColumnTransformer(
    transformers=[
        ('num', StandardScaler(), numeric_features),
        ('cat', OneHotEncoder(handle_unknown='ignore', sparse=False), categorical_cols)
    ]
)

X_train_processed = preprocessor.fit_transform(X_train)
X_test_processed = preprocessor.transform(X_test)

print(f"\nFeatures after encoding: {X_train_processed.shape[1]} columns")

6. Training Models with scikit-learn

6.1 Comparing Multiple Models

from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
import numpy as np

def train_and_evaluate_model(
    model_name: str,
    model, 
    X_train, y_train, 
    X_test, y_test,
    experiment_name: str
) -> dict:
    """
    Train and evaluate a model with full MLflow tracking.
    
    Returns:
        Dictionary of metrics
    """
    mlflow.set_experiment(experiment_name)
    
    with mlflow.start_run(run_name=model_name) as run:
        # Training
        model.fit(X_train, y_train)
        y_pred_train = model.predict(X_train)
        y_pred_test = model.predict(X_test)
        
        # Metrics
        metrics = {
            "train_r2": r2_score(y_train, y_pred_train),
            "test_r2": r2_score(y_test, y_pred_test),
            "test_rmse": np.sqrt(mean_squared_error(y_test, y_pred_test)),
            "test_mae": mean_absolute_error(y_test, y_pred_test),
        }
        
        # Log parameters
        mlflow.log_param("model_type", model_name)
        mlflow.log_params(model.get_params())
        
        # Log metrics
        mlflow.log_metrics(metrics)
        
        # Infer signature
        signature = infer_signature(X_train, y_pred_train)
        
        # Log model
        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path="model",
            signature=signature
        )
        
        print(f"{model_name:30s} → R² test: {metrics['test_r2']:.4f} | "
              f"RMSE: {metrics['test_rmse']:.2f}")
        
        return metrics, run.info.run_id

# Train multiple models for comparison
EXPERIMENT = "/Users/user@company.com/sales_regression_comparison"
results = {}

models_to_train = [
    ("LinearRegression",     LinearRegression()),
    ("Ridge(alpha=1.0)",     Ridge(alpha=1.0)),
    ("Lasso(alpha=0.1)",     Lasso(alpha=0.1)),
    ("DecisionTree(d=5)",    DecisionTreeRegressor(max_depth=5, random_state=42)),
    ("DecisionTree(d=10)",   DecisionTreeRegressor(max_depth=10, random_state=42)),
    ("RandomForest(n=100)",  RandomForestRegressor(n_estimators=100, random_state=42)),
    ("GradientBoosting",     GradientBoostingRegressor(n_estimators=100, random_state=42)),
]

print("=== Model Comparison ===")
for name, model in models_to_train:
    metrics, run_id = train_and_evaluate_model(
        name, model, 
        X_train_processed, y_train,
        X_test_processed, y_test,
        EXPERIMENT
    )
    results[name] = {"metrics": metrics, "run_id": run_id}

# Display summary
print("\n=== Performance Summary ===")
sorted_results = sorted(results.items(), 
                         key=lambda x: x[1]["metrics"]["test_r2"], 
                         reverse=True)
for name, data in sorted_results:
    m = data["metrics"]
    print(f"{name:30s} R²={m['test_r2']:.4f}  RMSE={m['test_rmse']:.2f}  MAE={m['test_mae']:.2f}")

7. MLflow Tracking — Experiment Tracking

7.1 MLflow Experiments Interface

graph TB
    subgraph "MLflow Experiment"
        E[Experiment:\nsales_regression_comparison]
        
        subgraph "Run 1: LinearRegression"
            P1[Params:\nmodel_type=Linear\nn_features=24]
            M1[Metrics:\ntrain_r2=0.912\ntest_r2=0.897\nRMSE=1245.32]
            A1[Artifacts:\nmodel/\nresiduals_plot.png]
        end
        
        subgraph "Run 2: RandomForest"
            P2[Params:\nn_estimators=100\nmax_depth=None]
            M2[Metrics:\ntrain_r2=0.987\ntest_r2=0.978\nRMSE=421.15]
            A2[Artifacts:\nmodel/\nfeature_importances.png]
        end
        
        subgraph "Run 3: GradientBoosting"
            P3[Params:\nn_estimators=100\nlearning_rate=0.1]
            M3[Metrics:\ntrain_r2=0.991\ntest_r2=0.985\nRMSE=326.78]
            A3[Artifacts:\nmodel/\ntrain_test_comparison.png]
        end
    end
    
    E --> P1 & P2 & P3

7.2 Comparing Runs Programmatically

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

# Get experiment ID
experiment = mlflow.get_experiment_by_name(EXPERIMENT)
experiment_id = experiment.experiment_id

# Retrieve all runs for the experiment
runs = client.search_runs(
    experiment_ids=[experiment_id],
    order_by=["metrics.test_r2 DESC"]
)

print("=== All Runs Ranked by Test R² ===")
print(f"{'Run Name':35s} {'R² Test':10s} {'RMSE':10s} {'Run ID':15s}")
print("-" * 75)

for run in runs:
    name = run.data.tags.get("mlflow.runName", "N/A")
    r2 = run.data.metrics.get("test_r2", 0)
    rmse = run.data.metrics.get("test_rmse", 0)
    print(f"{name:35s} {r2:10.4f} {rmse:10.2f} {run.info.run_id[:15]}")

# Find the best run
best_run = runs[0]
print(f"\nBest model: {best_run.data.tags.get('mlflow.runName')}")
print(f"Run ID: {best_run.info.run_id}")
print(f"R² Test: {best_run.data.metrics.get('test_r2', 0):.4f}")

8. MLflow Autologging

8.1 Simplified Autologging

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestRegressor

mlflow.set_experiment("/Users/user@company.com/sales_autolog_demo")

# Enable autologging — automatically logs params, metrics, model
mlflow.sklearn.autolog(
    log_input_examples=True,    # Records input examples
    log_model_signatures=True,  # Records signature automatically
    log_models=True,            # Records model artifact
    silent=False                # Show log messages
)

# Train with minimal code — MLflow captures EVERYTHING automatically
with mlflow.start_run(run_name="RandomForest_autolog"):
    model = RandomForestRegressor(
        n_estimators=150,
        max_depth=12,
        min_samples_split=5,
        random_state=42,
        n_jobs=-1
    )
    model.fit(X_train_processed, y_train)
    y_pred = model.predict(X_test_processed)

# MLflow automatically logged:
# - All hyperparameters (n_estimators, max_depth, etc.)
# - Metrics (training_score, etc.)
# - Model with its signature
# - Feature importance plot
# - Cross-validation scores if used

print("Autologging complete — check the Experiments interface!")

8.2 What Autologging Captures Automatically

ElementDescriptionLogged by
HyperparametersAll model parametersmlflow.log_param auto
Training metricsTraining score, CV scoresmlflow.log_metric auto
ModelSerialized artifactmlflow.sklearn.log_model auto
SignatureInput/output schemaInferred automatically
Feature ImportanceFor ensemble modelsPNG artifact auto
Confusion MatrixFor classifiersPNG artifact auto

9. Hyperparameter Tuning with Ray Tune

9.1 Distributed HPT Concept

graph TB
    subgraph "Ray Tune on Databricks"
        Tuner[Tuner Object\nSearch space config]
        
        subgraph "Parallel Workers (Spark Cluster)"
            T1[Trial 1\nn_estimators=50\nmax_depth=5\nF1=0.82]
            T2[Trial 2\nn_estimators=100\nmax_depth=10\nF1=0.87]
            T3[Trial 3\nn_estimators=200\nmax_depth=7\nF1=0.91]
            T4[Trial N\nn_estimators=150\nmax_depth=12\nF1=0.93]
        end
        
        Best[Best config found\nn_estimators=150, max_depth=12]
        
        Tuner -->|Launch in parallel| T1 & T2 & T3 & T4
        T1 & T2 & T3 & T4 -->|Report metrics| Best
    end
    
    subgraph "MLflow (integrated)"
        MLF[Experiment\nParent Run + Child Runs]
    end
    
    T1 & T2 & T3 & T4 -->|mlflow.log| MLF

9.2 HPT Configuration and Execution

import os
import mlflow
import numpy as np
from ray import tune
from ray.tune.integration.mlflow import MLflowLoggerCallback
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score

# Retrieve Databricks credentials for Ray workers
DATABRICKS_HOST = spark.conf.get("spark.databricks.workspaceUrl")
DATABRICKS_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext() \
    .apiToken().get()

def train_and_evaluate_classification(config: dict):
    """
    Training function for Ray Tune.
    Called for each hyperparameter combination.
    
    Args:
        config: Dictionary of hyperparameters for this trial
    """
    # Configure MLflow for the remote worker
    os.environ["DATABRICKS_HOST"] = DATABRICKS_HOST
    os.environ["DATABRICKS_TOKEN"] = DATABRICKS_TOKEN
    mlflow.set_tracking_uri("databricks")
    mlflow.set_experiment("/Users/user@company.com/hr_analytics_hpt")
    
    # Start a nested MLflow run (child run)
    with mlflow.start_run(nested=True, run_name=f"trial_rf_{config['n_estimators']}_{config['max_depth']}"):
        # Log hyperparameters for this trial
        mlflow.log_params(config)
        
        # Instantiate model with trial hyperparameters
        model = RandomForestClassifier(
            n_estimators=config["n_estimators"],
            max_depth=config["max_depth"],
            min_samples_split=config["min_samples_split"],
            min_samples_leaf=config["min_samples_leaf"],
            class_weight=config["class_weight"],
            random_state=42,
            n_jobs=-1
        )
        
        # Train and evaluate
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        y_pred_proba = model.predict_proba(X_test)[:, 1]
        
        # Calculate metrics
        accuracy = accuracy_score(y_test, y_pred)
        f1 = f1_score(y_test, y_pred, average='weighted')
        precision = precision_score(y_test, y_pred, average='weighted', zero_division=0)
        recall = recall_score(y_test, y_pred, average='weighted')
        
        # Log metrics to MLflow
        mlflow.log_metrics({
            "accuracy": accuracy,
            "f1_score": f1,
            "precision": precision,
            "recall": recall
        })
        
        # Log model
        mlflow.sklearn.log_model(model, "model")
    
    # Report metrics to Ray Tune to guide the search
    tune.report(
        accuracy=accuracy,
        f1_score=f1,
        precision=precision,
        recall=recall
    )

# Define hyperparameter search space
search_space = {
    "n_estimators": tune.choice([50, 100, 150, 200, 300]),
    "max_depth": tune.choice([5, 7, 10, 12, 15, None]),
    "min_samples_split": tune.choice([2, 5, 10]),
    "min_samples_leaf": tune.choice([1, 2, 4]),
    "class_weight": tune.choice(["balanced", None])
}

# Start parent MLflow run
mlflow.set_experiment("/Users/user@company.com/hr_analytics_hpt")

with mlflow.start_run(run_name="hyperparameter_tuning_rf") as parent_run:
    mlflow.log_param("algorithm", "RandomForestClassifier")
    mlflow.log_param("n_trials", 20)
    mlflow.log_param("search_method", "random")
    
    # Configure and launch Ray Tune
    tuner = tune.Tuner(
        train_and_evaluate_classification,
        param_space=search_space,
        tune_config=tune.TuneConfig(
            metric="f1_score",
            mode="max",
            num_samples=20  # Number of trials
        )
    )
    
    # Run optimization
    results = tuner.fit()
    
    # Retrieve best configuration
    best_result = results.get_best_result(metric="f1_score", mode="max")
    best_config = best_result.config
    best_metrics = best_result.metrics
    
    # Log best config in parent run
    mlflow.log_params({f"best_{k}": v for k, v in best_config.items()})
    mlflow.log_metrics({f"best_{k}": v for k, v in best_metrics.items()
                        if isinstance(v, (int, float))})
    
    print(f"\nBest configuration found:")
    for key, value in best_config.items():
        print(f"  {key}: {value}")
    print(f"\nBest metrics:")
    print(f"  F1 Score: {best_metrics['f1_score']:.4f}")
    print(f"  Accuracy: {best_metrics['accuracy']:.4f}")

10. Model Registry with Unity Catalog

10.1 Registering a Model in Unity Catalog

import mlflow
from mlflow.tracking import MlflowClient

# Configure MLflow to use Unity Catalog
mlflow.set_registry_uri("databricks-uc")

# Train and register directly in Unity Catalog
mlflow.set_experiment("/Users/user@company.com/sales_price_prediction")

with mlflow.start_run(run_name="rf_production_candidate") as run:
    
    # Train the best model (found via HPT)
    best_model = RandomForestRegressor(
        n_estimators=150,
        max_depth=12,
        min_samples_split=5,
        random_state=42,
        n_jobs=-1
    )
    best_model.fit(X_train_processed, y_train)
    y_pred = best_model.predict(X_test_processed)
    
    # Final metrics
    r2 = r2_score(y_test, y_pred)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    
    mlflow.log_params(best_model.get_params())
    mlflow.log_metrics({"r2_score": r2, "rmse": rmse})
    
    # Model signature
    signature = infer_signature(X_train_processed, y_pred)
    
    # Register in Unity Catalog
    # Format: catalog.schema.model_name
    result = mlflow.sklearn.log_model(
        sk_model=best_model,
        artifact_path="model",
        signature=signature,
        registered_model_name="analytics_db.default.sales_regression_model"
    )
    
    print(f"Model registered in Unity Catalog!")
    print(f"Model URI: {result.model_uri}")

10.2 Managing Versions and Stages

from mlflow.tracking import MlflowClient

client = MlflowClient()
model_name = "analytics_db.default.sales_regression_model"

# List all versions
versions = client.search_model_versions(f"name='{model_name}'")
for v in versions:
    print(f"Version {v.version}: {v.current_stage} | Run ID: {v.run_id[:8]}...")

# Promote latest version to "Staging"
latest_version = max([v.version for v in versions])
client.transition_model_version_stage(
    name=model_name,
    version=latest_version,
    stage="Staging",
    archive_existing_versions=False
)

print(f"Version {latest_version} promoted to Staging")

# After validation, promote to "Production"
client.transition_model_version_stage(
    name=model_name,
    version=latest_version,
    stage="Production",
    archive_existing_versions=True  # Archives old Production versions
)

print(f"Version {latest_version} in Production!")

# Add annotations
client.update_model_version(
    name=model_name,
    version=latest_version,
    description=f"RandomForest R²={r2:.4f} RMSE={rmse:.2f}. Validated on 2024-01-15."
)

# Add a tag
client.set_model_version_tag(
    name=model_name,
    version=latest_version,
    key="validated_by",
    value="alice@company.com"
)

10.3 Model Version Lifecycle

stateDiagram-v2
    [*] --> None : Model registration
    None --> Staging : Candidate for testing
    Staging --> Production : Validated after tests
    Staging --> Archived : Rejected
    Production --> Archived : Replaced by a new version
    Production --> Staging : Rollback required
    Archived --> [*] : Optional deletion

11. Model Serving — Deploying as REST API

11.1 Configure a Serving Endpoint

# Configure Model Serving via Databricks API
import requests
import json

workspace_url = "https://adb-xxxx.azuredatabricks.net"
token = dbutils.secrets.get(scope="kv-secrets", key="databricks-token")

# Create the serving endpoint
endpoint_config = {
    "name": "sales-regression-endpoint",
    "config": {
        "served_models": [
            {
                "name": "sales-regression-v1",
                "model_name": "analytics_db.default.sales_regression_model",
                "model_version": "1",
                "workload_size": "Small",
                "scale_to_zero_enabled": True  # Save costs when idle
            }
        ],
        "traffic_config": {
            "routes": [
                {
                    "served_model_name": "sales-regression-v1",
                    "traffic_percentage": 100
                }
            ]
        }
    }
}

response = requests.post(
    f"{workspace_url}/api/2.0/serving-endpoints",
    headers={"Authorization": f"Bearer {token}"},
    json=endpoint_config
)

print(f"Status: {response.status_code}")
print(f"Response: {json.dumps(response.json(), indent=2)}")

11.2 Endpoint Configuration Options

OptionDescriptionValues
Compute TypeCPU or GPUCPU (default), GPU_SMALL, GPU_LARGE
Workload SizeMax concurrencySmall (0-4), Medium (0-16), Large (0-64)
Scale to ZeroStop when idletrue/false
Traffic SplittingDistribution between versionsPercentages (total = 100%)
Environment VariablesVariables for the modelKey-value pairs

11.3 A/B Testing Between Two Versions

# A/B test configuration: 90% v1, 10% v2
ab_test_config = {
    "name": "sales-regression-endpoint",
    "config": {
        "served_models": [
            {
                "name": "sales-v1-stable",
                "model_name": "analytics_db.default.sales_regression_model",
                "model_version": "1",
                "workload_size": "Small"
            },
            {
                "name": "sales-v2-new",
                "model_name": "analytics_db.default.sales_regression_model",
                "model_version": "2",
                "workload_size": "Small"
            }
        ],
        "traffic_config": {
            "routes": [
                {"served_model_name": "sales-v1-stable", "traffic_percentage": 90},
                {"served_model_name": "sales-v2-new", "traffic_percentage": 10}
            ]
        }
    }
}

12. Predictions from a Deployed Model

12.1 Calling the Model via REST API

import os
import requests
import json
import pandas as pd

# Configuration
workspace_url = "https://adb-xxxx.azuredatabricks.net"
endpoint_name = "sales-regression-endpoint"
token = os.environ.get("DATABRICKS_TOKEN")

# Prepare inference data
inference_data = pd.DataFrame({
    "product_id": [101, 202, 303, 404],
    "category": ["Electronics", "Clothing", "Food", "Sports"],
    "region": ["North", "South", "East", "West"],
    "channel": ["Online", "Store", "Online", "Store"],
    "unit_price": [299.99, 49.99, 12.50, 89.99],
    "quantity": [2, 5, 10, 3],
    "discount_pct": [5.0, 10.0, 0.0, 15.0],
    "customer_age": [35, 28, 45, 52],
    "loyalty_score": [8.5, 6.2, 9.1, 7.3]
})

# Format expected by Databricks Model Serving (dataframe_split)
payload = {
    "dataframe_split": {
        "columns": inference_data.columns.tolist(),
        "data": inference_data.values.tolist()
    }
}

# REST call
response = requests.post(
    url=f"{workspace_url}/serving-endpoints/{endpoint_name}/invocations",
    headers={
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    },
    data=json.dumps(payload)
)

if response.status_code == 200:
    predictions = response.json()["predictions"]
    for i, pred in enumerate(predictions):
        row = inference_data.iloc[i]
        print(f"Record {i+1} ({row['category']} - {row['region']}): "
              f"${pred:,.2f}")
else:
    print(f"Error {response.status_code}: {response.text}")

12.2 Batch Inference on a Delta Table

import mlflow.pyfunc
from pyspark.sql import functions as F

# Load model from Unity Catalog for batch inference
model_uri = "models:/analytics_db.default.sales_regression_model/Production"
model = mlflow.pyfunc.load_model(model_uri)

# Load new data to score
new_sales_df = spark.table("analytics_db.default.new_transactions")

# Apply model via Spark UDF (for very large datasets)
@F.pandas_udf("double")
def predict_revenue(
    product_id: pd.Series,
    category: pd.Series,
    region: pd.Series,
    channel: pd.Series,
    unit_price: pd.Series,
    quantity: pd.Series,
    discount_pct: pd.Series,
    customer_age: pd.Series,
    loyalty_score: pd.Series
) -> pd.Series:
    """Vectorized Pandas UDF for distributed inference."""
    features = pd.DataFrame({
        "product_id": product_id, "category": category, "region": region,
        "channel": channel, "unit_price": unit_price, "quantity": quantity,
        "discount_pct": discount_pct, "customer_age": customer_age,
        "loyalty_score": loyalty_score
    })
    # Apply same preprocessing as during training
    features_processed = preprocessor.transform(features)
    return pd.Series(model.predict(features_processed))

# Apply the UDF on the Spark DataFrame
predictions_df = new_sales_df.withColumn(
    "predicted_revenue",
    predict_revenue(
        F.col("product_id"), F.col("category"), F.col("region"),
        F.col("channel"), F.col("unit_price"), F.col("quantity"),
        F.col("discount_pct"), F.col("customer_age"), F.col("loyalty_score")
    )
)

# Save predictions
predictions_df.write \
    .mode("overwrite") \
    .format("delta") \
    .saveAsTable("analytics_db.default.sales_predictions")

print(f"Predictions saved: {predictions_df.count():,} rows")
predictions_df.select("category", "region", "predicted_revenue") \
    .orderBy(F.desc("predicted_revenue")) \
    .show(10)

13. Automated Model Retraining

13.1 Auto-Retraining Notebook

# Notebook: AutoRetraining.py
# Triggered by File Arrival Trigger when new data is available

import mlflow
import mlflow.sklearn
from databricks.sdk import WorkspaceClient
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score, mean_squared_error
from glob import glob
import pandas as pd
import numpy as np
import os

# Configuration
VOLUME_PATH = "/Volumes/analytics_db/default/data-volume/*.csv"
EXPERIMENT_NAME = "/Users/user@company.com/sales_auto_retraining"
MODEL_NAME = "analytics_db.default.sales_regression_model"

mlflow.set_experiment(EXPERIMENT_NAME)
mlflow.set_registry_uri("databricks-uc")

# Load all available files in the volume
csv_files = glob(VOLUME_PATH)
print(f"Files found: {csv_files}")

if not csv_files:
    raise Exception("No files available in the volume!")

# Combine all data files
dfs = []
for f in csv_files:
    df = pd.read_csv(f)
    df["source_file"] = os.path.basename(f)
    dfs.append(df)

all_data = pd.concat(dfs, ignore_index=True)
all_data = all_data.dropna()

print(f"Total data loaded: {all_data.shape}")

# Preparation
X = all_data.drop(["revenue"], axis=1, errors="ignore")
y = all_data["revenue"]

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Apply preprocessing
X_train_proc = preprocessor.fit_transform(X_train)
X_test_proc = preprocessor.transform(X_test)

# Train with autologging
with mlflow.start_run(run_name=f"auto_retrain_n{len(csv_files)}_files") as run:
    
    model = RandomForestRegressor(
        n_estimators=100,
        max_depth=10,
        random_state=42,
        n_jobs=-1
    )
    
    model.fit(X_train_proc, y_train)
    y_pred = model.predict(X_test_proc)
    
    r2 = r2_score(y_test, y_pred)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    
    mlflow.log_params(model.get_params())
    mlflow.log_metrics({
        "r2_score": r2,
        "rmse": rmse,
        "n_training_files": len(csv_files),
        "n_training_rows": len(X_train)
    })
    
    signature = infer_signature(X_train_proc, y_pred)
    
    result = mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        signature=signature,
        registered_model_name=MODEL_NAME
    )
    
    print(f"New model registered!")
    print(f"  R² Score: {r2:.4f}")
    print(f"  RMSE: {rmse:.2f}")
    print(f"  Model URI: {result.model_uri}")
    
    # Auto-promote if metrics are good
    if r2 > 0.95:
        client = mlflow.tracking.MlflowClient()
        latest_version = client.get_latest_versions(MODEL_NAME)[0].version
        client.transition_model_version_stage(
            name=MODEL_NAME,
            version=latest_version,
            stage="Production"
        )
        print(f"Version {latest_version} auto-promoted to Production (R²={r2:.4f} > 0.95)")

13.2 Configure the Job with File Arrival Trigger

# Create automated job via Databricks SDK
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import (
    Task, NotebookTask, FileArrivalTriggerConfiguration,
    TriggerSettings, TriggerType
)

w = WorkspaceClient()

# Create the auto-retraining job
job = w.jobs.create(
    name="Auto_Retraining_Sales_Regression",
    tasks=[
        Task(
            task_key="Auto_Retraining",
            description="Retrain regression model on new data arrival",
            notebook_task=NotebookTask(
                notebook_path="/Users/user@company.com/AutoRetraining",
                source="WORKSPACE"
            ),
            existing_cluster_id="data-ml-cluster"
        )
    ],
    trigger=TriggerSettings(
        trigger_type=TriggerType.FILE_ARRIVAL,
        file_arrival=FileArrivalTriggerConfiguration(
            url="dbfs:/Volumes/analytics_db/default/data-volume/",
            min_time_between_triggers_seconds=3600,  # Minimum 1h between triggers
            wait_after_last_change_seconds=60
        )
    ),
    email_notifications={
        "on_failure": ["ml-team@company.com"],
        "on_success": ["ml-team@company.com"]
    }
)

print(f"Job created: {job.job_id}")

14. Databricks AutoML

14.1 AutoML — Overview

graph LR
    subgraph "AutoML Pipeline"
        Data[Raw Data\nDelta Table] --> Prep[Data Prep\nautomatic]
        Prep --> Select[Algorithm\nSelection]
        Select --> Train[Train\nmultiple models]
        Train --> HPT2[Hyperparameter\nAuto Tuning]
        HPT2 --> Eval[Evaluation\nand comparison]
        Eval --> Best[Best model\nregistered]
    end
    
    subgraph "Generated Notebooks"
        NB1[Trial 1: LogReg\nComplete notebook]
        NB2[Trial 2: RF\nComplete notebook]
        NB3[Trial 3: XGBoost\nComplete notebook]
        Best2[Best Notebook\nEditable]
    end
    
    Train --> NB1 & NB2 & NB3
    Best --> Best2

14.2 Launch AutoML via Python API

from databricks import automl
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

# Load data for classification
churn_df = spark.table("analytics_db.default.churn")

# Launch AutoML for classification
# AutoML automatically handles: cleaning, encoding, model selection, HPT
automl_run = automl.classify(
    dataset=churn_df,
    target_col="churn",
    
    # Optional configuration
    primary_metric="f1",
    time_budget_s=3600,          # 1 hour max
    max_trials=20,               # Max number of trials
    
    # Columns to exclude (e.g.: identifiers)
    exclude_cols=["customer_id"],
    
    # Test set fraction
    split_col=None,              # None = automatic split
    test_size=0.2,
    
    # MLflow experiment name
    experiment_dir="/Users/user@company.com/",
    
    # Verbosity
    verbosity="info"
)

# Results
print(f"Best run ID: {automl_run.best_trial.mlflow_run_id}")
print(f"Best model metrics:")
print(f"  F1 Score: {automl_run.best_trial.evaluation_metric_val:.4f}")

# Access generated notebooks
print(f"\nGenerated notebooks:")
for trial in automl_run.trials[:5]:
    print(f"  {trial.notebook_path} → F1={trial.evaluation_metric_val:.4f}")

14.3 AutoML Use Cases

Problem TypeAlgorithms TestedPrimary Metrics
ClassificationLogReg, RF, XGBoost, LightGBM, Decision TreeAccuracy, F1, AUC-ROC, Precision, Recall
RegressionLinear, Ridge, RF, XGBoost, LightGBMRMSE, MAE, R²
ForecastingProphet, ARIMA, Exponential SmoothingSMAPE, MAE, RMSE

15. ML Orchestration with Azure Data Factory

15.1 Invoke a Databricks Notebook from ADF

# In ADF Studio → New Pipeline → Add Databricks Notebook Activity

# Required Databricks Linked Service:
# - Workspace URL
# - Access Token (from Azure Key Vault)
# - Cluster: Job Cluster or Existing Interactive Cluster
// Databricks Notebook Activity configuration in ADF
{
  "name": "Train_Register_Model",
  "type": "DatabricksNotebook",
  "policy": {
    "timeout": "7.00:00:00",
    "retry": 1,
    "retryIntervalInSeconds": 30
  },
  "typeProperties": {
    "notebookPath": "/Users/user@company.com/training_and_registering_model",
    "baseParameters": {
      "experiment_name": {
        "value": "@pipeline().parameters.ExperimentName",
        "type": "Expression"
      },
      "model_version": {
        "value": "@string(pipeline().RunId)",
        "type": "Expression"
      }
    }
  },
  "linkedServiceName": {
    "referenceName": "DatabricksLinkedService",
    "type": "LinkedServiceReference"
  }
}

15.2 Complete ADF Pipeline for ML

sequenceDiagram
    participant T as ADF Trigger (Cron)
    participant P as ADF Pipeline
    participant DB as Databricks
    participant UC as Unity Catalog

    T->>P: Trigger (2 AM daily)
    P->>DB: Copy Activity: Ingest new data → ADLS
    DB-->>P: Data ingested
    P->>DB: Notebook Activity: ETL + Feature Engineering
    DB-->>P: Features ready in Feature Store
    P->>DB: Notebook Activity: Train + Register model
    DB->>UC: Register new model version
    UC-->>DB: Confirmed
    DB-->>P: Model metrics
    P->>DB: Notebook Activity: Validate and Promote if OK
    DB-->>P: Model in Production
    P-->>T: Pipeline complete — Confirmation email

16. Databricks Feature Store

16.1 Architecture and Concepts

graph TB
    subgraph "Source Data"
        R[Raw Data\nDelta Lake]
    end
    
    subgraph "Feature Engineering"
        FE[Spark / Python\nTransformations]
        ENG[Computed features\nscores, ratios, encodings]
    end
    
    subgraph "Feature Store (Delta Lake)"
        FT1[Feature Table:\nemployee_features\nKEY: emp_id]
        FT2[Feature Table:\nperformance_features\nKEY: emp_id]
    end
    
    subgraph "Training"
        TJ[Training Job]
        M[Trained model\nwith feature lineage]
    end
    
    subgraph "Inference"
        OFS[Online Feature Store\nResponse < 10ms]
        BI[Batch Inference\nDelta Table]
    end
    
    R --> FE --> ENG
    ENG --> FT1 & FT2
    FT1 & FT2 --> TJ --> M
    M --> OFS & BI

16.2 Create and Use a Feature Table

from databricks.feature_engineering import FeatureEngineeringClient
from pyspark.sql import functions as F
import pandas as pd

spark = SparkSession.builder.getOrCreate()
fe = FeatureEngineeringClient()

# Load source data
hr_analytics_df = spark.table("analytics_db.default.hr_data")

# Add unique identifier (required for Feature Store)
hr_analytics_with_id = hr_analytics_df.withColumn(
    "emp_id",
    F.monotonically_increasing_id().cast("long")
)

# Feature engineering — create computed features
hr_features_df = hr_analytics_with_id.withColumn(
    "overall_performance_score",
    (
        F.col("previous_year_rating") * 0.4 +
        F.col("kpis_met_above_80pct") * 0.3 +
        F.col("awards_received") * 0.2 +
        (F.col("avg_training_score") / 100) * 0.1
    )
).withColumn(
    "tenure_segment",
    F.when(F.col("years_of_service") <= 2, "Junior")
     .when(F.col("years_of_service") <= 5, "Mid-Level")
     .when(F.col("years_of_service") <= 10, "Senior")
     .otherwise("Veteran")
)

# Select features for Feature Store
selected_features = [
    "emp_id",
    "department", "region", "education",
    "gender", "recruitment_channel",
    "num_trainings", "age",
    "previous_year_rating", "years_of_service",
    "kpis_met_above_80pct", "awards_received",
    "avg_training_score", "overall_performance_score",
    "tenure_segment"
]

hr_features_selected = hr_features_df.select(selected_features)

# Create Feature Table in Unity Catalog
feature_table_name = "analytics_db.default.hr_employee_features"

fe.create_table(
    name=feature_table_name,
    primary_keys=["emp_id"],
    df=hr_features_selected,
    schema=hr_features_selected.schema,
    description="Computed employee features for promotion prediction"
)

print(f"Feature Table created: {feature_table_name}")
print(f"Features stored: {hr_features_selected.count()} records")

16.3 Train a Model with the Feature Store

from databricks.feature_engineering import FeatureEngineeringClient, FeatureLookup
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import mlflow

fe = FeatureEngineeringClient()

# Label data (prediction target)
labels_df = spark.table("analytics_db.default.hr_data") \
    .select("emp_id", "is_promoted")  # emp_id is the join key

# Define Feature Store lookup
feature_lookups = [
    FeatureLookup(
        table_name="analytics_db.default.hr_employee_features",
        feature_names=[
            "num_trainings", "age", "previous_year_rating",
            "years_of_service", "kpis_met_above_80pct",
            "awards_received", "avg_training_score",
            "overall_performance_score"
        ],
        lookup_key="emp_id"
    )
]

# Create training dataset by joining features + labels
training_set = fe.create_training_set(
    df=labels_df,
    feature_lookups=feature_lookups,
    label="is_promoted",
    exclude_columns=["emp_id"]  # Don't include key as feature
)

training_df = training_set.load_df()

# Convert to pandas for scikit-learn
training_pd = training_df.toPandas()
X = training_pd.drop("is_promoted", axis=1)
y = training_pd["is_promoted"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)

# Train with Feature Store (MLflow integrated)
mlflow.set_experiment("/Users/user@company.com/hr_promotion_prediction")

with mlflow.start_run(run_name="rf_with_feature_store") as run:
    model = RandomForestClassifier(
        n_estimators=100,
        class_weight="balanced",
        random_state=42
    )
    model.fit(X_train, y_train)
    
    # Log model with Feature Store (for lineage)
    fe.log_model(
        model=model,
        artifact_path="model",
        flavor=mlflow.sklearn,
        training_set=training_set,
        registered_model_name="analytics_db.default.hr_promotion_model"
    )
    
    accuracy = model.score(X_test, y_test)
    mlflow.log_metric("accuracy", accuracy)
    print(f"Accuracy: {accuracy:.4f}")
    print("Model registered with Feature Store lineage!")

17. Distributed Machine Learning with Spark MLlib

17.1 When to Use Spark MLlib vs scikit-learn?

Criterionscikit-learnSpark MLlib
Data size< 100 GB (in memory)> 100 GB (distributed)
ML typeStandard algorithmsLarge-scale distributed
EcosystemRich (sklearn, xgboost, etc.)Limited but distributed
PerformanceVery fast on small datasetsOptimal on large datasets
APINative Python + pandasPySpark (more verbose)

17.2 Distributed ML Pipeline with Spark MLlib

from pyspark.ml import Pipeline
from pyspark.ml.feature import (
    StringIndexer, OneHotEncoder, VectorAssembler,
    StandardScaler, Imputer
)
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import BinaryClassificationEvaluator, MulticlassClassificationEvaluator
from pyspark.ml.tuning import CrossValidator, ParamGridBuilder

# Load data
df = spark.table("analytics_db.default.hr_data") \
    .dropna()

# Categorical and numeric columns
cat_cols = ["department", "region", "education", "gender", "recruitment_channel"]
num_cols = ["num_trainings", "age", "previous_year_rating", 
            "years_of_service", "kpis_met_above_80pct",
            "awards_received", "avg_training_score"]

# Pipeline stages
stages = []

# 1. Index categorical columns
indexers = [
    StringIndexer(inputCol=col, outputCol=f"{col}_indexed", handleInvalid="keep")
    for col in cat_cols
]
stages.extend(indexers)

# 2. One-Hot Encoding
encoders = [
    OneHotEncoder(inputCol=f"{col}_indexed", outputCol=f"{col}_encoded")
    for col in cat_cols
]
stages.extend(encoders)

# 3. Assemble all features
feature_cols = [f"{col}_encoded" for col in cat_cols] + num_cols
assembler = VectorAssembler(inputCols=feature_cols, outputCol="features_raw")
stages.append(assembler)

# 4. Normalize features
scaler = StandardScaler(inputCol="features_raw", outputCol="features", 
                         withStd=True, withMean=True)
stages.append(scaler)

# 5. Index label
label_indexer = StringIndexer(inputCol="is_promoted", outputCol="label")
stages.append(label_indexer)

# 6. Distributed RandomForest Classifier
rf = RandomForestClassifier(
    featuresCol="features",
    labelCol="label",
    numTrees=100,
    maxDepth=10,
    seed=42
)
stages.append(rf)

# Create the Pipeline
pipeline = Pipeline(stages=stages)

# Train/test split
train_df, test_df = df.randomSplit([0.8, 0.2], seed=42)

# Train
model = pipeline.fit(train_df)

# Evaluate
predictions = model.transform(test_df)

evaluator_auc = BinaryClassificationEvaluator(labelCol="label", metricName="areaUnderROC")
evaluator_acc = MulticlassClassificationEvaluator(
    labelCol="label", predictionCol="prediction", metricName="accuracy"
)

auc = evaluator_auc.evaluate(predictions)
accuracy = evaluator_acc.evaluate(predictions)

print(f"AUC-ROC: {auc:.4f}")
print(f"Accuracy: {accuracy:.4f}")

# Register Spark model in MLflow
with mlflow.start_run(run_name="spark_mllib_rf"):
    mlflow.log_params({"num_trees": 100, "max_depth": 10, "algorithm": "RandomForestSpark"})
    mlflow.log_metrics({"auc_roc": auc, "accuracy": accuracy})
    mlflow.spark.log_model(model, "spark_model")

18. MLOps Best Practices on Databricks

18.1 Complete MLOps Pipeline

graph LR
    subgraph "Dev Environment"
        E1[Exploration\nEDA + Feature Eng] 
        E2[Experimentation\nMulti-model]
        E3[MLflow Tracking\nCompare runs]
    end
    
    subgraph "Staging Environment"
        S1[Validation\nIndependent dataset]
        S2[Tests\nPerformance + Fairness]
        S3[Model Registry\nStaging Stage]
    end
    
    subgraph "Production Environment"
        P1[Model Serving\nREST Endpoint]
        P2[Batch Inference\nDatabricks Jobs]
        P3[Monitoring\nData Drift + Performance]
    end
    
    E3 -->|Best model| S1
    S2 -->|Validation OK| S3
    S3 -->|Promotion| P1 & P2
    P3 -->|Drift detected| E1

18.2 MLOps Checklist

PhaseCheckImportance
DataVersioning with Delta LakeCritical
FeaturesFeature Store for reusabilityHigh
ExperimentsMLflow Tracking enabledCritical
ModelsSignatures inferred and loggedHigh
RegistryModels registered in Unity CatalogCritical
ValidationTests on independent holdout datasetCritical
MonitoringProduction metrics trackingHigh
RetrainingAutomated pipeline with triggerHigh
SecurityModel access control via UCHigh
DocumentationAnnotations and tags on versionsMedium

19. Summary and Tool Comparison

19.1 When to Use What?

ToolUse CaseWhen to Use
scikit-learnClassic MLData < 100 GB, rapid prototyping
Spark MLlibDistributed MLData > 100 GB, native Spark features
TensorFlow/PyTorchDeep learningImages, text, time series
XGBoost/LightGBMHigh-performance tabular dataML competitions, production
AutoMLQuick baselineNew problem, short deadline
Ray TuneLarge-scale HPTLarge search space, Spark cluster
Feature StoreReusable featuresMulti-project ML, large datasets

19.2 Reference ML Architecture on Databricks

graph TB
    subgraph "Data Zone"
        RAW[Bronze: Raw Data\nADLS Gen2]
        SILVER[Silver: Clean Data\nDelta Lake]
        GOLD[Gold: Engineered Features\nFeature Store]
    end
    
    subgraph "ML Zone"
        EXP[Experimentation\nNotebooks + AutoML]
        TRACK[MLflow Tracking\nRuns + Metrics]
        HPT3[HPT\nRay Tune]
        REG[Model Registry\nUnity Catalog]
    end
    
    subgraph "Production Zone"
        SERVE[Model Serving\nREST API]
        BATCH[Batch Jobs\nPeriodic predictions]
        MONITOR[Monitoring\nDrift + Performance]
        RETRAIN[Auto-Retraining\nFile Arrival Trigger]
    end
    
    RAW --> SILVER --> GOLD
    GOLD --> EXP
    EXP --> TRACK
    HPT3 --> EXP
    TRACK --> REG
    REG --> SERVE & BATCH
    SERVE & BATCH --> MONITOR
    MONITOR -->|Drift detected| RETRAIN
    RETRAIN --> EXP

20. Glossary

TermDefinition
AutoMLDatabricks feature that automates model selection, preprocessing, and HPT
AUC-ROCArea Under the ROC Curve — classification metric (1.0 = perfect)
Cross-ValidationEvaluation technique that divides data into K folds for robust estimation
Ensemble LearningTechnique combining multiple models (RandomForest, GradientBoosting)
Feature EngineeringProcess of creating new features from raw data
Feature StoreCentralized repository for storing, discovering, and reusing ML features
Feature TableDelta table in the Feature Store with a primary key
HPT (Hyperparameter Tuning)Search for best hyperparameters for an ML model
IQR (Interquartile Range)Q3 - Q1 — used to detect outliers
Lazy EvaluationSpark transformations only execute at the final action
Model DriftPerformance degradation of a model due to evolving data
MLflowOpen-source platform for managing the ML lifecycle (tracking, registry, serving)
MLflow AutologgingAutomatic capture of params, metrics, and models without explicit code
Model RegistryMLflow component for versioning and promoting models (Staging → Production)
Model ServingDeploying a model as a REST endpoint for real-time predictions
Model SignatureInput/output schema of an MLflow model for validation
Ray TuneDistributed framework for HPT, integrated in Databricks ML Runtime
RMSERoot Mean Square Error — regression metric
scikit-learnPython ML library, integrated in Databricks ML Runtime
Spark MLlibNative distributed ML library of Spark for large datasets
Trial RunAn individual execution in an HPT (with a specific config)
Unity CatalogDatabricks centralized governance including the Model Registry
VolumeNon-tabular storage in Unity Catalog (for CSV, images, etc.)

Search Terms

machine · azure · databricks · spark · data · engineering · analytics · model · mlflow · automl · api · architecture · autologging · distributed · feature · mllib · mlops · pipeline · runtime · adf · catalog · comparing · concepts · configuration

Interested in this course?

Contact us to book it or get a custom training plan for your team.