mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-27 01:40:33 -04:00
fix: pgvector reads leak their connection and lose most of their neighbours (#30142)
Two defects on the pgvector read path. A search, query or get that finds nothing returns before the rollback that ends its read-only transaction, so the session keeps the connection it checked out; retrieval fans out over worker threads and the session is thread-local, so every thread that runs an empty read holds a connection for the lifetime of that thread. Users see vector search die after a while with "QueuePool limit of size 5 overflow 10 reached" and no way back other than a restart. Rolling back before the three early returns puts the connection back: measured on PostgreSQL 16, twelve empty reads on a default-shaped pool left all 12 connections checked out before and 0 after. All collections also share one table under one vector index, and the WHERE collection_name filter is applied after the index walk, so a knowledge base holding a small share of the rows keeps only a small share of its neighbours, silently. pgvector 0.8 added iterative scans for this: the scan keeps going until enough rows pass the filter. This sets it per search so it cannot leak into other sessions, and only when the installed extension supports it, since setting it on pgvector 0.7 would make every search raise. PGVECTOR_ITERATIVE_SCAN turns it off or picks strict_order; it defaults to relaxed_order because the shipped behaviour is silently wrong results, and the cost is about a millisecond per search. Measured through PgvectorClient.search on PostgreSQL 17 and pgvector 0.8, 101500 rows of 384 dimensions with the knowledge base at 1.5% of the table, hnsw m=16, recall@10 against an exact scan over 40 queries: 0.070 at 4.9 ms before, 0.970 at 7.2 ms after. Fixes #30133 Fixes #30135
This commit is contained in:
@@ -741,6 +741,10 @@ else:
|
||||
except Exception:
|
||||
PGVECTOR_IVFFLAT_LISTS = 100
|
||||
|
||||
PGVECTOR_ITERATIVE_SCAN = os.getenv('PGVECTOR_ITERATIVE_SCAN', 'relaxed_order').strip().lower()
|
||||
if PGVECTOR_ITERATIVE_SCAN not in ('off', 'relaxed_order', 'strict_order'):
|
||||
PGVECTOR_ITERATIVE_SCAN = 'relaxed_order'
|
||||
|
||||
# openGauss
|
||||
OPENGAUSS_DB_URL = os.getenv('OPENGAUSS_DB_URL', DATABASE_URL)
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ from open_webui.config import (
|
||||
PGVECTOR_HNSW_M,
|
||||
PGVECTOR_INDEX_METHOD,
|
||||
PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH,
|
||||
PGVECTOR_ITERATIVE_SCAN,
|
||||
PGVECTOR_IVFFLAT_LISTS,
|
||||
PGVECTOR_PGCRYPTO,
|
||||
PGVECTOR_PGCRYPTO_KEY,
|
||||
@@ -154,6 +155,7 @@ class PgvectorClient(VectorDBBase):
|
||||
index_method, index_options = self._vector_index_configuration()
|
||||
self._ensure_vector_index(index_method, index_options)
|
||||
self._ensure_text_search_index()
|
||||
self.iterative_scan_sql = self._iterative_scan_setting(index_method)
|
||||
|
||||
self.session.execute(
|
||||
text(
|
||||
@@ -258,6 +260,20 @@ class PgvectorClient(VectorDBBase):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _iterative_scan_setting(self, index_method: str) -> Optional[str]:
|
||||
if PGVECTOR_ITERATIVE_SCAN == 'off':
|
||||
return None
|
||||
|
||||
version = self.session.execute(text("SELECT extversion FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
version_parts = [int(part) for part in (version or '').split('.') if part.isdigit()]
|
||||
if version_parts[:2] < [0, 8]:
|
||||
log.info('Iterative scan needs pgvector 0.8 or newer, the server has %s.', version or 'none')
|
||||
return None
|
||||
|
||||
# ivfflat only accepts relaxed_order
|
||||
mode = 'relaxed_order' if index_method == 'ivfflat' else PGVECTOR_ITERATIVE_SCAN
|
||||
return f'SET LOCAL {index_method}.iterative_scan = {mode}'
|
||||
|
||||
def _ensure_text_search_index(self) -> None:
|
||||
if PGVECTOR_PGCRYPTO:
|
||||
return
|
||||
@@ -524,6 +540,9 @@ class PgvectorClient(VectorDBBase):
|
||||
.order_by(query_vectors.c.qid, subq.c.distance)
|
||||
)
|
||||
|
||||
if self.iterative_scan_sql:
|
||||
self.session.execute(text(self.iterative_scan_sql))
|
||||
|
||||
result_proxy = self.session.execute(stmt)
|
||||
results = result_proxy.all()
|
||||
|
||||
@@ -533,6 +552,7 @@ class PgvectorClient(VectorDBBase):
|
||||
metadatas = [[] for _ in range(num_queries)]
|
||||
|
||||
if not results:
|
||||
self.session.rollback()
|
||||
return SearchResult(
|
||||
ids=ids,
|
||||
distances=distances,
|
||||
@@ -652,6 +672,7 @@ class PgvectorClient(VectorDBBase):
|
||||
results = query.all()
|
||||
|
||||
if not results:
|
||||
self.session.rollback()
|
||||
return None
|
||||
|
||||
ids = [[result.id for result in results]]
|
||||
@@ -691,6 +712,7 @@ class PgvectorClient(VectorDBBase):
|
||||
results = query.all()
|
||||
|
||||
if not results:
|
||||
self.session.rollback()
|
||||
return None
|
||||
|
||||
ids = [[result.id for result in results]]
|
||||
|
||||
Reference in New Issue
Block a user