scalable-pypeline 2.1.1__py2.py3-none-any.whl → 2.1.2__py2.py3-none-any.whl

pypeline/__init__.py

@@ -1 +1 @@
-__version__ = "2.1.1"
+__version__ = "2.1.2"

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)
 
 

pypeline/pipelines/middleware/deduplication_middleware.py

@@ -0,0 +1,91 @@
+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)
+            ):
+                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.dist-info → scalable_pypeline-2.1.2.dist-info}/METADATA

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

{scalable_pypeline-2.1.1.dist-info → scalable_pypeline-2.1.2.dist-info}/RECORD

@@ -1,7 +1,7 @@
-pypeline/__init__.py,sha256=zPJIgPGcoSNiD0qme18OnYJYE3A9VVytlhO-V5DaAW0,22
+pypeline/__init__.py,sha256=UiuBcRXPtXxPUBDdp0ZDvWl0U9Db1kMNfT3oAfhxqLg,22
 pypeline/barrier.py,sha256=oO964l9qOCOibweOHyNivmAvufdXOke9nz2tdgclouo,1172
 pypeline/constants.py,sha256=coiF8dMP25qIwoNYSnS7oy7hCd4-5yqPFmdPsN93Q1A,2892
-pypeline/dramatiq.py,sha256=LWsl0o0t5FdxewIl87ARZKrNK0ENoYJEJAEVDSNFa40,12272
+pypeline/dramatiq.py,sha256=TGwXWSInpCPFtYC5C-Omc2sEk5ecpsOG8xHH_mx9WTo,12451
 pypeline/extensions.py,sha256=BzOTnXhNxap3N7uIUUh_hO6dDwx08Vc_RJDE93_K0Lo,610
 pypeline/pipeline_config_schema.py,sha256=hK2_egtg-YFx_XJDs_NyrOTGKkel7W83X-G0sic52sM,10592
 pypeline/pipeline_settings_schema.py,sha256=84AuNFYsOUpoADsjEo_n9T6Ica-c21oK_V9s15I4lCg,20212
@@ -19,6 +19,7 @@ pypeline/pipelines/composition/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5
 pypeline/pipelines/composition/parallel_pipeline_composition.py,sha256=pTw9Xb9h4JnV4siFc3JStm5lB-i9djUADo3Kh5K3s7g,12976
 pypeline/pipelines/composition/pypeline_composition.py,sha256=ieTuQZ8zxTtvmPEkrWFbItjGtvO3JUotXcR-Jim2mss,7204
 pypeline/pipelines/middleware/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pypeline/pipelines/middleware/deduplication_middleware.py,sha256=IGO9Kc6NPMUaw1ytT3_ud2ITSZuh1-_WaU5_onazeh8,3556
 pypeline/pipelines/middleware/parallel_pipeline_middleware.py,sha256=kTp6niYoe2nXIiN6EGRfdpxrJyioo0GPxDkfefbGlEk,2821
 pypeline/pipelines/middleware/pypeline_middleware.py,sha256=kvt5A9OxDwpIo0PsH11Im62tH6VquUc6OFoZDw2Gxsk,8036
 pypeline/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -28,9 +29,9 @@ pypeline/utils/module_utils.py,sha256=-yEJIukDCoXnmlZVXB6Dww25tH6GdPE5SoFqv6pfdV
 pypeline/utils/pipeline_utils.py,sha256=kGP1QwCJikGC5QNRtzRXCDVewyRMpWIqERTNnxGLlSY,4795
 pypeline/utils/schema_utils.py,sha256=Fgl0y9Cuo_TZeEx_S3gaSVnLjn6467LTkjb2ek7Ms98,851
 tests/fixtures/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-scalable_pypeline-2.1.1.dist-info/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
-scalable_pypeline-2.1.1.dist-info/METADATA,sha256=wgJ50QCkHF7mvEqD1XN4X6x3XVbb7Ne2dso-NWe7Qsc,5926
-scalable_pypeline-2.1.1.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-scalable_pypeline-2.1.1.dist-info/entry_points.txt,sha256=uWs10ODfHSBKo2Cx_QaUjPHQTpZ3e77j9VlAdRRmMyg,119
-scalable_pypeline-2.1.1.dist-info/top_level.txt,sha256=C7dpkEOc_-nnsAQb28BfQknjD6XHRyS9ZrvVeoIbV7s,15
-scalable_pypeline-2.1.1.dist-info/RECORD,,
+scalable_pypeline-2.1.2.dist-info/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
+scalable_pypeline-2.1.2.dist-info/METADATA,sha256=M9cP1_2S3B289LGum8QtlruOBUfdnH7x4IDtcM3Swpk,5926
+scalable_pypeline-2.1.2.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+scalable_pypeline-2.1.2.dist-info/entry_points.txt,sha256=uWs10ODfHSBKo2Cx_QaUjPHQTpZ3e77j9VlAdRRmMyg,119
+scalable_pypeline-2.1.2.dist-info/top_level.txt,sha256=C7dpkEOc_-nnsAQb28BfQknjD6XHRyS9ZrvVeoIbV7s,15
+scalable_pypeline-2.1.2.dist-info/RECORD,,