diff --git a/tests/test_rate_runner_thread_safety.py b/tests/test_rate_runner_thread_safety.py new file mode 100644 index 000000000..906e91020 --- /dev/null +++ b/tests/test_rate_runner_thread_safety.py @@ -0,0 +1,85 @@ +"""Unit tests for RatedMultiThreadingInsertRunner.send_insert_task routing. + +These exercise the thread-safety routing without a live database: a thread-safe +client inserts through the shared object, while a non-thread-safe client is routed +through copy_for_thread() + init() so each worker owns its own connection. +""" + +from contextlib import contextmanager +from copy import deepcopy + +from vectordb_bench.backend.clients.api import VectorDB +from vectordb_bench.backend.runner.rate_runner import RatedMultiThreadingInsertRunner + + +class FakeDB(VectorDB): + def __init__(self, thread_safe: bool = True, name: str = "Fake"): + self.thread_safe = thread_safe + self.name = name + self.init_calls = 0 + self.inserted: list[int] = [] + + @contextmanager + def init(self): + self.init_calls += 1 + yield + + def insert_embeddings(self, embeddings: list, metadata: list, **kwargs): + self.inserted.append(len(embeddings)) + return len(embeddings), None + + def search_embedding(self, *args, **kwargs): + return [] + + def optimize(self, data_size: int | None = None): + return + + +class CapturingNonThreadSafeDB(FakeDB): + """Records the per-thread copy so the test can assert against it.""" + + def __init__(self): + super().__init__(thread_safe=False, name="Capturing") + self.thread_copy: CapturingNonThreadSafeDB | None = None + + def copy_for_thread(self) -> "VectorDB": + c = deepcopy(self) + self.thread_copy = c + return c + + +def _runner(db: VectorDB): + return RatedMultiThreadingInsertRunner(rate=10, db=db, dataset_iter=None) + + +def test_thread_safe_client_uses_shared_object(): + db = FakeDB(thread_safe=True) + _runner(db).send_insert_task(db, [[0.1], [0.2]], ["a", "b"]) + + # Shared client inserts directly, no per-thread copy/init from the runner. + assert db.inserted == [2] + assert db.init_calls == 0 + + +def test_non_thread_safe_client_routes_through_copy(): + db = CapturingNonThreadSafeDB() + _runner(db).send_insert_task(db, [[0.1], [0.2]], ["a", "b"]) + + copy = db.thread_copy + assert copy is not None + assert copy is not db + # Insert + init happen on the thread-local copy, never the parent. + assert copy.init_calls == 1 + assert copy.inserted == [2] + assert db.inserted == [] + assert db.init_calls == 0 + + +def test_default_copy_for_thread_is_a_distinct_deep_copy(): + db = FakeDB(thread_safe=False) + copy = db.copy_for_thread() + + assert copy is not db + assert isinstance(copy, FakeDB) + copy.inserted.append(1) + assert db.inserted == [] diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index bac90b7be..498fd76c0 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod from contextlib import contextmanager +from copy import deepcopy from enum import StrEnum from typing import ClassVar @@ -199,6 +200,17 @@ def prepare_filter(self, filters: Filter): (All search tests in a case use consistent filtering conditions.)""" return + def copy_for_thread(self) -> "VectorDB": + """Return a per-thread copy of this client for non-thread-safe backends. + + Runners call this (instead of branching on db.name) when thread_safe is + False, then init() the copy inside the worker so each thread owns its own + connection. Defaults to a deep copy; clients whose live connection can't be + deep-copied (e.g. an open DB-API socket) override this to shallow-copy and + drop their connection handles so init() re-establishes them per thread. + """ + return deepcopy(self) + @abstractmethod def __init__( self, diff --git a/vectordb_bench/backend/clients/doris/doris.py b/vectordb_bench/backend/clients/doris/doris.py index 01984d665..04d924ee7 100644 --- a/vectordb_bench/backend/clients/doris/doris.py +++ b/vectordb_bench/backend/clients/doris/doris.py @@ -1,6 +1,7 @@ import logging import os from contextlib import contextmanager +from copy import deepcopy from typing import Any import pandas as pd @@ -113,6 +114,14 @@ def _ensure_client_initialized(self): # Table might not exist yet; leave it to ready_to_load self.table = None + def copy_for_thread(self) -> "VectorDB": + # DorisVectorClient isn't thread-safe; hand each worker its own copy and force + # a fresh client/table so init() rebuilds them instead of sharing the parent's. + db_copy = deepcopy(self) + db_copy.client = None + db_copy.table = None + return db_copy + @contextmanager def init(self): try: diff --git a/vectordb_bench/backend/clients/oceanbase/oceanbase.py b/vectordb_bench/backend/clients/oceanbase/oceanbase.py index a34161037..b109d558d 100644 --- a/vectordb_bench/backend/clients/oceanbase/oceanbase.py +++ b/vectordb_bench/backend/clients/oceanbase/oceanbase.py @@ -3,6 +3,7 @@ import time from collections.abc import Generator from contextlib import contextmanager +from copy import copy from typing import Any import mysql.connector as mysql @@ -82,6 +83,14 @@ def _disconnect(self): self._conn.close() self._conn = None + def copy_for_thread(self) -> "VectorDB": + # mysql.connector holds an open socket that can't be deep-copied; shallow-copy + # and drop the connection so init() reconnects inside the worker thread. + db_copy = copy(self) + db_copy._conn = None + db_copy._cursor = None + return db_copy + @contextmanager def init(self) -> Generator[None, None, None]: try: diff --git a/vectordb_bench/backend/clients/seekdb/seekdb.py b/vectordb_bench/backend/clients/seekdb/seekdb.py index f4551a577..69a27bd4b 100644 --- a/vectordb_bench/backend/clients/seekdb/seekdb.py +++ b/vectordb_bench/backend/clients/seekdb/seekdb.py @@ -3,6 +3,7 @@ import struct from collections.abc import Generator from contextlib import contextmanager +from copy import copy from typing import Any import mysql.connector as mysql @@ -125,6 +126,14 @@ def _init_session_settings(self): # SeekDB uses OceanBase-style session vars (not plain hnsw_ef_search). self._cursor.execute(f"SET ob_hnsw_ef_search={ef_search}") + def copy_for_thread(self) -> "VectorDB": + # mysql.connector holds an open socket that can't be deep-copied; shallow-copy + # and drop the connection so init() reconnects inside the worker thread. + db_copy = copy(self) + db_copy._conn = None + db_copy._cursor = None + return db_copy + @contextmanager def init(self) -> Generator[None, None, None]: try: diff --git a/vectordb_bench/backend/clients/volc_mysql/volc_mysql.py b/vectordb_bench/backend/clients/volc_mysql/volc_mysql.py index 455286019..e01a3e3e2 100755 --- a/vectordb_bench/backend/clients/volc_mysql/volc_mysql.py +++ b/vectordb_bench/backend/clients/volc_mysql/volc_mysql.py @@ -5,6 +5,7 @@ import struct import tempfile from contextlib import contextmanager +from copy import copy from pathlib import Path import mysql.connector as mysql @@ -206,6 +207,15 @@ def _probe_binary_support(self) -> bool: log.debug("Failed to drop binary-probe temp table", exc_info=True) cur.close() + def copy_for_thread(self) -> "VectorDB": + # mysql.connector holds an open socket that can't be deep-copied; shallow-copy + # and drop the connection so init() reconnects inside the worker thread. + db_copy = copy(self) + db_copy.conn = None + db_copy.cursor = None + db_copy.admin_cursor = None + return db_copy + @contextmanager def init(self): """create and destory connections to database. diff --git a/vectordb_bench/backend/runner/rate_runner.py b/vectordb_bench/backend/runner/rate_runner.py index 91d0bb3ee..bee6b4a1b 100644 --- a/vectordb_bench/backend/runner/rate_runner.py +++ b/vectordb_bench/backend/runner/rate_runner.py @@ -3,7 +3,6 @@ import multiprocessing as mp import time from concurrent.futures import ThreadPoolExecutor -from copy import copy, deepcopy from vectordb_bench import config from vectordb_bench.backend.clients import api @@ -47,49 +46,15 @@ def _insert_embeddings(db: api.VectorDB, emb: list[list[float]], metadata: list[ msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times" raise RuntimeError(msg) from None - if db.name == "PgVector": - # pgvector is not thread-safe for concurrent insert, - # so we need to copy the db object, make sure each thread has its own connection - db_copy = deepcopy(db) - with db_copy.init(): - _insert_embeddings(db_copy, emb, metadata, retry_idx=0) - elif db.name == "Doris": - # DorisVectorClient is not thread-safe. Similar to pgvector, create a per-thread client - # by deep-copying the wrapper and forcing lazy re-init inside the thread. - db_copy = deepcopy(db) - # Ensure a fresh client/table will be created in this thread - try: - db_copy.client = None - db_copy.table = None - except Exception: - log.debug("Failed to reset Doris client or table on thread-local copy", exc_info=True) - with db_copy.init(): - _insert_embeddings(db_copy, emb, metadata, retry_idx=0) - elif db.name == "SeekDB": - # mysql.connector is not thread-safe; do not share one connection across workers. - # deepcopy() fails on an open _conn (socket is not picklable / not copy-safe in spawn workers). - db_copy = copy(db) - try: - db_copy._conn = None - db_copy._cursor = None - except Exception: - log.debug("Failed to reset SeekDB connection on thread-local copy", exc_info=True) - with db_copy.init(): - _insert_embeddings(db_copy, emb, metadata, retry_idx=0) - elif db.name == "VolcMySQL": - # mysql.connector is not thread-safe; do not share one connection across workers. - # deepcopy() fails on an open conn (socket is not picklable / not copy-safe in spawn workers). - db_copy = copy(db) - try: - db_copy.conn = None - db_copy.cursor = None - db_copy.admin_cursor = None - except Exception: - log.debug("Failed to reset VolcMySQL connection on thread-local copy", exc_info=True) + if db.thread_safe: + _insert_embeddings(db, emb, metadata, retry_idx=0) + else: + # Non-thread-safe clients can't share one connection across insert workers. + # Each client owns how to make a thread-local copy (see VectorDB.copy_for_thread); + # init() then re-establishes its connection inside this thread. + db_copy = db.copy_for_thread() with db_copy.init(): _insert_embeddings(db_copy, emb, metadata, retry_idx=0) - else: - _insert_embeddings(db, emb, metadata, retry_idx=0) @time_it def run_with_rate(self, q: mp.Queue):