scalable-pypeline 2.1.13__py2.py3-none-any.whl → 2.1.14__py2.py3-none-any.whl

pypeline/__init__.py

@@ -1 +1 @@
-__version__ = "2.1.13"
+__version__ = "2.1.14"

pypeline/dramatiq.py

@@ -2,6 +2,7 @@ import copy
 import typing
 import pika
 import logging
+import os
 
 import click
 from urllib.parse import urlparse
@@ -52,6 +53,7 @@ from pypeline.utils.dramatiq_utils import (
     register_lazy_actor,
     LazyActor,
 )
+from pypeline.utils.graceful_shutdown_util import enable_graceful_shutdown
 from pypeline.utils.module_utils import get_callable
 
 
@@ -92,6 +94,12 @@ def configure_default_broker(broker: Broker = None):
     broker.add_middleware(PypelineMiddleware(redis_url=REDIS_URL))
     broker.add_middleware(CurrentMessage())
     broker.add_middleware(GetActiveWorkerIdMiddleware())
+    if (
+        os.getenv("RESTRICT_WORKER_SHUTDOWN_WHILE_JOBS_RUNNING", "false").lower()
+        == "true"
+    ):
+        enable_graceful_shutdown(broker=broker, redis_url=REDIS_URL)
+
     register_actors_for_workers(broker)
     set_broker(broker)
 

pypeline/pipelines/factory.py

@@ -70,13 +70,16 @@ def dag_generator(
             )
             registered_actors[task] = lazy_actor
             if args and not kwargs:
-                message_group.append(registered_actors[task].message(*args))
+                msg = registered_actors[task].message(*args)
             elif kwargs and not args:
-                message_group.append(registered_actors[task].message(**kwargs))
+                msg = registered_actors[task].message(**kwargs)
             elif args and kwargs:
-                message_group.append(registered_actors[task].message(*args, **kwargs))
+                msg = registered_actors[task].message(*args, **kwargs)
             else:
-                message_group.append(registered_actors[task].message())
+                msg = registered_actors[task].message()
+            msg.options["task_ttl"] = pipeline_config["metadata"]["maxTtl"]
+            message_group.append(msg)
+
         messages.append(message_group)
     p = parallel_pipeline(messages)
 

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/utils/graceful_shutdown_util.py

@@ -0,0 +1,36 @@
+import threading
+import signal
+import os
+import redis
+import socket
+import sys
+import time
+import logging
+from pypeline.pipelines.middleware.graceful_shutdown_middleware import (
+    GraceFulShutdownMiddleware,
+)
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def enable_graceful_shutdown(broker, redis_url):
+    """Attach GracefulShutdownMiddleware and a SIGTERM handler to the current process."""
+    broker.add_middleware(GraceFulShutdownMiddleware(redis_url=redis_url))
+
+    if threading.current_thread().name == "MainThread":
+        key_prefix = "busy"
+        hostname = socket.gethostname()
+        pid = os.getpid()
+        busy_key = f"{key_prefix}:{hostname}-{pid}"
+        r = redis.Redis.from_url(redis_url)
+
+        def shutdown_handler(signum, frame):
+            logger.info(f"[Signal Handler] Received signal {signum}")
+            while r.get(busy_key):
+                logger.info(f"[Signal Handler] Busy ({busy_key}), waiting...")
+                time.sleep(30)
+            logger.info(f"[Signal Handler] Done. Exiting.")
+            sys.exit(0)
+
+        signal.signal(signal.SIGTERM, shutdown_handler)

{scalable_pypeline-2.1.13.dist-info → scalable_pypeline-2.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: scalable-pypeline
-Version: 2.1.13
+Version: 2.1.14
 Summary: PypeLine - Python pipelines for the Real World
 Home-page: https://gitlab.com/bravos2/pypeline
 Author: Bravos Power Corporation
@@ -33,6 +33,7 @@ Provides-Extra: workers
 Requires-Dist: networkx (>=2.4) ; extra == 'workers'
 Requires-Dist: dramatiq[rabbitmq] (==1.17.0) ; extra == 'workers'
 Requires-Dist: apscheduler (<4,>=3.10.4) ; extra == 'workers'
+Requires-Dist: tenacity (==9.1.2) ; extra == 'workers'
 
 ```
 ______ __   ________  _____  _     _____  _   _  _____ 

{scalable_pypeline-2.1.13.dist-info → scalable_pypeline-2.1.14.dist-info}/RECORD

@@ -1,7 +1,7 @@
-pypeline/__init__.py,sha256=59WYUQbpJnalweC5r75GQeXXOChhHcCgbKm0CW6XD1M,23
+pypeline/__init__.py,sha256=otJNk418s7IUoP3b0ZQy4k9fA8qhpkQusp3zKCHD2bQ,23
 pypeline/barrier.py,sha256=oO964l9qOCOibweOHyNivmAvufdXOke9nz2tdgclouo,1172
 pypeline/constants.py,sha256=EGSuLq4KhZ4bxrbtnUgKclELRyya5ipvv0WeybCzNAs,3049
-pypeline/dramatiq.py,sha256=nnNxm2akL6hjqfjwQglK76LspZrdrL_2E_rDN5ThwZ8,13432
+pypeline/dramatiq.py,sha256=D-E6_oJc9he-F2rvze-DCq4eBVY3Hq7V0pSC5crHrrU,13703
 pypeline/extensions.py,sha256=BzOTnXhNxap3N7uIUUh_hO6dDwx08Vc_RJDE93_K0Lo,610
 pypeline/pipeline_config_schema.py,sha256=kprvmfPfmNVP5MOqtWzDm4ON5NJWIKQC8GWbIy5EIuk,11183
 pypeline/pipeline_settings_schema.py,sha256=VgbMvTxQSqnvLX593RNoD9gTsZnJwzCc4wGpczV5CEI,20403
@@ -14,24 +14,26 @@ pypeline/flask/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuF
 pypeline/flask/api/pipelines.py,sha256=lw1ggsjp_Iha5MhyQGHtVW0akpVJnxIk0hn6NkC3c8s,9314
 pypeline/flask/api/schedules.py,sha256=8PKCMdPucaer8opchNlI5aDssK2UqT79hHpeg5BMtTA,1210
 pypeline/pipelines/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pypeline/pipelines/factory.py,sha256=p2b1XhhMYt_f5h2FAr6Zze1FfWz8SCyO9DEIS_EDDKE,3258
+pypeline/pipelines/factory.py,sha256=ftm6AwXIdkC8LEzy3c4CLPzXtoQonBKUo3yojbYTsUE,3309
 pypeline/pipelines/composition/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pypeline/pipelines/composition/parallel_pipeline_composition.py,sha256=pTw9Xb9h4JnV4siFc3JStm5lB-i9djUADo3Kh5K3s7g,12976
 pypeline/pipelines/composition/pypeline_composition.py,sha256=lY1RO7luLqDBLfoX99YbBatx1Qov1cxGoS5lRAjB_DQ,8534
 pypeline/pipelines/middleware/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pypeline/pipelines/middleware/get_active_worker_id_middleware.py,sha256=X4ZfRk3L8MD00DTsGHth7oOdy-W7LQV96T8vu5UC42A,755
+pypeline/pipelines/middleware/graceful_shutdown_middleware.py,sha256=k37zmFk9dOye05BoQP7KcB9MEQgvodI16kOJyYhRyAc,1764
 pypeline/pipelines/middleware/parallel_pipeline_middleware.py,sha256=kTp6niYoe2nXIiN6EGRfdpxrJyioo0GPxDkfefbGlEk,2821
 pypeline/pipelines/middleware/pypeline_middleware.py,sha256=ciXZD2rLMCLGI9wP5zB8g_5_zazt08LetQXuTJkedms,9147
 pypeline/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pypeline/utils/config_utils.py,sha256=rAIATyoW7kGETZ_Z2DqiXtGd7bJp5uPfcLtfNPOYsNs,2167
 pypeline/utils/dramatiq_utils.py,sha256=DUdgVywm1182A4i69XzH9EIh1EJ9zAHmJLtOaVSW7pw,3844
+pypeline/utils/graceful_shutdown_util.py,sha256=d3p8ZQ9-xasd_H0YUA5ZsjCQ9PVliSO7_4c1OYzhx-Q,1156
 pypeline/utils/module_utils.py,sha256=-yEJIukDCoXnmlZVXB6Dww25tH6GdPE5SoFqv6pfdVU,3682
 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.13.dist-info/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
-scalable_pypeline-2.1.13.dist-info/METADATA,sha256=qFVrWAUfDcyDVSxbUmeuc9SLjpyCRKUNr5twJRuuUXM,5927
-scalable_pypeline-2.1.13.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-scalable_pypeline-2.1.13.dist-info/entry_points.txt,sha256=uWs10ODfHSBKo2Cx_QaUjPHQTpZ3e77j9VlAdRRmMyg,119
-scalable_pypeline-2.1.13.dist-info/top_level.txt,sha256=C7dpkEOc_-nnsAQb28BfQknjD6XHRyS9ZrvVeoIbV7s,15
-scalable_pypeline-2.1.13.dist-info/RECORD,,
+scalable_pypeline-2.1.14.dist-info/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
+scalable_pypeline-2.1.14.dist-info/METADATA,sha256=6MPVldbryAie6UIDSDF3-9EJffhMyEqsiehraoaC2ZA,5982
+scalable_pypeline-2.1.14.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+scalable_pypeline-2.1.14.dist-info/entry_points.txt,sha256=uWs10ODfHSBKo2Cx_QaUjPHQTpZ3e77j9VlAdRRmMyg,119
+scalable_pypeline-2.1.14.dist-info/top_level.txt,sha256=C7dpkEOc_-nnsAQb28BfQknjD6XHRyS9ZrvVeoIbV7s,15
+scalable_pypeline-2.1.14.dist-info/RECORD,,