perf(router): cache classifier model instead of reloading it on every request

optillm/plugins/router_plugin.py

@@ -1,3 +1,5 @@
+import threading
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -47,12 +49,23 @@ def forward(self, input_ids, attention_mask, effort):
         logits = self.classifier(combined_input)
         return logits
 
-def load_optillm_model():
+# Cache the loaded router model/tokenizer/device across requests. Building it is
+# expensive and one-time: it deserializes the ~400M-param ModernBERT-large base
+# model, does a HuggingFace Hub round-trip to fetch the fine-tuned safetensors, and
+# constructs the tokenizer. The classifier is stateless at inference time (see
+# predict_approach: model.eval() + torch.no_grad), so the loaded objects are safe
+# to reuse. The lock makes the first concurrent load run exactly once under the
+# threaded Flask server; subsequent calls take the lock-free fast path.
+_model_cache = None
+_model_cache_lock = threading.Lock()
+
+
+def _load_optillm_model():
     device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
     # Load the base model
     base_model = AutoModel.from_pretrained(BASE_MODEL)
     # Create the OptILMClassifier
-    model = OptILMClassifier(base_model, num_labels=len(APPROACHES))  
+    model = OptILMClassifier(base_model, num_labels=len(APPROACHES))
     model.to(device)
     # Download the safetensors file
     safetensors_path = hf_hub_download(repo_id=OPTILLM_MODEL_NAME, filename="model.safetensors")
@@ -62,6 +75,17 @@ def load_optillm_model():
     tokenizer = AutoTokenizer.from_pretrained(OPTILLM_MODEL_NAME)
     return model, tokenizer, device
 
+
+def load_optillm_model():
+    # Double-checked locking: return the cached bundle immediately when present,
+    # and only serialize the very first (concurrent) load.
+    global _model_cache
+    if _model_cache is None:
+        with _model_cache_lock:
+            if _model_cache is None:
+                _model_cache = _load_optillm_model()
+    return _model_cache
+
 def preprocess_input(tokenizer, system_prompt, initial_query):
     combined_input = f"{system_prompt}\n\nUser: {initial_query}"
     encoding = tokenizer.encode_plus(

tests/test_plugins.py

@@ -428,6 +428,104 @@ def test_no_relative_import_errors():
                 raise
 
 
+# ---------------------------------------------------------------------------
+# Router plugin: the classifier model must be loaded once and cached, not
+# reloaded on every request. These tests mock the heavy loaders so they need
+# no network or real weights.
+# ---------------------------------------------------------------------------
+import threading as _threading
+import contextlib as _contextlib
+from unittest import mock as _mock
+
+
+class _FakeRouterConfig:
+    hidden_size = 768
+
+
+def _install_router_loader_mocks(stack, load_counter):
+    """Patch router_plugin's heavy loaders; load_counter grows once per base-model load."""
+    from optillm.plugins import router_plugin
+
+    def fake_from_pretrained(*args, **kwargs):
+        load_counter.append(1)
+        base = _mock.MagicMock()
+        base.config = _FakeRouterConfig()
+        return base
+
+    stack.enter_context(_mock.patch.object(
+        router_plugin.AutoModel, "from_pretrained", side_effect=fake_from_pretrained))
+    stack.enter_context(_mock.patch.object(
+        router_plugin, "hf_hub_download", return_value="/tmp/fake.safetensors"))
+    stack.enter_context(_mock.patch.object(
+        router_plugin, "load_model", return_value=None))
+    stack.enter_context(_mock.patch.object(
+        router_plugin.AutoTokenizer, "from_pretrained", return_value=_mock.MagicMock()))
+    return router_plugin
+
+
+def test_router_plugin_caches_model():
+    """load_optillm_model() must load the base model once across repeated calls."""
+    from optillm.plugins import router_plugin
+    router_plugin._model_cache = None
+    loads = []
+    try:
+        with _contextlib.ExitStack() as stack:
+            _install_router_loader_mocks(stack, loads)
+            first = router_plugin.load_optillm_model()
+            for _ in range(5):
+                assert router_plugin.load_optillm_model() is first, \
+                    "load_optillm_model() must return the cached bundle"
+        assert len(loads) == 1, (
+            f"expected exactly 1 base-model load across 6 calls, got {len(loads)} "
+            "(model is being reloaded per request instead of cached)")
+    finally:
+        router_plugin._model_cache = None
+
+
+def test_router_plugin_cache_bundle_shape():
+    """The cached value is the (model, tokenizer, device) triple callers unpack."""
+    from optillm.plugins import router_plugin
+    router_plugin._model_cache = None
+    try:
+        with _contextlib.ExitStack() as stack:
+            _install_router_loader_mocks(stack, [])
+            bundle = router_plugin.load_optillm_model()
+        assert isinstance(bundle, tuple) and len(bundle) == 3, "expected a 3-tuple bundle"
+        model, tokenizer, device = bundle
+        assert isinstance(model, router_plugin.OptILMClassifier)
+        assert tokenizer is not None and device is not None
+    finally:
+        router_plugin._model_cache = None
+
+
+def test_router_plugin_concurrent_single_load():
+    """Concurrent first-time access must still load the model exactly once."""
+    from optillm.plugins import router_plugin
+    router_plugin._model_cache = None
+    loads = []
+    results = []
+    try:
+        with _contextlib.ExitStack() as stack:
+            _install_router_loader_mocks(stack, loads)
+            barrier = _threading.Barrier(8)
+
+            def worker():
+                barrier.wait()  # maximize contention on the first load
+                results.append(router_plugin.load_optillm_model())
+
+            threads = [_threading.Thread(target=worker) for _ in range(8)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+        assert len(loads) == 1, (
+            f"expected exactly 1 load under concurrency, got {len(loads)} "
+            "(the double-checked lock is not serializing the first load)")
+        assert all(r is results[0] for r in results), "threads saw different cached bundles"
+    finally:
+        router_plugin._model_cache = None
+
+
 if __name__ == "__main__":
     print("Running plugin tests...")
 
@@ -455,6 +553,24 @@ def test_no_relative_import_errors():
     except Exception as e:
         print(f"❌ Memory plugin persistence test failed: {e}")
 
+    try:
+        test_router_plugin_caches_model()
+        print("✅ Router plugin model caching test passed")
+    except Exception as e:
+        print(f"❌ Router plugin model caching test failed: {e}")
+
+    try:
+        test_router_plugin_cache_bundle_shape()
+        print("✅ Router plugin cache bundle shape test passed")
+    except Exception as e:
+        print(f"❌ Router plugin cache bundle shape test failed: {e}")
+
+    try:
+        test_router_plugin_concurrent_single_load()
+        print("✅ Router plugin concurrent single-load test passed")
+    except Exception as e:
+        print(f"❌ Router plugin concurrent single-load test failed: {e}")
+
     try:
         test_genselect_plugin()
         print("✅ GenSelect plugin test passed")