Connect all ophyd-async devices in CLI

src/dodal/cli.py

@@ -1,12 +1,21 @@
+import asyncio
 import os
+from typing import Iterable, Mapping
 
 import click
-from bluesky.run_engine import RunEngine
-from ophyd_async.core import Device, NotConnected
+from bluesky.run_engine import RunEngine, call_in_bluesky_event_loop
+from ophyd_async.core import NotConnected, wait_for_connection
 from ophyd_async.plan_stubs import ensure_connected
 
 from dodal.beamlines import all_beamline_names, module_name_for_beamline
-from dodal.utils import make_all_devices
+from dodal.utils import (
+    AnyDevice,
+    OphydV1Device,
+    OphydV2Device,
+    connect_all_ophyd_async_devices,
+    filter_ophyd_devices,
+    make_all_devices,
+)
 
 from . import __version__
 
@@ -51,26 +60,37 @@ def connect(beamline: str, all: bool, sim_backend: bool) -> None:
 
     # We need to make a RunEngine to allow ophyd-async devices to connect.
     # See https://blueskyproject.io/ophyd-async/main/explanations/event-loop-choice.html
-    RE = RunEngine()
+    RunEngine()
 
     print(f"Attempting connection to {beamline} (using {full_module_path})")
-    devices, exceptions = make_all_devices(
+    created_devices, creation_exceptions = make_all_devices(
         full_module_path,
         include_skipped=all,
         fake_with_ophyd_sim=sim_backend,
+        wait_for_connection=False,
+    )
+    # Connect all ophyd-async devices to their signal backends, any failures will
+    # be included in the output to the user
+    ophyd_devices, ophyd_async_devices = filter_ophyd_devices(created_devices)
+    connected_devices, connection_exceptions = call_in_bluesky_event_loop(
+        connect_all_ophyd_async_devices(
+            ophyd_async_devices,
+            mock=sim_backend,
+        )
     )
+    # We assume that all ophyd v1 devices successfully connected, because there is
+    # no reliable way to tell
+    connected_devices = {**connected_devices, **ophyd_devices}
     sim_statement = " (sim mode)" if sim_backend else ""
-    ophyd_async_devices = [
-        device for device in devices.values() if isinstance(device, Device)
-    ]
-    RE(ensure_connected(*ophyd_async_devices))
 
-    print(f"{len(devices)} devices connected{sim_statement}:")
-    connected_devices = "\n".join(
-        sorted([f"\t{device_name}" for device_name in devices.keys()])
+    # Report a list of successfully connected devices
+    print(f"{len(connected_devices)} devices connected{sim_statement}:")
+    formatted_connected_devices = "\n".join(
+        sorted([f"\t{device_name}" for device_name in connected_devices.keys()])
     )
-    print(connected_devices)
+    print(formatted_connected_devices)
 
     # If exceptions have occurred, this will print details of the relevant PVs
+    exceptions = {**creation_exceptions, **connection_exceptions}
     if len(exceptions) > 0:
         raise NotConnected(exceptions)

src/dodal/utils.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib
 import inspect
 import os
@@ -37,6 +38,7 @@
 )
 from ophyd.device import Device as OphydV1Device
 from ophyd_async.core import Device as OphydV2Device
+from ophyd_async.core import NotConnected
 
 import dodal.log
 
@@ -301,6 +303,45 @@ def is_v1_device_type(obj: type[Any]) -> bool:
     return is_class and follows_protocols and not is_v2_device_type(obj)
 
 
+def filter_ophyd_devices(
+    devices: Mapping[str, AnyDevice],
+    allow_extra_types: bool = False,
+) -> tuple[Mapping[str, OphydV1Device], Mapping[str, OphydV2Device]]:
+    ophyd_devices = {}
+    ophyd_async_devices = {}
+    for name, device in devices.items():
+        if isinstance(devices, OphydV1Device):
+            ophyd_devices[name] = device
+        elif isinstance(device, OphydV2Device):
+            ophyd_async_devices[name] = device
+        elif not allow_extra_types:
+            raise ValueError(
+                f"{device} (named {name}) is not an ophyd or ophyd-async device"
+            )
+    return ophyd_devices, ophyd_async_devices
+
+
+async def connect_all_ophyd_async_devices(
+    devices: Mapping[str, OphydV2Device],
+    mock: bool,
+) -> tuple[Mapping[str, OphydV2Device], Mapping[str, Exception]]:
+    coros = {name: device.connect(mock=mock) for name, device in devices.items()}
+    results = await asyncio.gather(*coros.values(), return_exceptions=True)
+
+    successes = {}
+    exceptions = {}
+
+    for name, result in zip(coros, results, strict=False):
+        if isinstance(result, Exception):
+            exceptions[name] = result
+        else:
+            successes[name] = devices[name]
+    if exceptions:
+        raise NotConnected(exceptions)
+
+    return successes, exceptions
+
+
 def get_beamline_based_on_environment_variable() -> ModuleType:
     """
     Gets the dodal module for the current beamline, as specified by the