Add tensor_count property for ControlMessage (#2078)

For ControlMessage, msg.tensors().count is a common pattern, calling msg.tensors() might require a bit more cost than we think. Add a `tensor_count` property to avoid the overhead.

Closes #1876 

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/nv-morpheus/Morpheus/blob/main/docs/source/developer_guide/contributing.md).
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Yuchen Zhang (https://github.com/yczhang-nv)

Approvers:
  - David Gardner (https://github.com/dagardner-nv)

URL: #2078

Author: yczhang-nv

examples/log_parsing/inference.py

@@ -57,16 +57,16 @@ class TritonInferenceLogParsing(TritonInferenceWorker):
     """
 
     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 = cp.zeros((msg.tensor_count(), 3), dtype=cp.uint32)
+        seq_ids[:, 0] = cp.arange(0, msg.tensor_count(), dtype=cp.uint32)
         seq_ids[:, 2] = msg.tensors().get_tensor('seq_ids')[:, 2]
 
         memory = TensorMemory(
-            count=msg.tensors().count,
+            count=msg.tensor_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])),
+                'confidences': cp.zeros((msg.tensor_count(), self._inputs[list(self._inputs.keys())[0]].shape[1])),
+                'labels': cp.zeros((msg.tensor_count(), self._inputs[list(self._inputs.keys())[0]].shape[1])),
+                'input_ids': cp.zeros((msg.tensor_count(), msg.tensors().get_tensor('input_ids').shape[1])),
                 'seq_ids': seq_ids
             })
 
@@ -154,19 +154,19 @@ def _convert_one_response(output: ControlMessage, inf: ControlMessage, res: Tens
         seq_offset = seq_ids[0, 0].item()
         seq_count = seq_ids[-1, 0].item() + 1 - seq_offset
 
-        input_ids[batch_offset:inf.tensors().count + batch_offset, :] = inf.tensors().get_tensor('input_ids')
-        out_seq_ids[batch_offset:inf.tensors().count + batch_offset, :] = seq_ids
+        input_ids[batch_offset:inf.tensor_count() + batch_offset, :] = inf.tensors().get_tensor('input_ids')
+        out_seq_ids[batch_offset:inf.tensor_count() + batch_offset, :] = seq_ids
 
         resp_confidences = res.get_tensor('confidences')
         resp_labels = res.get_tensor('labels')
 
         # Two scenarios:
-        if (inf.payload().count == inf.tensors().count):
+        if (inf.payload().count == inf.tensor_count()):
             assert seq_count == res.count
-            confidences[batch_offset:inf.tensors().count + batch_offset, :] = resp_confidences
-            labels[batch_offset:inf.tensors().count + batch_offset, :] = resp_labels
+            confidences[batch_offset:inf.tensor_count() + batch_offset, :] = resp_confidences
+            labels[batch_offset:inf.tensor_count() + batch_offset, :] = resp_labels
         else:
-            assert inf.tensors().count == res.count
+            assert inf.tensor_count() == res.count
 
             mess_ids = seq_ids[:, 0].get().tolist()
 

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

@@ -19,6 +19,7 @@
 
 #include "morpheus/export.h"                  // for MORPHEUS_EXPORT
 #include "morpheus/messages/meta.hpp"         // for MessageMeta
+#include "morpheus/types.hpp"
 #include "morpheus/utilities/json_types.hpp"  // for json_t
 
 #include <pybind11/pytypes.h>  // for object, dict, list
@@ -197,6 +198,13 @@ class MORPHEUS_EXPORT ControlMessage
      */
     void tensors(const std::shared_ptr<TensorMemory>& tensor_memory);
 
+    /**
+     * @brief Get the length of tensors in the tensor memory.
+     *
+     * @return The length of tensors in the tensor memory.
+     */
+    TensorIndex tensor_count();
+
     /**
      * @brief Get the type of task associated with the control message.
      * @return An enum value indicating the task type.
@@ -262,6 +270,7 @@ class MORPHEUS_EXPORT ControlMessage
     ControlMessageType m_cm_type{ControlMessageType::NONE};
     std::shared_ptr<MessageMeta> m_payload{nullptr};
     std::shared_ptr<TensorMemory> m_tensors{nullptr};
+    TensorIndex m_tensor_count{0};
 
     morpheus::utilities::json_t m_tasks{};
     morpheus::utilities::json_t m_config{};

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

@@ -73,6 +73,7 @@ class ControlMessage():
     def task_type(self) -> ControlMessageType: ...
     @typing.overload
     def task_type(self, task_type: ControlMessageType) -> None: ...
+    def tensor_count(self) -> int: ...
     @typing.overload
     def tensors(self) -> TensorMemory: ...
     @typing.overload

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

@@ -290,6 +290,7 @@ PYBIND11_MODULE(messages, _module)
             py::arg("meta"))
         .def("tensors", pybind11::overload_cast<>(&ControlMessage::tensors))
         .def("tensors", pybind11::overload_cast<const std::shared_ptr<TensorMemory>&>(&ControlMessage::tensors))
+        .def("tensor_count", &ControlMessage::tensor_count)
         .def("remove_task", &ControlMessage::remove_task, py::arg("task_type"))
         .def("set_metadata", &ControlMessage::set_metadata, py::arg("key"), py::arg("value"))
         .def("task_type", pybind11::overload_cast<>(&ControlMessage::task_type))

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

@@ -59,9 +59,10 @@ ControlMessage::ControlMessage(const morpheus::utilities::json_t& _config) :
 
 ControlMessage::ControlMessage(const ControlMessage& other)
 {
-    m_cm_type = other.m_cm_type;
-    m_payload = other.m_payload;
-    m_tensors = other.m_tensors;
+    m_cm_type      = other.m_cm_type;
+    m_payload      = other.m_payload;
+    m_tensors      = other.m_tensors;
+    m_tensor_count = other.m_tensor_count;
 
     m_config = other.m_config;
     m_tasks  = other.m_tasks;
@@ -256,7 +257,13 @@ std::shared_ptr<TensorMemory> ControlMessage::tensors()
 
 void ControlMessage::tensors(const std::shared_ptr<TensorMemory>& tensors)
 {
-    m_tensors = tensors;
+    m_tensors      = tensors;
+    m_tensor_count = tensors ? tensors->count : 0;
+}
+
+TensorIndex ControlMessage::tensor_count()
+{
+    return m_tensor_count;
 }
 
 ControlMessageType ControlMessage::task_type()

python/morpheus/morpheus/_lib/src/stages/inference_client_stage.cpp

@@ -58,7 +58,7 @@ static ShapeType get_seq_ids(const std::shared_ptr<ControlMessage>& message)
     auto seq_ids         = message->tensors()->get_tensor("seq_ids");
     const auto item_size = seq_ids.dtype().item_size();
 
-    ShapeType host_seq_ids(message->tensors()->count);
+    ShapeType host_seq_ids(message->tensor_count());
     MRC_CHECK_CUDA(cudaMemcpy2D(host_seq_ids.data(),
                                 item_size,
                                 seq_ids.data(),
@@ -82,7 +82,7 @@ static TensorObject get_tensor(std::shared_ptr<ControlMessage> message, std::str
 
 static void reduce_outputs(std::shared_ptr<ControlMessage> const& message, TensorMap& output_tensors)
 {
-    if (message->payload()->count() == message->tensors()->count)
+    if (message->payload()->count() == message->tensor_count())
     {
         return;
     }

python/morpheus/morpheus/_lib/tests/messages/test_control_message.cpp

@@ -21,6 +21,7 @@
 #include "morpheus/messages/control.hpp"               // for ControlMessage
 #include "morpheus/messages/memory/tensor_memory.hpp"  // for TensorMemory
 #include "morpheus/messages/meta.hpp"                  // for MessageMeta
+#include "morpheus/types.hpp"
 #include "morpheus/utilities/json_types.hpp"           // for PythonByteContainer
 
 #include <gtest/gtest.h>       // for Message, TestPartResult, AssertionResult, TestInfo
@@ -298,7 +299,8 @@ TEST_F(TestControlMessage, SetAndGetTensorMemory)
 {
     auto msg = ControlMessage();
 
-    auto tensorMemory = std::make_shared<TensorMemory>(0);
+    TensorIndex count = 5;
+    auto tensorMemory = std::make_shared<TensorMemory>(count);
     // Optionally, modify tensorMemory here if it has any mutable state to test
 
     // Set the tensor memory
@@ -309,6 +311,7 @@ TEST_F(TestControlMessage, SetAndGetTensorMemory)
 
     // Verify that the retrieved tensor memory matches what was set
     EXPECT_EQ(tensorMemory, retrievedTensorMemory);
+    EXPECT_EQ(count, msg.tensor_count());
 }
 
 // Test setting TensorMemory to nullptr

python/morpheus/morpheus/messages/control_message.py

@@ -46,6 +46,7 @@ def __init__(self, config_or_message: typing.Union["ControlMessage", dict] = Non
 
         self._payload: MessageMeta = None
         self._tensors: TensorMemory = None
+        self._tensor_count: int = 0
 
         self._tasks: dict[str, deque] = defaultdict(deque)
         self._timestamps: dict[str, datetime] = {}
@@ -147,9 +148,13 @@ def payload(self, payload: MessageMeta = None) -> MessageMeta | None:
     def tensors(self, tensors: TensorMemory = None) -> TensorMemory | None:
         if tensors is not None:
             self._tensors = tensors
+            self._tensor_count = tensors.count
 
         return self._tensors
 
+    def tensor_count(self) -> int:
+        return self._tensor_count
+
     def task_type(self, new_task_type: ControlMessageType = None) -> ControlMessageType:
         if new_task_type is not None:
             self._type = new_task_type

python/morpheus/morpheus/stages/inference/identity_inference_stage.py

@@ -45,12 +45,12 @@ def __init__(self, inf_queue: ProducerConsumerQueue, c: Config):
         self._seq_length = c.feature_length
 
     def calc_output_dims(self, msg: ControlMessage) -> typing.Tuple:
-        return (msg.tensors().count, self._seq_length)
+        return (msg.tensor_count(), self._seq_length)
 
     def process(self, batch: ControlMessage, callback: typing.Callable[[TensorMemory], None]):
 
         def tmp(batch: ControlMessage, f):
-            count = batch.tensors().count
+            count = batch.tensor_count()
             f(TensorMemory(
                 count=count,
                 tensors={'probs': cp.zeros((count, self._seq_length), dtype=cp.float32)},

python/morpheus/morpheus/stages/inference/inference_stage.py

@@ -244,7 +244,7 @@ def set_output_fut(resp: TensorMemory, inner_batch, batch_future: mrc.Future):
                         nonlocal outstanding_requests
                         nonlocal batch_offset
                         mess = self._convert_one_response(output_message, inner_batch, resp, batch_offset)
-                        batch_offset += inner_batch.tensors().count
+                        batch_offset += inner_batch.tensor_count()
                         outstanding_requests -= 1
 
                         batch_future.set_result(mess)
@@ -359,13 +359,13 @@ def _convert_one_response(output: ControlMessage, inf: ControlMessage, res: Tens
         seq_count = seq_ids[-1, 0].item() + 1 - seq_offset
 
         # Two scenarios:
-        if (inf.payload().count == inf.tensors().count):
+        if (inf.payload().count == inf.tensor_count()):
             assert seq_count == res.count
 
             # In message and out message have same count. Just use probs as is
             probs[seq_offset:seq_offset + seq_count, :] = resp_probs
         else:
-            assert inf.tensors().count == res.count
+            assert inf.tensor_count() == res.count
 
             mess_ids = seq_ids[:, 0].get().tolist()