Skip to content

VisheshKamble/realtime-personalization-engine

Repository files navigation

Typing SVG

CI Python PyTorch FastAPI Kafka React License: MIT


A production-grade ML recommendation system that delivers personalized results in under 100ms.
Built with the same architecture used at YouTube, Pinterest, and Spotify.

Features · Architecture · ML Stack · Quickstart · API Docs · Benchmarks


✨ Features

Feature Details
Real-time event streaming Kafka topics partitioned by user_id — all user events hit the same partition
Streaming feature computation PyFlink micro-batch computes rolling click rate, dwell time, category affinity
Dual feature store Redis for 2ms online reads · Parquet/S3 for offline training
Two-tower retrieval model 64-dim PyTorch embeddings, trained with negative sampling
FAISS ANN search Top-200 candidate retrieval in ~10ms over 5000+ items
LinUCB contextual bandit Explore/exploit re-ranking — solves cold-start and novelty
Sub-100ms P99 SLO Full pipeline: feature fetch → ANN → re-rank → response
PSI drift detection Auto-triggers retraining when feature distribution shifts > 0.2
Prometheus metrics P50/P95/P99 latency, cold-start rate, request count — all instrumented
Live React dashboard 3-panel UI: user simulator · recommendation feed · live metrics
GitHub Actions CI Lint + tests + frontend build on every push

🏗 Architecture

┌─────────────────────────────────────────────────────────────────────────────┐
│                           CLIENT EVENTS                                     │
│                  click · dwell · purchase · skip                            │
└───────────────────────────────┬─────────────────────────────────────────────┘
                                │
                         ┌──────▼──────┐
                         │    KAFKA    │  partitioned by user_id
                         │   broker   │  topic: user-events
                         └──────┬──────┘
                                │
                   ┌────────────▼─────────────┐
                   │     FLINK PROCESSOR      │
                   │  micro-batch windows     │
                   │  session aggregation     │
                   │  category affinity       │
                   └──────┬──────────┬────────┘
                          │          │
              ┌───────────▼──┐  ┌────▼──────────┐
              │    REDIS     │  │    PARQUET     │
              │  (online)    │  │   (offline)    │
              │   ~2ms TTL   │  │  for training  │
              └──────┬───────┘  └────────────────┘
                     │
         ┌───────────▼────────────────────────────────┐
         │              ML SERVING PIPELINE            │
         │                                             │
         │  ① Feature fetch      Redis        ~2ms    │
         │  ② User embedding     PyTorch       ~5ms   │
         │  ③ ANN retrieval      FAISS        ~10ms   │
         │  ④ LinUCB re-rank     bandit        ~5ms   │
         │  ⑤ Response                                │
         │                       P50   ~45ms          │
         │                       P99  <100ms   ✓ SLO  │
         └───────────┬────────────────────────────────┘
                     │
          ┌──────────▼──────────┐
          │   FastAPI  :8000    │  /recommend  /event
          │ + Prometheus        │  /feedback   /docs
          └──────────┬──────────┘
                     │
          ┌──────────▼──────────┐
          │   React Dashboard   │  live feed · latency chart
          │   localhost:5173    │  user simulator · system status
          └─────────────────────┘

ML Stack

Two-Tower Retrieval Model

The same retrieval architecture used at YouTube, Pinterest, and Spotify.

User Tower                          Item Tower
──────────────────                  ──────────────────
user_id embedding (64d)             item_id embedding (64d)
+ category affinity (6d)            + category one-hot (6d)
+ behavioral features (3d)          + price + rating (2d)
        │                                   │
   Linear(73→128)                      Linear(72→128)
      ReLU + Dropout                   ReLU + Dropout
   Linear(128→64)                     Linear(128→64)
   L2 Normalize                        L2 Normalize
        │                                   │
        └──────── dot product ──────────────┘
                      │
               relevance score

Training: BCEWithLogitsLoss, negative sampling (2:1 ratio), early stopping, ReduceLROnPlateau scheduler.

FAISS ANN Index

All 5000 item embeddings → L2 normalized → IndexFlatIP
Query: user embedding → top-200 by inner product (~10ms)

LinUCB Contextual Bandit

Re-ranks the top-200 FAISS candidates using the Upper Confidence Bound algorithm. Balances exploitation (serve what worked before) with exploration (try new items). Each arm = one item. Context = 64-dim user embedding.

UCB score = θᵀx + α√(xᵀA⁻¹x)
Final score = 0.7 × cosine_score + 0.3 × UCB_score

Drift Detection

Population Stability Index (PSI) computed every 5 minutes across click_rate, avg_dwell_ms, purchase_rate, skip_rate. If max_PSI > 0.2, retraining triggers automatically.

PSI < 0.1  → stable        
PSI 0.1–0.2 → slight shift  
PSI > 0.2  → retrain       

Project Structure

realtime-personalization-engine/
│
├── data/
│   ├── seed/generate_mock_data.py     # 1K users · 5K items · 50K events
│   └── raw/                           # generated JSON (git-ignored)
│
├── docker/
│   ├── docker-compose.yml             # Kafka · Zookeeper · Redis · Prometheus
│   ├── Dockerfile                     # production API container
│   └── prometheus.yml                 # scrape config
│
├── event_pipeline/
│   ├── producer.py                    # Kafka event producer
│   ├── flink_processor.py             # streaming feature computation
│   └── consumer.py                    # debug consumer
│
├── feature_store/
│   ├── online_store.py                # Redis read/write with TTL
│   ├── offline_store.py               # Parquet read for training
│   └── feature_definitions.py         # shared schemas + normalization
│
├── models/
│   ├── two_tower.py                   # PyTorch two-tower model
│   ├── bandit.py                      # LinUCB contextual bandit
│   ├── dataset.py                     # training dataset with neg sampling
│   ├── faiss_index.py                 # FAISS ANN index wrapper
│   └── train.py                       # training loop + early stopping
│
├── serving/
│   ├── main.py                        # FastAPI app + CORS + Prometheus mount
│   ├── router.py                      # /recommend /event /feedback endpoints
│   ├── pipeline.py                    # feature→embed→FAISS→bandit→response
│   └── schemas.py                     # Pydantic request/response models
│
├── monitoring/
│   ├── metrics.py                     # Prometheus counters + histograms
│   ├── drift_detector.py              # PSI computation
│   └── retraining_trigger.py          # watches PSI, fires retraining
│
├── frontend/
│   └── src/
│       ├── App.jsx                    # 3-panel layout + auto-refresh
│       ├── components/
│       │   ├── UserSimulator.jsx      # user picker + event firing + pipeline viz
│       │   ├── RecommendFeed.jsx      # rec cards with score bars
│       │   └── MetricsDashboard.jsx   # latency chart + stats + system status
│       └── api/client.js              # axios API client
│
├── notebooks/
│   ├── 01_data_exploration.ipynb      # EDA: distributions, power law, dwell
│   ├── 02_model_training.ipynb        # train + PCA embedding viz
│   └── 03_latency_benchmarks.ipynb    # P50/P95/P99 + per-stage breakdown
│
├── scripts/
│   ├── build_index.py                 # builds FAISS index from trained model
│   └── load_test.py                   # Locust load test (SLO validation)
│
├── tests/
│   ├── test_feature_store.py          # online/offline store unit tests
│   ├── test_two_tower.py              # model + bandit unit tests
│   └── test_drift_detector.py         # PSI computation tests
│
├── .github/workflows/ci.yml           # lint + tests + frontend build
├── docker-compose.yml
├── requirements.txt
├── .env.example
└── README.md

Quickstart

Prerequisites

Tool Version Install
Python 3.10+ python.org
Docker Desktop Latest docker.com
Node.js 18+ nodejs.org

1 — Clone and set up environment

git clone https://github.com/VisheshKamble/realtime-personalization-engine.git
cd realtime-personalization-engine

python -m venv venv
source venv/bin/activate        # Windows: venv\Scripts\activate

pip install -r requirements.txt
cp .env.example .env

2 — Start infrastructure

docker-compose -f docker/docker-compose.yml up -d

# Verify all 4 containers are running
docker ps
NAMES        STATUS       PORTS
kafka        Up           0.0.0.0:9092->9092/tcp
zookeeper    Up           0.0.0.0:2181->2181/tcp
redis        Up           0.0.0.0:6379->6379/tcp
prometheus   Up           0.0.0.0:9090->9090/tcp

3 — Generate data and train model

Run once. Skip this if data/raw/ and data/processed/ already exist.

# Generate 1000 users, 5000 items, 50000 interaction events
python data/seed/generate_mock_data.py

# Train the two-tower model (~5 min on CPU)
python models/train.py

# Build FAISS ANN index from trained embeddings
python scripts/build_index.py

Expected training output:

Dataset: 87,432 samples | 1000 users | 5000 items
Epoch 01/15 | Train: 0.6821 | Val: 0.6543 | Acc: 61.2%
...
Epoch 12/15 | Train: 0.3901 | Val: 0.4012 | Acc: 77.4%
  ✓ Saved best model (val_loss=0.4012)
FAISS index built: 4987 items | dim=64

4 — Run the full stack

Open 4 terminals and run one command in each:

# Terminal 1 — Stream events into Kafka
python event_pipeline/producer.py

# Terminal 2 — Compute features from events
python event_pipeline/flink_processor.py

# Terminal 3 — Serve recommendations
python serving/main.py

# Terminal 4 — Launch dashboard
cd frontend && npm install && npm run dev
Service URL
Dashboard http://localhost:5173
API http://localhost:8000
Swagger docs http://localhost:8000/docs
Metrics http://localhost:8000/metrics
Prometheus http://localhost:9090

API Reference

POST /recommend

Get personalized recommendations for a user.

curl -X POST http://localhost:8000/recommend \
  -H "Content-Type: application/json" \
  -d '{"user_id": "u_1", "n": 10}'
{
  "user_id": "u_1",
  "recommendations": [
    { "item_id": "i_2341", "score": 0.9821, "source": "bandit", "rank": 1 },
    { "item_id": "i_891",  "score": 0.9743, "source": "bandit", "rank": 2 },
    { "item_id": "i_4102", "score": 0.9601, "source": "bandit", "rank": 3 }
  ],
  "latency_ms": 47.3,
  "is_cold_start": false,
  "experiment_id": "exp_001"
}

POST /event

Log a user interaction event (forwarded to Kafka, updates bandit online).

curl -X POST http://localhost:8000/event \
  -H "Content-Type: application/json" \
  -d '{"user_id": "u_1", "item_id": "i_2341", "event_type": "purchase"}'

POST /feedback

Send an explicit reward signal to update the bandit.

curl -X POST http://localhost:8000/feedback \
  -H "Content-Type: application/json" \
  -d '{"user_id": "u_1", "item_id": "i_2341", "reward": 1.0}'

GET /health

{ "status": "ok", "timestamp": 1719300000, "version": "1.0.0" }

Latency Benchmarks

Measured on a single-node MacBook M2 (CPU-only, no GPU).

Percentile Latency
P50 ~45ms
P95 ~72ms
P99 ~89ms
Max ~120ms

SLO: P99 < 100ms ✓

Per-stage breakdown

Stage Avg P99
Feature fetch (Redis) 1.8ms 4.2ms
User embedding (PyTorch) 4.1ms 8.3ms
ANN retrieval (FAISS) 9.3ms 14.1ms
LinUCB re-ranking 4.7ms 9.2ms
FastAPI overhead 5.2ms 9.8ms
Total ~25ms ~46ms

Throughput

Concurrent threads: 10 × 20 requests = 200 total
Elapsed: 1.3s
Throughput: ~154 req/sec
P99 under load: 87ms ✓

Testing

# Run all unit tests
PYTHONPATH=. pytest tests/ -v

# Expected output
tests/test_feature_store.py::test_flatten_category_affinity_basic    PASSED
tests/test_feature_store.py::test_online_store_write_read            PASSED
tests/test_two_tower.py::test_model_forward_shape                    PASSED
tests/test_two_tower.py::test_user_embedding_normalized              PASSED
tests/test_two_tower.py::test_bandit_prefers_rewarded_arm            PASSED
tests/test_drift_detector.py::test_psi_identical_distributions       PASSED
tests/test_drift_detector.py::test_check_drift_detects_shift         PASSED

Load Test

pip install locust
locust -f scripts/load_test.py --host=http://localhost:8000

Open http://localhost:8089 → set 50 users, 10 spawn/sec → watch P99 stay under 100ms.


Notebooks

Notebook Contents
01_data_exploration.ipynb Event distributions, power law analysis, dwell time segmentation, feature store verification
02_model_training.ipynb Training loop, loss curves, PCA embedding visualization, similar item retrieval test
03_latency_benchmarks.ipynb P50/P95/P99 end-to-end, per-stage breakdown, throughput under concurrency
pip install jupyter matplotlib seaborn ipykernel
python -m ipykernel install --user --name=rpe-venv --display-name "RPE (venv)"
jupyter notebook notebooks/

Select kernel RPE (venv) in each notebook.


Design Decisions

Why two-tower instead of matrix factorization? Two-tower supports arbitrary input features (behavioral signals, content features) and can embed new users/items at inference time. MF requires retraining for every new entity.

Why FAISS + re-rank instead of full model scoring? Scoring all 5000 items with the full model would take ~500ms. FAISS retrieves the top-200 in ~10ms, then the full model re-ranks only those. This two-stage pattern is used at YouTube and Pinterest.

Why LinUCB instead of pure ranking? Pure ranking always serves the same items to cold-start users. LinUCB explores new items with an upper confidence bound — over time, it converges on items that maximize reward for each user context.

Why Redis + Parquet dual store? Redis gives sub-5ms feature reads at serving time. Parquet gives rich columnar access for offline training. Both are written on every Flink window — no ETL pipeline needed.


Roadmap

  • Deploy API to Railway / Render (free tier)
  • Deploy frontend to Vercel
  • Grafana dashboard wired to Prometheus
  • A/B experiment framework (shadow deployment)
  • Replace mock data with a real public dataset (MovieLens / Amazon Reviews)
  • Model versioning with MLflow

Author

Vishesh Kamble

GitHub


License

MIT — see LICENSE


If this project helped you, please consider giving it a ⭐

About

Production-grade ML recommendation engine: Kafka → Flink → Two-Tower PyTorch + FAISS ANN + LinUCB Bandit → sub-100ms FastAPI serving with live React dashboard

Topics

Resources

License

Stars

1 star

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors