Case 1 — Build a Reproducible Training Pipeline with MLflow

Enterprise Scenario: A subscription SaaS company's churn model currently exists as an unversioned Jupyter notebook. Last week, an engineer retrained the model using an untracked local dataset, causing churn predictions for 12,000 customer accounts to shift unexpectedly. Management could not reproduce the weights, verify the training data version, or identify which hyperparameters were used. You have been tasked with refactoring this notebook into a deterministic, production-grade CLI pipeline (train.py) that logs full lineage to MLflow and registers the model for staging.

Technical Requirements & Deliverables

Deliverable 1: Modularity & CLI Ergonomics (train.py)

• Convert the model workflow into a clean CLI script using argparse.
• Supported parameters: --data-path, --n-estimators, --max-depth, --random-state, --test-size.
• Ensure deterministic execution: pin all random seeds (random_state) across NumPy, Scikit-Learn, and Python built-ins.

Deliverable 2: Experiment Tracking & Schema Signatures

• Wrap the execution in an mlflow.start_run() context within the "customer-churn-production" experiment.
• Log all CLI arguments using mlflow.log_params().
• Log evaluation metrics: ROC-AUC, PR-AUC, F1-Score, and Brier Score using mlflow.log_metric().
• Enforce strict input/output contract enforcement via mlflow.models.infer_signature(X_val, predictions).
• Log model evaluation artifacts: ROC curve plot and confusion matrix figure.

Deliverable 3: Model Registry & Metadata Lineage

• Register the trained model in the MLflow Model Registry under the name churn_predictor.
• Attach required governance tags: git_commit (captured via git rev-parse), data_hash (MD5/SHA256 of training data), author, and framework.
• Programmatically transition the registered model to the Staging stage.

Deliverable 4: Pinned Environment Spec

• Provide a strict requirements.txt with locked dependencies (scikit-learn==1.4.2, mlflow==2.11.3, pandas==2.2.1).

Starter Code Scaffold

import argparse
import hashlib
import json
import subprocess
import mlflow
import mlflow.sklearn
from mlflow.models import infer_signature
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score, average_precision_score, brier_score_loss, f1_score
from sklearn.model_selection import train_test_split

def get_git_commit() -> str:
    try:
        return subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
    except Exception:
        return "local-dev"

def compute_data_hash(filepath: str) -> str:
    hasher = hashlib.md5()
    with open(filepath, "rb") as f:
        hasher.update(f.read())
    return hasher.hexdigest()[:10]

def run_pipeline(args):
    mlflow.set_experiment("customer-churn-production")
    
    with mlflow.start_run() as run:
        # 1. Load Data
        df = pd.read_parquet(args.data_path)
        X = df.drop(columns=["churn", "customer_id"])
        y = df["churn"]
        
        X_train, X_val, y_train, y_val = train_test_split(
            X, y, test_size=args.test_size, random_state=args.random_state, stratify=y
        )
        
        # 2. Train Model
        clf = RandomForestClassifier(
            n_estimators=args.n_estimators,
            max_depth=args.max_depth,
            random_state=args.random_state
        )
        clf.fit(X_train, y_train)
        
        # 3. Log Parameters & Metrics
        # TODO: Log params, evaluate on X_val, log metrics
        
        # 4. Infer Signature & Register
        # TODO: infer_signature, log_model with registered_model_name='churn_predictor'
        pass

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Deterministic MLflow Churn Pipeline")
    parser.add_argument("--data-path", type=str, required=True)
    parser.add_argument("--n-estimators", type=int, default=100)
    parser.add_argument("--max-depth", type=int, default=6)
    parser.add_argument("--random-state", type=int, default=42)
    parser.add_argument("--test-size", type=float, default=0.2)
    run_pipeline(parser.parse_args())

100-Point Graded Rubric

• Criteria 1: Pipeline Structure & CLI Modularity (25 pts)
  • 25 pts: Clean, non-interactive CLI script with defensive argument validation and default fallbacks.
  • 15 pts: Script runs with arguments but lacks parameter validation or error handling.
  • 0 pts: Script fails to execute or requires manual code modifications.
• Criteria 2: Experiment Tracking & Metadata Lineage (25 pts)
  • 25 pts: Comprehensive parameter logging, multiple classification metrics (ROC-AUC, PR-AUC, Brier score), and diagnostic artifact plots.
  • 15 pts: Logs basic accuracy and parameters but omits PR-AUC/Brier score and diagnostic artifacts.
  • 0 pts: Fails to log parameters or metrics within an atomic MLflow run context.
• Criteria 3: Schema Signatures & Model Registry Governance (25 pts)
  • 25 pts: Infers strict input/output signatures; registers model in MLflow Registry with git_commit and data_hash tags.
  • 15 pts: Registers model in registry but omits schema signature or governance tags.
  • 0 pts: Fails to register model in registry.
• Criteria 4: Determinism & Auditability Verification (25 pts)
  • 25 pts: Sequential runs with identical seeds produce bit-for-bit identical metrics; includes pinned requirements.txt.
  • 15 pts: Minor variance in runs due to unpinned component seeds; incomplete lockfile.
  • 0 pts: Stochastic variation across runs; no dependency pinning.