perf(mcp): lazy-load vector backend — warm jrag search 2.6s → 0.15s#440
Merged
Conversation
… import `mcp_v2` eagerly imported `search_lancedb` at module load, which imports `sentence_transformers`/`torch` at its own top level (~2.8s). The warm `jrag watch` client reconstructs `SearchOutput` from `mcp_v2` (`watch/client._reconstruct`) without ever calling the vector backend, so every warm CLI read paid that ~2.8s import — masking the daemon's win and leaving warm `jrag search` at ~2.6s instead of near-instant. Convert the eager import to lazy: `run_search`/`TABLES` start as sentinels and are populated by `_ensure_vector_backend()` on the first `search_v2` call. The guard preserves existing semantics: - tests that monkeypatch `mcp_v2.run_search` (non-None) are not overwritten; - graph-only installs (ImportError) leave `run_search=None` -> lexical fallback. Measured on tests/bank-chat-system (130 Java files): - `SearchOutput` import: 2.85s -> 0.16s - warm `jrag search`: 2.6s -> 0.15s steady (17x); cold path unchanged - results byte-identical warm vs cold Tests: run_search monkeypatchers (185), lightweight watch (101), heavy warm-path (6), all 6 golden IPC byte-identity (warm == cold) — all green. Co-Authored-By: Claude <noreply@anthropic.com>
Address two code-review findings on the lazy-load change:
1. Critical — sentinel conflation. The guard `run_search is not None` could
not distinguish "unloaded" (`None` sentinel) from a test that forces
lexical mode via `monkeypatch.setattr(mcp_v2, "run_search", None)`. The
loader ran anyway, overwrote the patch with the real backend, and
`lexical_mode` flipped to False — regressing
`test_search_v2_lexical_dispatch_end_to_end` (Lance "Table not found").
Fix: a distinct `_NOT_LOADED = object()` sentinel for the unloaded state,
so `None` unambiguously means "lexical/disabled" and is authoritative in
both monkeypatch directions (None → lexical; non-None → vector).
2. Important — thread race. The MCP server dispatches `search_v2` via
`asyncio.to_thread` (server.py:659), so two concurrent first-batch
searches could race the unguarded sentinel mutation: Thread B sees
`_VECTOR_BACKEND_LOADED=True` but `run_search` still `None` and wrongly
takes the lexical path / reads `TABLES={}`. Fix: double-checked locking
with a `threading.Lock`, mirroring the existing `_st_lock` pattern.
The sentinel also lets us drop `_VECTOR_BACKEND_LOADED` entirely (the
sentinel itself tracks state), and a non-ImportError now propagates and
leaves `run_search` as `_NOT_LOADED` so the next call retries instead of
poisoning the process.
Verified: previously-failing lexical test passes; 206 lexical + run_search-
monkeypatcher tests pass; 12 heavy warm + golden IPC byte-identity tests
(warm == cold) pass; SearchOutput import stays ~0.14s.
Co-Authored-By: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
After #424 (
jrag watch), warmjrag searchwas still ~2.6s — not the near-instant latency the watch daemon was supposed to deliver. The daemon serves the query in ~12ms, but the client still paid ~2.8s on every invocation.Root cause
mcp_v2did an eager top-level import ofsearch_lancedb:search_lancedbimportssentence_transformers(+transformers/sklearn/torch) at its own top level — ~2.8s. The warmjrag watchclient reconstructsSearchOutputfrommcp_v2(watch/client._reconstruct) to keep rendering byte-identical to the cold path, without ever calling the vector backend. So every warm CLI read dragged in the entire ML stack just to rebuild one pydantic model.The daemon eliminated the model instantiation (~4.4s) but not the import (~2.8s) — that fixed floor is why the win didn't show.
Fix
Convert the eager import to lazy.
run_search/TABLESstart as sentinels and are populated by_ensure_vector_backend()on the firstsearch_v2call. The guard preserves all existing semantics:mcp_v2.run_search(set it to a non-Nonelambda) are not overwritten —_ensure_vector_backendtreats a non-Nonevalue as authoritative.run_searchstaysNone→search_v2takes the existing lexical fallback path.Measurement (
tests/bank-chat-system/, 130 Java files)SearchOutputimportjrag search(steady)jrag searchFirst warm run is ~1.1s (one-time daemon Lance connect); steady state is ~0.15s.
Tests
run_searchmonkeypatchers (test_mcp_v2,test_mcp_v2_compose,test_generated_surface,test_absence_mcp_integration): 185 passed, 1 skippedtests/watch: 101 passed, 8 skippedtest_warm+test_warm_model_reuse): 6 passedOut of scope
The residual client-side LadybugGraph open (
_load_graph_or_errorbeforeget_payloaddispatch, ~72ms) is the spec's documented §10.2 follow-up — minor next to the 2.8s fixed here; left for a separate change.🤖 Generated with Claude Code