feat(nidaq): bridge backend stream into the statistic groups during hardware scans

debye_bec/devices/nidaq/nidaq.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import threading
 import time
+import traceback
+from collections import defaultdict
 from typing import TYPE_CHECKING, Literal
 
 from bec_lib.logger import bec_logger
@@ -58,6 +61,22 @@ NIDAQ_SW_FORWARD_SOURCES: dict[str, dict[str, str]] = {
     "stream": {"energy": "energy_epics", "enc": "enc_epics"},
 }
 
+#: Backend channel name -> (group, group channel) for the unified bridge.
+#: Covers both backend naming variants: with RLE the per-oscillation statistics
+#: carry _mean/_std_dev/_max suffixes, without RLE the names are bare.
+NIDAQ_BRIDGE_ROUTING: dict[str, tuple[str, str]] = (
+    {f"ai{i}_mean": ("mean", f"ai{i}") for i in range(8)}
+    | {f"ai{i}": ("mean", f"ai{i}") for i in range(8)}
+    | {f"ai{i}_std_dev": ("std_dev", f"ai{i}") for i in range(8)}
+    | {f"ci{i}_mean": ("mean", f"ci{i}") for i in range(18)}
+    | {f"ci{i}": ("mean", f"ci{i}") for i in range(18)}
+    | {f"ci{i}_std_dev": ("std_dev", f"ci{i}") for i in range(18)}
+    | {f"di{i}_max": ("max", f"di{i}") for i in range(5)}
+    | {f"di{i}": ("max", f"di{i}") for i in range(5)}
+    | {channel: ("stream", channel) for channel in NIDAQ_STREAM_CHANNELS}
+    | {"xrd_ai0": ("stream", "xrd_ai0_mean")}
+)
+
 
 class NidaqError(Exception):
     """Nidaq specific error"""
@@ -179,7 +198,15 @@ class Nidaq(PSIDeviceBase, NidaqControl):
 
     USER_ACCESS = ["set_config"]
 
-    def __init__(self, prefix: str = "", *, name: str, scan_info: ScanInfo = None, **kwargs):
+    def __init__(
+        self,
+        prefix: str = "",
+        *,
+        name: str,
+        scan_info: ScanInfo = None,
+        unified_bridge: bool = False,
+        **kwargs,
+    ):
         super().__init__(name=name, prefix=prefix, scan_info=scan_info, **kwargs)
         self.scan_parameters: ScanServerScanInfo = None
         self.timeout_wait_for_signal = 5  # put 5s firsts
@@ -194,6 +221,10 @@ class Nidaq(PSIDeviceBase, NidaqControl):
             "nidaq_continuous_scan",
         ]
         self._group_put_errors: set[str] = set()
+        self.unified_bridge_enabled = unified_bridge
+        self._bridge_stop_event = threading.Event()
+        self._bridge_task = None
+        self._unknown_bridge_channels: set[str] = set()
 
     ########################################
     #           Beamline Methods           #
@@ -253,7 +284,11 @@ class Nidaq(PSIDeviceBase, NidaqControl):
                 self._guarded_group_put(group, batch, acquisition_group="monitored")
 
     def _guarded_group_put(
-        self, group: str, signals: dict, acquisition_group: str | None = None
+        self,
+        group: str,
+        signals: dict,
+        acquisition_group: str | None = None,
+        async_update: dict | None = None,
     ) -> None:
         """Put a batch into a group; report validation failures once.
 
@@ -264,7 +299,9 @@ class Nidaq(PSIDeviceBase, NidaqControl):
         """
         group_signal = getattr(self, group)
         try:
-            group_signal.put(signals, acquisition_group=acquisition_group)
+            group_signal.put(
+                signals, acquisition_group=acquisition_group, async_update=async_update
+            )
         except ValueError as exc:
             if group not in self._group_put_errors:
                 self._group_put_errors.add(group)
@@ -275,6 +312,102 @@ class Nidaq(PSIDeviceBase, NidaqControl):
                     f"{sorted(declared)}: {exc}"
                 )
 
+    ########################################
+    #  unified bridge (hardware-mode data) #
+    ########################################
+
+    def _start_unified_bridge(self) -> None:
+        """Re-emit the backend's device_async_readback stream through the groups.
+
+        Transitional measure until the backend emits the unified names natively:
+        during hardware scans the backend pushes channel-first names
+        (nidaq_ai0_mean, ...) on the legacy per-device stream. The bridge
+        translates them per NIDAQ_BRIDGE_ROUTING into the DynamicSignal groups
+        (nidaq_mean_ai0, ...) with each message's original async_update passed
+        through, so hardware data appears on device_async_signal (GUI) and in
+        the NeXus file under the unified names as well (double-write).
+        """
+        connector = getattr(self.device_manager, "connector", None)
+        if connector is None:
+            logger.warning(
+                f"Device {self.name}: unified bridge enabled but no device_manager "
+                "connector is available; bridge not started."
+            )
+            return
+        self._stop_unified_bridge()
+        self._group_put_errors.clear()
+        self._bridge_stop_event = threading.Event()
+        scan_id = getattr(self.scan_info.msg, "scan_id", None)
+        self._bridge_task = threading.Thread(
+            target=self._run_unified_bridge,
+            args=(scan_id,),
+            daemon=True,
+            name=f"{self.name}_unified_bridge",
+        )
+        self._bridge_task.start()
+
+    def _stop_unified_bridge(self) -> None:
+        """Signal the bridge task to stop."""
+        self._bridge_stop_event.set()
+        self._bridge_task = None
+
+    def _run_unified_bridge(self, scan_id: str | None) -> None:
+        """Consume this scan's backend stream and re-emit it via the groups."""
+        if not scan_id:
+            return
+        from bec_lib.endpoints import MessageEndpoints
+
+        connector = self.device_manager.connector
+        topic = MessageEndpoints.device_async_readback(scan_id=scan_id, device=self.name)
+        stop_event = self._bridge_stop_event
+        logger.info(f"Device {self.name}: unified bridge started for scan {scan_id}")
+        first_read = True
+        while not stop_event.is_set():
+            try:
+                # from_start resets the connector cursor on EVERY call; use it
+                # once to pick up pre-start entries, then continue via the
+                # connector's stored cursor to avoid re-emitting old bundles.
+                entries = connector.xread(topic, count=10, block=200, from_start=first_read)
+                first_read = False
+            except Exception:  # pylint: disable=broad-except
+                logger.error(
+                    f"Device {self.name}: reading backend stream failed: "
+                    f"{traceback.format_exc()}"
+                )
+                stop_event.wait(1)
+                continue
+            if not entries:
+                continue
+            for entry in entries:
+                msg = entry.get("data")
+                if msg is not None:
+                    self._forward_backend_message(msg)
+        logger.info(f"Device {self.name}: unified bridge stopped")
+
+    def _forward_backend_message(self, msg) -> None:
+        """Route one backend DeviceMessage into the statistic groups."""
+        device_prefix = f"{self.name}_"
+        buckets: dict[str, dict] = defaultdict(dict)
+        now = time.time()
+        for signal_name, data in msg.signals.items():
+            channel = signal_name.removeprefix(device_prefix)
+            route = NIDAQ_BRIDGE_ROUTING.get(channel)
+            if route is None:
+                if channel not in self._unknown_bridge_channels:
+                    self._unknown_bridge_channels.add(channel)
+                    logger.warning(
+                        f"Device {self.name}: skipping unknown backend channel '{signal_name}'"
+                    )
+                continue
+            group, group_channel = route
+            buckets[group][group_channel] = {
+                "value": data["value"],
+                "timestamp": data.get("timestamp") or now,
+            }
+        async_update = (msg.metadata or {}).get("async_update")
+        for group, signals in buckets.items():
+            self._guarded_group_put(group, signals, async_update=async_update)
+
     def set_config(
         self,
         sampling_rate: Literal[
@@ -433,6 +566,9 @@ class Nidaq(PSIDeviceBase, NidaqControl):
         if not self._check_if_scan_name_is_valid(self.scan_parameters):
             return None
 
+        if self.unified_bridge_enabled:
+            self._start_unified_bridge()
+
         if self.state.get() != NidaqState.STANDBY:
             status = CompareStatus(self.state, NidaqState.STANDBY)
             self.cancel_on_stop(status)
@@ -481,6 +617,7 @@ class Nidaq(PSIDeviceBase, NidaqControl):
 
     def on_unstage(self) -> DeviceStatus | StatusBase | None:
         """Called while unstaging the device. Check that the Nidaq goes into Standby"""
+        self._stop_unified_bridge()
 
         status = CompareStatus(self.state, NidaqState.STANDBY)
         self.cancel_on_stop(status)
@@ -552,4 +689,5 @@ class Nidaq(PSIDeviceBase, NidaqControl):
 
     def on_stop(self) -> None:
         """Called when the device is stopped."""
+        self._stop_unified_bridge()
         self.stop_call.put(1)