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

[SuperMarioYL]

Problem

router_plugin.run() calls load_optillm_model() on every request (router_plugin.py):

def run(system_prompt, initial_query, client, model, **kwargs):
    try:
        router_model, tokenizer, device = load_optillm_model()   # <-- every call

load_optillm_model() is not cheap. Each call:

• AutoModel.from_pretrained("answerdotai/ModernBERT-large") — deserializes a ~400M-param base model,
• hf_hub_download("codelion/optillm-modernbert-large", "model.safetensors") — a HuggingFace Hub round-trip (ETag check + potential download),
• load_model(...) — loads the fine-tuned state dict,
• AutoTokenizer.from_pretrained(...) — rebuilds the tokenizer.

The actual routing decision is a single forward pass (tens of ms), but it is preceded by a full transformer load on each request. So a step that should be O(1 forward pass) is O(load a transformer) every time the router approach is used.

Fix

Memoize the loaded (model, tokenizer, device) bundle behind a module-level singleton guarded by a threading.Lock (double-checked locking):

_model_cache = None
_model_cache_lock = threading.Lock()

def load_optillm_model():
    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

The heavy loading logic is unchanged — it just moves into _load_optillm_model() and now runs once.

Why this is safe

• Stateless at inference. predict_approach() runs the model under model.eval() + torch.no_grad() and never mutates it, so reusing one instance across requests is correctness-safe.
• Thread-safe first load. The lock serializes the first concurrent load so the expensive work happens exactly once under the threaded Flask server; steady-state calls take the lock-free fast path (the same idiom as CacheManager in optillm/inference.py, which caches models/tokenizers with a double-checked threading.Lock).
• No poisoned cache. The global is only assigned on a successful load, so if loading raises (e.g. a Hub/network error) the exception still propagates to run()'s existing except fallback and the next request retries the load.

The one intentional trade-off is that the model now stays resident for the process lifetime instead of being re-created and garbage-collected each request — which is the desired behavior for a long-running inference proxy that routes repeatedly.

Testing

Adds tests/test_router_plugin_caching.py (mocks the heavy loaders, so no network/weights required — CI-friendly):

• test_repeated_calls_load_model_once — 6 calls to load_optillm_model() load the base model exactly once and return the same cached bundle.
• test_returns_expected_bundle_shape — the cached value is the (OptILMClassifier, tokenizer, device) triple callers expect.
• test_concurrent_calls_load_model_once — 8 threads hitting the first load concurrently still load exactly once (double-checked lock holds).

The repeated-call test is red on main (base model loads 6×) and green with this change. Runs both under pytest and standalone, matching the repo's existing test convention.