scalable-pypeline 2.1.13__tar.gz → 2.1.14__tar.gz

{scalable-pypeline-2.1.13/scalable_pypeline.egg-info → scalable-pypeline-2.1.14}/PKG-INFO

@@ -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

scalable-pypeline-2.1.14/pypeline/__init__.py

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

{scalable-pypeline-2.1.13 → scalable-pypeline-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)
 

{scalable-pypeline-2.1.13 → scalable-pypeline-2.1.14}/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)
 

scalable-pypeline-2.1.14/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}")

scalable-pypeline-2.1.14/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 → scalable-pypeline-2.1.14/scalable_pypeline.egg-info}/PKG-INFO

@@ -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

{scalable-pypeline-2.1.13 → scalable-pypeline-2.1.14}/scalable_pypeline.egg-info/SOURCES.txt

@@ -26,11 +26,13 @@ pypeline/pipelines/composition/parallel_pipeline_composition.py
 pypeline/pipelines/composition/pypeline_composition.py
 pypeline/pipelines/middleware/__init__.py
 pypeline/pipelines/middleware/get_active_worker_id_middleware.py
+pypeline/pipelines/middleware/graceful_shutdown_middleware.py
 pypeline/pipelines/middleware/parallel_pipeline_middleware.py
 pypeline/pipelines/middleware/pypeline_middleware.py
 pypeline/utils/__init__.py
 pypeline/utils/config_utils.py
 pypeline/utils/dramatiq_utils.py
+pypeline/utils/graceful_shutdown_util.py
 pypeline/utils/module_utils.py
 pypeline/utils/pipeline_utils.py
 pypeline/utils/schema_utils.py

{scalable-pypeline-2.1.13 → scalable-pypeline-2.1.14}/scalable_pypeline.egg-info/requires.txt

@@ -30,3 +30,4 @@ gevent<22,>=21.12.0
 networkx>=2.4
 dramatiq[rabbitmq]==1.17.0
 apscheduler<4,>=3.10.4
+tenacity==9.1.2

{scalable-pypeline-2.1.13 → scalable-pypeline-2.1.14}/setup.py

@@ -1,5 +1,4 @@
-""" PypeLine Library Setup
-"""
+"""PypeLine Library Setup"""
 
 import re
 import ast
@@ -183,6 +182,7 @@ setup(
             "networkx>=2.4",
             "dramatiq[rabbitmq]==1.17.0",
             "apscheduler>=3.10.4,<4",
+            "tenacity==9.1.2",
         ],
         "dev": ["black"],
         "test": [

scalable-pypeline-2.1.13/pypeline/__init__.py

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