Swapping weights with no downtime

pip install scikit-learn pandas numpy joblib pytest

"""An in-process model host: swaps the live model's weights inside a single
running server, with no restart and no dropped requests.
"""
import os
import tempfile

import joblib
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 train_and_save(path: str, seed: int):
    """Writes a model to disk SAFELY: write to a temp file, then atomically
    rename it into place. A reader can never observe a half-written file.
    """
    X, y = make_training_data(seed=seed)
    model = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000)).fit(X, y)
    directory = os.path.dirname(path) or "."
    fd, tmp_path = tempfile.mkstemp(dir=directory)
    os.close(fd)
    joblib.dump(model, tmp_path)
    os.replace(tmp_path, path)  # atomic on POSIX and on Windows (same volume)


class ModelHost:
    def __init__(self, path: str):
        self.path = path
        self._model = joblib.load(path)

    def reload(self):
        """Loads the NEW model fully, off to the side, then swaps in one step."""
        new_model = joblib.load(self.path)   # slow work happens BEFORE the swap
        self._model = new_model              # single attribute write -- atomic

    def predict(self, applicant: dict) -> int:
        row = pd.DataFrame([applicant], columns=COLUMNS)
        return int(self._model.predict(row)[0])


train_and_save("model_v1.joblib", seed=0)
host = ModelHost("model_v1.joblib")
print("initial prediction:", host.predict({"income": 60.0, "years": 7.0, "age": 34.0}))

train_and_save("model_v1.joblib", seed=1)  # a fresh version, same filename, written atomically
host.reload()
print("prediction after hot reload:", host.predict({"income": 60.0, "years": 7.0, "age": 34.0}))

import threading
import time
from host import ModelHost, train_and_save

train_and_save("model_v2.joblib", seed=0)
host = ModelHost("model_v2.joblib")

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

def hammer():
    while not stop.is_set():
        try:
            results.append(host.predict(applicant))
        except Exception as e:
            errors.append(repr(e))

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

time.sleep(0.05)
train_and_save("model_v2.joblib", seed=1)  # write the new version to disk, atomically
host.reload()                               # load it fully, then swap in one step
time.sleep(0.05)

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

print(f"total predictions during the reload: {len(results)}")
print(f"errors: {len(errors)}")

def unsafe_writer():
    """The unsafe version: writes straight over the live file, no temp file, no rename."""
    X2, y2 = make_training_data(seed=1)
    new_model = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000)).fit(X2, y2)
    joblib.dump(new_model, PATH)   # overwrites the file readers are actively reading

import glob
from host import ModelHost, train_and_save


def test_reload_actually_changes_the_served_predictions_where_they_differ():
    train_and_save("model_test.joblib", seed=0)
    host = ModelHost("model_test.joblib")
    applicant = {"income": 7.6, "years": 7.6, "age": 53.0}  # a real point where the two models disagree
    before = host.predict(applicant)

    train_and_save("model_test.joblib", seed=1)
    host.reload()
    after = host.predict(applicant)

    assert (before, after) == (0, 1)


def test_atomic_save_leaves_no_temp_file_behind():
    before = set(glob.glob("tmp*"))
    train_and_save("model_test3.joblib", seed=0)
    after = set(glob.glob("tmp*"))
    assert after == before  # the temp file was renamed away, not left sitting around

pytest test_host.py -q