Add ZMQ API message for requesting the current state of the RunEngine Metadata

bluesky_queueserver/manager/manager.py

@@ -1,6 +1,7 @@
 import asyncio
 import copy
 import enum
+import json
 import logging
 import time as ttime
 import uuid
@@ -1586,6 +1587,18 @@ async def _worker_request_plans_and_devices_list(self):
             plans_and_devices_list, err_msg = None, "Timeout occurred while processing the request"
         return plans_and_devices_list, err_msg
 
+    async def _worker_request_runengine_metadata(self):
+        try:
+            runengine_metadata = await self._comm_to_worker.send_msg("request_runengine_metadata")
+            err_msg = ""
+            if runengine_metadata is None:
+                err_msg = "Failed to obtain the RE metadata from the worker"
+
+        except CommTimeoutError:
+            runengine_metadata, err_msg = None, "Timeout occurred while processing the request"
+
+        return runengine_metadata, err_msg
+
     async def _worker_request_task_results(self):
         try:
             tt = self._comm_to_worker_timeout_long
@@ -3351,6 +3364,48 @@ async def _re_runs_handler(self, request):
 
         return {"success": success, "msg": msg, "run_list": run_list, "run_list_uid": run_list_uid}
 
+    async def _re_metadata_handler(self, request):
+        """
+        Returns the runengine metadata from the RE worker process
+        """
+
+        # Clients may ask for RE metadata frequently to update some state,
+        # so make this message debug only.
+        logger.debug("Returning the runengine metadata dictionary ...")
+
+        success, msg, re_metadata = True, "", {}
+        try:
+            self._check_request_for_unsupported_params(request=request, param_names=[])
+
+            if self._environment_exists:
+                re_metadata, msg = await self._worker_request_runengine_metadata()
+
+                if re_metadata is None:
+                    success, re_metadata = False, {}
+                else:
+                    # Make sure the metadata is serializable given the encoding
+                    try:
+                        if self._zmq_encoding == ZMQEncoding.JSON:
+                            json.dumps(re_metadata)
+                        else:
+                            msgpack.packb(re_metadata)
+                    except Exception as ex:
+                        success, msg, re_metadata = (
+                            False,
+                            f"Failed to serialize RE metadata with {self._zmq_encoding.name}: {ex}",
+                            {},
+                        )
+            else:
+                success, msg, re_metadata = (
+                    False,
+                    "Environment does not exist. Cannot retrieve RE metadata.",
+                    {},
+                )
+        except Exception as ex:
+            success, msg, re_metadata = False, f"Error: {ex}", {}
+
+        return {"success": success, "msg": msg, "re_metadata": re_metadata}
+
     def _lock_key_invalid_msg(self):
         """
         Format error message, which reports an invalid lock key.
@@ -3684,6 +3739,7 @@ async def _zmq_execute(self, msg):
             "re_abort": "_re_abort_handler",
             "re_halt": "_re_halt_handler",
             "re_runs": "_re_runs_handler",
+            "re_metadata": "_re_metadata_handler",
             "lock": "_lock_handler",
             "lock_info": "_lock_info_handler",
             "unlock": "_unlock_handler",

bluesky_queueserver/manager/qserver_cli.py

@@ -150,6 +150,7 @@ class QServerExitCodes(enum.Enum):
 qserver re stop            # STOP execution of a paused plan
 qserver re abort           # ABORT execution of a paused plan
 qserver re halt            # HALT execution of a paused plan
+qserver re metadata        # Request RunEngine metadata
 
 qserver re runs            # Get the list of active runs (runs generated by the currently running plans)
 qserver re runs active     # Get the list of active runs
@@ -1135,7 +1136,7 @@ def create_msg(params, *, lock_key):
     elif command == "re":
         if len(params) < 1:
             raise CommandParameterError(f"Request '{command}' must include at least one parameter")
-        supported_params = ("pause", "resume", "stop", "abort", "halt", "runs")
+        supported_params = ("pause", "resume", "stop", "abort", "halt", "runs", "metadata")
         if params[0] in supported_params:
             if params[0] == "pause":
                 method, prms = msg_re_pause(params)

bluesky_queueserver/manager/tests/test_qserver_cli.py

@@ -131,6 +131,8 @@ def test_qserver_cli_and_manager(re_manager):  # noqa: F811
 
     assert sp_call(["qserver", "re", "resume"]) == SUCCESS
 
+    assert sp_call(["qserver", "re", "metadata"]) == SUCCESS
+
     assert wait_for_condition(
         time=60, condition=condition_queue_processing_finished
     ), "Timeout while waiting for process to finish"

bluesky_queueserver/manager/tests/test_zmq_api_base.py

@@ -5971,6 +5971,101 @@ def test_zmq_api_re_runs_1(re_manager_pc_copy, tmp_path, test_with_manager_resta
     assert wait_for_condition(time=5, condition=condition_environment_closed)
 
 
+# fmt: off
+@pytest.mark.parametrize("test_with_manager_restart", [False, True])
+# fmt: on
+def test_zmq_api_re_metadata_1(re_manager_pc_copy, tmp_path, test_with_manager_restart):  # noqa: F811
+    """
+    Tests `re_metadata` functionality. First checks if the request correctly fails if the environment is not open.
+    Then checks to make sure initial metadata could be retrieved. Then, executes a plan that adds `scan_id` and
+    `versions` keys to the metadata, and makes sure we can read those back, as well as the original key/value.
+    """
+
+    _, pc_path = re_manager_pc_copy
+
+    resp, _ = zmq_request("re_metadata")
+    assert resp["success"] is False, f"{resp =}"
+    assert resp["msg"] == "Environment does not exist. Cannot retrieve RE metadata."
+
+    # Open the environment
+    resp, _ = zmq_request("environment_open")
+    assert resp["success"] is True, f"{resp =}"
+    assert wait_for_condition(time=5, condition=condition_environment_created)
+
+    # Get the initial run_engine metadata
+    resp, _ = zmq_request("re_metadata")
+    assert resp["success"] is True, f"{resp =}"
+
+    def check_initial_metadata(re_md):
+        """Function that checks if initial metadata matches what we expect"""
+
+        assert "metadata_key" in re_md, pprint.pformat(re_md)
+        assert re_md["metadata_key"] == "metadata_value"
+        assert "versions" in re_md, pprint.pformat(re_md)
+        assert "bluesky" in re_md["versions"], pprint.pformat(re_md)
+
+    re_md = resp["re_metadata"]
+    check_initial_metadata(re_md)
+    assert len(re_md) == 2, "Only two metadata keys should be present initially!"
+
+    # Add the first 'count' plan
+    resp, _ = zmq_request("queue_item_add", {"item": _plan1, "user": _user, "user_group": _user_group})
+    assert resp["success"] is True
+
+    # Execute the count plan. This should add some additional metadata items
+    resp, _ = zmq_request("queue_start")
+    assert resp["success"] is True
+    assert wait_for_condition(20, condition=condition_manager_idle)
+
+    resp, _ = zmq_request("re_metadata")
+    assert resp["success"] is True, f"{resp =}"
+    re_md = resp["re_metadata"]
+
+    # Make sure we can retrieve scan ID, and it's an int
+    # the `count` plan should add this key, and it should start at `1`
+    assert "scan_id" in re_md, pprint.pformat(re_md)
+    assert re_md["scan_id"] == 1, "Scan ID for first plan should be 1"
+
+    # Make sure our original metadata key is still there
+    check_initial_metadata(re_md)
+
+    resp, _ = zmq_request("environment_close")
+    assert resp["success"] is True, f"{resp =}"
+    assert wait_for_condition(time=5, condition=condition_environment_closed)
+
+
+_add_datetime_to_md = """
+from datetime import datetime
+RE.md['date'] = datetime.now()
+"""
+
+
+def test_zmq_api_re_metadata_non_serializable_md(re_manager_pc_copy, tmp_path):  # noqa: F811
+
+    _, pc_path = re_manager_pc_copy
+    # Add a non-serializable type to the runengine metadata
+    append_code_to_last_startup_file(pc_path, additional_code=_add_datetime_to_md)
+
+    encoding = use_zmq_encoding_for_tests()
+
+    resp, _ = zmq_request("environment_open")
+    assert resp["success"] is True, f"{resp =}"
+    assert wait_for_condition(time=5, condition=condition_environment_created)
+
+    resp, _ = zmq_request("re_metadata")
+    assert resp["success"] is False
+
+    # Check that the error message is correct depending on encoding
+    if encoding == "json":
+        assert resp["msg"].startswith("Failed to serialize RE metadata with JSON:")
+    elif encoding == "msgpack":
+        assert resp["msg"].startswith("Failed to serialize RE metadata with MSGPACK:")
+
+    resp, _ = zmq_request("environment_close")
+    assert resp["success"] is True, f"{resp =}"
+    assert wait_for_condition(time=5, condition=condition_environment_closed)
+
+
 # fmt: off
 @pytest.mark.parametrize("em_lock_code", [False, True])
 # fmt: on

bluesky_queueserver/manager/worker.py

@@ -977,6 +977,13 @@ def _request_plans_and_devices_list_handler(self):
         self._existing_plans_and_devices_changed = False
         return msg_out
 
+    def _request_runengine_metadata_handler(self):
+        """
+        Returns the current state of the Run Engine metadata as a dictionary
+        """
+
+        return dict(self._RE.md)
+
     def _command_close_env_handler(self):
         """
         Close RE Worker environment in orderly way.
@@ -1341,6 +1348,7 @@ def _worker_prepare_for_startup(self):
         self._comm_to_manager.add_method(
             self._request_plans_and_devices_list_handler, "request_plans_and_devices_list"
         )
+        self._comm_to_manager.add_method(self._request_runengine_metadata_handler, "request_runengine_metadata")
         self._comm_to_manager.add_method(self._request_task_results_handler, "request_task_results")
         self._comm_to_manager.add_method(self._command_close_env_handler, "command_close_env")
         self._comm_to_manager.add_method(self._command_confirm_exit_handler, "command_confirm_exit")

bluesky_queueserver/profile_collection_sim/05-run-engine.py

@@ -16,7 +16,7 @@
 except Exception:
     pass
 
-RE = RunEngine()
+RE = RunEngine({"metadata_key": "metadata_value"})
 
 bec = BestEffortCallback()
 if not is_ipython_mode() or not ipython_matplotlib: