Support ControlMessage output in the C++ impl of DeserializeStage

examples/llm/cli.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -32,7 +32,7 @@
               callback=parse_log_level,
               help="Specify the logging level to use.")
 @click.option('--use_cpp',
-              default=False,
+              default=True,
               type=bool,
               help=("Whether or not to use C++ node and message types or to prefer python. "
                     "Only use as a last resort if bugs are encountered"))

examples/llm/vdb_upload/pipeline.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,6 +15,7 @@
 import time
 
 from morpheus.config import Config
+from morpheus.config import CppConfig
 from morpheus.config import PipelineModes
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.monitor_stage import MonitorStage
@@ -48,6 +49,8 @@ def pipeline(num_threads: int,
              vector_db_resource_name: str,
              triton_server_url: str):
 
+    # WebScraperStage requires C++ mode to be disabled
+    CppConfig.set_should_use_cpp(False)
     config = Config()
     config.mode = PipelineModes.NLP
 

morpheus/_lib/include/morpheus/stages/deserialize.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -17,9 +17,11 @@
 
 #pragma once
 
+#include "morpheus/messages/control.hpp"
 #include "morpheus/messages/meta.hpp"
 #include "morpheus/messages/multi.hpp"
-#include "morpheus/types.hpp"  // for TensorIndex
+#include "morpheus/types.hpp"                  // for TensorIndex
+#include "morpheus/utilities/python_util.hpp"  // for show_warning_message
 
 #include <boost/fiber/context.hpp>
 #include <boost/fiber/future/future.hpp>
@@ -30,6 +32,7 @@
 #include <mrc/segment/builder.hpp>
 #include <mrc/segment/object.hpp>
 #include <mrc/types.hpp>
+#include <nlohmann/json.hpp>
 #include <pymrc/node.hpp>
 #include <rxcpp/rx.hpp>
 // IWYU pragma: no_include "rxcpp/sources/rx-iterate.hpp"
@@ -38,11 +41,13 @@
 #include <memory>
 #include <string>
 #include <thread>
+#include <utility>  // for pair
 #include <vector>
 
 namespace morpheus {
 /****** Component public implementations *******************/
 /****** DeserializationStage********************************/
+using namespace std::literals::string_literals;
 
 /**
  * @addtogroup stages
@@ -51,15 +56,25 @@ namespace morpheus {
  */
 
 #pragma GCC visibility push(default)
-/**
- * @brief Slices incoming Dataframes into smaller `batch_size`'d chunks. This stage accepts the `MessageMeta` output
- * from `FileSourceStage`/`KafkaSourceStage` stages breaking them up into into `MultiMessage`'s. This should be one of
- * the first stages after the `Source` object.
- */
-class DeserializeStage : public mrc::pymrc::PythonNode<std::shared_ptr<MessageMeta>, std::shared_ptr<MultiMessage>>
+using cm_task_t = std::pair<std::string, nlohmann::json>;
+
+void make_output_message(std::shared_ptr<MultiMessage>& full_message,
+                         TensorIndex start,
+                         TensorIndex stop,
+                         cm_task_t* task,
+                         std::shared_ptr<MultiMessage>& windowed_message);
+
+void make_output_message(std::shared_ptr<MultiMessage>& full_message,
+                         TensorIndex start,
+                         TensorIndex stop,
+                         cm_task_t* task,
+                         std::shared_ptr<ControlMessage>& windowed_message);
+
+template <typename OutputT>
+class DeserializeStage : public mrc::pymrc::PythonNode<std::shared_ptr<MessageMeta>, std::shared_ptr<OutputT>>
 {
   public:
-    using base_t = mrc::pymrc::PythonNode<std::shared_ptr<MessageMeta>, std::shared_ptr<MultiMessage>>;
+    using base_t = mrc::pymrc::PythonNode<std::shared_ptr<MessageMeta>, std::shared_ptr<OutputT>>;
     using typename base_t::sink_type_t;
     using typename base_t::source_type_t;
     using typename base_t::subscribe_fn_t;
@@ -70,22 +85,27 @@ class DeserializeStage : public mrc::pymrc::PythonNode<std::shared_ptr<MessageMe
      * @param batch_size Number of messages to be divided into each batch
      * @param ensure_sliceable_index Whether or not to call `ensure_sliceable_index()` on all incoming `MessageMeta`
      */
-    DeserializeStage(TensorIndex batch_size, bool ensure_sliceable_index = true);
+    DeserializeStage(TensorIndex batch_size,
+                     bool ensure_sliceable_index     = true,
+                     std::unique_ptr<cm_task_t> task = nullptr) :
+      base_t(base_t::op_factory_from_sub_fn(build_operator())),
+      m_batch_size(batch_size),
+      m_ensure_sliceable_index(ensure_sliceable_index),
+      m_task(std::move(task)){};
 
   private:
-    /**
-     * TODO(Documentation)
-     */
     subscribe_fn_t build_operator();
 
     TensorIndex m_batch_size;
     bool m_ensure_sliceable_index{true};
+    std::unique_ptr<cm_task_t> m_task{nullptr};
 };
 
 /****** DeserializationStageInterfaceProxy******************/
 /**
  * @brief Interface proxy, used to insulate python bindings.
  */
+template <typename OutputT>
 struct DeserializeStageInterfaceProxy
 {
     /**
@@ -96,11 +116,87 @@ struct DeserializeStageInterfaceProxy
      * @param batch_size : Number of messages to be divided into each batch
      * @return std::shared_ptr<mrc::segment::Object<DeserializeStage>>
      */
-    static std::shared_ptr<mrc::segment::Object<DeserializeStage>> init(mrc::segment::Builder& builder,
-                                                                        const std::string& name,
-                                                                        TensorIndex batch_size,
-                                                                        bool ensure_sliceable_index);
+    static std::shared_ptr<mrc::segment::Object<DeserializeStage<OutputT>>> init(mrc::segment::Builder& builder,
+                                                                                 const std::string& name,
+                                                                                 TensorIndex batch_size,
+                                                                                 bool ensure_sliceable_index,
+                                                                                 const pybind11::object& task_type,
+                                                                                 const pybind11::object& task_payload);
 };
+
+template <typename OutputT>
+typename DeserializeStage<OutputT>::subscribe_fn_t DeserializeStage<OutputT>::build_operator()
+{
+    return [this](rxcpp::observable<sink_type_t> input, rxcpp::subscriber<source_type_t> output) {
+        return input.subscribe(rxcpp::make_observer<sink_type_t>(
+            [this, &output](sink_type_t x) {
+                if (!x->has_sliceable_index())
+                {
+                    if (m_ensure_sliceable_index)
+                    {
+                        auto old_index_name = x->ensure_sliceable_index();
+
+                        if (old_index_name.has_value())
+                        {
+                            // Generate a warning
+                            LOG(WARNING) << MORPHEUS_CONCAT_STR(
+                                "Incoming MessageMeta does not have a unique and monotonic index. Updating index "
+                                "to be unique. Existing index will be retained in column '"
+                                << *old_index_name << "'");
+                        }
+                    }
+                    else
+                    {
+                        utilities::show_warning_message(
+                            "Detected a non-sliceable index on an incoming MessageMeta. Performance when taking slices "
+                            "of messages may be degraded. Consider setting `ensure_sliceable_index==True`",
+                            PyExc_RuntimeWarning);
+                    }
+                }
+
+                // Make one large MultiMessage
+                auto full_message = std::make_shared<MultiMessage>(x, 0, x->count());
+
+                // Loop over the MessageMeta and create sub-batches
+                for (TensorIndex i = 0; i < x->count(); i += this->m_batch_size)
+                {
+                    std::shared_ptr<OutputT> windowed_message{nullptr};
+                    make_output_message(
+                        full_message, i, std::min(i + this->m_batch_size, x->count()), m_task.get(), windowed_message);
+                    output.on_next(std::move(windowed_message));
+                }
+            },
+            [&](std::exception_ptr error_ptr) {
+                output.on_error(error_ptr);
+            },
+            [&]() {
+                output.on_completed();
+            }));
+    };
+}
+
+template <typename OutputT>
+std::shared_ptr<mrc::segment::Object<DeserializeStage<OutputT>>> DeserializeStageInterfaceProxy<OutputT>::init(
+    mrc::segment::Builder& builder,
+    const std::string& name,
+    TensorIndex batch_size,
+    bool ensure_sliceable_index,
+    const pybind11::object& task_type,
+    const pybind11::object& task_payload)
+{
+    std::unique_ptr<cm_task_t> task{nullptr};
+
+    if (!task_type.is_none() && !task_payload.is_none())
+    {
+        task = std::make_unique<cm_task_t>(pybind11::cast<std::string>(task_type),
+                                           mrc::pymrc::cast_from_pyobject(task_payload));
+    }
+
+    auto stage =
+        builder.construct_object<DeserializeStage<OutputT>>(name, batch_size, ensure_sliceable_index, std::move(task));
+
+    return stage;
+}
 #pragma GCC visibility pop
 /** @} */  // end of group
 }  // namespace morpheus

morpheus/_lib/include/morpheus/utilities/python_util.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -27,6 +27,7 @@ using PyObject = _object;  // NOLINT(readability-identifier-naming)
 
 namespace morpheus::utilities {
 
+#pragma GCC visibility push(default)
 /**
  * @brief Shows a python warning using the `warnings.warn` module. These warnings can be suppressed and work different
  * than `logger.warn()`
@@ -38,5 +39,6 @@ namespace morpheus::utilities {
 void show_warning_message(const std::string& deprecation_message,
                           PyObject* category  = nullptr,
                           ssize_t stack_level = 1);
+#pragma GCC visibility pop
 
 }  // namespace morpheus::utilities

morpheus/_lib/src/stages/deserialize.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -24,14 +24,15 @@
 #include "mrc/segment/object.hpp"
 #include "mrc/types.hpp"
 
+#include "morpheus/messages/control.hpp"
 #include "morpheus/types.hpp"
-#include "morpheus/utilities/python_util.hpp"
 #include "morpheus/utilities/string_util.hpp"
 
 #include <glog/logging.h>
 #include <mrc/segment/builder.hpp>
 #include <pyerrors.h>
 #include <pymrc/node.hpp>
+#include <pymrc/utils.hpp>  // for cast_from_pyobject
 #include <rxcpp/rx.hpp>
 
 #include <algorithm>  // for min
@@ -43,65 +44,33 @@
 #include <utility>
 
 namespace morpheus {
-// Component public implementations
-// ************ DeserializationStage **************************** //
-DeserializeStage::DeserializeStage(TensorIndex batch_size, bool ensure_sliceable_index) :
-  PythonNode(base_t::op_factory_from_sub_fn(build_operator())),
-  m_batch_size(batch_size),
-  m_ensure_sliceable_index(ensure_sliceable_index)
-{}
 
-DeserializeStage::subscribe_fn_t DeserializeStage::build_operator()
+void make_output_message(std::shared_ptr<MultiMessage>& full_message,
+                         TensorIndex start,
+                         TensorIndex stop,
+                         cm_task_t* task,
+                         std::shared_ptr<MultiMessage>& windowed_message)
 {
-    return [this](rxcpp::observable<sink_type_t> input, rxcpp::subscriber<source_type_t> output) {
-        return input.subscribe(rxcpp::make_observer<sink_type_t>(
-            [this, &output](sink_type_t x) {
-                if (!x->has_sliceable_index())
-                {
-                    if (m_ensure_sliceable_index)
-                    {
-                        auto old_index_name = x->ensure_sliceable_index();
-
-                        if (old_index_name.has_value())
-                        {
-                            // Generate a warning
-                            LOG(WARNING) << MORPHEUS_CONCAT_STR(
-                                "Incoming MessageMeta does not have a unique and monotonic index. Updating index "
-                                "to be unique. Existing index will be retained in column '"
-                                << *old_index_name << "'");
-                        }
-                    }
-                    else
-                    {
-                        utilities::show_warning_message(
-                            "Detected a non-sliceable index on an incoming MessageMeta. Performance when taking slices "
-                            "of messages may be degraded. Consider setting `ensure_sliceable_index==True`",
-                            PyExc_RuntimeWarning);
-                    }
-                }
-
-                // Make one large MultiMessage
-                auto full_message = std::make_shared<MultiMessage>(x, 0, x->count());
-
-                // Loop over the MessageMeta and create sub-batches
-                for (TensorIndex i = 0; i < x->count(); i += this->m_batch_size)
-                {
-                    auto next = full_message->get_slice(i, std::min(i + this->m_batch_size, x->count()));
-
-                    output.on_next(std::move(next));
-                }
-            },
-            [&](std::exception_ptr error_ptr) { output.on_error(error_ptr); },
-            [&]() { output.on_completed(); }));
-    };
+    DCHECK_EQ(task, nullptr) << "Task is not supported for MultiMessage";
+    auto sliced_msg = full_message->get_slice(start, stop);
+    windowed_message.swap(sliced_msg);
 }
 
-// ************ DeserializationStageInterfaceProxy ************* //
-std::shared_ptr<mrc::segment::Object<DeserializeStage>> DeserializeStageInterfaceProxy::init(
-    mrc::segment::Builder& builder, const std::string& name, TensorIndex batch_size, bool ensure_sliceable_index)
+void make_output_message(std::shared_ptr<MultiMessage>& full_message,
+                         TensorIndex start,
+                         TensorIndex stop,
+                         cm_task_t* task,
+                         std::shared_ptr<ControlMessage>& windowed_message)
 {
-    auto stage = builder.construct_object<DeserializeStage>(name, batch_size, ensure_sliceable_index);
+    auto window      = full_message->copy_ranges({{start, stop}}, stop - start);
+    auto new_message = std::make_shared<ControlMessage>();
+    new_message->payload(window->meta);
+    if (task)
+    {
+        new_message->add_task(task->first, task->second);
+    }
 
-    return stage;
+    windowed_message.swap(new_message);
 }
+
 }  // namespace morpheus

morpheus/_lib/stages/__init__.pyi

@@ -15,7 +15,8 @@ import mrc.core.segment
 __all__ = [
     "AddClassificationsStage",
     "AddScoresStage",
-    "DeserializeStage",
+    "DeserializeControlMessageStage",
+    "DeserializeMultiMessageStage",
     "FileSourceStage",
     "FilterDetectionsStage",
     "FilterSource",
@@ -37,7 +38,10 @@ class AddClassificationsStage(mrc.core.segment.SegmentObject):
 class AddScoresStage(mrc.core.segment.SegmentObject):
     def __init__(self, builder: mrc.core.segment.Builder, name: str, idx2label: typing.Dict[int, str]) -> None: ...
     pass
-class DeserializeStage(mrc.core.segment.SegmentObject):
+class DeserializeControlMessageStage(mrc.core.segment.SegmentObject):
+    def __init__(self, builder: mrc.core.segment.Builder, name: str, batch_size: int, ensure_sliceable_index: bool = True, task_type: object = None, task_payload: object = None) -> None: ...
+    pass
+class DeserializeMultiMessageStage(mrc.core.segment.SegmentObject):
     def __init__(self, builder: mrc.core.segment.Builder, name: str, batch_size: int, ensure_sliceable_index: bool = True) -> None: ...
     pass
 class FileSourceStage(mrc.core.segment.SegmentObject):