diff --git a/mock_tests/test_batch.py b/mock_tests/test_batch.py index ad7357f04..4d9ce6805 100644 --- a/mock_tests/test_batch.py +++ b/mock_tests/test_batch.py @@ -1,9 +1,12 @@ -from typing import AsyncGenerator, Generator, List +import asyncio +from typing import AsyncGenerator, Generator, List, Optional import grpc import pytest import pytest_asyncio import weaviate +from weaviate.collections.batch.async_ import _BgTasks +from weaviate.collections.batch.base import _BgThreads from weaviate.proto.v1 import batch_pb2, weaviate_pb2_grpc from .conftest import MOCK_IP, MOCK_PORT, MOCK_PORT_GRPC, mock_class, HTTPServer @@ -160,3 +163,61 @@ def test_ssb_stream_reports_has_errors( batch.add_object({"name": f"Object {i}"}) assert len(failed_object_stream.batch.failed_objects) == 2 assert failed_object_stream.batch.results.objs.has_errors + + +def _interrupt_the_join(monkeypatch: pytest.MonkeyPatch) -> None: + """Turn the join at the end of a `stream()` block into a Ctrl-C. + + `_ContextManagerSync.__exit__` waits for the background threads before it hands the errors + over, so a Ctrl-C arriving during that wait is what issue #1300 is about. The real join + runs first, so the stream itself still completes and the expected error count is exact. + """ + real_join = _BgThreads.join + + def join(threads: _BgThreads, timeout: Optional[float] = None) -> None: + real_join(threads, timeout) + raise KeyboardInterrupt + + monkeypatch.setattr(_BgThreads, "join", join) + + +def _cancel_the_gather(monkeypatch: pytest.MonkeyPatch) -> None: + """The async equivalent: a notebook interrupt reaches an awaiting cell as a cancellation.""" + real_gather = _BgTasks.gather + + async def gather(tasks: _BgTasks, timeout: Optional[float] = None) -> None: + await real_gather(tasks, timeout) + raise asyncio.CancelledError + + monkeypatch.setattr(_BgTasks, "gather", gather) + + +def test_ssb_stream_keeps_errors_reachable_after_an_interrupt( + failed_object_stream: weaviate.collections.Collection, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _interrupt_the_join(monkeypatch) + + with pytest.raises(KeyboardInterrupt): + with failed_object_stream.batch.stream() as batch: + for i in range(4): + batch.add_object({"name": f"Object {i}"}) + + assert len(failed_object_stream.batch.failed_objects) == 2 + assert failed_object_stream.batch.results.objs.has_errors + + +@pytest.mark.asyncio +async def test_ssb_stream_keeps_errors_reachable_after_an_interrupt_async( + failed_object_stream_async: weaviate.collections.CollectionAsync, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _cancel_the_gather(monkeypatch) + + with pytest.raises(asyncio.CancelledError): + async with failed_object_stream_async.batch.stream() as batch: + for i in range(4): + await batch.add_object({"name": f"Object {i}"}) + + assert len(failed_object_stream_async.batch.failed_objects) == 2 + assert failed_object_stream_async.batch.results.objs.has_errors diff --git a/test/collection/test_batch.py b/test/collection/test_batch.py index 79053710e..06ef740ef 100644 --- a/test/collection/test_batch.py +++ b/test/collection/test_batch.py @@ -1,9 +1,23 @@ +import asyncio +import os +import signal +import sys +import threading +import time import uuid +from types import SimpleNamespace +from typing import Awaitable, Callable, Iterator import pytest -from weaviate.collections.batch.base import _RateLimitedBatching +from weaviate.collections.batch.async_ import _BatchBaseAsync +from weaviate.collections.batch.base import ( + _BatchBase, + _BatchDataWrapper, + _RateLimitedBatching, +) from weaviate.collections.batch.grpc_batch import _validate_props +from weaviate.collections.batch.sync import _BatchBaseSync from weaviate.collections.classes.batch import ( MAX_STORED_RESULTS, BatchObject, @@ -12,6 +26,7 @@ BatchReferenceReturn, ErrorObject, ErrorReference, + Shard, ) from weaviate.exceptions import WeaviateInsertInvalidPropertyError @@ -141,3 +156,200 @@ def test_validate_props_raises_for_top_level_vector() -> None: def test_validate_props_raises_for_nested_vector() -> None: with pytest.raises(WeaviateInsertInvalidPropertyError): _validate_props({"vector": [0.1, 0.2]}, nested=True) + + +def _collect_one_of_everything(collected: _BatchDataWrapper) -> _BatchDataWrapper: + """Fill a batch's internal results the way a partially failed run would.""" + err_obj = _error_object(0) + err_ref = _error_reference(0) + collected.failed_objects.append(err_obj) + collected.failed_references.append(err_ref) + collected.results.objs += BatchObjectReturn(_all_responses=[err_obj], errors={0: err_obj}) + collected.results.refs += BatchReferenceReturn(errors={0: err_ref}) + collected.imported_shards.add(Shard(collection="Test")) + return collected + + +def _assert_published(published: _BatchDataWrapper, collected: _BatchDataWrapper) -> None: + """Assert that everything gathered so far reached the public accessors, as a snapshot.""" + assert len(published.failed_objects) == 1 + assert len(published.failed_references) == 1 + assert published.results.objs.has_errors + assert published.results.refs.has_errors + assert published.imported_shards == {Shard(collection="Test")} + + # The background workers outlive an interrupted shutdown, so what was published has to be + # a snapshot: a user who prints `len(batch.failed_objects)` and then builds a frame from + # the same property in the next cell must not get two different answers. + collected.failed_objects.append(_error_object(1)) + collected.failed_references.append(_error_reference(1)) + collected.imported_shards.add(Shard(collection="Other")) + # `BatchObjectReturn.__add__` mutates the left-hand side, so this also catches a `results` + # that was handed over by reference instead of rebuilt. + collected.results.objs += BatchObjectReturn(errors={1: _error_object(1)}) + collected.results.refs += BatchReferenceReturn(errors={1: _error_reference(1)}) + assert len(published.failed_objects) == 1 + assert len(published.failed_references) == 1 + assert published.imported_shards == {Shard(collection="Test")} + assert len(published.results.objs.errors) == 1 + assert len(published.results.refs.errors) == 1 + + +def _bare_batch(published: _BatchDataWrapper) -> _BatchBase: + """Build a `_BatchBase` carrying only the state that `_shutdown` touches. + + `_BatchBase.__init__` needs a live connection and starts background threads, neither of + which belongs in a unit test, so the handful of attributes `_shutdown` uses are set + directly. The `_BatchBase__` prefixes are what Python's name mangling turns the + `__`-private names into. + """ + batch = object.__new__(_BatchBase) + # The wrapper the public `batch.failed_objects` / `batch.results` accessors read. + batch._BatchBase__results_for_wrapper_backup = published + # The wrapper the batch itself collects into while it runs. + batch._BatchBase__results_for_wrapper = _BatchDataWrapper() + batch._BatchBase__results_lock = threading.Lock() + batch._BatchBase__shut_background_thread_down = threading.Event() + # A thread that was never started is never alive, so `_shutdown`'s wait loop is a no-op. + batch._BatchBase__bg_threads = threading.Thread(target=lambda: None) + return batch + + +def _bare_stream_batch( + published: _BatchDataWrapper, join: Callable[[float], None] +) -> _BatchBaseSync: + """Build a `_BatchBaseSync` (the `batch.stream()` colour) carrying only what `_wait` reads.""" + batch = object.__new__(_BatchBaseSync) + batch._BatchBaseSync__results_for_wrapper_backup = published + batch._BatchBaseSync__results_for_wrapper = _BatchDataWrapper() + batch._BatchBaseSync__results_lock = threading.Lock() + batch._BatchBaseSync__connection = SimpleNamespace(timeout_config=SimpleNamespace(insert=1)) + batch._BatchBaseSync__bg_threads = SimpleNamespace(join=join) + return batch + + +def _bare_stream_batch_async( + published: _BatchDataWrapper, gather: Callable[..., Awaitable[None]] +) -> _BatchBaseAsync: + """Build a `_BatchBaseAsync` (the only colour the async client exposes) for `_wait`.""" + batch = object.__new__(_BatchBaseAsync) + batch._BatchBaseAsync__results_for_wrapper_backup = published + batch._BatchBaseAsync__results_for_wrapper = _BatchDataWrapper() + batch._BatchBaseAsync__connection = SimpleNamespace(timeout_config=SimpleNamespace(insert=1)) + batch._BatchBaseAsync__bg_tasks = SimpleNamespace(gather=gather) + return batch + + +def _interrupt_main_with_sigint() -> None: + """Send a real SIGINT to this process and wait for the interpreter to act on it. + + `os.kill` only trips a flag in the C signal handler; CPython turns that into a + `KeyboardInterrupt` in the main thread at the next bytecode boundary, which the sleep + loop below is guaranteed to reach. That is the same path a Ctrl-C takes in a notebook. + """ + os.kill(os.getpid(), signal.SIGINT) + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + time.sleep(0.01) + raise AssertionError("SIGINT did not raise KeyboardInterrupt in the main thread") + + +@pytest.fixture +def default_sigint_handler() -> Iterator[None]: + """Make SIGINT raise `KeyboardInterrupt` regardless of what the test runner installed.""" + previous = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGINT, signal.default_int_handler) + try: + yield + finally: + signal.signal(signal.SIGINT, previous) + + +needs_sigint = pytest.mark.skipif( + sys.platform == "win32", + reason="os.kill cannot deliver SIGINT to the current process on Windows", +) + + +@needs_sigint +def test_shutdown_publishes_results_when_interrupted(default_sigint_handler: None) -> None: + published = _BatchDataWrapper() + batch = _bare_batch(published) + collected = _collect_one_of_everything(batch._BatchBase__results_for_wrapper) + batch.flush = _interrupt_main_with_sigint # type: ignore[method-assign] + + with pytest.raises(KeyboardInterrupt): + batch._shutdown() + + _assert_published(published, collected) + # The interrupt unwound out of `flush()` before the batch asked its daemon threads to + # stop, so `_shutdown` has to ask on the way out: otherwise they keep uploading the batch + # the user just aborted, and a notebook that retries leaks a pair of them per attempt. + assert batch._BatchBase__shut_background_thread_down.is_set() + + +def test_shutdown_publishes_results_on_a_clean_exit() -> None: + published = _BatchDataWrapper() + batch = _bare_batch(published) + collected = _collect_one_of_everything(batch._BatchBase__results_for_wrapper) + batch.flush = lambda: None # type: ignore[method-assign] + + batch._shutdown() + + _assert_published(published, collected) + assert batch._BatchBase__shut_background_thread_down.is_set() + + +@needs_sigint +def test_stream_wait_publishes_results_when_interrupted(default_sigint_handler: None) -> None: + published = _BatchDataWrapper() + batch = _bare_stream_batch(published, join=lambda _timeout: _interrupt_main_with_sigint()) + collected = _collect_one_of_everything(batch._BatchBaseSync__results_for_wrapper) + + with pytest.raises(KeyboardInterrupt): + batch._wait() + + _assert_published(published, collected) + + +def test_stream_wait_publishes_results_on_a_clean_exit() -> None: + published = _BatchDataWrapper() + batch = _bare_stream_batch(published, join=lambda _timeout: None) + collected = _collect_one_of_everything(batch._BatchBaseSync__results_for_wrapper) + + batch._wait() + + _assert_published(published, collected) + + +@pytest.mark.parametrize("interrupt", [KeyboardInterrupt, asyncio.CancelledError]) +def test_stream_wait_async_publishes_results_when_interrupted( + interrupt: type[BaseException], +) -> None: + # A Ctrl-C in a notebook cell that is awaiting reaches the coroutine as a CancelledError + # thrown in at the await point; `asyncio.run` in a script surfaces the KeyboardInterrupt + # itself. Both are BaseExceptions, so only a `finally` catches them. + async def gather(timeout: float) -> None: + raise interrupt() + + published = _BatchDataWrapper() + batch = _bare_stream_batch_async(published, gather=gather) + collected = _collect_one_of_everything(batch._BatchBaseAsync__results_for_wrapper) + + with pytest.raises(interrupt): + asyncio.run(batch._wait()) + + _assert_published(published, collected) + + +def test_stream_wait_async_publishes_results_on_a_clean_exit() -> None: + async def gather(timeout: float) -> None: + return None + + published = _BatchDataWrapper() + batch = _bare_stream_batch_async(published, gather=gather) + collected = _collect_one_of_everything(batch._BatchBaseAsync__results_for_wrapper) + + asyncio.run(batch._wait()) + + _assert_published(published, collected) diff --git a/weaviate/collections/batch/async_.py b/weaviate/collections/batch/async_.py index c63ec2106..c27aa85f7 100644 --- a/weaviate/collections/batch/async_.py +++ b/weaviate/collections/batch/async_.py @@ -19,6 +19,7 @@ _BatchDataWrapper, _BatchStreamRequest, _ClusterBatchAsync, + _publish_results, ) from weaviate.collections.batch.grpc_batch import _BatchGRPC from weaviate.collections.classes.batch import ( @@ -191,16 +192,11 @@ async def _wait(self) -> None: raise WeaviateBatchStreamError( "Background batch tasks did not terminate after forced shutdown." ) from e - - # copy the results to the public results - self.__results_for_wrapper_backup.results = self.__results_for_wrapper.results - self.__results_for_wrapper_backup.failed_objects = self.__results_for_wrapper.failed_objects - self.__results_for_wrapper_backup.failed_references = ( - self.__results_for_wrapper.failed_references - ) - self.__results_for_wrapper_backup.imported_shards = ( - self.__results_for_wrapper.imported_shards - ) + finally: + # publish the results, also when the gather above was cut short by an interrupt or + # a cancellation (e.g. Ctrl-C in a notebook): the errors gathered so far are the + # only record of what did not make it into Weaviate + _publish_results(self.__results_for_wrapper, self.__results_for_wrapper_backup) async def _shutdown(self) -> None: self.__is_stopped.set() diff --git a/weaviate/collections/batch/base.py b/weaviate/collections/batch/base.py index 1d736303c..ffea63e9e 100644 --- a/weaviate/collections/batch/base.py +++ b/weaviate/collections/batch/base.py @@ -243,6 +243,40 @@ class _BatchDataWrapper: imported_shards: Set[Shard] = field(default_factory=set) +def _publish_results(collected: _BatchDataWrapper, published: _BatchDataWrapper) -> None: + """Copy what a batch has collected onto the wrapper its public accessors read from. + + Snapshots rather than aliases: when a batch is torn down by an interrupt its background + workers are still running, and publishing the live containers would hand the user an error + list that keeps growing while they inspect it. `BatchObjectReturn.__add__` and + `BatchReferenceReturn.__add__` both mutate the left-hand side in place, so the two returns + have to be rebuilt rather than reassigned. The errors themselves are shared, not deep + copied - nothing ever mutates one. + + Args: + collected: The wrapper the batch collects into while it runs. + published: The wrapper `batch.failed_objects` and friends read from. + """ + objs, refs = collected.results.objs, collected.results.refs + results = BatchResult() + results.objs = BatchObjectReturn( + _all_responses=list(objs._all_responses), + elapsed_seconds=objs.elapsed_seconds, + errors=dict(objs.errors), + uuids=dict(objs.uuids), + has_errors=objs.has_errors, + ) + results.refs = BatchReferenceReturn( + elapsed_seconds=refs.elapsed_seconds, + errors=dict(refs.errors), + has_errors=refs.has_errors, + ) + published.results = results + published.failed_objects = list(collected.failed_objects) + published.failed_references = list(collected.failed_references) + published.imported_shards = set(collected.imported_shards) + + @dataclass class _DynamicBatching: pass @@ -384,22 +418,24 @@ def _wait(self): def _shutdown(self) -> None: """Shutdown the current batch and wait for all requests to be finished.""" - self.flush() - - # we are done, shut bg threads down and end the event loop - self.__shut_background_thread_down.set() - while self.__bg_threads.is_alive(): - time.sleep(0.01) - - # copy the results to the public results - self.__results_for_wrapper_backup.results = self.__results_for_wrapper.results - self.__results_for_wrapper_backup.failed_objects = self.__results_for_wrapper.failed_objects - self.__results_for_wrapper_backup.failed_references = ( - self.__results_for_wrapper.failed_references - ) - self.__results_for_wrapper_backup.imported_shards = ( - self.__results_for_wrapper.imported_shards - ) + try: + self.flush() + + # we are done, shut bg threads down and end the event loop + self.__shut_background_thread_down.set() + while self.__bg_threads.is_alive(): + time.sleep(0.01) + finally: + # an interrupt (e.g. Ctrl-C in a notebook) unwinds out of flush() before the line + # above, so ask the daemon threads to stop here too: otherwise they keep uploading + # the batch the user just aborted, and one pair of them leaks per attempt + self.__shut_background_thread_down.set() + # publish the results, also when the wait above was cut short by that interrupt: + # the errors gathered so far are the only record of what did not make it into + # Weaviate, and the batch has already told the user to inspect + # `batch.failed_objects` for them + with self.__results_lock: + _publish_results(self.__results_for_wrapper, self.__results_for_wrapper_backup) def __batch_send(self) -> None: refresh_time: float = 0.01 diff --git a/weaviate/collections/batch/sync.py b/weaviate/collections/batch/sync.py index 6cf8c1edc..df58472b7 100644 --- a/weaviate/collections/batch/sync.py +++ b/weaviate/collections/batch/sync.py @@ -16,6 +16,7 @@ _BatchStreamRequest, _BgThreads, _ClusterBatch, + _publish_results, ) from weaviate.collections.batch.grpc_batch import _BatchGRPC from weaviate.collections.classes.batch import ( @@ -141,16 +142,12 @@ def _wait(self) -> None: raise WeaviateBatchStreamError( "Background batch threads did not terminate after forced shutdown." ) from e - - # copy the results to the public results - self.__results_for_wrapper_backup.results = self.__results_for_wrapper.results - self.__results_for_wrapper_backup.failed_objects = self.__results_for_wrapper.failed_objects - self.__results_for_wrapper_backup.failed_references = ( - self.__results_for_wrapper.failed_references - ) - self.__results_for_wrapper_backup.imported_shards = ( - self.__results_for_wrapper.imported_shards - ) + finally: + # publish the results, also when the join above was cut short by an interrupt + # (e.g. Ctrl-C in a notebook): the errors gathered so far are the only record of + # what did not make it into Weaviate + with self.__results_lock: + _publish_results(self.__results_for_wrapper, self.__results_for_wrapper_backup) def _shutdown(self) -> None: # Shutdown the current batch and wait for all requests to be finished