--- name: databricks-core-workflow-b-b212124d0111 description: | Execute Databricks secondary workflow: MLflow model training and deployment. Use when building ML pipelines, training models, or deploying to production. Trigger with phrases like "databricks ML", "mlflow training", "databricks model", "feature store", "model registry". allowed-tools: Read, Write, Edit, Bash(databricks:*), Grep version: 1.0.0 license: MIT author: Jeremy Longshore --- # Databricks Core Workflow B: MLflow Training ## Overview Build ML pipelines with MLflow experiment tracking, model registry, and deployment. ## Prerequisites - Completed `databricks-install-auth` setup - Familiarity with `databricks-core-workflow-a` (data pipelines) - MLflow and scikit-learn installed - Unity Catalog for model registry (recommended) ## Instructions ### Step 1: Feature Engineering with Feature Store ```python # src/ml/features.py from databricks.feature_engineering import FeatureEngineeringClient from pyspark.sql import SparkSession, DataFrame from pyspark.sql.functions import col, avg, count, datediff, current_date def create_customer_features(spark: SparkSession, orders_table: str) -> DataFrame: """Create customer features from order history.""" orders_df = spark.table(orders_table) features_df = ( orders_df .groupBy("customer_id") .agg( count("*").alias("total_orders"), avg("amount").alias("avg_order_value"), sum("amount").alias("lifetime_value"), max("order_date").alias("last_order_date"), min("order_date").alias("first_order_date"), ) .withColumn( "days_since_last_order", datediff(current_date(), col("last_order_date")) ) .withColumn( "customer_tenure_days", datediff(current_date(), col("first_order_date")) ) ) return features_df def register_feature_table( spark: SparkSession, df: DataFrame, feature_table_name: str, primary_keys: list[str], description: str, ) -> None: """Register DataFrame as Feature Store table.""" fe = FeatureEngineeringClient() fe.create_table( name=feature_table_name, primary_keys=primary_keys, df=df, description=description, tags={"team": "data-science", "domain": "customer"}, ) ``` ### Step 2: MLflow Experiment Tracking ```python # src/ml/training.py import mlflow import mlflow.sklearn from sklearn.model_selection import train_test_split from sklearn.ensemble import RandomForestClassifier from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score import pandas as pd def train_churn_model( spark, feature_table: str, label_column: str = "churned", experiment_name: str = "/Experiments/churn-prediction", ) -> str: """ Train churn prediction model with MLflow tracking. Returns: run_id: MLflow run ID """ # Set experiment mlflow.set_experiment(experiment_name) # Load features fe = FeatureEngineeringClient() df = spark.table(feature_table).toPandas() # Prepare data feature_cols = [c for c in df.columns if c not in [label_column, "customer_id"]] X = df[feature_cols] y = df[label_column] X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.2, random_state=42 ) # Train with MLflow tracking with mlflow.start_run() as run: # Log parameters params = { "n_estimators": 100, "max_depth": 10, "min_samples_split": 5, "random_state": 42, } mlflow.log_params(params) # Train model model = RandomForestClassifier(**params) model.fit(X_train, y_train) # Evaluate y_pred = model.predict(X_test) metrics = { "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), } mlflow.log_metrics(metrics) # Log model with signature from mlflow.models import infer_signature signature = infer_signature(X_train, y_pred) mlflow.sklearn.log_model( model, artifact_path="model", signature=signature, input_example=X_train.head(5), registered_model_name="churn-prediction-model", ) # Log feature importance importance_df = pd.DataFrame({ "feature": feature_cols, "importance": model.feature_importances_ }).sort_values("importance", ascending=False) mlflow.log_table(importance_df, "feature_importance.json") print(f"Run ID: {run.info.run_id}") print(f"Metrics: {metrics}") return run.info.run_id ``` ### Step 3: Model Registry and Versioning ```python # src/ml/registry.py from mlflow import MlflowClient from mlflow.entities.model_registry import ModelVersion def promote_model( model_name: str, version: int, stage: str = "Production", archive_existing: bool = True, ) -> ModelVersion: """ Promote model version to specified stage. Args: model_name: Registered model name version: Model version number stage: Target stage (Staging, Production, Archived) archive_existing: Archive current production model """ client = MlflowClient() # Archive existing production model if archive_existing and stage == "Production": for mv in client.search_model_versions(f"name='{model_name}'"): if mv.current_stage == "Production": client.transition_model_version_stage( name=model_name, version=mv.version, stage="Archived", ) # Promote new version model_version = client.transition_model_version_stage( name=model_name, version=version, stage=stage, ) # Add description client.update_model_version( name=model_name, version=version, description=f"Promoted to {stage} on {pd.Timestamp.now()}", ) return model_version def compare_model_versions(model_name: str) -> pd.DataFrame: """Compare metrics across model versions.""" client = MlflowClient() versions = client.search_model_versions(f"name='{model_name}'") comparisons = [] for v in versions: run = client.get_run(v.run_id) comparisons.append({ "version": v.version, "stage": v.current_stage, "accuracy": run.data.metrics.get("accuracy"), "f1": run.data.metrics.get("f1"), "created": v.creation_timestamp, }) return pd.DataFrame(comparisons) ``` ### Step 4: Model Serving and Inference ```python # src/ml/serving.py from databricks.sdk import WorkspaceClient from databricks.sdk.service.serving import ( EndpointCoreConfigInput, ServedEntityInput, TrafficConfig, Route, ) def deploy_model_endpoint( model_name: str, endpoint_name: str, model_version: str = None, # None = latest scale_to_zero: bool = True, workload_size: str = "Small", ) -> str: """ Deploy model to Databricks Model Serving endpoint. Returns: endpoint_url: Serving endpoint URL """ w = WorkspaceClient() # Create or update endpoint endpoint = w.serving_endpoints.create_and_wait( name=endpoint_name, config=EndpointCoreConfigInput( served_entities=[ ServedEntityInput( name=f"{model_name}-entity", entity_name=model_name, entity_version=model_version, workload_size=workload_size, scale_to_zero_enabled=scale_to_zero, ) ], traffic_config=TrafficConfig( routes=[ Route( served_model_name=f"{model_name}-entity", traffic_percentage=100, ) ] ), ), ) return endpoint.name def batch_inference( spark, model_uri: str, input_table: str, output_table: str, feature_columns: list[str], ) -> None: """Run batch inference with logged model.""" import mlflow.pyfunc # Load model model = mlflow.pyfunc.spark_udf(spark, model_uri) # Read input data input_df = spark.table(input_table) # Run inference predictions_df = input_df.withColumn( "prediction", model(*[col(c) for c in feature_columns]) ) # Write predictions predictions_df.write \ .format("delta") \ .mode("overwrite") \ .saveAsTable(output_table) ``` ## Output - Feature table in Unity Catalog - MLflow experiment with tracked runs - Registered model with versions - Model serving endpoint ## Error Handling | Error | Cause | Solution | |-------|-------|----------| | `Model not found` | Wrong model name/version | Verify in Model Registry | | `Feature mismatch` | Schema changed | Retrain with updated features | | `Endpoint timeout` | Cold start | Disable scale-to-zero for latency | | `Memory error` | Large batch | Reduce batch size or increase cluster | ## Examples ### Complete ML Pipeline Job ```python # jobs/train_pipeline.py from src.ml import features, training, registry # 1. Create features features_df = features.create_customer_features(spark, "catalog.silver.orders") features.register_feature_table( spark, features_df, "catalog.ml.customer_features", ["customer_id"], "Customer behavior features" ) # 2. Train model run_id = training.train_churn_model( spark, "catalog.ml.customer_features", experiment_name="/Experiments/churn-v2" ) # 3. Compare and promote comparison = registry.compare_model_versions("churn-prediction-model") best_version = comparison.sort_values("f1", ascending=False).iloc[0]["version"] registry.promote_model("churn-prediction-model", best_version, "Production") ``` ## Resources - [MLflow on Databricks](https://docs.databricks.com/mlflow/index.html) - [Feature Engineering](https://docs.databricks.com/machine-learning/feature-store/index.html) - [Model Serving](https://docs.databricks.com/machine-learning/model-serving/index.html) - [Unity Catalog ML](https://docs.databricks.com/machine-learning/manage-model-lifecycle/index.html) ## Next Steps For common errors, see `databricks-common-errors`. --- Source attribution: Dicklesworthstone Original project: https://github.com/Dicklesworthstone/pi_agent_rust/blob/f28fc46ee0db8e710d95b8b55592cd38d60f2b26/tests/ext_conformance/artifacts/plugins-community/plugins/saas-packs/databricks-pack/skills/databricks-core-workflow-b/SKILL.md License: MIT License information: https://opensource.org/license/mit Original upstream instructions; downloaded from Dave AI. ## Upstream license notice MIT License (with OpenAI/Anthropic Rider) Copyright (c) 2026 Jeffrey Emanuel Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: ADDITIONAL RIDER / RESTRICTION (OpenAI / Anthropic): This rider is part of the "conditions" of this License. In the event of any conflict between this rider and any other portion of this License, this rider controls. "Restricted Parties" means OpenAI, L.L.C.; Anthropic, PBC; any of their respective Affiliates; and any person or entity acting directly or indirectly on behalf of, for the benefit of, or under the direction of any of the foregoing (including any officer, director, employee, contractor, agent, consultant, service provider, or representative). Notwithstanding any other provision of this License, no rights are granted to any Restricted Party. Any purported license, sublicense, assignment, transfer, or other permission to any Restricted Party is null and void absent the express prior written permission of Jeffrey Emanuel. You may not provide, disclose, distribute, sublicense, sell, lease, lend, host, make available, or otherwise permit access to the Software or any derivative work of the Software (as defined in applicable copyright law) (collectively, "Derivative Works") to or for any Restricted Party. For purposes of this rider, "use" includes, without limitation: copying, modifying, merging, publishing, distributing, sublicensing, selling, transferring, making available, hosting, deploying, executing, benchmarking, testing, analyzing, indexing, or incorporating the Software or any Derivative Works into any dataset, training corpus, evaluation harness, or pipeline for machine learning or other automated systems. This rider applies to the Software and all Derivative Works. As a condition of use, you agree that this rider is a precondition to exercising any rights under this License, and you agree that any distribution of the Software or any Derivative Works must include this rider provision unmodified. Any breach of this rider automatically and immediately terminates the permissions granted by this License. Upon termination, you must immediately cease all use and distribution of the Software and any Derivative Works and destroy all copies under your control. You agree that a breach of this rider would cause irreparable harm and that Jeffrey Emanuel may seek injunctive or other equitable relief to enforce this rider, in addition to any other remedies available at law. To the maximum extent permitted by applicable law, the prevailing party in any action to enforce this rider shall be entitled to recover reasonable attorneys' fees and costs. For purposes of this rider, "Affiliate" means any entity that directly or indirectly controls, is controlled by, or is under common control with the specified party. "Control" means ownership of more than 50% of the voting securities or other ownership interest, or the power to direct management or policies by contract or otherwise. The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. # Dave AI resource metadata Creator / publisher: Dicklesworthstone Original project: https://github.com/Dicklesworthstone/pi_agent_rust/blob/f28fc46ee0db8e710d95b8b55592cd38d60f2b26/tests/ext_conformance/artifacts/plugins-community/plugins/saas-packs/databricks-pack/skills/databricks-core-workflow-b/SKILL.md Attribution basis: Source repository record Fields of work: Data & Analytics, Development & Engineering Roles: Data Analyst, Data Scientist, ML / AI Engineer Specialties: API, Backend Tools: Databricks, Mlflow Country applicability: Country not specified Classification basis: Source catalog labels Instruction SHA-256: f8ed86087f45dd83771892931e71c39531a37d2546335e355a4fb535de30596c