Implementing Feature Flags for ML Model A/B Testing

Shipping a new ML model to production is not like shipping a UI feature. If a button breaks, users complain and you roll back. If your recommendation model degrades, you might not notice for days — by then, revenue has bled. Feature flags give you the control plane you need: route traffic, compare models side by side, and roll back in seconds without a redeploy.

This article walks through a production-grade implementation of feature flags for ML model A/B testing in Python, using a simple but extensible pattern.

Why Feature Flags, Not Blue/Green Deploys

Blue/green deploys solve infrastructure rollout. Feature flags solve logical rollout — they let you separate "code is deployed" from "code is active." For ML, this distinction matters for several reasons:

• You want to run model A and model B simultaneously on the same infrastructure
• You want to split traffic by user segment, not just percentage
• You want to roll back a model without touching your deployment pipeline

A feature flag for ML A/B testing typically needs three things: a consistent assignment function (same user → same model), percentage-based traffic splits, and observable metadata per request.

The Core Assignment Function

Start with a flag definition and a deterministic variant assignment:

from dataclasses import dataclass
from typing import Optional
import hashlib

@dataclass
class ModelFlag:
    name: str
    control_model: str       # "model_a"
    treatment_model: str     # "model_b"
    traffic_split: float     # 0.0 to 1.0 → fraction sent to treatment
    enabled: bool = True
    segment_key: Optional[str] = None  # restrict to a user segment

def assign_variant(flag: ModelFlag, user_id: str, segment: Optional[str] = None) -> str:
    """Deterministic: same user_id always gets the same variant for this flag."""
    if not flag.enabled:
        return flag.control_model

    if flag.segment_key and segment != flag.segment_key:
        return flag.control_model

    # Hash user_id + flag name: consistent and flag-specific
    raw = f"{flag.name}:{user_id}".encode()
    bucket = int(hashlib.md5(raw).hexdigest(), 16) % 1000
    threshold = int(flag.traffic_split * 1000)

    return flag.treatment_model if bucket < threshold else flag.control_model

The hashlib.md5 trick is intentional: it is fast, produces a uniform distribution, and makes assignment stable across restarts. Do not use random.random() — it breaks consistency between requests.

Loading Flags at Runtime Without a Redeploy

Hard-coding flags defeats the purpose. You want to change splits without redeploying. The simplest production-ready approach: a JSON config file (or S3 object) polled on a background thread.

import json
import threading
import time
from pathlib import Path

class FlagStore:
    def __init__(self, config_path: str, poll_interval: int = 30):
        self._path = Path(config_path)
        self._flags: dict[str, ModelFlag] = {}
        self._lock = threading.RLock()
        self._load()
        self._start_polling(poll_interval)

    def _load(self):
        try:
            raw = json.loads(self._path.read_text())
            with self._lock:
                self._flags = {
                    name: ModelFlag(**cfg)
                    for name, cfg in raw.items()
                }
        except Exception as e:
            print(f"[FlagStore] reload failed: {e}")  # keep stale config, do not crash

    def _start_polling(self, interval: int):
        def loop():
            while True:
                time.sleep(interval)
                self._load()
        t = threading.Thread(target=loop, daemon=True)
        t.start()

    def get(self, flag_name: str) -> Optional[ModelFlag]:
        with self._lock:
            return self._flags.get(flag_name)

The RLock matters: the GIL does not protect dict mutations during _load. The daemon=True thread dies cleanly with the process.

Your flags.json looks like:

{
  "recommendation_v2": {
    "name": "recommendation_v2",
    "control_model": "rec_model_v1",
    "treatment_model": "rec_model_v2",
    "traffic_split": 0.1,
    "enabled": true
  }
}

Start at 10%, watch metrics, ramp to 50%, then 100%. To kill the experiment: set enabled: false or traffic_split: 0.0 and reload.

Integrating With Your Inference Layer

The flag system is useless without routing. Assuming you have a model registry:

class ModelRouter:
    def __init__(self, flag_store: FlagStore, registry: dict):
        self.flags = flag_store
        self.registry = registry  # {"model_name": model_instance}

    def predict(self, flag_name: str, user_id: str, payload: dict,
                segment: str = None) -> dict:
        flag = self.flags.get(flag_name)
        if flag is None:
            raise ValueError(f"Unknown flag: {flag_name}")

        assigned_model = assign_variant(flag, user_id, segment)
        model = self.registry[assigned_model]
        result = model.predict(payload)

        # Attach observability metadata — your analytics pipeline needs this
        return {
            "prediction": result,
            "_meta": {
                "flag": flag_name,
                "model": assigned_model,
                "user_id": user_id,
                "variant": "treatment" if assigned_model == flag.treatment_model else "control",
            }
        }

Log _meta to your data warehouse on every request. Without it, you cannot compute experiment results after the fact.

What to Measure and When to Stop

A/B testing a model without a stopping rule is wishful thinking. Before you start the experiment, commit to these in writing:

1. Primary metric: what does "better" mean? Click-through rate, conversion, latency p99, error rate?
2. Minimum detectable effect: what delta is worth shipping? Do not stop at the first statistically significant fluctuation.
3. Sample size: run a power analysis (statsmodels.stats.power) upfront to know how many users you need before the result is reliable.

A common trap: checking significance daily and stopping when p < 0.05. This inflates false positive rates. Either pre-commit to a fixed runtime, or use sequential testing (Bayesian or SPRT) if early stopping is a hard requirement.

The free security and deployment checklists we publish include an ML deployment checklist that covers experiment hygiene alongside infrastructure hardening — the two concerns are more entangled than they look.

Operational Concerns That Tutorials Skip

Carryover effects: if users experience model B today and you roll back to A tomorrow, their behavior is already influenced. Plan your rollback window accordingly and do not reuse experiment groups too quickly.

Cache invalidation: if you cache model outputs per user, include the assigned model name in the cache key. Otherwise a user rolling from treatment to control gets stale treatment outputs.

Logging volume: at scale, logging every _meta block adds up. Sampling at 1–5% is fine for aggregate metrics. Use structured logs so you can filter by flag and variant without regex gymnastics.

Flag hygiene: once an experiment concludes, remove the flag and ship the winning model directly. Stale flags accumulate technical debt faster than stale comments.

The Takeaway

Feature flags for ML A/B testing are not complicated, but they require discipline: stable assignment, observable metadata, pre-defined stopping rules, and cleanup. The pattern above is intentionally thin — no external service, no SDK dependency. For small to medium traffic, a polled JSON file and a 60-line router get you 90% of the value of heavier solutions like LaunchDarkly or Statsig.

If you need more sophisticated segmentation — multi-variate tests, contextual bandits, online learning — the structure scales: swap assign_variant for a bandit policy and log rewards back to the flag store.

---

I run AYI NEDJIMI Consultants, a cybersecurity consulting firm. We publish free security hardening checklists — PDF and Excel.