scalable-pypeline 2.1.31__py2.py3-none-any.whl

pypeline/pipelines/middleware/get_active_worker_id_middleware.py

@@ -0,0 +1,22 @@
+import contextvars
+from dramatiq import Middleware
+
+
+class GetActiveWorkerIdMiddleware(Middleware):
+    _ACTIVE_WORKER_ID: contextvars.ContextVar["Optional[Message[Any]]"] = (
+        contextvars.ContextVar("_ACTIVE_WORKER_ID", default=None)
+    )
+
+    @classmethod
+    def get_active_worker_id(cls):
+        return cls._ACTIVE_WORKER_ID.get()
+
+    def before_process_message(self, broker, message):
+        """Sets the active worker ID, verifying broker has a broker_id first."""
+        broker_id = getattr(broker, "broker_id", None)
+        if not broker_id:
+            return
+        self._ACTIVE_WORKER_ID.set(broker_id)
+
+    def after_process_message(self, broker, message, *, result=None, exception=None):
+        self._ACTIVE_WORKER_ID.set(None)

pypeline/pipelines/middleware/graceful_shutdown_middleware.py

@@ -0,0 +1,50 @@
+import os
+import socket
+import logging
+import redis
+
+from dramatiq.middleware import Middleware
+from tenacity import retry, stop_after_attempt, wait_exponential, after_log
+
+logger = logging.getLogger(__name__)
+
+
+class GraceFulShutdownMiddleware(Middleware):
+    def __init__(self, redis_url, key_prefix="busy"):
+        self.redis = redis.Redis.from_url(redis_url)
+        self.hostname = socket.gethostname()
+        self.pid = os.getpid()
+        self.key_prefix = key_prefix
+        self.key = f"{self.key_prefix}:{self.hostname}-{self.pid}"
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=2, min=2, max=10),
+        after=after_log(logger, logging.WARNING),
+        reraise=True,
+    )
+    def _set_busy_flag(self, message_ttl):
+        self.redis.set(self.key, "1", ex=message_ttl)
+        logger.debug(f"[GracefulShutdownMiddleware] Set busy flag: {self.key}")
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=2, min=2, max=10),
+        after=after_log(logger, logging.WARNING),
+        reraise=True,
+    )
+    def _clear_busy_flag(self):
+        self.redis.delete(self.key)
+        logger.debug(f"[GracefulShutdownMiddleware] Cleared busy flag: {self.key}")
+
+    def before_process_message(self, broker, message):
+        try:
+            self._set_busy_flag(message_ttl=message.options["task_ttl"])
+        except Exception as e:
+            logger.error(f"[GracefulShutdownMiddleware] Failed to set busy flag: {e}")
+
+    def after_process_message(self, broker, message, *, result=None, exception=None):
+        try:
+            self._clear_busy_flag()
+        except Exception as e:
+            logger.error(f"[GracefulShutdownMiddleware] Failed to clear busy flag: {e}")

pypeline/pipelines/middleware/parallel_pipeline_middleware.py

@@ -0,0 +1,60 @@
+import copy
+from dramatiq.middleware import Middleware
+
+from pypeline.barrier import LockingParallelBarrier
+from pypeline.constants import PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+
+
+class ParallelPipeline(Middleware):
+    def __init__(self, redis_url):
+        self.redis_url = redis_url
+
+    def after_process_message(self, broker, message, *, result=None, exception=None):
+        from dramatiq.message import Message
+
+        if exception is None:
+            group_completion_uuid = message.options.get("group_completion_uuid")
+            if group_completion_uuid:
+                locking_parallel_barrier = LockingParallelBarrier(
+                    self.redis_url,
+                    task_key=group_completion_uuid,
+                    lock_key=f"{group_completion_uuid}-lock",
+                )
+                try:
+                    locking_parallel_barrier.acquire_lock(
+                        timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+                    )
+                    remaining_tasks = locking_parallel_barrier.decrement_task_count()
+                finally:
+                    locking_parallel_barrier.release_lock()
+                if remaining_tasks <= 0:
+                    execution_graph = message.options.get("execution_graph")
+
+                    for i in range(len(execution_graph)):
+                        message_group = execution_graph[i]
+
+                        # Check if the current group matches the group_completion_uuid
+                        if (
+                            message_group[0]["options"]["group_completion_uuid"]
+                            == group_completion_uuid
+                        ):
+                            # Check if there is a next group
+                            if i + 1 < len(execution_graph):
+                                next_group = execution_graph[i + 1]
+
+                                completion_uuid = next_group[0]["options"][
+                                    "group_completion_uuid"
+                                ]
+                                locking_parallel_barrier = LockingParallelBarrier(
+                                    self.redis_url,
+                                    task_key=completion_uuid,
+                                    lock_key=f"{completion_uuid}-lock",
+                                )
+                                locking_parallel_barrier.set_task_count(len(next_group))
+                                execution_graph_copy = copy.deepcopy(execution_graph)
+
+                                for next_message in next_group:
+                                    next_message["options"][
+                                        "execution_graph"
+                                    ] = execution_graph_copy
+                                    broker.enqueue(Message(**next_message))

pypeline/pipelines/middleware/pypeline_middleware.py

@@ -0,0 +1,202 @@
+from copy import copy
+
+import networkx as nx
+from dramatiq import Middleware
+
+from pypeline.barrier import LockingParallelBarrier
+from pypeline.constants import PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+from pypeline.utils.module_utils import get_callable
+from pypeline.utils.pipeline_utils import get_execution_graph
+from pypeline.utils.dramatiq_utils import register_lazy_actor
+
+
+class PypelineMiddleware(Middleware):
+
+    def __init__(self, redis_url):
+        self.redis_url = redis_url
+
+    def after_process_message(self, broker, message, *, result=None, exception=None):
+
+        if exception is not None:
+            return
+
+        if "pipeline" not in message.options:
+            return
+
+        pipeline = message.options["pipeline"]
+        max_retries = message.options.get("max_retries", None)
+        pipeline_config = pipeline["config"]
+        task_replacements = message.options["task_replacements"]
+        execution_id = message.options["execution_id"]
+        task_definitions = pipeline_config["taskDefinitions"]
+        task_name = message.options["task_name"]
+        task_key = f"{execution_id}-{task_name}"
+
+        # Signal to other jobs that current task is finished
+        locking_parallel_barrier = LockingParallelBarrier(
+            self.redis_url,
+            task_key=task_key,
+            lock_key=f"{message.options['base_case_execution_id']}-lock",
+        )
+        try:
+            locking_parallel_barrier.acquire_lock(timeout=10)
+            _ = locking_parallel_barrier.decrement_task_count()
+        finally:
+            locking_parallel_barrier.release_lock()
+
+        graph = get_execution_graph(pipeline_config)
+        children_tasks = pipeline_config["dagAdjacency"].get(task_name, [])
+        messages = []
+        for child in children_tasks:
+            child_ancestors = sorted(graph.predecessors(child))
+            child_ancestors_complete = {a: False for a in child_ancestors}
+
+            for scenario in message.options["scenarios"]:
+                if scenario["execution_id"] == execution_id:
+                    tasks_to_run_in_scenario = scenario["tasksToRunInScenario"]
+
+            for ancestor in child_ancestors:
+                if ancestor in tasks_to_run_in_scenario:
+                    current_scenario_ancestor_task_key = f"{execution_id}-{ancestor}"
+                    locking_parallel_barrier = LockingParallelBarrier(
+                        self.redis_url,
+                        task_key=current_scenario_ancestor_task_key,
+                        lock_key=f"{message.options['base_case_execution_id']}-lock",
+                    )
+                    try:
+                        locking_parallel_barrier.acquire_lock(
+                            timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+                        )
+                        if not locking_parallel_barrier.task_exists():
+                            child_ancestors_complete[ancestor] = False
+                        elif locking_parallel_barrier.get_task_count() <= 0:
+                            child_ancestors_complete[ancestor] = True
+                    finally:
+                        locking_parallel_barrier.release_lock()
+                else:
+                    base_scenario_ancestor_task_key = (
+                        f"{message.options['base_case_execution_id']}-{ancestor}"
+                    )
+                    locking_parallel_barrier = LockingParallelBarrier(
+                        self.redis_url,
+                        task_key=base_scenario_ancestor_task_key,
+                        lock_key=f"{message.options['base_case_execution_id']}-lock",
+                    )
+                    try:
+                        locking_parallel_barrier.acquire_lock(
+                            timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
+                        )
+                        if not locking_parallel_barrier.task_exists():
+                            child_ancestors_complete[ancestor] = False
+                        elif locking_parallel_barrier.get_task_count() <= 0:
+                            child_ancestors_complete[ancestor] = True
+                    finally:
+                        locking_parallel_barrier.release_lock()
+
+            if any(complete is False for complete in child_ancestors_complete.values()):
+                continue
+
+            if (
+                message.options["base_case_execution_id"]
+                == message.options["execution_id"]
+            ):
+                for scenario in message.options["scenarios"]:
+                    child_ancestors = list(graph.predecessors(child))
+                    child_has_other_ancestors_in_scenario = False
+
+                    for ancestor in child_ancestors:
+                        if ancestor in scenario["tasksToRunInScenario"]:
+                            child_has_other_ancestors_in_scenario = True
+                            break
+
+                    if (
+                        child in scenario["tasksToRunInScenario"]
+                        and task_name in child_ancestors
+                        and task_name not in scenario["tasksToRunInScenario"]
+                        and not child_has_other_ancestors_in_scenario
+                    ):
+                        task_key = f"{scenario['execution_id']}-{child}"
+                        locking_parallel_barrier = LockingParallelBarrier(
+                            self.redis_url,
+                            task_key=task_key,
+                            lock_key=f"{message.options['base_case_execution_id']}-lock",
+                        )
+                        locking_parallel_barrier.set_task_count(1)
+                        handler = task_definitions[child]["handlers"][
+                            task_replacements.get(child, 0)
+                        ]
+                        server_type = task_definitions[child].get("serverType", None)
+
+                        lazy_actor = register_lazy_actor(
+                            broker,
+                            get_callable(handler),
+                            pipeline_config["metadata"],
+                            server_type,
+                        )
+                        scenario_message = lazy_actor.message()
+                        scenario_message.options["pipeline"] = pipeline
+                        if max_retries is not None:
+                            scenario_message.options["max_retries"] = max_retries
+                        scenario_message.options["task_replacements"] = (
+                            task_replacements
+                        )
+                        scenario_message.options["execution_id"] = scenario[
+                            "execution_id"
+                        ]
+
+                        scenario_message.options["task_name"] = child
+                        scenario_message.options["base_case_execution_id"] = (
+                            message.options["base_case_execution_id"]
+                        )
+                        scenario_message.options["scenarios"] = message.options[
+                            "scenarios"
+                        ]
+                        if "settings" in message.kwargs:
+                            scenario_message.kwargs["settings"] = copy(
+                                message.kwargs["settings"]
+                            )
+                            scenario_message.kwargs["settings"]["execution_id"] = (
+                                scenario["execution_id"]
+                            )
+                        messages.append(scenario_message)
+
+            # If we've made it here all ancestors of this child are complete, and it's time to run.
+            task_key = f"{execution_id}-{child}"
+            locking_parallel_barrier = LockingParallelBarrier(
+                self.redis_url,
+                task_key=task_key,
+                lock_key=f"{message.options['base_case_execution_id']}-lock",
+            )
+            locking_parallel_barrier.set_task_count(1)
+            handler = task_definitions[child]["handlers"][
+                task_replacements.get(child, 0)
+            ]
+            server_type = task_definitions[child].get("serverType", None)
+            lazy_actor = register_lazy_actor(
+                broker,
+                get_callable(handler),
+                pipeline_config["metadata"],
+                server_type,
+            )
+
+            child_message = lazy_actor.message()
+            child_message.options["pipeline"] = pipeline
+            if max_retries is not None:
+                child_message.options["max_retries"] = max_retries
+            child_message.options["task_replacements"] = task_replacements
+            child_message.options["execution_id"] = execution_id
+            child_message.options["task_name"] = child
+            child_message.options["base_case_execution_id"] = message.options[
+                "base_case_execution_id"
+            ]
+            child_message.options["scenarios"] = message.options["scenarios"]
+            if "settings" in message.kwargs:
+                child_message.kwargs["settings"] = message.kwargs["settings"]
+                child_message.kwargs["settings"]["execution_id"] = message.options[
+                    "execution_id"
+                ]
+
+            messages.append(child_message)
+
+        for new_message in messages:
+            broker.enqueue(new_message)