Update ControlMessage to hold arbitrary Python objects & update MessageMeta to copy & slice

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

@@ -21,6 +21,7 @@
 
 #include <nlohmann/json.hpp>   // for json, basic_json
 #include <pybind11/pytypes.h>  // for object, dict, list, none
+#include <pymrc/utilities/json_values.hpp>
 
 #include <chrono>    // for system_clock, time_point
 #include <map>       // for map
@@ -328,6 +329,20 @@ class ControlMessage
      */
     void task_type(ControlMessageType task_type);
 
+    /**
+     * @brief Set a Python object at a specific path
+     * @param path the path in the JSON object where the value should be set
+     * @param value the Python object to set
+     */
+    void set_py_object(const std::string& path, const pybind11::object& value);
+
+    /**
+     * @brief Get the Python object at a specific path
+     * @param path Path to the specified object
+     * @return The Python representation of the object at the specified path
+     */
+    pybind11::object get_py_object(const std::string& path) const;
+
     /**
      * @brief Sets a timestamp for a specific key.
      *
@@ -374,6 +389,7 @@ class ControlMessage
 
     nlohmann::json m_tasks{};
     nlohmann::json m_config{};
+    mrc::pymrc::JSONValues m_py_objects;
 
     std::map<std::string, time_point_t> m_timestamps{};
 };

morpheus/_lib/include/morpheus/messages/meta.hpp

@@ -122,6 +122,23 @@ class MessageMeta
      */
     virtual std::optional<std::string> ensure_sliceable_index();
 
+    /**
+     * @brief Creates a deep copy of DataFrame with the specified ranges.
+     *
+     * @param ranges the tensor index ranges to copy
+     * @return std::shared_ptr<MessageMeta> the deep copy of the specified ranges
+     */
+    virtual std::shared_ptr<MessageMeta> copy_ranges(const std::vector<RangeType>& ranges) const;
+
+    /**
+     * @brief Get a slice of the underlying DataFrame by creating a deep copy
+     *
+     * @param start the tensor index of the start of the copy
+     * @param stop the tensor index of the end of the copy
+     * @return std::shared_ptr<MessageMeta> the deep copy of the speicifed slice
+     */
+    virtual std::shared_ptr<MessageMeta> get_slice(TensorIndex start, TensorIndex stop) const;
+
     /**
      * @brief Create MessageMeta cpp object from a python object
      *
@@ -297,6 +314,23 @@ struct MessageMetaInterfaceProxy
      * @return std::string The name of the column with the old index or nullopt if no changes were made.
      */
     static std::optional<std::string> ensure_sliceable_index(MessageMeta& self);
+
+    /**
+     * @brief Creates a deep copy of DataFrame with the specified ranges.
+     *
+     * @param ranges the tensor index ranges to copy
+     * @return std::shared_ptr<MessageMeta> the deep copy of the specified ranges
+     */
+    static std::shared_ptr<MessageMeta> copy_ranges(MessageMeta& self, const std::vector<RangeType>& ranges);
+
+    /**
+     * @brief Get a slice of the underlying DataFrame by creating a deep copy
+     *
+     * @param start the tensor index of the start of the copy
+     * @param stop the tensor index of the end of the copy
+     * @return std::shared_ptr<MessageMeta> the deep copy of the speicifed slice
+     */
+    static std::shared_ptr<MessageMeta> get_slice(MessageMeta& self, TensorIndex start, TensorIndex stop);
 };
 
 #pragma GCC visibility pop

morpheus/_lib/messages/__init__.pyi

@@ -54,6 +54,7 @@ class ControlMessage():
         Retrieve timestamps matching a regex filter within a given group.
         """
     def get_metadata(self, key: object = None, default_value: object = None) -> object: ...
+    def get_py_object(self, path: str) -> object: ...
     def get_tasks(self) -> dict: ...
     def get_timestamp(self, key: str, fail_if_nonexist: bool = False) -> object: 
         """
@@ -70,6 +71,7 @@ class ControlMessage():
     def payload(self, meta: object) -> None: ...
     def remove_task(self, task_type: str) -> dict: ...
     def set_metadata(self, key: str, value: object) -> None: ...
+    def set_py_object(self, path: str, value: object) -> None: ...
     def set_timestamp(self, key: str, timestamp: object) -> None: 
         """
         Set a timestamp for a given key and group.
@@ -182,6 +184,7 @@ class InferenceMemoryNLP(InferenceMemory, TensorMemory):
 class MessageMeta():
     def __init__(self, df: object) -> None: ...
     def copy_dataframe(self) -> object: ...
+    def copy_ranges(self, ranges: typing.List[typing.Tuple[int, int]]) -> MessageMeta: ...
     def ensure_sliceable_index(self) -> typing.Optional[str]: ...
     def get_column_names(self) -> typing.List[str]: ...
     @typing.overload
@@ -192,6 +195,7 @@ class MessageMeta():
     def get_data(self, columns: str) -> object: ...
     @typing.overload
     def get_data(self, columns: typing.List[str]) -> object: ...
+    def get_slice(self, start: int, stop: int) -> MessageMeta: ...
     def has_sliceable_index(self) -> bool: ...
     @staticmethod
     def make_from_file(arg0: str) -> MessageMeta: ...

morpheus/_lib/messages/module.cpp

@@ -250,6 +250,8 @@ PYBIND11_MODULE(messages, _module)
         .def("mutable_dataframe", &MessageMetaInterfaceProxy::mutable_dataframe, py::return_value_policy::move)
         .def("has_sliceable_index", &MessageMetaInterfaceProxy::has_sliceable_index)
         .def("ensure_sliceable_index", &MessageMetaInterfaceProxy::ensure_sliceable_index)
+        .def("copy_ranges", &MessageMetaInterfaceProxy::copy_ranges, py::return_value_policy::move, py::arg("ranges"))
+        .def("get_slice", &MessageMetaInterfaceProxy::get_slice, py::return_value_policy::move, py::arg("start"), py::arg("stop"))
         .def_static("make_from_file", &MessageMetaInterfaceProxy::init_cpp);
 
     py::class_<MultiMessage, std::shared_ptr<MultiMessage>>(_module, "MultiMessage")
@@ -427,7 +429,9 @@ PYBIND11_MODULE(messages, _module)
         .def("set_metadata", &ControlMessageProxy::set_metadata, py::arg("key"), py::arg("value"))
         .def("task_type", pybind11::overload_cast<>(&ControlMessage::task_type))
         .def(
-            "task_type", pybind11::overload_cast<ControlMessageType>(&ControlMessage::task_type), py::arg("task_type"));
+            "task_type", pybind11::overload_cast<ControlMessageType>(&ControlMessage::task_type), py::arg("task_type"))
+        .def("set_py_object", &ControlMessage::set_py_object, py::arg("path"), py::arg("value"))
+        .def("get_py_object", &ControlMessage::get_py_object, py::arg("path"));
 
     py::class_<LoaderRegistry, std::shared_ptr<LoaderRegistry>>(_module, "DataLoaderRegistry")
         .def_static("contains", &LoaderRegistry::contains, py::arg("name"))

morpheus/_lib/src/messages/control.cpp

@@ -156,6 +156,16 @@ nlohmann::json ControlMessage::remove_task(const std::string& task_type)
     throw std::runtime_error("No tasks of type " + task_type + " found");
 }
 
+void ControlMessage::set_py_object(const std::string& path, const pybind11::object& value)
+{
+    m_py_objects = std::move(m_py_objects.set_value(path, value));
+}
+
+pybind11::object ControlMessage::get_py_object(const std::string& path) const {
+    auto abs_path = "/" + path;
+    return m_py_objects.get_python(abs_path);
+}
+
 void ControlMessage::set_timestamp(const std::string& key, time_point_t timestamp_ns)
 {
     // Insert or update the timestamp in the map

morpheus/_lib/src/messages/meta.cpp

@@ -27,6 +27,8 @@
 
 #include <cuda_runtime.h>               // for cudaMemcpy, cudaMemcpy2D, cudaMemcpyKind
 #include <cudf/column/column_view.hpp>  // for column_view
+#include <cudf/concatenate.hpp>
+#include <cudf/copying.hpp>
 #include <cudf/io/types.hpp>
 #include <cudf/types.hpp>  // for type_id, data_type, size_type
 #include <glog/logging.h>
@@ -84,7 +86,9 @@ TableInfo MessageMeta::get_info(const std::vector<std::string>& column_names) co
 
 void MessageMeta::set_data(const std::string& col_name, TensorObject tensor)
 {
-    this->set_data({col_name}, {tensor});
+    // This causes a segfault in copy ctor of TensorObject, when shared_ptr<MemoryDescriptor> increases the ref count
+    // this->set_data({col_name}, {tensor});
+    this->set_data({col_name}, std::vector<TensorObject>{tensor});
 }
 
 void MessageMeta::set_data(const std::vector<std::string>& column_names, const std::vector<TensorObject>& tensors)
@@ -111,16 +115,13 @@ void MessageMeta::set_data(const std::vector<std::string>& column_names, const s
         const auto tensor_type    = DType(tensors[i].dtype());
         const auto tensor_type_id = tensor_type.cudf_type_id();
         const auto row_stride     = tensors[i].stride(0);
-
         CHECK(tensors[i].count() == cv.size() &&
               (table_type_id == tensor_type_id ||
                (table_type_id == cudf::type_id::BOOL8 && tensor_type_id == cudf::type_id::UINT8)));
-
         const auto item_size = tensors[i].dtype().item_size();
 
         // Dont use cv.data<>() here since that does not account for the size of each element
         auto data_start = const_cast<uint8_t*>(cv.head<uint8_t>()) + cv.offset() * item_size;
-
         if (row_stride == 1)
         {
             // column major just use cudaMemcpy
@@ -193,6 +194,41 @@ bool MessageMeta::has_sliceable_index() const
     return table.has_sliceable_index();
 }
 
+std::shared_ptr<MessageMeta> MessageMeta::copy_ranges(const std::vector<RangeType>& ranges) const
+{
+    // copy ranges into a sequntial list of values
+    // https://github.com/rapidsai/cudf/issues/11223
+    std::vector<TensorIndex> cudf_ranges;
+    for (const auto& p : ranges)
+    {
+        // Append the message offset to the range here
+        cudf_ranges.push_back(p.first);
+        cudf_ranges.push_back(p.second);
+    }
+    auto table_info   = this->get_info();
+    auto column_names = table_info.get_column_names();
+    auto metadata     = cudf::io::table_metadata{};
+
+    metadata.schema_info.reserve(column_names.size() + 1);
+    metadata.schema_info.emplace_back("");
+
+    for (auto column_name : column_names)
+    {
+        metadata.schema_info.emplace_back(column_name);
+    }
+
+    auto table_view                     = table_info.get_view();
+    auto sliced_views                   = cudf::slice(table_view, cudf_ranges);
+    cudf::io::table_with_metadata table = {cudf::concatenate(sliced_views), std::move(metadata)};
+
+    return MessageMeta::create_from_cpp(std::move(table), 1);
+}
+
+std::shared_ptr<MessageMeta> MessageMeta::get_slice(TensorIndex start, TensorIndex stop) const
+{
+    return this->copy_ranges({{start, stop}});
+}
+
 std::optional<std::string> MessageMeta::ensure_sliceable_index()
 {
     auto table = this->get_mutable_info();
@@ -462,6 +498,20 @@ std::optional<std::string> MessageMetaInterfaceProxy::ensure_sliceable_index(Mes
     return self.ensure_sliceable_index();
 }
 
+std::shared_ptr<MessageMeta> MessageMetaInterfaceProxy::copy_ranges(MessageMeta& self, const std::vector<RangeType>& ranges)
+{
+    pybind11::gil_scoped_release no_gil;
+
+    return self.copy_ranges(ranges);
+}
+
+std::shared_ptr<MessageMeta> MessageMetaInterfaceProxy::get_slice(MessageMeta& self, TensorIndex start, TensorIndex stop)
+{
+    pybind11::gil_scoped_release no_gil;
+
+    return self.get_slice(start, stop);
+}
+
 SlicedMessageMeta::SlicedMessageMeta(std::shared_ptr<MessageMeta> other,
                                      TensorIndex start,
                                      TensorIndex stop,
@@ -492,4 +542,6 @@ std::optional<std::string> SlicedMessageMeta::ensure_sliceable_index()
     throw std::runtime_error{"Unable to set a new index on the DataFrame from a partial view of the columns/rows."};
 }
 
+
+
 }  // namespace morpheus

morpheus/_lib/tests/CMakeLists.txt

@@ -105,6 +105,7 @@ add_morpheus_test(
     messages/test_control_message.cpp
     messages/test_dev_doc_ex3.cpp
     messages/test_sliced_message_meta.cpp
+    messages/test_message_meta.cpp
 )
 
 add_morpheus_test(

morpheus/_lib/tests/messages/test_control_message.cpp

@@ -33,6 +33,8 @@
 #include <stdexcept>  // for runtime_error
 #include <string>     // for operator<=>, string, char_traits, basic_string
 #include <vector>     // for vector
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h> // IWYU pragma: keep
 
 using namespace morpheus;
 using namespace morpheus::test;
@@ -41,6 +43,11 @@ using clock_type_t = std::chrono::system_clock;
 
 using TestControlMessage = morpheus::test::TestMessages;  // NOLINT(readability-identifier-naming)
 
+namespace py = pybind11;
+using namespace pybind11::literals;
+using namespace std::string_literals;
+
+
 TEST_F(TestControlMessage, InitializationTest)
 {
     auto msg_one = ControlMessage();
@@ -334,3 +341,29 @@ TEST_F(TestControlMessage, GetTensorMemoryWhenNoneSet)
     // Verify that the retrieved tensor memory is nullptr
     EXPECT_EQ(nullptr, retrievedTensorMemory);
 }
+
+TEST_F(TestControlMessage, SetAndGetPyObject)
+{
+    auto msg = ControlMessage();
+
+    std::array<std::string, 3> alphabet = {"a", "b", "c"};
+    auto py_dict = py::dict("this"_a     = py::dict("is"_a = "a test"s),
+                    "alphabet"_a = py::cast(alphabet),
+                    "ncc"_a      = 1701,
+                    "cost"_a     = 47.47);
+
+    // <path, expected_result>
+    std::vector<std::pair<std::string, py::object>> tests = {{"", py_dict},
+                                                             {"this", py::dict("is"_a = "a test"s)},
+                                                             {"this/is", py::str("a test"s)},
+                                                             {"alphabet", py_dict["alphabet"]},
+                                                             {"ncc", py::int_(1701)},
+                                                             {"cost", py::float_(47.47)}};
+
+    for (auto& [path, expected_object] : tests)
+    {
+        msg.set_py_object(path, expected_object);
+        auto object = msg.get_py_object(path);
+        EXPECT_TRUE(object.equal(expected_object));
+    }
+}