remove from mlflow_drift_stage

morpheus/stages/postprocess/ml_flow_drift_stage.py

@@ -25,7 +25,6 @@
 from morpheus.config import Config
 from morpheus.config import PipelineModes
 from morpheus.messages import ControlMessage
-from morpheus.messages import MultiResponseMessage
 from morpheus.pipeline.pass_thru_type_mixin import PassThruTypeMixin
 from morpheus.pipeline.single_port_stage import SinglePortStage
 
@@ -124,16 +123,13 @@ def accepted_types(self) -> typing.Tuple:
             Accepted input types.
 
         """
-        return (MultiResponseMessage, ControlMessage)
+        return (ControlMessage,)
 
     def supports_cpp_node(self):
         return False
 
-    def _calc_drift(self, x: MultiResponseMessage | ControlMessage):
-        if isinstance(x, MultiResponseMessage):
-            probs_tensor = x.get_probs_tensor()
-        elif isinstance(x, ControlMessage):
-            probs_tensor = x.tensors().get_tensor("probs")
+    def _calc_drift(self, msg: ControlMessage):
+        probs_tensor = msg.tensors().get_tensor("probs")
 
         # All probs in a batch will be calculated
         shifted = cp.abs(probs_tensor - 0.5) + 0.5
@@ -142,10 +138,7 @@ def _calc_drift(self, x: MultiResponseMessage | ControlMessage):
         for label in range(len(self._labels), shifted.shape[1]):
             self._labels.append(str(label))
 
-        if isinstance(x, MultiResponseMessage):
-            count = x.count
-        elif isinstance(x, ControlMessage):
-            count = x.payload().count
+        count = msg.payload().count
 
         for i in list(range(0, count, self._batch_size)):
             start = i
@@ -159,7 +152,7 @@ def _calc_drift(self, x: MultiResponseMessage | ControlMessage):
 
             mlflow.log_metrics(metrics)
 
-        return x
+        return msg
 
     def _build_single(self, builder: mrc.Builder, input_node: mrc.SegmentObject) -> mrc.SegmentObject:
 

tests/stages/test_ml_flow_drift_stage.py

@@ -23,19 +23,10 @@
 
 import morpheus._lib.messages as _messages
 from morpheus.messages import ControlMessage
-from morpheus.messages import MultiResponseMessage
-from morpheus.messages import ResponseMemory
 from morpheus.messages.message_meta import MessageMeta
 from morpheus.stages.postprocess.ml_flow_drift_stage import MLFlowDriftStage
 
 
-def _make_multi_response_message(df, probs):
-    df_ = df[0:len(probs)]
-    mem = ResponseMemory(count=len(df_), tensors={'probs': probs})
-
-    return MultiResponseMessage(meta=MessageMeta(df_), count=len(df_), memory=mem)
-
-
 def _make_control_message(df, probs):
     df_ = df[0:len(probs)]
     cm = ControlMessage()
@@ -51,7 +42,6 @@ def test_constructor(config):
     assert stage.name == "mlflow_drift"
 
     accepted_union = typing.Union[stage.accepted_types()]
-    assert typing_utils.issubtype(MultiResponseMessage, accepted_union)
     assert typing_utils.issubtype(ControlMessage, accepted_union)
 
 
@@ -63,7 +53,7 @@ def test_calc_drift(config, filter_probs_df):
         stage = MLFlowDriftStage(config, labels=labels, batch_size=1)
 
     probs = cp.array([[0.1, 0.5, 0.3], [0.2, 0.3, 0.4]])
-    mock_multi_response_message = _make_multi_response_message(filter_probs_df, probs)
+
     mock_control_message = _make_control_message(filter_probs_df, probs)
 
     expected_metrics = [{
@@ -72,16 +62,9 @@ def test_calc_drift(config, filter_probs_df):
         'a': 0.8, 'b': 0.7, 'c': 0.6, 'total': 0.7000000000000001
     }]
 
-    multi_response_message_metrics = []
-    with patch("morpheus.stages.postprocess.ml_flow_drift_stage.mlflow.log_metrics") as mock_log_metrics:
-        stage._calc_drift(mock_multi_response_message)
-        for call_arg in mock_log_metrics.call_args_list:
-            multi_response_message_metrics.append(call_arg[0][0])
-    assert multi_response_message_metrics == expected_metrics
-
     control_message_metrics = []
     with patch("morpheus.stages.postprocess.ml_flow_drift_stage.mlflow.log_metrics") as mock_log_metrics:
         stage._calc_drift(mock_control_message)
         for call_arg in mock_log_metrics.call_args_list:
             control_message_metrics.append(call_arg[0][0])
-    assert control_message_metrics == multi_response_message_metrics
+    assert control_message_metrics == expected_metrics