Propgate worker errors back to user via message bus (#99)

Author: callumforrester

src/blueapi/cli/amq.py

@@ -1,5 +1,5 @@
 import threading
-from typing import Any, Callable, List, Mapping, Optional, TypeVar, Union
+from typing import Any, Callable, Mapping, Optional, TypeVar
 
 from blueapi.messaging import MessageContext, MessagingTemplate
 from blueapi.service.model import (
@@ -9,11 +9,14 @@
     PlanResponse,
     TaskResponse,
 )
-from blueapi.worker import TaskEvent
+from blueapi.worker import ProgressEvent, WorkerEvent
 
 T = TypeVar("T")
 
-_Json = Union[List[Any], Mapping[str, Any]]
+
+class BlueskyRemoteError(Exception):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
 
 
 class AmqClient:
@@ -26,19 +29,33 @@ def run_plan(
         self,
         name: str,
         params: Mapping[str, Any],
-        on_event: Optional[Callable[[TaskEvent], None]] = None,
+        on_event: Optional[Callable[[WorkerEvent], None]] = None,
+        on_progress_event: Optional[Callable[[ProgressEvent], None]] = None,
         timeout: Optional[float] = None,
     ) -> str:
         complete = threading.Event()
 
-        def on_event_wrapper(ctx: MessageContext, event: TaskEvent) -> None:
+        def on_event_wrapper(ctx: MessageContext, event: WorkerEvent) -> None:
             if on_event is not None:
                 on_event(event)
-            if event.is_task_terminated():
+
+            if event.is_complete():
                 complete.set()
+                if event.is_error():
+                    raise BlueskyRemoteError(str(event.errors) or "Unknown error")
+
+        def on_progress_event_wrapper(
+            ctx: MessageContext, event: ProgressEvent
+        ) -> None:
+            if on_progress_event is not None:
+                on_progress_event(event)
 
         self.app.subscribe(
-            self.app.destinations.topic("public.worker.event.task"), on_event_wrapper
+            self.app.destinations.topic("public.worker.event"), on_event_wrapper
+        )
+        self.app.subscribe(
+            self.app.destinations.topic("public.worker.event.progress"),
+            on_progress_event_wrapper,
         )
         # self.app.send("worker.run", {"name": name, "params": params})
         task_response = self.app.send_and_recieve(

src/blueapi/cli/cli.py

@@ -1,18 +1,16 @@
-import itertools
 import json
 import logging
 from pathlib import Path
-from typing import Dict, Mapping, Optional
+from typing import Optional
 
 import click
-from tqdm import tqdm
 
 from blueapi import __version__
 from blueapi.config import StompConfig
 from blueapi.messaging import StompMessagingTemplate
-from blueapi.worker import StatusView, TaskEvent
 
 from .amq import AmqClient
+from .updates import CliEventRenderer
 
 
 @click.group(invoke_without_command=True)
@@ -91,52 +89,11 @@ def get_devices(ctx) -> None:
 @click.pass_context
 def run_plan(ctx, name: str, parameters: str) -> None:
     client: AmqClient = ctx.obj["client"]
-
-    renderer = ProgressBarRenderer()
-
-    def handle_event(event: TaskEvent) -> None:
-        renderer.update(event.statuses)
-        if event.is_task_terminated():
-            print("")
-            print("")
-            print("")
-            print("DONE")
-
-    client.run_plan(name, json.loads(parameters), handle_event, timeout=120.0)
-
-
-_BAR_FMT = "{desc}: |{bar}| {percentage:3.0f}% [{elapsed}/{remaining}]"
-
-
-class ProgressBarRenderer:
-    _bars: Dict[str, tqdm]
-    _count: itertools.count
-
-    def __init__(self) -> None:
-        self._bars = {}
-        self._count = itertools.count()
-
-    def update(self, status_view: Mapping[str, StatusView]) -> None:
-        for name, view in status_view.items():
-            if name not in self._bars:
-                pos = next(self._count)
-                self._bars[name] = tqdm(
-                    position=pos,
-                    total=1.0,
-                    initial=0.0,
-                    bar_format=_BAR_FMT,
-                )
-            self._update(name, view)
-
-    def _update(self, name: str, view: StatusView) -> None:
-        bar = self._bars[name]
-        if (
-            view.current is not None
-            and view.target is not None
-            and view.initial is not None
-            and view.percentage is not None
-            and view.time_elapsed is not None
-        ):
-            bar.desc = view.display_name
-            bar.update(view.percentage - bar.n)
-            bar.unit = view.unit
+    renderer = CliEventRenderer()
+    client.run_plan(
+        name,
+        json.loads(parameters),
+        renderer.on_worker_event,
+        renderer.on_progress_event,
+        timeout=120.0,
+    )

src/blueapi/cli/updates.py

@@ -0,0 +1,78 @@
+import itertools
+from typing import Dict, Mapping, Optional, Union
+
+from tqdm import tqdm
+
+from blueapi.worker import ProgressEvent, StatusView, WorkerEvent
+
+_BAR_FMT = "{desc}: |{bar}| {percentage:3.0f}% [{elapsed}/{remaining}]"
+
+
+class ProgressBarRenderer:
+    _bars: Dict[str, tqdm]
+    _count: itertools.count
+
+    def __init__(self) -> None:
+        self._bars = {}
+        self._count = itertools.count()
+
+    def update(self, status_view: Mapping[str, StatusView]) -> None:
+        for name, view in status_view.items():
+            if name not in self._bars:
+                pos = next(self._count)
+                self._bars[name] = tqdm(
+                    position=pos,
+                    total=1.0,
+                    initial=0.0,
+                    bar_format=_BAR_FMT,
+                )
+            self._update(name, view)
+
+    def _update(self, name: str, view: StatusView) -> None:
+        bar = self._bars[name]
+        if (
+            view.current is not None
+            and view.target is not None
+            and view.initial is not None
+            and view.percentage is not None
+            and view.time_elapsed is not None
+        ):
+            bar.desc = view.display_name
+            bar.update(view.percentage - bar.n)
+            bar.unit = view.unit
+
+
+class CliEventRenderer:
+    _task_name: Optional[str]
+    _pbar_renderer: ProgressBarRenderer
+
+    def __init__(
+        self,
+        task_name: Optional[str] = None,
+        pbar_renderer: Optional[ProgressBarRenderer] = None,
+    ) -> None:
+        self._task_name = task_name
+        if pbar_renderer is None:
+            pbar_renderer = ProgressBarRenderer()
+        self._pbar_renderer = pbar_renderer
+
+    def on_progress_event(self, event: ProgressEvent) -> None:
+        if self._relates_to_task(event):
+            self._pbar_renderer.update(event.statuses)
+
+    def on_worker_event(self, event: WorkerEvent) -> None:
+        if self._relates_to_task(event):
+            print(str(event.state))
+
+    def _relates_to_task(self, event: Union[WorkerEvent, ProgressEvent]) -> bool:
+        if self._task_name is None:
+            return True
+        elif isinstance(event, WorkerEvent):
+            return (
+                event.task_status is not None
+                and event.task_status.task_name == self._task_name
+            )
+        elif isinstance(event, ProgressEvent):
+            return event.task_name == self._task_name
+        else:
+            return False

src/blueapi/service/app.py

@@ -1,13 +1,13 @@
 import logging
 import uuid
 from pathlib import Path
-from typing import Optional
+from typing import Mapping, Optional
 
 from blueapi.config import ApplicationConfig
-from blueapi.core import BlueskyContext, DataEvent
+from blueapi.core import BlueskyContext, EventStream
 from blueapi.messaging import MessageContext, MessagingTemplate, StompMessagingTemplate
 from blueapi.utils import ConfigLoader
-from blueapi.worker import RunEngineWorker, RunPlan, TaskEvent, Worker, WorkerEvent
+from blueapi.worker import RunEngineWorker, RunPlan, Worker
 
 from .model import (
     DeviceModel,
@@ -35,9 +35,20 @@ def __init__(self, config: ApplicationConfig) -> None:
 
     def run(self) -> None:
         logging.basicConfig(level=self._config.logging.level)
-        self._worker.worker_events.subscribe(self._on_worker_event)
-        self._worker.task_events.subscribe(self._on_task_event)
-        self._worker.data_events.subscribe(self._on_data_event)
+
+        self._publish_event_streams(
+            {
+                self._worker.worker_events: self._template.destinations.topic(
+                    "public.worker.event"
+                ),
+                self._worker.progress_events: self._template.destinations.topic(
+                    "public.worker.event.progress"
+                ),
+                self._worker.data_events: self._template.destinations.topic(
+                    "public.worker.event.data"
+                ),
+            }
+        )
 
         self._template.connect()
 
@@ -47,20 +58,14 @@ def run(self) -> None:
 
         self._worker.run_forever()
 
-    def _on_worker_event(self, event: WorkerEvent) -> None:
-        self._template.send(
-            self._template.destinations.topic("public.worker.event"), event
-        )
-
-    def _on_task_event(self, event: TaskEvent) -> None:
-        self._template.send(
-            self._template.destinations.topic("public.worker.event.task"), event
-        )
+    def _publish_event_streams(
+        self, streams_to_destinations: Mapping[EventStream, str]
+    ) -> None:
+        for stream, destination in streams_to_destinations.items():
+            self._publish_event_stream(stream, destination)
 
-    def _on_data_event(self, event: DataEvent) -> None:
-        self._template.send(
-            self._template.destinations.topic("public.worker.event.data"), event
-        )
+    def _publish_event_stream(self, stream: EventStream, destination: str) -> None:
+        stream.subscribe(lambda event: self._template.send(destination, event))
 
     def _on_run_request(self, message_context: MessageContext, task: RunPlan) -> None:
         name = str(uuid.uuid1())

src/blueapi/startup/example.py

@@ -1,14 +1,16 @@
 from ophyd.sim import Syn2DGauss, SynGauss, SynSignal
 
 from blueapi.plans import *  # noqa: F401, F403
-from blueapi.service.simmotor import SynAxisWithMotionEvents
+
+from .simmotor import BrokenSynAxis, SynAxisWithMotionEvents
 
 x = SynAxisWithMotionEvents(name="x", delay=1.0, events_per_move=8)
 y = SynAxisWithMotionEvents(name="y", delay=3.0, events_per_move=24)
 z = SynAxisWithMotionEvents(name="z", delay=2.0, events_per_move=16)
 theta = SynAxisWithMotionEvents(
     name="theta", delay=0.2, events_per_move=12, egu="degrees"
 )
+x_err = BrokenSynAxis(name="x_err", timeout=1.0)
 sample_pressure = SynAxisWithMotionEvents(
     name="sample_pressure", delay=30.0, events_per_move=128, egu="MPa", value=0.101
 )

src/blueapi/service/simmotor.py src/blueapi/startup/simmotor.pysrc/blueapi/service/simmotor.py renamed to src/blueapi/startup/simmotor.py

@@ -3,7 +3,7 @@
 from typing import Callable, Optional
 
 from ophyd.sim import SynAxis
-from ophyd.status import MoveStatus
+from ophyd.status import MoveStatus, Status
 
 
 class SynAxisWithMotionEvents(SynAxis):
@@ -81,3 +81,14 @@ def sleep_and_finish():
         threading.Thread(target=sleep_and_finish, daemon=True).start()
 
         return st
+
+
+class BrokenSynAxis(SynAxis):
+    _timeout: float
+
+    def __init__(self, *, timeout: float, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._timeout = timeout
+
+    def set(self, value: float) -> Status:
+        return Status(timeout=self._timeout)

src/blueapi/worker/__init__.py

@@ -1,18 +1,18 @@
-from .event import RunnerState, StatusView, TaskEvent, WorkerEvent
+from .event import ProgressEvent, StatusView, TaskStatus, WorkerEvent, WorkerState
 from .multithread import run_worker_in_own_thread
 from .reworker import RunEngineWorker
-from .task import RunPlan, Task, TaskState
+from .task import RunPlan, Task
 from .worker import Worker
 
 __all__ = [
     "run_worker_in_own_thread",
     "RunEngineWorker",
     "Task",
-    "TaskState",
     "Worker",
     "RunPlan",
     "WorkerEvent",
-    "RunnerState",
-    "TaskEvent",
+    "WorkerState",
     "StatusView",
+    "ProgressEvent",
+    "TaskStatus",
 ]