perf: stop query_collection blocking the event loop (#27824)

RAG vector search runs in a thread pool, but then calls `future.result()` on the event loop thread, so the whole worker freezes until every collection answers. Every other user's token stream stops for that long. It's the default retrieval path.

Now `asyncio.gather` over `asyncio.to_thread`, matching what `routers/retrieval.py:2779` already does for the same call.

Measured with 3 queries across 4 collections, 60 ms search, and a second request wanting a turn every 5 ms:

| | before | after |
|---|---|---|
| RAG call | 62.0 ms | 61.2 ms |
| other request's turns | 0 | 7 |
| its worst stall | 62.5 ms | 16.0 ms |

Same results, same order, same `(result, error)` contract. Cancellation now lands mid-search instead of after every thread finishes. Threads move from an unbounded per-call pool to the loop's bounded shared one.
This commit is contained in:
Classic298 authored and GitHub committed 2026-07-31 17:24:31 -04:00
1 parent 810378c0b8
commit 466e05801b
1 file changed
+7 -8
+7 -8
View File
@@ -6,7 +6,6 @@ import logging
import os import os
import re import re
import time import time
from concurrent.futures import ThreadPoolExecutor
from typing import Awaitable, Optional, Union from typing import Awaitable, Optional, Union
from urllib.parse import quote from urllib.parse import quote
@@ -739,13 +738,13 @@ async def query_collection(
query_embeddings = await embedding_function(queries, prefix=RAG_EMBEDDING_QUERY_PREFIX) query_embeddings = await embedding_function(queries, prefix=RAG_EMBEDDING_QUERY_PREFIX)
log.debug(f'query_collection: processing {len(queries)} queries across {len(collection_names)} collections') log.debug(f'query_collection: processing {len(queries)} queries across {len(collection_names)} collections')
with ThreadPoolExecutor() as executor: task_results = await asyncio.gather(
future_results = [] *[
for query_embedding in query_embeddings: asyncio.to_thread(process_query_collection, collection_name, query_embedding)
for collection_name in collection_names: for query_embedding in query_embeddings
result = executor.submit(process_query_collection, collection_name, query_embedding) for collection_name in collection_names
future_results.append(result) ]
task_results = [future.result() for future in future_results] )
for result, err in task_results: for result, err in task_results:
if err is not None: if err is not None: