Repository navigation
fix(vector_stores): finish in-flight expired file batch cleanup on provider shutdown - #6226
Conversation
asyncio.create_task(self._cleanup_expired_file_batches()) discarded its result, so the event loop's weak reference allowed the cleanup task to be garbage- collected before completion, silently dropping the cleanup mid-run. Store it in a dedicated _background_tasks set with a self-cleaning done-callback, and cancel those tasks on shutdown() alongside the existing file-batch tasks.
|
@kratos0718 how can we reproduce this? |
|
Fair question, and the honest answer is that the GC angle in my description is the RUF006 rationale rather than something I can trigger on demand. I tried and couldn't make CPython collect it reliably. What is deterministic is the shutdown path, which I think is the real bug. async def _cleanup_expired_file_batches(self):
for b in ("expired-1", "expired-2", "expired-3"):
await asyncio.sleep(0.05) # stands in for kvstore.delete
deleted.append(b)
async def create_file_batch(self, b):
task = asyncio.create_task(self._process_file_batch_async(b))
self._file_batch_tasks[b] = task # tracked
asyncio.create_task(self._cleanup_expired_file_batches()) # not tracked
await asyncio.sleep(0.06)
async def shutdown(self):
for _, t in list(self._file_batch_tasks.items()):
if not t.done():
t.cancel()
try: await t
except asyncio.CancelledError: passSo one expired batch gets deleted and the other two stay in the store, with no exception and no log line. The One thing worth deciding: my PR tracks the task and cancels it in |
|
@mattf gentle ping on this one. The reproduction for the shutdown path is in my comment above. If you tell me whether you'd rather shutdown() cancel the cleanup task or await it with a timeout, I'll push that as a follow-up commit. |
There was a problem hiding this comment.
Community review. This does not clear the merge gate.
I read the diff in openai_vector_store_mixin.py and ran it against the PR branch. I think it is correct.
What I ran (PR branch, C:/ox, Python venv, mocked adapter built on OpenAIVectorStoreMixin):
- Called
openai_create_vector_store_file_batchwithcleanup_interval_seconds = 0and a cleanup stub that waits on an event. While it waits,_background_taskshas 1 entry. After the stub finishes, the set is empty, so the done callback discards it. - Started a second cleanup that never finishes, then called
shutdown(). The task ends up cancelled and the set is empty. - On origin/main the same script fails with
AttributeError: 'M' object has no attribute '_background_tasks', as expected, since nothing tracks the cleanup task there. tests/unit/providers/utils/memory/test_openai_vector_store_mixin.pypasses on the PR branch (16 passed).
I also checked that every adapter that subclasses the mixin calls super().__init__(...) (faiss, sqlite_vec, chroma, elasticsearch, infinispan, milvus, neo4j, oci26ai, pgvector, qdrant, weaviate), so _background_tasks always exists. The test mocks in test_openai_vector_store_mixin.py and test_vector_store_config_registration.py do the same.
What I did not do: I did not reproduce the task being garbage collected on main. That is the documented asyncio behavior (the loop holds only a weak reference), but I did not trigger it, so I cannot say how likely it is with the real cleanup coroutine.
Notes, none blocking:
- The PR adds no test. A small test like the probe above (assert the set is populated, then emptied, and that
shutdown()cancels a pending cleanup) would lock in the shutdown change. - The PR is marked BEHIND main. I applied its diff to current origin/main with
git apply --checkand it applies cleanly.
mattf
changed the title
bdac1a3 to
d5d12f1
Compare
…ovider shutdown The cleanup task was cancelled in shutdown() alongside the file batch tasks, which still stopped the cleanup mid-delete: the observable state (expired batches left behind) was identical to leaving the task untracked. Awaiting the in-flight cleanup runs instead makes shutdown a barrier for the work its docstring says it covers. A cancelled cleanup task is treated as finished, while shutdown's own cancellation (e.g. the stack's 5-second provider shutdown deadline) propagates. Adds a deterministic reproducer driving the real mixin with a fake KV store whose deletes have a fixed latency. Signed-off-by: Matthew Farrellee <[email protected]>
What
openai_create_vector_store_file_batchschedules the throttled expired-batch cleanup with a bareasyncio.create_task(...)whose result is discarded.shutdown()walks_file_batch_tasksonly, so an in-flight cleanup run is invisible to it:await mixin.shutdown()returns while the cleanup is still running — shutdown is not a barrier for the work its docstring says it covers ("Clean up mixin resources including background tasks").deletewith no exception and no log line, leaving expired batches half-cleaned: the survivors sit in the store (and in theopenai_file_batchesregistry) until the next run, which is throttled to once a day. On a non-exit re-initialization the same task survives as a zombie mutating reloaded state.Why await, not cancel: cancelling the task in
shutdown()still stops the cleanup early — the observable state (expired batches left behind) is identical to the bug. The new tests fail against a track-and-cancel fix and pass against track-and-await.Fix
Track the cleanup task in a self-cleaning
self._background_tasksset and await those tasks at the end ofshutdown(), alongside the existing file-batch task handling. The await is bounded by the stack's existing 5-second provider shutdown deadline (asyncio.wait_for(impl.shutdown(), timeout=5)instack.py); a cancelled cleanup task is treated as finished, while shutdown's own cancellation (the deadline firing) propagates.Test plan
New:
tests/unit/providers/utils/memory/test_openai_vector_store_mixin_cleanup_task.py— drives the real mixin (entry point → untracked spawn → realshutdown()) with a deterministic fake KV store whose deletes have a fixed 50 ms latency:test_shutdown_is_a_barrier_for_cleanup—shutdown()must not return while the in-flight cleanup is unfinished.test_process_exit_does_not_leave_cleanup_half_done— a nestedasyncio.run(loop teardown = server exit) must not leave expired batches half-cleaned.Notes