fix(python): run AsyncTable.search embeddings on a dedicated executor (#3459)

## Summary
  
`AsyncTable.search()` computes the query embedding with
`loop.run_in_executor(None, ...)`, which uses asyncio's **default**
`ThreadPoolExecutor`. That pool is shared with all other
`run_in_executor(None, ...)` work, so a slow embedding call — a heavy
local model or an HTTP request to an embeddings API — ties up those
threads and starves unrelated async I/O under concurrent load.
  
This moves the (potentially blocking) embedding call onto a **dedicated
executor**, isolating it from the default pool.
  
  Closes #3310.
  
  ## Problem

  `python/lancedb/table.py`, `AsyncTable.search()`:

  ```python
  return (
      await loop.run_in_executor(
None, # asyncio's default executor, shared with other blocking I/O
          embedding.function.compute_query_embeddings_with_retry,
          query,
      )   
  )[0]
  ```
  
Under load, concurrent searches whose embeddings block (or any other
code using the default executor) contend for the same small thread pool.
  
  ## Change

- Add a dedicated
`ThreadPoolExecutor(thread_name_prefix="lancedb-embedding")` in
`background_loop.py`, exposed via `embedding_executor()`.
- Use it in `AsyncTable.search()`'s `make_embedding` instead of the
default executor.
- Reset the executor in the existing `_reset_after_fork` hook — its
worker threads don't survive `fork()`, same as the background event
loop. It's recreated lazily, so this is cheap.

  ## Design notes
  
The issue asked whether maintainers preferred a configurable executor, a
dedicated internal one, or another approach (no response in the thread).
I went with a **dedicated internal executor**: it fixes the starvation
with no public API change and stays consistent with the existing `LOOP`
singleton. Making the pool size configurable would be an easy follow-up
if preferred.
  
Scope is limited to `search()`. The broader "embedding functions need
real async support" (including `add()`) is tracked separately in #3268.
  
  ## Testing
  
- Added `test_async_search_runs_embedding_on_dedicated_executor`:
patches the embedding function to record the executing thread during an
async search and asserts it runs on a `lancedb-embedding` thread.
Verified it **fails** against the previous `run_in_executor(None, ...)`
and passes with the fix.
- `ruff format`, `ruff check`, and `pyright` pass on the changed files.
This commit is contained in:
Armaan Sandhu
2026-06-05 10:27:16 +05:30
committed by GitHub
parent c13ebc6796
commit 3868965413
3 changed files with 67 additions and 3 deletions

View File

@@ -4,6 +4,7 @@
import os
import sys
import threading
import warnings
from datetime import date, datetime, timedelta
from time import sleep
@@ -2837,3 +2838,38 @@ def test_sanitize_data_metadata_not_stripped():
assert result_schema.metadata is not None
assert result_schema.metadata[b"existing_key"] == b"existing_value"
assert result_schema.metadata[b"new_key"] == b"new_value"
@pytest.mark.asyncio
async def test_async_search_runs_embedding_on_dedicated_executor(
mem_db_async: AsyncConnection,
):
# Regression test for #3310: AsyncTable.search() must run the (potentially
# blocking) query-embedding call on the dedicated embedding executor, not
# asyncio's default executor -- which is shared with other blocking I/O and
# can be starved by a slow embedding call under concurrent load.
func = MockTextEmbeddingFunction.create()
class Schema(LanceModel):
text: str = func.SourceField()
vector: Vector(func.ndims()) = func.VectorField()
table = await mem_db_async.create_table("embed_executor", schema=Schema)
await table.add([{"text": "hello world"}])
captured_threads: List[str] = []
original = MockTextEmbeddingFunction.generate_embeddings
def record_thread(self, texts):
captured_threads.append(threading.current_thread().name)
return original(self, texts)
# Patch only around the search so we capture the query-embedding call, not
# the add-time source-embedding call.
with patch.object(MockTextEmbeddingFunction, "generate_embeddings", record_thread):
await (await table.search("a query string")).limit(1).to_list()
assert captured_threads, "search did not invoke the embedding function"
assert all(name.startswith("lancedb-embedding") for name in captured_threads), (
f"embedding ran off the dedicated executor: {captured_threads}"
)