From 69392586c8ea10b11676cdbc555bfb8731f33c88 Mon Sep 17 00:00:00 2001 From: SuperMarioYL Date: Sun, 12 Jul 2026 11:27:06 +0800 Subject: [PATCH 1/2] perf(router): cache classifier model instead of reloading it per request MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit router_plugin.run() called load_optillm_model() on every request, which deserializes the ~400M-param ModernBERT-large base model, does a HuggingFace Hub round-trip for the fine-tuned safetensors, and rebuilds the tokenizer — all just to run a single forward pass for the routing decision. That turns an O(1 forward pass) routing step into O(load a transformer) per call. Memoize the loaded (model, tokenizer, device) bundle behind a module-level singleton guarded by a threading.Lock (double-checked locking), the same idiom already used by CacheManager in optillm/inference.py. The classifier is stateless at inference (predict_approach runs under model.eval() + torch.no_grad), so reusing the loaded objects across requests is correctness-safe. A failed load is not cached (the global is only assigned on success), so run()'s existing except-fallback still retries on the next call. Add tests/test_router_plugin_caching.py: repeated calls load the base model exactly once, the cached value is the expected (model, tokenizer, device) triple, and concurrent first-time access loads exactly once. All tests mock the heavy loaders, so they need no network or weights. --- optillm/plugins/router_plugin.py | 28 +++++- tests/test_router_plugin_caching.py | 146 ++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 2 deletions(-) create mode 100644 tests/test_router_plugin_caching.py diff --git a/optillm/plugins/router_plugin.py b/optillm/plugins/router_plugin.py index 3c9b1381..2a6f5a21 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_router_plugin_caching.py b/tests/test_router_plugin_caching.py new file mode 100644 index 00000000..b006f03d --- /dev/null +++ b/tests/test_router_plugin_caching.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +""" +Tests for the router plugin's model caching. + +The router plugin classifies each incoming query to pick an optimization +approach. Loading its classifier is expensive and one-time (a ~400M-param +ModernBERT-large deserialize + a HuggingFace Hub round-trip + tokenizer build), +so it must be loaded once and reused across requests rather than reloaded on +every call. These tests pin that behavior down without downloading the real +model by mocking the underlying loaders. +""" + +import os +import sys +import threading + +# Try to import pytest, but don't fail if it's not available (matches the +# repo convention in tests/test_plugins.py so this runs under CI both ways). +try: + import pytest +except ImportError: + pytest = None + +from unittest import mock + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from optillm.plugins import router_plugin + + +class _FakeConfig: + hidden_size = 768 + + +def _install_loader_mocks(stack, from_pretrained_counter): + """Patch the heavy loaders so load_optillm_model() runs without any network + or real model weights. `from_pretrained_counter` counts base-model loads.""" + + def fake_from_pretrained(*args, **kwargs): + from_pretrained_counter.append(1) + base = mock.MagicMock() + base.config = _FakeConfig() + 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()) + ) + + +def _reset_cache(): + router_plugin._model_cache = None + + +def test_repeated_calls_load_model_once(): + """Calling load_optillm_model() N times must load the base model exactly once.""" + _reset_cache() + loads = [] + import contextlib + with contextlib.ExitStack() as stack: + _install_loader_mocks(stack, loads) + + first = router_plugin.load_optillm_model() + for _ in range(5): + again = router_plugin.load_optillm_model() + # Same cached bundle object returned every time. + assert again 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)" + ) + _reset_cache() + + +def test_returns_expected_bundle_shape(): + """The cached value is the (model, tokenizer, device) triple callers expect.""" + _reset_cache() + import contextlib + loads = [] + with contextlib.ExitStack() as stack: + _install_loader_mocks(stack, loads) + 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 + assert device is not None + _reset_cache() + + +def test_concurrent_calls_load_model_once(): + """Under concurrent first-time access the expensive load must run only once.""" + _reset_cache() + import contextlib + loads = [] + results = [] + with contextlib.ExitStack() as stack: + _install_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)" + ) + # Every thread observed the same cached bundle. + assert all(r is results[0] for r in results), "threads saw different cached bundles" + _reset_cache() + + +if __name__ == "__main__": + for name, fn in [ + ("repeated calls load model once", test_repeated_calls_load_model_once), + ("returns expected bundle shape", test_returns_expected_bundle_shape), + ("concurrent calls load model once", test_concurrent_calls_load_model_once), + ]: + try: + fn() + print(f"✅ {name}") + except Exception as e: + print(f"❌ {name}: {e}") + raise + + print("\nDone!") From 1fdaa65c0d49288a997ead65845ad61bf866e063 Mon Sep 17 00:00:00 2001 From: SuperMarioYL Date: Sun, 12 Jul 2026 11:42:35 +0800 Subject: [PATCH 2/2] test(router): wire caching tests into CI-collected test_plugins.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dedicated tests/test_router_plugin_caching.py was not in the unit-tests job's fixed file list in .github/workflows/test.yml, so CI never executed it. Fold the three caching tests (single-load-on-repeat, bundle shape, concurrent single-load) into tests/test_plugins.py — which the unit-tests job already runs via `pytest tests/test_plugins.py` — and drop the standalone file, so the regression is actually guarded in CI. --- tests/test_plugins.py | 116 ++++++++++++++++++++++ tests/test_router_plugin_caching.py | 146 ---------------------------- 2 files changed, 116 insertions(+), 146 deletions(-) delete mode 100644 tests/test_router_plugin_caching.py diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 0f769d5a..e7f6b744 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") diff --git a/tests/test_router_plugin_caching.py b/tests/test_router_plugin_caching.py deleted file mode 100644 index b006f03d..00000000 --- a/tests/test_router_plugin_caching.py +++ /dev/null @@ -1,146 +0,0 @@ -#!/usr/bin/env python3 -""" -Tests for the router plugin's model caching. - -The router plugin classifies each incoming query to pick an optimization -approach. Loading its classifier is expensive and one-time (a ~400M-param -ModernBERT-large deserialize + a HuggingFace Hub round-trip + tokenizer build), -so it must be loaded once and reused across requests rather than reloaded on -every call. These tests pin that behavior down without downloading the real -model by mocking the underlying loaders. -""" - -import os -import sys -import threading - -# Try to import pytest, but don't fail if it's not available (matches the -# repo convention in tests/test_plugins.py so this runs under CI both ways). -try: - import pytest -except ImportError: - pytest = None - -from unittest import mock - -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -from optillm.plugins import router_plugin - - -class _FakeConfig: - hidden_size = 768 - - -def _install_loader_mocks(stack, from_pretrained_counter): - """Patch the heavy loaders so load_optillm_model() runs without any network - or real model weights. `from_pretrained_counter` counts base-model loads.""" - - def fake_from_pretrained(*args, **kwargs): - from_pretrained_counter.append(1) - base = mock.MagicMock() - base.config = _FakeConfig() - 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()) - ) - - -def _reset_cache(): - router_plugin._model_cache = None - - -def test_repeated_calls_load_model_once(): - """Calling load_optillm_model() N times must load the base model exactly once.""" - _reset_cache() - loads = [] - import contextlib - with contextlib.ExitStack() as stack: - _install_loader_mocks(stack, loads) - - first = router_plugin.load_optillm_model() - for _ in range(5): - again = router_plugin.load_optillm_model() - # Same cached bundle object returned every time. - assert again 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)" - ) - _reset_cache() - - -def test_returns_expected_bundle_shape(): - """The cached value is the (model, tokenizer, device) triple callers expect.""" - _reset_cache() - import contextlib - loads = [] - with contextlib.ExitStack() as stack: - _install_loader_mocks(stack, loads) - 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 - assert device is not None - _reset_cache() - - -def test_concurrent_calls_load_model_once(): - """Under concurrent first-time access the expensive load must run only once.""" - _reset_cache() - import contextlib - loads = [] - results = [] - with contextlib.ExitStack() as stack: - _install_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)" - ) - # Every thread observed the same cached bundle. - assert all(r is results[0] for r in results), "threads saw different cached bundles" - _reset_cache() - - -if __name__ == "__main__": - for name, fn in [ - ("repeated calls load model once", test_repeated_calls_load_model_once), - ("returns expected bundle shape", test_returns_expected_bundle_shape), - ("concurrent calls load model once", test_concurrent_calls_load_model_once), - ]: - try: - fn() - print(f"✅ {name}") - except Exception as e: - print(f"❌ {name}: {e}") - raise - - print("\nDone!")