pyworkflow-engine 0.1.11__tar.gz → 0.1.13__tar.gz

{pyworkflow_engine-0.1.11 → pyworkflow_engine-0.1.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.1.11
+Version: 0.1.13
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT

{pyworkflow_engine-0.1.11 → pyworkflow_engine-0.1.13}/docs/concepts/step-context.mdx

@@ -129,7 +129,7 @@ async def my_workflow(user_id: str):
 ```
 
 <Note>
-  `set_step_context()` is an async function because it persists the context to storage and records a `CONTEXT_UPDATED` event.
+  `set_step_context()` is an async function because it persists the context to storage.
 </Note>
 
 ### Reading Context (Workflow and Steps)
@@ -201,14 +201,12 @@ async def process_order(order_id: str):
 
 ## Context Persistence and Replay
 
-Step Context is event-sourced for durability:
+Step Context is persisted for durability:
 
-1. **Persistence**: When you call `set_step_context()`, the context is:
-   - Stored in the `WorkflowRun.context` field
-   - Recorded as a `CONTEXT_UPDATED` event
+1. **Persistence**: When you call `set_step_context()`, the context is stored in the `WorkflowRun.context` field.
 
 2. **Replay**: When a workflow resumes after suspension:
-   - Context is restored from the event log
+   - Context is restored from `WorkflowRun.context`
    - Steps receive the same context they had during original execution
 
 ```python

{pyworkflow_engine-0.1.11 → pyworkflow_engine-0.1.13}/examples/local/durable/13_step_context.py

@@ -20,7 +20,6 @@ from pyworkflow import (
     StepContext,
     configure,
     get_step_context,
-    get_workflow_events,
     get_workflow_run,
     has_step_context,
     reset_config,
@@ -190,11 +189,6 @@ async def main():
     # Check stored context
     print(f"Stored context: {run.context}")
 
-    # Inspect events for CONTEXT_UPDATED
-    events = await get_workflow_events(run_id)
-    context_events = [e for e in events if e.type.value == "context.updated"]
-    print(f"Context update events: {len(context_events)}")
-
     # Example 2: Context with sleep (persistence test)
     print("\n--- Example 2: Context Persistence Across Sleep ---")
     run_id_2 = await start(

{pyworkflow_engine-0.1.11 → pyworkflow_engine-0.1.13}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.1.11"
+version = "0.1.13"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"

{pyworkflow_engine-0.1.11 → pyworkflow_engine-0.1.13}/pyworkflow/__init__.py

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

{pyworkflow_engine-0.1.11 → pyworkflow_engine-0.1.13}/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,19 @@ 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")
+    # Configure singleton locking for Redis brokers
+    # This enables distributed locking to prevent duplicate task execution
+    is_redis_broker = broker_url.startswith("redis://") or broker_url.startswith("rediss://")
+    if is_redis_broker:
+        app.conf.update(
+            singleton_backend_url=broker_url,
+            singleton_key_prefix="pyworkflow:lock:",
+            singleton_lock_expiry=3600,  # 1 hour TTL (safety net)
+        )
+
+    # 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 +219,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_engine-0.1.13/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()

pyworkflow_engine-0.1.13/pyworkflow/celery/singleton.py

@@ -0,0 +1,368 @@
+"""
+Singleton task implementation for PyWorkflow.
+
+Provides Redis-based distributed locking to prevent duplicate task execution.
+Self-contained implementation (no external dependencies beyond redis).
+
+Based on:
+- steinitzu/celery-singleton library concepts
+- FlowHunt's battle-tested refinements for retry-safe lock management
+"""
+
+import inspect
+import json
+from hashlib import md5
+from typing import Any
+from uuid import uuid4
+
+from celery import Task
+from celery.exceptions import WorkerLostError
+from loguru import logger
+
+
+def generate_lock_key(
+    task_name: str,
+    task_args: list[Any] | tuple[Any, ...] | None = None,
+    task_kwargs: dict[str, Any] | None = None,
+    key_prefix: str = "pyworkflow:lock:",
+) -> str:
+    """
+    Generate a unique lock key for a task based on its name and arguments.
+
+    Uses MD5 hash to keep key length reasonable while ensuring uniqueness.
+    """
+    str_args = json.dumps(task_args or [], sort_keys=True, default=str)
+    str_kwargs = json.dumps(task_kwargs or {}, sort_keys=True, default=str)
+    task_hash = md5((task_name + str_args + str_kwargs).encode()).hexdigest()
+    return key_prefix + task_hash
+
+
+class SingletonConfig:
+    """Configuration for singleton task behavior."""
+
+    def __init__(self, app: Any):
+        self.app = app
+
+    @property
+    def backend_url(self) -> str | None:
+        return self.app.conf.get("singleton_backend_url")
+
+    @property
+    def key_prefix(self) -> str:
+        return self.app.conf.get("singleton_key_prefix", "pyworkflow:lock:")
+
+    @property
+    def lock_expiry(self) -> int:
+        return self.app.conf.get("singleton_lock_expiry", 3600)
+
+    @property
+    def raise_on_duplicate(self) -> bool:
+        return self.app.conf.get("singleton_raise_on_duplicate", False)
+
+
+class RedisLockBackend:
+    """Redis backend for distributed locking."""
+
+    def __init__(self, url: str):
+        import redis
+
+        self.redis = redis.from_url(url, decode_responses=True)
+
+    def lock(self, lock_key: str, task_id: str, expiry: int | None = None) -> bool:
+        """Acquire lock atomically. Returns True if acquired."""
+        return bool(self.redis.set(lock_key, task_id, nx=True, ex=expiry))
+
+    def unlock(self, lock_key: str) -> None:
+        """Release the lock."""
+        self.redis.delete(lock_key)
+
+    def get(self, lock_key: str) -> str | None:
+        """Get the task ID holding the lock."""
+        return self.redis.get(lock_key)
+
+
+class DuplicateTaskError(Exception):
+    """Raised when attempting to queue a duplicate singleton task."""
+
+    def __init__(self, message: str, task_id: str):
+        self.task_id = task_id
+        super().__init__(message)
+
+
+class SingletonWorkflowTask(Task):
+    """
+    Base class for singleton workflow tasks with distributed locking.
+
+    Features:
+    - Redis-based lock prevents duplicate execution
+    - Support for unique_on with nested dict/list access (e.g., "data.run_id")
+    - Retry-safe: locks NOT released on failure (prevents duplicate during retries)
+    - Lock released on success or when max retries exceeded
+    - Time-based lock expiry as safety net
+
+    Configuration:
+        unique_on: List of argument names to use for uniqueness (e.g., ["run_id", "step_id"])
+                   Supports nested access with dot notation (e.g., ["data.run_id"])
+        raise_on_duplicate: If True, raise DuplicateTaskError instead of returning existing result
+        lock_expiry: Lock TTL in seconds (default: 3600 = 1 hour)
+
+    Example:
+        @celery_app.task(
+            base=SingletonWorkflowTask,
+            unique_on=["run_id", "step_id"],
+        )
+        def my_task(run_id: str, step_id: str, data: dict):
+            ...
+    """
+
+    abstract = True
+
+    # Singleton configuration (can be overridden per-task)
+    unique_on: list[str] | str | None = None
+    raise_on_duplicate: bool | None = None
+    lock_expiry: int | None = None
+
+    # Lock behavior
+    release_lock_on_success: bool = True
+    release_lock_on_failure: bool = False  # Keep lock during retries
+
+    # Celery task settings
+    max_retries: int | None = None
+    acks_on_failure_or_timeout: bool = True
+
+    # Cached instances (class-level, shared across task instances)
+    _singleton_backend: RedisLockBackend | None = None
+    _singleton_config: SingletonConfig | None = None
+
+    @property
+    def singleton_config(self) -> SingletonConfig:
+        if self._singleton_config is None:
+            self._singleton_config = SingletonConfig(self.app)
+        return self._singleton_config
+
+    @property
+    def singleton_backend(self) -> RedisLockBackend | None:
+        if self._singleton_backend is None:
+            url = self.singleton_config.backend_url
+            if not url:
+                # Try broker URL if it's Redis
+                broker = self.app.conf.broker_url or ""
+                if broker.startswith("redis://") or broker.startswith("rediss://"):
+                    url = broker
+            if url:
+                self._singleton_backend = RedisLockBackend(url)
+        return self._singleton_backend
+
+    @property
+    def _lock_expiry(self) -> int:
+        if self.lock_expiry is not None:
+            return self.lock_expiry
+        return self.singleton_config.lock_expiry
+
+    @property
+    def _raise_on_duplicate(self) -> bool:
+        if self.raise_on_duplicate is not None:
+            return self.raise_on_duplicate
+        return self.singleton_config.raise_on_duplicate
+
+    def generate_lock(
+        self,
+        task_name: str,
+        task_args: list[Any] | tuple[Any, ...] | None = None,
+        task_kwargs: dict[str, Any] | None = None,
+    ) -> str:
+        """Generate lock key, supporting nested attribute access via unique_on."""
+        unique_on = self.unique_on
+        task_args = task_args or []
+        task_kwargs = task_kwargs or {}
+
+        if unique_on:
+            if isinstance(unique_on, str):
+                unique_on = [unique_on]
+
+            # Bind arguments to function signature
+            sig = inspect.signature(self.run)
+            bound = sig.bind(*task_args, **task_kwargs).arguments
+
+            unique_args: list[Any] = []
+            for key in unique_on:
+                keys = key.split(".")
+                if keys[0] not in bound:
+                    raise ValueError(f"Key '{keys[0]}' not found in task arguments")
+
+                value = bound[keys[0]]
+                # Navigate nested structure (supports one level of nesting)
+                if len(keys) == 2:
+                    nested_key = keys[1]
+                    if isinstance(value, dict):
+                        if nested_key not in value:
+                            raise ValueError(f"Key '{nested_key}' not found in dict")
+                        unique_args.append(value[nested_key])
+                    elif isinstance(value, (list, tuple)):
+                        unique_args.append(value[int(nested_key)])
+                    elif hasattr(value, nested_key):
+                        unique_args.append(getattr(value, nested_key))
+                    else:
+                        raise ValueError(f"Key '{key}' has unsupported type")
+                elif len(keys) == 1:
+                    unique_args.append(value)
+                else:
+                    raise ValueError(f"Key '{key}' has too many levels (max 2)")
+
+            return generate_lock_key(
+                task_name,
+                unique_args,
+                {},
+                key_prefix=self.singleton_config.key_prefix,
+            )
+        else:
+            return generate_lock_key(
+                task_name,
+                list(task_args),
+                task_kwargs,
+                key_prefix=self.singleton_config.key_prefix,
+            )
+
+    def acquire_lock(self, lock_key: str, task_id: str) -> bool:
+        """Attempt to acquire lock. Returns True if successful."""
+        backend = self.singleton_backend
+        if backend is None:
+            return True  # No Redis = no locking
+        return backend.lock(lock_key, task_id, expiry=self._lock_expiry)
+
+    def release_lock(
+        self,
+        task_args: list[Any] | tuple[Any, ...] | None = None,
+        task_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        """Release the lock for this task."""
+        backend = self.singleton_backend
+        if backend is None:
+            return
+        lock_key = self.generate_lock(self.name, task_args, task_kwargs)
+        backend.unlock(lock_key)
+
+    def get_existing_task_id(self, lock_key: str) -> str | None:
+        """Get task ID holding the lock, if any."""
+        backend = self.singleton_backend
+        if backend is None:
+            return None
+        return backend.get(lock_key)
+
+    def apply_async(
+        self,
+        args: list[Any] | tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        task_id: str | None = None,
+        **options: Any,
+    ) -> Any:
+        """Override apply_async to implement singleton behavior."""
+        args = args or []
+        kwargs = kwargs or {}
+        task_id = task_id or str(uuid4())
+
+        backend = self.singleton_backend
+        if backend is None:
+            # No Redis = normal behavior
+            return super().apply_async(args, kwargs, task_id=task_id, **options)
+
+        lock_key = self.generate_lock(self.name, args, kwargs)
+
+        # Try to acquire lock and run
+        if self.acquire_lock(lock_key, task_id):
+            try:
+                return super().apply_async(args, kwargs, task_id=task_id, **options)
+            except Exception:
+                # Release lock if apply_async fails
+                backend.unlock(lock_key)
+                raise
+
+        # Lock not acquired - check for existing task
+        existing_task_id = self.get_existing_task_id(lock_key)
+        if existing_task_id:
+            logger.debug(
+                "Singleton: duplicate task blocked",
+                task=self.name,
+                existing_task_id=existing_task_id,
+            )
+            if self._raise_on_duplicate:
+                raise DuplicateTaskError(f"Duplicate of task {existing_task_id}", existing_task_id)
+            return self.AsyncResult(existing_task_id)
+
+        # Race condition: lock disappeared, retry
+        if self.acquire_lock(lock_key, task_id):
+            try:
+                return super().apply_async(args, kwargs, task_id=task_id, **options)
+            except Exception:
+                backend.unlock(lock_key)
+                raise
+
+        # Still can't acquire - return existing or submit anyway
+        existing_task_id = self.get_existing_task_id(lock_key)
+        if existing_task_id:
+            return self.AsyncResult(existing_task_id)
+
+        # Fallback: submit anyway (rare edge case)
+        logger.warning(f"Singleton lock unstable, submitting anyway: {self.name}")
+        return super().apply_async(args, kwargs, task_id=task_id, **options)
+
+    def on_success(
+        self, retval: Any, task_id: str, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> None:
+        """Release lock on successful task completion."""
+        if self.release_lock_on_success:
+            self.release_lock(task_args=args, task_kwargs=kwargs)
+
+    def on_failure(
+        self,
+        exc: Exception,
+        task_id: str,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+        einfo: Any,
+    ) -> None:
+        """
+        Retry-aware lock management on failure.
+
+        - If task will retry: Keep lock
+        - If max retries exceeded: Release lock
+        """
+        max_retries_exceeded = False
+        if hasattr(self, "request") and self.request:
+            current_retries = getattr(self.request, "retries", 0)
+            max_retries = self.max_retries if self.max_retries is not None else 3
+            max_retries_exceeded = current_retries >= max_retries
+
+        if self.release_lock_on_failure or max_retries_exceeded:
+            self.release_lock(task_args=args, task_kwargs=kwargs)
+            if max_retries_exceeded:
+                logger.warning(
+                    f"Task {self.name} failed after {current_retries} retries. Lock released.",
+                    task_id=task_id,
+                    error=str(exc),
+                )
+
+        # Log appropriately
+        if isinstance(exc, WorkerLostError):
+            logger.warning("Task interrupted due to worker loss", task_id=task_id)
+        else:
+            logger.error(
+                f"Task {self.name} failed: {exc}",
+                task_id=task_id,
+                traceback=einfo.traceback if einfo else None,
+            )
+
+    def on_retry(
+        self,
+        exc: Exception,
+        task_id: str,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+        einfo: Any,
+    ) -> None:
+        """Lock is retained during retry."""
+        logger.warning(
+            f"Task {self.name} retrying (lock retained)",
+            task_id=task_id,
+            retry_count=self.request.retries,
+        )