Blue-green model deployments

pip install scikit-learn pandas numpy pytest

"""A blue-green router: two fully warmed-up environments, one active pointer,
switched with a single atomic assignment.
"""
import time

import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

COLUMNS = ["income", "years", "age"]


def make_training_data(n=400, seed=0):
    rng = np.random.RandomState(seed)
    X = pd.DataFrame({
        "income": rng.uniform(5, 80, n).round(1),
        "years": rng.uniform(0, 10, n).round(1),
        "age": rng.randint(21, 65, n).astype(float),
    })
    score = 0.05 * X["income"] + 0.35 * X["years"] + 0.01 * X["age"] - 3.5
    y = (score + rng.normal(0, 0.8, n) > 0).astype(int)
    return X, y


def load_model(seed) -> object:
    """Standing in for a slow, real model load (reading a large file from disk)."""
    time.sleep(0.05)  # simulated load latency -- this is the cost warming up avoids
    X, y = make_training_data(seed=seed)
    return make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000)).fit(X, y)


class BlueGreenRouter:
    def __init__(self):
        self.models = {}
        self.active = "blue"

    def deploy(self, slot: str, seed: int):
        """Loads and warms up a slot WITHOUT sending it any traffic yet."""
        self.models[slot] = load_model(seed)

    def switch_to(self, slot: str):
        assert slot in self.models, f"{slot} was never deployed -- refusing to switch to it"
        self.active = slot  # a single attribute write -- atomic under the GIL

    def predict(self, applicant: dict) -> tuple[int, str]:
        model = self.models[self.active]
        row = pd.DataFrame([applicant], columns=COLUMNS)
        return int(model.predict(row)[0]), self.active


router = BlueGreenRouter()

start = time.perf_counter()
router.deploy("blue", seed=0)
router.switch_to("blue")
print(f"blue deployed and warm in {time.perf_counter() - start:.3f}s")

answer, slot = router.predict({"income": 60.0, "years": 7.0, "age": 34.0})
print(f"serving from '{slot}': {answer}")

start = time.perf_counter()
router.deploy("green", seed=1)  # deployed AHEAD of the switch, while blue keeps serving
print(f"green deployed and warm in {time.perf_counter() - start:.3f}s (blue kept serving throughout)")

router.switch_to("green")
answer, slot = router.predict({"income": 60.0, "years": 7.0, "age": 34.0})
print(f"serving from '{slot}': {answer}")

import threading
import time
from bluegreen import BlueGreenRouter

router = BlueGreenRouter()
router.deploy("blue", seed=0)
router.deploy("green", seed=1)

applicant = {"income": 60.0, "years": 7.0, "age": 34.0}
results = []
stop = threading.Event()

def hammer():
    while not stop.is_set():
        _, slot = router.predict(applicant)
        results.append(slot)

threads = [threading.Thread(target=hammer) for _ in range(8)]
for t in threads:
    t.start()

time.sleep(0.05)
router.switch_to("green")   # the atomic cutover, while 8 threads are mid-flight
time.sleep(0.05)

stop.set()
for t in threads:
    t.join()

print(f"total requests during the switch: {len(results)}")
print(f"served by blue:  {results.count('blue')}")
print(f"served by green: {results.count('green')}")
print(f"served by neither (would mean a bug): {len(results) - results.count('blue') - results.count('green')}")

def unsafe_switch_to(slot):
    """A tempting-looking but unsafe 'clear then set' switch -- two separate steps."""
    router.active = None   # step 1: "clear" the old one first
    time.sleep(0.001)      # any gap here is a real production risk
    router.active = slot   # step 2: set the new one

import pytest
from bluegreen import BlueGreenRouter


def test_switching_to_an_undeployed_slot_is_refused():
    router = BlueGreenRouter()
    router.deploy("blue", seed=0)
    with pytest.raises(AssertionError):
        router.switch_to("green")


def test_switch_and_instant_rollback():
    router = BlueGreenRouter()
    router.deploy("blue", seed=0)
    router.deploy("green", seed=1)
    applicant = {"income": 60.0, "years": 7.0, "age": 34.0}

    router.switch_to("green")
    _, slot = router.predict(applicant)
    assert slot == "green"

    router.switch_to("blue")  # rollback -- blue was kept warm the whole time
    _, slot = router.predict(applicant)
    assert slot == "blue"


def test_the_idle_slot_is_never_used_for_real_traffic():
    router = BlueGreenRouter()
    router.deploy("blue", seed=0)
    router.deploy("green", seed=1)
    applicant = {"income": 60.0, "years": 7.0, "age": 34.0}
    for _ in range(20):
        _, slot = router.predict(applicant)
        assert slot == "blue"  # never green, until switch_to is called

pytest test_bluegreen.py -q