ml-pipeline-setup

name	ml-pipeline-setup
description	MLflow and ML Model patterns for Databricks including experiment creation, model training, batch inference, and Unity Catalog integration. Use when implementing ML pipelines, training models with Feature Store, or deploying batch inference jobs. Includes 19 non-negotiable rules covering experiment paths, dataset logging, UC model registration, NaN handling, label binarization, feature engineering workflows, and signature-driven preprocessing.
clients	["ide_cli","genie_code"]
bundle_resource	jobs
deploy_verb	bundle_deploy
deploy_note	Feature/training/batch-inference jobs deploy via `bundle deploy --target dev` (runDatabricksCli on Genie Code); UC model registration in the per-user prefixed catalog/schema. On Genie Code, write generated training/inference .py under the cloned repo root (`{REPO_ROOT}` = `state_file_root` from `skills/vibecoding-state`, e.g. `src/{project}_ml/...`), not a bare relative path — relative paths resolve against the page CWD (see `skills/genie-code-environment` §8).
coverage	full
metadata	{"author":"prashanth subrahmanyam","version":"2.0","domain":"ml","role":"orchestrator","pipeline_stage":8,"pipeline_stage_name":"ml","next_stages":["genai-agents-setup"],"workers":[],"common_dependencies":["databricks-asset-bundles","databricks-expert-agent","databricks-python-imports","naming-tagging-standards","databricks-autonomous-operations"],"consumes":["plans/manifests/ml-manifest.yaml"],"consumes_fallback":"Gold table inventory (self-discovery from catalog)","last_verified":"2026-06-05","volatility":"high","upstream_sources":[{"name":"databricks-agent-skills","repo":"databricks/databricks-agent-skills","paths":"[Truncated]","relationship":"extended","last_synced":"2026-08-30","sync_commit":"ca92a6c"},{"name":"databricks-docs-feature-store","url":"https://docs.databricks.com/aws/en/machine-learning/feature-store/","relationship":"upstream","last_synced":"2026-06-05"},{"name":"databricks-docs-feature-tables-uc","url":"https://docs.databricks.com/aws/en/machine-learning/feature-store/uc/feature-tables-uc","relationship":"upstream","last_synced":"2026-06-05"},{"name":"databricks-docs-feature-store-python-api","url":"https://docs.databricks.com/aws/en/machine-learning/feature-store/python-api","relationship":"upstream","last_synced":"2026-06-05"}]}

MLflow & ML Models Patterns

Phase 0: Read Plan (5 minutes)

Before starting implementation, check for a planning manifest that defines what to build.

import yaml
from pathlib import Path

manifest_path = Path("plans/manifests/ml-manifest.yaml")

if manifest_path.exists():
    with open(manifest_path) as f:
        manifest = yaml.safe_load(f)
    
    # Extract implementation checklist from manifest
    feature_tables = manifest.get('feature_tables', [])
    models = manifest.get('models', [])
    experiments = manifest.get('experiments', [])
    print(f"Plan: {len(feature_tables)} feature tables, {len(models)} models, {len(experiments)} experiments")
    
    # Each model has: name, domain, model_type, algorithm, feature_table,
    #                  label_column, label_type, business_questions
    # Each feature table has: name, primary_keys, source_gold_tables, features
else:
    # Fallback: self-discovery from Gold tables
    print("No manifest found — falling back to Gold table self-discovery")
    # Discover Gold fact tables, infer feature columns, create one model per domain

If manifest exists: Use it as the implementation checklist. Every feature table, model, and experiment is pre-defined with configuration details. Track completion against the manifest's summary counts.

If manifest doesn't exist: Fall back to self-discovery — inventory Gold fact tables, infer feature columns from numeric columns, and create one model per domain. This works but may miss specific label derivations and business context the planning phase would have defined.

---

Quick Start (4-6 hours)

Goal: Build production-ready ML pipelines with MLflow 3.1+, Unity Catalog Model Registry, and Databricks Feature Engineering for training-serving consistency.

What You'll Create:

1. features/create_feature_tables.py - Feature tables in Unity Catalog
2. {domain}/train_{model_name}.py - Training pipelines with Feature Engineering
3. inference/batch_inference_all_models.py - Batch scoring with fe.score_batch
4. Asset Bundle jobs for orchestration

Fast Track:

Client note: IDE runs these in a terminal; Genie Code runs the databricks bundle … commands via runDatabricksCli (be on the bundle's page; generated files anchor to {REPO_ROOT}). See skills/genie-code-environment.

# 1. Create Feature Tables
databricks bundle run ml_feature_pipeline_job -t dev

# 2. Train all models (parallel)
databricks bundle run ml_training_pipeline_job -t dev

# 3. Run batch inference
databricks bundle run ml_inference_pipeline_job -t dev

---

Overview

Production-grade patterns for implementing ML pipelines on Databricks using MLflow, Unity Catalog, and Feature Store. Based on production experience with 25 models across 5 domains, achieving 96% inference success rate and 93% reduction in debugging time.

Pattern Origin: December 2025 (Updated: February 6, 2026 - v5.0)

---

When to Use This Skill

Use this skill when:

• Implementing ML pipelines on Databricks with MLflow tracking
• Training models with Feature Store integration
• Deploying batch inference jobs
• Registering models to Unity Catalog
• Troubleshooting MLflow experiment, model registration, or inference errors
• Setting up Databricks Asset Bundle jobs for ML workflows
• Creating feature tables in Unity Catalog with proper primary keys and NaN handling

Critical for:

• Ensuring training and inference consistency via fe.score_batch
• Preventing common MLflow signature errors
• Handling data quality issues (NaN, label binarization, single-class data)
• Configuring serverless ML jobs correctly

---

Working Memory Management

This orchestrator covers Phase 0 (plan reading) plus multiple implementation sections (feature tables, training, inference, deployment). To maintain coherence without context pollution:

After each major section, persist a brief summary note capturing:

• Phase 0 output: Manifest found (yes/no), model count, feature table count, experiment names from manifest or discovery
• Feature tables output: Feature table names and paths, primary key columns, NaN handling decisions
• Training output: Experiment names, model URIs, MLflow signature details, label binarization strategy
• Inference output: Batch inference notebook paths, fe.score_batch config, output table names
• Jobs output: Job YAML file paths, environment config, databricks.yml sync status

What to keep in working memory: Only the current section's reference skill, the model/feature inventory (from Phase 0), and the previous section's summary note. Discard intermediate outputs (full DataFrames, training logs, model artifacts) — they are in MLflow and reproducible.

---

Architecture Overview

┌─────────────────────────────────────────────────────────────────┐
│                          Gold Layer                              │
│   (fact_tables, dim_tables - source for feature engineering)    │
└───────────────────────┬─────────────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────────────┐
│              Feature Tables (Unity Catalog)                      │
│  ┌───────────────────┐ ┌──────────────────┐ ┌─────────────────┐ │
│  │ cost_features     │ │ security_features│ │ performance_    │ │
│  │ PK: workspace_id, │ │ PK: user_id,     │ │   features      │ │
│  │     usage_date    │ │     event_date   │ │ PK: warehouse_id│ │
│  └───────────────────┘ └──────────────────┘ │     query_date  │ │
│                                             └─────────────────┘ │
└───────────────────────┬─────────────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────────────┐
│                    Training Pipelines                            │
│  ┌─────────────────────────────────────────────────────────────┐│
│  │  FeatureLookup → create_training_set → train → fe.log_model ││
│  │  (Embeds feature metadata for inference consistency)        ││
│  └─────────────────────────────────────────────────────────────┘│
└───────────────────────┬─────────────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────────────┐
│             Unity Catalog Model Registry (MLflow 3.1+)           │
│         catalog.{feature_schema}.{model_name}                    │
│         (Model + Feature Lookup Metadata embedded)               │
└───────────────────────┬─────────────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────────────┐
│                      Inference Layer                             │
│  ┌──────────────────────────────────────────────────────────┐   │
│  │  fe.score_batch(model_uri, df_with_lookup_keys_only)     │   │
│  │  → Automatically retrieves features from Feature Tables  │   │
│  │  → Guarantees training-serving consistency               │   │
│  └──────────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────┘

---

Directory Structure

src/{project}_ml/
├── features/
│   └── create_feature_tables.py     # Feature table creation
├── cost/
│   ├── train_budget_forecaster.py
│   ├── train_cost_anomaly_detector.py
│   └── train_chargeback_attribution.py
├── security/
│   └── train_security_threat_detector.py
├── performance/
│   └── train_query_performance_forecaster.py
├── reliability/
│   └── train_job_failure_predictor.py
├── quality/
│   └── train_data_drift_detector.py
├── inference/
│   └── batch_inference_all_models.py  # Uses fe.score_batch
└── README.md

resources/ml/
├── ml_feature_pipeline_job.yml       # Feature table creation
├── ml_training_pipeline_job.yml      # Training orchestrator
└── ml_inference_pipeline_job.yml     # Batch inference

---

Critical Rules (Quick Reference)

#	Rule	Pattern	Why It Fails Otherwise
0	Pin Package Versions Consistently	Pin the same mlflow / sklearn / xgboost versions in training AND inference (exact pins); the mlflow==3.7.0 value in this skill's templates is a repo-template baseline, not an official Databricks requirement — match the version you actually train with	Version mismatch warnings, deserialization failures, autologging behavior drift
1	Experiment Path	/Shared/{project}_ml_{model_name}	/Users/... fails silently if subfolder doesn't exist
2	Dataset Logging	Inside mlflow.start_run() context	Won't associate with run, invisible in UI
3	Exit Signal	dbutils.notebook.exit("SUCCESS")	Job status unclear, may show SUCCESS on failure
4	UC Model Logging	Prefer fe.log_model() with infer_input_example=True; supply output_schema only when the training_set has no label column (or the model returns a non-default output shape) — current Feature Engineering docs treat output_schema as a fallback, not a universal UC requirement	Unity Catalog model fails to register or returns the wrong output spec at inference
5	Feature Engineering Workflow	FeatureLookup + create_training_set + fe.log_model	Feature skew between training and inference
6	NaN Handling at Source	Clean NaN/Inf at feature table creation with clean_numeric()	sklearn GradientBoosting fails at inference; XGBoost handles NaN but sklearn doesn't
7	Label Binarization	Convert 0-1 rates to binary for classifiers	XGBoostError: base_score must be in (0,1)
8	Single-Class Check	Verify label distribution before training	Classifier can't train on all-same labels
9	Exclude Labels	Use exclude_columns=[LABEL_COLUMN] in create_training_set	Label included as feature causes inference failure
10	Label Type Casting	Cast to INT (classification) or DOUBLE (regression) before training	Type mismatch in model output
11	Double Type Casting	Cast ALL numeric features to DOUBLE in feature tables	MLflow signatures reject DecimalType
12	Lookup Keys Match PKs	lookup_key MUST match Feature Table primary keys EXACTLY	Unable to find feature errors
13	Use fe.score_batch	Use fe.score_batch NOT manual feature joins for inference	Automatic feature retrieval ensures training-serving consistency
14	Feature Registry	Query feature table schemas dynamically	Hardcoded feature lists drift out of sync
15	Custom Inference	Separate task for TF-IDF/NLP models that need runtime features	fe.score_batch() can't compute runtime features
16	Helper Functions Inline	ALWAYS inline helper functions (don't import modules)	ModuleNotFoundError in serverless Asset Bundle notebooks
17	Bundle Path Setup	Use sys.path.insert(0, _bundle_root) pattern	Module imports fail in serverless
18	Standardized Templates	Copy-and-customize from skill templates (don't roll custom)	Custom implementations miss edge cases

---

Core Patterns (Quick Examples)

Experiment Setup

# ✅ CORRECT: Always use /Shared/ path
experiment_name = f"/Shared/{project}_ml_{model_name}"
mlflow.set_experiment(experiment_name)

See: Experiment Patterns for full details

Model Registration with Feature Store

from databricks.feature_engineering import FeatureEngineeringClient
from mlflow.types import ColSpec, DataType, Schema

fe = FeatureEngineeringClient()
mlflow.set_registry_uri("databricks-uc")

# Regression: output_schema = Schema([ColSpec(DataType.double)])
# Classification: output_schema = Schema([ColSpec(DataType.long)])

fe.log_model(
    model=model,
    artifact_path="model",
    flavor=mlflow.sklearn,  # required
    training_set=training_set,
    registered_model_name=f"{catalog}.{schema}.{model_name}",
    infer_input_example=True,  # preferred per current Feature Engineering docs
    output_schema=output_schema,  # required only when training_set has no label, or to override inferred shape
)

See: Model Registry for full patterns by model type

Feature Table Creation with NaN Cleaning

# ✅ CORRECT: Clean NaN/Inf at source for sklearn compatibility
from pyspark.sql.functions import F, isnan
from pyspark.sql.types import DoubleType

def clean_numeric(col_name):
    return F.when(
        F.col(col_name).isNull() | isnan(F.col(col_name)) |
        (F.col(col_name) == float('inf')) | (F.col(col_name) == float('-inf')),
        F.lit(0.0)
    ).otherwise(F.col(col_name))

# Apply to ALL DoubleType columns
for field in df.schema.fields:
    if isinstance(field.dataType, DoubleType):
        df = df.withColumn(field.name, clean_numeric(field.name))

See: Data Quality Patterns for full patterns

Batch Inference with fe.score_batch

from databricks.feature_engineering import FeatureEngineeringClient

fe = FeatureEngineeringClient()

# ✅ CORRECT: scoring_df has ONLY lookup keys (features auto-retrieved)
scoring_df = spark.table(feature_table).select(*lookup_keys).distinct()

predictions_df = fe.score_batch(
    model_uri=model_uri,
    df=scoring_df  # Contains only lookup keys
)

See: Feature Engineering Workflow for full inference patterns

Asset Bundle Job Configuration

resources:
  jobs:
    ml_training_job:
      environments:
        - environment_key: default
          spec:
            environment_version: "4"
            dependencies:
              - "mlflow==3.7.0"  # Pin exact version
              - "xgboost==2.0.3"
      tasks:
        - task_key: train_model
          notebook_task:
            notebook_path: ../../src/ml/models/train.py
            base_parameters:  # NOT parameters with --flags!
              catalog: ${var.catalog}
              model_name: my_model
      # ❌ DO NOT define experiments in Asset Bundle

See: DAB Integration for full patterns

---

Reference Files

Detailed documentation is organized in the references/ directory:

Experiment Patterns

Complete experiment setup, tracking, metric logging, dataset logging, hyperparameter tuning. Covers /Shared/ vs /Users/ paths, run context requirements, helper function inlining, exit signals, and common errors.

Model Registry

Model registration, versioning, aliases, deployment patterns, serving endpoints. Covers Unity Catalog integration with fe.log_model(), output schema patterns by model type (regression, classification, anomaly detection), signature-driven preprocessing, and model loading from UC. Documents both output_schema (primary) and infer_signature (alternative) approaches.

DAB Integration

Asset Bundle integration, notebook patterns, inline helpers, parameter passing. Covers training and inference job configuration, package version pinning, base_parameters vs argparse, serverless environment setup, and common deployment errors.

Feature Store Patterns

Feature table creation, feature lookup configuration, column conflict resolution, Feature Registry pattern for dynamic schema querying. Covers training set creation, feature lookups, and inference patterns.

Data Quality Patterns

NaN/Inf handling at feature table source, label binarization for XGBoost classifiers, single-class data detection, feature column exclusion. Covers sklearn vs XGBoost compatibility, preprocessing requirements, and training/inference checklists.

Troubleshooting

Comprehensive error reference table, schema verification patterns, SCD2 vs regular dimension table handling, pre-development checklist. Covers common MLflow errors, signature issues, and debugging workflows.

Requirements Template

Fill-in-the-blank requirements template for ML projects. Includes project context (catalog, schemas), feature table inventory (primary keys, features, source tables), model inventory (type, algorithm, label column, label type), and label type reference (regression vs classification vs anomaly detection casting).

Feature Engineering Workflow