diff --git a/pymilvus/milvus_client/async_milvus_client.py b/pymilvus/milvus_client/async_milvus_client.py index b748d0c62..d532a708a 100644 --- a/pymilvus/milvus_client/async_milvus_client.py +++ b/pymilvus/milvus_client/async_milvus_client.py @@ -2,6 +2,7 @@ import copy import time import types +from contextlib import suppress from typing import Dict, List, Optional, Type, Union from pymilvus.client import type_info @@ -2310,6 +2311,7 @@ async def optimize( If not provided, uses system default. wait (bool): Whether to wait for optimization to complete. Defaults to True. If False, returns an OptimizeTask for async tracking. + If True, a wait timeout or caller cancellation stops the client-side task. timeout (Optional[float]): Maximum time in seconds to wait for optimization. Only applies when wait=True. **kwargs: Additional arguments. @@ -2356,7 +2358,14 @@ async def optimize( task.start() if wait: - return await task.result(timeout=timeout) + try: + return await task.result(timeout=timeout) + finally: + # No task handle is returned on this path; stop and retrieve unfinished work. + if not task.done(): + task.cancel() + with suppress(MilvusException): + await task.result() return task diff --git a/pymilvus/milvus_client/async_optimize_task.py b/pymilvus/milvus_client/async_optimize_task.py index e429694ac..f1b515264 100644 --- a/pymilvus/milvus_client/async_optimize_task.py +++ b/pymilvus/milvus_client/async_optimize_task.py @@ -84,17 +84,23 @@ def check_cancelled(self) -> None: raise MilvusException(message="Optimization task was cancelled") async def result(self, timeout: Optional[float] = None) -> OptimizeResult: + """Wait without cancelling optimization on wait timeout or caller cancellation. + + Call cancel() to stop the background optimization explicitly. + """ if not self._task: raise MilvusException(message="Task has not been started") try: if timeout is not None: - return await asyncio.wait_for(self._task, timeout=timeout) - return await self._task + return await asyncio.wait_for(asyncio.shield(self._task), timeout=timeout) + return await asyncio.shield(self._task) except asyncio.TimeoutError as e: raise MilvusException(message="Timeout waiting for optimization to complete") from e except asyncio.CancelledError as e: - raise MilvusException(message="Optimization task was cancelled") from e + if self._task.cancelled(): + raise MilvusException(message="Optimization task was cancelled") from e + raise def set_progress(self, stage: ProgressStage) -> None: if self._cancelled: diff --git a/tests/unit/test_async_milvus_client_ops.py b/tests/unit/test_async_milvus_client_ops.py index 0873acda8..b9a5363ef 100644 --- a/tests/unit/test_async_milvus_client_ops.py +++ b/tests/unit/test_async_milvus_client_ops.py @@ -1,3 +1,4 @@ +import asyncio import inspect from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -444,6 +445,97 @@ async def test_upsert_exception_propagates(self): class TestAsyncClientOptimize: + @pytest.mark.asyncio + @pytest.mark.parametrize("timeout", [0, 0.001]) + async def test_wait_true_timeout_cleans_up_owned_task(self, timeout): + client, _ = _make_client() + tasks = [] + start = AsyncOptimizeTask.start + + def record_start(task): + tasks.append(task) + start(task) + + async def execute(**kwargs): + await asyncio.Event().wait() + + with patch.object(AsyncOptimizeTask, "start", record_start), patch.object( + client, "_execute_optimize", execute + ): + try: + with pytest.raises(MilvusException, match="Timeout waiting"): + await client.optimize("col", wait=True, timeout=timeout) + assert len(tasks) == 1 + assert tasks[0].done() + assert tasks[0].cancelled() + finally: + for task in tasks: + task.cancel() + await asyncio.gather(task._task, return_exceptions=True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("timeout", [None, 10]) + async def test_wait_true_caller_cancellation_cleans_up_owned_task(self, timeout): + client, _ = _make_client() + started = asyncio.Event() + tasks = [] + + async def execute(task, **kwargs): + tasks.append(task) + started.set() + await asyncio.Event().wait() + + with patch.object(client, "_execute_optimize", execute): + waiter = asyncio.create_task(client.optimize("col", wait=True, timeout=timeout)) + await started.wait() + try: + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert tasks[0].done() + assert tasks[0].cancelled() + finally: + tasks[0].cancel() + await asyncio.gather(waiter, tasks[0]._task, return_exceptions=True) + + @pytest.mark.asyncio + async def test_wait_timeout_preserves_optimization_and_other_waiter(self): + client, handler = _make_client() + compacting = asyncio.Event() + finish_compaction = asyncio.Event() + + async def compaction_state(*args, **kwargs): + compacting.set() + await finish_compaction.wait() + return MagicMock(state=2) + + handler.compact.return_value = 99 + handler.get_compaction_state.side_effect = compaction_state + handler.get_load_state.return_value = LoadState.Loaded + handler._get_schema.return_value = ({"fields": []}, 0) + handler.list_indexes.return_value = [] + + with patch.object(client, "refresh_load", new_callable=AsyncMock) as refresh_load: + task = await client.optimize("col", wait=False) + await compacting.wait() + other_waiter = asyncio.create_task(task.result()) + await asyncio.sleep(0) + try: + with pytest.raises(MilvusException, match="Timeout waiting"): + await task.result(timeout=0) + assert not task.done() + assert not other_waiter.done() + finish_compaction.set() + result = await other_waiter + assert result.status == "success" + assert result.compaction_id == 99 + assert await task.result() is result + handler.compact.assert_awaited_once() + refresh_load.assert_awaited_once_with("col", timeout=None) + finally: + finish_compaction.set() + await asyncio.gather(task._task, other_waiter, return_exceptions=True) + @pytest.mark.asyncio async def test_is_collection_loaded_true(self): client, handler = _make_client() diff --git a/tests/unit/test_async_optimize_task.py b/tests/unit/test_async_optimize_task.py index 8d3dc25f0..81aaa1123 100644 --- a/tests/unit/test_async_optimize_task.py +++ b/tests/unit/test_async_optimize_task.py @@ -63,6 +63,143 @@ def test_check_cancelled_when_cancelled_raises(self): class TestAsyncOptimizeTaskResult: + @pytest.mark.asyncio + @pytest.mark.parametrize("timeout", [None, 10]) + async def test_cancelled_result_waiter_preserves_background_task(self, timeout): + started = asyncio.Event() + finish = asyncio.Event() + expected = MagicMock(collection_name="col") + + async def execute(**kwargs): + started.set() + await finish.wait() + return expected + + task = _make_task(execute) + task.start() + await started.wait() + waiter = asyncio.create_task(task.result(timeout=timeout)) + other_waiter = asyncio.create_task(task.result()) + await asyncio.sleep(0) + try: + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert not task.done() + assert not task.cancelled() + assert not other_waiter.done() + finish.set() + assert await other_waiter is expected + finally: + finish.set() + await asyncio.gather(task._task, waiter, other_waiter, return_exceptions=True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("timeout", [None, 10]) + async def test_external_wait_for_preserves_background_task(self, timeout): + started = asyncio.Event() + finish = asyncio.Event() + expected = MagicMock(collection_name="col") + + async def execute(**kwargs): + started.set() + await finish.wait() + return expected + + task = _make_task(execute) + task.start() + await started.wait() + try: + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(task.result(timeout=timeout), timeout=0.001) + assert not task.done() + assert not task.cancelled() + finish.set() + assert await task.result() is expected + finally: + finish.set() + await asyncio.gather(task._task, return_exceptions=True) + + @pytest.mark.asyncio + @pytest.mark.skipif(not hasattr(asyncio, "timeout"), reason="Requires Python 3.11+") + async def test_external_timeout_context_preserves_background_task(self): + started = asyncio.Event() + finish = asyncio.Event() + expected = MagicMock(collection_name="col") + + async def execute(**kwargs): + started.set() + await finish.wait() + return expected + + task = _make_task(execute) + task.start() + await started.wait() + try: + with pytest.raises(asyncio.TimeoutError): + async with asyncio.timeout(0.001): + await task.result() + assert not task.done() + assert not task.cancelled() + finish.set() + assert await task.result() is expected + finally: + finish.set() + await asyncio.gather(task._task, return_exceptions=True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("timeout", [0, 0.001]) + @pytest.mark.parametrize("fails", [False, True]) + async def test_wait_timeout_preserves_background_result(self, timeout, fails): + started = asyncio.Event() + finish = asyncio.Event() + expected = MagicMock(collection_name="col") + + async def execute(**kwargs): + started.set() + await finish.wait() + if fails: + raise MilvusException(message="Index rebuild failed") + return expected + + task = _make_task(execute) + task.start() + await started.wait() + try: + with pytest.raises(MilvusException, match="Timeout waiting"): + await task.result(timeout=timeout) + assert not task.done() + assert not task.cancelled() + assert ProgressStage.CANCELLED not in task.progress_history() + finish.set() + if fails: + with pytest.raises(MilvusException, match="Index rebuild failed"): + await task.result() + else: + assert await task.result() is expected + finally: + finish.set() + await asyncio.gather(task._task, return_exceptions=True) + + @pytest.mark.asyncio + async def test_explicit_cancel_interrupts_timed_wait(self): + started = asyncio.Event() + + async def execute(**kwargs): + started.set() + await asyncio.Event().wait() + + task = _make_task(execute) + task.start() + await started.wait() + waiter = asyncio.create_task(task.result(timeout=10)) + await asyncio.sleep(0) + assert task.cancel() + with pytest.raises(MilvusException, match="cancelled"): + await waiter + assert task.done() + assert task.cancelled() + def test_result_without_start_raises(self): task = _make_task() @@ -84,19 +221,6 @@ async def run(): asyncio.run(run()) - def test_result_timeout_raises(self): - async def slow_execute(**kwargs): - await asyncio.sleep(10) - return MagicMock() - - async def run(): - task = AsyncOptimizeTask("col", None, None, slow_execute) - task.start() - with pytest.raises(MilvusException): - await task.result(timeout=0.01) - - asyncio.run(run()) - def test_result_cancelled_raises(self): async def run(): execute_fn = AsyncMock(return_value=MagicMock())