Flit ML: Complete Project Implementation Plan

[whitehackr]

Flit ML Project: Complete Phase Implementation Plan

Overview

Build ML components for the flit ecosystem, focusing on predictive models for financial risk assessment for a BNPL product. This is the first ML project at Flit, requiring research-first approach followed by production infrastructure.

Phase -1: Data Infrastructure (whitehackr/flit-data-platform#9)

Objectives

• Generate and store 3 months of synthetic BNPL data
• Set up data warehouse for ML research
• Build data pipeline for ongoing data collection

Deliverables

• Data Generation Pipeline

  • Airflow DAG to collect data from simtom API
  • Generate realistic 3-month historical dataset
  • Data quality validation and monitoring

• Data Warehouse Setup

  • BigQuery database with proper schema
  • Organized tables (transactions, users, risk_events, etc.)
  • Partitioned for fast analytical queries

• Data Pipeline

  • Automated daily data collection from simtom
  • Data cleaning and validation with Great Expectations
  • Backup and versioning strategy

Technical Tasks

• Simtom Enhancement: Create PR to add date range parameters (/stream/bnpl?start_date=2024-06-01&end_date=2024-09-01)
• Database Setup: Design BigQuery schema, create tables
• Airflow Setup: DAG for data collection, scheduling
• Data Generation: Script to generate 3 months of historical data
• Validation: Great Expectations suite for data quality

Tooling Required

• Database: google-cloud-bigquery, pandas-gbq
• Orchestration: apache-airflow
• Validation: great-expectations
• Simtom Enhancement: Date range API feature

---

Phase 0: Research & Discovery

Objectives

• Understand BNPL data patterns and business problem
• Define ML problem clearly (classification vs regression, target variable)
• Establish baseline performance metrics
• Identify best performing model architectures

Deliverables

• Data Understanding Report

  • Data schema documentation
  • Statistical summary of all features
  • Data quality assessment (missing values, outliers)

• Problem Definition Document

  • Clear target variable definition (default probability? risk score?)
  • Success metrics (precision, recall, AUC, business metrics)
  • Model performance thresholds for production

• Model Experimentation Results

  • Baseline model performance (simple logistic regression)
  • Comparison of 5-7 different algorithms
  • Feature importance analysis
  • Recommended models for production

Technical Tasks (Some of these notebooks could be combined depending on logical workflow, and io overhead)

• 01_data_exploration.ipynb: Connect to BigQuery, basic data inspection
• 02_eda_analysis.ipynb: Deep statistical analysis, visualizations
• 03_feature_engineering.ipynb: Create derived features, handle categorical data
• 04_baseline_models.ipynb: Simple models (logistic regression, decision tree)
• 05_advanced_models.ipynb: XGBoost, Random Forest, Neural Networks
• 06_model_comparison.ipynb: Cross-validation, performance comparison
• 07_final_recommendations.ipynb: Model selection with business justification

Tooling Required

• Jupyter Notebook: Interactive development
• Additional packages: matplotlib, seaborn, plotly, mlflow
• Model libraries: lightgbm, catboost
• Data processing: pandas, numpy (already included)

---

Phase 1: Production Infrastructure

Objectives

• Build production-ready ML serving infrastructure
• Implement model versioning and deployment pipeline
• Create monitoring and observability

Deliverables

• Model Serving API

  • FastAPI service with prediction endpoints
  • Model loading and caching
  • Input validation and error handling

• Model Registry System

  • Model versioning and storage
  • A/B testing capabilities
  • Model rollback functionality

• Monitoring Dashboard

  • Prediction latency metrics
  • Model performance monitoring
  • Data drift detection

Technical Tasks

• Core Architecture: Base classes, model registry, plugin system
• API Development: FastAPI endpoints, async request handling
• Model Deployment: Model loading, caching, version management
• Data Pipeline: Real-time feature engineering, validation
• Monitoring: Metrics collection, alerting, dashboards
• Testing: Unit tests, integration tests, load tests

Tooling Required

• FastAPI: Already included
• Model Storage: joblib, pickle, or mlflow model registry
• Monitoring: prometheus, grafana or simple logging
• Caching: redis (optional for model caching)
• Container: docker for deployment

---

Phase 2: Real-time Processing

Objectives

• Handle streaming data from simtom
• Implement real-time feature engineering
• Build batch prediction capabilities

Deliverables

• Streaming Data Pipeline

  • Real-time data ingestion from simtom
  • Feature engineering on streaming data
  • Batch processing for historical data

• Real-time Prediction Service

  • Low-latency prediction API (<100ms)
  • Async processing for high throughput
  • Queue management for spike handling

Technical Tasks

• Stream Processing: Async data consumption from simtom API
• Feature Store: Real-time feature computation and storage
• Batch Processing: Historical data processing for model retraining
• Queue Management: Handle prediction request spikes

Tooling Required

• Streaming: asyncio, httpx (already included)
• Message Queue: celery + redis or simple async queues
• Feature Store: redis or in-memory with persistence
• Batch Processing: pandas for data processing

---

Phase 3: Model Operations (MLOps)

Objectives

• Automated model retraining pipeline
• A/B testing framework
• Model performance monitoring

Deliverables

• Automated Training Pipeline

  • Scheduled model retraining
  • Data validation before training
  • Automated model evaluation and deployment

• A/B Testing Framework

  • Traffic splitting between model versions
  • Statistical significance testing
  • Automated winner selection

• Performance Monitoring

  • Model drift detection
  • Performance degradation alerts
  • Business metrics tracking

Technical Tasks

• Training Automation: Scheduled training jobs, data validation
• A/B Testing: Traffic routing, experiment management
• Monitoring: Data drift detection, performance tracking
• Alerting: Automated alerts for model issues

Tooling Required

• Scheduling: celery beat or cron jobs
• Experiment Management: Custom A/B testing or mlflow
• Monitoring: evidently for data drift, custom metrics
• Alerting: slack webhooks or email notifications

---

Phase 4: Advanced Features

Objectives

• Model interpretability and explainability
• Advanced model architectures
• Integration with flit ecosystem

Deliverables

• Model Explainability

  • SHAP values for predictions
  • Feature importance explanations
  • Model decision boundaries

• Advanced Models

  • Ensemble methods
  • Deep learning models (if beneficial)
  • Time-series models for temporal patterns

• Ecosystem Integration

  • Integration with flit-data-platform
  • Connection to production flit services
  • Business metrics dashboard

Technical Tasks

• Explainability: SHAP implementation, visualization
• Advanced Models: Ensemble methods, neural networks
• Integration: APIs for flit ecosystem, data connectors
• Business Metrics: Revenue impact tracking, risk assessment

Tooling Required

• Explainability: shap, lime, eli5
• Deep Learning: pytorch or tensorflow (if needed)
• Visualization: streamlit for dashboards
• Integration: Custom APIs, database connectors

---

Phase 5: Production Deployment (1-2 weeks)

Objectives

• Deploy to production environment
• Load testing and performance optimization
• Documentation and handover

Deliverables

• Production Deployment

  • Railway deployment configuration
  • Environment management
  • Load balancing and scaling

• Documentation

  • API documentation
  • Model documentation
  • Operational runbooks

• Performance Validation

  • Load testing results
  • Performance benchmarks
  • Monitoring setup verification

Technical Tasks

• Deployment: Railway configuration, environment setup
• Testing: Load testing, stress testing
• Documentation: API docs, model cards, operational guides
• Handover: Knowledge transfer, operational procedures

Tooling Required

• Deployment: railway CLI, docker
• Load Testing: locust or wrk
• Documentation: mkdocs or simple markdown
• Monitoring: Production monitoring setup

---

Technology Stack Summary

Core ML Stack

• Python: 3.11+ (already set)
• ML Libraries: scikit-learn, xgboost, pandas, numpy (already included)
• API: FastAPI + uvicorn (already included)
• Validation: Pydantic (already included)

Additional Requirements by Phase

• Phase -1: google-cloud-bigquery, pandas-gbq, apache-airflow, great-expectations
• Phase 0: jupyter, matplotlib, seaborn, plotly, mlflow
• Phase 1: redis (optional), prometheus (optional)
• Phase 2: celery (optional), message queue
• Phase 3: evidently, experiment tracking
• Phase 4: shap, streamlit, pytorch (optional)
• Phase 5: locust, mkdocs

Infrastructure

• Development: Poetry + virtual env (already set)
• Database: BigQuery
• Deployment: Railway
• Storage: GCP Cloud Storage
• Monitoring: Simple logging initially, then proper monitoring

---

Dependencies

1. Simtom API Enhancement: Need to contribute date range feature to simtom project before Phase -1
2. BigQuery Setup: GCP project and BigQuery dataset creation
3. Airflow Environment: Local Airflow setup or cloud-managed Airflow

Success Criteria

• Phase -1: 3 months of quality BNPL data in BigQuery
• Phase 0: Clear model recommendations with >75% baseline accuracy
• Phase 1: Production API with <100ms latency, 99%+ uptime
• Phase 2: Handle 100+ predictions/second
• Phase 3: Automated retraining and A/B testing
• Phase 4: Model explainability and advanced features
• Phase 5: Full production deployment on Railway

---

This issue will be updated as we progress through each phase. Each phase will have its own sub-issues for detailed tracking.

[whitehackr]

Phase 1 Implementation Details - Core Architecture

Key Components to Implement

1. Base Model Interface (flit_ml/core/base.py)

class BaseModel(ABC):
    @abstractmethod
    async def predict(self, request: PredictionRequest) -> PredictionResponse
    
    @abstractmethod  
    async def train(self, training_data: List[Dict[str, Any]]) -> None

2. Model Registry (flit_ml/core/registry.py)

• Plugin registration: @register_model("bnpl_risk")
• Model loading and caching
• Version management
• A/B testing routing

3. Pydantic Models (flit_ml/core/schemas.py)

• ModelConfig: Type-safe model configuration
• PredictionRequest: API request format
• PredictionResponse: API response format
• TrainingRequest: Training job configuration

4. FastAPI Service (flit_ml/api/main.py)

• /predict endpoint with async handling
• /models endpoint for model management
• /health endpoint for monitoring
• Automatic request validation

5. Example Implementation Pattern

@register_model("bnpl_risk")
class BNPLRiskModel(BaseModel):
    async def predict(self, request: PredictionRequest) -> PredictionResponse:
        # Feature validation
        # Model inference
        # Response formatting

This architecture enables:

• Easy model addition without changing core API
• Async prediction handling for high throughput
• Type safety throughout the pipeline
• Built-in monitoring and health checks
• A/B testing capabilities

Note: These details captured from initial architecture planning session before pivoting to research-first approach.

[whitehackr]

Architectural Decision: ML Service Design Patterns

Context

Flit will have multiple diverse ML models over time:

• BNPL Risk Assessment (binary classification)
• Push Personalization (recommendation system)
• Shipping/Order Routing (optimization)
• Fraud Detection (anomaly detection)
• Credit Scoring (regression/classification)

Architecture Options Evaluated

1. Direct Model Classes (Simple)

class BNPLRiskModel:
    def predict(self, data): return risk_score

class PushPersonalizationModel:  
    def recommend(self, user_id): return offers

# Usage: Different interfaces per model
bnpl_score = bnpl_model.predict(user_data)
offers = push_model.recommend(user_id)

Pros:

• Simple and direct
• Easy to understand and debug
• No abstraction overhead
• Fast initial development

Cons:

• Inconsistent interfaces across models
• Difficult to add cross-cutting concerns (monitoring, logging, A/B testing)
• Hard to standardize deployment and operations
• Each model needs custom API endpoints
• Monitoring and alerting logic duplicated

---

2. Service-Oriented Architecture (SOA)

# Separate microservice for each model type
BNPL_SERVICE_URL = "https://bnpl-risk.flit.com/predict"
PUSH_SERVICE_URL = "https://personalization.flit.com/recommend"

# Usage: HTTP calls between services
bnpl_result = requests.post(BNPL_SERVICE_URL, json=data)
push_result = requests.post(PUSH_SERVICE_URL, json=user_context)

Pros:

• Clear separation of concerns
• Independent deployment and scaling
• Technology diversity (different languages/frameworks per service)
• Team autonomy

Cons:

• Network overhead and latency
• Complex service discovery and communication
• Distributed system complexity (failures, timeouts, circuit breakers)
• Operational overhead (monitoring multiple services)
• Data consistency challenges

---

3. Plugin Architecture (Recommended)

@register_model("bnpl_risk")
class BNPLRiskModel(BaseModel):
    async def predict(self, request: PredictionRequest) -> PredictionResponse:
        return {"risk_score": 0.85, "confidence": 0.92}

@register_model("push_personalization")  
class PushModel(BaseModel):
    async def predict(self, request: PredictionRequest) -> PredictionResponse:
        return {"recommended_offers": ["offer_1", "offer_2"]}

# Usage: Consistent interface
result = await model_registry.predict("bnpl_risk", features)
result = await model_registry.predict("push_personalization", user_context)

Pros:

• Consistent Interface: Same API contract across all models
• Cross-cutting Concerns: Monitoring, logging, caching, A/B testing implemented once
• Operational Simplicity: Single deployment, monitoring, and alerting system
• Team Scaling: New ML engineers follow established patterns
• Flexible Deployment: Can run multiple models in one service or separate them later
• Easy A/B Testing: Built-in support for model versioning and traffic splitting
• Type Safety: Pydantic models ensure consistent request/response formats

Cons:

• Higher initial complexity
• Abstraction layer adds some overhead
• All models must conform to same interface
• Shared deployment means one model can affect others

---

4. Model-as-a-Function (Serverless)

# Each model as separate cloud function
def bnpl_risk_function(request):
    model = load_model()
    return model.predict(request.data)

def push_personalization_function(request):
    model = load_model()  
    return model.recommend(request.user_id)

Pros:

• Zero-cost when not in use
• Automatic scaling
• No infrastructure management
• Independent deployment

Cons:

• Cold start latency
• Limited execution time
• Vendor lock-in
• Difficult local development
• Model loading overhead on each invocation

---

Decision: Plugin Architecture

Chosen for flit-ml based on:

1. Diverse Model Types: BNPL, recommendations, routing require different logic but benefit from consistent operations
2. Team Scaling: As team grows, consistent patterns reduce onboarding time
3. Operational Efficiency: Single service to deploy, monitor, and maintain
4. A/B Testing: Built-in support for testing different model versions
5. Future Flexibility: Can split into microservices later if needed

Key Design Principles:

• Interface Segregation: Models only implement what they need
• Dependency Injection: Easy to test and mock
• Configuration-Driven: Models configured via Pydantic schemas
• Async-First: Handle high throughput prediction requests

Implementation Strategy:

• Start with plugin architecture in Phase 1
• Build base abstractions that work for all model types
• Add cross-cutting concerns (monitoring, A/B testing) in Phase 3
• Migrate to microservices only if scaling requires it

This architecture balances simplicity with extensibility, enabling rapid ML model development while maintaining operational sanity.

[whitehackr]

Plugin Architecture → Microservices Evolution Path

Phase 1-3: Monolithic Plugin Service

# Single flit-ml service with all models
@register_model("bnpl_risk")
class BNPLRiskModel(BaseModel): ...

@register_model("push_personalization")
class PushModel(BaseModel): ...

@register_model("shipping_routing")
class RoutingModel(BaseModel): ...

# All models served from single FastAPI service
app.post("/predict/{model_name}") 

Future: Selective Microservices Extraction

When to Split:

• Model has different scaling requirements (e.g., BNPL gets 1000x more traffic)
• Model needs specialized infrastructure (e.g., GPU for deep learning)
• Team wants independent deployment cycles
• Model has different SLA requirements

How to Split:

1. Extract Model: Move model class to separate service

 # New bnpl-risk-service
@register_model("bnpl_risk")
class BNPLRiskModel(BaseModel): ...

# Main flit-ml service (other models remain)
@register_model("push_personalization")
@register_model("shipping_routing")

2. Update Registry: Route requests to appropriate service

class DistributedModelRegistry:
    def __init__(self):
        self.local_models = {"push_personalization", "shipping_routing"}
        self.remote_models = {
            "bnpl_risk": "https://bnpl-risk-service.flit.com"
        }

    async def predict(self, model_name: str, request: PredictionRequest):
        if model_name in self.local_models:
            return await self.local_predict(model_name, request)
        else:
            return await self.remote_predict(model_name, request)

3. Maintain Interface: Client code unchanged

# Same API call works for both local and remote models
result = await model_registry.predict("bnpl_risk", features)  # → Remote service
result = await model_registry.predict("push_personalization", features)  # → Local model

Benefits of This Evolution Path

Start Simple:

• Single deployment and monitoring
• Faster development and testing
• Lower operational complexity

Scale Selectively:

• Extract only models that need it
• Maintain shared infrastructure for common models
• Gradual learning curve for distributed systems

Preserve Consistency:

• Same BaseModel interface across all services
• Shared monitoring, logging, and A/B testing patterns
• Type-safe communication via Pydantic models

Example Migration Timeline (Remember we're setting up our ML function from scratch):

• Month 1-6: All models in single service
• Month 7: Extract BNPL model due to high traffic
• Month 12: Extract recommendations due to GPU requirements
• Month 18: Keep routing model local (low traffic, simple needs)

This approach gives us microservices benefits without microservices complexity until we actually need them.

[whitehackr]

Phase -1 Status: Accelerated Completion + Architecture Discovery

Completed Beyond Requirements

• ✅ 1.9M transactions (vs 3 months requested) - Sept 2024 to Aug 2025 coverage
• ✅ 42 validated fields with comprehensive dbt testing framework
• ✅ Programmatic JSON extraction handling API schema evolution automatically
• ✅ Data quality framework established (20 tests passing, unique transaction ID resolution)
• ✅ BigQuery staging layer production-ready

Critical Discovery: Real-time Architecture Required

During data platform development, discovered that BNPL approval decisions require sub-second response times. This fundamentally changes our ML approach:

Production Reality:
Transaction Event → Streaming → Python ML Service → Approval/Denial

Impact on Project Phases:

• Phase -1: Complete (exceeded objectives)
• Phase 0: Modify for real-time research approach using staging data
• Phases 1-2: Implement streaming ML architecture, not batch processing
• Batch components: Unnecessary for primary BNPL use case (Airflow, scheduled generation)

Data Foundation Ready

ML research can begin immediately with:

# Access via BigQuery
df = client.query("""
    SELECT * FROM `flit-data-platform.flit_staging.stg_bnpl_raw_transactions`
    WHERE transaction_timestamp >= '2024-09-01'
""").to_dataframe()

# 42 fields including target variable 'will_default'
# All ML features ready for immediate experimentation  

Recommended Path Forward

1. Mark Phase -1 complete - data foundation exceeds requirements
2. Modify Phase 0 - focus on real-time ML research and model development
3. Plan streaming deployment - Python services for sub-second decisions
4. Defer batch analytics - return to data-platform repo post-ML deployment for downstream work

Next Step: Proceed to ML research phase using established data foundation.