diff --git a/docs/tutorials.md b/docs/tutorials.md index c921046..2cfe950 100644 --- a/docs/tutorials.md +++ b/docs/tutorials.md @@ -10,4 +10,5 @@ tutorials/installation tutorials/saving tutorials/search tutorials/streaming +tutorials/live-visualization ``` diff --git a/docs/tutorials/live-visualization.md b/docs/tutorials/live-visualization.md new file mode 100644 index 0000000..14e5294 --- /dev/null +++ b/docs/tutorials/live-visualization.md @@ -0,0 +1,118 @@ +# Visualize Live Bluesky Runs + +`TiledBestEffortStreamer` reproduces the output of +`bluesky.callbacks.best_effort.BestEffortCallback` — heading, live table, +baseline readings, live plots, and peak statistics — from Tiled live streams +instead of a RunEngine subscription. It works wherever the catalog is reachable: +in the acquisition process, or in a separate UI process while plans run under a +Queue Server. + +```python +from bluesky_tiled_plugins.streaming import TiledBestEffortStreamer + +with TiledBestEffortStreamer(container) as bec: + bec.run_forever() +``` + +`container` is the direct parent `Container` of the `BlueskyRun` nodes that a +`TiledWriter` writes into, as in the [streaming tutorial](./streaming.md). Set +`batch_size=1` on the writer so each Event is delivered as it is taken, and +enable the server's streaming cache per the +[deployment guide](../how-to/deploy.md). + +Live plots require matplotlib, available through the `plotting` extra: + +``` +pip install "bluesky-tiled-plugins[plotting]" +``` + +Without it, plots are disabled with a warning, but the text output still works. + +## Render on the right thread + +Tiled delivers updates on worker threads, so the streamer only queues them; +nothing is rendered until the thread that owns matplotlib asks for it. Use +whichever fits the application: + +- `run_forever(interval=0.1)` blocks, rendering updates and keeping the figures + responsive until `KeyboardInterrupt`. Use this in a dedicated console. +- `process_events(timeout=0.0)` renders everything queued and returns. Call it + periodically from an existing event loop or GUI timer, or after `RE(plan)` + when running in-process: + +```python +from bluesky_tiled_plugins import TiledWriter + +RE.subscribe(TiledWriter(client, batch_size=1)) +bec = TiledBestEffortStreamer(client).start() + +RE(scan([det], motor, -1, 1, 5)) +bec.process_events(timeout=1.0) +``` + +## BEC-compatible behavior + +The familiar toggles and results work the same way: + +```python +bec.disable_table() +bec.enable_plots() +bec.noplot_streams.append("monitor") + +bec.peaks["cen"]["det"] # like BestEffortCallback.peaks +bec.peaks.fwhm["det"] +``` + +Peak statistics are computed per hinted field when a run's table closes, from +the same data as the line plots. Hinted columns and plot dimensions follow the +Start document's `hints` and each stream's descriptor, with the same fallbacks +as the BEC (plan `motors`, then `time`). + +## Runs are rendered one at a time + +To keep the table output readable, the streamer renders a single run at a time; +updates for later runs wait their turn in arrival order. Two timing parameters +govern the transitions: + +- After the Stop document for a run arrives, the streamer waits up to + `stop_timeout` seconds (default 2) for the remaining table rows, compared + against `stop["num_events"]`. Streams holding only external array data are + recognized and not waited for. +- A run that never receives a Stop (for example, a failed run that the writer + backed up without one) does not block forever: once the next run has been + waiting `stop_timeout`, the streamer notes the missing Stop and moves on. + +By default, only runs started after the streamer connects are shown. Pass +`replay=True` to also render earlier runs still retained by the server's +streaming cache. + +## Build your own consumer + +`subscribe_to_hinted_streams` is the layer underneath: it subscribes to every +event stream of accepted runs and delivers table updates annotated with the +run's Start document, the stream's descriptor, and the resolved hinted fields — +everything needed to drive a custom table or plot widget. + +```python +from bluesky_tiled_plugins.streaming import HintedStreamUpdate, subscribe_to_hinted_streams + + +def on_update(update: HintedStreamUpdate) -> None: + print(update.run_uid, update.stream_name, update.fields) + print(update.data()) # seq_num, time, and the hinted columns + + +with subscribe_to_hinted_streams(container, on_update): + ... +``` + +`update.table` keeps every column `TiledWriter` wrote; `update.data()` narrows +it to `seq_num`, `time`, and the hinted fields. The callback runs on Tiled +worker threads, possibly concurrently — synchronize any shared state, or queue +updates for a GUI thread as `TiledBestEffortStreamer` does. Start documents and +descriptors are cached for recent runs (scaled by `retained_runs`); evicted +streams are refetched transparently. + +For selecting specific streams and data keys, raw native updates, and run +filtering, use `subscribe_to_streams` directly — see the +[streaming tutorial](./streaming.md). diff --git a/pyproject.toml b/pyproject.toml index 106307b..50578c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,9 +54,12 @@ ragged = [ "awkward >=2.4.3", "ragged", ] +plotting = [ + "matplotlib", +] # All optional features. all = [ - "bluesky-tiled-plugins[ragged]", + "bluesky-tiled-plugins[ragged, plotting]", ] @@ -78,6 +81,7 @@ test = [ "fakeredis", "h5py", "jinja2", + "matplotlib", "mongomock", "ophyd", "pre-commit", @@ -138,6 +142,12 @@ filterwarnings = [ # 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", + # Newer Alembic warns when alembic.ini lacks `path_separator`; the config is + # shipped by tiled (third-party issue). + "ignore:No path_separator found in configuration:DeprecationWarning", + # websockets 17.1 deprecates non-context-manager `connect()`, which tiled's + # subscription client still uses (third-party issue). + "ignore:connect\\(\\) must be used as a context manager:DeprecationWarning", ] log_level = "INFO" testpaths = ["tests"] diff --git a/src/bluesky_tiled_plugins/streaming/__init__.py b/src/bluesky_tiled_plugins/streaming/__init__.py new file mode 100644 index 0000000..b84ff75 --- /dev/null +++ b/src/bluesky_tiled_plugins/streaming/__init__.py @@ -0,0 +1,17 @@ +from .streaming import ( + BlueskyStreamUpdate, + BlueskyStreamSubscription, + subscribe_to_streams, +) +from .hints import HintedStreamUpdate, subscribe_to_hinted_streams +from .best_effort import TiledBestEffortStreamer + + +__all__ = [ + "BlueskyStreamUpdate", + "BlueskyStreamSubscription", + "subscribe_to_streams", + "HintedStreamUpdate", + "subscribe_to_hinted_streams", + "TiledBestEffortStreamer", +] diff --git a/src/bluesky_tiled_plugins/streaming/best_effort.py b/src/bluesky_tiled_plugins/streaming/best_effort.py new file mode 100644 index 0000000..8b6c01e --- /dev/null +++ b/src/bluesky_tiled_plugins/streaming/best_effort.py @@ -0,0 +1,603 @@ +"""A best-effort callback that consumes Tiled live streams instead of Bluesky documents.""" + +from __future__ import annotations + +import logging +import queue +import time +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any + +import numpy as np +import pandas as pd +from tiled.client.container import Container +from tiled.client.stream import LiveChildCreated + +from .hints import SEQ_NUM as _SEQ_NUM +from .hints import TIME as _TIME +from .hints import HintedStreamUpdate, _dimensions, subscribe_to_hinted_streams + +logger = logging.getLogger(__name__) + + +def peak_stats(x: np.ndarray, y: np.ndarray) -> dict[str, Any]: + """Coarse peak statistics (max, min, com, cen, fwhm) of ``y`` versus ``x``.""" + i_max, i_min = int(np.argmax(y)), int(np.argmin(y)) + total = float(np.sum(y)) + half = y[i_min] + (y[i_max] - y[i_min]) / 2 + above = np.flatnonzero(y >= half) + left, right = x[above[0]], x[above[-1]] + return { + "max": (x[i_max], y[i_max]), + "min": (x[i_min], y[i_min]), + "com": float(np.sum(x * y) / total) if total else np.nan, + "cen": (left + right) / 2, + "fwhm": abs(right - left), + } + + +class PeakResults: + """Peak statistics keyed like ``BestEffortCallback.peaks``. + + Each statistic is a dict keyed by field name, reachable both as an + attribute (``peaks.cen["det"]``) and by subscription + (``peaks["cen"]["det"]``), mirroring + ``bluesky.callbacks.best_effort.PeakResults``. + """ + + ATTRS = ("com", "cen", "max", "min", "fwhm") + + def __init__(self) -> None: + for attr in self.ATTRS: + setattr(self, attr, {}) + + def clear(self) -> None: + for attr in self.ATTRS: + getattr(self, attr).clear() + + def __getitem__(self, key: str) -> dict[str, Any]: + if key not in self.ATTRS: + raise KeyError(f"Keys are: {self.ATTRS}") + return getattr(self, key) + + def __repr__(self) -> str: + lines = [] + for attr in self.ATTRS: + lines.append(f"{attr}:") + lines.extend( + f" {field}: {value}" for field, value in getattr(self, attr).items() + ) + return "\n".join(lines) + + +class _TablePrinter: + """Text table laid out like ``bluesky.callbacks.LiveTable``.""" + + def __init__( + self, + columns: list[str], + data_keys: dict[str, Any], + out: Callable[[str], None], + min_width: int = 12, + ): + self.columns = columns + self.data_keys = data_keys + self.out = out + # LiveTable: seq_num is 10 wide with a trailing pad; other columns are padded on both sides. + self.widths = [10] + [max(min_width - 2, len(c)) for c in (_TIME, *columns)] + self.border = ( + "+" + "-" * 11 + "".join("+" + "-" * (w + 2) for w in self.widths[1:]) + "+" + ) + + def _line(self, cells: Iterable[str]) -> str: + first, *rest = cells + return f"|{first:>{self.widths[0]}} |" + "".join( + f" {c:>{w}} |" for c, w in zip(rest, self.widths[1:]) + ) + + def header(self) -> None: + self.out(self.border) + self.out(self._line((_SEQ_NUM, _TIME, *self.columns))) + self.out(self.border) + + def row(self, row: pd.Series) -> None: + when = datetime.fromtimestamp(row[_TIME]).strftime("%H:%M:%S.%f")[:10] + cells = [ + str(int(row[_SEQ_NUM])), + when, + *(self._format(c, row.get(c)) for c in self.columns), + ] + self.out(self._line(cells)) + + def footer(self) -> None: + self.out(self.border) + + def _format(self, key: str, value: Any) -> str: + if value is None or (isinstance(value, float) and np.isnan(value)): + return "" + if isinstance(value, (float, np.floating)): + precision = self.data_keys.get(key, {}).get("precision") + return f"{value:.{3 if precision is None else int(precision)}f}" + return str(value) + + +class _LinePlotter: + """1D live plot of several y fields against one x field, like ``LivePlot``.""" + + def __init__( + self, fig_name: str, x_key: str, y_keys: list[str], label: str, x_offset: float + ): + import matplotlib.pyplot as plt + + self.x_key, self.y_keys, self.x_offset = x_key, y_keys, x_offset + self.fig = plt.figure(fig_name) + # Overplot on an existing figure with the same layout, as BEC does by default. + if len(self.fig.axes) != len(y_keys): + self.fig.clf() + self.fig.subplots(len(y_keys), 1, sharex=True, squeeze=False) + self.axes = self.fig.axes + self.x: list[float] = [] + self.y: dict[str, list[float]] = {k: [] for k in y_keys} + self.lines = {} + for ax, key in zip(self.axes, y_keys): + ax.set_ylabel(key) + (self.lines[key],) = ax.plot([], [], "o-", label=label) + ax.legend(loc="best") + self.axes[-1].set_xlabel(x_key) + + def update(self, df: pd.DataFrame) -> None: + self.x.extend((df[self.x_key] - self.x_offset).tolist()) + for key in self.y_keys: + self.y[key].extend(df[key].tolist()) + self.lines[key].set_data(self.x, self.y[key]) + for ax in self.axes: + ax.relim() + ax.autoscale_view() + self.fig.canvas.draw_idle() + + def peaks(self) -> dict[str, dict[str, Any]]: + x = np.asarray(self.x, dtype=float) + return { + key: peak_stats(x, np.asarray(self.y[key], dtype=float)) + for key in self.y_keys + if self.y[key] + } + + +class _ScatterPlotter: + """2D live scatter of y fields colored by value, like ``LiveScatter``.""" + + def __init__( + self, fig_name: str, x_key: str, y_key: str, z_keys: list[str], label: str + ): + import matplotlib.pyplot as plt + + self.x_key, self.y_key, self.z_keys = x_key, y_key, z_keys + self.fig = plt.figure(fig_name) + self.fig.clf() + axes = self.fig.subplots(1, len(z_keys), squeeze=False)[0] + self.xy: list[tuple[float, float]] = [] + self.z: dict[str, list[float]] = {k: [] for k in z_keys} + self.collections = {} + for ax, key in zip(axes, z_keys): + self.collections[key] = ax.scatter([], [], c=[], cmap="viridis") + ax.set(title=f"{key} [{label}]", xlabel=x_key, ylabel=y_key) + self.fig.colorbar(self.collections[key], ax=ax) + + def update(self, df: pd.DataFrame) -> None: + self.xy.extend(zip(df[self.x_key].tolist(), df[self.y_key].tolist())) + offsets = np.asarray(self.xy, dtype=float) + for key, collection in self.collections.items(): + self.z[key].extend(df[key].tolist()) + values = np.asarray(self.z[key], dtype=float) + collection.set_offsets(offsets) + collection.set_array(values) + collection.set_clim(np.nanmin(values), np.nanmax(values)) + collection.axes.update_datalim(offsets) + collection.axes.autoscale_view() + self.fig.canvas.draw_idle() + + +@dataclass +class _StreamState: + name: str + descriptor: dict[str, Any] + rows: int = 0 + table: _TablePrinter | None = None + plotter: _LinePlotter | _ScatterPlotter | None = None + + +@dataclass +class _RunState: + uid: str + start: dict[str, Any] + streams: dict[str, _StreamState] = field(default_factory=dict) + stop: dict[str, Any] | None = None + stop_received: float = 0.0 + baseline_readings: int = 0 + # BEC holds the end-of-run baseline until the run's table is closed. + end_baseline: list[str] = field(default_factory=list) + # Whether a stream that has not produced table updates is expected to + # (True), holds only external array data (False), or is unknown yet. + expects_table: dict[str, bool] = field(default_factory=dict) + + +class TiledBestEffortStreamer: + """Best-effort visualization of Bluesky runs streamed live from Tiled. + + Unlike ``bluesky.callbacks.best_effort.BestEffortCallback``, this is not + subscribed to a RunEngine. It subscribes to a Tiled container that a + ``TiledWriter`` writes into and reconstructs the heading, live table, + live plots, baseline readings, and peak statistics from: + + * the run node creation (``metadata["start"]``), + * each event stream node's metadata (the descriptor) and its live table data, + * the run node metadata update that adds ``metadata["stop"]``. + + Tiled invokes callbacks on worker threads, so updates are queued and only + rendered when :meth:`process_events` (or :meth:`run_forever`) is called, + which must happen on the thread that owns the matplotlib GUI. + + Parameters + ---------- + container : tiled.client.container.Container + Direct parent of the ``BlueskyRun`` nodes created by ``TiledWriter``. + noplot_streams : iterable of str + Streams that get neither a live table nor a plot. + replay : bool + Also show runs whose start time precedes this callback's connection and + that are still retained in Tiled's streaming cache. + stop_timeout : float + Seconds to wait after the stop arrives for the remaining stream data + before finishing the run. + retained_runs : int + Stopped runs whose start and descriptors stay cached; see + ``subscribe_to_hinted_streams``. + out : callable + Receives each line of text output. Defaults to ``print``. + """ + + def __init__( + self, + container: Container, + *, + noplot_streams: Iterable[str] = ("baseline",), + replay: bool = False, + stop_timeout: float = 2.0, + retained_runs: int = 4, + max_size: int = 1_000_000, + out: Callable[[str], None] = print, + ): + self._container = container + self.noplot_streams = list(noplot_streams) + self._replay = replay + self._stop_timeout = stop_timeout + self._retained_runs = retained_runs + self._max_size = max_size + self._out = out + self._queue: queue.Queue[tuple[Any, ...]] = queue.Queue() + # Runs are rendered one at a time; updates for later runs wait in _deferred. + self._run: _RunState | None = None + self._deferred: dict[str, list[tuple[Any, ...]]] = {} + self._deferred_start_at: dict[str, float] = {} + self._data_subscription = None + self._connected_at = 0.0 + self._heading_enabled = True + self._table_enabled = True + self._baseline_enabled = True + self._plots_enabled = True + self.peaks = PeakResults() + + # -- BEC-compatible toggles ------------------------------------------------ + + def enable_heading(self) -> None: + self._heading_enabled = True + + def disable_heading(self) -> None: + self._heading_enabled = False + + def enable_table(self) -> None: + self._table_enabled = True + + def disable_table(self) -> None: + self._table_enabled = False + + def enable_baseline(self) -> None: + self._baseline_enabled = True + + def disable_baseline(self) -> None: + self._baseline_enabled = False + + def enable_plots(self) -> None: + self._plots_enabled = True + + def disable_plots(self) -> None: + self._plots_enabled = False + + # -- Subscription lifecycle ------------------------------------------------ + + def start(self) -> TiledBestEffortStreamer: + """Connect the Tiled websocket subscriptions.""" + self._connected_at = time.time() + self._data_subscription = subscribe_to_hinted_streams( + self._container, + self._on_data, + run_filter=self._on_run_created, + on_run_stop=self._on_run_stop, + retained_runs=self._retained_runs, + max_size=self._max_size, + ) + return self + + def disconnect(self) -> None: + """Disconnect all Tiled subscriptions.""" + if self._data_subscription is not None: + self._data_subscription.disconnect() + self._data_subscription = None + + def __enter__(self) -> TiledBestEffortStreamer: + return self.start() + + def __exit__(self, *args: object) -> None: + self.disconnect() + + # -- Tiled worker-thread callbacks: only enqueue ----------------------------- + + def _on_run_created(self, update: LiveChildCreated) -> bool: + start = dict(update.metadata.get("start", {})) + if not self._replay and start.get("time", 0.0) < self._connected_at: + return False + # Enqueued before subscribe_to_streams subscribes to the run's children, + # so the start always precedes the run's data in the queue. + self._queue.put(("start", update.key, start)) + return True + + def _on_data(self, update: HintedStreamUpdate) -> None: + self._queue.put(("data", update)) + + def _on_run_stop(self, run_uid: str, stop: Mapping[str, Any]) -> None: + self._queue.put(("stop", run_uid, dict(stop))) + + # -- Main-thread processing ---------------------------------------------------- + + def process_events(self, timeout: float = 0.0) -> None: + """Render queued Tiled updates, waiting up to ``timeout`` seconds for more.""" + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + try: + event = ( + self._queue.get(timeout=remaining) + if remaining > 0 + else self._queue.get_nowait() + ) + except queue.Empty: + break + self._dispatch(event) + self._finish_stopped_run() + + def _dispatch(self, event: tuple[Any, ...]) -> None: + kind = event[0] + run_uid = event[1].run_uid if kind == "data" else event[1] + if self._run is not None and run_uid == self._run.uid: + if kind != "start": # a replayed start of the active run is a no-op + self._handle(event) + elif self._run is None and kind == "start" and not self._deferred: + self._handle(event) + elif kind == "start": + # Runs are rendered one at a time, in arrival order; this run + # waits for its turn. + self._deferred.setdefault(run_uid, []).append(event) + self._deferred_start_at.setdefault(run_uid, time.monotonic()) + elif run_uid in self._deferred: + self._deferred[run_uid].append(event) + # Updates of runs with no start here were filtered out; drop them. + + def _handle(self, event: tuple[Any, ...]) -> None: + try: + getattr(self, f"_handle_{event[0]}")(*event[1:]) + except Exception: + logger.exception("Failed to process Tiled update %r", event[0]) + + def run_forever(self, interval: float = 0.1) -> None: + """Process updates and keep matplotlib figures responsive until interrupted.""" + try: + import matplotlib.pyplot as plt + except ImportError: + plt = None + else: + plt.ion() + try: + while True: + self.process_events(timeout=interval) + if plt is not None: + for num in plt.get_fignums(): + plt.figure(num).canvas.flush_events() + except KeyboardInterrupt: + pass + + def _handle_start(self, run_uid: str, start: dict[str, Any]) -> None: + self._run = _RunState(run_uid, start) + if self._heading_enabled: + when = datetime.fromtimestamp(start.get("time", time.time())) + self._out("") + self._out( + f"Transient Scan ID: {start.get('scan_id')} Time: {when:%Y-%m-%d %H:%M:%S}" + ) + self._out(f"Persistent Unique Scan ID: '{run_uid}'") + + def _handle_data(self, update: HintedStreamUpdate) -> None: + run = self._run + stream = run.streams.get(update.stream_name) + if stream is None: + stream = run.streams[update.stream_name] = self._new_stream(run, update) + + df = update.table + stream.rows += len(df) + if stream.name == "baseline": + self._print_baseline(run, stream, df) + if stream.table is not None: + for _, row in df.iterrows(): + stream.table.row(row) + if stream.plotter is not None: + stream.plotter.update(df) + + def _handle_stop(self, run_uid: str, stop: dict[str, Any]) -> None: + self._run.stop, self._run.stop_received = stop, time.monotonic() + + def _new_stream(self, run: _RunState, update: HintedStreamUpdate) -> _StreamState: + stream_name = update.stream_name + stream = _StreamState(stream_name, update.descriptor) + if self._heading_enabled: + self._out(f"New stream: {stream_name!r}") + if stream_name in self.noplot_streams: + return stream + + data_keys = update.descriptor.get("data_keys", {}) + columns = list(update.fields) + dim_fields = [ + fields[0] for fields, name in _dimensions(run.start) if name == stream_name + ] + + if self._table_enabled: + stream.table = _TablePrinter(columns, data_keys, self._out) + stream.table.header() + if self._plots_enabled and dim_fields: + stream.plotter = self._make_plotter( + run, dim_fields, [c for c in columns if c not in dim_fields] + ) + return stream + + def _make_plotter(self, run: _RunState, dim_fields: list[str], y_keys: list[str]): + if not y_keys: + return None + try: + import matplotlib.pyplot as plt + except ImportError: + logger.warning( + "matplotlib is not installed, so live plots are disabled; install " + 'the plotting extra: pip install "bluesky-tiled-plugins[plotting]"' + ) + self.disable_plots() + return None + + plt.ion() + label = str(run.start.get("scan_id", run.uid[:8])) + if len(dim_fields) == 1: + x_key = dim_fields[0] + x_offset = run.start.get("time", 0.0) if x_key == _TIME else 0.0 + return _LinePlotter( + f"{' '.join(sorted(y_keys))} vs {x_key}", x_key, y_keys, label, x_offset + ) + if len(dim_fields) == 2: + x_key, y_key = dim_fields + return _ScatterPlotter( + f"{' '.join(sorted(y_keys))} vs {x_key} {y_key}", + x_key, + y_key, + y_keys, + label, + ) + return None + + def _print_baseline( + self, run: _RunState, stream: _StreamState, df: pd.DataFrame + ) -> None: + if not self._baseline_enabled: + return + data_keys = stream.descriptor.get("data_keys", {}) + keys = [k for k in df.columns if k in data_keys] + border = "+" + "-" * 32 + "+" + "-" * 32 + "+" + for _, row in df.iterrows(): + run.baseline_readings += 1 + if run.baseline_readings == 1: + lines, out = ["Start-of-run baseline readings:"], self._out + else: + run.end_baseline.clear() + lines, out = ["End-of-run baseline readings:"], run.end_baseline.append + lines.append(border) + lines.extend(f"| {key:>30} | {row[key]!s:<30} |" for key in keys) + lines.append(border) + for line in lines: + out(line) + + def _finish_stopped_run(self) -> None: + while (run := self._run) is not None: + now = time.monotonic() + if run.stop is None: + # A failed run may never receive a stop document (for example, + # TiledWriter backs it up without writing one). Once a newer + # run has been waiting for stop_timeout, move on -- like the + # BEC, where a new start document simply begins a new table. + waiting_since = min(self._deferred_start_at.values(), default=now) + if now - waiting_since < self._stop_timeout: + return + self._out( + f"Run {run.uid[:8]} received no stop document; moving on to the next run." + ) + elif ( + not self._run_complete(run) + and now - run.stop_received < self._stop_timeout + ): + return + self._finish_run(run) + self._run = None + if self._deferred: + next_uid = next(iter(self._deferred)) + self._deferred_start_at.pop(next_uid, None) + for event in self._deferred.pop(next_uid): + self._dispatch(event) + + def _run_complete(self, run: _RunState) -> bool: + """Whether every table-producing stream has delivered all of its rows.""" + for name, n in (run.stop.get("num_events", {}) or {}).items(): + if not n: + continue + stream = run.streams.get(name) + if stream is not None: + if stream.rows < n: + return False + elif self._stream_expects_table(run, name) is not False: + # An unseen stream blocks until it is known to hold only + # external (array) data, which never produces table updates. + return False + return True + + def _stream_expects_table(self, run: _RunState, stream_name: str) -> bool | None: + """Whether an unseen stream will produce table updates; None if unknown yet.""" + if stream_name not in run.expects_table: + try: + node = self._container[run.uid] + node = getattr(node, "base", node) + data_keys = dict(node[stream_name].metadata).get("data_keys", {}) + except Exception: + return None # not in Tiled yet; its data may still arrive + run.expects_table[stream_name] = any( + "external" not in spec for spec in data_keys.values() + ) + return run.expects_table[stream_name] + + def _finish_run(self, run: _RunState) -> None: + for stream in run.streams.values(): + if stream.table is not None: + stream.table.footer() + if isinstance(stream.plotter, _LinePlotter): + for fld, stats in stream.plotter.peaks().items(): + for attr in PeakResults.ATTRS: + self.peaks[attr][fld] = stats[attr] + if self._table_enabled: + plan = run.start.get("plan_name", "") + self._out( + f"generator {plan} ['{run.uid[:8]}'] (scan num: {run.start.get('scan_id')})" + ) + for line in run.end_baseline: + self._out(line) + if run.stop is None: + return + status = run.stop.get("exit_status") + if status != "success": + self._out( + f"Run {run.uid[:8]} exited with status {status!r}: {run.stop.get('reason', '')}" + ) diff --git a/src/bluesky_tiled_plugins/streaming/hints.py b/src/bluesky_tiled_plugins/streaming/hints.py new file mode 100644 index 0000000..3b77dde --- /dev/null +++ b/src/bluesky_tiled_plugins/streaming/hints.py @@ -0,0 +1,176 @@ +"""Subscribe to the hinted fields of live Bluesky event streams in Tiled.""" + +from __future__ import annotations + +import functools +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any + +import pandas as pd +from .streaming import ( + BlueskyStreamSubscription, + BlueskyStreamUpdate, + subscribe_to_streams, +) +from tiled.client.container import Container +from tiled.client.stream import LiveChildCreated, LiveTableData + +# Columns that TiledWriter adds to every internal table row. +SEQ_NUM = "seq_num" +TIME = "time" + + +def _hinted_fields(descriptor: dict[str, Any]) -> list[str]: + """Return the hinted data keys of a stream, mirroring ``bluesky.callbacks.best_effort``. + + TiledWriter drops ``object_keys`` from the stream metadata but records each + key's ``object_name``, so objects are reconstructed from ``data_keys``. + """ + data_keys = descriptor.get("data_keys", {}) + hints = descriptor.get("hints", {}) or {} + by_object: dict[str, list[str]] = {} + for key, spec in data_keys.items(): + by_object.setdefault(spec.get("object_name", key), []).append(key) + + fields: list[str] = [] + for obj, keys in by_object.items(): + # Like bluesky, fall back to all of the object's keys both when the + # object has no hints and when its hints have no "fields". + obj_fields = hints.get(obj, {}).get("fields", keys) + fields.extend(f for f in obj_fields if f in data_keys and f not in fields) + return fields + + +def _dimensions(start: dict[str, Any]) -> list[tuple[list[str], str]]: + """Return the plan's ``(fields, stream_name)`` dimensions with BEC's defaults.""" + dims = (start.get("hints") or {}).get("dimensions") + if dims: + return [(list(fields), stream) for fields, stream in dims] + if motors := start.get("motors"): + return [([motor], "primary") for motor in motors] + return [([TIME], "primary")] + + +def _is_scalar(spec: dict[str, Any]) -> bool: + """Whether a data key holds internal scalar readings.""" + return "external" not in spec and not spec.get("shape") + + +@dataclass(frozen=True) +class HintedStreamUpdate: + """New rows of a Bluesky event stream together with its hints. + + Parameters + ---------- + run_uid : str + Uid of the run, the key of its ``BlueskyRun`` node. + stream_name : str + Name of the event stream. + start : dict + The run's start document. + descriptor : dict + The event stream node's metadata (the descriptor as written by TiledWriter). + fields : tuple of str + Scalar dimension fields of this stream followed by its hinted fields. + table : pandas.DataFrame + All new rows, with every column TiledWriter wrote (``seq_num``, ``time``, + data keys, and ``ts_*`` timestamps). + """ + + run_uid: str + stream_name: str + start: dict[str, Any] + descriptor: dict[str, Any] + fields: tuple[str, ...] + table: pd.DataFrame + + def data(self) -> pd.DataFrame: + """Return the ``seq_num``, ``time``, and hinted columns of the new rows.""" + return self.table.loc[:, [SEQ_NUM, TIME, *self.fields]] + + +def subscribe_to_hinted_streams( + container: Container, + callback: Callable[[HintedStreamUpdate], None], + *, + run_filter: Callable[[LiveChildCreated], bool] | None = None, + on_run_stop: Callable[[str, Mapping[str, Any]], None] | None = None, + retained_runs: int = 4, + start: int | None = 0, + max_size: int = 1_000_000, +) -> BlueskyStreamSubscription: + """Subscribe to live table data of every event stream, annotated with its hints. + + Hints are only known once a run and its streams exist, so every stream is + subscribed and each stream's descriptor is fetched from Tiled on its first + update. Only TiledWriter's internal tables are delivered; array nodes are ignored. + + Parameters + ---------- + container : tiled.client.container.Container + Direct parent of the ``BlueskyRun`` nodes created by ``TiledWriter``. + callback : Callable[[HintedStreamUpdate], None] + Called on a Tiled worker thread, possibly concurrently, for each table update. + run_filter : Callable[[LiveChildCreated], bool], optional + Predicate applied to each new run. It runs before any of the run's data is delivered. + on_run_stop : Callable[[str, Mapping], None], optional + Called with the run uid and stop document of each accepted run. It may + precede the run's last data updates; see ``subscribe_to_streams``. + retained_runs : int, optional + Scales the bounded cache of start documents and descriptors (roughly + the number of recent runs whose streams stay cached). Data arriving + for an evicted stream is still delivered, but refetches the start + document and descriptor from Tiled. + start, max_size + Passed through to ``subscribe_to_streams``. + """ + if retained_runs < 0: + raise ValueError(f"retained_runs must be non-negative, not {retained_runs}.") + + # Thread-safe bounded cache; concurrent first updates of one stream may + # both fetch, which is harmless. + @functools.lru_cache(maxsize=32 * max(1, retained_runs)) + def stream_info( + run_uid: str, stream_name: str + ) -> tuple[dict[str, Any], dict[str, Any], tuple[str, ...]]: + run = container[run_uid] + run = getattr(run, "base", run) + start_doc = dict(run.metadata.get("start", {})) + descriptor = dict(run[stream_name].metadata) + data_keys = descriptor.get("data_keys", {}) + dim_fields = [f[0] for f, name in _dimensions(start_doc) if name == stream_name] + fields = tuple( + dict.fromkeys( + f + for f in (*dim_fields, *_hinted_fields(descriptor)) + if f in data_keys and _is_scalar(data_keys[f]) + ) + ) + return start_doc, descriptor, fields + + def on_update(update: BlueskyStreamUpdate) -> None: + if not isinstance(update.update, LiveTableData): + return + start_doc, descriptor, fields = stream_info(update.run_uid, update.stream_name) + callback( + HintedStreamUpdate( + run_uid=update.run_uid, + stream_name=update.stream_name, + start=start_doc, + descriptor=descriptor, + fields=fields, + # The raw update keeps seq_num and time, which BlueskyStreamUpdate.data() drops. + table=update.update.data(), + ) + ) + + return subscribe_to_streams( + container, + on_update, + streams=None, + run_filter=run_filter, + on_run_stop=on_run_stop, + start=start, + max_size=max_size, + ) diff --git a/src/bluesky_tiled_plugins/streaming.py b/src/bluesky_tiled_plugins/streaming/streaming.py similarity index 88% rename from src/bluesky_tiled_plugins/streaming.py rename to src/bluesky_tiled_plugins/streaming/streaming.py index 5f78d8d..ec88776 100644 --- a/src/bluesky_tiled_plugins/streaming.py +++ b/src/bluesky_tiled_plugins/streaming/streaming.py @@ -15,6 +15,7 @@ LiveArrayData, LiveArrayRef, LiveChildCreated, + LiveChildMetadataUpdated, LiveTableData, Subscription, TableSubscription, @@ -96,6 +97,7 @@ def __init__( start: int | None, max_size: int, run_filter: Callable[[LiveChildCreated], bool] | None, + on_run_stop: Callable[[str, Mapping[str, Any]], None] | None = None, ) -> None: """ Create and start a managed stream subscription. @@ -125,6 +127,9 @@ def __init__( ``1_000_000``. run_filter : Callable[[LiveChildCreated], bool] or None Optional raw Tiled run predicate. + on_run_stop : Callable[[str, Mapping[str, Any]], None] or None + Optional function called once with the run uid and stop document + of each accepted run. See :func:`subscribe_to_streams`. Raises ------ @@ -142,6 +147,8 @@ def __init__( self._streams = None if streams is None else dict(streams) self._callback = callback self._run_filter = run_filter + self._on_run_stop = on_run_stop + self._accepted_runs: set[str] = set() self._start = start self._max_size = max_size self._lock = threading.Lock() @@ -156,7 +163,11 @@ def __init__( established = False try: - self._subscribe_container(container, self._handle_container_child) + self._subscribe_container( + container, + self._handle_container_child, + self._handle_run_metadata_updated, + ) established = True finally: if not established: @@ -228,11 +239,15 @@ def _subscribe_container( self, node: Container, callback: Callable[[LiveChildCreated], None], + on_metadata_updated: Callable[[LiveChildMetadataUpdated], None] | None = None, ) -> None: def attach(subscription: Subscription) -> None: - cast(ContainerSubscription, subscription).child_created.add_callback( - callback - ) + container_subscription = cast(ContainerSubscription, subscription) + container_subscription.child_created.add_callback(callback) + if on_metadata_updated is not None: + container_subscription.child_metadata_updated.add_callback( + on_metadata_updated + ) self._subscribe(node, callback, attach) @@ -317,9 +332,22 @@ def _handle_container_child(self, update: LiveChildCreated) -> None: "Run filter failed for Tiled run %s at %s", run_uid, child.uri ) return + with self._lock: + self._accepted_runs.add(run_uid) callback = functools.partial(self._handle_run_child, run_uid) self._subscribe_child_container(child, callback) + def _handle_run_metadata_updated(self, update: LiveChildMetadataUpdated) -> None: + stop = update.metadata.get("stop") + if stop is None or self._on_run_stop is None: + return + with self._lock: + if self._closed or update.key not in self._accepted_runs: + return + # Forgetting the run makes the hook fire once and keeps the set bounded. + self._accepted_runs.discard(update.key) + self._on_run_stop(update.key, stop) + def _handle_run_child( self, run_uid: str, @@ -476,6 +504,7 @@ def subscribe_to_streams( metadata_filter: Callable[[Mapping[str, Any]], bool] | None = None, required_specs: str | Iterable[str] | None = None, run_filter: Callable[[LiveChildCreated], bool] | None = None, + on_run_stop: Callable[[str, Mapping[str, Any]], None] | None = None, start: int | None = 0, max_size: int = 1_000_000, ) -> BlueskyStreamSubscription: @@ -501,6 +530,11 @@ def subscribe_to_streams( One spec name or every spec name required for a run to match. run_filter : Callable[[LiveChildCreated], bool] or None, optional Predicate applied to the raw Tiled run-creation update. + on_run_stop : Callable[[str, Mapping[str, Any]], None] or None, optional + Function called as ``on_run_stop(run_uid, stop_document)`` once for each + run that passed the filters, when its stop document is recorded. Runs + rejected by the filters never trigger it, nor do later metadata updates + of the run without ``"stop"``. start : int or None, optional Tiled sequence number supplied to every owned subscription. The default, ``0``, replays records retained by the streaming cache. ``None`` receives @@ -536,6 +570,15 @@ def subscribe_to_streams( columns remain one native table update; this function does not join or align updates across nodes or streams. + ``on_run_stop`` arrives on the root container subscription, so it follows + the run's creation but has no ordering relative to the run's data, which + arrives on other subscriptions. It may run before the last data callbacks + of the same run; to wait for all data, compare against + ``stop["num_events"]`` or wait briefly. Like ``callback``, it runs on a + Tiled thread and its exceptions are not caught. With ``start=0``, a + finished run still in the streaming cache replays both its creation and + its stop. + With Tiled 0.2.18, the server can stream ragged nodes, but ``RaggedClient`` has no subscription API and the client has no ragged live schema or update model. ``BytesClient`` likewise has no subscription API, @@ -562,4 +605,5 @@ def subscribe_to_streams( start=start, max_size=max_size, run_filter=_compose_run_filter(metadata_filter, required_specs, run_filter), + on_run_stop=on_run_stop, ) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 5a7a5b1..74b80aa 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -969,3 +969,128 @@ def test_subscribe_to_streams_rejects_empty_selection( ): with pytest.raises(ValueError, match=message): subscribe_to_streams(streaming_client, lambda update: None, streams=streams) + + +def _write_run(writer, run_uid, metadata=None): + descriptor_uid = uuid.uuid4().hex + writer("start", _start_document(run_uid, metadata=metadata)) + writer( + "descriptor", + _descriptor_document( + run_uid, descriptor_uid, "primary", _scalar_data_keys("x") + ), + ) + writer("event", _event_document(descriptor_uid, 1, {"x": 1.0})) + writer("stop", _stop_document(run_uid, {"primary": 1})) + + +def test_on_run_stop_called_once_per_run(streaming_client): + run_uids = [uuid.uuid4().hex, uuid.uuid4().hex] + stops: queue.Queue[tuple[str, dict]] = queue.Queue() + writer = TiledWriter(streaming_client, batch_size=1) + + with subscribe_to_streams( + streaming_client, + lambda update: None, + streams={"primary": "x"}, + on_run_stop=lambda uid, stop: stops.put((uid, dict(stop))), + ): + for run_uid in run_uids: + _write_run(writer, run_uid) + received = [stops.get(timeout=5) for _ in run_uids] + with pytest.raises(queue.Empty): + stops.get(timeout=1) + + assert [uid for uid, _ in received] == run_uids + for run_uid, stop in received: + assert stop["run_start"] == run_uid + assert stop["exit_status"] == "success" + + +@pytest.mark.parametrize( + "filters", + [ + pytest.param( + {"metadata_filter": lambda start: start.get("group") == "accepted"}, + id="metadata_filter", + ), + pytest.param({"required_specs": "Accepted"}, id="required_specs"), + pytest.param( + { + "run_filter": lambda update: update.metadata["start"].get("group") + == "accepted" + }, + id="run_filter", + ), + ], +) +def test_on_run_stop_skips_filtered_runs(streaming_client, filters): + accepted_run_uid = uuid.uuid4().hex + rejected_run_uid = uuid.uuid4().hex + stops: queue.Queue[str] = queue.Queue() + writer = TiledWriter(streaming_client, batch_size=1) + + with subscribe_to_streams( + streaming_client, + lambda update: None, + streams={"primary": "x"}, + on_run_stop=lambda uid, stop: stops.put(uid), + **filters, + ): + _write_run(writer, rejected_run_uid, metadata={"group": "rejected"}) + _write_run( + writer, + accepted_run_uid, + metadata={"group": "accepted", "tiled_specs": ["Accepted"]}, + ) + assert stops.get(timeout=5) == accepted_run_uid + with pytest.raises(queue.Empty): + stops.get(timeout=1) + + +def test_on_run_stop_not_called_after_disconnect(streaming_client): + run_uid = uuid.uuid4().hex + descriptor_uid = uuid.uuid4().hex + stops: queue.Queue[str] = queue.Queue() + writer = TiledWriter(streaming_client, batch_size=1) + + subscription = subscribe_to_streams( + streaming_client, + lambda update: None, + streams={"primary": "x"}, + start=None, + on_run_stop=lambda uid, stop: stops.put(uid), + ) + try: + writer("start", _start_document(run_uid)) + writer( + "descriptor", + _descriptor_document( + run_uid, descriptor_uid, "primary", _scalar_data_keys("x") + ), + ) + writer("event", _event_document(descriptor_uid, 1, {"x": 1.0})) + finally: + subscription.disconnect() + writer("stop", _stop_document(run_uid, {"primary": 1})) + + with pytest.raises(queue.Empty): + stops.get(timeout=1) + + +def test_on_run_stop_ignores_later_metadata_updates(streaming_client): + run_uid = uuid.uuid4().hex + stops: queue.Queue[str] = queue.Queue() + writer = TiledWriter(streaming_client, batch_size=1) + + with subscribe_to_streams( + streaming_client, + lambda update: None, + streams={"primary": "x"}, + on_run_stop=lambda uid, stop: stops.put(uid), + ): + _write_run(writer, run_uid) + assert stops.get(timeout=5) == run_uid + streaming_client[run_uid].update_metadata({"notes": ["edited later"]}) + with pytest.raises(queue.Empty): + stops.get(timeout=1) diff --git a/tests/test_tiled_best_effort_streaming.py b/tests/test_tiled_best_effort_streaming.py new file mode 100644 index 0000000..3082ac2 --- /dev/null +++ b/tests/test_tiled_best_effort_streaming.py @@ -0,0 +1,140 @@ +import time + +import matplotlib +import pandas as pd +import pytest + +matplotlib.use("Agg") + +from bluesky import RunEngine # noqa: E402 +from bluesky.plans import count, scan # noqa: E402 +from bluesky.preprocessors import SupplementalData # noqa: E402 +from bluesky_tiled_plugins import TiledWriter # noqa: E402 +from ophyd.sim import det, motor, motor1 # noqa: E402 +from tiled.client import from_uri # noqa: E402 +from tiled.server import SimpleTiledServer # noqa: E402 + +from bluesky_tiled_plugins.streaming import ( + TiledBestEffortStreamer, + subscribe_to_hinted_streams, +) # noqa: E402 + + +@pytest.fixture(scope="module") +def client(): + server = SimpleTiledServer() + client = from_uri(server.uri) + yield client + client.context.close() + server.close() + + +@pytest.fixture +def run_engine(client): + RE = RunEngine() + RE.subscribe(TiledWriter(client)) + return RE + + +def _wait_for_finish(bec, uid, timeout=15.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + bec.process_events(timeout=0.1) + if bec._run is None and not bec._deferred and bec._queue.empty(): + return + raise TimeoutError(f"Run {uid} was not finished by the callback") + + +def test_scan_table_plot_and_peaks(client, run_engine): + lines = [] + with TiledBestEffortStreamer(client, out=lines.append) as bec: + (uid,) = run_engine(scan([det], motor, -1, 1, 5)) + bec.process_events(timeout=0.5) + _wait_for_finish(bec, uid) + + text = "\n".join(lines) + print(text) + assert f"Persistent Unique Scan ID: '{uid}'" in text + assert "motor" in text and "det" in text + data_rows = [ + line + for line in lines + if line.startswith("|") and line.split("|")[1].strip().isdigit() + ] + assert len(data_rows) == 5 + assert f"generator scan ['{uid[:8]}']" in text + # Peaks are keyed like BestEffortCallback.peaks: stat-first, with both + # attribute and subscription access. + assert "det" in bec.peaks["max"] + assert bec.peaks.max["det"][0] == pytest.approx(0.0) + assert bec.peaks.cen["det"] == pytest.approx(0.0, abs=0.5) + + +@pytest.mark.parametrize("retained_runs", [4, 0]) +def test_subscribe_to_hinted_streams(client, run_engine, retained_runs): + updates = [] + stops = [] + subscription = subscribe_to_hinted_streams( + client, + updates.append, + on_run_stop=lambda uid, stop: stops.append(uid), + retained_runs=retained_runs, + ) + try: + (uid,) = run_engine(scan([det], motor, -1, 1, 5)) + deadline = time.monotonic() + 15 + while ( + sum(len(u.table) for u in updates if u.run_uid == uid) < 5 + or uid not in stops + ) and time.monotonic() < deadline: + time.sleep(0.1) + finally: + subscription.disconnect() + + assert uid in stops + primary = [u for u in updates if u.run_uid == uid and u.stream_name == "primary"] + assert primary + assert primary[0].fields == ("motor", "det") + assert primary[0].start["uid"] == uid + data = pd.concat([u.data() for u in primary]) + assert list(data.columns) == ["seq_num", "time", "motor", "det"] + assert sorted(data["seq_num"]) == [1, 2, 3, 4, 5] + + +def test_baseline_and_time_dimension(client, run_engine): + sd = SupplementalData(baseline=[motor1]) + run_engine.preprocessors.append(sd) + lines = [] + with TiledBestEffortStreamer(client, out=lines.append) as bec: + (uid,) = run_engine(count([det], num=3)) + _wait_for_finish(bec, uid) + + text = "\n".join(lines) + print(text) + assert "Start-of-run baseline readings:" in text + assert "End-of-run baseline readings:" in text + assert "motor1" in text + assert f"generator count ['{uid[:8]}']" in text + + +def test_takeover_after_missing_stop(client, run_engine): + """A run without a stop document (e.g. backed up by the writer) must not + block the runs that follow it.""" + import uuid + + lines = [] + writer = TiledWriter(client) + with TiledBestEffortStreamer(client, out=lines.append, stop_timeout=0.5) as bec: + # A run that never gets a stop document. + broken_uid = str(uuid.uuid4()) + writer("start", {"uid": broken_uid, "time": time.time(), "scan_id": 998}) + bec.process_events(timeout=1.0) + assert f"Persistent Unique Scan ID: '{broken_uid}'" in lines + + (uid,) = run_engine(count([det])) + _wait_for_finish(bec, uid) + + text = "\n".join(lines) + assert f"Run {broken_uid[:8]} received no stop document" in text + assert f"Persistent Unique Scan ID: '{uid}'" in text + assert f"generator count ['{uid[:8]}']" in text