diff --git a/optillm/plugins/router_plugin.py b/optillm/plugins/router_plugin.py index 3c9b138..2a6f5a2 100644 --- a/optillm/plugins/router_plugin.py +++ b/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( diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 0f769d5..e7f6b74 100644 --- a/tests/test_plugins.py +++ b/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")