feat: Add support for device manager beamline modules (#1267)

Author: tpoliaw

helm/blueapi/config_schema.json

@@ -64,6 +64,39 @@
             "title": "CORSConfig",
             "type": "object"
         },
+        "DeviceManagerSource": {
+            "additionalProperties": false,
+            "properties": {
+                "module": {
+                    "description": "Module to be imported",
+                    "title": "Module",
+                    "type": "string"
+                },
+                "kind": {
+                    "const": "deviceManager",
+                    "default": "deviceManager",
+                    "title": "Kind",
+                    "type": "string"
+                },
+                "mock": {
+                    "default": false,
+                    "description": "If true, ophyd_async device connections are mocked",
+                    "title": "Mock",
+                    "type": "boolean"
+                },
+                "name": {
+                    "default": "devices",
+                    "description": "Name of the device manager in the module",
+                    "title": "Name",
+                    "type": "string"
+                }
+            },
+            "required": [
+                "module"
+            ],
+            "title": "DeviceManagerSource",
+            "type": "object"
+        },
         "DeviceSource": {
             "additionalProperties": false,
             "properties": {
@@ -131,6 +164,7 @@
                         "discriminator": {
                             "mapping": {
                                 "deviceFunctions": "#/$defs/DeviceSource",
+                                "deviceManager": "#/$defs/DeviceManagerSource",
                                 "dodal": "#/$defs/DodalSource",
                                 "planFunctions": "#/$defs/PlanSource"
                             },
@@ -145,6 +179,9 @@
                             },
                             {
                                 "$ref": "#/$defs/DodalSource"
+                            },
+                            {
+                                "$ref": "#/$defs/DeviceManagerSource"
                             }
                         ]
                     },

helm/blueapi/values.schema.json

@@ -553,6 +553,38 @@
                     },
                     "additionalProperties": false
                 },
+                "DeviceManagerSource": {
+                    "title": "DeviceManagerSource",
+                    "type": "object",
+                    "required": [
+                        "module"
+                    ],
+                    "properties": {
+                        "kind": {
+                            "title": "Kind",
+                            "default": "deviceManager",
+                            "const": "deviceManager"
+                        },
+                        "mock": {
+                            "title": "Mock",
+                            "description": "If true, ophyd_async device connections are mocked",
+                            "default": false,
+                            "type": "boolean"
+                        },
+                        "module": {
+                            "title": "Module",
+                            "description": "Module to be imported",
+                            "type": "string"
+                        },
+                        "name": {
+                            "title": "Name",
+                            "description": "Name of the device manager in the module",
+                            "default": "devices",
+                            "type": "string"
+                        }
+                    },
+                    "additionalProperties": false
+                },
                 "DeviceSource": {
                     "title": "DeviceSource",
                     "type": "object",
@@ -640,6 +672,9 @@
                                     },
                                     {
                                         "$ref": "#/$defs/DodalSource"
+                                    },
+                                    {
+                                        "$ref": "#/$defs/DeviceManagerSource"
                                     }
                                 ]
                             }

src/blueapi/config.py

@@ -47,6 +47,7 @@ class SourceKind(str, Enum):
     PLAN_FUNCTIONS = "planFunctions"
     DEVICE_FUNCTIONS = "deviceFunctions"
     DODAL = "dodal"
+    DEVICE_MANAGER = "deviceManager"
 
 
 class Source(BlueapiBaseModel):
@@ -72,6 +73,18 @@ class DodalSource(Source):
     )
 
 
+class DeviceManagerSource(Source):
+    kind: Literal[SourceKind.DEVICE_MANAGER] = Field(
+        SourceKind.DEVICE_MANAGER, init=False
+    )
+    mock: bool = Field(
+        description="If true, ophyd_async device connections are mocked", default=False
+    )
+    name: str = Field(
+        default="devices", description="Name of the device manager in the module"
+    )
+
+
 class TcpUrl(AnyUrl):
     _constraints = UrlConstraints(allowed_schemes=["tcp"])
 
@@ -121,7 +134,7 @@ class EnvironmentConfig(BlueapiBaseModel):
 
     sources: list[
         Annotated[
-            PlanSource | DeviceSource | DodalSource,
+            PlanSource | DeviceSource | DodalSource | DeviceManagerSource,
             Field(discriminator="kind"),
         ]
     ] = [

src/blueapi/core/__init__.py

@@ -15,6 +15,7 @@
 )
 from .context import BlueskyContext
 from .event import EventPublisher, EventStream
+from .protocols import DeviceManager
 
 OTLP_EXPORT_ENABLED = environ.get("OTLP_EXPORT_ENABLED") == "true"
 
@@ -28,6 +29,7 @@
     "EventPublisher",
     "EventStream",
     "DataEvent",
+    "DeviceManager",
     "WatchableStatus",
     "is_bluesky_compatible_device",
     "is_bluesky_plan_generator",

src/blueapi/core/context.py

@@ -11,7 +11,7 @@
 from bluesky.run_engine import RunEngine
 from dodal.common.beamlines.beamline_utils import get_path_provider, set_path_provider
 from dodal.utils import AnyDevice, make_all_devices
-from ophyd_async.core import NotConnectedError
+from ophyd_async.core import NotConnectedError, PathProvider
 from pydantic import BaseModel, GetCoreSchemaHandler, GetJsonSchemaHandler, create_model
 from pydantic.fields import FieldInfo
 from pydantic.json_schema import JsonSchemaValue, SkipJsonSchema
@@ -21,11 +21,13 @@
 from blueapi.client.numtracker import NumtrackerClient
 from blueapi.config import (
     ApplicationConfig,
+    DeviceManagerSource,
     DeviceSource,
     DodalSource,
     EnvironmentConfig,
     PlanSource,
 )
+from blueapi.core.protocols import DeviceManager
 from blueapi.utils import (
     BlueapiPlanModelConfig,
     is_function_sourced_from_module,
@@ -115,6 +117,7 @@ class BlueskyContext:
         default_factory=lambda: RunEngine(context_managers=[])
     )
     numtracker: NumtrackerClient | None = field(default=None, init=False, repr=False)
+    path_provider: PathProvider | None = None
     plans: dict[str, Plan] = field(default_factory=dict)
     devices: dict[str, Device] = field(default_factory=dict)
     plan_functions: dict[str, PlanGenerator] = field(default_factory=dict)
@@ -135,10 +138,12 @@ def __post_init__(self, configuration: ApplicationConfig | None):
                 )
 
             path_provider = StartDocumentPathProvider()
+            # TODO: Remove this when device manager is rolled out
             set_path_provider(path_provider)
 
             self.run_engine.subscribe(path_provider.run_start, "start")
             self.run_engine.subscribe(path_provider.run_stop, "stop")
+            self.path_provider = path_provider
 
             # local reference so it's available in _update_scan_num
             numtracker = self.numtracker
@@ -194,6 +199,13 @@ def with_config(self, config: EnvironmentConfig) -> None:
                     self.with_device_module(mod)
                 case DodalSource(mock=mock):
                     self.with_dodal_module(mod, mock=mock)
+                case DeviceManagerSource(mock=mock, name=name):
+                    manager = getattr(mod, name)
+                    if not isinstance(manager, DeviceManager):
+                        raise ValueError(
+                            f"{name} in module {mod} is not a device manager"
+                        )
+                    self.with_device_manager(manager, mock)
 
     def with_plan_module(self, module: ModuleType) -> None:
         """
@@ -227,6 +239,30 @@ def plan_2(...) -> MsgGenerator:
             ):
                 self.register_plan(obj)
 
+    def with_device_manager(self, manager: DeviceManager, mock: bool = False):
+        fixtures = {"path_provider": self.path_provider} if self.path_provider else {}
+        build_result = manager.build_and_connect(mock=mock, fixtures=fixtures)
+
+        for device in build_result.devices.values():
+            self.register_device(device)
+
+        if errs := build_result.build_errors:
+            LOGGER.warning(
+                f"{errs} errors while building devices",
+                exc_info=ExceptionGroup(
+                    "Errors while building devices", list(errs.values())
+                ),
+            )
+        if errs := build_result.connection_errors:
+            LOGGER.warning(
+                f"{len(errs)} errors while connecting devices",
+                exc_info=NotConnectedError(errs),
+            )
+        return build_result.devices, {
+            **build_result.build_errors,
+            **build_result.connection_errors,
+        }
+
     def with_device_module(self, module: ModuleType) -> None:
         self.with_dodal_module(module)
 

src/blueapi/core/protocols.py

@@ -0,0 +1,27 @@
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class DeviceConnectResult(Protocol):
+    devices: dict[str, Any]
+    build_errors: dict[str, Exception]
+    connection_errors: dict[str, Exception]
+
+
+@runtime_checkable
+class DeviceBuildResult(Protocol):
+    devices: dict[str, Any]
+    errors: dict[str, Exception]
+
+    def connect(self, timeout: float) -> DeviceConnectResult: ...
+
+
+@runtime_checkable
+class DeviceManager(Protocol):
+    def build_and_connect(
+        self,
+        *,
+        mock: bool = False,
+        timeout: float | None = None,
+        fixtures: dict[str, Any] | None = None,
+    ) -> DeviceConnectResult: ...