From 5a8c4a03487c9d307ef6277dee0b7a1054990551 Mon Sep 17 00:00:00 2001 From: Dmitry Voropaev Date: Thu, 24 Sep 2026 07:19:01 +0300 Subject: [PATCH] fix(batch): keep batch errors reachable after an interrupted shutdown Every batch colour copies the results it collected during the run into the wrapper that `client.batch.failed_objects`, `collection.batch.failed_objects` and `batch.results` read, and every one of them does it only after the wait for the background workers has returned: `_BatchBase._shutdown()` for `dynamic()`, `fixed_size()` and `rate_limit()`, `_BatchBaseSync._wait()` and `_BatchBaseAsync._wait()` for `stream()`, which is the only mode the async client exposes. A Ctrl-C during that wait - the usual way a bulk import is stopped in a notebook, reaching an awaiting async cell as a `CancelledError` - unwinds before the copy, so every error collected so far is dropped and the accessors come back empty. The batch has already told the user to go and look at them: "Failed to send N objects in a batch of M. Please inspect client.batch.failed_objects or collection.batch.failed_objects for the failed objects." Move the copy into a `finally:` in all three places, so the results are published on the interrupted path as well as the happy one. Two details that only become visible once the errors are reachable: - The copy handed over the live containers, and on the interrupted path the background workers are still running, so the "final" error list kept growing while the user looked at it: `len()` in one notebook cell and a frame built from the same property in the next disagreed. The shared `_publish_results()` helper now publishes snapshots. `BatchObjectReturn.__add__` and `BatchReferenceReturn.__add__` mutate the left-hand side, so `results` has to be rebuilt rather than reassigned; the errors themselves are shared, not copied. - `_BatchBase._shutdown()` only asked its daemon threads to stop after `flush()` returned, so an interrupt left them uploading the batch the user had just aborted and a `BgBatchScheduler` thread survived every aborted attempt. The `finally:` now sets the shutdown event first. Objects still queued at that moment are no longer sent; a request already in flight still completes. Fixes #1300 --- mock_tests/test_batch.py | 63 +++++++- test/collection/test_batch.py | 214 ++++++++++++++++++++++++++- weaviate/collections/batch/async_.py | 16 +- weaviate/collections/batch/base.py | 68 +++++++-- weaviate/collections/batch/sync.py | 17 +-- 5 files changed, 340 insertions(+), 38 deletions(-) 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