From f037c6568f8b98ca3d4291210deaac86c2a1e910 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 7 Nov 2025 16:18:00 +0000 Subject: [PATCH] Add support for device manager beamline modules --- helm/blueapi/config_schema.json | 37 +++++ helm/blueapi/values.schema.json | 35 +++++ src/blueapi/config.py | 15 ++- src/blueapi/core/__init__.py | 2 + src/blueapi/core/context.py | 38 +++++- src/blueapi/core/protocols.py | 27 ++++ tests/unit_tests/core/test_context.py | 127 +++++++++++++++++- .../plans_and_devices.yaml | 4 + 8 files changed, 279 insertions(+), 6 deletions(-) create mode 100644 src/blueapi/core/protocols.py diff --git a/helm/blueapi/config_schema.json b/helm/blueapi/config_schema.json index a115ac9758..7104b9050a 100644 --- a/helm/blueapi/config_schema.json +++ b/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" } ] }, diff --git a/helm/blueapi/values.schema.json b/helm/blueapi/values.schema.json index e7f7ae3869..250dc38f10 100644 --- a/helm/blueapi/values.schema.json +++ b/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" } ] } diff --git a/src/blueapi/config.py b/src/blueapi/config.py index d6eb7853b7..a05a26c6ae 100644 --- a/src/blueapi/config.py +++ b/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"), ] ] = [ diff --git a/src/blueapi/core/__init__.py b/src/blueapi/core/__init__.py index 15e3b2602e..192a72adf4 100644 --- a/src/blueapi/core/__init__.py +++ b/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", diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index af5ede3cea..c4602ffd4e 100644 --- a/src/blueapi/core/context.py +++ b/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) diff --git a/src/blueapi/core/protocols.py b/src/blueapi/core/protocols.py new file mode 100644 index 0000000000..bf84fe7ece --- /dev/null +++ b/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: ... diff --git a/tests/unit_tests/core/test_context.py b/tests/unit_tests/core/test_context.py index f57e6a2f8c..9ece7896f7 100644 --- a/tests/unit_tests/core/test_context.py +++ b/tests/unit_tests/core/test_context.py @@ -1,10 +1,10 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path -from types import NoneType -from typing import Generic, TypeVar, Union -from unittest.mock import ANY, patch +from types import ModuleType, NoneType +from typing import Any, Generic, TypeVar, Union +from unittest.mock import ANY, MagicMock, Mock, patch import pytest from bluesky.protocols import ( @@ -18,7 +18,9 @@ from bluesky.run_engine import RunEngine from bluesky.utils import MsgGenerator from dodal.common import PlanGenerator, inject +from ophyd import Device from ophyd_async.core import ( + PathProvider, StandardDetector, StaticPathProvider, UUIDFilenameProvider, @@ -31,6 +33,7 @@ from pytest import LogCaptureFixture from blueapi.config import ( + DeviceManagerSource, DeviceSource, DodalSource, EnvironmentConfig, @@ -39,6 +42,7 @@ ) from blueapi.core import BlueskyContext, is_bluesky_compatible_device from blueapi.core.context import DefaultFactory, generic_bounds, qualified_name +from blueapi.core.protocols import DeviceConnectResult, DeviceManager from blueapi.utils.connect_devices import _establish_device_connections SIM_MOTOR_NAME = "sim" @@ -718,3 +722,118 @@ def demo_plan(foo: int | None) -> MsgGenerator: "foo": {"title": "Foo", "anyOf": [{"type": "integer"}, {"type": "null"}]} } assert "foo" in schema.get("required", []) + + +@dataclass +class StaticDeviceManager: + devices: dict[str, Any] = field(default_factory=dict) + build_errors: dict[str, Exception] = field(default_factory=dict) + connection_errors: dict[str, Exception] = field(default_factory=dict) + + def build_and_connect( + self, + *, + mock: bool = False, + timeout: float | None = None, + fixtures: dict[str, Any] | None = None, + ) -> DeviceConnectResult: + return self + + +def test_empty_device_manager(empty_context: BlueskyContext): + sdm = StaticDeviceManager() + empty_context.with_device_manager(sdm) + + assert empty_context.devices == {} + + +def test_single_device_manager(empty_context: BlueskyContext): + foo = Mock(spec=Device, name="foo") + foo.name = "foo" + sdm = StaticDeviceManager(devices={"foo": foo}) + empty_context.with_device_manager(sdm) + + assert empty_context.devices == {"foo": foo} + + +def test_device_manager_build_error(empty_context: BlueskyContext): + exc = ValueError("broken foo") + sdm = StaticDeviceManager(build_errors={"foo": exc}) + res = empty_context.with_device_manager(sdm) + + assert empty_context.devices == {} + assert res[1] == {"foo": exc} + + +def test_device_manager_connection_error(empty_context: BlueskyContext): + exc = ValueError("broken foo") + sdm = StaticDeviceManager(connection_errors={"foo": exc}) + res = empty_context.with_device_manager(sdm) + + assert empty_context.devices == {} + assert res[1] == {"foo": exc} + + +def test_device_manager_errors_are_combined(empty_context: BlueskyContext): + exc1 = ValueError("broken foo") + exc2 = ValueError("disconnected bar") + sdm = StaticDeviceManager( + build_errors={"foo": exc1}, connection_errors={"bar": exc2} + ) + res = empty_context.with_device_manager(sdm) + + assert empty_context.devices == {} + assert res[1] == {"foo": exc1, "bar": exc2} + + +def test_mock_passed_to_device_manager(empty_context: BlueskyContext): + dm = MagicMock(spec=DeviceManager) + bc = dm.build_and_connect.return_value + bc.devices = {} + bc.build_errors = {} + bc.connection_errors = {} + empty_context.with_device_manager(dm, mock=True) + + dm.build_and_connect.assert_called_once_with(mock=True, fixtures={}) + + +def test_path_provider_passed_to_device_manager(empty_context: BlueskyContext): + dm = MagicMock(spec=DeviceManager) + bc = dm.build_and_connect.return_value + bc.devices = {} + bc.build_errors = {} + bc.connection_errors = {} + pp = Mock(spec=PathProvider) + empty_context.path_provider = pp + empty_context.with_device_manager(dm) + + dm.build_and_connect.assert_called_once_with( + mock=False, fixtures={"path_provider": pp} + ) + + +def test_device_manager_environment_source(empty_context: BlueskyContext): + foo = Mock(spec=Device, name="foo") + foo.name = "foo" + stm = StaticDeviceManager(devices={"foo": foo}) + dev_mod = Mock(spec=ModuleType) + dev_mod.devices = stm + env = Mock(spec=EnvironmentConfig) + env.metadata = None + env.sources = [DeviceManagerSource(module="foo.bar", mock=True)] + with patch("blueapi.core.context.import_module") as imp_mod: + imp_mod.side_effect = lambda mod: dev_mod if mod == "foo.bar" else None + empty_context.with_config(env) + assert empty_context.devices == {"foo": foo} + + +def test_non_device_manager_errors(empty_context: BlueskyContext): + dev_mod = Mock(spec=ModuleType, name="dev_mod") + dev_mod.devices = "not-a-device-manager" + env = Mock(spec=EnvironmentConfig) + env.metadata = None + env.sources = [DeviceManagerSource(module="foo.bar", mock=True)] + with patch("blueapi.core.context.import_module") as imp_mod: + imp_mod.side_effect = lambda mod: dev_mod if mod == "foo.bar" else None + with pytest.raises(ValueError, match="not a device manager"): + empty_context.with_config(env) diff --git a/tests/unit_tests/valid_example_config/plans_and_devices.yaml b/tests/unit_tests/valid_example_config/plans_and_devices.yaml index 59b87c2688..d12ec638a0 100644 --- a/tests/unit_tests/valid_example_config/plans_and_devices.yaml +++ b/tests/unit_tests/valid_example_config/plans_and_devices.yaml @@ -4,3 +4,7 @@ env: module: my_plan_library.tomography.plans - kind: dodal module: dodal.beamlines.i04 + mock: true + - kind: deviceManager + module: dodal.beamlines.ixx + mock: True