perf(mcp): lazy-load vector backend (warm search skips torch import)#441
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>
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.
Summary
jrag searchskips the heavytorchimport entirely, cutting first-response latency.Commits
ce8e294perf(mcp): lazy-load vector backend so warmjrag searchskips torch importb844cb6fix(mcp): sentinel + lock for lazy vector backend (review follow-up)d7cbcb0bump version to 0.10.2Note
PyPI was published from this branch ahead of merge — merging brings
masterback in line with the released0.10.2.🤖 Generated with Claude Code