flowcept 0.9.17__py3-none-any.whl → 0.9.19__py3-none-any.whl

flowcept/commons/daos/mq_dao/mq_dao_base.py

@@ -7,6 +7,7 @@ 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 (
@@ -29,6 +30,8 @@ class MQDao(object):
 
     ENCODER = GenericJSONEncoder if JSON_SERIALIZER == "complex" else None
     # TODO we don't have a unit test to cover complex dict!
+    MQ_THREAD_SET_ID = "started_mq_thread_execution"
+    MQ_FLUSH_COMPLETE_SET_ID = "pending_mq_flush_complete"
 
     @staticmethod
     def build(*args, **kwargs) -> "MQDao":
@@ -51,20 +54,6 @@ class MQDao(object):
         else:
             raise NotImplementedError
 
-    @staticmethod
-    def _get_set_name(exec_bundle_id=None):
-        """Get the set name.
-
-        :param exec_bundle_id: A way to group one or many interceptors, and
-         treat each group as a bundle to control when their time_based
-         threads started and ended.
-        :return:
-        """
-        set_id = "started_mq_thread_execution"
-        if exec_bundle_id is not None:
-            set_id += "_" + str(exec_bundle_id)
-        return set_id
-
     def __init__(self, adapter_settings=None):
         self.logger = FlowceptLogger()
         self.started = False
@@ -103,22 +92,36 @@ class MQDao(object):
 
     def register_time_based_thread_init(self, interceptor_instance_id: str, exec_bundle_id=None):
         """Register the time."""
-        set_name = MQDao._get_set_name(exec_bundle_id)
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_THREAD_SET_ID, exec_bundle_id)
         # self.logger.info(
         #     f"Register start of time_based MQ flush thread {set_name}.{interceptor_instance_id}"
         # )
         self._keyvalue_dao.add_key_into_set(set_name, interceptor_instance_id)
+        flush_set_name = KeyValueDAO.get_set_name(MQDao.MQ_FLUSH_COMPLETE_SET_ID, exec_bundle_id)
+        self._keyvalue_dao.add_key_into_set(flush_set_name, interceptor_instance_id)
 
     def register_time_based_thread_end(self, interceptor_instance_id: str, exec_bundle_id=None):
         """Register time."""
-        set_name = MQDao._get_set_name(exec_bundle_id)
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_THREAD_SET_ID, exec_bundle_id)
         self.logger.info(f"Registering end of time_based MQ flush thread {set_name}.{interceptor_instance_id}")
         self._keyvalue_dao.remove_key_from_set(set_name, interceptor_instance_id)
         self.logger.info(f"Done registering time_based MQ flush thread {set_name}.{interceptor_instance_id}")
 
     def all_time_based_threads_ended(self, exec_bundle_id=None):
         """Get all time."""
-        set_name = MQDao._get_set_name(exec_bundle_id)
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_THREAD_SET_ID, exec_bundle_id)
+        return self._keyvalue_dao.set_is_empty(set_name)
+
+    def register_flush_complete(self, interceptor_instance_id: str, exec_bundle_id=None):
+        """Register a flush-complete signal for an interceptor."""
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_FLUSH_COMPLETE_SET_ID, exec_bundle_id)
+        self.logger.info(f"Registering flush completion {set_name}.{interceptor_instance_id}")
+        self._keyvalue_dao.remove_key_from_set(set_name, interceptor_instance_id)
+        self.logger.info(f"Done registering flush completion {set_name}.{interceptor_instance_id}")
+
+    def all_flush_complete_received(self, exec_bundle_id=None):
+        """Return True when all interceptors in the bundle reported flush completion."""
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_FLUSH_COMPLETE_SET_ID, exec_bundle_id)
         return self._keyvalue_dao.set_is_empty(set_name)
 
     def set_campaign_id(self, campaign_id=None):
@@ -172,11 +175,14 @@ class MQDao(object):
             if self._time_based_flushing_started:
                 self.buffer.stop()
                 self._time_based_flushing_started = False
+                self.logger.debug("MQ time-based flushed for the last time!")
             else:
                 self.logger.error("MQ time-based flushing is not started")
         else:
             self.buffer = list()
 
+        self.logger.debug("Buffer closed.")
+
     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, check_safe_stops, bundle_exec_id)
@@ -190,10 +196,12 @@ class MQDao(object):
 
     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()
-        self.logger.debug("Flushed MQ for the last time!")
-        if check_safe_stops:
+        if check_safe_stops and MQ_ENABLED:
+            self.logger.debug(
+                f"Sending flush-complete msg. Bundle: {bundle_exec_id}; interceptor id: {interceptor_instance_id}"
+            )
+            self._send_mq_dao_flush_complete(interceptor_instance_id, bundle_exec_id)
             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
@@ -210,6 +218,15 @@ class MQDao(object):
         # self.logger.info("Control msg sent: " + str(msg))
         self.send_message(msg)
 
+    def _send_mq_dao_flush_complete(self, interceptor_instance_id, exec_bundle_id=None):
+        msg = {
+            "type": "flowcept_control",
+            "info": "mq_flush_complete",
+            "interceptor_instance_id": interceptor_instance_id,
+            "exec_bundle_id": exec_bundle_id,
+        }
+        self.send_message(msg)
+
     def send_document_inserter_stop(self, exec_bundle_id=None):
         """Send the document."""
         # These control_messages are handled by the document inserter

flowcept/commons/daos/mq_dao/mq_dao_kafka.py

@@ -42,7 +42,7 @@ class MQDaoKafka(MQDao):
     def message_listener(self, message_handler: Callable):
         """Get message listener."""
         try:
-            while True:
+            while self._consumer is not None:
                 msg = self._consumer.poll(1.0)
                 if msg is None:
                     continue
@@ -59,7 +59,7 @@ class MQDaoKafka(MQDao):
         except Exception as e:
             self.logger.exception(e)
         finally:
-            self._consumer.close()
+            self.unsubscribe()
 
     def send_message(self, message: dict, channel=MQ_CHANNEL, serializer=msgpack.dumps):
         """Send the message."""

flowcept/commons/daos/mq_dao/mq_dao_redis.py

@@ -1,5 +1,6 @@
 """MQ redis module."""
 
+from threading import Thread
 from typing import Callable
 import redis
 
@@ -14,12 +15,15 @@ from flowcept.configs import MQ_CHANNEL, MQ_HOST, MQ_PORT, MQ_PASSWORD, MQ_URI,
 class MQDaoRedis(MQDao):
     """MQ redis class."""
 
-    MESSAGE_TYPES_IGNORE = {"psubscribe"}
+    MESSAGE_TYPES_IGNORE = {"psubscribe", "subscribe", "pong"}
 
     def __init__(self, adapter_settings=None):
         super().__init__(adapter_settings)
 
         self._consumer = None
+        self._ping_thread = None
+        self._ping_stop = False
+
         use_same_as_kv = MQ_SETTINGS.get("same_as_kvdb", False)
         if use_same_as_kv:
             if KVDB_ENABLED:
@@ -37,6 +41,26 @@ class MQDaoRedis(MQDao):
         """
         self._consumer = self._producer.pubsub()
         self._consumer.psubscribe(MQ_CHANNEL)
+        self._start_ping_thread()
+
+    def _start_ping_thread(self, interval: int = 30):
+        """Start a background thread to ping Redis pubsub periodically."""
+        if self._ping_thread and self._ping_thread.is_alive():
+            return
+
+        self._ping_stop = False
+
+        def _pinger():
+            while not self._ping_stop:
+                try:
+                    if self._consumer is not None:
+                        self._consumer.ping()
+                except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
+                    self.logger.critical(f"Redis PubSub ping failed: {e}")
+                sleep(interval)
+
+        self._ping_thread = Thread(target=_pinger, daemon=True)
+        self._ping_thread.start()
 
     def unsubscribe(self):
         """
@@ -75,8 +99,15 @@ class MQDaoRedis(MQDao):
                     current_trials = 0
             except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
                 current_trials += 1
-                self.logger.critical(f"Redis connection lost: {e}. Reconnecting in 3 seconds...")
+                self.logger.critical(f"Redis connection lost: {e}. Trying to reconnect in 3 seconds...")
                 sleep(3)
+                try:
+                    self.subscribe()
+                    self.logger.warning(f"Redis reconnected after {current_trials} trials.")
+                    current_trials = 0
+                except Exception as e:
+                    self.logger.critical(f"Redis error when trying to reconnect: {e}.")
+
             except Exception as e:
                 self.logger.exception(e)
                 continue

flowcept/commons/flowcept_dataclasses/task_object.py

@@ -134,7 +134,10 @@ class TaskObject:
     """User-defined tags attached to the task."""
 
     agent_id: str = None
-    """Identifier of the agent responsible for executing this task (if any)."""
+    """Identifier of the agent that executed (or is going to execute) this task."""
+
+    source_agent_id: str = None
+    """Identifier of the agent that sent this task to be executed (if any)."""
 
     _DEFAULT_ENRICH_VALUES = {
         "node_name": NODE_NAME,

flowcept/configs.py

@@ -9,7 +9,7 @@ from flowcept.version import __version__
 PROJECT_NAME = "flowcept"
 
 DEFAULT_SETTINGS = {
-    "version": __version__,
+    "flowcept_version": __version__,
     "log": {"log_file_level": "disable", "log_stream_level": "disable"},
     "project": {"dump_buffer": {"enabled": True}},
     "telemetry_capture": {},
@@ -81,7 +81,7 @@ FLOWCEPT_USER = settings["experiment"].get("user", "blank_user")
 
 MQ_INSTANCES = settings["mq"].get("instances", None)
 MQ_SETTINGS = settings["mq"]
-MQ_ENABLED = os.getenv("MQ_ENABLED", settings["mq"].get("enabled", True))
+MQ_ENABLED = os.getenv("MQ_ENABLED", str(settings["mq"].get("enabled", True))).strip().lower() in _TRUE_VALUES
 MQ_TYPE = os.getenv("MQ_TYPE", settings["mq"].get("type", "redis"))
 MQ_CHANNEL = os.getenv("MQ_CHANNEL", settings["mq"].get("channel", "interception"))
 MQ_PASSWORD = settings["mq"].get("password", None)
@@ -103,6 +103,11 @@ KVDB_PORT = int(os.getenv("KVDB_PORT", settings["kv_db"].get("port", "6379")))
 KVDB_URI = os.getenv("KVDB_URI", settings["kv_db"].get("uri", None))
 KVDB_ENABLED = settings["kv_db"].get("enabled", False)
 
+if MQ_ENABLED and not KVDB_ENABLED:
+    raise ValueError(
+        "Invalid configuration: MQ is enabled but kv_db is disabled. "
+        "Enable kv_db.enabled (and KVDB) when MQ is enabled."
+    )
 
 DATABASES = settings.get("databases", {})
 
@@ -155,14 +160,22 @@ DB_INSERTER_SLEEP_TRIALS_STOP = db_buffer_settings.get("stop_trials_sleep", 0.01
 ###########################
 
 DB_FLUSH_MODE = settings["project"].get("db_flush_mode", "offline")
-# DEBUG_MODE = settings["project"].get("debug", False)
 PERF_LOG = settings["project"].get("performance_logging", False)
 JSON_SERIALIZER = settings["project"].get("json_serializer", "default")
 REPLACE_NON_JSON_SERIALIZABLE = settings["project"].get("replace_non_json_serializable", True)
 ENRICH_MESSAGES = settings["project"].get("enrich_messages", True)
 
+if DB_FLUSH_MODE == "online" and not MQ_ENABLED:
+    raise ValueError(
+        "Invalid configuration: project.db_flush_mode is 'online' but MQ is disabled. "
+        "Enable mq.enabled (or MQ_ENABLED=true) or set project.db_flush_mode to 'offline'."
+    )
+
+# Default: enable dump buffer only when running in offline flush mode.
 _DEFAULT_DUMP_BUFFER_ENABLED = DB_FLUSH_MODE == "offline"
 DUMP_BUFFER_ENABLED = (
+    # Env var "DUMP_BUFFER" overrides settings.yaml.
+    # Falls back to settings project.dump_buffer.enabled, then to the default above.
     os.getenv(
         "DUMP_BUFFER", str(settings["project"].get("dump_buffer", {}).get("enabled", _DEFAULT_DUMP_BUFFER_ENABLED))
     )
@@ -170,6 +183,7 @@ DUMP_BUFFER_ENABLED = (
     .lower()
     in _TRUE_VALUES
 )
+# Path is only read from settings.yaml; env override is not supported here.
 DUMP_BUFFER_PATH = settings["project"].get("dump_buffer", {}).get("path", "flowcept_buffer.jsonl")
 
 TELEMETRY_CAPTURE = settings.get("telemetry_capture", None)

flowcept/flowcept_api/flowcept_controller.py

@@ -320,7 +320,11 @@ class Flowcept(object):
             file_path = DUMP_BUFFER_PATH
         assert file_path is not None, "Please indicate file_path either in the argument or in the config file."
         if not os.path.exists(file_path):
-            raise FileNotFoundError(f"Flowcept buffer file '{file_path}' was not found.")
+            raise FileNotFoundError(
+                f"Flowcept buffer file '{file_path}' was not found. "
+                f"Check your settings to see if you're dumping the data to a file and check if you"
+                f"have started Flowcept."
+            )
 
         with open(file_path, "rb") as f:
             lines = [ln for ln in f.read().splitlines() if ln]

flowcept/flowceptor/adapters/mlflow/interception_event_handler.py

@@ -13,7 +13,38 @@ class InterceptionEventHandler(FileSystemEventHandler):
         self.callback_function = callback_function
         self.interceptor_instance = interceptor_instance
 
+    def _matches_watch_target(self, path):
+        """Return True when a path matches the watch target.
+
+        If the target is a directory, any event under that directory matches.
+        If the target is a file, only the exact file path matches.
+        """
+        if not path:
+            return False
+        target = Path(self.file_path_to_watch).resolve()
+        candidate = Path(path).resolve()
+        if target.is_dir():
+            try:
+                candidate.relative_to(target)
+                return True
+            except ValueError:
+                return False
+        return candidate == target
+
+    def _maybe_callback(self, event):
+        """Invoke the callback when an event matches the watch target."""
+        paths = [getattr(event, "src_path", None), getattr(event, "dest_path", None)]
+        if any(self._matches_watch_target(path) for path in paths):
+            self.callback_function(self.interceptor_instance)
+
     def on_modified(self, event):
         """Get on modified."""
-        if Path(event.src_path).resolve() == Path(self.file_path_to_watch).resolve():
-            self.callback_function(self.interceptor_instance)
+        self._maybe_callback(event)
+
+    def on_created(self, event):
+        """Get on created."""
+        self._maybe_callback(event)
+
+    def on_moved(self, event):
+        """Get on moved."""
+        self._maybe_callback(event)

flowcept/flowceptor/adapters/mlflow/mlflow_interceptor.py

@@ -50,9 +50,10 @@ class MLFlowInterceptor(BaseInterceptor):
         interesting change, it calls self.intercept; otherwise, let it
         go....
         """
+        intercepted = 0
         runs = self.dao.get_finished_run_uuids()
         if not runs:
-            return
+            return intercepted
         for run_uuid_tuple in runs:
             run_uuid = run_uuid_tuple[0]
             if not self.state_manager.has_element_id(run_uuid):
@@ -63,6 +64,8 @@ class MLFlowInterceptor(BaseInterceptor):
                     continue
                 task_msg = self.prepare_task_msg(run_data).to_dict()
                 self.intercept(task_msg)
+                intercepted += 1
+        return intercepted
 
     def start(self, bundle_exec_id, check_safe_stops) -> "MLFlowInterceptor":
         """Start it."""
@@ -74,10 +77,20 @@ class MLFlowInterceptor(BaseInterceptor):
     def stop(self, check_safe_stops: bool = True) -> bool:
         """Stop it."""
         sleep(1)
-        super().stop(check_safe_stops)
         self.logger.debug("Interceptor stopping...")
-        self._observer.stop()
-        self._observer_thread.join()
+        # Flush any late writes before stopping the observer.
+        try:
+            intercepted = self.callback()
+            if intercepted == 0:
+                sleep(self.settings.watch_interval_sec)
+                self.callback()
+        except Exception as e:
+            self.logger.exception(e)
+        super().stop(check_safe_stops)
+        if self._observer is not None:
+            self._observer.stop()
+            if self._observer_thread is not None:
+                self._observer_thread.join()
         self.logger.debug("Interceptor stopped.")
         return True
 
@@ -98,4 +111,5 @@ class MLFlowInterceptor(BaseInterceptor):
         watch_dir = os.path.dirname(self.settings.file_path) or "."
         self._observer.schedule(event_handler, watch_dir, recursive=True)
         self._observer.start()
+        sleep(0.2)
         self.logger.info(f"Watching directory {watch_dir} with file {self.settings.file_path} ")

flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py

@@ -122,4 +122,5 @@ class TensorboardInterceptor(BaseInterceptor):
 
         self._observer.schedule(event_handler, self.settings.file_path, recursive=True)
         self._observer.start()
+        sleep(0.2)
         self.logger.debug(f"Watching {self.settings.file_path}")

flowcept/flowceptor/consumers/agent/base_agent_context_manager.py

@@ -20,6 +20,12 @@ class BaseAppContext:
 
     tasks: List[Dict]
 
+    def reset_context(self):
+        """
+        Method to reset the variables in the context.
+        """
+        self.tasks = []
+
 
 class BaseAgentContextManager(BaseConsumer):
     """
@@ -39,14 +45,13 @@ class BaseAgentContextManager(BaseConsumer):
 
     agent_id = None
 
-    def __init__(self):
+    def __init__(self, allow_mq_disabled: bool = False):
         """
         Initializes the agent and resets its context state.
         """
         self._started = False
-        super().__init__()
-        self.context = None
-        self.reset_context()
+        super().__init__(allow_mq_disabled=allow_mq_disabled)
+        # self.context = BaseAppContext(tasks=[])
         self.agent_id = BaseAgentContextManager.agent_id
 
     def message_handler(self, msg_obj: Dict) -> bool:
@@ -77,12 +82,6 @@ class BaseAgentContextManager(BaseConsumer):
 
         return True
 
-    def reset_context(self):
-        """
-        Resets the internal context, clearing all stored task data.
-        """
-        self.context = BaseAppContext(tasks=[])
-
     @asynccontextmanager
     async def lifespan(self, app):
         """

flowcept/flowceptor/consumers/base_consumer.py

@@ -13,18 +13,28 @@ class BaseConsumer(object):
 
     This class provides a standard interface and shared logic for subscribing to
     message queues and dispatching messages to a handler.
+
+    Note
+    ----
+    The MQ-disabled path is only intended for agent consumers that can operate
+    from an offline buffer file. General consumers that require MQ should keep
+    the default behavior (raise when MQ_ENABLED is False).
     """
 
-    def __init__(self):
+    def __init__(self, allow_mq_disabled: bool = False):
         """Initialize the message queue DAO and logger."""
+        self.logger = FlowceptLogger()
+        self._main_thread: Optional[Thread] = None
+
         if not MQ_ENABLED:
+            if allow_mq_disabled:
+                self._mq_dao = None
+                self.logger.warning("MQ is disabled; starting consumer without a message queue.")
+                return
             raise Exception("MQ is disabled in the settings. You cannot consume messages.")
 
         self._mq_dao = MQDao.build()
 
-        self.logger = FlowceptLogger()
-        self._main_thread: Optional[Thread] = None
-
     @abstractmethod
     def message_handler(self, msg_obj: Dict) -> bool:
         """
@@ -62,6 +72,9 @@ class BaseConsumer(object):
         BaseConsumer
             The current instance (to allow chaining).
         """
+        if self._mq_dao is None:
+            self.logger.warning("MQ is disabled; skipping message consumption start.")
+            return self
         if target is None:
             target = self.default_thread_target
         self._mq_dao.subscribe()
@@ -85,6 +98,9 @@ class BaseConsumer(object):
         --------
         start : Starts the consumer and optionally spawns a background thread to run this method.
         """
+        if self._mq_dao is None:
+            self.logger.warning("MQ is disabled; no message listener will run.")
+            return
         self.logger.debug("Going to wait for new messages!")
         self._mq_dao.message_listener(self.message_handler)
         self.logger.debug("Broke main message listening loop!")
@@ -96,4 +112,6 @@ class BaseConsumer(object):
         """
         Stop consuming messages by unsubscribing from the message queue.
         """
+        if self._mq_dao is None:
+            return
         self._mq_dao.unsubscribe()

flowcept/flowceptor/consumers/document_inserter.py

@@ -197,6 +197,24 @@ class DocumentInserter(BaseConsumer):
                     f"{'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
                 )
             return "continue"
+        elif message["info"] == "mq_flush_complete":
+            exec_bundle_id = message.get("exec_bundle_id", None)
+            interceptor_instance_id = message.get("interceptor_instance_id")
+            self.logger.info(
+                f"DocInserter id {id(self)}. Received mq_flush_complete message "
+                f"from the interceptor {'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
+            )
+            if self.check_safe_stops:
+                self.logger.info(
+                    f"Begin register_flush_complete "
+                    f"{'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
+                )
+                self._mq_dao.register_flush_complete(interceptor_instance_id, exec_bundle_id)
+                self.logger.info(
+                    f"Done register_flush_complete "
+                    f"{'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
+                )
+            return "continue"
         elif message["info"] == "stop_document_inserter":
             exec_bundle_id = message.get("exec_bundle_id", None)
             if self._bundle_exec_id == exec_bundle_id:
@@ -297,7 +315,10 @@ class DocumentInserter(BaseConsumer):
             return self
         if self.check_safe_stops:
             trial = 0
-            while not self._mq_dao.all_time_based_threads_ended(bundle_exec_id):
+            while not (
+                self._mq_dao.all_time_based_threads_ended(bundle_exec_id)
+                and self._mq_dao.all_flush_complete_received(bundle_exec_id)
+            ):
                 self.logger.debug(
                     f"# time_based_threads for bundle_exec_id {bundle_exec_id} is"
                     f"{self._mq_dao._keyvalue_dao.set_count(bundle_exec_id)}"