Use own kafka callback

Author: Tom-Willemsen

doc/architectural_decisions/009-kafka-streaming.md

@@ -23,7 +23,8 @@ central Kafka instance is not currently considered "reliable" in an experiment c
 streaming the documents will allow testing to be done. Kafka will eventually be deployed in a "reliable"
 way accessible to each instrument.
 - We will encode messages from bluesky using `msgpack` (with the `msgpack-numpy` extension), because:
-  - It is the default encoder used by the upstream `bluesky-kafka` integration
+  - It is the default encoder used by the upstream `bluesky-kafka` integration (though we are not using
+that integration directly, as it has been soft-deprecated upstream)
   - It is a schema-less encoder, meaning we do not have to write/maintain fixed schemas for all the
 documents allowed by `event-model`
   - It has reasonable performance in terms of encoding speed and message size
@@ -41,7 +42,7 @@ Encoding bluesky documents into JSON and then wrapping them in the
 was considered.
 
 We chose `msgpack` instead of json strings + flatbuffers because:
-- It is more standard in the bluesky community (e.g. it is the default used in `bluesky-kafka`)
+- It is more standard in the bluesky community
 - Bluesky events will be streamed to a dedicated topic, which is unlikely to be confused with data
 using any other schema.
 

doc/dev/kafka.md

@@ -1,6 +1,6 @@
 # Kafka
 
-`ibex_bluesky_core` uses [the `bluesky-kafka` library](https://github.com/bluesky/bluesky-kafka) to send documents
+`ibex_bluesky_core` uses {py:obj}`~ibex_bluesky_core.callbacks.KafkaCallback` to send documents
 emitted by the {py:obj}`~bluesky.run_engine.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.

pyproject.toml

@@ -41,7 +41,7 @@ classifiers = [
 
 dependencies = [
   "bluesky",  # Bluesky framework
-  "bluesky-kafka",  # Bluesky-kafka integration
+  "confluent-kafka",  # Kafka producer
   "ophyd-async[ca] == 0.14.0",  # Device abstraction. When changing, also change in doc/conf.py
   "lmfit",  # Fitting
   "matplotlib",  # Plotting

src/ibex_bluesky_core/callbacks/__init__.py

@@ -33,6 +33,7 @@
     LiveFit,
     LiveFitLogger,
 )
+from ibex_bluesky_core.callbacks._kafka import KafkaCallback
 from ibex_bluesky_core.callbacks._plotting import LivePColorMesh, LivePlot, PlotPNGSaver, show_plot
 from ibex_bluesky_core.callbacks._utils import get_default_output_path
 from ibex_bluesky_core.fitting import FitMethod
@@ -49,6 +50,7 @@
     "DocLoggingCallback",
     "HumanReadableFileCallback",
     "ISISCallbacks",
+    "KafkaCallback",
     "LiveFit",
     "LiveFitLogger",
     "LivePColorMesh",

src/ibex_bluesky_core/callbacks/_kafka.py

@@ -0,0 +1,80 @@
+import logging
+import os
+import socket
+from typing import Any
+
+import msgpack_numpy
+from bluesky.callbacks import CallbackBase
+from confluent_kafka import Producer
+
+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 KafkaCallback(CallbackBase):
+    """Forward all bluesky documents to Kafka.
+
+    Documents are sent to kafka encoded using the MsgPack format with
+    the ``msgpack_numpy`` extension to allow efficiently encoding arrays.
+
+    .. note::
+
+        This callback is automatically configured by
+        :py:obj:`ibex_bluesky_core.run_engine.get_run_engine`, and does not need
+        to be configured manually.
+    """
+
+    def __init__(
+        self,
+        *,
+        bootstrap_servers: list[str] | None = None,
+        topic: str | None = None,
+        key: str,
+        kafka_config: dict[str, Any],
+    ) -> None:
+        super().__init__()
+
+        self._topic = topic or get_kafka_topic_name()
+        self._key = key
+
+        if "bootstrap.servers" in kafka_config:
+            raise ValueError(
+                "Do not specify bootstrap.servers in kafka config, use bootstrap_servers argument."
+            )
+
+        if bootstrap_servers is None:
+            bootstrap_servers = [
+                os.environ.get("IBEX_BLUESKY_CORE_KAFKA_BROKER", DEFAULT_KAFKA_BROKER)
+            ]
+
+        kafka_config["bootstrap.servers"] = ",".join(bootstrap_servers)
+
+        self._producer = Producer(kafka_config)
+
+    def __call__(
+        self, name: str, doc: dict[str, Any], validate: bool = False
+    ) -> tuple[str, dict[str, Any]]:
+        try:
+            data = msgpack_numpy.dumps([name, doc])
+            self._producer.produce(topic=self._topic, key=self._key, value=data)
+        except Exception:
+            # If we can't produce to kafka, log and carry on. We don't want
+            # kafka failures to kill a scan - kafka is currently considered
+            # 'non-critical'.
+            logger.exception("Failed to publish Kafka message")
+
+        return name, doc

src/ibex_bluesky_core/run_engine/__init__.py

@@ -3,46 +3,27 @@
 import asyncio
 import functools
 import logging
-import os
-import socket
 from collections.abc import Generator
 from functools import cache
 from threading import Event, Lock
 from typing import Any, cast
 
 import bluesky.preprocessors as bpp
-import msgpack_numpy
 from bluesky.run_engine import RunEngine, RunEngineResult
 from bluesky.utils import DuringTask, Msg, RunEngineControlException, RunEngineInterrupted
-from bluesky_kafka import Publisher
 
-from ibex_bluesky_core.callbacks import DocLoggingCallback
+from ibex_bluesky_core.callbacks import DocLoggingCallback, KafkaCallback
 from ibex_bluesky_core.plan_stubs import CALL_QT_AWARE_MSG_KEY, CALL_SYNC_MSG_KEY
 from ibex_bluesky_core.preprocessors import add_rb_number_processor
 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
 from ibex_bluesky_core.version import version
 
-__all__ = ["get_kafka_topic_name", "get_run_engine", "run_plan"]
+__all__ = ["get_run_engine", "run_plan"]
 
 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.
@@ -120,12 +101,9 @@ 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),
+    kafka_callback = KafkaCallback(
         key="doc",
-        serializer=msgpack_numpy.dumps,
-        producer_config={
+        kafka_config={
             "enable.idempotence": True,
             "log_level": 0,
             "log.connection.close": False,

tests/callbacks/test_kafka.py

@@ -0,0 +1,35 @@
+import re
+from unittest import mock
+
+import pytest
+
+from ibex_bluesky_core.callbacks._kafka import KafkaCallback, get_kafka_topic_name
+
+
+def test_get_kafka_topic_name():
+    with mock.patch("ibex_bluesky_core.callbacks._kafka.os.environ.get", return_value="FOO"):
+        assert get_kafka_topic_name() == "FOO_bluesky"
+
+    with mock.patch("ibex_bluesky_core.callbacks._kafka.os.environ.get", return_value="NDXBAR"):
+        assert get_kafka_topic_name() == "BAR_bluesky"
+
+    with mock.patch("ibex_bluesky_core.callbacks._kafka.os.environ.get", return_value="NDHBAZ"):
+        assert get_kafka_topic_name() == "BAZ_bluesky"
+
+
+def test_init_kafka_callback_with_duplicate_bootstrap_servers():
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Do not specify bootstrap.servers in kafka config, use bootstrap_servers argument."
+        ),
+    ):
+        KafkaCallback(bootstrap_servers=["abc"], kafka_config={"bootstrap.servers": "foo"}, key="")
+
+
+def test_exceptions_suppressed():
+    cb = KafkaCallback(bootstrap_servers=["abc"], kafka_config={}, key="")
+    with mock.patch(
+        "ibex_bluesky_core.callbacks._kafka.msgpack_numpy.dumps", side_effect=ValueError
+    ):
+        cb("start", {})

tests/test_run_engine.py

@@ -3,7 +3,6 @@
 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
@@ -12,7 +11,7 @@
 from bluesky.run_engine import RunEngineResult
 from bluesky.utils import Msg, RequestAbort, RunEngineInterrupted
 
-from ibex_bluesky_core.run_engine import _DuringTask, get_kafka_topic_name, get_run_engine, run_plan
+from ibex_bluesky_core.run_engine import _DuringTask, get_run_engine, run_plan
 from ibex_bluesky_core.version import version
 
 
@@ -147,14 +146,3 @@ 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"