Case 1 — Build a Reproducible Training Pipeline with MLflow

# 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 ```python 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 metri