fix: route insert-runner thread safety through a client capability - #849
Conversation
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: serhiizghama The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
|
/assign @jamesgao-jpg |
| 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) |
There was a problem hiding this comment.
[P1] Avoid deep-copying live Psycopg handles
Adbpg and VectorChord both declare thread_safe=False but do not override copy_for_thread. run_with_rate keeps the parent init context open while these workers run, so the default deepcopy(self) receives a live psycopg.Connection. I reproduced this with Psycopg 3.3.4: deepcopy(connection) raises TypeError: no default reduce due to non-trivial cinit. As a result, the first fixed-rate insert task for either backend fails before the copy reaches init().
Please give both clients a shallow-copy override that clears conn and cursor before worker init reopens them; PgVector uses the same handle pattern and is worth covering too. The fake regression client only contains copyable lists, so it does not exercise this boundary.
|
|
||
|
|
||
| def _runner(db: VectorDB): | ||
| return RatedMultiThreadingInsertRunner(rate=10, db=db, dataset_iter=None) |
There was a problem hiding this comment.
[P2] Construct this runner with a valid current batch size
Current main defaults the insert batch size to 100 and rejects rate % batch_size != 0. This helper passes rate=10 without batch_size, so in the current merge result every helper-using test raises ValueError before reaching send_insert_task.
Please use a valid rate and batch-size pair here. The green PR workflow does not catch this because make unittest currently runs only tests/test_dataset.py::TestDataSet::test_download_small, not this new test file.
Replace the db.name branching in RatedMultiThreadingInsertRunner with a VectorDB.copy_for_thread() hook driven by the existing thread_safe flag, so non-thread-safe clients get a per-thread copy without the runner hardcoding which databases those are.
The mysql.connector-backed clients (SeekDB, VolcMySQL, OceanBase) can't be deep-copied while holding an open socket, and Doris needs its client/table cleared; each now overrides copy_for_thread accordingly. OceanBase was previously missing from the runner's name list entirely, so its fixed-rate insert workers shared one connection.
PgVector, Adbpg and VectorChord all keep a live psycopg connection open on the parent while run_with_rate spins up insert workers. The default copy_for_thread() deep-copies the client, and psycopg connections have no __reduce__ for their non-trivial __cinit__ state, so the copy raises before the worker ever calls init(). Shallow-copy and drop conn/cursor instead, matching the pattern already used by Doris/OceanBase/SeekDB.
…idation Upstream's insert-batch-size work (zilliztech#813) landed after this branch and now rejects rate % batch_size != 0. The helper's implicit rate=10 with the new default batch_size=100 tripped that check before send_insert_task ever ran.
29b454c to
8544074
Compare
|
Both fixed. Gave PgVector, Adbpg and VectorChord their own shallow copy_for_thread (same pattern as Doris/OceanBase/SeekDB) so init() reconnects instead of deepcopy choking on the live psycopg handle. Also rebased on main to pick up the insert-batch-size work and fixed the test helper's rate/batch_size pair — you were right that it was raising before reaching send_insert_task. |
|
/lgtm |
3 participants
Fixes #810. The fixed-rate insert runner decided per-thread client handling by checking
db.nameagainst a hardcoded list (PgVector,Doris,SeekDB,VolcMySQL), so clients that declarethread_safe = Falsebut aren't on that list —OceanBase,VectorChord,Adbpg,LanceDB— fell through to the default path and shared one non-thread-safe connection across insert workers.OceanBasewas never in the list at all.I moved the decision onto the existing
thread_safecapability. The runner now branches ondb.thread_safeand asks the client for a per-thread copy via a newVectorDB.copy_for_thread(). The default deep-copies the instance, which already works for the psycopg-family clients and forLanceDB(it has a custom__deepcopy__); themysql.connector-backed clients override it to shallow-copy and drop their open socket soinit()reconnects inside the worker. Behavior for the databases that were already special-cased is unchanged — the capability just also covers the ones the name list missed, and adding a new non-thread-safe client no longer means remembering to edit the runner.Added a unit test that drives
send_insert_taskwith a fake client and checks that thread-safe clients insert through the shared object while non-thread-safe ones go through a copy plusinit().