KKBox Churn Prediction Pipeline

This repository contains an end-to-end Machine Learning Operations (MLOps) pipeline designed to predict customer churn for KKBox, a prominent subscription-based music streaming service. The system is architected using ZenML for orchestration, MLflow for experiment tracking and model registry, and FastAPI for real-time model serving.

The primary objective is to identify users who exhibit a high probability of canceling their subscription, enabling proactive retention strategies.

• Business Context
• System Architecture
• Technology Stack
• Repository Structure
• Installation and Setup
• Usage Guide
• Cloud Deployment Strategy
• Limitations and Future Work

Customer churn is a critical Key Performance Indicator (KPI) for subscription-driven business models. Identifying "at-risk" users allows the business to optimize marketing spend through targeted interventions.

• Behavioral Attrition: Churning users typically exhibit a significant decline in service usage (music listening) 14 to 30 days prior to subscription expiration.
• Payment Method Friction: Users utilizing manual payment methods experience substantially higher churn rates compared to those on auto-renewal plans.
• Data Integrity: Approximately 60% of demographic age data is missing or invalid, necessitating robust imputation strategies within the preprocessing pipeline.

The project implements a complete MLOps lifecycle, structured into discrete, scalable pipeline steps orchestrated by ZenML.

1. Data Ingestion: Utilizes Polars for highly efficient, memory-optimized data loading.
2. Data Validation: Employs Pandera for schema enforcement, ensuring data integrity and flagging anomalies before processing.
3. Feature Engineering: Constructs Recency, Frequency, Monetary (RFM), and behavioral features. Implements memory downcasting to manage large datasets.
4. Model Training: Utilizes XGBoost with strict temporal constraints to prevent data leakage (e.g., removing future-dated transactions).
5. Evaluation: Calculates standard classification metrics (AUC, Precision, Recall) and logs feature importance alongside evaluation results to MLflow.
6. Model Registry: Automatically promotes models exceeding the minimum AUC threshold (0.75) to the MLflow Model Registry for production serving.

kkbox-churn-prediction/
├── data/                    # Raw and processed datasets (Git ignored)
├── mlruns/                  # MLflow local tracking directory
├── notebooks/               # Exploratory Data Analysis (EDA)
├── serving/
│   ├── main.py              # FastAPI application entrypoint
│   ├── request_testing.py   # Client script for API testing
│   └── schemas.py           # Pydantic data validation schemas
├── src/
│   ├── pipelines/           # ZenML pipeline definitions
│   ├── steps/               # Individual pipeline components
│   └── config.py            # Centralized system configuration
├── config.py                # Global data paths and thresholds
├── run_pipeline.py          # Primary execution script for training
├── requirements.txt         # Project dependencies
└── setup_venv.sh            # Environment initialization script

• Python 3.8 or higher
• Git
• Virtual environment tool (venv, conda)

1. Clone the repository:

   git clone <repository_url>
   cd churn-prediction-pipeline

2. Initialize the virtual environment: You can use the provided bash script or set it up manually:

   bash setup_venv.sh
   # OR manually
   python -m venv venv
   source venv/bin/activate  # On Windows: venv\Scripts\activate
   pip install -r requirements.txt

3. Data Preparation: Ensure the required KKBox datasets (members_v3.csv, transactions_v2.csv, user_logs_v2.csv, train_v2.csv) are placed within the data/ directory as configured in config.py.

The system requires multiple services running concurrently for a complete local execution. It is recommended to use separate terminal sessions.

Start the MLflow server to track experiments and host the model registry.

mlflow server --backend-store-uri mlruns --default-artifact-root mlruns --host 0.0.0.0 --port 5000

Trigger the ZenML pipeline to ingest data, train the model, and register it if performance criteria are met.

python run_pipeline.py

Note: Verify the model registration via the MLflow UI at http://localhost:5000.

Start the FastAPI server to expose the registered model for predictions.

python -m serving.main

Simulate client requests to verify the serving infrastructure.

python -m serving.request_testing

ZenML abstracts infrastructure, allowing seamless transitions from local development to cloud environments via stack configurations.

1. Register Infrastructure Components:

   zenml artifact-store register s3_store --flavor=s3 --path=s3://your-bucket-name
   zenml container-registry register ecr_registry --flavor=aws --uri=<account>.dkr.ecr.<region>.amazonaws.com
   zenml orchestrator register sagemaker_orch --flavor=sagemaker

2. Configure and Activate Stack:

   zenml stack register aws_production_stack -a s3_store -c ecr_registry -o sagemaker_orch
   zenml stack set aws_production_stack

3. Execute:

   python run_pipeline.py

• Hardware Constraints: Due to local memory limitations, the default pipeline configuration processes only Version 2 data (March 2017). This truncates historical behavioral features.
• Scalability: The Polars ingestion layer is designed for scalability. Processing the full 30GB dataset requires migrating execution to high-memory cloud instances (e.g., AWS EC2 or SageMaker) and updating data paths in config.py.