MLFlowForge

End-to-end MLOps pipeline for credit card fraud detection.
XGBoost + PyTorch · MLflow experiment tracking + model registry · Evidently drift monitoring · FastAPI serving · Airflow orchestration · SHAP explainability · Docker Compose · GitHub Actions CI/CD

---

Architecture

┌─────────────────────────────────────────────────────────────────────────────┐
│                        MLFlowForge Pipeline                                 │
│                                                                             │
│  ┌──────────┐   ┌──────────┐   ┌───────────────┐   ┌────────────────────┐  │
│  │  Ingest  │──▶│ Validate │──▶│   Feature Eng  │──▶│ XGBoost Training   │  │
│  │ (Kaggle) │   │   (GE)   │   │ SMOTE · Scaler │   │ MLflow · SHAP      │  │
│  └──────────┘   └──────────┘   └───────────────┘   └────────────────────┘  │
│                                                              │              │
│                                                     ┌────────────────────┐  │
│                                                     │  PyTorch MLP       │  │
│                                                     │  MLflow · SHAP     │  │
│                                                     └────────────────────┘  │
│                                                              │              │
│  ┌─────────────────────┐   ┌────────────────┐   ┌───────────────────────┐  │
│  │   Evidently Drift   │◀──│    Registry    │◀──│  Evaluate & Register  │  │
│  │   Wasserstein · PSI │   │  Staging→Prod  │   │  Quality Gate 0.70 AP │  │
│  └─────────────────────┘   └────────────────┘   └───────────────────────┘  │
│              │                                                              │
│  ┌───────────────────────┐   ┌─────────────────────────────────────────┐   │
│  │  Retrain Trigger      │   │  FastAPI Serving                        │   │
│  │  (event-driven)       │   │  /predict · /predict/batch · /reload    │   │
│  └───────────────────────┘   └─────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────────────────┘
           Orchestrated by Airflow DAG (fraud_detection_pipeline)

---

Tech Stack

Component	Technology
Orchestration	Apache Airflow 2.8
Experiment tracking	MLflow 2.10
Model registry	MLflow Model Registry (Staging → Production → Archived)
ML models	XGBoost 2.x + PyTorch 2.x tabular MLP
Feature engineering	scikit-learn + imbalanced-learn (SMOTE)
Serving	FastAPI + Uvicorn (loaded from MLflow registry)
Drift monitoring	Evidently AI (DataDriftPreset) + Wasserstein + PSI
Explainability	SHAP (TreeExplainer for XGBoost, GradientExplainer for PyTorch)
Dataset	Kaggle Credit Card Fraud Detection (284k transactions)
CI/CD	GitHub Actions
Infrastructure	Docker Compose (PostgreSQL · MLflow · FastAPI · Airflow)

---

Dataset

Kaggle Credit Card Fraud Detection — 284,807 transactions with 492 frauds (~0.172%).

Features: V1–V28 (PCA-transformed), Amount, Time, Class (0=legit, 1=fraud).

The pipeline automatically falls back to a synthetic dataset if Kaggle credentials are unavailable (e.g. in CI).

---

Quick Start

1. Prerequisites

• Docker + Docker Compose
• Python 3.11+
• (Optional) Kaggle API credentials

2. Clone and configure

git clone https://github.com/rohanmukka/MLFlowForge.git
cd MLFlowForge
cp .env.example .env
# Fill in KAGGLE_USERNAME, KAGGLE_KEY, and any other values

3. Start infrastructure

docker compose up -d postgres mlflow
# Wait for MLflow to be healthy (~20s)
docker compose ps

4. Install Python dependencies

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

5. Run the full pipeline manually

# Step 1 — Ingest data
python -m pipeline.ingest

# Step 2 — Feature engineering (SMOTE, scaling, splits)
python -m pipeline.features

# Step 3 — Train XGBoost
python -m pipeline.train_xgboost

# Step 3b — Train PyTorch MLP (optional, for comparison)
python -m pipeline.train_pytorch

# Step 4 — SHAP explanations (pass run_id from step 3)
python -m explainability.shap_analysis xgboost <run_id>

# Step 5 — Promote to Production (pass version number)
python -m pipeline.evaluate promote --version 1 --stage Production

# Step 6 — Drift detection
python -m monitoring.drift_detector

# Step 7 — Start the inference API
uvicorn serving.api:app --host 0.0.0.0 --port 8000

6. Run via Airflow

docker compose up -d
# Access Airflow UI at http://localhost:8080 (admin/admin)
# Trigger the fraud_detection_pipeline DAG

---

API Reference

POST /predict

Score a single transaction.

{
  "transaction": {
    "V1": -1.36, "V2": -0.07, ..., "V28": -0.02,
    "hour_sin": 0.707, "hour_cos": 0.707,
    "log_amount": 4.49
  },
  "threshold": 0.5
}

Response:

{
  "fraud_probability": 0.042631,
  "is_fraud": false,
  "model_version": "3",
  "latency_ms": 1.83
}

POST /predict/batch

Score up to 1,000 transactions in a single call.

GET /health

Returns model load status and rolling p99 latency.

POST /reload

Hot-swap the Production model from the registry (requires X-Admin-Key header).

---

Pipeline Details

Feature Engineering (pipeline/features.py)

Transformation	Input	Output
Cyclic time encoding	Time (seconds)	hour_sin, hour_cos
Log transform	Amount	log_amount
Drop raw columns	Time, Amount	—
SMOTE	Training split	Resampled training split
StandardScaler	All features	Scaled features

SMOTE is applied only to the training split to prevent leakage.

Model Training

XGBoost (pipeline/train_xgboost.py)

• hist tree method for speed on large tabular data
• scale_pos_weight handles residual imbalance post-SMOTE
• Early stopping on validation PR-AUC (aucpr)
• All hyperparameters logged to MLflow

PyTorch MLP (pipeline/train_pytorch.py)

• Architecture: [Linear → BatchNorm → ReLU → Dropout] × N_layers → Sigmoid
• BCEWithLogitsLoss with pos_weight for class imbalance
• Cosine annealing LR schedule + early stopping on validation Average Precision

MLflow Model Registry

None → Staging → Production → Archived

Quality gate before Production promotion: average_precision >= 0.70 on test set. The previous Production version is automatically archived on promotion.

Drift Detection (monitoring/drift_detector.py)

Three-signal drift detection:

Signal	Method	Default threshold
Per-feature distribution shift	Wasserstein-1 distance	0.1
Population Stability Index	PSI	0.2
Column-level drift share	Evidently DataDriftPreset	30% of columns

Drift detected if any signal triggers. An HTML report and JSON summary are written to reports/.

Automated Retraining (monitoring/retrain_trigger.py)

Event-driven (not scheduled). Triggered by drift_detected=True from the Airflow DAG's drift task. Runs:

1. Feature engineering on new/existing data
2. XGBoost retraining
3. (Optional) PyTorch retraining
4. Quality-gated Production promotion

All retraining runs produce a structured JSON event log in reports/retrain_events/.

SHAP Explainability (explainability/shap_analysis.py)

Per-run artifacts logged to MLflow under shap/:

File	Description
summary_bar_<model>.png	Global mean |SHAP| feature importance
summary_beeswarm_<model>.png	Direction + magnitude per sample
waterfall_top_fraud_<model>.png	Local explanation for highest-confidence fraud
dependence_<feature>_<model>.png	Top-3 feature dependence plots
shap_values_<model>.npy	Raw SHAP values array

---

CI/CD (.github/workflows/ci.yml)

Job	Trigger	Steps
unit-tests	Every PR	pytest (ingest, features, API) + coverage report
smoke-test	Every PR (after unit-tests)	Full pipeline on 1k synthetic rows
build-docker	Merge to main	Build + push API + MLflow images to GHCR
lint	Every PR	ruff lint check (informational)

---

Design Decisions

Why XGBoost + PyTorch?

XGBoost provides strong out-of-the-box performance and fast inference for tabular fraud data. The PyTorch MLP is included for comparison and to demonstrate framework-agnostic MLflow logging. In production, XGBoost typically wins on tabular imbalanced data.

Why event-driven retraining?

Scheduled retraining wastes compute when the distribution is stable. The Evidently + Wasserstein + PSI combination provides multi-signal confidence before triggering expensive retraining.

Why SMOTE only on training data?

Applying SMOTE to val/test splits would inflate minority metrics by introducing synthetic fraud samples that the model could trivially identify as similar to training examples — creating an optimistic but unrealistic benchmark.

Why three drift signals?

Each signal has different sensitivity:

• Wasserstein: captures distributional shift in the full range
• PSI: industry-standard banking metric, intuitive thresholds
• Evidently: additional statistical tests (KS, Z-score) with visual reports

Using all three reduces false positives while increasing recall for real drift.

---

Metrics

Expected performance on the Kaggle credit card fraud dataset (test set):

Model	ROC-AUC	Average Precision	F1 (threshold=0.5)
XGBoost	~0.98	~0.85	~0.83
PyTorch MLP	~0.97	~0.79	~0.76

Exact values vary with SMOTE random state and hyperparameters.

---

Project Structure

MLFlowForge/
├── pipeline/
│   ├── ingest.py           # Data download + Great Expectations validation
│   ├── features.py         # Feature engineering + SMOTE + scaling
│   ├── train_xgboost.py    # XGBoost training + MLflow logging
│   ├── train_pytorch.py    # PyTorch MLP training + MLflow logging
│   └── evaluate.py         # Metrics + model registry lifecycle
├── monitoring/
│   ├── drift_detector.py   # Evidently + Wasserstein + PSI drift detection
│   └── retrain_trigger.py  # Event-driven retraining orchestration
├── serving/
│   ├── api.py              # FastAPI inference endpoints
│   └── model_loader.py     # Thread-safe MLflow registry model loader
├── explainability/
│   └── shap_analysis.py    # SHAP summary, beeswarm, waterfall, dependence plots
├── dags/
│   └── fraud_pipeline.py   # Airflow DAG (8 tasks)
├── tests/
│   ├── test_ingest.py
│   ├── test_features.py
│   ├── test_api.py
│   └── test_smoke_train.py
├── config/
│   ├── hyperparams.yaml    # All model hyperparameters
│   └── postgres-init.sql
├── .github/workflows/
│   └── ci.yml              # GitHub Actions CI/CD
├── docker-compose.yml
├── Dockerfile.mlflow
├── Dockerfile.api
├── requirements.txt
└── .env.example

---

License

MIT