Removing support for MultiMessage from stages

docs/source/developer_guide/guides/9_control_messages.md

@@ -79,7 +79,9 @@ msg_meta == retrieved_payload # True
 
 ### Conversion from `MultiMessage` to `ControlMessage`
 
-Starting with version 24.06, the `MultiMessage` type will be deprecated, and all usage should transition to `ControlMessage`. Each `MultiMessage` functionality has a corresponding equivalent in `ControlMessage`, as illustrated below.
+**The `MultiMessage` type was deprecated in 24.06 and has been completely removed in version 24.10.**
+
+When upgrading to 24.10, all uses of `MultiMessage` need to be converted to `ControlMessage`. Each `MultiMessage` functionality has a corresponding equivalent in `ControlMessage`, as illustrated below.
 ```python
 import cudf
 from morpheus.messages import MultiMessage, ControlMessage
@@ -88,12 +90,17 @@ data = cudf.DataFrame()
 msg_meta = MessageMeta(data)
 ```
 
-| **Functionality**                                              | **MultiMessage**                      | **ControlMessage**                                                  |
-| -------------------------------------------------------------- | ------------------------------------- | ------------------------------------------------------------------- |
-| Initialization                                                 | `multi_msg = MultiMessage(msg_meta)`  | `control_msg = ControlMessage()`<br>`control_msg.payload(msg_meta)` |
-| Get columns from `cudf.DataFrame`                              | `multi_msg.get_meta(col_name)`        | `control_msg.payload().get_data(col_name)`                          |
-| Set columns values to `cudf.DataFrame`                         | `multi_msg.set_meta(col_name, value)` | `control_msg.payload().set_data(col_name, value)`                   |
-| Get sliced `cudf.DataFrame` for given start and stop positions | `multi_msg.get_slice(start, stop)`    | `control_msg.payload().get_slice(start, stop)`                      |
-| Copy the `cudf.DataFrame` for given ranges of rows             | `multi_msg.copy_ranges(ranges)`       | `control_msg.payload().copy_ranges(ranges)`                         |
+| **Functionality**                                              | **MultiMessage**                           | **ControlMessage**                                                  |
+| -------------------------------------------------------------- | ------------------------------------------ | ------------------------------------------------------------------- |
+| Initialization                                                 | `multi_msg = MultiMessage(msg_meta)`       | `control_msg = ControlMessage()`<br>`control_msg.payload(msg_meta)` |
+| Get `cudf.DataFrame`                                           | `multi_msg.get_meta()`                     | `control_msg.payload().get_data()`                                  |
+| Get columns from `cudf.DataFrame`                              | `multi_msg.get_meta(col_name)`             | `control_msg.payload().get_data(col_name)`                          |
+| Set columns values to `cudf.DataFrame`                         | `multi_msg.set_meta(col_name, value)`      | `control_msg.payload().set_data(col_name, value)`                   |
+| Get sliced `cudf.DataFrame` for given start and stop positions | `multi_msg.get_slice(start, stop)`         | `control_msg.payload().get_slice(start, stop)`                      |
+| Copy the `cudf.DataFrame` for given ranges of rows             | `multi_msg.copy_ranges(ranges)`            | `control_msg.payload().copy_ranges(ranges)`                         |
+|                                                                | **MultiTensorMessage**                     | **ControlMessage**                                                  |
+| Get the inference tensor `cupy.ndarray`                        | `multi_tensor_msg.tensor()`                | `control_msg.tensors()`                                              |
+| Get a specific inference tensor                                 | `multi_tensor_msg.get_tensor(tensor_name)` | `control_msg.tensors().get_tensor(tensor_name)`                                   |
+
 
-Note that the `get_slice()` and `copy_ranges()` functions in `ControlMessage` return the `MessageMeta` after slicing, whereas these functions in `MultiMessage` return a new `MultiMessage` instance.
+Note that the `get_slice()` and `copy_ranges()` functions in `ControlMessage` return the `MessageMeta` after slicing, whereas these functions in `MultiMessage` return a new `MultiMessage` instance.

examples/abp_pcap_detection/abp_pcap_preprocessing.py

@@ -21,14 +21,12 @@
 
 import cudf
 
+import morpheus._lib.messages as _messages
 from morpheus.cli.register_stage import register_stage
 from morpheus.common import TypeId
 from morpheus.config import Config
 from morpheus.config import PipelineModes
-from morpheus.messages import InferenceMemoryFIL
-from morpheus.messages import MultiInferenceFILMessage
-from morpheus.messages import MultiInferenceMessage
-from morpheus.messages import MultiMessage
+from morpheus.messages import ControlMessage
 from morpheus.stages.preprocess.preprocess_base_stage import PreprocessBaseStage
 
 
@@ -81,20 +79,22 @@ def supports_cpp_node(self):
         return False
 
     @staticmethod
-    def pre_process_batch(x: MultiMessage, fea_len: int, fea_cols: typing.List[str],
-                          req_cols: typing.List[str]) -> MultiInferenceFILMessage:
+    def pre_process_batch(msg: ControlMessage, fea_len: int, fea_cols: typing.List[str],
+                          req_cols: typing.List[str]) -> ControlMessage:
+        meta = msg.payload()
         # Converts the int flags field into a binary string
-        flags_bin_series = x.get_meta("flags").to_pandas().apply(lambda x: format(int(x), "05b"))
+        flags_bin_series = meta.get_data("flags").to_pandas().apply(lambda x: format(int(x), "05b"))
 
         # Expand binary string into an array
-        df = cudf.DataFrame(np.vstack(flags_bin_series.str.findall("[0-1]")).astype("int8"), index=x.get_meta().index)
+        df = cudf.DataFrame(np.vstack(flags_bin_series.str.findall("[0-1]")).astype("int8"),
+                            index=meta.get_data().index)
 
         # adding [ack, psh, rst, syn, fin] details from the binary flag
         rename_cols_dct = {0: "ack", 1: "psh", 2: "rst", 3: "syn", 4: "fin"}
         df = df.rename(columns=rename_cols_dct)
 
         df["flags_bin"] = flags_bin_series
-        df["timestamp"] = x.get_meta("timestamp").astype("int64")
+        df["timestamp"] = meta.get_data("timestamp").astype("int64")
 
         def round_time_kernel(timestamp, rollup_time, secs):
             for i, time in enumerate(timestamp):
@@ -113,8 +113,8 @@ def round_time_kernel(timestamp, rollup_time, secs):
         df["rollup_time"] = cudf.to_datetime(df["rollup_time"], unit="us").dt.strftime("%Y-%m-%d %H:%M")
 
         # creating flow_id "src_ip:src_port=dst_ip:dst_port"
-        df["flow_id"] = (x.get_meta("src_ip") + ":" + x.get_meta("src_port").astype("str") + "=" +
-                         x.get_meta("dest_ip") + ":" + x.get_meta("dest_port").astype("str"))
+        df["flow_id"] = (meta.get_data("src_ip") + ":" + meta.get_data("src_port").astype("str") + "=" +
+                         meta.get_data("dest_ip") + ":" + meta.get_data("dest_port").astype("str"))
         agg_dict = {
             "ack": "sum",
             "psh": "sum",
@@ -125,7 +125,7 @@ def round_time_kernel(timestamp, rollup_time, secs):
             "flow_id": "count",
         }
 
-        df["data_len"] = x.get_meta("data_len").astype("int16")
+        df["data_len"] = meta.get_data("data_len").astype("int16")
 
         # group by operation
         grouped_df = df.groupby(["rollup_time", "flow_id"]).agg(agg_dict)
@@ -175,22 +175,24 @@ def round_time_kernel(timestamp, rollup_time, secs):
         count = data.shape[0]
 
         for col in req_cols:
-            x.set_meta(col, merged_df[col])
+            meta.set_data(col, merged_df[col])
 
         del merged_df
 
         seq_ids = cp.zeros((count, 3), dtype=cp.uint32)
-        seq_ids[:, 0] = cp.arange(x.mess_offset, x.mess_offset + count, dtype=cp.uint32)
+        seq_ids[:, 0] = cp.arange(0, count, dtype=cp.uint32)
         seq_ids[:, 2] = fea_len - 1
 
         # Create the inference memory. Keep in mind count here could be > than input count
-        memory = InferenceMemoryFIL(count=count, input__0=data, seq_ids=seq_ids)
+        memory = _messages.InferenceMemoryFIL(count=count, input__0=data, seq_ids=seq_ids)
 
-        infer_message = MultiInferenceFILMessage.from_message(x, memory=memory)
+        infer_message = ControlMessage(msg)
+        infer_message.payload(meta)
+        infer_message.tensors(memory)
 
         return infer_message
 
-    def _get_preprocess_fn(self) -> typing.Callable[[MultiMessage], MultiInferenceMessage]:
+    def _get_preprocess_fn(self) -> typing.Callable[[ControlMessage], ControlMessage]:
         return partial(AbpPcapPreprocessingStage.pre_process_batch,
                        fea_len=self._fea_length,
                        fea_cols=self.features,

examples/llm/agents/common.py

@@ -21,7 +21,6 @@
 from langchain.llms.openai import OpenAI
 
 from morpheus.config import Config
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.monitor_stage import MonitorStage
 from morpheus.stages.output.in_memory_sink_stage import InMemorySinkStage
@@ -67,8 +66,7 @@ def build_common_pipeline(config: Config, pipe: LinearPipeline, task_payload: di
     Construct the elements of the pipeline common to the simple and kafka agent pipelines.
     This method should be called after the source stage has been set.
     """
-    pipe.add_stage(
-        DeserializeStage(config, message_type=ControlMessage, task_type="llm_engine", task_payload=task_payload))
+    pipe.add_stage(DeserializeStage(config, task_type="llm_engine", task_payload=task_payload))
 
     pipe.add_stage(MonitorStage(config, description="Source rate", unit='questions'))
 

examples/llm/completion/pipeline.py

@@ -20,7 +20,6 @@
 from morpheus.config import Config
 from morpheus.config import PipelineModes
 from morpheus.io.deserializers import read_file_to_df
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.monitor_stage import MonitorStage
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
@@ -116,8 +115,7 @@ def pipeline(num_threads: int,
 
     pipe.set_source(InMemorySourceStage(config, dataframes=[source_df], repeat=repeat_count))
 
-    pipe.add_stage(
-        DeserializeStage(config, message_type=ControlMessage, task_type="llm_engine", task_payload=completion_task))
+    pipe.add_stage(DeserializeStage(config, task_type="llm_engine", task_payload=completion_task))
 
     pipe.add_stage(MonitorStage(config, description="Source rate", unit='questions'))
 

examples/llm/rag/standalone_pipeline.py

@@ -18,7 +18,6 @@
 
 from morpheus.config import Config
 from morpheus.config import PipelineModes
-from morpheus.messages import ControlMessage
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.monitor_stage import MonitorStage
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
@@ -114,8 +113,7 @@ def standalone(num_threads,
 
     pipe.set_source(InMemorySourceStage(config, dataframes=source_dfs, repeat=repeat_count))
 
-    pipe.add_stage(
-        DeserializeStage(config, message_type=ControlMessage, task_type="llm_engine", task_payload=completion_task))
+    pipe.add_stage(DeserializeStage(config, task_type="llm_engine", task_payload=completion_task))
 
     pipe.add_stage(MonitorStage(config, description="Source rate", unit='questions'))
 

examples/log_parsing/inference.py

@@ -19,12 +19,11 @@
 import tritonclient.grpc as tritonclient
 from scipy.special import softmax
 
+import morpheus._lib.messages as _messages
 from morpheus.cli.register_stage import register_stage
 from morpheus.config import Config
 from morpheus.config import PipelineModes
-from morpheus.messages import MultiInferenceMessage
-from morpheus.messages import MultiInferenceNLPMessage
-from morpheus.messages import MultiResponseMessage
+from morpheus.messages import ControlMessage
 from morpheus.messages import TensorMemory
 from morpheus.pipeline.stage_schema import StageSchema
 from morpheus.stages.inference.triton_inference_stage import TritonInferenceStage
@@ -58,28 +57,27 @@ 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) -> MultiResponseMessage:
-        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]
+    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 = TensorMemory(
-            count=x.count,
+        memory = _messages.TensorMemory(
+            count=msg.tensors().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])),
+                '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
             })
 
-        return MultiResponseMessage(meta=x.meta,
-                                    mess_offset=x.mess_offset,
-                                    mess_count=x.mess_count,
-                                    memory=memory,
-                                    offset=0,
-                                    count=x.count)
+        resp = ControlMessage(msg)
+        resp.payload(msg.payload())
+        resp.tensors(memory)
 
-    def _build_response(self, batch: MultiInferenceMessage, result: tritonclient.InferResult) -> TensorMemory:
+        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()}
@@ -140,44 +138,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, inf: MultiInferenceNLPMessage,
-                              res: TensorMemory) -> MultiResponseMessage:
-        memory = output.memory
+    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')
 
-        seq_ids = inf.get_id_tensor()
+        seq_ids = inf.tensors().get_tensor('seq_ids')
 
-        seq_offset = seq_ids[0, 0].item() - output.mess_offset
-        seq_count = (seq_ids[-1, 0].item() + 1 - seq_offset) - output.mess_offset
+        seq_offset = seq_ids[0, 0].item()
+        seq_count = seq_ids[-1, 0].item() + 1 - seq_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
+        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')
 
         # Two scenarios:
-        if (inf.mess_count == inf.count):
+        if (inf.payload().count == inf.tensors().count):
             assert seq_count == res.count
-            confidences[inf.offset:inf.offset + inf.count, :] = resp_confidences
-            labels[inf.offset:inf.offset + inf.count, :] = resp_labels
+            confidences[0:inf.tensors().count, :] = resp_confidences
+            labels[0:inf.tensors().count, :] = resp_labels
         else:
-            assert inf.count == res.count
+            assert inf.tensors().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)
+        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

@@ -27,8 +27,8 @@
 from morpheus.cli.register_stage import register_stage
 from morpheus.config import Config
 from morpheus.config import PipelineModes
+from morpheus.messages import ControlMessage
 from morpheus.messages import MessageMeta
-from morpheus.messages import MultiResponseMessage
 from morpheus.pipeline.single_port_stage import SinglePortStage
 from morpheus.pipeline.stage_schema import StageSchema
 
@@ -77,18 +77,17 @@ def supports_cpp_node(self):
         return False
 
     def accepted_types(self) -> typing.Tuple:
-        return (MultiResponseMessage, )
+        return (ControlMessage, )
 
     def compute_schema(self, schema: StageSchema):
         schema.output_schema.set_type(MessageMeta)
 
-    def _postprocess(self, x: MultiResponseMessage):
-
-        infer_pdf = pd.DataFrame(x.get_tensor('seq_ids').get()).astype(int)
+    def _postprocess(self, msg: ControlMessage):
+        infer_pdf = pd.DataFrame(msg.tensors().get_tensor('seq_ids').get()).astype(int)
         infer_pdf.columns = ["doc", "start", "stop"]
-        infer_pdf["confidences"] = x.get_tensor('confidences').tolist()
-        infer_pdf["labels"] = x.get_tensor('labels').tolist()
-        infer_pdf["token_ids"] = x.get_tensor('input_ids').tolist()
+        infer_pdf["confidences"] = msg.tensors().get_tensor('confidences').tolist()
+        infer_pdf["labels"] = msg.tensors().get_tensor('labels').tolist()
+        infer_pdf["token_ids"] = msg.tensors().get_tensor('input_ids').tolist()
 
         infer_pdf["confidences"] = infer_pdf.apply(lambda row: row["confidences"][row["start"]:row["stop"]], axis=1)