ML Pipeline with MLflow

This lesson covers the ML lifecycle for power prediction: training models, tracking experiments, managing versions in the model registry, and deploying the champion model for batch inference. By the end, you'll understand how MLflow fits into your architecture and what decisions you make as Tech Lead.

MLflow in Databricks — The ML Operating System

MLflow is integrated directly into Databricks and handles four concerns:^[1]

Training a Wind Power Model

Here's the pattern for training with MLflow tracking and Feature Store integration:

import mlflow
from mlflow.models import infer_signature
from databricks.feature_engineering import FeatureEngineeringClient, FeatureLookup
import lightgbm as lgb
from sklearn.metrics import mean_absolute_error, mean_squared_error
import numpy as np

# 1. Set up experiment
mlflow.set_experiment("/Shared/wind-power-prediction")
fe = FeatureEngineeringClient()

# 2. Create training set from Feature Store (point-in-time correct)
labels_df = spark.sql("""
    SELECT asset_id, timestamp, power_kw as actual_power_kw
    FROM prod.silver.power_ts
    WHERE timestamp BETWEEN '2025-01-01' AND '2026-06-30'
      AND power_kw IS NOT NULL
""")

training_set = fe.create_training_set(
    df=labels_df,
    feature_lookups=[
        FeatureLookup(
            table_name="prod.features.wind_power_features",
            lookup_key=["asset_id", "timestamp"]
        )
    ],
    label="actual_power_kw"
)

df = training_set.load_df().toPandas()

# 3. Split (temporal — never random for time series!)
split_date = "2026-04-01"
train = df[df["timestamp"] < split_date]
test = df[df["timestamp"] >= split_date]

feature_cols = [c for c in train.columns 
                if c not in ["actual_power_kw", "asset_id", "timestamp"]]
X_train, y_train = train[feature_cols], train["actual_power_kw"]
X_test, y_test = test[feature_cols], test["actual_power_kw"]

# 4. Train with MLflow tracking
with mlflow.start_run(run_name="lgbm_v1_all_features") as run:
    params = {
        "n_estimators": 500,
        "max_depth": 8,
        "learning_rate": 0.05,
        "num_leaves": 63,
        "min_child_samples": 50,
        "subsample": 0.8,
    }
    mlflow.log_params(params)
    mlflow.log_param("feature_count", len(feature_cols))
    mlflow.log_param("train_size", len(train))
    mlflow.log_param("test_size", len(test))
    mlflow.log_param("split_date", split_date)

    model = lgb.LGBMRegressor(**params)
    model.fit(X_train, y_train)

    # Evaluate
    y_pred = model.predict(X_test)
    mae = mean_absolute_error(y_test, y_pred)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    nmae = mae / df["actual_power_kw"].mean()  # normalised MAE

    mlflow.log_metrics({
        "mae_kw": mae,
        "rmse_kw": rmse,
        "nmae": nmae,
        "r2": model.score(X_test, y_test),
    })

    # Log model with Feature Store (enables automatic feature lookup at inference)
    fe.log_model(
        model=model,
        artifact_path="model",
        flavor=mlflow.lightgbm,
        training_set=training_set,
        registered_model_name="prod.ml.wind_power_model"
    )

The Model Registry

Every call to log_model with a registered_model_name creates a new version in the registry. The registry tracks all versions and supports aliases for deployment:

from mlflow import MlflowClient

client = MlflowClient()

# Promote a version to production using aliases
client.set_registered_model_alias(
    name="prod.ml.wind_power_model",
    alias="champion",
    version=3  # the version that passed validation
)

# Set a challenger for A/B testing or shadow scoring
client.set_registered_model_alias(
    name="prod.ml.wind_power_model",
    alias="challenger",
    version=4
)

Batch Inference Pipeline

For weather-to-power, batch inference runs hourly (after fresh weather arrives). It loads the champion model, looks up features from the Feature Store, and writes predictions to Gold:

import mlflow
from databricks.feature_engineering import FeatureEngineeringClient

fe = FeatureEngineeringClient()

# 1. Load champion model
model_uri = "models:/prod.ml.wind_power_model@champion"

# 2. Prepare scoring data (asset × future timestamps to predict)
scoring_df = spark.sql("""
    SELECT DISTINCT asset_id, timestamp
    FROM prod.features.wind_power_features
    WHERE timestamp > current_timestamp()
      AND timestamp <= current_timestamp() + INTERVAL 72 HOURS
""")

# 3. Score using Feature Store (automatic feature lookup!)
predictions = fe.score_batch(
    model_uri=model_uri,
    df=scoring_df
)
# Result has: asset_id, timestamp, prediction (power_kw)

# 4. Add metadata and write to Gold
gold_df = predictions.select(
    "asset_id",
    "timestamp",
    F.col("prediction").alias("forecast_power_kw"),
    F.current_timestamp().alias("produced_at"),
    F.lit("prod.ml.wind_power_model@champion").alias("model_version")
)

gold_df.write.format("delta") \
    .mode("append") \
    .saveAsTable("prod.gold.forecast_ts")

Model Validation Gate

Before promoting a new model version to champion, validate it against the current champion:

def validate_challenger(challenger_version: int, threshold_mae_improvement: float = 0.05):
    """
    Compare challenger model against current champion on recent data.
    Returns True if challenger should be promoted.
    """
    # Score recent actuals with both models
    recent = spark.sql("""
        SELECT asset_id, timestamp, power_kw as actual
        FROM prod.silver.power_ts
        WHERE timestamp > current_timestamp() - INTERVAL 7 DAYS
          AND power_kw IS NOT NULL
    """)
    
    champion_preds = fe.score_batch(
        model_uri="models:/prod.ml.wind_power_model@champion",
        df=recent.select("asset_id", "timestamp")
    )
    
    challenger_preds = fe.score_batch(
        model_uri=f"models:/prod.ml.wind_power_model/{challenger_version}",
        df=recent.select("asset_id", "timestamp")
    )
    
    # Compare MAE
    champion_mae = compute_mae(recent, champion_preds)
    challenger_mae = compute_mae(recent, challenger_preds)
    
    improvement = (champion_mae - challenger_mae) / champion_mae
    
    mlflow.log_metrics({
        "validation_champion_mae": champion_mae,
        "validation_challenger_mae": challenger_mae,
        "validation_improvement_pct": improvement * 100,
    })
    
    return improvement >= threshold_mae_improvement

Per-Asset vs. Global Models

An architectural decision you'll face: one model for all assets, or one model per asset/park?

Check Your Understanding

Primary Source

What's Next

The model produces predictions. The next lesson covers the Gold layer and downstream serving — how to structure the forecast output, compute confidence intervals, aggregate turbine→park, and serve results to trading systems.

Questions about model selection (LightGBM vs. neural nets for wind), ensemble strategies, or retraining frequency? Ask your teacher.

Rechtliches: Impressum