Stream bluesky documents to kafka

doc/dev/kafka.md

@@ -0,0 +1,14 @@
+# Kafka
+
+`ibex_bluesky_core` uses [the `bluesky-kafka` library](https://github.com/bluesky/bluesky-kafka) to send documents
+emitted by the `RunEngine` to kafka. The kafka callback is automatically added by
+{py:obj}`ibex_bluesky_core.run_engine.get_run_engine`, and so no user configuration is required - the callback is always
+enabled.
+
+Documents are encoded using [the `msgpack` format](https://msgpack.org/index.html).
+
+The kafka broker to send to can be controlled using the `IBEX_BLUESKY_CORE_KAFKA_BROKER` environment variable, if
+an instrument needs to override the default. The kafka topic will be `<INSTRUMENT>_bluesky`, where `INSTRUMENT` is the
+instrument name with any NDX or NDH prefix stripped.
+
+The message key will always be `doc` for bluesky documents; specifying a non-null key enforces message ordering.

pyproject.toml

@@ -41,6 +41,7 @@ classifiers = [
 
 dependencies = [
   "bluesky",  # Bluesky framework
+  "bluesky-kafka",  # Bluesky-kafka integration
   "ophyd-async[ca] == 0.12.3",  # Device abstraction
   "matplotlib",  # Plotting
   "lmfit",  # Fitting

src/ibex_bluesky_core/log.py

@@ -18,7 +18,7 @@
 # Find the log directory, if already set in the environment, else use the default
 log_location = os.environ.get("IBEX_BLUESKY_CORE_LOGS", DEFAULT_LOG_FOLDER)
 
-INTERESTING_LOGGER_NAMES = ["ibex_bluesky_core", "bluesky", "ophyd_async"]
+INTERESTING_LOGGER_NAMES = ["ibex_bluesky_core", "bluesky", "ophyd_async", "bluesky_kafka"]
 
 
 @cache

src/ibex_bluesky_core/run_engine/__init__.py

@@ -9,15 +9,21 @@
 from typing import Any, cast
 
 import bluesky.preprocessors as bpp
+import msgpack
 from bluesky.run_engine import RunEngine, RunEngineResult
 from bluesky.utils import DuringTask, Msg, RunEngineControlException, RunEngineInterrupted
 
 from ibex_bluesky_core.callbacks import DocLoggingCallback
 from ibex_bluesky_core.preprocessors import add_rb_number_processor
 
-__all__ = ["get_run_engine", "run_plan"]
+__all__ = ["get_kafka_topic_name", "get_run_engine", "run_plan"]
 
 
+import os
+import socket
+
+from bluesky_kafka import Publisher
+
 from ibex_bluesky_core.plan_stubs import CALL_QT_AWARE_MSG_KEY, CALL_SYNC_MSG_KEY
 from ibex_bluesky_core.run_engine._msg_handlers import call_qt_aware_handler, call_sync_handler
 from ibex_bluesky_core.utils import is_matplotlib_backend_qt
@@ -26,6 +32,21 @@
 logger = logging.getLogger(__name__)
 
 
+DEFAULT_KAFKA_BROKER = "livedata.isis.cclrc.ac.uk:31092"
+
+
+def get_kafka_topic_name() -> str:
+    """Get the name of the bluesky kafka topic for this machine."""
+    computer_name = os.environ.get("COMPUTERNAME", socket.gethostname()).upper()
+    computer_name = computer_name.upper()
+    if computer_name.startswith(("NDX", "NDH")):
+        name = computer_name[3:]
+    else:
+        name = computer_name
+
+    return f"{name}_bluesky"
+
+
 class _DuringTask(DuringTask):
     def block(self, blocking_event: Event) -> None:
         """On windows, event.wait() on the main thread is not interruptible by a CTRL-C.
@@ -103,6 +124,15 @@ def get_run_engine() -> RunEngine:
     log_callback = DocLoggingCallback()
     RE.subscribe(log_callback)
 
+    kafka_callback = Publisher(
+        topic=get_kafka_topic_name(),
+        bootstrap_servers=os.environ.get("IBEX_BLUESKY_CORE_KAFKA_BROKER", DEFAULT_KAFKA_BROKER),
+        key="doc",
+        serializer=msgpack.dumps,
+        producer_config={"enable.idempotence": True},
+    )
+    RE.subscribe(kafka_callback)
+
     RE.register_command(CALL_SYNC_MSG_KEY, call_sync_handler)
     RE.register_command(CALL_QT_AWARE_MSG_KEY, call_qt_aware_handler)
 

tests/test_run_engine.py

@@ -3,6 +3,7 @@
 import threading
 from collections.abc import Generator
 from typing import Any
+from unittest import mock
 from unittest.mock import MagicMock
 
 import bluesky.plan_stubs as bps
@@ -11,7 +12,7 @@
 from bluesky.run_engine import RunEngineResult
 from bluesky.utils import Msg, RequestAbort, RunEngineInterrupted
 
-from ibex_bluesky_core.run_engine import _DuringTask, get_run_engine, run_plan
+from ibex_bluesky_core.run_engine import _DuringTask, get_kafka_topic_name, get_run_engine, run_plan
 from ibex_bluesky_core.version import version
 
 
@@ -146,3 +147,14 @@ def plan():
     result = run_plan(plan())
     assert result.plan_result == "happy_path_result"
     assert result.exit_status == "success"
+
+
+def test_get_kafka_topic_name():
+    with mock.patch("ibex_bluesky_core.run_engine.os.environ.get", return_value="FOO"):
+        assert get_kafka_topic_name() == "FOO_bluesky"
+
+    with mock.patch("ibex_bluesky_core.run_engine.os.environ.get", return_value="NDXBAR"):
+        assert get_kafka_topic_name() == "BAR_bluesky"
+
+    with mock.patch("ibex_bluesky_core.run_engine.os.environ.get", return_value="NDHBAZ"):
+        assert get_kafka_topic_name() == "BAZ_bluesky"