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

pyworkflow/__init__.py

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

pyworkflow/celery/app.py

@@ -194,6 +194,16 @@ def create_celery_app(
         worker_task_log_format="[%(asctime)s: %(levelname)s/%(processName)s] [%(task_name)s(%(task_id)s)] %(message)s",
     )
 
+    # 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.

pyworkflow/celery/singleton.py

@@ -0,0 +1,370 @@
+"""
+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: lock released in on_retry callback to allow retry to acquire it
+    - 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  # Only release on max retries exceeded
+
+    # 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:
+        """Release lock during retry to allow retry task to acquire it."""
+        # Release lock so retry can acquire it via apply_async()
+        self.release_lock(task_args=args, task_kwargs=kwargs)
+        logger.warning(
+            f"Task {self.name} retrying (lock released for retry)",
+            task_id=task_id,
+            retry_count=self.request.retries,
+        )

pyworkflow/celery/tasks.py

@@ -19,12 +19,12 @@ from typing import TYPE_CHECKING, Any
 if TYPE_CHECKING:
     from pyworkflow.context.step_context import StepContext
 
-from celery import Task
-from celery.exceptions import MaxRetriesExceededError, Retry, WorkerLostError
+from celery.exceptions import MaxRetriesExceededError, Retry
 from loguru import logger
 
 from pyworkflow.celery.app import celery_app
 from pyworkflow.celery.loop import run_async
+from pyworkflow.celery.singleton import SingletonWorkflowTask
 from pyworkflow.core.exceptions import (
     CancellationError,
     ContinueAsNewSignal,
@@ -33,6 +33,7 @@ from pyworkflow.core.exceptions import (
     SuspensionSignal,
 )
 from pyworkflow.core.registry import WorkflowMetadata, get_workflow
+from pyworkflow.core.validation import validate_step_parameters
 from pyworkflow.core.workflow import execute_workflow_with_context
 from pyworkflow.engine.events import (
     EventType,
@@ -73,58 +74,15 @@ def _calculate_exponential_backoff(
     return delay * jitter
 
 
-class WorkflowTask(Task):
-    """Base task class for workflow execution with custom error handling."""
-
-    # Allow unlimited Celery-level retries - our code controls the actual limit
-    # via the max_retries parameter passed to execute_step_task
-    max_retries = None
-    # Prevent message requeue loops when task fails
-    acks_on_failure_or_timeout = True
-
-    def on_failure(self, exc, task_id, args, kwargs, einfo):
-        """
-        Handle task failure.
-
-        Detects worker loss and handles recovery appropriately:
-        - WorkerLostError: Infrastructure failure, may trigger recovery
-        - Other exceptions: Application failure
-        """
-        is_worker_loss = isinstance(exc, WorkerLostError)
-        if is_worker_loss:
-            logger.warning(
-                f"Task {self.name} interrupted due to worker loss",
-                task_id=task_id,
-                error=str(exc),
-            )
-            # Note: Recovery is handled when the task is requeued and picked up
-            # by another worker. See _handle_workflow_recovery() for logic.
-        else:
-            logger.error(
-                f"Task {self.name} failed: {str(exc)}",
-                task_id=task_id,
-                error=str(exc),
-                traceback=einfo.traceback if einfo else None,
-            )
-
-    def on_retry(self, exc, task_id, args, kwargs, einfo):
-        """Handle task retry."""
-        logger.warning(
-            f"Task {self.name} retrying",
-            task_id=task_id,
-            error=str(exc),
-            retry_count=self.request.retries,
-        )
-
-
 @celery_app.task(
     name="pyworkflow.execute_step",
-    base=WorkflowTask,
+    base=SingletonWorkflowTask,
     bind=True,
     queue="pyworkflow.steps",
+    unique_on=["run_id", "step_id"],
 )
 def execute_step_task(
-    self: WorkflowTask,
+    self: SingletonWorkflowTask,
     step_name: str,
     args_json: str,
     kwargs_json: str,
@@ -212,10 +170,28 @@ def execute_step_task(
         )
         raise FatalError(f"Step '{step_name}' not found in registry")
 
+    # Ignore processing step if already completed (idempotency)
+    events = run_async(storage.get_events(run_id))
+    already_completed = any(
+        evt.type == EventType.STEP_COMPLETED and evt.data.get("step_id") == step_id
+        for evt in events
+    )
+    if already_completed:
+        logger.warning(
+            "Step already completed by another task, skipping execution",
+            run_id=run_id,
+            step_id=step_id,
+            step_name=step_name,
+        )
+        return None
+
     # Deserialize arguments
     args = deserialize_args(args_json)
     kwargs = deserialize_kwargs(kwargs_json)
 
+    # Validate parameters before execution on worker (defense in depth)
+    validate_step_parameters(step_meta.original_func, args, kwargs, step_name)
+
     # Set up step context if provided (read-only mode)
     step_context_token = None
     readonly_token = None
@@ -345,7 +321,7 @@ def execute_step_task(
             # Use exponential backoff for unexpected errors
             countdown = _calculate_exponential_backoff(self.request.retries)
             logger.warning(
-                f"Step failed (unexpected): {step_name}, retrying in {countdown:.1f}s...",
+                f"Step failed (unexpected): {step_name}, retrying in {countdown:.1f}s...: {str(e)}",
                 run_id=run_id,
                 step_id=step_id,
                 error=str(e),
@@ -608,8 +584,9 @@ def _resolve_context_class(class_name: str) -> type["StepContext"] | None:
 
 @celery_app.task(
     name="pyworkflow.start_workflow",
-    base=WorkflowTask,
+    base=SingletonWorkflowTask,
     queue="pyworkflow.workflows",
+    unique_on=["run_id"],
 )
 def start_workflow_task(
     workflow_name: str,
@@ -678,8 +655,9 @@ def start_workflow_task(
 
 @celery_app.task(
     name="pyworkflow.start_child_workflow",
-    base=WorkflowTask,
+    base=SingletonWorkflowTask,
     queue="pyworkflow.workflows",
+    unique_on=["child_run_id"],
 )
 def start_child_workflow_task(
     workflow_name: str,
@@ -1719,12 +1697,14 @@ async def _start_workflow_on_worker(
 
 @celery_app.task(
     name="pyworkflow.resume_workflow",
-    base=WorkflowTask,
+    base=SingletonWorkflowTask,
     queue="pyworkflow.schedules",
+    unique_on=["run_id"],
 )
 def resume_workflow_task(
     run_id: str,
     storage_config: dict[str, Any] | None = None,
+    triggered_by_hook_id: str | None = None,
 ) -> Any | None:
     """
     Resume a suspended workflow.
@@ -1735,6 +1715,9 @@ def resume_workflow_task(
     Args:
         run_id: Workflow run ID to resume
         storage_config: Storage backend configuration
+        triggered_by_hook_id: Optional hook ID that triggered this resume.
+                              Used to prevent spurious resumes when a workflow
+                              has already moved past the triggering hook.
 
     Returns:
         Workflow result if completed, None if suspended again
@@ -1748,13 +1731,18 @@ def resume_workflow_task(
         f"RESUME_WORKFLOW_TASK ENTRY: {run_id}",
         run_id=run_id,
         celery_task_id=resume_workflow_task.request.id,
+        triggered_by_hook_id=triggered_by_hook_id,
     )
 
     # Get storage backend
     storage = _get_storage_backend(storage_config)
 
     # Resume workflow directly on worker
-    result = run_async(_resume_workflow_on_worker(run_id, storage, storage_config))
+    result = run_async(
+        _resume_workflow_on_worker(
+            run_id, storage, storage_config, triggered_by_hook_id=triggered_by_hook_id
+        )
+    )
 
     if result is not None:
         logger.info(f"Workflow completed on worker: {run_id}")
@@ -1766,8 +1754,9 @@ def resume_workflow_task(
 
 @celery_app.task(
     name="pyworkflow.execute_scheduled_workflow",
-    base=WorkflowTask,
+    base=SingletonWorkflowTask,
     queue="pyworkflow.schedules",
+    # No unique_on - scheduled workflows create new runs each time, no deduplication needed
 )
 def execute_scheduled_workflow_task(
     schedule_id: str,
@@ -1960,15 +1949,81 @@ async def _complete_pending_sleeps(
     return updated_events
 
 
+def _is_hook_still_relevant(hook_id: str, events: list[Any]) -> bool:
+    """
+    Check if a hook is still relevant for resuming the workflow.
+
+    A hook is "still relevant" if there are no newer hooks created after
+    this hook was received. This prevents spurious resumes when:
+    1. resume_hook() is called multiple times for the same hook
+    2. The workflow moved past the first resume and created a new hook
+    3. The duplicate resume task runs but the workflow is now waiting on a different hook
+
+    Args:
+        hook_id: The hook ID that triggered the resume
+        events: List of workflow events
+
+    Returns:
+        True if the hook is still relevant, False if workflow has moved past it
+    """
+    from pyworkflow.engine.events import EventType
+
+    # Sort events by sequence to process in order
+    sorted_events = sorted(events, key=lambda e: e.sequence or 0)
+
+    # Find the sequence number of HOOK_RECEIVED for this hook
+    hook_received_sequence = None
+    for event in sorted_events:
+        if event.type == EventType.HOOK_RECEIVED and event.data.get("hook_id") == hook_id:
+            hook_received_sequence = event.sequence
+            break
+
+    if hook_received_sequence is None:
+        # Hook was never received - shouldn't happen, but allow resume
+        logger.warning(
+            f"Hook {hook_id} was not found in HOOK_RECEIVED events, allowing resume",
+            hook_id=hook_id,
+        )
+        return True
+
+    # Check if there's a HOOK_CREATED event after this hook was received
+    # (indicating the workflow has moved past this hook and created a new one)
+    for event in sorted_events:
+        if event.type == EventType.HOOK_CREATED:
+            event_sequence = event.sequence or 0
+            if event_sequence > hook_received_sequence:
+                # There's a newer hook - this resume is stale
+                newer_hook_id = event.data.get("hook_id")
+                logger.debug(
+                    f"Found newer hook {newer_hook_id} (seq {event_sequence}) "
+                    f"after triggered hook {hook_id} (received at seq {hook_received_sequence})",
+                    hook_id=hook_id,
+                    newer_hook_id=newer_hook_id,
+                )
+                return False
+
+    # No newer hooks created - this resume is still relevant
+    return True
+
+
 async def _resume_workflow_on_worker(
     run_id: str,
     storage: StorageBackend,
     storage_config: dict[str, Any] | None = None,
+    triggered_by_hook_id: str | None = None,
 ) -> Any | None:
     """
     Internal function to resume workflow on Celery worker.
 
     This mirrors the logic from testing.py but runs on workers.
+
+    Args:
+        run_id: Workflow run ID to resume
+        storage: Storage backend
+        storage_config: Storage configuration for task dispatch
+        triggered_by_hook_id: Optional hook ID that triggered this resume.
+                              If provided, we verify the hook is still relevant
+                              before resuming to prevent spurious resumes.
     """
     from pyworkflow.core.exceptions import WorkflowNotFoundError
 
@@ -2003,6 +2058,22 @@ async def _resume_workflow_on_worker(
         )
         return None
 
+    # If this resume was triggered by a specific hook, verify the hook is still relevant.
+    # A hook is "stale" if the workflow has already moved past it (created a newer hook).
+    # This prevents spurious resumes from duplicate resume_hook() calls.
+    if triggered_by_hook_id:
+        events = await storage.get_events(run_id)
+        hook_still_relevant = _is_hook_still_relevant(triggered_by_hook_id, events)
+        if not hook_still_relevant:
+            logger.info(
+                f"Hook {triggered_by_hook_id} is no longer relevant (workflow moved past it), "
+                "skipping spurious resume",
+                run_id=run_id,
+                workflow_name=run.workflow_name,
+                triggered_by_hook_id=triggered_by_hook_id,
+            )
+            return None
+
     # Check for cancellation flag
     cancellation_requested = await storage.check_cancellation_flag(run_id)