diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4561fa5..f8f3a33 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,7 +54,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.13"] + python-version: ["3.11", "3.13"] services: mongodb: @@ -106,7 +106,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.13"] + python-version: ["3.11", "3.13"] steps: - uses: actions/checkout@v7 @@ -137,7 +137,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.13"] + python-version: ["3.11", "3.13"] runs-on: [windows-latest, macos-latest] steps: diff --git a/docs/tutorials/installation.md b/docs/tutorials/installation.md index 019933f..52451c8 100644 --- a/docs/tutorials/installation.md +++ b/docs/tutorials/installation.md @@ -8,7 +8,7 @@ :::{tab-item} pip :sync: pip -You will need Python 3.10 or later. You can check your version of Python by +You will need Python 3.11 or later. You can check your version of Python by typing into a terminal: ```sh diff --git a/docs/tutorials/saving.md b/docs/tutorials/saving.md index c0a8d7f..178b30b 100644 --- a/docs/tutorials/saving.md +++ b/docs/tutorials/saving.md @@ -13,7 +13,7 @@ server dependencies. :::{tab-item} pip :sync: pip -You will need Python 3.10 or later. You can check your version of Python by +You will need Python 3.11 or later. You can check your version of Python by typing into a terminal: ```sh diff --git a/pyproject.toml b/pyproject.toml index 30d23f6..106307b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ description = "Integrations for using Tiled with Bluesky" readme = { file = "README.md", content-type = "text/markdown" } license = "BSD-3-Clause" license-files = ["LICENSE"] -requires-python = ">=3.10" +requires-python = ">=3.11" classifiers = [ "Development Status :: 1 - Planning", "Intended Audience :: Science/Research", @@ -24,7 +24,6 @@ classifiers = [ "Programming Language :: Python", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -36,13 +35,14 @@ dynamic = ["version"] dependencies = [ "dask", "event-model", + "httpx2", # the generation used by the tiled client since 0.2.21 "mongoquery", "ndindex", # imported by tiled.ndslice "numpy", "pandas", "pyarrow", "pytz", - "tiled[minimal-client] >=0.2.19", + "tiled[minimal-client] >=0.2.21", "tzlocal", "xarray", ] @@ -88,7 +88,7 @@ test = [ "redis-json-dict", "suitcase-mongo", "tifffile", - "tiled[server] >=0.2.19", + "tiled[server] >=0.2.21", ] dev = [{ include-group = "test" }] docs = [ @@ -133,6 +133,11 @@ filterwarnings = [ # so they remain readable; zarr warns that these dtypes have no stable Zarr V3 # specification yet. This warning is emitted by the (in-process) test server. "ignore:The data type .* does not have a Zarr V3 specification", + # SQLAlchemy 2.1 deprecates the `noload` loader strategy, which tiled's catalog + # ORM uses for the Webhook.node relationship. Raising this warning as an error + # inside `configure_mappers()` poisons every mapper in the (in-process) test + # server; third-party issue, to be fixed in tiled. + "ignore:The ``noload`` loader strategy is deprecated:DeprecationWarning", ] log_level = "INFO" testpaths = ["tests"] @@ -144,7 +149,7 @@ report.exclude_also = ['\.\.\.', 'if typing.TYPE_CHECKING:'] [tool.mypy] files = ["src", "tests"] -python_version = "3.10" +python_version = "3.11" warn_unused_configs = true strict = false enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"] @@ -174,7 +179,7 @@ ignore = [ [tool.pylint] -py-version = "3.10" +py-version = "3.11" ignore-paths = [".*/_version.py"] reports.output-format = "colorized" similarities.ignore-imports = "yes" diff --git a/src/bluesky_tiled_plugins/__init__.py b/src/bluesky_tiled_plugins/__init__.py index 3f2833e..69bc4b0 100644 --- a/src/bluesky_tiled_plugins/__init__.py +++ b/src/bluesky_tiled_plugins/__init__.py @@ -6,6 +6,7 @@ BlueskyStreamUpdate, subscribe_to_streams, ) +from .writing._buffer import BufferingWrapper from .writing.tiled_writer import TiledWriter, TiledInserter __all__ = [ @@ -13,6 +14,7 @@ "BlueskyRun", "BlueskyStreamSubscription", "BlueskyStreamUpdate", + "BufferingWrapper", "CatalogOfBlueskyRuns", "TiledInserter", "TiledWriter", diff --git a/src/bluesky_tiled_plugins/clients/bluesky_run.py b/src/bluesky_tiled_plugins/clients/bluesky_run.py index 7be9ccf..83da862 100644 --- a/src/bluesky_tiled_plugins/clients/bluesky_run.py +++ b/src/bluesky_tiled_plugins/clients/bluesky_run.py @@ -1,7 +1,7 @@ import codecs import copy import functools -import httpx +import httpx2 import io import json import keyword @@ -534,13 +534,13 @@ def validate( msg = "Remote validation failed: " + "; ".join(notes) raise ValidationException(msg, self.item["id"]) - except httpx.HTTPStatusError as e: + except httpx2.HTTPStatusError as e: # Backcompatibility: if the server does not support validation endpoint, # it will return 404 Not Found error; in this case, attempt to validate # the data structure locally by the client itself (requires multiple # round-trips to the server, but better than nothing). - if response.status_code == httpx.codes.NOT_FOUND: + if response.status_code == httpx2.codes.NOT_FOUND: warnings.warn( "Tiled server does not support remote validation. " "Attempting to validate the data structure by the client.", diff --git a/src/bluesky_tiled_plugins/writing/_buffer.py b/src/bluesky_tiled_plugins/writing/_buffer.py new file mode 100644 index 0000000..8b888b4 --- /dev/null +++ b/src/bluesky_tiled_plugins/writing/_buffer.py @@ -0,0 +1,161 @@ +""" +This module vendors BufferingWrapper from bluesky.callbacks.buffer, with the +addition of the `flush` method, so that beamline profiles can use the read-after-write +synchronization and avoids a bluesky dependency. + +Once a bluesky release includes `BufferingWrapper.flush`, this copy can be +replaced by a re-export (or deprecated). The implementation is self-contained +(standard library only); keep it in sync with the upstream module. +""" + +import atexit +import logging +import threading +import time +from collections.abc import Callable +from queue import Empty, Full, Queue + +logger = logging.getLogger(__name__) + + +class BufferingWrapper: + """A wrapper for callbacks that processes documents in a separate thread. + + This class allows a callback to be executed in a background thread, processing + documents as they are received. This prevent the blocking of RunEngine on any + slow I/O operations by the callback. It handles graceful shutdown on exit or signal + termination, ensuring that no new documents are accepted after shutdown has been + initiated. + + The wrapped callback should be thread-safe and not subscribed to the RE directly. + If it maintains shared mutable state, it must protect it using internal locking. + + Because documents are processed asynchronously, a run may not be fully written + by the time ``RE(plan)`` returns. Code that reads a run back immediately after + it completes (e.g. ``db[-1]`` or ``db[uid]``) must call :meth:`flush` first to + wait for the buffer to drain; otherwise it may observe a missing or incomplete + run, or silently pick up the previous one. + + Parameters + ---------- + target : callable + The instance of a callback that will be called with the documents. + It should accept two parameters: `name` and `doc`. + queue_size : int, optional + The maximum size of the internal queue. Default is 1,000,000. + + Usage + ----- + tw = TiledWriter(client) + buff_tw = BufferingWrapper(tw) + RE.subscribe(buff_tw) + """ + + def __init__(self, target: Callable, queue_size: int = 1_000_000): + self._wrapped_callback = target + self._queue: Queue = Queue(maxsize=queue_size) + self._stop_event = threading.Event() + self._shutdown_lock = threading.Lock() + + self._thread = threading.Thread(target=self._process_queue, daemon=True) + self._thread.start() + + atexit.register(self.shutdown) + + def __call__(self, name, doc): + if self._stop_event.is_set(): + raise RuntimeError("Cannot accept new data after shutdown.") + try: + self._queue.put((name, doc)) + except Full as e: + logger.exception( + f"The buffer is full. The {self._wrapped_callback.__class__.__name__} can not keep up " + f"with the incoming data. " + f"Consider increasing the queue size or optimizing the callback processing: {e}" + ) + raise RuntimeError( + f"The buffer is full. The {self._wrapped_callback.__class__.__name__} can not keep up " + f"with the incoming data. " + "Consider increasing the queue size or optimizing the callback processing." + ) from e + except Exception as e: + logger.exception(f"Failed to put document {name} in queue: {e}") + raise RuntimeError(f"Failed to put document {name} in queue: {e}") from e + + def _process_queue(self): + while True: + try: + item = self._queue.get(timeout=1) + except Empty: + if self._stop_event.is_set(): + break + continue + try: + if item is None: + break # Received sentinel value to stop processing + self._wrapped_callback(*item) # Delegate to wrapped callback + except Exception as e: + logger.exception( + f"Exception in {self._wrapped_callback.__class__.__name__}: {e}" + ) + finally: + # Account for every retrieved item (including the sentinel and + # items whose processing raised), so that `flush` can rely on + # the queue's unfinished-task counter. + self._queue.task_done() + + def flush(self, timeout: "float | None" = None) -> bool: + """Block until all documents enqueued so far have been processed. + + This provides a synchronization barrier for read-after-write patterns: + a plan or script that reads a run back (e.g. ``db[uid]``) immediately + after ``RE(plan)`` returns must call ``flush()`` first, because the + wrapped callback processes documents asynchronously and may not have + written the run yet by the time the plan finishes. + + Documents enqueued by another thread *while* flushing extend the wait: + the method returns once the queue has no unfinished items. + + Parameters + ---------- + timeout : float, optional + Maximum time to wait, in seconds. By default, wait indefinitely. + + Returns + ------- + bool + True if the buffer was fully drained; False if the timeout expired + (or the worker thread is no longer running) with items remaining. + """ + if timeout is not None: + deadline = time.monotonic() + timeout + with self._queue.all_tasks_done: + while self._queue.unfinished_tasks: + if not self._thread.is_alive(): + # The worker has exited (e.g. after shutdown); any items + # still in the queue will never be processed. + return not self._queue.unfinished_tasks + if timeout is None: + # Wake up periodically to recheck worker liveness. + self._queue.all_tasks_done.wait(timeout=1) + else: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._queue.all_tasks_done.wait(timeout=min(remaining, 1)) + return True + + def shutdown(self, wait: bool = True): + if self._stop_event.is_set(): + return + self._stop_event.set() + self._queue.put(None) + + atexit.unregister(self.shutdown) + + if wait: + self._thread.join() + + logger.info( + f"{self._wrapped_callback.__class__.__name__} shut down gracefully." + ) diff --git a/src/bluesky_tiled_plugins/writing/tiled_writer.py b/src/bluesky_tiled_plugins/writing/tiled_writer.py index a95c63e..be50eb0 100644 --- a/src/bluesky_tiled_plugins/writing/tiled_writer.py +++ b/src/bluesky_tiled_plugins/writing/tiled_writer.py @@ -9,7 +9,7 @@ from typing import Any, Optional, cast import warnings -import httpx +import httpx2 import numpy import pyarrow from event_model import ( @@ -129,9 +129,7 @@ def concatenate_stream_datums(*docs: StreamDatum): "All StreamDatum documents must reference the same stream_resource." ) docs = tuple(sorted(docs, key=lambda doc: doc["indices"]["start"])) - for d1, d2 in zip( - docs[:-1], docs[1:] - ): # TODO: use itertools.pairwise(docs) in python 3.10+ + for d1, d2 in itertools.pairwise(docs): if d1["indices"]["stop"] != d2["indices"]["start"]: raise ValueError("StreamDatum documents must be consecutive.") @@ -908,9 +906,9 @@ def _append_with_retry(self, df_client: DataFrameClient, table, desc_name: str): try: df_client.append_partition(0, table) return - except (httpx.TransportError, httpx.HTTPStatusError) as err: + except (httpx2.TransportError, httpx2.HTTPStatusError) as err: if ( - isinstance(err, httpx.HTTPStatusError) + isinstance(err, httpx2.HTTPStatusError) and err.response.status_code < 500 and err.response.status_code != 429 ): diff --git a/tests/test_buffering.py b/tests/test_buffering.py new file mode 100644 index 0000000..3ec12bd --- /dev/null +++ b/tests/test_buffering.py @@ -0,0 +1,161 @@ +"""Tests for the vendored `BufferingWrapper` (see writing/_buffer.py). + +Kept in sync with the upstream tests in bluesky (`tests/test_buffering.py`), +minus the bluesky-specific fixtures. +""" + +import threading +import time + +import pytest + +from bluesky_tiled_plugins import BufferingWrapper + + +class SlowDummyCallback: + """Simulates a slow (blocking) callback for testing threaded wrappers.""" + + def __init__(self, delay=0.1): + self.delay = delay + self.called = [] + + def __call__(self, name, doc): + time.sleep(self.delay) + self.called.append((name, doc)) + + +class CallbackWithException: + """A callback that raises an exception when called.""" + + def __call__(self, name, doc): + raise RuntimeError("This callback always raises an exception.") + + +class BlockingCallback: + """A callback that blocks until released, for exercising flush timeouts.""" + + def __init__(self): + self.release = threading.Event() + self.called = [] + + def __call__(self, name, doc): + self.release.wait() + self.called.append((name, doc)) + + +@pytest.fixture +def make_wrapper(): + """Factory for `BufferingWrapper` instances that are shut down on teardown.""" + created: list[BufferingWrapper] = [] + + def factory(*args, **kwargs): + wrapper = BufferingWrapper(*args, **kwargs) + created.append(wrapper) + return wrapper + + yield factory + + for wrapper in created: + wrapper.shutdown() + + +def test_calls_are_delegated(make_wrapper): + cb = SlowDummyCallback(delay=0) + buff_cb = make_wrapper(cb) + buff_cb("start", {"uid": "abc"}) + assert buff_cb.flush(timeout=5) is True + assert cb.called == [("start", {"uid": "abc"})] + + +def test_shutdown_refuses_new_documents(make_wrapper): + cb = SlowDummyCallback(delay=0) + buff_cb = make_wrapper(cb) + buff_cb("one", {}) + buff_cb.shutdown(wait=True) + with pytest.raises(RuntimeError): + buff_cb("two", {}) + assert ("one", {}) in cb.called + assert ("two", {}) not in cb.called + + +def test_exception_in_callback_does_not_crash_worker(make_wrapper): + buff_cb = make_wrapper(CallbackWithException()) + buff_cb("event", {}) + assert buff_cb.flush(timeout=5) is True + assert buff_cb._thread.is_alive() + buff_cb("event", {}) # still accepting documents + + +def test_flush_waits_for_processing(make_wrapper): + """flush() returns only after all enqueued documents are processed.""" + cb = SlowDummyCallback(delay=0.02) + buff_cb = make_wrapper(cb) + + for i in range(10): + buff_cb("event", {"seq_num": i}) + + assert buff_cb.flush() is True + assert len(cb.called) == 10 # all documents delivered, in order + assert [doc["seq_num"] for _, doc in cb.called] == list(range(10)) + + +def test_flush_empty_buffer_returns_immediately(make_wrapper): + """flush() on an empty buffer is an idempotent no-op, however many times + it is called, with or without a timeout.""" + buff_cb = make_wrapper(SlowDummyCallback(delay=0)) + for _ in range(3): + assert buff_cb.flush(timeout=1) is True + assert buff_cb.flush() is True + assert buff_cb._thread.is_alive() + + +def test_flush_timeout(make_wrapper): + """flush() returns False when the buffer can not drain within the timeout.""" + cb = BlockingCallback() + buff_cb = make_wrapper(cb) + + buff_cb("event", {"seq_num": 0}) + assert buff_cb.flush(timeout=0.2) is False # worker is blocked + + cb.release.set() + assert buff_cb.flush(timeout=5) is True + assert len(cb.called) == 1 + + +def test_flush_with_raising_callback(make_wrapper): + """Documents whose processing raises still count as processed: flush + must not hang on them.""" + buff_cb = make_wrapper(CallbackWithException()) + for _ in range(5): + buff_cb("event", {}) + assert buff_cb.flush(timeout=5) is True + assert buff_cb._thread.is_alive() + + +def test_flush_after_shutdown(make_wrapper): + """After a clean shutdown everything has been processed; flush returns True.""" + cb = SlowDummyCallback(delay=0.01) + buff_cb = make_wrapper(cb) + for i in range(5): + buff_cb("event", {"seq_num": i}) + buff_cb.shutdown(wait=True) + + assert buff_cb.flush(timeout=1) is True + assert buff_cb.flush(timeout=1) is True # repeated calls remain safe + assert len(cb.called) == 5 + + +def test_read_after_write_pattern(make_wrapper): + """The canonical use: write documents, flush, then read back.""" + store = {} + + def writer(name, doc): + if name == "stop": + store[doc["run_start"]] = "complete" + + buff_cb = make_wrapper(writer) + buff_cb("start", {"uid": "abc"}) + buff_cb("stop", {"run_start": "abc"}) + + assert buff_cb.flush(timeout=5) is True + assert store.get("abc") == "complete" # guaranteed visible after flush diff --git a/tests/test_tiled_writer.py b/tests/test_tiled_writer.py index c2def29..3b53c24 100644 --- a/tests/test_tiled_writer.py +++ b/tests/test_tiled_writer.py @@ -6,7 +6,7 @@ from typing import cast from urllib.parse import parse_qs, urlparse -import httpx +import httpx2 from bluesky import RunEngine import bluesky.plan_stubs as bps @@ -1013,7 +1013,7 @@ def patched_event(name, doc): assert lines[6]["name"] == "stop" -class _LoseAppendResponseOnce(httpx.BaseTransport): +class _LoseAppendResponseOnce(httpx2.BaseTransport): """Forward the first PATCH /table/partition request to the wrapped transport -- so the server fully processes the append -- but raise ReadTimeout to the caller, simulating a lost response for an append that actually succeeded.""" @@ -1030,7 +1030,7 @@ def handle_request(self, request): and "/table/partition/" in str(request.url) ): self.tripped = True - raise httpx.ReadTimeout("Simulated lost response", request=request) + raise httpx2.ReadTimeout("Simulated lost response", request=request) return response