Add a get_timestamps method to ControlMessage

python/morpheus/morpheus/_lib/include/morpheus/messages/control.hpp

@@ -235,6 +235,13 @@ class MORPHEUS_EXPORT ControlMessage
      */
     std::optional<time_point_t> get_timestamp(const std::string& key, bool fail_if_nonexist = false);
 
+    /**
+     * @brief Return a reference to the timestamps map
+     *
+     * @return A const map reference containing timestamps
+     */
+    const std::map<std::string, time_point_t>& get_timestamps() const;
+
     /**
      * @brief Retrieves timestamps for all keys that match a regex pattern.
      *
@@ -340,6 +347,13 @@ struct MORPHEUS_EXPORT ControlMessageProxy
      */
     static pybind11::object get_timestamp(ControlMessage& self, const std::string& key, bool fail_if_nonexist = false);
 
+    /**
+     * @brief Return all timestamps
+     *
+     * @return A Python dictionary of timestamps
+     */
+    static pybind11::dict get_timestamps(ControlMessage& self);
+
     /**
      * @brief Retrieves timestamps for all keys that match a regex pattern from the ControlMessage object.
      *

python/morpheus/morpheus/_lib/messages/__init__.pyi

@@ -53,6 +53,7 @@ class ControlMessage():
         """
         Retrieve the timestamp for a given group and key. Returns None if the timestamp does not exist and fail_if_nonexist is False.
         """
+    def get_timestamps(self) -> dict: ...
     def has_metadata(self, key: str) -> bool: ...
     def has_task(self, task_type: str) -> bool: ...
     def list_metadata(self) -> list: ...

python/morpheus/morpheus/_lib/messages/module.cpp

@@ -274,6 +274,7 @@ PYBIND11_MODULE(messages, _module)
              "fail_if_nonexist is False.",
              py::arg("key"),
              py::arg("fail_if_nonexist") = false)
+        .def("get_timestamps", &ControlMessageProxy::get_timestamps)
         .def("set_timestamp",
              &ControlMessageProxy::set_timestamp,
              "Set a timestamp for a given key and group.",

python/morpheus/morpheus/_lib/src/messages/control.cpp

@@ -177,6 +177,11 @@ void ControlMessage::set_timestamp(const std::string& key, time_point_t timestam
     m_timestamps[key] = timestamp_ns;
 }
 
+const std::map<std::string, time_point_t>& ControlMessage::get_timestamps() const
+{
+    return m_timestamps;
+}
+
 std::map<std::string, time_point_t> ControlMessage::filter_timestamp(const std::string& regex_filter)
 {
     std::map<std::string, time_point_t> matching_timestamps;
@@ -365,6 +370,11 @@ py::list ControlMessageProxy::list_metadata(ControlMessage& self)
     return py_keys;
 }
 
+py::dict ControlMessageProxy::get_timestamps(ControlMessage& self)
+{
+    return py::cast(self.get_timestamps());
+}
+
 py::dict ControlMessageProxy::filter_timestamp(ControlMessage& self, const std::string& regex_filter)
 {
     auto cpp_map = self.filter_timestamp(regex_filter);

python/morpheus/morpheus/messages/control_message.py

@@ -167,6 +167,9 @@ def get_timestamp(self, key: str, fail_if_nonexist: bool = False) -> datetime |
                 raise ValueError("Timestamp for the specified key does not exist.") from e
             return None
 
+    def get_timestamps() -> dict[str, datetime]:
+        return self._timestamps
+
     def filter_timestamp(self, regex_filter: str) -> dict[str, datetime]:
         re_obj = re.compile(regex_filter)
 

python/morpheus_llm/morpheus_llm/stages/llm/llm_engine_stage.py

@@ -93,6 +93,10 @@ def _copy_tasks_and_metadata(self,
             for tv in task_value:
                 dst.add_task(task, tv)
 
+        timestamps = src.get_timestamps()
+        for (ts_key, ts) in timestamps.items():
+            dst.set_timestamp(key=ts_key, timestamp=ts)
+
     def _cast_to_cpp_control_message(self, py_message: ControlMessage, *,
                                      cpp_messages_lib: types.ModuleType) -> ControlMessage:
         """

tests/morpheus/messages/test_control_message.py

@@ -250,6 +250,22 @@ def test_filter_timestamp():
     assert result[f"{group}::key2"] == timestamp2, "The timestamp for key2 should match."
 
 
+@pytest.mark.gpu_and_cpu_mode
+def test_get_timestamps():
+    # Create a ControlMessage instance
+    msg = messages.ControlMessage()
+
+    # Setup test data
+    timestamp1 = datetime.datetime.now()
+    timestamp2 = timestamp1 + datetime.timedelta(seconds=1)
+    msg.set_timestamp("key1", timestamp1)
+    msg.set_timestamp("key2", timestamp2)
+
+    # Assert both keys are in the result and have correct timestamps
+    timestamps = msg.get_timestamps()
+    assert timestamps == {"key1": timestamp1, "key2": timestamp2}
+
+
 @pytest.mark.gpu_and_cpu_modetest_tensor_manipulation_after_retrieval
 def test_get_timestamp_fail_if_nonexist():
     # Create a ControlMessage instance