fix(async): preserve optimization after result wait timeout - #3783
fix(async): preserve optimization after result wait timeout#3783Kuang-xianxin wants to merge 2 commits into
Conversation
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: Kuang-xianxin 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 |
|
Welcome @Kuang-xianxin! It looks like this is your first PR to milvus-io/pymilvus 🎉 |
Assisted-by: Codex Signed-off-by: Kuang-xianxin <243476082+Kuang-xianxin@users.noreply.github.com>
df639b2 to
d22d7af
Compare
|
Tick the box to add this pull request to the merge queue (same as
|
| try: | ||
| if timeout is not None: | ||
| return await asyncio.wait_for(self._task, timeout=timeout) | ||
| return await asyncio.wait_for(asyncio.shield(self._task), timeout=timeout) |
There was a problem hiding this comment.
pymilvus/milvus_client/async_optimize_task.py line:93
Medium ---- With wait=True the caller no longer has a handle after a wait timeout: optimize(wait=True, timeout=X) raises "Timeout waiting..." while the detached background task keeps running, cannot be cancelled or re-awaited, and typically ends with its own remaining_timeout()=0 exception that nobody retrieves ("Task exception was never retrieved" at GC). Was losing the old abort-on-timeout behavior for wait=True intended? Consider cancelling the task on that path or keeping the task reachable/documenting the detach.
There was a problem hiding this comment.
Addressed in 258e5d1. optimize(wait=True) now cancels and awaits unfinished owned work in finally, retrieving the task's cancellation exception while preserving the original timeout or caller cancellation. This keeps the no-handle path from detaching the optimization. The behavior is documented on wait.
New regressions cover zero/positive wait timeouts and caller cancellation with/without a result timeout; they verify that the owned task is done and cancelled before optimize exits. Together with the waiter-isolation cases, 9 new cases failed on d22d7af; all 148 related tests pass after the fix. Ruff and Black pass on all four changed files. No live Milvus server was used.
| raise MilvusException(message="Optimization task was cancelled") | ||
|
|
||
| async def result(self, timeout: Optional[float] = None) -> OptimizeResult: | ||
| """Wait for the result without cancelling optimization when the wait times out.""" |
There was a problem hiding this comment.
pymilvus/milvus_client/async_optimize_task.py line:87
Low ---- Follow-up: only the timed branch shields the background task. The timeout=None branch (line 94) still awaits self._task directly, so a caller cancellation - e.g. await asyncio.timeout(X): await task.result() or asyncio.wait_for(task.result(), X) - still cancels the optimization, and result() converts the caller CancelledError into a MilvusException. The exact failure this PR fixes is therefore still reachable through the standard asyncio timeout wrappers; consider shielding both branches.
There was a problem hiding this comment.
Addressed in 258e5d1. Both branches now shield the background task. result() propagates cancellation of its caller as CancelledError, so asyncio.wait_for(task.result(), ...) and asyncio.timeout(...) produce their normal timeout result while the optimization remains available to other waiters. Explicit cancellation of the optimization still reports the existing Milvus cancellation error.
The new regressions cover direct caller cancellation and wait_for around both timed/untimed result(), plus the Python 3.11+ timeout context. They verify subsequent successful retrieval and isolation of another waiter. These cases failed on the previous head and pass in the 148-test related suite.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #3783 +/- ##
==========================================
+ Coverage 94.22% 94.24% +0.01%
==========================================
Files 77 77
Lines 16204 16204
==========================================
+ Hits 15268 15271 +3
+ Misses 936 933 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Preserve background tasks when result waiters are cancelled, propagate caller cancellation, and cancel and retrieve the owned task when optimize(wait=True) exits early. Add nine regressions for external timeout wrappers and owned-task cleanup. Assisted-by: Codex Signed-off-by: Kuang-xianxin <243476082+Kuang-xianxin@users.noreply.github.com>
3 participants
Problem
AsyncOptimizeTask.result(timeout=...)passes its owned background task directly toasyncio.wait_for. A short wait cancels the optimization, affects other result waiters, and prevents a later retry. For the previewoptimize(wait=False)API this can stop the client workflow during compaction before index rebuild and load refresh.Cancellation also needs to distinguish ownership: a caller holding an
AsyncOptimizeTaskcan wait again or explicitly cancel it, whileoptimize(wait=True)does not return a handle when it exits early.Change
result()waits. A local wait timeout raises the existing timeout error; caller cancellation propagates asCancelledError, allowing standardasyncio.wait_forandasyncio.timeoutwrappers to work without cancelling the shared optimization.task.cancel()behavior and the optimization's own timeout forwarding.optimize(wait=True), cancel and await unfinished owned work before returning a wait timeout or propagating caller cancellation. This prevents detaching an inaccessible task and retrieves its cancellation exception.Regression coverage includes later success/failure after a short wait, concurrent waiters, explicit task cancellation, external timeout wrappers, and zero/positive timeouts or cancellation of
optimize(wait=True). A publicoptimize(wait=False)workflow test uses mocked RPC responses and verifies load refresh after a timed-out result wait.Validation
Python 3.13.14 on Windows:
asyncio.timeoutcontext test is skipped on Python versions before 3.11;wait_forcases cover older versions.Found through source inspection. Searches did not find a matching optimization/result-timeout fix; the existing async refresh-load PRs concern separate RPC options.