vellum-workflow-server 1.8.2__py3-none-any.whl → 1.10.7__py3-none-any.whl

workflow_server/start.py

@@ -33,6 +33,7 @@ class CustomGunicornLogger(glogging.Logger):
         logger = logging.getLogger("gunicorn.access")
         logger.addFilter(HealthCheckFilter())
         logger.addFilter(SignalFilter())
+        logger.addFilter(StatusIsAvailableFilter())
 
 
 class HealthCheckFilter(logging.Filter):
@@ -45,6 +46,11 @@ class SignalFilter(logging.Filter):
         return "SIGTERM" not in record.getMessage()
 
 
+class StatusIsAvailableFilter(logging.Filter):
+    def filter(self, record: Any) -> bool:
+        return "/status/is_available" not in record.getMessage()
+
+
 def start() -> None:
     if not is_development():
         start_oom_killer_worker()
@@ -58,8 +64,8 @@ def start() -> None:
         "workers": int(os.getenv("GUNICORN_WORKERS", 2)),
         "threads": int(os.getenv("GUNICORN_THREADS", 9 if ENABLE_PROCESS_WRAPPER else 6)),
         # Aggressively try to avoid memory leaks when using non process mode
-        "max_requests": 120 if ENABLE_PROCESS_WRAPPER else 20,
-        "max_requests_jitter": 30 if ENABLE_PROCESS_WRAPPER else 10,
+        "max_requests": int(os.getenv("GUNICORN_MAX_REQUESTS", 120 if ENABLE_PROCESS_WRAPPER else 20)),
+        "max_requests_jitter": int(os.getenv("GUNICORN_MAX_REQUESTS_JITTER", 30 if ENABLE_PROCESS_WRAPPER else 10)),
         "worker_class": "gthread",
         "timeout": max_workflow_runtime_seconds,
         "logger_class": CustomGunicornLogger,

workflow_server/utils/exit_handler.py

@@ -1,15 +1,43 @@
+from datetime import datetime
 import logging
 import multiprocessing
 import signal
+from time import sleep
 from typing import Any
 
+from workflow_server.config import IS_ASYNC_MODE, is_development
+from workflow_server.utils.system_utils import get_active_process_count
+
 logger = logging.getLogger(__name__)
 process_killed_switch = multiprocessing.Event()
 
 
+def _wait_for_workers() -> None:
+    # Would be annoying to have this on for dev since would prevent reload restarts. Also disabling this
+    # for non async mode for now since it shouldn't be needed anyway cus we keep the requests open.
+    if is_development() and not IS_ASYNC_MODE:
+        return
+
+    start_time = datetime.now()
+    loops = 0
+
+    while get_active_process_count() > 0:
+        if loops % 30 == 0:
+            logger.info("Waiting for workflow processes to finish...")
+
+        # TODO needa pass in max workflow time here for VPC
+        if (datetime.now() - start_time).total_seconds() > 1800:
+            logger.warning("Max elapsed time waiting for workflow processes to complete exceeded, shutting down")
+            exit(1)
+
+        sleep(1)
+        loops += 1
+
+
 def gunicorn_exit_handler(_worker: Any) -> None:
+    logger.info("Received gunicorn kill signal")
     process_killed_switch.set()
-    logger.warning("Received gunicorn kill signal")
+    _wait_for_workers()
 
 
 def exit_handler(_signal: int, _frame: Any) -> None:
@@ -19,6 +47,7 @@ def exit_handler(_signal: int, _frame: Any) -> None:
     """
     process_killed_switch.set()
     logger.warning("Received kill signal")
+    _wait_for_workers()
     exit(1)