feat(device manager): add device dependency resolution

bec_lib/bec_lib/atlas_models.py

@@ -43,6 +43,7 @@ class _DeviceModelCore(BaseModel):
     connectionTimeout: float = 5.0
     description: str = ""
     deviceTags: set[str] = set()
+    needs: list[str] = []
     onFailure: Literal["buffer", "retry", "raise"] = "retry"
     readOnly: bool = False
     softwareTrigger: bool = False
@@ -119,6 +120,7 @@ class DeviceHashModel(BaseModel, frozen=True):
     connectionTimeout: HashInclusion = HashInclusion.EXCLUDE
     deviceConfig: DictHashInclusion = DictHashInclusion(field_inclusion=HashInclusion.VARIANT)
     deviceTags: HashInclusion = HashInclusion.EXCLUDE
+    needs: HashInclusion = HashInclusion.VARIANT
     readoutPriority: HashInclusion = HashInclusion.EXCLUDE
     description: HashInclusion = HashInclusion.EXCLUDE
     readOnly: HashInclusion = HashInclusion.EXCLUDE

bec_lib/bec_lib/config_helper.py

@@ -542,6 +542,7 @@ def _save_config_to_file(self, file_path: str, raise_on_error: bool = True) -> b
             '# - description: A string description of the device. Default is "" (empty string).\n'
             "# - deviceConfig: A dictionary of configuration parameters specific to the device class. Default is None.\n"
             "# - deviceTags: A list/set of tags associated with the device. Default is an empty list/set.\n"
+            "# - needs: A list of other device names that this device depends on. Default is [].\n"
             '# - onFailure: The action to take on failure. Possible values are ["buffer", "retry", "raise"]. Default is "retry".\n'
             "# - readOnly: A boolean indicating if the device is read-only. Default is false.\n"
             "# - softwareTrigger: A boolean indicating if the device uses software triggering. Default is false.\n"
@@ -563,6 +564,9 @@ def _save_config_to_file(self, file_path: str, raise_on_error: bool = True) -> b
             "#   enabled: true\n"
             "#   connectionTimeout: 20\n"
             "#   readoutPriority: baseline\n"
+            "#   needs:\n"
+            "#     - device2\n"
+            "#     - device3\n"
             "#   onFailure: retry\n"
             "#   readOnly: false\n"
             "#   softwareTrigger: false\n"

bec_lib/tests/test_device_hashing.py

@@ -18,6 +18,7 @@
     "deviceClass": "TestDeviceClass",
     "readoutPriority": "baseline",
     "enabled": True,
+    "needs": [],
 }
 
 
@@ -246,11 +247,12 @@ def test_device_equality_according_to_model(hash_model: DeviceHashModel, equal:
 @pytest.mark.parametrize(
     "hash_model, expected",
     [
-        (DeviceHashModel(), {"deviceConfig": {"a": "b", "c": "d", "l": "m"}}),
-        (DeviceHashModel(deviceConfig=DictHashInclusion()), {}),
+        (DeviceHashModel(), {"deviceConfig": {"a": "b", "c": "d", "l": "m"}, "needs": []}),
+        (DeviceHashModel(deviceConfig=DictHashInclusion(), needs=HashInclusion.EXCLUDE), {}),
         (
             DeviceHashModel(
                 deviceConfig=DictHashInclusion(),
+                needs=HashInclusion.EXCLUDE,
                 enabled=HashInclusion.VARIANT,
                 softwareTrigger=HashInclusion.VARIANT,
             ),
@@ -259,6 +261,7 @@ def test_device_equality_according_to_model(hash_model: DeviceHashModel, equal:
         (
             DeviceHashModel(
                 deviceConfig=DictHashInclusion(),
+                needs=HashInclusion.EXCLUDE,
                 userParameter=DictHashInclusion(
                     field_inclusion=HashInclusion.INCLUDE,
                     inclusion_keys={"a"},
@@ -275,7 +278,7 @@ def test_device_equality_according_to_model(hash_model: DeviceHashModel, equal:
                     remainder_inclusion=HashInclusion.EXCLUDE,
                 )
             ),
-            {"deviceConfig": {"a": "b", "c": "d", "l": "m"}},
+            {"deviceConfig": {"a": "b", "c": "d", "l": "m"}, "needs": []},
         ),
     ],
 )
@@ -426,5 +429,6 @@ def test_hashable_device_set_or_adds_variants():
     combined_device: HashableDevice = combined.pop()
     assert combined_device.variants != set()
     assert HashableDevice.model_validate(combined_device.variants.pop())._variant_info() == {
-        "deviceConfig": {"param": "other_value"}
+        "deviceConfig": {"param": "other_value"},
+        "needs": [],
     }

bec_server/bec_server/device_server/device_server.py

@@ -575,6 +575,7 @@ def _trigger_device(self, instr: messages.DeviceInstructionMessage) -> None:
         devices = instr.content["device"]
         if not isinstance(devices, list):
             devices = [devices]
+        devices = self.device_manager.get_device_order(devices)
         self.requests_handler.add_request(instr, num_status_objects=len(devices))
         for dev in devices:
             obj = self.device_manager.devices.get(dev)
@@ -611,6 +612,8 @@ def _complete_device(self, instr: messages.DeviceInstructionMessage) -> None:
             if not isinstance(devices, list):
                 devices = [devices]
 
+        devices = self.device_manager.get_device_order(devices)
+
         self.requests_handler.add_request(instr, num_status_objects=len(devices))
         num_status_objects = 0
         for dev in devices:
@@ -662,6 +665,8 @@ def _pre_scan(self, instr: messages.DeviceInstructionMessage) -> None:
         if not isinstance(devices, list):
             devices = [devices]
 
+        devices = self.device_manager.get_device_order(devices)
+
         self.requests_handler.add_request(instr, num_status_objects=len(devices))
         num_status_objects = 0
         for dev in devices:
@@ -745,6 +750,8 @@ def _read_device(self, instr: messages.DeviceInstructionMessage, new_status=True
         if not isinstance(devices, list):
             devices = [devices]
 
+        devices = self.device_manager.get_device_order(devices)
+
         if not new_status:
             self._read_and_update_devices(devices, instr.metadata)
             return
@@ -757,6 +764,7 @@ def _read_and_update_devices(self, devices: list[str], metadata: dict) -> list:
         start = time.time()
         pipe = self.connector.pipeline()
         signal_container = []
+        devices = self.device_manager.get_device_order(devices)
         for dev in devices:
             device_root = dev.split(".")[0]
             self.device_manager.devices.get(device_root).metadata = metadata
@@ -788,6 +796,7 @@ def _read_config_and_update_devices(self, devices: list[str], metadata: dict) ->
         start = time.time()
         pipe = self.connector.pipeline()
         signal_container = []
+        devices = self.device_manager.get_device_order(devices)
         for dev in devices:
             self.device_manager.devices.get(dev).metadata = metadata
             obj = self.device_manager.devices.get(dev).obj
@@ -853,6 +862,8 @@ def _stage_device(
         if not isinstance(devices, list):
             devices = [devices]
 
+        devices = self.device_manager.get_device_order(devices)
+
         self.requests_handler.add_request(instr, num_status_objects=len(devices))
 
         num_status_objects = 0
@@ -913,6 +924,8 @@ def _unstage_device(self, instr: messages.DeviceInstructionMessage) -> None:
         if not isinstance(devices, list):
             devices = [devices]
 
+        devices = self.device_manager.get_device_order(devices)
+
         self.requests_handler.add_request(instr, num_status_objects=len(devices))
         num_status_objects = 0
         for dev in devices:

bec_server/bec_server/device_server/devices/config_update_handler.py

@@ -146,14 +146,19 @@ def parse_config_request(
                     self._remove_config(msg, cancel_event)
                 case _:
                     pass
-
+            # After any config change, resolve dependencies. It will raise if dependencies are not met.
+            self.update_session_config(msg)
+            self.device_manager.resolve_device_dependencies(
+                self.device_manager.current_session["devices"]
+            )
         except CancelledError:
             error_msg = "Request was cancelled"
             accepted = False
             logger.info(
                 f"Config request {msg.metadata.get('RID')} was cancelled. The config will be flushed."
             )
             self._flush_config()
+
         except Exception:
             error_msg = traceback.format_exc()
             accepted = False
@@ -290,6 +295,43 @@ def _remove_config(
             self.device_manager.reset_device(device)
             self.device_manager.devices.pop(dev)
 
+    def update_session_config(self, msg: messages.DeviceConfigMessage) -> None:
+        """
+        Updates the current session config with the new config from the message.
+
+        Args:
+            msg (BECMessage.DeviceConfigMessage): Config message containing the new config
+
+        """
+        action = msg.action
+        match action:
+            case "update":
+                # Update the session config
+                for dev in msg.content["config"]:
+                    dev_config = self.device_manager.devices[dev]._config
+                    session_device_config = next(
+                        (
+                            d
+                            for d in self.device_manager.current_session["devices"]
+                            if d["name"] == dev
+                        ),
+                        None,
+                    )
+                    if session_device_config:
+                        session_device_config.update(dev_config)
+            case "add":
+                # Add new devices to the session config
+                for dev, dev_config in msg.content["config"].items():
+                    self.device_manager.current_session["devices"].append(dev_config)
+            case "remove":
+                # Remove devices from the session config
+                for dev in msg.content["config"]:
+                    self.device_manager.current_session["devices"] = [
+                        d
+                        for d in self.device_manager.current_session["devices"]
+                        if d["name"] != dev
+                    ]
+
     def handle_failed_device_inits(self):
         if self.device_manager.failed_devices:
             msg = messages.DeviceConfigMessage(