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:
1 parent
810378c0b8
commit
466e05801b
1 file changed
+7
-8
@@ -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:
|
||||||
|
|||||||
Reference in new issue
Block a user