Table of Contents
- Why Databricks for Machine Learning?
- Databricks ML Runtime
- MLflow — ML Lifecycle Management Platform
- Exploratory Data Analysis (EDA)
- Data Preprocessing and Preparation
- Training Models with scikit-learn
- MLflow Tracking — Experiment Tracking
- MLflow Autologging
- Hyperparameter Tuning with Ray Tune
- Model Registry with Unity Catalog
- Model Serving — Deploying as REST API
- Predictions from a Deployed Model
- Automated Model Retraining
- Databricks AutoML
- ML Orchestration with Azure Data Factory
- Databricks Feature Store
- Distributed Machine Learning with Spark MLlib
- MLOps Best Practices on Databricks
- Summary and Tool Comparison
- Glossary
1. Why Databricks for Machine Learning?
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
| Capability | Description | Benefit |
|---|
| Distributed Scale | Apache Spark + Distributed ML | Train on millions of rows |
| Built-in Libraries | scikit-learn, TF, PyTorch, XGBoost, MLlib | No complex installation |
| Native MLflow | Tracking, Registry, Serving integrated | Full lifecycle on one platform |
| AutoML | No-code/low-code with automatic HPT | Rapid prototyping |
| Feature Store | Centralized feature repository | Reusability and consistency |
| Delta Lake | Data versioning | Model reproducibility |
| GPU Support | CUDA, cuDNN via ML Runtime GPU | High-performance Deep Learning |
| Unity Catalog | Model governance | Compliance and auditability |
1.3 MLflow vs Competitors
| Tool | Type | Strengths | Limitations |
|---|
| MLflow | Open Source | Native Databricks integration, multi-framework | Less rich UI |
| Amazon SageMaker | Cloud (AWS) | Managed, full AWS integration | AWS lock-in |
| Google Vertex AI | Cloud (GCP) | Advanced AutoML, BigQuery integration | GCP lock-in |
| Azure Machine Learning | Cloud (Azure) | Native Azure integration | Separate from Databricks |
| Weights & Biases | SaaS | Rich visualizations, collaboration | Costly for large teams |
| Kubeflow | Open Source (K8s) | Portable, Kubernetes native | Operational complexity |
| ClearML | Open Source | Full MLOps, orchestration | Less 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.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
| Element | Description | Logged by |
|---|
| Hyperparameters | All model parameters | mlflow.log_param auto |
| Training metrics | Training score, CV scores | mlflow.log_metric auto |
| Model | Serialized artifact | mlflow.sklearn.log_model auto |
| Signature | Input/output schema | Inferred automatically |
| Feature Importance | For ensemble models | PNG artifact auto |
| Confusion Matrix | For classifiers | PNG 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
# 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
| Option | Description | Values |
|---|
| Compute Type | CPU or GPU | CPU (default), GPU_SMALL, GPU_LARGE |
| Workload Size | Max concurrency | Small (0-4), Medium (0-16), Large (0-64) |
| Scale to Zero | Stop when idle | true/false |
| Traffic Splitting | Distribution between versions | Percentages (total = 100%) |
| Environment Variables | Variables for the model | Key-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)")
# 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 Type | Algorithms Tested | Primary Metrics |
|---|
| Classification | LogReg, RF, XGBoost, LightGBM, Decision Tree | Accuracy, F1, AUC-ROC, Precision, Recall |
| Regression | Linear, Ridge, RF, XGBoost, LightGBM | RMSE, MAE, R² |
| Forecasting | Prophet, ARIMA, Exponential Smoothing | SMAPE, 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?
| Criterion | scikit-learn | Spark MLlib |
|---|
| Data size | < 100 GB (in memory) | > 100 GB (distributed) |
| ML type | Standard algorithms | Large-scale distributed |
| Ecosystem | Rich (sklearn, xgboost, etc.) | Limited but distributed |
| Performance | Very fast on small datasets | Optimal on large datasets |
| API | Native Python + pandas | PySpark (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
| Phase | Check | Importance |
|---|
| Data | Versioning with Delta Lake | Critical |
| Features | Feature Store for reusability | High |
| Experiments | MLflow Tracking enabled | Critical |
| Models | Signatures inferred and logged | High |
| Registry | Models registered in Unity Catalog | Critical |
| Validation | Tests on independent holdout dataset | Critical |
| Monitoring | Production metrics tracking | High |
| Retraining | Automated pipeline with trigger | High |
| Security | Model access control via UC | High |
| Documentation | Annotations and tags on versions | Medium |
19.1 When to Use What?
| Tool | Use Case | When to Use |
|---|
| scikit-learn | Classic ML | Data < 100 GB, rapid prototyping |
| Spark MLlib | Distributed ML | Data > 100 GB, native Spark features |
| TensorFlow/PyTorch | Deep learning | Images, text, time series |
| XGBoost/LightGBM | High-performance tabular data | ML competitions, production |
| AutoML | Quick baseline | New problem, short deadline |
| Ray Tune | Large-scale HPT | Large search space, Spark cluster |
| Feature Store | Reusable features | Multi-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
| Term | Definition |
|---|
| AutoML | Databricks feature that automates model selection, preprocessing, and HPT |
| AUC-ROC | Area Under the ROC Curve — classification metric (1.0 = perfect) |
| Cross-Validation | Evaluation technique that divides data into K folds for robust estimation |
| Ensemble Learning | Technique combining multiple models (RandomForest, GradientBoosting) |
| Feature Engineering | Process of creating new features from raw data |
| Feature Store | Centralized repository for storing, discovering, and reusing ML features |
| Feature Table | Delta 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 Evaluation | Spark transformations only execute at the final action |
| Model Drift | Performance degradation of a model due to evolving data |
| MLflow | Open-source platform for managing the ML lifecycle (tracking, registry, serving) |
| MLflow Autologging | Automatic capture of params, metrics, and models without explicit code |
| Model Registry | MLflow component for versioning and promoting models (Staging → Production) |
| Model Serving | Deploying a model as a REST endpoint for real-time predictions |
| Model Signature | Input/output schema of an MLflow model for validation |
| Ray Tune | Distributed framework for HPT, integrated in Databricks ML Runtime |
| RMSE | Root Mean Square Error — regression metric |
| scikit-learn | Python ML library, integrated in Databricks ML Runtime |
| Spark MLlib | Native distributed ML library of Spark for large datasets |
| Trial Run | An individual execution in an HPT (with a specific config) |
| Unity Catalog | Databricks centralized governance including the Model Registry |
| Volume | Non-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