Repository: bluesky/bluesky (python)

Issue:
Add metrics on how long waits took to the RunEngine
In general we would like to be able to get tracing and metric information out of the RunEngine. 

The specific MX DLS use case is that we would like to have good data on where the plans are waiting for hardware to finish doing things. 

## Expected Behavior
As a plan runs through information about how long `wait` messages took should be easily externally available This should be invisible to the plan and turned on via some flag before running the plan e.g. in the RunEngine.

## Possible Solution
From some discussion it seems that https://opentelemetry.io/docs/languages/python/ is the agreed module to use for this. Further discussion may be needed for specifics


Pull request description:
Add OpenTelemetry tracing to some RunEngine methods
## Description

Submitted together with https://github.com/bluesky/ophyd/pull/1189

- adds `opentelemetry-api` as a dependency
- wraps some potentially long-running `RunEngine` methods (`wait()`,`set()`,`complete()`...) in tracing spans
- creates spans at `_open_run` which are ended on `_close_run` (or abort, halt etc)
- attaches some metadata about each of the above

## Motivation and Context
Fixes #1682 

## How Has This Been Tested?
The `api` component of `OpenTelemetry` is designed to not really do anything on its own - I welcome suggestions about tests to include here.
This was tested by setting up a `TracingProvider` in Hyperion and exporting the results from running our unit tests to a Jaeger container.

Making use of these features in Hyperion is currently in [this branch](https://github.com/DiamondLightSource/hyperion/tree/bluesky_1682_improve_tracing), and an example of how to set up a tracing provider can be found [here](https://github.com/DiamondLightSource/hyperion/blob/bluesky_1682_improve_tracing/src/hyperion/tracing.py).



![screenshot-2024-04-17-16:04:56](https://github.com/bluesky/bluesky/assets/115003895/87bb1212-534d-4ae3-954c-4b2840171401)

![screenshot-2024-04-17-16:06:23](https://github.com/bluesky/bluesky/assets/115003895/a00e49ec-a243-41f8-af99-6b07add4754c)

![screenshot-2024-04-17-16:30:56](https://github.com/bluesky/bluesky/assets/115003895/a5f00a22-3b05-4397-bd1e-311f368edd4a)




Patch:
diff --git a/docs/otel-tracing-demo.py b/docs/otel-tracing-demo.py
new file mode 100644
index 00000000..b2ef7224
--- /dev/null
+++ b/docs/otel-tracing-demo.py
@@ -0,0 +1,31 @@
+from opentelemetry import trace
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+from ophyd.sim import det1
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from bluesky.run_engine import RunEngine
+from bluesky.tracing import trace_plan, tracer
+
+# Service name is required for most backends,
+# and although it's not necessary for console export,
+# it's good to set service name anyways.
+resource = Resource(attributes={SERVICE_NAME: "bluesky-docs-example"})
+
+traceProvider = TracerProvider(resource=resource)
+processor = BatchSpanProcessor(ConsoleSpanExporter())
+traceProvider.add_span_processor(processor)
+trace.set_tracer_provider(traceProvider)
+
+RE = RunEngine()
+
+
+@bpp.run_decorator()
+@trace_plan(tracer, "demo plan")
+def test_tracing_plan():
+    yield from bps.abs_set(det1._motor, 5)
+
+
+RE(test_tracing_plan())
diff --git a/docs/otel-tracing.rst b/docs/otel-tracing.rst
new file mode 100644
index 00000000..4ccc7f93
--- /dev/null
+++ b/docs/otel-tracing.rst
@@ -0,0 +1,150 @@
+Tracing with Opentelemetry
+==========================
+
+.. role:: py(code)
+   :language: python
+
+Bluesky is instrumented with [OpenTelemetry](https://opentelemetry.io/) tracing span hooks on runs, and on potentially long-running RunEngine methods such as :py:`wait()` and :py:`collect()`. This can allow you to analyze where your plans are spending a lot of time.
+
+By itself, this doesn't do anything. To collect tracing output you should configure an exporter in your application or plan, and a collector to send the data to. You can find an example of how to do that [here](https://opentelemetry.io/docs/languages/python/exporters/#usage). Since the :py:`@start_as_current_span` decorator from the opentelemetry library doesn't work for generators, we also provide a :py:`@trace_plan` decorator in :py:`bluesky.tracing`
+
+Tracing messages from the :py:`RunEngine` are named :py:`"Bluesky RunEngine <method name>"`, e.g. :py:`"Bluesky RunEngine wait"`. Traces for runs are tagged with the :py:`success`, :py:`exit_status` as attributes, as well as the :py:`reason` if one is available. Traces for methods are tagged with the content of the message (:py:`msg.command`, :py:`msg.args`, :py:`msg.kwargs`, and :py:`msg.obj`). Traces for :py:`wait()` also log the :py:`group` if one was given, or set :py:`no_group_given` true if none was. Ophyd also has traces on :py:`Status` objects to easily record how long they live for.
+
+Examples:
+---------
+
+Using the following script, which runs a simple scan and sends traces to the console:
+
+.. literalinclude:: otel-tracing-demo.py
+   :language: python
+
+We obtain this trace data. Note that the innermost spans are closed first, and therefore printed to the console first. You can see that the :code:`"parent_id"` of the :code:`Status` span corresponds to the :code:`"span_id"` of the :code:`set()` span, and so on.
+
+.. code-block:: JSON
+
+   {
+      "name": "Ophyd Status",
+      "context": {
+         "trace_id": "0x6a5abd70e7f74967975ecf64a863a12d",
+         "span_id": "0x0f94466a988f72ea",
+         "trace_state": "[]"
+      },
+      "kind": "SpanKind.INTERNAL",
+      "parent_id": "0x4d6e0e523cfdfec8",
+      "start_time": "2024-06-18T08:40:56.378644Z",
+      "end_time": "2024-06-18T08:40:56.378896Z",
+      "status": {
+         "status_code": "UNSET"
+      },
+      "attributes": {
+         "status_type": "MoveStatus",
+         "settle_time": 0,
+         "no_timeout_given": true,
+         "object_repr": "MoveStatus(done=False, pos=motor1, elapsed=0.0, success=False, settle_time=0.0)",
+         "device_name": "motor1",
+         "device_type": "SynAxis",
+         "kwargs": "{}",
+         "target": 5,
+         "start_time": 1718700056.37859,
+         "start_pos ": 0,
+         "unit": "mm",
+         "positioner_name": "motor1",
+         "positioner": "SynAxis(prefix='', name='motor1', read_attrs=['readback', 'setpoint'], configuration_attrs=['velocity', 'acceleration'])",
+         "finish_time": 1718700056.3788693,
+         "finish_pos": 5.0
+      },
+      "events": [],
+      "links": [],
+      "resource": {
+         "attributes": {
+               "service.name": "bluesky-docs-example"
+         },
+         "schema_url": ""
+      }
+   }
+   {
+      "name": "Bluesky RunEngine set",
+      "context": {
+         "trace_id": "0x6a5abd70e7f74967975ecf64a863a12d",
+         "span_id": "0x4d6e0e523cfdfec8",
+         "trace_state": "[]"
+      },
+      "kind": "SpanKind.INTERNAL",
+      "parent_id": "0x6d4272c02f6990f5",
+      "start_time": "2024-06-18T08:40:56.378468Z",
+      "end_time": "2024-06-18T08:40:56.378995Z",
+      "status": {
+         "status_code": "UNSET"
+      },
+      "attributes": {
+         "msg.command": "set",
+         "msg.args": [
+               5
+         ],
+         "msg.kwargs": "{\"group\": null}",
+         "msg.obj": "SynAxis(prefix='', name='motor1', read_attrs=['readback', 'setpoint'], configuration_attrs=['velocity', 'acceleration'])"
+      },
+      "events": [],
+      "links": [],
+      "resource": {
+         "attributes": {
+               "service.name": "bluesky-docs-example"
+         },
+         "schema_url": ""
+      }
+   }
+   {
+      "name": "demo plan",
+      "context": {
+         "trace_id": "0x6a5abd70e7f74967975ecf64a863a12d",
+         "span_id": "0x6d4272c02f6990f5",
+         "trace_state": "[]"
+      },
+      "kind": "SpanKind.INTERNAL",
+      "parent_id": null,
+      "start_time": "2024-06-18T08:40:56.378255Z",
+      "end_time": "2024-06-18T08:40:56.379044Z",
+      "status": {
+         "status_code": "UNSET"
+      },
+      "attributes": {},
+      "events": [],
+      "links": [],
+      "resource": {
+         "attributes": {
+               "service.name": "bluesky-docs-example"
+         },
+         "schema_url": ""
+      }
+   }
+   {
+      "name": "Bluesky RunEngine run",
+      "context": {
+         "trace_id": "0xedc3dae09fb45d4982e0af29f053141c",
+         "span_id": "0xb54c64dcdebaa81a",
+         "trace_state": "[]"
+      },
+      "kind": "SpanKind.INTERNAL",
+      "parent_id": null,
+      "start_time": "2024-06-18T08:40:56.377673Z",
+      "end_time": "2024-06-18T08:40:56.379335Z",
+      "status": {
+         "status_code": "UNSET"
+      },
+      "attributes": {
+         "msg.command": "open_run",
+         "msg.args": [],
+         "msg.kwargs": "{}",
+         "msg.no_obj_given": true,
+         "exit_status": "success",
+         "reason": ""
+      },
+      "events": [],
+      "links": [],
+      "resource": {
+         "attributes": {
+               "service.name": "bluesky-docs-example"
+         },
+         "schema_url": ""
+      }
+   }
diff --git a/docs/userindex.rst b/docs/userindex.rst
index 09a6d791..ff28c885 100644
--- a/docs/userindex.rst
+++ b/docs/userindex.rst
@@ -19,6 +19,7 @@ User Documentation
    run_engine_api
    utils
    magics
+   otel-tracing
    from-pyepics-to-bluesky
    comparison-with-spec
    hardware-interfaces
diff --git a/pyproject.toml b/pyproject.toml
index d56bbf16..6e08a733 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,9 +18,11 @@ dependencies = [
     "cycler",
     "event-model>=1.19.8",
     "historydict",
+    "importlib-resources;python_version<'3.9'",
     "msgpack",
     "msgpack-numpy",
-    "numpy",
+    "numpy<2.0.0",
+    "opentelemetry-api",
     "toolz",
     "tqdm>=4.44",
     "typing-extensions>=4.0.0;python_version<'3.11'",
diff --git a/src/bluesky/protocols.py b/src/bluesky/protocols.py
index fa0a34ca..5408b3a6 100644
--- a/src/bluesky/protocols.py
+++ b/src/bluesky/protocols.py
@@ -30,6 +30,11 @@ from typing_extensions import TypedDict
 # Squashes warning
 Dtype = Dtype  # type: ignore
 
+try:
+    from typing import ParamSpec
+except ImportError:
+    from typing_extensions import ParamSpec  # type: ignore
+
 
 # TODO: these are not placed in Events by RE yet
 class ReadingOptional(TypedDict, total=False):
@@ -67,6 +72,7 @@ StreamAsset = Union[
 
 
 T = TypeVar("T")
+P = ParamSpec("P")
 SyncOrAsync = Union[T, Awaitable[T]]
 SyncOrAsyncIterator = Union[Iterator[T], AsyncIterator[T]]
 
diff --git a/src/bluesky/run_engine.py b/src/bluesky/run_engine.py
index 8d30c2a3..0f97ed0c 100644
--- a/src/bluesky/run_engine.py
+++ b/src/bluesky/run_engine.py
@@ -3,6 +3,7 @@ import concurrent
 import copy
 import functools
 import inspect
+import json
 import sys
 import threading
 import typing
@@ -17,6 +18,8 @@ from itertools import count
 from warnings import warn
 
 from event_model import DocumentNames
+from opentelemetry import trace
+from opentelemetry.trace import Span
 
 from bluesky._vendor.super_state_machine.errors import TransitionError
 from bluesky._vendor.super_state_machine.extras import PropertyMachine
@@ -34,9 +37,11 @@ from .protocols import (
     Stageable,
     Status,
     Stoppable,
+    T,
     Triggerable,
     check_supports,
 )
+from .tracing import tracer
 from .utils import (
     AsyncInput,
     CallbackRegistry,
@@ -58,6 +63,8 @@ from .utils import (
     warn_if_msg_args_or_kwargs,
 )
 
+_SPAN_NAME_PREFIX = "Bluesky RunEngine"
+
 current_task: typing.Callable[[typing.Optional[asyncio.AbstractEventLoop]], typing.Optional[asyncio.Task]]
 try:
     from asyncio import current_task
@@ -427,6 +434,8 @@ class RunEngine:
         # When set, RunEngine.__call__ should stop blocking.
         self._blocking_event = threading.Event()
 
+        self._run_tracing_spans: list[Span] = []
+
         # When cleared, RunEngine._run will pause until set.
         self._run_permit = None
 
@@ -1316,6 +1325,7 @@ class RunEngine:
         self._reason = reason
 
         self._exit_status = "abort"
+        self._destroy_open_run_tracing_spans()
 
         was_paused = self._state == "paused"
         self._state = "aborting"
@@ -1419,6 +1429,7 @@ class RunEngine:
         if self._state.is_idle:
             raise TransitionError("RunEngine is already idle.")
         print("Halting: skipping cleanup and marking exit_status as 'abort'...")
+        self._destroy_open_run_tracing_spans()
         self._interrupted = True
         was_paused = self._state == "paused"
         self._state = "halting"
@@ -1447,6 +1458,12 @@ class RunEngine:
             else:
                 self.log.debug("No 'stop' method available on %r", obj)
 
+    def _destroy_open_run_tracing_spans(self):
+        while len(self._run_tracing_spans):
+            _span = self._run_tracing_spans.pop()
+            _span.set_attribute("exit_status", "aborted")
+            _span.end()
+
     async def _run(self):
         """Pull messages from the plan, process them, send results back.
 
@@ -1812,6 +1829,11 @@ class RunEngine:
         where **kwargs are any additional metadata that should go into
         the RunStart document
         """
+        _span = tracer.start_span(f"{_SPAN_NAME_PREFIX} run")
+        _set_span_msg_attributes(_span, msg)
+
+        self._run_tracing_spans.append(_span)
+
         # TODO extract this from the Msg
         run_key = msg.run
         if run_key in self._run_bundlers:
@@ -1872,8 +1894,20 @@ class RunEngine:
             raise IllegalMessageSequence(ims_msg)
         ret = await current_run.close_run(msg)
         del self._run_bundlers[run_key]
+        self._close_run_trace(msg)
         return ret
 
+    def _close_run_trace(self, msg):
+        exit_status = msg.exit_status if hasattr(msg, "exit_status") else self._exit_status
+        reason = msg.reason if hasattr(msg, "reason") else self._reason
+        try:
+            _span: Span = self._run_tracing_spans.pop()
+            _span.set_attribute("exit_status", exit_status)
+            _span.set_attribute("reason", reason)
+            _span.end()
+        except IndexError:
+            logger.warning("No open traces left to close!")
+
     async def _create(self, msg):
         """Trigger the run engine to start bundling future obj.read() calls for
          an Event document
@@ -2145,6 +2179,7 @@ class RunEngine:
 
         return ret
 
+    @tracer.start_as_current_span(f"{_SPAN_NAME_PREFIX} complete")
     async def _complete(self, msg):
         """
         Tell a flyer, 'stop collecting, whenever you are ready'.
@@ -2161,6 +2196,7 @@ class RunEngine:
 
         where <GROUP> is a hashable identifier.
         """
+        _set_span_msg_attributes(trace.get_current_span(), msg)
         kwargs = dict(msg.kwargs)
         group = kwargs.pop("group", None)
         obj = check_supports(msg.obj, Flyable)
@@ -2187,6 +2223,7 @@ class RunEngine:
         self._status_objs[group].add(ret)
         return ret
 
+    @tracer.start_as_current_span(f"{_SPAN_NAME_PREFIX} collect")
     async def _collect(self, msg):
         """
         Collect data cached by a flyer and emit documents
@@ -2196,6 +2233,7 @@ class RunEngine:
             Msg('collect', flyer_object)
             Msg('collect', flyer_object, stream=True, return_payload=False, name="a_name")
         """
+        _set_span_msg_attributes(trace.get_current_span(), msg)
         run_key = msg.run
         if (
             current_run := self._run_bundlers.get(run_key, key_absence_sentinel := object())
@@ -2218,6 +2256,7 @@ class RunEngine:
         """
         return type(self)
 
+    @tracer.start_as_current_span(f"{_SPAN_NAME_PREFIX} set")
     async def _set(self, msg):
         """
         Set a device and cache the returned status object.
@@ -2231,6 +2270,7 @@ class RunEngine:
 
         where arguments are passed through to `obj.set(*args, **kwargs)`.
         """
+        _set_span_msg_attributes(trace.get_current_span(), msg)
         obj = check_supports(msg.obj, Movable)
         kwargs = dict(msg.kwargs)
         group = kwargs.pop("group", None)
@@ -2275,6 +2315,7 @@ class RunEngine:
         if self.waiting_hook is not None:
             self.waiting_hook(*args, **kwargs)
 
+    @tracer.start_as_current_span(f"{_SPAN_NAME_PREFIX} wait")
     async def _wait(self, msg):
         """Block progress until every object that was triggered or set
         with the keyword argument `group=<GROUP>` is done. Returns a boolean that is
@@ -2288,6 +2329,7 @@ class RunEngine:
 
         where ``<GROUP>`` is any hashable key and ``<MOVE_ON>`` is a boolean.
         """
+        _set_span_msg_attributes(trace.get_current_span(), msg)
         done = False  # boolean that tracks whether waiting is complete
         if msg.args:
             (group,) = msg.args
@@ -2295,6 +2337,10 @@ class RunEngine:
         else:
             group = msg.kwargs["group"]
             move_on = msg.kwargs.get("move_on", False)
+        if group:
+            trace.get_current_span().set_attribute("group", group)
+        else:
+            trace.get_current_span().set_attribute("no_group_given", True)
         futs = list(self._groups.pop(group, []))
         if futs:
             status_objs = self._status_objs.pop(group)
@@ -2773,6 +2819,13 @@ http://nsls-ii.github.io/bluesky/plans_intro.html#combining-plans
 """
 
 
+def _set_span_msg_attributes(span, msg):
+    span.set_attribute("msg.command", msg.command)
+    span.set_attribute("msg.args", msg.args)
+    span.set_attribute("msg.kwargs", json.dumps(msg.kwargs, default=repr))
+    span.set_attribute("msg.obj", repr(msg.obj)) if msg.obj else span.set_attribute("msg.no_obj_given", True)
+
+
 def _default_md_validator(md):
     if "sample" in md and not (hasattr(md["sample"], "keys") or isinstance(md["sample"], str)):
         raise ValueError(
@@ -2830,9 +2883,6 @@ def in_bluesky_event_loop() -> bool:
         return loop is _bluesky_event_loop
 
 
-T = typing.TypeVar("T")
-
-
 def call_in_bluesky_event_loop(coro: typing.Awaitable[T], timeout: typing.Optional[float] = None) -> T:
     if _bluesky_event_loop is None or not _bluesky_event_loop.is_running():
         # Quell "coroutine never awaited" warnings
diff --git a/src/bluesky/tracing.py b/src/bluesky/tracing.py
new file mode 100644
index 00000000..b1ee3379
--- /dev/null
+++ b/src/bluesky/tracing.py
@@ -0,0 +1,23 @@
+import functools
+from typing import Callable, cast
+
+from opentelemetry.trace import Tracer, get_tracer
+
+from .protocols import P
+from .utils import MsgGenerator
+
+tracer = get_tracer(__name__)
+
+
+def trace_plan(tracer: Tracer, span_name: str) -> Callable[[Callable[P, MsgGenerator]], Callable[P, MsgGenerator]]:
+    """Wraps a generator function in tracer.start_as_current_span(span_name)"""
+
+    def wrap(f: Callable[P, MsgGenerator]) -> Callable[P, MsgGenerator]:
+        @functools.wraps(f)
+        def wrap_f(*args: P.args, **kwargs: P.kwargs):
+            with tracer.start_as_current_span(span_name):
+                yield from f(*args, **kwargs)
+
+        return cast(Callable[P, MsgGenerator], wrap_f)
+
+    return wrap
diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py
index bdcd13c3..fabb780b 100644
--- a/src/bluesky/utils/__init__.py
+++ b/src/bluesky/utils/__init__.py
@@ -24,6 +24,7 @@ from typing import (
     AsyncIterator,
     Callable,
     Dict,
+    Generator,
     List,
     Optional,
     Tuple,
@@ -88,6 +89,9 @@ class Msg(namedtuple("Msg_base", ["command", "obj", "args", "kwargs", "run"])):
         )
 
 
+MsgGenerator = Generator[Msg, Any, None]
+
+
 class RunEngineControlException(Exception):
     """Exception for signaling within the RunEngine."""
 


Test patch:
diff --git a/src/bluesky/tests/test_tracing.py b/src/bluesky/tests/test_tracing.py
new file mode 100644
index 00000000..835a2699
--- /dev/null
+++ b/src/bluesky/tests/test_tracing.py
@@ -0,0 +1,24 @@
+from typing import Generator
+
+from opentelemetry.trace import get_current_span
+
+from bluesky.plan_stubs import sleep
+from bluesky.tracing import trace_plan, tracer
+
+
+@trace_plan(tracer, "test_plan")
+def _test_plan():
+    yield from sleep(0)
+
+
+def test_trace_plan_wrapper_gives_back_generator():
+    assert isinstance(_test_plan(), Generator)
+
+
+def test_trace_plan_wrapper_opens_span_but_doesnt_do_anything_without_trace_provider(RE):
+    @trace_plan(tracer, "test_plan_2")
+    def test_plan_2():
+        yield from sleep(0)
+        assert not get_current_span().is_recording()
+
+    RE(test_plan_2())