diff --git a/test/test_batch_stream_async.py b/test/test_batch_stream_async.py new file mode 100644 index 000000000..2c3c45e42 --- /dev/null +++ b/test/test_batch_stream_async.py @@ -0,0 +1,127 @@ +"""Unit tests for the async batch-stream failure handling (no cluster needed).""" + +import asyncio +import logging +from typing import Optional + +import pytest + +from weaviate.collections.batch.async_ import _BatchBaseAsync, _BgTasks +from weaviate.collections.batch.base import _BatchDataWrapper +from weaviate.collections.batch.batch_wrapper import _ContextManagerAsync +from weaviate.exceptions import WeaviateBatchStreamError + + +class _NotAnException(BaseException): + """Escapes the wrappers' `except Exception` like grpc.aio's CancelledError does.""" + + +def _bare_batch(**mangled) -> _BatchBaseAsync: + batch = object.__new__(_BatchBaseAsync) + for name, value in mangled.items(): + setattr(batch, f"_BatchBaseAsync__{name}", value) + return batch + + +class _FakeTimeouts: + insert = 1 + + +class _FakeConnection: + timeout_config = _FakeTimeouts() + + +def test_wait_raises_the_background_exception_and_keeps_partial_results() -> None: + # a background failure must not come back as a success, and a user catching it + # must still see what failed + class FakeBgTasks: + async def gather(self, timeout=None) -> None: + return None + + partial = _BatchDataWrapper() + partial.failed_objects = ["sentinel-failure"] # type: ignore[list-item] + backup = _BatchDataWrapper() + batch = _bare_batch( + bg_exception=RuntimeError("boom"), + bg_tasks=FakeBgTasks(), + connection=_FakeConnection(), + results_for_wrapper=partial, + results_for_wrapper_backup=backup, + ) + + with pytest.raises(RuntimeError, match="boom"): + asyncio.run(batch._wait()) + assert backup.failed_objects == ["sentinel-failure"] + + +def test_flush_raises_when_a_task_dies_without_setting_bg_exception() -> None: + # loop_wrapper/recv_wrapper only catch Exception, so a BaseException kills a task with + # __bg_exception unset; flush() must notice the dead task instead of spinning forever + async def run() -> None: + async def dies() -> None: + raise _NotAnException("boom") + + loop_task = asyncio.get_running_loop().create_task(asyncio.sleep(3600)) + recv_task = asyncio.get_running_loop().create_task(dies()) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert recv_task.done() + + batch = _bare_batch( + bg_exception=None, + bg_tasks=_BgTasks(recv=recv_task, loop=loop_task), + batch_objects=[object()], + batch_references=[], + ) + try: + await asyncio.wait_for(batch.flush(), timeout=2) # a regression hangs here + finally: + loop_task.cancel() + + with pytest.raises(WeaviateBatchStreamError, match="stream has ended"): + asyncio.run(run()) + + +class _FakeBatch: + """Stands in for _BatchBaseAsync behind the context manager.""" + + def __init__(self, wait_error: Optional[BaseException] = None) -> None: + self.wait_error = wait_error + self.wait_called = False + + async def _start(self) -> None: + pass + + async def _shutdown(self) -> None: + pass + + async def _wait(self) -> None: + self.wait_called = True + if self.wait_error is not None: + raise self.wait_error + + +def test_aexit_raises_a_background_failure_on_a_clean_block() -> None: + fake = _FakeBatch(WeaviateBatchStreamError("bg died")) + + async def run() -> None: + async with _ContextManagerAsync(fake): # type: ignore[arg-type] + pass + + with pytest.raises(WeaviateBatchStreamError, match="bg died"): + asyncio.run(run()) + + +def test_aexit_keeps_the_users_exception_over_a_background_failure(caplog) -> None: + # the block's own exception must not be replaced by the background failure + fake = _FakeBatch(WeaviateBatchStreamError("bg died")) + + async def run() -> None: + async with _ContextManagerAsync(fake): # type: ignore[arg-type] + raise ValueError("user code") + + with caplog.at_level(logging.WARNING, logger="weaviate-client"): + with pytest.raises(ValueError, match="user code"): + asyncio.run(run()) + assert fake.wait_called # still drained/awaited + assert "bg died" in caplog.text diff --git a/test/test_batch_stream_sync.py b/test/test_batch_stream_sync.py new file mode 100644 index 000000000..ff3a76e93 --- /dev/null +++ b/test/test_batch_stream_sync.py @@ -0,0 +1,126 @@ +"""Unit tests for the sync batch-stream failure handling (no cluster needed).""" + +import logging +import time +from typing import Optional + +import pytest + +from weaviate.collections.batch.base import _BatchDataWrapper +from weaviate.collections.batch.batch_wrapper import _ContextManagerSync +from weaviate.collections.batch.sync import _BatchBaseSync +from weaviate.exceptions import WeaviateBatchStreamError + + +def _bare_batch(**mangled) -> _BatchBaseSync: + batch = object.__new__(_BatchBaseSync) + for name, value in mangled.items(): + setattr(batch, f"_BatchBaseSync__{name}", value) + return batch + + +class _FakeThreads: + def __init__(self, alive: bool = False) -> None: + self.alive = alive + + def join(self, timeout=None) -> None: + return None + + def is_alive(self) -> bool: + return self.alive + + +class _FakeTimeouts: + insert = 1 + + +class _FakeConnection: + timeout_config = _FakeTimeouts() + + +def _batch_for_wait(**mangled) -> _BatchBaseSync: + defaults = { + "bg_exception": None, + "bg_threads": _FakeThreads(alive=False), + "connection": _FakeConnection(), + "results_for_wrapper": _BatchDataWrapper(), + "results_for_wrapper_backup": _BatchDataWrapper(), + "batch_objects": [], + "batch_references": [], + } + defaults.update(mangled) + return _bare_batch(**defaults) + + +def test_wait_raises_the_background_exception_and_keeps_partial_results() -> None: + partial = _BatchDataWrapper() + partial.failed_objects = ["sentinel-failure"] # type: ignore[list-item] + backup = _BatchDataWrapper() + batch = _batch_for_wait( + bg_exception=RuntimeError("boom"), + results_for_wrapper=partial, + results_for_wrapper_backup=backup, + ) + + with pytest.raises(RuntimeError, match="boom"): + batch._wait() + assert backup.failed_objects == ["sentinel-failure"] + + +def test_wait_names_unsent_data_when_the_threads_are_gone() -> None: + batch = _batch_for_wait(batch_objects=[object(), object()], batch_references=[object()]) + with pytest.raises( + WeaviateBatchStreamError, match="ended with 2 objects and 1 references unsent" + ): + batch._wait() + + +def test_start_raises_the_stored_background_exception_at_once() -> None: + # a background thread that died right away (e.g. on a closed connection) used to leave + # _start polling for 60 s and then blaming the network (#2139) + batch = _bare_batch(bg_exception=None, bg_threads=None) + + def start_dead_threads() -> None: + setattr(batch, "_BatchBaseSync__bg_threads", _FakeThreads(alive=False)) + setattr(batch, "_BatchBaseSync__bg_exception", RuntimeError("connection closed")) + + setattr(batch, "_BatchBaseSync__start_bg_threads", start_dead_threads) + started = time.monotonic() + with pytest.raises(RuntimeError, match="connection closed"): + batch._start() + assert time.monotonic() - started < 5 + + +class _FakeBatch: + def __init__(self, wait_error: Optional[BaseException] = None) -> None: + self.wait_error = wait_error + self.wait_called = False + + def _start(self) -> None: + pass + + def _shutdown(self) -> None: + pass + + def _wait(self) -> None: + self.wait_called = True + if self.wait_error is not None: + raise self.wait_error + + +def test_exit_raises_a_background_failure_on_a_clean_block() -> None: + # the sync colour used to swallow this entirely + fake = _FakeBatch(WeaviateBatchStreamError("bg died")) + with pytest.raises(WeaviateBatchStreamError, match="bg died"): + with _ContextManagerSync(fake): # type: ignore[type-var] + pass + + +def test_exit_keeps_the_users_exception_over_a_background_failure(caplog) -> None: + fake = _FakeBatch(WeaviateBatchStreamError("bg died")) + with caplog.at_level(logging.WARNING, logger="weaviate-client"): + with pytest.raises(ValueError, match="user code"): + with _ContextManagerSync(fake): # type: ignore[type-var] + raise ValueError("user code") + assert fake.wait_called + assert "bg died" in caplog.text diff --git a/weaviate/collections/batch/async_.py b/weaviate/collections/batch/async_.py index c63ec2106..e3d255ea3 100644 --- a/weaviate/collections/batch/async_.py +++ b/weaviate/collections/batch/async_.py @@ -202,6 +202,16 @@ async def _wait(self) -> None: self.__results_for_wrapper.imported_shards ) + if self.__bg_exception is not None: + # surface the background failure instead of returning partial results as a success + raise self.__bg_exception + n_objs, n_refs = len(self.__batch_objects), len(self.__batch_references) + if n_objs + n_refs > 0 and not self.__all_tasks_alive(): + # the tasks are gone with data still queued: the batch did NOT complete + raise WeaviateBatchStreamError( + f"batch stream ended with {n_objs} objects and {n_refs} references unsent" + ) + async def _shutdown(self) -> None: self.__is_stopped.set() @@ -532,6 +542,8 @@ async def flush(self) -> None: # bg thread is sending objs+refs automatically, so simply wait for everything to be done while len(self.__batch_objects) > 0 or len(self.__batch_references) > 0: await asyncio.sleep(0.01) + # a dead task means nothing drains the queues: raise, don't spin forever + self.__check_bg_tasks_alive() async def _add_object( self, @@ -628,4 +640,4 @@ def __check_bg_tasks_alive(self) -> None: if self.__all_tasks_alive(): return - raise self.__bg_exception or Exception("Batch tasks died unexpectedly") + raise self.__bg_exception or WeaviateBatchStreamError("the batch stream has ended") diff --git a/weaviate/collections/batch/batch_wrapper.py b/weaviate/collections/batch/batch_wrapper.py index a3a3598d6..0976cea11 100644 --- a/weaviate/collections/batch/batch_wrapper.py +++ b/weaviate/collections/batch/batch_wrapper.py @@ -508,7 +508,13 @@ def __init__(self, current_batch: T): def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.__current_batch._shutdown() - self.__current_batch._wait() + try: + self.__current_batch._wait() + except Exception as e: + if exc_type is None: + raise + # the exception leaving the block wins; the background failure is only logged + logger.warning(f"batch stream failed while the block raised {exc_type.__name__}: {e}") def __enter__(self) -> P: self.__current_batch._start() @@ -521,7 +527,13 @@ def __init__(self, current_batch: _BatchBaseAsync): async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: await self.__current_batch._shutdown() - await self.__current_batch._wait() + try: + await self.__current_batch._wait() + except Exception as e: + if exc_type is None: + raise + # the exception leaving the block wins; the background failure is only logged + logger.warning(f"batch stream failed while the block raised {exc_type.__name__}: {e}") async def __aenter__(self) -> Q: await self.__current_batch._start() diff --git a/weaviate/collections/batch/sync.py b/weaviate/collections/batch/sync.py index 6cf8c1edc..36f52fff8 100644 --- a/weaviate/collections/batch/sync.py +++ b/weaviate/collections/batch/sync.py @@ -125,7 +125,10 @@ def _start(self) -> None: logger.info("Provisioned stream to the server for batch processing") now = time.time() while not self.__all_threads_alive(): - # wait for the recv threads to be started + if self.__bg_exception is not None: + # a background thread already died (e.g. on a closed connection): raise its + # error now instead of polling for 60 s and blaming the network + raise self.__bg_exception time.sleep(0.01) if time.time() - now > 60: raise WeaviateBatchStreamError( @@ -152,6 +155,16 @@ def _wait(self) -> None: self.__results_for_wrapper.imported_shards ) + if self.__bg_exception is not None: + # surface the background failure instead of returning partial results as a success + raise self.__bg_exception + n_objs, n_refs = len(self.__batch_objects), len(self.__batch_references) + if n_objs + n_refs > 0 and not self.__all_threads_alive(): + # the threads are gone with data still queued: the batch did NOT complete + raise WeaviateBatchStreamError( + f"batch stream ended with {n_objs} objects and {n_refs} references unsent" + ) + def _shutdown(self) -> None: # Shutdown the current batch and wait for all requests to be finished self.__is_stopped.set() @@ -644,4 +657,4 @@ def __check_bg_threads_alive(self) -> None: if self.__all_threads_alive(): return - raise self.__bg_exception or Exception("Batch thread died unexpectedly") + raise self.__bg_exception or WeaviateBatchStreamError("the batch stream has ended") diff --git a/weaviate/exceptions.py b/weaviate/exceptions.py index ce0fe6f7e..5440af6d5 100644 --- a/weaviate/exceptions.py +++ b/weaviate/exceptions.py @@ -418,5 +418,8 @@ def __init__(self, pb: version.Version, grpc: version.Version) -> None: ) -class _BatchStreamShutdownError(Exception): - """Internal exception to signal that the batch stream was shutdown.""" +class _BatchStreamShutdownError(WeaviateBatchStreamError): + """Internal exception to signal that the batch stream was shutdown (gRPC ABORTED).""" + + def __init__(self, message: str = "the server aborted the batch stream") -> None: + super().__init__(message)