Case 1 — Build a Reproducible Training Pipeline with MLflow

# Case 1 — Build a Reproducible Training Pipeline with MLflow > *"Your churn model lives in a notebook. Someone retrained it last week; nobody knows the data version. Convert this to a reproducible MLflow pipeline that anyone can rerun and produce the same result."* --- ## Step 1 — Convert the notebook to a script `train.py`: ```python import argparse import json import subprocess from datetime import datetime import joblib import mlflow import mlflow.sklearn import pandas as pd from sklearn.metrics import roc_auc_score, average_precision_score from sklearn.model_selection import train_test_split from xgboost import XGBClassifier from features.feature_engineering import build_features def get_git_commit(): return subprocess.check_output(['git', 'rev-parse', 'HEAD']).strip().decode() def main(args): mlflow.set_experiment(args.experiment_name) with mlflow.start_run(run_name=args.run_name): # Log inputs mlflow.log_params(vars(args)) mlflow.set_tag('git_commit', get_git_commit()) mlflow.set_tag('data_version', args.data_version) mlflow.set_tag('trained_at', datetime.utcnow().isoformat()) # Load data data = pd.read_parquet(args.data_path) X = data.drop(columns=['target']) y = data['target'] # Feature engineering (same function as serving) X_features = X.apply(build_features, axis=1) # Time-based split X_train, X_test, y_train, y_test = train_test_split( X_features, y, test_size=0.2, random_state=42, stratify=y ) # Train model = XGBClassifier( n_estimators=args.n_estimators, learning_rate=args.learning_rate, max_depth=args.max_depth, random_state=42, scale_pos_weight=(y_train == 0).sum() / (y_train == 1).sum(), ) model.fit(X_train, y_train, eval_set=[(X_test, y_test)]) # Metrics y_pred = model.predict_proba(X_test)[:, 1] auc = roc_auc_score(y_test, y_pred) pr_auc = average_precision_score(y_test, y_pred) mlflow.log_metric('test_auc', auc) mlflow.log_metric('test_pr_auc', pr_auc) # Log model mlflow.sklearn.log_model(model, 'model') # Log feature importance importances = dict(zip(X_features.columns, model.feature_importances_)) with open('/tmp/feature_importances.json', 'w') as f: json.dump(importances, f, indent=2) mlflow.log_artifact('/tmp/feature_importances.json') print(f'Run complete. AUC={auc:.4f}, PR-AUC={pr_auc:.4f}') if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--data-path', required=True) parser.add_argument('--data-version', required=True) parser.add_argument('--experiment-name', default='churn-detection') parser.add_argument('--run-name', default='default') parser.add_argument('--n-estimators', type=int, default=1000) parser.add_argument('--learning-rate', type=float, default=0.05) parser.add_argument('--max-depth', type=int, default=6) args = parser.parse_args() main(args) ``` Run via: ```bash python train.py --data-path data/v3.parquet --data-version v3 --run-name xgb_v2 ``` Every run logs everything to MLflow. ## Step 2 — Lock the environment `requirements.txt` (pinned): ``` xgboost==2.0.3 scikit-learn==1.4.0 pandas==2.1.4 numpy==1.26.0 mlflow==2.10.0 pyarrow==15.0.0 ``` Dockerfile: ```dockerfile FROM python:3.11.5-slim WORKDIR /app COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt COPY . . CMD ["python", "train.py"] ``` Build: ```bash docker build -t churn-trainer:v1 . docker run -v $(pwd)/data:/app/data churn-trainer:v1 \ python train.py --data-path data/v3.parquet --data-version v3 ``` Anyone can now build the image and run with the same dependencies. ## Step 3 — Add data versioning with DVC ```bash dvc init dvc add data/v3.parquet git add data/v3.parquet.dvc .dvcignore git commit -m "Add training data v3" # Configure remote dvc remote add -d storage s3://my-mlops-data dvc push # upload data file to S3 ``` Now anyone can: ```bash git checkout dvc pull # fetch the right data version from S3 docker build && docker run ``` And produce the same model. ## Step 4 — Track results in MLflow UI Start the MLflow UI: ```bash mlflow ui ``` Browser opens to localhost:5000. You see: - All runs across all experiments. - Sortable by metrics. - Side-by-side comparison. - Model artifacts downloadable. For team use, host MLflow on a server with a persistent backend (Postgres + S3 for artifacts): ```bash mlflow server \ --backend-store-uri postgresql://user:pass@db:5432/mlflow \ --default-artifact-root s3://my-mlflow-artifacts \ --host 0.0.0.0 ``` ## Step 5 — Promote the best run to the Model Registry ```python import mlflow # Find the best run runs = mlflow.search_runs(experiment_names=['churn-detection'], o