ml-pipeline-setup

# MLflow & ML Models Patterns ## Phase 0: Read Plan (5 minutes) **Before starting implementation, check for a planning manifest that defines what to build.** ```python 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`. ```bash # 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 ```python # ✅ CORRECT: Always use /Shared/ path experiment_name = f"/Shared/{project}_ml_{model_name}" mlflow.set_experiment(experiment_name) ``` **See:** [Experiment Patterns](references/experiment-patterns.md) for full details ### Model Registration with Feature Store ```python 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](references/model-registry.md) for full patterns by model type ### Feature Table Creation with NaN Cleaning ```python # ✅ 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](references/data-quality-patterns.md) for full patterns ### Batch Inference with fe.score_batch ```python 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](references/feature-engineering-workflow.md) for full inference patterns ### Asset Bundle Job Configuration ```yaml 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](references/dab-integration.md) for full patterns --- ## Reference Files Detailed documentation is organized in the `references/` directory: ### **[Experiment Patterns](references/experiment-patterns.md)** 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](references/model-registry.md)** 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](references/dab-integration.md)** 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](references/feature-store-patterns.md)** 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](references/data-quality-patterns.md)** 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](references/troubleshooting.md)** 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](references/requirements-template.md)** 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](references/feature-engineering-workflow.md)**