mirror of
https://github.com/open-webui/open-webui.git
synced 2026-08-25 15:04:49 -06:00
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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user