databricks-ml-pipeline

# 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: ```python # 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: ```python # 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: ```python # 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: ```python 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:** ```python # 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(