Skip to content

Emart29/ml-feature-store

BranchesTags

Repository files navigation

ML Feature Store

Python FastAPI PostgreSQL Redis Streamlit Docker License

A production-grade lightweight feature store for applied ML engineering. Handles offline storage with point-in-time correct joins (PostgreSQL), online serving with sub-10ms latency (Redis), feature versioning, drift detection (Evidently AI), a CLI, Python SDK, and Streamlit registry browser — all in one repo.

The Problem

Machine-learning models run on "features" — the prepared input signals they make decisions from (a customer's average spend, days since last login, a risk score). The catch: each feature has to be computed identically in two very different settings — once in bulk over historical data when the model is trained, and again instantly, one record at a time, when the model is live and serving real users. When those two drift apart even slightly, the model quietly performs worse in production than it did in testing. On top of that, teams keep rebuilding the same features for every new project, and small mistakes in how training data is assembled can let a model accidentally "see the future" and look better than it really is.

What This Does

This is a single place to define a feature once and reuse it everywhere, consistently and safely. You write a feature's logic one time and the store handles both worlds:

  • an offline store (PostgreSQL) for building training datasets, with point-in-time-correct joins so there's no data leakage,
  • an online store (Redis) that serves the freshest feature values in under 10 milliseconds when a live model asks,
  • versioning so you can document, inspect, and roll any feature back to a previous version,
  • drift detection on every ingestion run, so you're warned when incoming data starts to look different,
  • and a browsable UI plus a Python SDK (store.get_features(entity_id, [...])) so anyone can find and use features without reinventing them.

Tech Stack

Layer Tool
Offline store PostgreSQL + SQLAlchemy (async)
Online store Redis — target <10ms p99
Ingestion Polars
Serving API FastAPI + Uvicorn
CLI Click
Scheduler APScheduler
Drift detection Evidently AI
Monitoring Prometheus + Grafana
UI Streamlit
SDK Python (sync + async)

Architecture

Data Sources (CSV, PostgreSQL, Kafka)
         │
         ▼
Ingestion Pipeline (Polars)
    ├── Schema validation
    ├── Feature transforms  (@feature decorator)
    └── Drift detection     (Evidently AI)
         │
    ┌────┴────┐
    ▼         ▼
Offline Store   Online Store
(PostgreSQL)    (Redis, <10ms)
    │               │
    └──────┬─────────┘
           ▼
  Feature Registry (metadata, versions, lineage)
           │
           ▼
  Serving API (FastAPI)  ←  Python SDK
           │
           ▼
  Streamlit UI  (browse, search, version history, drift)

Services

Service URL Description
Feature Store API http://localhost:8000 Feature retrieval, ingestion, registry
API Docs (Swagger) http://localhost:8000/docs Interactive endpoint documentation
Prometheus metrics http://localhost:8000/metrics Latency, throughput, error counters
Streamlit UI http://localhost:8501 Browse features, stats, version diffs, drift

Quick Start — Docker

Prerequisites: Docker Desktop installed and running, plus the shared infrastructure stack.

# 1. Start shared infrastructure (PostgreSQL + Redis + MinIO + Kafka + Prometheus + Grafana)
git clone https://github.com/Emart29/ml-platform-infra
cd ml-platform-infra && docker compose up -d && cd ..

# 2. Clone and start the feature store
git clone https://github.com/Emart29/ml-feature-store
cd ml-feature-store
docker compose up

Docker will automatically:

  1. Create the database schema in the shared PostgreSQL instance
  2. Download and ingest the Heart Disease demo dataset (303 patients, 6 features)
  3. Sync features from PostgreSQL to Redis
  4. Start the Feature Store API on port 8000
  5. Start the Streamlit registry browser on port 8501

Open http://localhost:8501 — the feature registry is live.

# Stop containers (keep data)
docker compose down

# Stop and delete all data
docker compose down -v

Local Setup

Requires Python 3.10+.

# 1. Start PostgreSQL + Redis (uses the shared infra stack)
git clone https://github.com/Emart29/ml-platform-infra
cd ml-platform-infra && docker compose up -d && cd ..

# 2. Install
git clone https://github.com/Emart29/ml-feature-store
cd ml-feature-store
python -m venv .venv
.venv\Scripts\activate      # Windows
# source .venv/bin/activate  # Linux / macOS
pip install -r requirements.txt
pip install -e .

# 3. Configure (optional — defaults connect to localhost)
copy .env.example .env

# 4. Create tables + ingest demo data + start
python examples/heart_disease/ingest.py
featurestore sync --full
featurestore serve           # API on :8000 (keep this running)
featurestore ui              # Streamlit on :8501 (new terminal)

Heart Disease Demo

Demonstrates the full feature store lifecycle end-to-end: ingestion → versioning → training → serving → drift detection.

# Full scripted demo (requires API running on :8000)
python examples/heart_disease/demo.py

# Step by step
python examples/heart_disease/ingest.py          # register + ingest 6 features
featurestore sync --full                         # push feature values to Redis
python examples/heart_disease/train.py           # train GBM via SDK point-in-time join
uvicorn examples.heart_disease.serve:app --port 8001  # prediction API
curl -X POST http://localhost:8001/predict/1     # predict for patient 1

Or open the Jupyter notebook for an interactive walkthrough:

jupyter lab examples/heart_disease/notebook.ipynb

The notebook covers: registry exploration → point-in-time correct training data → model training → online serving latency benchmark (p50/p95/p99) → batch inference → feature statistics and drift history.

CLI

# List registered features
featurestore list
featurestore list --search risk --entity-type patient --status active

# Inspect a feature (metadata, statistics, run history)
featurestore describe composite_risk_score

# Compare two versions (code diff + statistics)
featurestore diff composite_risk_score 1 2

# View statistics history
featurestore stats composite_risk_score
featurestore stats composite_risk_score --plot --runs 10

# Push features from a CSV (with custom transform file)
featurestore push age_normalized \
  --source-type csv \
  --source-path examples/heart_disease/data/heart.csv \
  --entity-col patient_id \
  --features-file examples/heart_disease/features.py

# Dry-run to validate without writing
featurestore push age_normalized --source-type csv --source-path data.csv --dry-run

# Sync offline store → Redis
featurestore sync                            # incremental (since last watermark)
featurestore sync --full                     # full resync
featurestore sync --feature age_normalized   # one feature only

# Register a new feature interactively
featurestore register \
  --name conversion_rate \
  --entity-type user \
  --dtype float \
  --description "7-day rolling conversion rate" \
  --owner ml-team \
  --tags conversion,rolling

# Start servers
featurestore serve                           # API on :8000
featurestore serve --reload                  # with hot-reload (dev)
featurestore ui                              # Streamlit on :8501

Python SDK

from sdk.client import FeatureStoreClient

store = FeatureStoreClient(url="http://localhost:8000")

# Single entity — online store (Redis, <10ms)
features = store.get_features(
    entity_id="42",
    features=["age_normalized", "cholesterol_risk_score", "composite_risk_score"],
    entity_type="patient",
)
# {"age_normalized": 0.63, "cholesterol_risk_score": 1.21, ...}

# Point-in-time correct training data — no leakage
import pandas as pd

labels_df = pd.DataFrame({
    "patient_id":     ["1", "2", "3"],
    "diagnosis_date": ["2024-01-31T00:00:00Z"] * 3,
    "target":         [1, 0, 1],
})
train_df = store.get_historical_features(
    entity_df=labels_df,
    features=["age_normalized", "cholesterol_risk_score", "composite_risk_score"],
    entity_column="patient_id",
    timestamp_column="diagnosis_date",
    entity_type="patient",
)
# pandas DataFrame: original columns + feature columns, aligned by position

# Batch online serving — one Redis round-trip for all entities
results = store.get_features_batch(
    requests=[
        {"entity_id": str(i), "entity_type": "patient", "features": ["age_normalized"]}
        for i in range(1, 11)
    ]
)
# [{"entity_id": "1", "features": {"age_normalized": 0.63}, "missing_features": []}, ...]

# Trigger ingestion from code
result = store.push_features(
    "age_normalized",
    source_config={
        "type": "csv",
        "path": "examples/heart_disease/data/heart.csv",
        "entity_column": "patient_id",
        "timestamp_column": "diagnosis_date",
    },
    wait=True,
)
print(result.records_processed, result.drift_detected)

Async client — for FastAPI handlers and other async contexts:

from sdk.client import AsyncFeatureStoreClient

async with AsyncFeatureStoreClient(url="http://localhost:8000") as store:
    features = await store.get_features("42", ["age_normalized"], entity_type="patient")
    batch = await store.get_features_batch(requests=[...])

Define Features

Features are plain Python functions decorated with @feature. The function receives a Polars DataFrame (the raw source data) and returns a DataFrame with three columns: entity_id, value, computed_at.

import polars as pl
from ingestion.transforms import feature

@feature(
    name="age_normalized",
    entity_type="patient",
    data_type="float",
    description="Patient age normalized to [0, 1] using min-max scaling",
    owner="ml-team",
    tags=["demographic", "normalized"],
)
def age_normalized(df: pl.DataFrame) -> pl.DataFrame:
    age = pl.col("age").cast(pl.Float64)
    age_norm = (age - age.min()) / (age.max() - age.min())
    return df.select([
        pl.col("patient_id").alias("entity_id"),
        age_norm.alias("value"),
        pl.lit(None).cast(pl.Datetime("us", "UTC")).alias("computed_at"),
    ])

Every time the transform function body changes, a new version is automatically created (SHA-256 hash of source code). Use featurestore diff to compare code and statistics between versions.

Six features are defined for the Heart Disease demo in examples/heart_disease/features.py: age_normalized, cholesterol_risk_score, heart_rate_reserve, chest_pain_severity, composite_risk_score, thalassemia_risk.

Configuration

Settings are read from environment variables or a .env file (copy .env.example to get started):

Variable Default Description
POSTGRES_URL postgresql+asyncpg://postgres:postgres@localhost:5432/ml_platform PostgreSQL connection string
REDIS_URL redis://localhost:6379 Redis connection string
SYNC_INTERVAL_MINUTES 15 Scheduled offline → online sync interval
FEATURE_VALUE_TTL_DAYS 7 Redis TTL for online feature values
LOG_LEVEL INFO Logging level (DEBUG, INFO, WARNING)

Project Layout

ml-feature-store/
├── store/
│   ├── offline.py          # PostgreSQL tables + point-in-time joins
│   ├── online.py           # Redis serving + TTL management
│   ├── registry.py         # Feature metadata, versions, search, lineage
│   └── sync.py             # Scheduled offline → online sync (APScheduler)
├── ingestion/
│   ├── pipeline.py         # Orchestrates load → transform → validate → store
│   ├── transforms.py       # @feature decorator + FeatureTransformRegistry
│   └── validation.py       # Schema checks + Evidently AI drift detection
├── serving/
│   └── api.py              # FastAPI: online, offline, registry, ingestion endpoints
├── cli/
│   └── main.py             # featurestore CLI (Click)
├── sdk/
│   ├── client.py           # FeatureStoreClient + AsyncFeatureStoreClient
│   └── exceptions.py       # Typed exceptions
├── ui/
│   ├── app.py              # Streamlit registry browser (5 pages)
│   └── api_client.py       # HTTP client for the UI
├── db/
│   ├── models.py           # SQLAlchemy ORM models
│   └── base.py             # Async engine + session factory
├── examples/
│   └── heart_disease/
│       ├── features.py     # 6 @feature definitions for cardiac data
│       ├── ingest.py       # Download UCI dataset, register + ingest
│       ├── train.py        # Train GBM on feature-store training data
│       ├── serve.py        # Real-time prediction API (FastAPI, port 8001)
│       ├── demo.py         # Scripted end-to-end demo
│       └── notebook.ipynb  # Jupyter walkthrough (7 steps)
├── scripts/
│   └── docker_init.py      # One-shot Docker setup (tables + ingest + sync)
├── tests/
│   └── test_api.py         # Integration test suite (requires live API)
├── docker-compose.yml
├── Dockerfile
├── pyproject.toml
└── .env.example

API Reference

Method Endpoint Description
GET /health Health check with Redis + PostgreSQL latency
GET /features/online Get features for one entity from Redis
POST /features/online/batch Get features for multiple entities (pipeline)
POST /features/offline/historical Point-in-time correct feature join
GET /registry/features List all features (search, filter, paginate)
POST /registry/features Register a new feature
GET /registry/features/{name} Feature detail with stats + version history
GET /registry/features/{name}/versions All versions for a feature
GET /registry/features/{name}/statistics Latest distribution statistics
GET /registry/features/{name}/drift Drift detection run history
POST /ingestion/run Trigger an ingestion run
GET /ingestion/runs List ingestion run history
POST /ingestion/sync Trigger offline → online sync
GET /metrics Prometheus metrics

Full interactive docs: http://localhost:8000/docs

About

Production-grade ML feature store: point-in-time correct joins, Redis online serving <10ms, feature versioning, drift detection, CLI, Python SDK, and Streamlit registry. Heart Disease demo included.

Topics

evidentlyfastapifeature-storemachine-learningml-platform-inframlopspostgresqlpythonredisstreamlit

Resources

Readme
Activity

Stars

1 star

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages