From 1047d4edd360cb537b3b47448ecd9c99e9ed0767 Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Sat, 3 Oct 2026 16:24:14 +0200 Subject: [PATCH] fix(waveform): track async subscriptions per scan and device --- .../widgets/plots/waveform/waveform.py | 134 ++++++++++---- tests/unit_tests/test_waveform.py | 172 +++++++++++++++++- 2 files changed, 274 insertions(+), 32 deletions(-) diff --git a/bec_widgets/widgets/plots/waveform/waveform.py b/bec_widgets/widgets/plots/waveform/waveform.py index b8b689f2..1e0745bf 100644 --- a/bec_widgets/widgets/plots/waveform/waveform.py +++ b/bec_widgets/widgets/plots/waveform/waveform.py @@ -148,7 +148,8 @@ class Waveform(PlotBase): self._async_curves = [] # Async streams already subscribed in the current scan. Keyed by # (device, stream) (stream is None for the old per-device endpoint); the value - # is the list of signal entries that share that single subscription. + # holds the signal entries that share that single subscription ("signals") and the + # exact "endpoint" and "cb_info" it was registered with. self._async_streams_setup: dict = {} self._history_curves = [] self._slice_index = None @@ -1656,26 +1657,41 @@ class Waveform(PlotBase): self.update_scan_info_from_source(self.scan_item) self._slice_index = None # Reset the slice index self._update_curve_visibility() - self._mode = self._categorise_device_curves() - - # First trigger to sync and async data. Async curve subscriptions - # were already set up by _categorise_device_curves above (it calls - # _setup_async_curve per async curve for live scans); repeating it - # here cleared each curve's data a second time and re-attempted - # the stream registration for every new scan. - if self._mode == "sync": - self.sync_signal_update.emit() - logger.info("Scan status: Sync mode") - elif self._mode == "async": - self.async_signal_update.emit() - logger.info("Scan status: Async mode") - else: - self.sync_signal_update.emit() - self.async_signal_update.emit() - logger.info("Scan status: Mixed mode") - logger.warning("Mixed mode - integrity of x axis cannot be guaranteed.") + self._start_scan_updates() + elif self.scan_item is None: + # The first status of this scan arrived before its ScanItem existed, so the curves + # could not be categorised and no live async subscription was set up. Do it as + # soon as the ScanItem is available, although the scan ID did not change. + self.scan_item = self.queue.scan_storage.find_scan_by_ID(self.scan_id) + if self.scan_item is not None: + self.update_scan_info_from_source(self.scan_item) + self._start_scan_updates() self.setup_dap_for_scan() + def _start_scan_updates(self): + """ + Categorise the device curves for the current scan, which also sets up the live async + subscriptions, and trigger the first sync and async updates. + """ + self._mode = self._categorise_device_curves() + + # First trigger to sync and async data. Async curve subscriptions + # were already set up by _categorise_device_curves above (it calls + # _setup_async_curve per async curve for live scans); repeating it + # here cleared each curve's data a second time and re-attempted + # the stream registration for every new scan. + if self._mode == "sync": + self.sync_signal_update.emit() + logger.info("Scan status: Sync mode") + elif self._mode == "async": + self.async_signal_update.emit() + logger.info("Scan status: Async mode") + else: + self.sync_signal_update.emit() + self.async_signal_update.emit() + logger.info("Scan status: Mixed mode") + logger.warning("Mixed mode - integrity of x axis cannot be guaranteed.") + @SafeSlot(dict, dict) def on_scan_progress(self, msg: dict, meta: dict): """ @@ -1842,7 +1858,6 @@ class Waveform(PlotBase): if async_signal_found: stream_key = (name, stream) # AsyncMultiSignal sub-signals share `stream` new_endpoint = MessageEndpoints.device_async_signal(self.scan_id, name, stream) - old_endpoint = MessageEndpoints.device_async_signal(self.old_scan_id, name, stream) legacy_endpoint = False else: # No BEC async signal (AsyncSignal/AsyncMultiSignal/DynamicSignal) was found @@ -1850,7 +1865,6 @@ class Waveform(PlotBase): # device_async_readback endpoint. stream_key = (name, None) # old endpoint is keyed by device only new_endpoint = MessageEndpoints.device_async_readback(self.scan_id, name) - old_endpoint = MessageEndpoints.device_async_readback(self.old_scan_id, name) legacy_endpoint = True endpoint_str = getattr(new_endpoint, "endpoint", new_endpoint) @@ -1858,9 +1872,9 @@ class Waveform(PlotBase): shared = self._async_streams_setup.get(stream_key) if shared is not None: # Another curve already subscribed to this exact stream this scan. - shared.append(signal) + shared["signals"].append(signal) logger.info( - f"Async signals {shared} share a single subscription on endpoint " + f"Async signals {shared['signals']} share a single subscription on endpoint " f"'{endpoint_str}'; reusing it instead of subscribing again." ) return @@ -1875,13 +1889,66 @@ class Waveform(PlotBase): # "waveform widget will be removed in a future release." # ) - self._async_streams_setup[stream_key] = [signal] - self.bec_dispatcher.disconnect_slot(self.on_async_readback, old_endpoint) + # Record the exact endpoint and subscription context next to the subscribers: the + # subscription is released with them (_release_async_streams), whatever scan_id holds by + # then, and on_async_readback matches readbacks to the curves of this stream. + cb_info = {"scan_id": self.scan_id, "device": name, "stream": stream_key[1]} + self._async_streams_setup[stream_key] = { + "signals": [signal], + "endpoint": new_endpoint, + "cb_info": cb_info, + } self.bec_dispatcher.connect_slot( - self.on_async_readback, new_endpoint, from_start=True, cb_info={"scan_id": self.scan_id} + self.on_async_readback, new_endpoint, from_start=True, cb_info=dict(cb_info) ) logger.info(f"Setup async curve {name} (signal '{signal}', endpoint '{endpoint_str}')") + def _release_async_streams(self, streams: dict) -> None: + """ + Disconnect async stream subscriptions that the current curve categorisation no longer + uses, e.g. those of the previous scan or of a live scan after switching to history. + + Each subscription is released with the exact endpoint and cb_info it was registered + with, so it is found regardless of the scan IDs tracked by the widget by now. + + Args: + streams(dict): Previous ``_async_streams_setup`` records. + """ + for stream_key, stream in streams.items(): + current = self._async_streams_setup.get(stream_key) + if current is not None and current["cb_info"] == stream["cb_info"]: + continue # the same subscription is still in use + self.bec_dispatcher.disconnect_slot( + self.on_async_readback, stream["endpoint"], cb_info=stream["cb_info"] + ) + + def _async_stream_curves(self, cb_info: dict) -> list[Curve]: + """ + Get the async curves fed by the subscription described by cb_info. + + Devices can expose signals with the same entry name, so curves are matched by device + and stream, not only by the signal entries of the readback. + + Args: + cb_info(dict): Subscription context set up by _setup_async_curve. + + Returns: + list[Curve]: The curves of the stream; all async curves if cb_info carries no + device, in which case the readback is matched by signal entry only. + """ + device = cb_info.get("device") + if device is None: + return self._async_curves + stream = self._async_streams_setup.get((device, cb_info.get("stream"))) + if stream is None: + return [] # the stream is no longer subscribed + return [ + curve + for curve in self._async_curves + if curve.config.signal.device == device + and curve.config.signal.signal in stream["signals"] + ] + @SafeSlot(dict, dict) def on_async_readback(self, msg, metadata): """ @@ -1899,14 +1966,15 @@ class Waveform(PlotBase): Args: msg(dict): Message with the async data. - metadata(dict): Message metadata, including the subscription's scan ID in ``cb_info``. + metadata(dict): Message metadata. ``metadata["cb_info"]`` is the subscription + context set up by _setup_async_curve (scan ID, device and stream). """ cb_info = metadata.get("cb_info") if not isinstance(cb_info, dict): - logger.info("Async readback has no subscription context.") + logger.info("Async readback has no subscription context, ignoring it.") return - scan_id = cb_info.get("scan_id") - if scan_id != self.scan_id: + if cb_info.get("scan_id") != self.scan_id: + # A late readback of a previous scan must not overwrite the current scan's data logger.info("Scan ID mismatch, ignoring async readback.") return @@ -1916,7 +1984,7 @@ class Waveform(PlotBase): return max_shape = metadata.get("async_update", {}).get("max_shape", []) plot_mode = self.x_axis_mode["name"] - for curve in self._async_curves: + for curve in self._async_stream_curves(cb_info): x_data = None # Reset x_data y_data = None # Reset y_data # Get the curve data @@ -2437,6 +2505,7 @@ class Waveform(PlotBase): # Reset sync/async curve lists self._async_curves.clear() self._sync_curves.clear() + previous_async_streams = self._async_streams_setup self._async_streams_setup = {} found_async = False found_sync = False @@ -2465,6 +2534,9 @@ class Waveform(PlotBase): "device not found in readout priority list." ) continue + # Subscriptions of the previous source that were not set up again (another scan, a + # history scan, removed curves) would keep feeding stale readbacks; release them. + self._release_async_streams(previous_async_streams) # Determine the mode of the scan if found_async and found_sync: mode = "mixed" diff --git a/tests/unit_tests/test_waveform.py b/tests/unit_tests/test_waveform.py index a17cf885..ac0a60e3 100644 --- a/tests/unit_tests/test_waveform.py +++ b/tests/unit_tests/test_waveform.py @@ -622,6 +622,176 @@ def test_on_scan_status_ignored_without_device_curves(qtbot, mocked_client_with_ assert setup_dap_spy.call_count == calls_before +def _async_readback_subscriptions(wf) -> list[tuple[str, object]]: + """(endpoint, dispatcher wrapper) of every subscription of wf.on_async_readback.""" + subscriptions = [ + (topic, slot) + for slot in wf.bec_dispatcher._registered_slots.values() + if slot.cb == wf.on_async_readback + for topic in slot.topics + ] + return sorted(subscriptions, key=lambda subscription: subscription[0]) + + +def _async_readback_endpoints(wf) -> list[str]: + return [endpoint for endpoint, _ in _async_readback_subscriptions(wf)] + + +def _deliver_async_readback(wf, endpoint: str, values: dict, update_type: str = "replace"): + """Deliver an async message like the dispatcher: through the wrapper subscribed to endpoint.""" + [wrapper] = [slot for topic, slot in _async_readback_subscriptions(wf) if topic == endpoint] + msg = {"signals": {signal: {"value": value} for signal, value in values.items()}} + wrapper(msg, {"async_update": {"type": update_type}}) + + +def _async_readback_endpoint(scan_id: str, device: str) -> str: + from bec_lib.endpoints import MessageEndpoints + + return MessageEndpoints.device_async_readback(scan_id, device).endpoint + + +def test_async_subscription_bound_when_scan_item_arrives_after_scan_status( + qtbot, mocked_client, monkeypatch +): + """ + Regression for #1297: the first scan status arrived before the ScanItem existed, so no + async subscription could be set up. A later status of the same scan must bind the live + async endpoint, so that async readbacks reach the plot. + """ + wf = create_widget(qtbot, Waveform, client=mocked_client) + curve = wf.plot(arg1="async_device") + scan_storage = {} + monkeypatch.setattr( + wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_storage.get(scan_id) + ) + + wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {}) + assert wf.scan_id == "scan-1" + assert _async_readback_endpoints(wf) == [] + + scan_storage["scan-1"] = create_dummy_scan_item(scan_id="scan-1") + wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {}) + + endpoint = _async_readback_endpoint("scan-1", "async_device") + assert _async_readback_endpoints(wf) == [endpoint] + _deliver_async_readback(wf, endpoint, {"async_device": [4, 5, 6]}) + np.testing.assert_array_equal(curve.get_data()[1], [4, 5, 6]) + + +def test_async_subscription_follows_live_history_live_switch( + qtbot, mocked_client, monkeypatch, scan_history_factory +): + """ + Regression for #1298: switching from a live scan to a history scan must release the live + async subscription (registered under the live scan ID), and the next live scan must end up + with exactly one subscription, its own. + """ + wf = create_widget(qtbot, Waveform, client=mocked_client) + wf.plot(arg1="async_device") + scan_storage = { + scan_id: create_dummy_scan_item(scan_id=scan_id) for scan_id in ("live-1", "live-2") + } + monkeypatch.setattr( + wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_storage.get(scan_id) + ) + inject_scan_history(wf, scan_history_factory, ("hist-1", 1)) + history_item = wf.client.history.get_by_scan_id("hist-1") + # Stored scans carry the readout priority of the scan in their BEC metadata + history_item.metadata["bec"]["readout_priority"] = { + "monitored": ["bpm4i"], + "async": ["async_device"], + } + monkeypatch.setattr(wf.client.history, "get_by_scan_id", lambda scan_id: history_item) + + wf.on_scan_status({"scan_id": "live-1", "status": "open"}, {}) + assert _async_readback_endpoints(wf) == [_async_readback_endpoint("live-1", "async_device")] + + wf.update_with_scan_history(scan_id="hist-1") + assert wf.scan_id == "hist-1" + assert _async_readback_endpoints(wf) == [] + + wf.on_scan_status({"scan_id": "live-2", "status": "open"}, {}) + assert _async_readback_endpoints(wf) == [_async_readback_endpoint("live-2", "async_device")] + + +def test_async_readback_of_previous_scan_is_ignored(qtbot, mocked_client, monkeypatch): + """ + Regression for #1299: an async readback of the previous scan that is delivered after the + scan changed (already in flight) must not overwrite the data of the current scan. + """ + wf = create_widget(qtbot, Waveform, client=mocked_client) + curve = wf.plot(arg1="async_device") + scan_storage = { + scan_id: create_dummy_scan_item(scan_id=scan_id) for scan_id in ("scan-1", "scan-2") + } + monkeypatch.setattr( + wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_storage.get(scan_id) + ) + + wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {}) + [(_, previous_subscription)] = _async_readback_subscriptions(wf) + wf.on_scan_status({"scan_id": "scan-2", "status": "open"}, {}) + _deliver_async_readback( + wf, _async_readback_endpoint("scan-2", "async_device"), {"async_device": [1, 2, 3]} + ) + + previous_subscription( + {"signals": {"async_device": {"value": [99]}}}, {"async_update": {"type": "replace"}} + ) + + np.testing.assert_array_equal(curve.get_data()[1], [1, 2, 3]) + + +@pytest.mark.parametrize( + "async_signal", [False, True], ids=["device_async_readback", "device_async_signal"] +) +def test_async_readback_updates_only_curves_of_its_device( + qtbot, mocked_client, monkeypatch, async_signal +): + """ + Regression for #1300: two async devices expose a signal with the same entry name. A + readback of one device must update that device's curve only. + """ + from bec_lib.device import ReadoutPriority + from bec_lib.endpoints import MessageEndpoints + + from bec_widgets.tests.utils import FakeDevice + + for name in ("async_a", "async_b"): + device = FakeDevice(name, readout_priority=ReadoutPriority.ASYNC) + device._info["signals"] = { + "value": { + "kind_str": "hinted", + "component_name": "value", + "obj_name": "value", + "signal_class": "AsyncSignal", + } + } + monkeypatch.setitem(mocked_client.device_manager.devices, name, device) + wf = create_widget(qtbot, Waveform, client=mocked_client) + if async_signal: + # Both signals are published as BEC async signal streams named after the signal + monkeypatch.setattr(wf, "_check_async_signal_found", lambda name, signal: (True, signal)) + curve_a = wf.plot(device_y="async_a", signal_y="value", label="curve_a") + curve_b = wf.plot(device_y="async_b", signal_y="value", label="curve_b") + scan_item = create_dummy_scan_item(scan_id="scan-1") + monkeypatch.setattr(wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_item) + + wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {}) + + def endpoint(device: str) -> str: + if async_signal: + return MessageEndpoints.device_async_signal("scan-1", device, "value").endpoint + return MessageEndpoints.device_async_readback("scan-1", device).endpoint + + assert _async_readback_endpoints(wf) == [endpoint("async_a"), endpoint("async_b")] + _deliver_async_readback(wf, endpoint("async_b"), {"value": [7, 8]}) + _deliver_async_readback(wf, endpoint("async_a"), {"value": [1, 2, 3]}) + + np.testing.assert_array_equal(curve_a.get_data()[1], [1, 2, 3]) + np.testing.assert_array_equal(curve_b.get_data()[1], [7, 8]) + + def test_request_dap_skips_unchanged_static_parent(qtbot, mocked_client_with_dap, monkeypatch): """ DAP curves whose parent is a static custom curve must not be resubmitted by @@ -1419,7 +1589,7 @@ def test_on_async_readback_through_dispatcher(qtbot, mocked_client): wrapper = next( slot for slot in wf.bec_dispatcher._registered_slots.values() - if slot.cb == wf.on_async_readback and slot.cb_info == {"scan_id": wf.scan_id} + if slot.cb == wf.on_async_readback and slot.cb_info["scan_id"] == wf.scan_id ) msg = {"signals": {"async_device": {"value": [123, 456]}}} metadata = {"async_update": {"type": "replace"}}