Support streaming RE waiting hook watcher updates over a new ZMQ topic

pixi.toml

@@ -3,9 +3,10 @@ authors = ["Dmitri Gavrilov <gavrilov.dvs@gmail.com>"]
 channels = ["conda-forge"]
 name = "bluesky-queueserver"
 platforms = ["linux-64"]
-version = "0.0.24"
+version = "0.0.25"
 
 [tasks]
+lint = "pre-commit run --all-files"
 
 [dependencies]
 python = "*"

src/bluesky_queueserver/__init__.py

@@ -12,6 +12,8 @@
 from .manager.output_streaming import (  # noqa: E402, F401
     ReceiveConsoleOutput,
     ReceiveConsoleOutputAsync,
+    ReceiveProgressInfo,
+    ReceiveProgressInfoAsync,
     ReceiveSystemInfo,
     ReceiveSystemInfoAsync,
 )

src/bluesky_queueserver/manager/config.py

@@ -185,6 +185,8 @@ def parse_configs(config_path):
     "zmq_info_addr": "network/zmq_info_addr",
     "zmq_encoding": "network/zmq_encoding",
     "zmq_publish_console": "network/zmq_publish_console",
+    "zmq_publish_info": "network/zmq_publish_info",
+    "zmq_publish_progress": "network/zmq_publish_progress",
     "redis_addr": "network/redis_addr",
     "redis_name_prefix": "network/redis_name_prefix",
     "ignore_invalid_plans": "startup/ignore_invalid_plans",
@@ -355,6 +357,18 @@ def __init__(self, *, parser, args):
             value_cli=self._args_existing("zmq_publish_console"),
         )
 
+        self._settings["zmq_publish_info"] = self._get_param_boolean(
+            value_default=args.zmq_publish_info,
+            value_config=self._get_value_from_config("zmq_publish_info"),
+            value_cli=self._args_existing("zmq_publish_info"),
+        )
+
+        self._settings["zmq_publish_progress"] = self._get_param_boolean(
+            value_default=args.zmq_publish_progress,
+            value_config=self._get_value_from_config("zmq_publish_progress"),
+            value_cli=self._args_existing("zmq_publish_progress"),
+        )
+
         redis_addr = self._get_param(
             value_default=self._args.redis_addr,
             value_config=self._get_value_from_config("redis_addr"),

src/bluesky_queueserver/manager/config_schemas/config_schema.yml

@@ -21,6 +21,10 @@ properties:
         type: string
       zmq_publish_console:
         type: boolean
+      zmq_publish_info:
+        type: boolean
+      zmq_publish_progress:
+        type: boolean
       redis_addr:
         type: string
       redis_name_prefix:

src/bluesky_queueserver/manager/output_streaming.py

@@ -115,8 +115,29 @@ def push_info_to_msg_queue(*, key, msg, msg_queue):
     msg_queue.put(msg)
 
 
+def push_progress_to_msg_queue(*, msg, msg_queue):
+    """
+    Format a progress message and put it into the message queue. The message is published
+    to the ``progress`` channel on the ``QS_Progress`` 0MQ topic.
+
+    Parameters
+    ----------
+    msg : dict
+        The progress payload dictionary.
+    msg_queue : multiprocessing.Queue
+        Reference to the queue used for collecting messages.
+
+    Returns
+    -------
+    None
+    """
+    msg = {"channel": "progress", "time": ttime.time(), "msg": msg}
+    msg_queue.put(msg)
+
+
 _default_zmq_console_topic = "QS_Console"
 _default_zmq_info_topic = "QS_Info"
+_default_zmq_progress_topic = "QS_Progress"
 
 
 class PublishZMQStreamOutput:
@@ -135,8 +156,12 @@ class PublishZMQStreamOutput:
         The messages added to the queue will be automatically published to 0MQ socket.
     console_output_on : boolean
         Enable/disable printing console output to the terminal
-    zmq_publish_on : boolean
-        Enable/disable publishing console output to 0MQ socket
+    zmq_publish_console : boolean
+        Enable/disable publishing console output to 0MQ socket (``QS_Console`` topic).
+    zmq_publish_info : boolean
+        Enable/disable publishing info/status messages to 0MQ socket (``QS_Info`` topic).
+    zmq_publish_progress : boolean
+        Enable/disable publishing progress messages to 0MQ socket (``QS_Progress`` topic).
     zmq_publish_addr : str, None
         Address of 0MQ PUB socket for the publishing server. If ``None``, then
         the default address ``tcp://*:60625`` is used.
@@ -146,6 +171,8 @@ class PublishZMQStreamOutput:
         Name of the 0MQ topic where the console messages are published.
     zmq_topic_info : str
         Name of the 0MQ topic where the system information messages are published.
+    zmq_topic_progress : str
+        Name of the 0MQ topic where the progress messages are published.
     name : str
         Name of the thread where the messages are published.
     """
@@ -155,11 +182,14 @@ def __init__(
         *,
         msg_queue,
         console_output_on=True,
-        zmq_publish_on=True,
+        zmq_publish_console=True,
+        zmq_publish_info=True,
+        zmq_publish_progress=True,
         zmq_publish_addr=None,
         encoding="json",
         zmq_topic_console=_default_zmq_console_topic,
         zmq_topic_info=_default_zmq_info_topic,
+        zmq_topic_progress=_default_zmq_progress_topic,
         name="RE Console Output Publisher",
     ):
         self._thread_running = False  # Set True to exit the thread
@@ -168,7 +198,9 @@ def __init__(
         self._polling_timeout = 0.1  # in sec.
 
         self._console_output_on = console_output_on
-        self._zmq_publish_on = zmq_publish_on
+        self._zmq_publish_console = zmq_publish_console
+        self._zmq_publish_info = zmq_publish_info
+        self._zmq_publish_progress = zmq_publish_progress
 
         self._encoding = process_zmq_encoding_name(encoding)
 
@@ -177,22 +209,24 @@ def __init__(
         self._zmq_publish_addr = zmq_publish_addr
         self._zmq_topic_console = zmq_topic_console
         self._zmq_topic_info = zmq_topic_info
+        self._zmq_topic_progress = zmq_topic_progress
 
+        zmq_publish_on = zmq_publish_console or zmq_publish_info or zmq_publish_progress
         self._socket = None
-        if self._zmq_publish_on:
+        if zmq_publish_on:
             try:
                 context = zmq.Context()
                 self._socket = context.socket(zmq.PUB)
                 self._socket.bind(self._zmq_publish_addr)
             except Exception as ex:
                 logger.error(
-                    "Failed to create 0MQ socket at %s. Console output will not be published. Exception: %s",
+                    "Failed to create 0MQ socket at %s. Output will not be published. Exception: %s",
                     self._zmq_publish_addr,
                     ex,
                 )
 
-        if self._socket and self._zmq_publish_on:
-            logging.info("Publishing console output to 0MQ socket at %s", zmq_publish_addr)
+        if self._socket and zmq_publish_on:
+            logging.info("Publishing output to 0MQ socket at %s", zmq_publish_addr)
 
     def start(self):
         """
@@ -237,13 +271,15 @@ def _publish(self, payload):
             sys.__stdout__.write(payload["msg"])
             sys.__stdout__.flush()
 
-        if self._zmq_publish_on and self._socket:
-            if channel == "console":
+        if self._socket:
+            if channel == "console" and self._zmq_publish_console:
                 topic = self._zmq_topic_console
-            elif channel == "info":
+            elif channel == "info" and self._zmq_publish_info:
                 topic = self._zmq_topic_info
+            elif channel == "progress" and self._zmq_publish_progress:
+                topic = self._zmq_topic_progress
             else:
-                logger.error("Failed to publish the message: unsupported 0MQ channel %s.")
+                return
             payload = {k: payload[k] for k in ("time", "msg")}
             if self._encoding == ZMQEncoding.JSON:
                 payload_json = json.dumps(payload)
@@ -427,8 +463,26 @@ def __init__(
         )
 
 
+class ReceiveProgressInfo(_ReceiveZMQStreamOutput):
+    """
+    The class defaults are set to receive 0MQ messages with progress information
+    (RunEngine waiting/watcher updates).
+    """
+
+    def __init__(
+        self, *, zmq_subscribe_addr=None, encoding="json", zmq_topic=_default_zmq_progress_topic, timeout=1000
+    ):
+        super().__init__(
+            zmq_subscribe_addr=zmq_subscribe_addr,
+            encoding=encoding,
+            zmq_topic=zmq_topic,
+            timeout=timeout,
+        )
+
+
 ReceiveConsoleOutput.__doc__ += _ReceiveZMQStreamOutput.__doc__
 ReceiveSystemInfo.__doc__ += _ReceiveZMQStreamOutput.__doc__
+ReceiveProgressInfo.__doc__ += _ReceiveZMQStreamOutput.__doc__
 
 
 class _ReceiveZMQStreamOutputAsync:
@@ -717,8 +771,26 @@ def __init__(
         )
 
 
+class ReceiveProgressInfoAsync(_ReceiveZMQStreamOutputAsync):
+    """
+    The class defaults are set to receive 0MQ messages with progress information
+    (RunEngine waiting/watcher updates).
+    """
+
+    def __init__(
+        self, *, zmq_subscribe_addr=None, encoding="json", zmq_topic=_default_zmq_progress_topic, timeout=1000
+    ):
+        super().__init__(
+            zmq_subscribe_addr=zmq_subscribe_addr,
+            encoding=encoding,
+            zmq_topic=zmq_topic,
+            timeout=timeout,
+        )
+
+
 ReceiveConsoleOutputAsync.__doc__ += _ReceiveZMQStreamOutputAsync.__doc__
 ReceiveSystemInfoAsync.__doc__ += _ReceiveZMQStreamOutputAsync.__doc__
+ReceiveProgressInfoAsync.__doc__ += _ReceiveZMQStreamOutputAsync.__doc__
 
 
 def qserver_console_monitor_cli():
@@ -744,9 +816,10 @@ def formatter(prog):
         dest="zmq_info_addr",
         type=str,
         default=None,
-        help="The address of RE Manager socket used for publishing console output. The parameter overrides "
-        "the address set using QSERVER_ZMQ_INFO_ADDRESS environment variable. The default value is used "
-        "if the address is not set using the parameter or the environment variable. Address format: "
+        help="The address of RE Manager socket used for publishing console output, status info, and "
+        "progress updates. The parameter overrides the address set using QSERVER_ZMQ_INFO_ADDRESS "
+        "environment variable. The default value is used if the address is not set using the parameter "
+        "or the environment variable. Address format: "
         f"'tcp://127.0.0.1:60625' (default: {default_zmq_info_address}).",
     )