databricks-ml-pipeline

name	databricks-ml-pipeline
description	End-to-end machine learning pipelines on Databricks including data exploration, feature engineering, model training with hyperparameter optimization, MLflow experiment tracking, model registration to Unity Catalog, and deployment as DABs. Use when building ML workflows, training models, or deploying ML pipelines.
allowed-tools	["Bash","Read","Write","Edit","Grep","Glob"]
model	claude-sonnet-4-5-20250929
user-invocable	true

Databricks ML Pipeline Builder

Build complete machine learning pipelines from data exploration through model deployment. Orchestrates testing, Unity Catalog setup, and bundle deployment for production ML workflows.

When to Use This Skill

• Building end-to-end ML pipelines
• Training models with hyperparameter optimization
• Setting up MLflow experiment tracking
• Registering models to Unity Catalog
• Deploying ML training pipelines
• Scheduling periodic model retraining
• Feature engineering workflows

Complete ML Workflow

This skill orchestrates other skills to deliver complete ML pipelines:

1. Setup (databricks-unity-catalog skill)

   • Create catalog schema for ML assets
   • Set up feature store schema
   • Prepare model registry

2. Exploration (databricks-testing skill)

   • Load and profile data on cluster
   • Test feature engineering logic
   • Validate data quality

3. Training (databricks-testing skill)

   • Test training code interactively
   • Hyperparameter tuning
   • MLflow experiment tracking
   • Model validation

4. Registration

   • Register best model to Unity Catalog
   • Add model metadata and tags
   • Set model aliases (Champion/Challenger)

5. Deployment (databricks-bundle-deploy skill)

   • Package as Databricks Asset Bundle
   • Deploy to dev/staging/prod
   • Schedule periodic retraining

Phase 1: Setup & Data Exploration

Step 1: Create ML Schemas

Use databricks-unity-catalog skill to set up catalog structure:

# Create schema for ML models
create_schema(
    catalog_name="ml_dev",
    schema_name="churn_prediction",
    comment="Churn prediction model v2.0. Features, training data, and model registry."
)

# Verify
get_schema(full_schema_name="ml_dev.churn_prediction")

Step 2: Load and Profile Data

Use databricks-testing skill to explore data on cluster:

# Test data loading
databricks_command(
    cluster_id="0123-456789-abc123",
    language="python",
    code="""
# Load customer data
df = spark.table("ml_dev.bronze.customer_transactions")

print(f"Total records: {df.count()}")
print(f"Date range: {df.select('transaction_date').agg({'transaction_date': 'min'}).collect()[0][0]} to {df.select('transaction_date').agg({'transaction_date': 'max'}).collect()[0][0]}")

# Profile data
df.describe().show()
df.groupBy('customer_status').count().show()

# Check for nulls
from pyspark.sql import functions as F
null_counts = df.select([F.sum(F.when(F.col(c).isNull(), 1).otherwise(0)).alias(c) for c in df.columns])
null_counts.show()
"""
)

Phase 2: Feature Engineering

Step 3: Test Feature Transformations

Use databricks-testing skill with stateful context:

# Create context for iterative development
context_id = create_context(
    cluster_id="0123-456789-abc123",
    language="python"
)

# Load data (persists in context)
execute_command_with_context(
    cluster_id="0123-456789-abc123",
    context_id=context_id,
    code="""
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# Load customer transactions
df = spark.table("ml_dev.bronze.customer_transactions")
print(f"Loaded {df.count()} transactions")
"""
)

# Create features (uses df from previous step)
execute_command_with_context(
    cluster_id="0123-456789-abc123",
    context_id=context_id,
    code="""
# Time-based features
w = Window.partitionBy("customer_id").orderBy("transaction_date")

features_df = df.groupBy("customer_id").agg(
    F.count("*").alias("transaction_count"),
    F.sum("amount").alias("total_spend"),
    F.avg("amount").alias("avg_transaction"),
    F.datediff(F.current_date(), F.max("transaction_date")).alias("days_since_last"),
    F.count(F.when(F.col("transaction_date") >= F.date_sub(F.current_date(), 30), 1)).alias("recent_transactions")
)

# Add churn label (no transactions in last 90 days)
features_df = features_df.withColumn(
    "churned",
    F.when(F.col("days_since_last") > 90, 1).otherwise(0)
)

print("Features created:")
features_df.show(10)

# Check class balance
features_df.groupBy("churned").count().show()
"""
)

# Save features (uses features_df from previous step)
execute_command_with_context(
    cluster_id="0123-456789-abc123",
    context_id=context_id,
    code="""
# Save to feature table
features_df.write \\
    .format("delta") \\
    .mode("overwrite") \\
    .saveAsTable("ml_dev.churn_prediction.customer_features")

print("Features saved to ml_dev.churn_prediction.customer_features")
"""
)

# Cleanup context
destroy_context(cluster_id="0123-456789-abc123", context_id=context_id)

Phase 3: Model Training

Step 4: Test Training Code

Use databricks-testing skill to validate training logic:

databricks_command(
    cluster_id="0123-456789-abc123",
    language="python",
    code="""
# MAGIC %pip install mlflow scikit-learn

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

# Load features
features_df = spark.table("ml_dev.churn_prediction.customer_features").toPandas()

# Prepare data
feature_cols = ["transaction_count", "total_spend", "avg_transaction", "days_since_last", "recent_transactions"]
X = features_df[feature_cols]
y = features_df["churned"]

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

# Train baseline model
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# Evaluate
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
precision = precision_score(y_test, y_pred)
recall = recall_score(y_test, y_pred)
f1 = f1_score(y_test, y_pred)

print(f"Baseline Model Performance:")
print(f"  Accuracy: {accuracy:.4f}")
print(f"  Precision: {precision:.4f}")
print(f"  Recall: {recall:.4f}")
print(f"  F1 Score: {f1:.4f}")
"""
)

Phase 4: MLflow Experiment Tracking

Complete Training Notebook with MLflow

Template for ML training notebook:

# Databricks notebook source
# MAGIC %md
# MAGIC # Churn Prediction Model Training
# MAGIC
# MAGIC Trains Random Forest model with hyperparameter optimization and MLflow tracking

# COMMAND ----------
# MAGIC %pip install mlflow scikit-learn

# COMMAND ----------
# Widget parameterization
try:
    catalog = dbutils.widgets.get("catalog")
except:
    catalog = "ml_dev"

try:
    schema = dbutils.widgets.get("schema")
except:
    schema = "churn_prediction"

try:
    experiment_name = dbutils.widgets.get("experiment_name")
except:
    experiment_name = f"/Users/{spark.sql('SELECT current_user()').collect()[0][0]}/experiments/churn_model"

print(f"Training with parameters:")
print(f"  Catalog: {catalog}")
print(f"  Schema: {schema}")
print(f"  Experiment: {experiment_name}")

# COMMAND ----------
import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score

# Set MLflow experiment
mlflow.set_experiment(experiment_name)

# COMMAND ----------
# Load features
features_df = spark.table(f"{catalog}.{schema}.customer_features").toPandas()

# Prepare data
feature_cols = ["transaction_count", "total_spend", "avg_transaction", "days_since_last", "recent_transactions"]
X = features_df[feature_cols]
y = features_df["churned"]

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

print(f"Training set: {len(X_train)} samples")
print(f"Test set: {len(X_test)} samples")
print(f"Churn rate: {y.mean():.2%}")

# COMMAND ----------
# Hyperparameter tuning with MLflow
params_grid = {
    "n_estimators": [50, 100, 200],
    "max_depth": [5, 10, 20, None],
    "min_samples_split": [2, 5, 10]
}

best_score = 0
best_run_id = None
best_params = None

for n_est in params_grid["n_estimators"]:
    for depth in params_grid["max_depth"]:
        for split in params_grid["min_samples_split"]:
            with mlflow.start_run(run_name=f"RF_n{n_est}_d{depth}_s{split}"):
                # Log parameters
                mlflow.log_param("n_estimators", n_est)
                mlflow.log_param("max_depth", depth if depth else "None")
                mlflow.log_param("min_samples_split", split)

                # Train model
                model = RandomForestClassifier(
                    n_estimators=n_est,
                    max_depth=depth,
                    min_samples_split=split,
                    random_state=42
                )

                # Cross-validation
                cv_scores = cross_val_score(model, X_train, y_train, cv=5, scoring='f1')
                cv_mean = cv_scores.mean()
                cv_std = cv_scores.std()

                # Log CV metrics
                mlflow.log_metric("cv_f1_mean", cv_mean)
                mlflow.log_metric("cv_f1_std", cv_std)

                # Train on full training set
                model.fit(X_train, y_train)

                # Test set evaluation
                y_pred = model.predict(X_test)
                y_pred_proba = model.predict_proba(X_test)[:, 1]

                accuracy = accuracy_score(y_test, y_pred)
                precision = precision_score(y_test, y_pred)
                recall = recall_score(y_test, y_pred)
                f1 = f1_score(y_test, y_pred)
                auc = roc_auc_score(y_test, y_pred_proba)

                # Log test metrics
                mlflow.log_metric("test_accuracy", accuracy)
                mlflow.log_metric("test_precision", precision)
                mlflow.log_metric("test_recall", recall)
                mlflow.log_metric("test_f1", f1)
                mlflow.log_metric("test_auc", auc)

                # Log model
                mlflow.sklearn.log_model(model, "model")

                # Track best model by F1 score
                if f1 > best_score:
                    best_score = f1
                    best_run_id = mlflow.active_run().info.run_id
                    best_params = {
                        "n_estimators": n_est,
                        "max_depth": depth,
                        "min_samples_split": split
                    }

print(f"\\nBest Model:")
print(f"  Run ID: {best_run_id}")
print(f"  F1 Score: {best_score:.4f}")
print(f"  Parameters: {best_params}")

# COMMAND ----------
# Register best model to Unity Catalog
model_name = f"{catalog}.{schema}.churn_prediction_model"

model_uri = f"runs:/{best_run_id}/model"
registered_model = mlflow.register_model(model_uri, model_name)

print(f"Model registered as: {model_name}")
print(f"Version: {registered_model.version}")

# COMMAND ----------
# Set model alias
from mlflow.tracking import MlflowClient

client = MlflowClient()
client.set_registered_model_alias(