flowcept 0.8.10__py3-none-any.whl → 0.8.12__py3-none-any.whl

flowcept/commons/daos/mq_dao/mq_dao_base.py

@@ -1,15 +1,12 @@
 """MQ base module."""
 
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from typing import Union, List, Callable
 import csv
 import msgpack
 from time import time
 import flowcept.commons
 from flowcept.commons.autoflush_buffer import AutoflushBuffer
-
-from flowcept.commons.daos.keyvalue_dao import KeyValueDAO
-
 from flowcept.commons.utils import chunked
 from flowcept.commons.flowcept_logger import FlowceptLogger
 from flowcept.configs import (
@@ -20,12 +17,15 @@ from flowcept.configs import (
     MQ_CHUNK_SIZE,
     MQ_TYPE,
     MQ_TIMING,
+    KVDB_ENABLED,
+    MQ_ENABLED,
+    DUMP_BUFFER_PATH,
 )
 
 from flowcept.commons.utils import GenericJSONEncoder
 
 
-class MQDao(ABC):
+class MQDao(object):
     """MQ base class."""
 
     ENCODER = GenericJSONEncoder if JSON_SERIALIZER == "complex" else None
@@ -34,6 +34,9 @@ class MQDao(ABC):
     @staticmethod
     def build(*args, **kwargs) -> "MQDao":
         """Build it."""
+        if not MQ_ENABLED:
+            return MQDao()
+
         if MQ_TYPE == "redis":
             from flowcept.commons.daos.mq_dao.mq_dao_redis import MQDaoRedis
 
@@ -67,7 +70,12 @@ class MQDao(ABC):
         self.logger = FlowceptLogger()
         self.started = False
         self._adapter_settings = adapter_settings
-        self._keyvalue_dao = KeyValueDAO()
+        if KVDB_ENABLED:
+            from flowcept.commons.daos.keyvalue_dao import KeyValueDAO
+
+            self._keyvalue_dao = KeyValueDAO()
+        else:
+            self._keyvalue_dao = None
         self._time_based_flushing_started = False
         self.buffer: Union[AutoflushBuffer, List] = None
         if MQ_TIMING:
@@ -88,11 +96,21 @@ class MQDao(ABC):
     def bulk_publish(self, buffer):
         """Publish it."""
         # self.logger.info(f"Going to flush {len(buffer)} to MQ...")
-        if MQ_CHUNK_SIZE > 1:
-            for chunk in chunked(buffer, MQ_CHUNK_SIZE):
-                self._bulk_publish(chunk)
+        if flowcept.configs.DB_FLUSH_MODE == "offline":
+            if DUMP_BUFFER_PATH is not None:
+                import orjson
+
+                with open(DUMP_BUFFER_PATH, "wb", buffering=1_048_576) as f:
+                    for obj in buffer:
+                        f.write(orjson.dumps(obj))
+                        f.write(b"\n")
+                self.logger.info(f"Saved Flowcept messages into {DUMP_BUFFER_PATH}.")
         else:
-            self._bulk_publish(buffer)
+            if MQ_CHUNK_SIZE > 1:
+                for chunk in chunked(buffer, MQ_CHUNK_SIZE):
+                    self._bulk_publish(chunk)
+            else:
+                self._bulk_publish(buffer)
 
     def register_time_based_thread_init(self, interceptor_instance_id: str, exec_bundle_id=None):
         """Register the time."""
@@ -138,7 +156,7 @@ class MQDao(ABC):
         """
         self._keyvalue_dao.delete_key("current_campaign_id")
 
-    def init_buffer(self, interceptor_instance_id: str, exec_bundle_id=None):
+    def init_buffer(self, interceptor_instance_id: str, exec_bundle_id=None, check_safe_stops=True):
         """Create the buffer."""
         if not self.started:
             if flowcept.configs.DB_FLUSH_MODE == "online":
@@ -147,7 +165,8 @@ class MQDao(ABC):
                     max_size=MQ_BUFFER_SIZE,
                     flush_interval=MQ_INSERTION_BUFFER_TIME,
                 )
-                self.register_time_based_thread_init(interceptor_instance_id, exec_bundle_id)
+                if check_safe_stops:
+                    self.register_time_based_thread_init(interceptor_instance_id, exec_bundle_id)
                 self._time_based_flushing_started = True
             else:
                 self.buffer = list()
@@ -164,9 +183,9 @@ class MQDao(ABC):
             self.bulk_publish(self.buffer)
             self.buffer = list()
 
-    def _stop_timed(self, interceptor_instance_id: str, bundle_exec_id: int = None):
+    def _stop_timed(self, interceptor_instance_id: str, check_safe_stops: bool = True, bundle_exec_id: int = None):
         t1 = time()
-        self._stop(interceptor_instance_id, bundle_exec_id)
+        self._stop(interceptor_instance_id, check_safe_stops, bundle_exec_id)
         t2 = time()
         self._flush_events.append(["final", t1, t2, t2 - t1, "n/a"])
 
@@ -175,14 +194,14 @@ class MQDao(ABC):
             writer.writerow(["type", "start", "end", "duration", "size"])
             writer.writerows(self._flush_events)
 
-    def _stop(self, interceptor_instance_id: str, bundle_exec_id: int = None):
-        """Stop it."""
-        msg0 = "MQ publisher received stop signal! bundle: "
-        self.logger.debug(msg0 + f"{bundle_exec_id}; interceptor id: {interceptor_instance_id}")
+    def _stop(self, interceptor_instance_id: str = None, check_safe_stops: bool = True, bundle_exec_id: int = None):
+        """Stop MQ publisher."""
+        self.logger.debug(f"MQ pub received stop sign: bundle={bundle_exec_id}, interceptor={interceptor_instance_id}")
         self._close_buffer()
-        msg = "Flushed MQ for last time! Send stop msg. bundle: "
-        self.logger.debug(msg + f"{bundle_exec_id}; interceptor id: {interceptor_instance_id}")
-        self._send_mq_dao_time_thread_stop(interceptor_instance_id, bundle_exec_id)
+        self.logger.debug("Flushed MQ for the last time!")
+        if check_safe_stops:
+            self.logger.debug(f"Sending stop msg. Bundle: {bundle_exec_id}; interceptor id: {interceptor_instance_id}")
+            self._send_mq_dao_time_thread_stop(interceptor_instance_id, bundle_exec_id)
         self.started = False
 
     def _send_mq_dao_time_thread_stop(self, interceptor_instance_id, exec_bundle_id=None):
@@ -197,10 +216,10 @@ class MQDao(ABC):
         # self.logger.info("Control msg sent: " + str(msg))
         self.send_message(msg)
 
-    def send_document_inserter_stop(self):
+    def send_document_inserter_stop(self, exec_bundle_id=None):
         """Send the document."""
         # These control_messages are handled by the document inserter
-        msg = {"type": "flowcept_control", "info": "stop_document_inserter"}
+        msg = {"type": "flowcept_control", "info": "stop_document_inserter", "exec_bundle_id": exec_bundle_id}
         self.send_message(msg)
 
     @abstractmethod
@@ -223,20 +242,12 @@ class MQDao(ABC):
         """Subscribe to the interception channel."""
         raise NotImplementedError()
 
+    @abstractmethod
+    def unsubscribe(self):
+        """Subscribe to the interception channel."""
+        raise NotImplementedError()
+
     @abstractmethod
     def liveness_test(self) -> bool:
-        """Check whether the base KV store's connection is ready. This is enough for Redis MQ too. Other MQs need
-        further liveness implementation.
-        """
-        try:
-            response = self._keyvalue_dao.redis_conn.ping()
-            if response:
-                return True
-            else:
-                return False
-        except ConnectionError as e:
-            self.logger.exception(e)
-            return False
-        except Exception as e:
-            self.logger.exception(e)
-            return False
+        """Checks if the MQ system is alive."""
+        raise NotImplementedError()

flowcept/commons/daos/mq_dao/mq_dao_kafka.py

@@ -108,12 +108,29 @@ class MQDaoKafka(MQDao):
     def liveness_test(self):
         """Get the livelyness of it."""
         try:
-            if not super().liveness_test():
-                self.logger.error("KV Store not alive!")
-                return False
             admin_client = AdminClient(self._kafka_conf)
             kafka_metadata = admin_client.list_topics(timeout=5)
             return MQ_CHANNEL in kafka_metadata.topics
         except Exception as e:
             self.logger.exception(e)
             return False
+
+    def unsubscribe(self):
+        """Unsubscribes from Kafka topic and closes consumer if open."""
+        if self._consumer is None:
+            self.logger.warning("No Kafka consumer to unsubscribe.")
+            return
+
+        try:
+            self._consumer.unsubscribe()
+            self.logger.info("Unsubscribed from Kafka topics.")
+        except RuntimeError as e:
+            self.logger.debug(f"Consumer already closed while unsubscribing: {e}")
+        except Exception as e:
+            self.logger.exception(f"Error while unsubscribing from Kafka: {e}")
+        finally:
+            try:
+                self._consumer.close()
+            except Exception as e:
+                self.logger.debug(f"Error closing consumer after unsubscribe: {e}")
+            self._consumer = None

flowcept/commons/daos/mq_dao/mq_dao_mofka.py

@@ -104,3 +104,7 @@ class MQDaoMofka(MQDao):
     def liveness_test(self):
         """Test Mofka Liveness."""
         return True
+
+    def unsubscribe(self):
+        """Unsubscribes from Mofka topic."""
+        raise NotImplementedError()

flowcept/commons/daos/mq_dao/mq_dao_redis.py

@@ -7,7 +7,8 @@ import msgpack
 from time import time, sleep
 
 from flowcept.commons.daos.mq_dao.mq_dao_base import MQDao
-from flowcept.configs import MQ_CHANNEL
+from flowcept.commons.daos.redis_conn import RedisConn
+from flowcept.configs import MQ_CHANNEL, MQ_HOST, MQ_PORT, MQ_PASSWORD, MQ_URI, MQ_SETTINGS, KVDB_ENABLED
 
 
 class MQDaoRedis(MQDao):
@@ -17,16 +18,32 @@ class MQDaoRedis(MQDao):
 
     def __init__(self, adapter_settings=None):
         super().__init__(adapter_settings)
-        self._producer = self._keyvalue_dao.redis_conn  # if MQ is redis, we use the same KV for the MQ
+
         self._consumer = None
+        use_same_as_kv = MQ_SETTINGS.get("same_as_kvdb", False)
+        if use_same_as_kv:
+            if KVDB_ENABLED:
+                self._producer = self._keyvalue_dao.redis_conn
+            else:
+                raise Exception("You have same_as_kvdb in your settings, but kvdb is disabled.")
+        else:
+            self._producer = RedisConn.build_redis_conn_pool(
+                host=MQ_HOST, port=MQ_PORT, password=MQ_PASSWORD, uri=MQ_URI
+            )
 
     def subscribe(self):
         """
         Subscribe to interception channel.
         """
-        self._consumer = self._keyvalue_dao.redis_conn.pubsub()
+        self._consumer = self._producer.pubsub()
         self._consumer.psubscribe(MQ_CHANNEL)
 
+    def unsubscribe(self):
+        """
+        Unsubscribe to interception channel.
+        """
+        self._consumer.unsubscribe(MQ_CHANNEL)
+
     def message_listener(self, message_handler: Callable):
         """Get message listener with automatic reconnection."""
         max_retrials = 10
@@ -37,13 +54,22 @@ class MQDaoRedis(MQDao):
                 for message in self._consumer.listen():
                     if message and message["type"] in MQDaoRedis.MESSAGE_TYPES_IGNORE:
                         continue
+
+                    if not isinstance(message["data"], (bytes, bytearray)):
+                        self.logger.warning(
+                            f"Skipping message with unexpected data type: {type(message['data'])} - {message['data']}"
+                        )
+                        continue
+
                     try:
                         msg_obj = msgpack.loads(message["data"], strict_map_key=False)
+                        # self.logger.debug(f"In mq dao redis, received msg!  {msg_obj}")
                         if not message_handler(msg_obj):
                             should_continue = False  # Break While loop
                             break  # Break For loop
                     except Exception as e:
-                        self.logger.error(f"Failed to process message: {e}")
+                        self.logger.error(f"Failed to process message {message}")
+                        self.logger.exception(e)
 
                     current_trials = 0
             except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
@@ -103,7 +129,14 @@ class MQDaoRedis(MQDao):
     def liveness_test(self):
         """Get the livelyness of it."""
         try:
-            return super().liveness_test()
+            response = self._producer.ping()
+            if response:
+                return True
+            else:
+                return False
+        except ConnectionError as e:
+            self.logger.exception(e)
+            return False
         except Exception as e:
             self.logger.exception(e)
             return False

flowcept/commons/daos/redis_conn.py

@@ -0,0 +1,47 @@
+"""RedisConn module."""
+
+from redis import Redis, ConnectionPool
+
+
+class RedisConn:
+    """RedisConn DAO class."""
+
+    @staticmethod
+    def build_redis_conn_pool(host: str = None, port: str = None, password: str = None, uri: str = None) -> Redis:
+        """
+        Create a Redis connection using either a URI or host/port.
+
+        If `uri` is provided, it will be used to initialize the Redis connection.
+        Otherwise, the connection will fall back to using `host` and `port`.
+
+        Parameters
+        ----------
+        host : str, optional
+            Redis host address. Used only if `uri` is not provided.
+        port : str, optional
+            Redis port. Used only if `uri` is not provided.
+        uri : str, optional
+            Full Redis URI. Takes precedence over `host` and `port` if defined.
+        password : str, optional
+            Password for authenticating with Redis.
+
+        Returns
+        -------
+        Redis
+            An instance of the Redis client with a configured connection pool.
+        """
+        pool_kwargs = {
+            "db": 0,
+            "password": password,
+            "decode_responses": False,
+            "max_connections": 10000,
+            "socket_keepalive": True,
+            "retry_on_timeout": True,
+        }
+
+        if uri:
+            pool = ConnectionPool.from_url(uri, **pool_kwargs)
+        else:
+            pool = ConnectionPool(host=host, port=port, **pool_kwargs)
+
+        return Redis(connection_pool=pool)

flowcept/commons/flowcept_dataclasses/task_object.py

@@ -24,13 +24,14 @@ class TaskObject:
     utc_timestamp: float = None
     adapter_id: AnyStr = None
     user: AnyStr = None
+    data: Any = None
     used: Dict[AnyStr, Any] = None  # Used parameter and files
     campaign_id: AnyStr = None
     generated: Dict[AnyStr, Any] = None  # Generated results and files
     submitted_at: float = None
     started_at: float = None
     ended_at: float = None
-    registered_at: float = None
+    registered_at: float = None  # Leave this for dates generated at the DocInserter
     telemetry_at_start: Telemetry = None
     telemetry_at_end: Telemetry = None
     workflow_name: AnyStr = None
@@ -52,6 +53,8 @@ class TaskObject:
     address: AnyStr = None
     dependencies: List = None
     dependents: List = None
+    tags: List = None
+    agent_id: str = None
 
     _DEFAULT_ENRICH_VALUES = {
         "node_name": NODE_NAME,
@@ -103,20 +106,16 @@ class TaskObject:
         if self.utc_timestamp is None:
             self.utc_timestamp = flowcept.commons.utils.get_utc_now()
 
-        if self.node_name is None and NODE_NAME is not None:
-            self.node_name = NODE_NAME
-
-        if self.login_name is None and LOGIN_NAME is not None:
-            self.login_name = LOGIN_NAME
-
-        if self.public_ip is None and PUBLIC_IP is not None:
-            self.public_ip = PUBLIC_IP
-
-        if self.private_ip is None and PRIVATE_IP is not None:
-            self.private_ip = PRIVATE_IP
+        for key, fallback_value in TaskObject._DEFAULT_ENRICH_VALUES.items():
+            if getattr(self, key) is None and fallback_value is not None:
+                setattr(self, key, fallback_value)
 
-        if self.hostname is None and HOSTNAME is not None:
-            self.hostname = HOSTNAME
+    @staticmethod
+    def enrich_task_dict(task_dict: dict):
+        """Enrich the task."""
+        for key, fallback_value in TaskObject._DEFAULT_ENRICH_VALUES.items():
+            if (key not in task_dict or task_dict[key] is None) and fallback_value is not None:
+                task_dict[key] = fallback_value
 
     def to_dict(self):
         """Convert to dictionary."""
@@ -139,16 +138,40 @@ class TaskObject:
         return msgpack.dumps(self.to_dict())
 
     @staticmethod
-    def enrich_task_dict(task_dict: dict):
-        """Enrich the task."""
-        for key, fallback_value in TaskObject._DEFAULT_ENRICH_VALUES.items():
-            if (key not in task_dict or task_dict[key] is None) and fallback_value is not None:
-                task_dict[key] = fallback_value
-
-    # @staticmethod
-    # def deserialize(serialized_data) -> 'TaskObject':
-    #     dict_obj = msgpack.loads(serialized_data)
-    #     obj = TaskObject()
-    #     for k, v in dict_obj.items():
-    #         setattr(obj, k, v)
-    #     return obj
+    def from_dict(task_obj_dict: Dict[AnyStr, Any]) -> "TaskObject":
+        """Create a TaskObject from a dictionary.
+
+        Parameters
+        ----------
+        task_obj_dict : Dict[AnyStr, Any]
+            Dictionary containing task attributes.
+
+        Returns
+        -------
+        TaskObject
+            A TaskObject instance populated with available data.
+        """
+        task = TaskObject()
+
+        for key, value in task_obj_dict.items():
+            if hasattr(task, key):
+                if key == "status" and isinstance(value, str):
+                    setattr(task, key, Status(value))
+                else:
+                    setattr(task, key, value)
+
+        return task
+
+    def __str__(self):
+        """Return a user-friendly string representation of the TaskObject."""
+        return self.__repr__()
+
+    def __repr__(self):
+        """Return an unambiguous string representation of the TaskObject."""
+        attrs = ["task_id", "workflow_id", "campaign_id", "activity_id", "started_at", "ended_at"]
+        optionals = ["subtype", "parent_task_id", "agent_id"]
+        for opt in optionals:
+            if getattr(self, opt) is not None:
+                attrs.append(opt)
+        attr_str = ", ".join(f"{attr}={repr(getattr(self, attr))}" for attr in attrs)
+        return f"TaskObject({attr_str})"

flowcept/commons/flowcept_dataclasses/workflow_object.py

@@ -5,7 +5,7 @@ import msgpack
 from omegaconf import OmegaConf, DictConfig
 
 from flowcept.version import __version__
-from flowcept.commons.utils import get_utc_now
+from flowcept.commons.utils import get_utc_now, get_git_info
 from flowcept.configs import (
     settings,
     FLOWCEPT_USER,
@@ -38,6 +38,7 @@ class WorkflowObject:
     sys_name: str = None
     extra_metadata: str = None
     used: Dict = None
+    code_repository: Dict = None
     generated: Dict = None
 
     def __init__(self, workflow_id=None, name=None, used=None, generated=None):
@@ -93,6 +94,13 @@ class WorkflowObject:
             )
             self.extra_metadata = _extra_metadata
 
+        if self.code_repository is None:
+            try:
+                self.code_repository = get_git_info()
+            except Exception as e:
+                print(e)
+                pass
+
         if self.flowcept_version is None:
             self.flowcept_version = __version__
 

flowcept/commons/settings_factory.py

@@ -7,9 +7,6 @@ from flowcept.commons.flowcept_dataclasses.base_settings_dataclasses import (
     BaseSettings,
     KeyValue,
 )
-from flowcept.flowceptor.adapters.zambeze.zambeze_dataclasses import (
-    ZambezeSettings,
-)
 from flowcept.flowceptor.adapters.mlflow.mlflow_dataclasses import (
     MLFlowSettings,
 )
@@ -22,7 +19,6 @@ from flowcept.flowceptor.adapters.dask.dask_dataclasses import (
 
 
 SETTINGS_CLASSES = {
-    Vocabulary.Settings.ZAMBEZE_KIND: ZambezeSettings,
     Vocabulary.Settings.MLFLOW_KIND: MLFlowSettings,
     Vocabulary.Settings.TENSORBOARD_KIND: TensorboardSettings,
     Vocabulary.Settings.DASK_KIND: DaskSettings,
@@ -30,6 +26,8 @@ SETTINGS_CLASSES = {
 
 
 def _build_base_settings(kind: str, settings_dict: dict) -> BaseSettings:
+    if kind not in SETTINGS_CLASSES:
+        return settings_dict
     settings_obj = SETTINGS_CLASSES[kind](**settings_dict)
     return settings_obj