bec-server 3.146.2__tar.gz → 3.147.0__tar.gz

{bec_server-3.146.2 → bec_server-3.147.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bec-server
-Version: 3.146.2
+Version: 3.147.0
 Summary: BEC server
 Project-URL: Bug Tracker, https://github.com/bec-project/bec/issues
 Project-URL: Homepage, https://github.com/bec-project/bec

{bec_server-3.146.2 → bec_server-3.147.0}/bec_server/device_server/bec_message_handler.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterable, TypedDict
 
+import numpy as np
 from ophyd import OphydObject
 from ophyd_devices.utils import bec_signals as bms
 
@@ -16,6 +17,11 @@ if TYPE_CHECKING:  # pragma: no cover
     from bec_server.device_server.devices.devicemanager import DeviceManagerDS
 
 
+class _AsyncSignalIndexState(TypedDict):
+    scan_id: str
+    indices: dict[str, int]
+
+
 class BECMessageHandler:
     """
     A message handler for the BEC device server that processes and emits messages related to device signals.
@@ -32,6 +38,7 @@ class BECMessageHandler:
         self.device_manager = device_manager
         self.connector: RedisConnector = device_manager.connector
         self.devices = device_manager.devices
+        self._async_signal_index_state: dict[tuple[str, str], _AsyncSignalIndexState] = {}
 
     def emit(self, obj: OphydObject, message: messages.BECMessage):
         """
@@ -156,15 +163,71 @@ class BECMessageHandler:
                 f"Signal metadata is None for device {device_name} and signal {signal_name}."
             )
             return
+
+        indices = self._get_async_signal_indices(
+            obj=obj, scan_id=scan_id, signal_names=message.signals.keys()
+        )
+
+        # we augment the message metadata with the async signal indices so that
+        # downstream consumers can know which message corresponds to which async signal update
+        message.metadata["async_indices"] = indices
+
+        pipe = self.connector.pipeline()
         self.connector.xadd(
             MessageEndpoints.device_async_signal(
                 scan_id=scan_id, device=device_name, signal=signal_name
             ),
             {"data": message},
             max_size=obj.signal_metadata.get("max_size", 1000),
+            pipe=pipe,
             expire=obj.signal_metadata.get("expire", 600),  # default to 10 minutes (600 seconds)
         )
 
+        # NOTE: We are storing the shape information. This means that
+        # - a single value update will have shape [] (scalar)
+        # - an empty array will have shape [0]
+        # - a 1D array of length N will have shape [N]
+        # - a 2D array of shape (M, N) will have shape [M, N]
+        index_message = messages.DeviceAsyncSignalIndexMessage(
+            scan_id=scan_id,
+            device=device_name,
+            signal=signal_name,
+            shapes={
+                current_signal_name: list(np.shape(signal_data["value"]))
+                for current_signal_name, signal_data in message.signals.items()
+            },
+            indices=indices,
+            async_update=messages.DeviceAsyncUpdate.model_validate(
+                message.metadata["async_update"]
+            ),
+        )
+        self.connector.rpush(
+            MessageEndpoints.device_async_signal_index(
+                scan_id=scan_id, device=device_name, signal=signal_name
+            ),
+            index_message,
+            pipe=pipe,
+            expire=3600,  # expire after 1 hour; note that this is longer than the async signal message itself
+        )
+        pipe.execute()
+
+    def _get_async_signal_indices(
+        self, obj: bms.DynamicSignal, scan_id: str, signal_names: Iterable[str]
+    ) -> dict[str, int]:
+        state_key = (obj.root.name, obj.name)
+        state = self._async_signal_index_state.get(state_key)
+        if state is None or state["scan_id"] != scan_id:
+            state = _AsyncSignalIndexState(scan_id=scan_id, indices={})
+            self._async_signal_index_state[state_key] = state
+
+        out = {}
+        indices = state["indices"]
+        for signal_name in signal_names:
+            current_index = indices.get(signal_name, -1) + 1
+            indices[signal_name] = current_index
+            out[signal_name] = current_index
+        return out
+
     def _get_scan_status_message(self) -> messages.ScanStatusMessage | None:
         """
         Get the current scan status message from the device manager.

{bec_server-3.146.2 → bec_server-3.147.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "bec-server"
-version = "3.146.2"
+version = "3.147.0"
 description = "BEC server"
 requires-python = ">=3.11"
 classifiers = [
@@ -132,6 +132,7 @@ Homepage = "https://github.com/bec-project/bec"
 
 
 
+
 
 
 [tool.hatch.build.targets.wheel]

{bec_server-3.146.2 → bec_server-3.147.0}/tests/tests_device_server/test_device_manager_ds.py

@@ -24,13 +24,10 @@ def dm_with_devices_and_status(dm_with_devices):
     Fixture that adds the scan_info message to the device manager
     """
     device_manager = dm_with_devices
-    with mock.patch.object(
-        device_manager.scan_info, "msg", new_callable=mock.PropertyMock
-    ) as mock_scan_info_msg:
-        mock_scan_info_msg.return_value = messages.ScanStatusMessage(
-            scan_id="12345", status="open", info={"num_points": 10, "RID": "RID123"}
-        )
-        yield device_manager
+    device_manager.scan_info.msg = messages.ScanStatusMessage(
+        scan_id="12345", status="open", info={"num_points": 10, "RID": "RID123"}
+    )
+    yield device_manager
 
 
 class ControllerMock:
@@ -421,7 +418,8 @@ def test_device_manager_ds_obj_callback_progress_signal_disabled_device(dm_with_
     [
         None,
         messages.DeviceMessage(
-            signals={"async_signal": {"value": np.random.rand(10), "timestamp": time.time()}}
+            signals={"async_signal": {"value": np.random.rand(10), "timestamp": time.time()}},
+            metadata={"async_update": {"type": "add", "max_shape": [None]}},
         ),
         "some string",
     ],
@@ -446,7 +444,8 @@ def test_device_manager_ds_obj_callback_async_signal_incomplete_info(dm_with_dev
     dm_with_devices.devices.bec_signals_device.metadata = metadata
 
     msg = messages.DeviceMessage(
-        signals={"async_signal": {"value": np.random.rand(10), "timestamp": time.time()}}
+        signals={"async_signal": {"value": np.random.rand(10), "timestamp": time.time()}},
+        metadata={"async_update": {"type": "add", "max_shape": [None]}},
     )
     with mock.patch.object(device_manager.connector, "xadd") as mock_xadd:
         device_manager._obj_callback_bec_message_signal(obj=device.async_signal, value=msg)