diff --git a/python/morpheus/morpheus/messages/control_message.py b/python/morpheus/morpheus/messages/control_message.py index b179fb464e..5bfa48ca73 100644 --- a/python/morpheus/morpheus/messages/control_message.py +++ b/python/morpheus/morpheus/messages/control_message.py @@ -40,6 +40,8 @@ class ControlMessage(MessageBase, cpp_class=_messages.ControlMessage): def __init__(self, config: dict = None): super().__init__() + self._config: dict = {"metadata": {}} + self._payload: MessageMeta = None self._tensors: TensorMemory = None @@ -72,6 +74,24 @@ def config(self, config: dict = None) -> dict: return self._config + def copy(self) -> "ControlMessage": + config = self._config.copy() + config["type"] = self.task_type().name + + tasks = [] + for (task_type, task_queue) in self.get_tasks().items(): + for task in task_queue: + tasks.append({"type": task_type, "properties": task}) + + config["tasks"] = tasks + + new_cm = ControlMessage(config) + new_cm._payload = self._payload + new_cm._tensors = self._tensors + new_cm._timestamps = self._timestamps.copy() + + return new_cm + def has_task(self, task_type: str) -> bool: """ Return True if the control message has at least one task of the given type