scalable-pypeline 2.1.1__tar.gz → 2.1.3__tar.gz

{scalable-pypeline-2.1.1/scalable_pypeline.egg-info → scalable-pypeline-2.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: scalable-pypeline
-Version: 2.1.1
+Version: 2.1.3
 Summary: PypeLine - Python pipelines for the Real World
 Home-page: https://gitlab.com/bravos2/pypeline
 Author: Bravos Power Corporation

scalable-pypeline-2.1.3/pypeline/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.1.3"

{scalable-pypeline-2.1.1 → scalable-pypeline-2.1.3}/pypeline/dramatiq.py

@@ -34,6 +34,9 @@ from pypeline.constants import (
 )
 from pypeline.pipelines.middleware.parallel_pipeline_middleware import ParallelPipeline
 from pypeline.pipelines.middleware.pypeline_middleware import PypelineMiddleware
+from pypeline.pipelines.middleware.deduplication_middleware import (
+    DeduplicationMiddleware,
+)
 from pypeline.utils.config_utils import (
     retrieve_latest_schedule_config,
     get_service_config_for_worker,
@@ -72,6 +75,7 @@ def configure_default_broker(broker: Broker = None):
     rabbit_broker.add_middleware(PypelineMiddleware(redis_url=REDIS_URL))
     rabbit_broker.add_middleware(CurrentMessage())
     register_actors_for_workers(rabbit_broker)
+    rabbit_broker.add_middleware(DeduplicationMiddleware(redis_url=REDIS_URL))
     set_broker(rabbit_broker)
 
 

scalable-pypeline-2.1.3/pypeline/pipelines/middleware/deduplication_middleware.py

@@ -0,0 +1,94 @@
+import dramatiq
+import signal
+from dramatiq.middleware import Middleware
+from pypeline.barrier import LockingParallelBarrier
+from pypeline.constants import DEFAULT_TASK_TTL
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DeduplicationMiddleware(Middleware):
+    def __init__(self, redis_url="redis://localhost:6379/0"):
+        self.redis_url = redis_url
+        self.active_locks = {}
+
+    def before_process_message(self, broker, message):
+        task_id = message.message_id
+        task_key = f"dramatiq:task_counter:{task_id}"
+        lock_key = f"dramatiq:lock:{task_id}"
+        try:
+            # Try to acquire a lock for the task
+            locking_parallel_barrier = LockingParallelBarrier(
+                self.redis_url,
+                task_key=task_key,
+                lock_key=lock_key,
+            )
+            if (
+                locking_parallel_barrier.get_task_count() > 0
+                or not locking_parallel_barrier.acquire_lock(timeout=DEFAULT_TASK_TTL)
+            ):
+                logger.info(f"Found duplicate task {task_id}.  Skipping...")
+                raise dramatiq.middleware.SkipMessage(
+                    f"Task {task_id} is already being processed."
+                )
+
+            locking_parallel_barrier.set_task_count(1)
+            # Store the lock reference in the message and track it globally
+            message.options["dedupe_task_key"] = task_key
+            message.options["dedupe_lock_key"] = lock_key
+            self.active_locks[lock_key] = locking_parallel_barrier
+        except dramatiq.middleware.SkipMessage:
+            raise dramatiq.middleware.SkipMessage(
+                f"Task {task_id} is already being processed."
+            )
+        except Exception as e:
+            logger.exception(e)
+            raise e
+
+    def after_process_message(self, broker, message, *, result=None, exception=None):
+        """Releases lock for the message that just finished."""
+        dedupe_task_key = message.options.get("dedupe_task_key", None)
+        dedupe_lock_key = message.options.get("dedupe_lock_key", None)
+        if not dedupe_lock_key or not dedupe_task_key:
+            logger.warning(
+                "unexpected in after_process_message: dedupe task or lock key not in message"
+            )
+            return
+        if dedupe_lock_key in self.active_locks:
+            try:
+                lock = self.active_locks[dedupe_lock_key]
+                lock.decrement_task_count()
+                lock.release_lock()
+                del self.active_locks[dedupe_lock_key]
+            except Exception as e:
+                logger.info(
+                    f"Exception while trying to release lock {dedupe_lock_key}: {e}"
+                )
+                raise e
+        else:
+            lock = LockingParallelBarrier(
+                self.redis_url,
+                task_key=dedupe_task_key,
+                lock_key=dedupe_lock_key,
+            )
+            lock.decrement_task_count()
+            lock.release_lock()
+
+    def before_worker_shutdown(self, *args):
+        self.release_all_locks()
+
+    def before_worker_thread_shutdown(self, *args):
+        self.release_all_locks()
+
+    def release_all_locks(self, *args):
+        """Release all locks when the worker shuts down."""
+        for lock_key, lock in self.active_locks.items():
+            try:
+                lock.decrement_task_count()
+                lock.release_lock()
+            except Exception as e:
+                logger.info(f"Exception while trying to release lock {lock_key}: {e}")
+                raise e
+        self.active_locks.clear()

{scalable-pypeline-2.1.1 → scalable-pypeline-2.1.3/scalable_pypeline.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: scalable-pypeline
-Version: 2.1.1
+Version: 2.1.3
 Summary: PypeLine - Python pipelines for the Real World
 Home-page: https://gitlab.com/bravos2/pypeline
 Author: Bravos Power Corporation

{scalable-pypeline-2.1.1 → scalable-pypeline-2.1.3}/scalable_pypeline.egg-info/SOURCES.txt

@@ -25,6 +25,7 @@ pypeline/pipelines/composition/__init__.py
 pypeline/pipelines/composition/parallel_pipeline_composition.py
 pypeline/pipelines/composition/pypeline_composition.py
 pypeline/pipelines/middleware/__init__.py
+pypeline/pipelines/middleware/deduplication_middleware.py
 pypeline/pipelines/middleware/parallel_pipeline_middleware.py
 pypeline/pipelines/middleware/pypeline_middleware.py
 pypeline/utils/__init__.py

scalable-pypeline-2.1.1/pypeline/__init__.py

@@ -1 +0,0 @@
-__version__ = "2.1.1"