fix ci

examples/log_parsing/inference.py

@@ -24,9 +24,6 @@
 from morpheus.config import Config
 from morpheus.config import PipelineModes
 from morpheus.messages import ControlMessage
-from morpheus.messages import MultiInferenceMessage
-from morpheus.messages import MultiInferenceNLPMessage
-from morpheus.messages import MultiResponseMessage
 from morpheus.messages import TensorMemory
 from morpheus.pipeline.stage_schema import StageSchema
 from morpheus.stages.inference.triton_inference_stage import TritonInferenceStage
@@ -60,47 +57,26 @@ class TritonInferenceLogParsing(TritonInferenceWorker):
         Determines whether a logits calculation is needed for the value returned by the Triton inference response.
     """
 
-    def build_output_message(self, x: MultiInferenceMessage | ControlMessage) -> MultiResponseMessage | ControlMessage:
-        if isinstance(x, MultiInferenceMessage):
-            seq_ids = cp.zeros((x.count, 3), dtype=cp.uint32)
-            seq_ids[:, 0] = cp.arange(x.mess_offset, x.mess_offset + x.count, dtype=cp.uint32)
-            seq_ids[:, 2] = x.get_tensor('seq_ids')[:, 2]
-
-            memory = TensorMemory(
-                count=x.count,
-                tensors={
-                    'confidences': cp.zeros((x.count, self._inputs[list(self._inputs.keys())[0]].shape[1])),
-                    'labels': cp.zeros((x.count, self._inputs[list(self._inputs.keys())[0]].shape[1])),
-                    'input_ids': cp.zeros((x.count, x.get_tensor('input_ids').shape[1])),
-                    'seq_ids': seq_ids
-                })
-
-            return MultiResponseMessage(meta=x.meta,
-                                        mess_offset=x.mess_offset,
-                                        mess_count=x.mess_count,
-                                        memory=memory,
-                                        offset=0,
-                                        count=x.count)
-        if isinstance(x, ControlMessage):
-            seq_ids = cp.zeros((x.tensors().count, 3), dtype=cp.uint32)
-            seq_ids[:, 0] = cp.arange(0, x.tensors().count, dtype=cp.uint32)
-            seq_ids[:, 2] = x.tensors().get_tensor('seq_ids')[:, 2]
-
-            memory = _messages.TensorMemory(
-                count=x.tensors().count,
-                tensors={
-                    'confidences': cp.zeros((x.tensors().count, self._inputs[list(self._inputs.keys())[0]].shape[1])),
-                    'labels': cp.zeros((x.tensors().count, self._inputs[list(self._inputs.keys())[0]].shape[1])),
-                    'input_ids': cp.zeros((x.tensors().count, x.tensors().get_tensor('input_ids').shape[1])),
-                    'seq_ids': seq_ids
-                })
-
-            resp = ControlMessage(x)
-            resp.payload(x.payload())
-            resp.tensors(memory)
-            return resp
-
-    def _build_response(self, batch: MultiInferenceMessage, result: tritonclient.InferResult) -> TensorMemory:
+    def build_output_message(self, msg: ControlMessage) -> ControlMessage:
+        seq_ids = cp.zeros((msg.tensors().count, 3), dtype=cp.uint32)
+        seq_ids[:, 0] = cp.arange(0, msg.tensors().count, dtype=cp.uint32)
+        seq_ids[:, 2] = msg.tensors().get_tensor('seq_ids')[:, 2]
+
+        memory = _messages.TensorMemory(
+            count=msg.tensors().count,
+            tensors={
+                'confidences': cp.zeros((msg.tensors().count, self._inputs[list(self._inputs.keys())[0]].shape[1])),
+                'labels': cp.zeros((msg.tensors().count, self._inputs[list(self._inputs.keys())[0]].shape[1])),
+                'input_ids': cp.zeros((msg.tensors().count, msg.tensors().get_tensor('input_ids').shape[1])),
+                'seq_ids': seq_ids
+            })
+
+        resp = ControlMessage(msg)
+        resp.payload(msg.payload())
+        resp.tensors(memory)
+        return resp
+
+    def _build_response(self, batch: ControlMessage, result: tritonclient.InferResult) -> TensorMemory:
 
         outputs = {output.mapped_name: result.as_numpy(output.name) for output in self._outputs.values()}
         outputs = {key: softmax(val, axis=2) for key, val in outputs.items()}
@@ -161,83 +137,46 @@ def supports_cpp_node(self) -> bool:
         return False
 
     def compute_schema(self, schema: StageSchema):
-        schema.output_schema.set_type(MultiResponseMessage)
+        schema.output_schema.set_type(ControlMessage)
 
     @staticmethod
-    def _convert_one_response(output: MultiResponseMessage | ControlMessage, inf: MultiInferenceNLPMessage | ControlMessage,
-                              res: TensorMemory) -> MultiResponseMessage | ControlMessage:
-        if isinstance(output, MultiResponseMessage):
-            memory = output.memory
-
-            out_seq_ids = memory.get_tensor('seq_ids')
-            input_ids = memory.get_tensor('input_ids')
-            confidences = memory.get_tensor('confidences')
-            labels = memory.get_tensor('labels')
-
-            seq_ids = inf.get_id_tensor()
-
-            seq_offset = seq_ids[0, 0].item() - output.mess_offset
-            seq_count = (seq_ids[-1, 0].item() + 1 - seq_offset) - output.mess_offset
-
-            input_ids[inf.offset:inf.count + inf.offset, :] = inf.get_tensor('input_ids')
-            out_seq_ids[inf.offset:inf.count + inf.offset, :] = seq_ids
-
-            resp_confidences = res.get_tensor('confidences')
-            resp_labels = res.get_tensor('labels')
-
-            # Two scenarios:
-            if (inf.mess_count == inf.count):
-                assert seq_count == res.count
-                confidences[inf.offset:inf.offset + inf.count, :] = resp_confidences
-                labels[inf.offset:inf.offset + inf.count, :] = resp_labels
-            else:
-                assert inf.count == res.count
-
-                mess_ids = seq_ids[:, 0].get().tolist()
-
-                for i, idx in enumerate(mess_ids):
-                    confidences[idx, :] = cp.maximum(confidences[idx, :], resp_confidences[i, :])
-                    labels[idx, :] = cp.maximum(labels[idx, :], resp_labels[i, :])
-
-            return MultiResponseMessage.from_message(inf, memory=memory, offset=inf.offset, count=inf.mess_count)
-
-        if isinstance(output, ControlMessage):
-            memory = output.tensors()
+    def _convert_one_response(output: ControlMessage, inf: ControlMessage, res: TensorMemory) -> ControlMessage:
+        memory = output.tensors()
 
-            out_seq_ids = memory.get_tensor('seq_ids')
-            input_ids = memory.get_tensor('input_ids')
-            confidences = memory.get_tensor('confidences')
-            labels = memory.get_tensor('labels')
+        out_seq_ids = memory.get_tensor('seq_ids')
+        input_ids = memory.get_tensor('input_ids')
+        confidences = memory.get_tensor('confidences')
+        labels = memory.get_tensor('labels')
 
-            seq_ids = inf.tensors().get_tensor('seq_ids')
+        seq_ids = inf.tensors().get_tensor('seq_ids')
 
-            seq_offset = seq_ids[0, 0].item()
-            seq_count = seq_ids[-1, 0].item() + 1 - seq_offset
+        seq_offset = seq_ids[0, 0].item()
+        seq_count = seq_ids[-1, 0].item() + 1 - seq_offset
 
-            input_ids[0:inf.tensors().count, :] = inf.tensors().get_tensor('input_ids')
-            out_seq_ids[0:inf.tensors().count, :] = seq_ids
+        input_ids[0:inf.tensors().count, :] = inf.tensors().get_tensor('input_ids')
+        out_seq_ids[0:inf.tensors().count, :] = seq_ids
 
-            resp_confidences = res.get_tensor('confidences')
-            resp_labels = res.get_tensor('labels')
+        resp_confidences = res.get_tensor('confidences')
+        resp_labels = res.get_tensor('labels')
 
-            # Two scenarios:
-            if (inf.payload().count == inf.tensors().count):
-                assert seq_count == res.count
-                confidences[0:inf.tensors().count, :] = resp_confidences
-                labels[0:inf.tensors().count, :] = resp_labels
-            else:
-                assert inf.tensors().count == res.count
+        # Two scenarios:
+        if (inf.payload().count == inf.tensors().count):
+            assert seq_count == res.count
+            confidences[0:inf.tensors().count, :] = resp_confidences
+            labels[0:inf.tensors().count, :] = resp_labels
+        else:
+            assert inf.tensors().count == res.count
 
-                mess_ids = seq_ids[:, 0].get().tolist()
+            mess_ids = seq_ids[:, 0].get().tolist()
 
-                for i, idx in enumerate(mess_ids):
-                    confidences[idx, :] = cp.maximum(confidences[idx, :], resp_confidences[i, :])
-                    labels[idx, :] = cp.maximum(labels[idx, :], resp_labels[i, :])
+            for i, idx in enumerate(mess_ids):
+                confidences[idx, :] = cp.maximum(confidences[idx, :], resp_confidences[i, :])
+                labels[idx, :] = cp.maximum(labels[idx, :], resp_labels[i, :])
 
-            resp = ControlMessage(inf)
-            resp.payload(inf.payload())
-            resp.tensors(memory)
-            return resp
+        resp = ControlMessage(inf)
+        resp.payload(inf.payload())
+        resp.tensors(memory)
+        return resp
 
     def _get_inference_worker(self, inf_queue: ProducerConsumerQueue) -> TritonInferenceLogParsing:
         return TritonInferenceLogParsing(inf_queue=inf_queue,

examples/log_parsing/postprocessing.py

@@ -110,7 +110,6 @@ def _postprocess(self, x: MultiResponseMessage | ControlMessage):
                     else:
                         parsed_df[col_name] = ext_parsed[label]
 
-
             # decode cleanup
             parsed_df = self.__decode_cleanup(parsed_df)
             parsed_df["doc"] = parsed_dfs.index

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

@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * 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");
  * you may not use this file except in compliance with the License.
@@ -17,29 +17,23 @@
 
 #pragma once
 
-#include "morpheus/export.h"
-#include "morpheus/messages/control.hpp"
-#include "morpheus/messages/meta.hpp"
-#include "morpheus/messages/multi.hpp"
-#include "morpheus/types.hpp"                  // for TensorIndex
-#include "morpheus/utilities/python_util.hpp"  // for show_warning_message
-#include "morpheus/utilities/string_util.hpp"  // for MORPHEUS_CONCAT_STR
+#include "morpheus/export.h"              // for MORPHEUS_EXPORT
+#include "morpheus/messages/control.hpp"  // for ControlMessage
+#include "morpheus/messages/meta.hpp"     // for MessageMeta
+#include "morpheus/types.hpp"             // for TensorIndex
 
-#include <glog/logging.h>
-#include <mrc/segment/builder.hpp>
-#include <mrc/segment/object.hpp>
-#include <nlohmann/json.hpp>
-#include <pybind11/pytypes.h>  // for object
-#include <pyerrors.h>          // for PyExc_RuntimeWarning
-#include <pymrc/node.hpp>
-#include <rxcpp/rx.hpp>
+#include <boost/fiber/context.hpp>  // for operator<<
+#include <mrc/segment/builder.hpp>  // for Builder
+#include <mrc/segment/object.hpp>   // for Object
+#include <nlohmann/json.hpp>        // for basic_json, json
+#include <pybind11/pytypes.h>       // for object
+#include <pymrc/node.hpp>           // for PythonNode
+#include <rxcpp/rx.hpp>             // for decay_t, trace_activity, from, observable_member
 
-#include <algorithm>  // IWYU pragma: keep for std::min
-#include <exception>  // for exception_ptr
-#include <memory>
-#include <sstream>  // IWYU pragma: keep for glog
-#include <string>
-#include <utility>  // for pair
+#include <memory>   // for shared_ptr, unique_ptr
+#include <string>   // for string
+#include <thread>   // for operator<<
+#include <utility>  // for move, pair
 
 namespace morpheus {
 /****** Component public implementations *******************/
@@ -72,7 +66,8 @@ class MORPHEUS_EXPORT DeserializeStage
      * @brief Construct a new Deserialize Stage object
      *
      * @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`
+     * @param ensure_sliceable_index Whether or not to call
+     * `ensure_sliceable_index()` on all incoming `MessageMeta`
      * @param task Optional task to be added to all outgoing `ControlMessage`s
      */
     DeserializeStage(TensorIndex batch_size,
@@ -98,26 +93,27 @@ class MORPHEUS_EXPORT DeserializeStage
 struct MORPHEUS_EXPORT DeserializeStageInterfaceProxy
 {
     /**
-     * @brief Create and initialize a DeserializationStage that emits ControlMessage's, and return the result.
-     * If `task_type` is not None, `task_payload` must also be not None, and vice versa.
+     * @brief Create and initialize a DeserializationStage that emits
+     * ControlMessage's, and return the result. If `task_type` is not None,
+     * `task_payload` must also be not None, and vice versa.
      *
      * @param builder : Pipeline context object reference
      * @param name : Name of a stage reference
      * @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`
+     * @param ensure_sliceable_index Whether or not to call
+     * `ensure_sliceable_index()` on all incoming `MessageMeta`
      * @param task_type : Optional task type to be added to all outgoing messages
-     * @param task_payload : Optional json object describing the task to be added to all outgoing messages
+     * @param task_payload : Optional json object describing the task to be added
+     * to all outgoing messages
      * @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,
-        const pybind11::object& task_type,
-        const pybind11::object& task_payload);
+    static std::shared_ptr<mrc::segment::Object<DeserializeStage>> 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);
 };
 
-
 /** @} */  // end of group
 }  // namespace morpheus