wip - async signals

Author: wakonig

bec_lib/bec_lib/data_api/data_api.py

@@ -416,3 +416,17 @@ def unsubscribe(
         """
         for plugin in self.plugins:
             plugin.unsubscribe(subscription_id, scan_id, callback)
+
+
+if __name__ == "__main__":
+    from bec_lib.client import BECClient
+
+    def my_callback(data, metadata):
+        print("Received data:", data)
+        print("With metadata:", metadata)
+
+    client = BECClient()
+    data_api = DataAPI(client)
+    sub = data_api.create_subscription("test_scan")
+    sub.add_device(device_name="waveform", device_entry="waveform_data")
+    sub.set_callback(my_callback)

bec_lib/bec_lib/data_api/plugins.py

@@ -295,13 +295,28 @@ def _device_entry_is_async_signal(self, device_name: str, device_entry: str) ->
         Returns:
             True if the device entry is an async signal, False otherwise.
         """
+        async_signal_info = self._get_async_signal_info(device_name, device_entry)
+        return async_signal_info is not None
+
+    def _get_async_signal_info(self, device_name: str, device_entry: str) -> dict | None:
+        """
+        Get the async signal information for the given device and entry.
+
+        Args:
+            device_name: Name of the device.
+            device_entry: Specific entry of the device.
+        Returns:
+            The async signal information dict if found, None otherwise.
+        """
         if not self.client.device_manager:
-            return False
-        async_signals = self.client.device_manager.get_bec_signals("AsyncSignal")
-        for entry_name, _, entry_data in async_signals:
-            if entry_name == device_entry and entry_data.get("device_name") == device_name:
-                return True
-        return False
+            return None
+        async_signals = self.client.device_manager.get_bec_signals(
+            ["AsyncSignal", "AsyncMultiSignal", "DynamicSignal"]
+        )
+        for dev_name, _, entry_info in async_signals:
+            if entry_info.get("obj_name") == device_entry and dev_name == device_name:
+                return entry_info
+        return None
 
     def subscribe(
         self,
@@ -665,9 +680,16 @@ def _subscribe_to_async_signal(
                 async_sub.callback_refs.append(callback_ref)
         else:
             # Create new redis connector subscription
+            async_signal_info = self._get_async_signal_info(device_name, device_entry)
+            if async_signal_info is None:
+                raise ValueError(
+                    f"Cannot subscribe to async signal '{device_name}' entry '{device_entry}': signal not found."
+                )
             connector_id = self.client.connector.register(
                 MessageEndpoints.device_async_signal(
-                    scan_id=scan_id, device=device_name, signal=device_entry
+                    scan_id=scan_id,
+                    device=device_name,
+                    signal=async_signal_info.get("storage_name"),
                 ),
                 cb=self._async_signal_sync_callback,
                 from_start=True,
@@ -783,11 +805,7 @@ def _get_expected_device_count(self, callback_ref: CallbackRef, scan_id: str) ->
                 count += len(self._monitored_subscriptions[scan_id][callback_ref].devices)
 
         # Count async signals
-        for (
-            sub_scan_id,
-            device_name,
-            device_entry,
-        ), async_sub in self._async_subscriptions.items():
+        for (sub_scan_id, _, _), async_sub in self._async_subscriptions.items():
             if sub_scan_id == scan_id and callback_ref in async_sub.callback_refs:
                 count += 1
 
@@ -854,16 +872,3 @@ def _check_and_emit_synchronized_data(self, callback_ref: CallbackRef, scan_id:
 
         # Update the min_length to track what we've already emitted
         callback_buffer.min_length = min_length
-
-
-"""
-NOTES
-
-- AsyncSignal subscriptions should be shared between multiple subscribers to avoid redundant subscriptions.
-- Whenever the redis connector triggers the callback, we broadcast to all subscribers of that device/entry/scan combination.
-- When the user subscribed to an AsyncSignal, we check if there is already a subscription for that device/entry/scan combination.
-- If yes, we just add the callback to the list of callbacks for that subscription.
-- When the user subscribes to multiple AsyncSignals, we synchronize the data length and only broadcast the data of equal length. Same for mixtures 
-    of monitored devices and AsyncSignals.
-
-"""

bec_lib/tests/test_data_api.py

@@ -308,8 +308,8 @@ def test_device_entry_is_monitored(self, mock_client, scan_item_with_monitored_d
     def test_device_entry_is_async_signal(self, mock_client):
         """Test detection of async signal device entries."""
         mock_client.device_manager.get_bec_signals.return_value = [
-            ("async_sig1", None, {"device_name": "detector1"}),
-            ("async_sig2", None, {"device_name": "detector2"}),
+            ("detector1", None, {"obj_name": "async_sig1", "storage_name": "async_sig1"}),
+            ("detector2", None, {"obj_name": "async_sig2", "storage_name": "async_sig2"}),
         ]
 
         plugin = BECLiveDataPlugin(mock_client)
@@ -355,12 +355,12 @@ def callback2(data, metadata):
     def test_subscribe_to_async_signal(self, mock_client, mock_callback):
         """Test subscription to async signal."""
         mock_client.device_manager.get_bec_signals.return_value = [
-            ("async_sig1", None, {"device_name": "detector1"})
+            ("detector1", None, {"obj_name": "async_sig1", "storage_name": "async_sig1"})
         ]
 
         # Setup scan item
         scan_item = ScanItem(
-            queue_id="test_queue", scan_number=[1], scan_id=["test_scan_id"], status="open"
+            queue_id="test_queue", scan_number=1, scan_id="test_scan_id", status="open"
         )
         scan_item.status_message = messages.ScanStatusMessage(
             scan_id="test_scan_id", status="open", info={}
@@ -391,7 +391,7 @@ def test_subscribe_to_async_signal(self, mock_client, mock_callback):
     def test_subscribe_to_async_signal_shared_subscription(self, mock_client):
         """Test that multiple callbacks share the same async signal subscription."""
         mock_client.device_manager.get_bec_signals.return_value = [
-            ("async_sig1", None, {"device_name": "detector1"})
+            ("detector1", None, {"obj_name": "async_sig1", "storage_name": "async_sig1"})
         ]
 
         scan_item = ScanItem(
@@ -520,7 +520,7 @@ def test_data_synchronization_mixed_sources(
     ):
         """Test data synchronization between monitored and async sources."""
         mock_client.device_manager.get_bec_signals.return_value = [
-            ("async_sig1", None, {"device_name": "detector1"})
+            ("detector1", None, {"obj_name": "async_sig1", "storage_name": "async_sig1"})
         ]
 
         plugin = BECLiveDataPlugin(mock_client)
@@ -588,7 +588,7 @@ def test_unsubscribe_monitored_device(
     def test_unsubscribe_async_signal(self, mock_client, mock_callback):
         """Test unsubscribing from async signal."""
         mock_client.device_manager.get_bec_signals.return_value = [
-            ("async_sig1", None, {"device_name": "detector1"})
+            ("detector1", None, {"obj_name": "async_sig1", "storage_name": "async_sig1"})
         ]
 
         scan_item = ScanItem(
@@ -653,7 +653,7 @@ def test_can_provide_monitored(self, mock_client, scan_item_with_monitored_devic
     def test_can_provide_async_signal(self, mock_client):
         """Test can_provide returns True for async signals."""
         mock_client.device_manager.get_bec_signals.return_value = [
-            ("async_sig1", None, {"device_name": "detector1"})
+            ("detector1", None, {"obj_name": "async_sig1", "storage_name": "async_sig1"})
         ]
 
         scan_item = ScanItem(