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
2026-07-31 23:24:31 +02:00
committed by GitHub
parent 810378c0b8
commit 466e05801b
+7 -8
View File
@@ -6,7 +6,6 @@ import logging
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Awaitable, Optional, Union
from urllib.parse import quote
@@ -739,13 +738,13 @@ async def query_collection(
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')
with ThreadPoolExecutor() as executor:
future_results = []
for query_embedding in query_embeddings:
for collection_name in collection_names:
result = executor.submit(process_query_collection, collection_name, query_embedding)
future_results.append(result)
task_results = [future.result() for future in future_results]
task_results = await asyncio.gather(
*[
asyncio.to_thread(process_query_collection, collection_name, query_embedding)
for query_embedding in query_embeddings
for collection_name in collection_names
]
)
for result, err in task_results:
if err is not None: