Ch 1 — ML Infrastructure¶

"""
mlflow_training.py — MLflow experiment tracking example.

Requirements:
    pip install mlflow scikit-learn pandas matplotlib
"""
from __future__ import annotations

import os

import matplotlib.pyplot as plt
import mlflow
import mlflow.sklearn
import pandas as pd
from mlflow.models import infer_signature
from sklearn.datasets import load_breast_cancer
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
from sklearn.model_selection import train_test_split


def train_and_track(
    n_estimators: int = 200,
    max_depth: int = 5,
    learning_rate: float = 0.1,
    subsample: float = 0.8,
    experiment_name: str = "breast-cancer-gbt",
) -> str:
    """
    Train a Gradient Boosting classifier with full MLflow tracking.

    Returns:
        The MLflow run ID for downstream reference.
    """
    mlflow.set_tracking_uri(
        os.environ.get("MLFLOW_TRACKING_URI", "http://localhost:5000")
    )
    mlflow.set_experiment(experiment_name)

    data = load_breast_cancer()
    X_train, X_test, y_train, y_test = train_test_split(
        data.data,
        data.target,
        test_size=0.2,
        random_state=42,
        stratify=data.target,
    )

    with mlflow.start_run() as run:
        # ── Log hyperparameters ──────────────────────────────────────────────
        mlflow.log_params(
            {
                "n_estimators": n_estimators,
                "max_depth": max_depth,
                "learning_rate": learning_rate,
                "subsample": subsample,
                "dataset": "sklearn:breast_cancer",
                "n_features": data.data.shape[1],
                "n_train_samples": len(X_train),
            }
        )

        # ── Train ────────────────────────────────────────────────────────────
        model = GradientBoostingClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            learning_rate=learning_rate,
            subsample=subsample,
            random_state=42,
        )
        model.fit(X_train, y_train)

        # ── Evaluate ─────────────────────────────────────────────────────────
        y_pred = model.predict(X_test)
        y_proba = model.predict_proba(X_test)[:, 1]
        metrics = {
            "accuracy": accuracy_score(y_test, y_pred),
            "f1": f1_score(y_test, y_pred),
            "roc_auc": roc_auc_score(y_test, y_proba),
        }
        mlflow.log_metrics(metrics)

        # ── Log artefacts ────────────────────────────────────────────────────
        fig, ax = plt.subplots(figsize=(8, 5))
        importances = pd.Series(
            model.feature_importances_, index=data.feature_names
        )
        importances.sort_values().tail(10).plot.barh(ax=ax)
        ax.set_title("Top 10 Feature Importances")
        fig.tight_layout()
        mlflow.log_figure(fig, "feature_importances.png")
        plt.close(fig)

        # ── Register model ───────────────────────────────────────────────────
        signature = infer_signature(X_train, model.predict(X_train))
        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path="model",
            signature=signature,
            input_example=X_test[:3],
            registered_model_name="breast-cancer-gbt",
        )

        print(f"Run ID: {run.info.run_id}  |  Metrics: {metrics}")
        return run.info.run_id


def load_production_model(model_name: str = "breast-cancer-gbt") -> object:
    """Load the model currently in the Production stage of the registry."""
    return mlflow.sklearn.load_model(f"models:/{model_name}/Production")


if __name__ == "__main__":
    for lr in [0.05, 0.1, 0.2]:
        for depth in [3, 5, 7]:
            train_and_track(n_estimators=200, max_depth=depth, learning_rate=lr)

"""
wandb_training.py — W&B experiment tracking with hyperparameter sweeps.

Requirements:
    pip install wandb torch transformers datasets
"""
from __future__ import annotations

import wandb


def train_one_run(config: dict | None = None) -> None:
    """Train a single model run, optionally driven by a W&B sweep."""
    with wandb.init(config=config) as run:
        cfg = wandb.config

        # Build model with sweep-provided hyperparameters
        model = build_transformer_model(
            model_name=cfg.model_name,
            dropout=cfg.dropout,
        )
        optimizer = build_optimizer(model, lr=cfg.learning_rate, wd=cfg.weight_decay)
        train_loader, val_loader = get_dataloaders(batch_size=cfg.batch_size)

        for epoch in range(cfg.epochs):
            model.train()
            for step, batch in enumerate(train_loader):
                loss, acc = train_step(model, optimizer, batch)
                wandb.log(
                    {
                        "train/loss": loss,
                        "train/accuracy": acc,
                        "epoch": epoch,
                        "global_step": epoch * len(train_loader) + step,
                    }
                )

            val_metrics = evaluate(model, val_loader)
            wandb.log(
                {
                    "val/loss": val_metrics["loss"],
                    "val/accuracy": val_metrics["accuracy"],
                    "val/f1": val_metrics["f1"],
                    "epoch": epoch,
                }
            )

        # Save best checkpoint as a W&B artefact
        artifact = wandb.Artifact(
            name=f"model-{run.id}",
            type="model",
            metadata={"val_accuracy": val_metrics["accuracy"]},
        )
        artifact.add_file("checkpoint_best.pt")
        run.log_artifact(artifact)

import wandb

sweep_config: dict = {
    "method": "bayes",
    "metric": {"name": "val/accuracy", "goal": "maximize"},
    "parameters": {
        "model_name": {"value": "bert-base-uncased"},
        "learning_rate": {
            "distribution": "log_uniform_values",
            "min": 1e-5,
            "max": 1e-3,
        },
        "batch_size": {"values": [16, 32, 64]},
        "dropout": {"distribution": "uniform", "min": 0.0, "max": 0.5},
        "weight_decay": {"distribution": "uniform", "min": 0.0, "max": 0.1},
        "epochs": {"value": 5},
    },
    "early_terminate": {"type": "hyperband", "min_iter": 2},
}

sweep_id = wandb.sweep(sweep_config, project="transformer-classification")
# Launch 30 agents to explore the parameter space
wandb.agent(sweep_id, function=train_one_run, count=30)

"""
feast_features.py — Point-in-time correct feature store with Feast.

Requirements:
    pip install feast[redis]>=0.38
"""
from datetime import timedelta

import pandas as pd
from feast import Entity, FeatureService, FeatureStore, FeatureView, Field
from feast.infra.offline_stores.file_source import FileSource
from feast.types import Float32, Int64, String

# ── Entity ───────────────────────────────────────────────────────────────────
user = Entity(name="user_id", description="Unique user identifier")

# ── Data source ──────────────────────────────────────────────────────────────
user_stats_source = FileSource(
    path="data/user_stats.parquet",          # can also be BigQuery, Redshift
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# ── Feature view ─────────────────────────────────────────────────────────────
user_stats_fv = FeatureView(
    name="user_stats",
    entities=[user],
    ttl=timedelta(days=30),
    schema=[
        Field(name="purchase_count_7d", dtype=Int64),
        Field(name="avg_order_value_30d", dtype=Float32),
        Field(name="days_since_last_login", dtype=Int64),
        Field(name="preferred_category", dtype=String),
    ],
    online=True,
    source=user_stats_source,
)

# ── Feature service (named set for a model) ──────────────────────────────────
recommendation_svc = FeatureService(
    name="recommendation_v1",
    features=[
        user_stats_fv[["purchase_count_7d", "avg_order_value_30d"]],
    ],
)

from datetime import datetime
from feast import FeatureStore

store = FeatureStore(repo_path="feature_repo/")

# Materialise to online store (incremental, safe to run repeatedly)
store.materialize_incremental(end_date=datetime.utcnow())

# ── Online serving (< 10 ms with Redis) ──────────────────────────────────────
online_features = store.get_online_features(
    features=[
        "user_stats:purchase_count_7d",
        "user_stats:avg_order_value_30d",
        "user_stats:days_since_last_login",
    ],
    entity_rows=[
        {"user_id": "user_42"},
        {"user_id": "user_99"},
    ],
).to_dict()

# ── Point-in-time correct historical features for training ───────────────────
entity_df = pd.DataFrame(
    {
        "user_id": ["user_42", "user_99", "user_07"],
        "event_timestamp": pd.to_datetime(
            ["2024-01-15", "2024-01-16", "2024-01-14"]
        ),
        "label": [1, 0, 1],
    }
)

training_df = store.get_historical_features(
    entity_df=entity_df,
    features=[
        "user_stats:purchase_count_7d",
        "user_stats:avg_order_value_30d",
        "user_stats:days_since_last_login",
    ],
).to_df()