pyworkflow-engine 0.1.11__py3-none-any.whl → 0.1.12__py3-none-any.whl

pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.1.11"
+__version__ = "0.1.12"
 
 # Configuration
 from pyworkflow.config import (

pyworkflow/celery/app.py

@@ -15,10 +15,28 @@ garbage collector and Celery's saferepr module. It does not affect functionality
 import os
 
 from celery import Celery
+from celery.signals import worker_init, worker_process_init, worker_shutdown
 from kombu import Exchange, Queue
 
 from pyworkflow.observability.logging import configure_logging
 
+# Track if logging has been configured in this process
+_logging_configured = False
+
+
+def _configure_worker_logging() -> None:
+    """Configure logging for the current worker process."""
+    global _logging_configured
+    if not _logging_configured:
+        from loguru import logger as loguru_logger
+
+        # Enable pyworkflow logging (may have been disabled by CLI)
+        loguru_logger.enable("pyworkflow")
+
+        log_level = os.getenv("PYWORKFLOW_LOG_LEVEL", "INFO").upper()
+        configure_logging(level=log_level)
+        _logging_configured = True
+
 
 def discover_workflows(modules: list[str] | None = None) -> None:
     """
@@ -118,6 +136,14 @@ def create_celery_app(
         accept_content=["json"],
         timezone="UTC",
         enable_utc=True,
+        # Broker transport options - prevent task redelivery
+        # See: https://github.com/celery/celery/issues/5935
+        broker_transport_options={
+            "visibility_timeout": 3600,  # 12 hours - prevent Redis from re-queueing tasks
+        },
+        result_backend_transport_options={
+            "visibility_timeout": 3600,
+        },
         # Task routing
         task_default_queue="pyworkflow.default",
         task_default_exchange="pyworkflow",
@@ -154,7 +180,7 @@ def create_celery_app(
         task_reject_on_worker_lost=True,
         worker_prefetch_multiplier=1,  # Fair task distribution
         # Retry settings
-        task_autoretry_for=(Exception,),
+        task_autoretry_for=(),
         task_retry_backoff=True,
         task_retry_backoff_max=600,  # 10 minutes max
         task_retry_jitter=True,
@@ -168,8 +194,9 @@ def create_celery_app(
         worker_task_log_format="[%(asctime)s: %(levelname)s/%(processName)s] [%(task_name)s(%(task_id)s)] %(message)s",
     )
 
-    # Configure logging
-    configure_logging(level="INFO")
+    # Note: Logging is configured via Celery signals (worker_init, worker_process_init)
+    # to ensure proper initialization AFTER process forking.
+    # See on_worker_init() and on_worker_process_init() below.
 
     # Auto-discover workflows from environment variable or configured modules
     discover_workflows()
@@ -182,6 +209,63 @@ def create_celery_app(
 celery_app = create_celery_app()
 
 
+# ========== Celery Worker Signals ==========
+# These signals ensure proper initialization in forked worker processes
+
+
+@worker_init.connect
+def on_worker_init(**kwargs):
+    """
+    Called when the main worker process starts (before forking).
+
+    For prefork pool, this runs in the parent process.
+    For solo/threads pool, this is the main initialization point.
+    """
+    _configure_worker_logging()
+
+
+@worker_process_init.connect
+def on_worker_process_init(**kwargs):
+    """
+    Called when a worker child process is initialized (after forking).
+
+    This is critical for prefork pool:
+    - loguru's background thread doesn't survive fork()
+    - We need a persistent event loop for connection pool reuse
+    """
+    _configure_worker_logging()
+
+    # Initialize persistent event loop for this worker
+    from pyworkflow.celery.loop import init_worker_loop
+
+    init_worker_loop()
+
+
+@worker_shutdown.connect
+def on_worker_shutdown(**kwargs):
+    """
+    Called when the worker is shutting down.
+
+    Cleans up:
+    - Storage backend connections (PostgreSQL connection pools, etc.)
+    - The persistent event loop
+    """
+    from loguru import logger
+
+    from pyworkflow.celery.loop import close_worker_loop, run_async
+    from pyworkflow.storage.config import disconnect_all_cached
+
+    try:
+        # Clean up storage connections using the persistent loop
+        run_async(disconnect_all_cached())
+    except Exception as e:
+        # Log but don't fail shutdown
+        logger.warning(f"Error during storage cleanup on shutdown: {e}")
+    finally:
+        # Close the persistent event loop
+        close_worker_loop()
+
+
 def get_celery_app() -> Celery:
     """
     Get the global Celery application instance.

pyworkflow/celery/loop.py

@@ -0,0 +1,108 @@
+"""
+Persistent event loop management for Celery workers.
+
+This module provides a single, persistent event loop per worker process.
+Using a persistent loop allows asyncpg connection pools to be reused across
+tasks, avoiding the overhead of creating/destroying pools for each task.
+
+Usage:
+    from pyworkflow.celery.loop import run_async
+
+    # Instead of: result = asyncio.run(some_coroutine())
+    # Use:        result = run_async(some_coroutine())
+"""
+
+import asyncio
+import threading
+from collections.abc import Coroutine
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+# Per-worker persistent event loop
+# Created in worker_process_init, closed in worker_shutdown
+_worker_loop: asyncio.AbstractEventLoop | None = None
+_loop_lock = threading.Lock()
+
+
+def init_worker_loop() -> None:
+    """
+    Initialize the persistent event loop for this worker process.
+
+    Called from worker_process_init signal handler.
+    """
+    global _worker_loop
+
+    with _loop_lock:
+        if _worker_loop is None or _worker_loop.is_closed():
+            _worker_loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(_worker_loop)
+
+
+def close_worker_loop() -> None:
+    """
+    Close the persistent event loop for this worker process.
+
+    Called from worker_shutdown signal handler.
+    """
+    global _worker_loop
+
+    with _loop_lock:
+        if _worker_loop is not None and not _worker_loop.is_closed():
+            try:
+                # Run any pending cleanup
+                _worker_loop.run_until_complete(_worker_loop.shutdown_asyncgens())
+            except Exception:
+                pass
+            finally:
+                _worker_loop.close()
+                _worker_loop = None
+
+
+def get_worker_loop() -> asyncio.AbstractEventLoop:
+    """
+    Get the persistent event loop for this worker process.
+
+    If no loop exists (e.g., running outside Celery worker), creates one.
+
+    Returns:
+        The worker's event loop
+    """
+    global _worker_loop
+
+    with _loop_lock:
+        if _worker_loop is None or _worker_loop.is_closed():
+            # Not in a Celery worker or loop was closed - create a new one
+            _worker_loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(_worker_loop)
+        return _worker_loop
+
+
+def run_async(coro: Coroutine[Any, Any, T]) -> T:
+    """
+    Run a coroutine on the persistent worker event loop.
+
+    This is a drop-in replacement for asyncio.run() that reuses
+    the same event loop across tasks, allowing connection pools
+    to be shared.
+
+    Args:
+        coro: The coroutine to run
+
+    Returns:
+        The result of the coroutine
+
+    Example:
+        # Instead of:
+        result = asyncio.run(storage.get_run(run_id))
+
+        # Use:
+        result = run_async(storage.get_run(run_id))
+    """
+    loop = get_worker_loop()
+    return loop.run_until_complete(coro)
+
+
+def is_loop_running() -> bool:
+    """Check if the worker loop exists and is not closed."""
+    return _worker_loop is not None and not _worker_loop.is_closed()