pyworkflow-engine 0.1.21__py3-none-any.whl → 0.1.23__py3-none-any.whl

pyworkflow/__init__.py

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

pyworkflow/celery/app.py

@@ -151,6 +151,8 @@ def create_celery_app(
     sentinel_master_name: str | None = None,
     broker_transport_options: dict[str, Any] | None = None,
     result_backend_transport_options: dict[str, Any] | None = None,
+    worker_max_memory_per_child: int | None = None,
+    worker_max_tasks_per_child: int | None = None,
 ) -> Celery:
     """
     Create and configure a Celery application for PyWorkflow.
@@ -162,6 +164,8 @@ def create_celery_app(
         sentinel_master_name: Redis Sentinel master name. Priority: parameter > PYWORKFLOW_CELERY_SENTINEL_MASTER env var > "mymaster"
         broker_transport_options: Additional transport options for the broker (merged with defaults)
         result_backend_transport_options: Additional transport options for the result backend (merged with defaults)
+        worker_max_memory_per_child: Max memory per worker child process (KB). Priority: parameter > PYWORKFLOW_WORKER_MAX_MEMORY env var > None (no limit)
+        worker_max_tasks_per_child: Max tasks per worker child before recycling. Priority: parameter > PYWORKFLOW_WORKER_MAX_TASKS env var > None (no limit)
 
     Returns:
         Configured Celery application
@@ -170,6 +174,8 @@ def create_celery_app(
         PYWORKFLOW_CELERY_BROKER: Celery broker URL (used if broker_url param not provided)
         PYWORKFLOW_CELERY_RESULT_BACKEND: Result backend URL (used if result_backend param not provided)
         PYWORKFLOW_CELERY_SENTINEL_MASTER: Sentinel master name (used if sentinel_master_name param not provided)
+        PYWORKFLOW_WORKER_MAX_MEMORY: Max memory per worker child (KB) (used if worker_max_memory_per_child param not provided)
+        PYWORKFLOW_WORKER_MAX_TASKS: Max tasks per worker child (used if worker_max_tasks_per_child param not provided)
 
     Examples:
         # Default configuration (uses env vars if set, otherwise localhost Redis)
@@ -202,6 +208,14 @@ def create_celery_app(
         or "redis://localhost:6379/1"
     )
 
+    # Worker memory limits (KB) - prevents memory leaks from accumulating
+    # Priority: parameter > env var > None (no limit by default)
+    max_memory_env = os.getenv("PYWORKFLOW_WORKER_MAX_MEMORY")
+    max_memory = worker_max_memory_per_child or (int(max_memory_env) if max_memory_env else None)
+
+    max_tasks_env = os.getenv("PYWORKFLOW_WORKER_MAX_TASKS")
+    max_tasks = worker_max_tasks_per_child or (int(max_tasks_env) if max_tasks_env else None)
+
     # Detect broker and backend types
     is_sentinel_broker = is_sentinel_url(broker_url)
     is_sentinel_backend = is_sentinel_url(result_backend)
@@ -310,6 +324,10 @@ def create_celery_app(
         # Logging
         worker_log_format="[%(asctime)s: %(levelname)s/%(processName)s] %(message)s",
         worker_task_log_format="[%(asctime)s: %(levelname)s/%(processName)s] [%(task_name)s(%(task_id)s)] %(message)s",
+        # Worker memory management - prevents memory leaks from accumulating
+        # When set, workers are recycled after exceeding these limits
+        worker_max_memory_per_child=max_memory,  # KB, None = no limit
+        worker_max_tasks_per_child=max_tasks,  # None = no limit
     )
 
     # Configure singleton locking for Redis or Sentinel brokers

pyworkflow/celery/tasks.py

@@ -171,10 +171,9 @@ 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
+    # Use has_event() for efficient EXISTS check instead of loading all events
+    already_completed = run_async(
+        storage.has_event(run_id, EventType.STEP_COMPLETED.value, step_id=step_id)
     )
     if already_completed:
         logger.warning(
@@ -379,9 +378,9 @@ async def _record_step_completion_and_resume(
 
     Called by execute_step_task after successful step execution.
 
-    Only schedules resume if WORKFLOW_SUSPENDED event exists, indicating
-    the workflow has fully suspended. This prevents race conditions where
-    a step completes before the workflow has suspended.
+    IMPORTANT: This function waits for WORKFLOW_SUSPENDED event before recording
+    STEP_COMPLETED to prevent race conditions where both events get the same
+    sequence number. The workflow must fully suspend before we record completion.
 
     Idempotency: If STEP_COMPLETED already exists for this step_id, skip
     recording and resume scheduling (another task already handled it).
@@ -397,10 +396,9 @@ async def _record_step_completion_and_resume(
         await storage.connect()
 
     # Idempotency check: skip if step already completed
-    events = await 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
+    # Use has_event() for efficient EXISTS check instead of loading all events
+    already_completed = await storage.has_event(
+        run_id, EventType.STEP_COMPLETED.value, step_id=step_id
     )
     if already_completed:
         logger.info(
@@ -411,43 +409,64 @@ async def _record_step_completion_and_resume(
         )
         return
 
-    # Record STEP_COMPLETED event
-    completion_event = create_step_completed_event(
-        run_id=run_id,
-        step_id=step_id,
-        result=serialize(result),
-        step_name=step_name,
-    )
-    await storage.record_event(completion_event)
+    # Wait for WORKFLOW_SUSPENDED event before recording STEP_COMPLETED
+    # This prevents race conditions where both events get the same sequence number
+    # Use has_event() for memory-efficient polling instead of loading all events
+    max_wait_attempts = 50  # 50 * 10ms = 500ms max wait
+    wait_interval = 0.01  # 10ms between checks
 
-    # Refresh events to include the one we just recorded
-    events = await storage.get_events(run_id)
+    for _attempt in range(max_wait_attempts):
+        has_suspended = await storage.has_event(
+            run_id, EventType.WORKFLOW_SUSPENDED.value, step_id=step_id
+        )
+        if has_suspended:
+            break
 
-    # Check if workflow has suspended (WORKFLOW_SUSPENDED event exists)
-    # Only schedule resume if workflow has properly suspended
-    has_suspended = any(evt.type == EventType.WORKFLOW_SUSPENDED for evt in events)
+        # Wait and check again
+        await asyncio.sleep(wait_interval)
 
-    if has_suspended:
-        # Workflow has suspended, safe to schedule resume
-        schedule_workflow_resumption(
-            run_id, datetime.now(UTC), storage_config, triggered_by="step_completed"
-        )
-        logger.info(
-            "Step completed and workflow resumption scheduled",
-            run_id=run_id,
-            step_id=step_id,
-            step_name=step_name,
+        # Also check if step was already completed by another task during wait
+        already_completed = await storage.has_event(
+            run_id, EventType.STEP_COMPLETED.value, step_id=step_id
         )
+        if already_completed:
+            logger.info(
+                "Step already completed by another task during wait, skipping",
+                run_id=run_id,
+                step_id=step_id,
+                step_name=step_name,
+            )
+            return
     else:
-        # Workflow hasn't suspended yet - don't schedule resume
-        # The suspension handler will check for step completion and schedule resume
-        logger.info(
-            "Step completed but workflow not yet suspended, skipping resume scheduling",
+        # Timeout waiting for suspension - log warning but proceed anyway
+        # This handles edge cases where the workflow completes without suspending
+        logger.warning(
+            "Timeout waiting for WORKFLOW_SUSPENDED event, proceeding with completion",
             run_id=run_id,
             step_id=step_id,
             step_name=step_name,
         )
 
+    # Record STEP_COMPLETED event
+    completion_event = create_step_completed_event(
+        run_id=run_id,
+        step_id=step_id,
+        result=serialize(result),
+        step_name=step_name,
+    )
+    await storage.record_event(completion_event)
+
+    # Schedule workflow resumption
+    schedule_workflow_resumption(
+        run_id, datetime.now(UTC), storage_config, triggered_by="step_completed"
+    )
+    logger.info(
+        "Step completed and workflow resumption scheduled",
+        run_id=run_id,
+        step_id=step_id,
+        step_name=step_name,
+    )
+
 
 async def _record_step_failure_and_resume(
     storage_config: dict[str, Any] | None,
@@ -464,9 +483,9 @@ async def _record_step_failure_and_resume(
     Called by execute_step_task after step failure (when retries are exhausted).
     The workflow will fail when it replays and sees the failure event.
 
-    Only schedules resume if WORKFLOW_SUSPENDED event exists, indicating
-    the workflow has fully suspended. This prevents race conditions where
-    a step fails before the workflow has suspended.
+    IMPORTANT: This function waits for WORKFLOW_SUSPENDED event before recording
+    STEP_FAILED to prevent race conditions where both events get the same
+    sequence number. The workflow must fully suspend before we record failure.
 
     Idempotency: If STEP_COMPLETED or terminal STEP_FAILED already exists
     for this step_id, skip recording and resume scheduling.
@@ -481,17 +500,18 @@ async def _record_step_failure_and_resume(
         await storage.connect()
 
     # Idempotency check: skip if step already completed or terminally failed
-    events = await storage.get_events(run_id)
-    already_handled = any(
-        (evt.type == EventType.STEP_COMPLETED and evt.data.get("step_id") == step_id)
-        or (
-            evt.type == EventType.STEP_FAILED
-            and evt.data.get("step_id") == step_id
-            and not evt.data.get("is_retryable", True)
-        )
-        for evt in events
+    # Use has_event() for efficient EXISTS check instead of loading all events
+    # Note: For STEP_FAILED with is_retryable check, we use has_event for STEP_COMPLETED
+    # and separately check STEP_FAILED (non-retryable failures are rare, so this is still efficient)
+    already_completed = await storage.has_event(
+        run_id, EventType.STEP_COMPLETED.value, step_id=step_id
+    )
+    # For terminal failures, we check separately (is_retryable=false in data)
+    # This is less common, so checking completion first is the fast path
+    already_failed_terminal = await storage.has_event(
+        run_id, EventType.STEP_FAILED.value, step_id=step_id, is_retryable="False"
     )
-    if already_handled:
+    if already_completed or already_failed_terminal:
         logger.info(
             "Step already completed/failed by another task, skipping",
             run_id=run_id,
@@ -500,6 +520,46 @@ async def _record_step_failure_and_resume(
         )
         return
 
+    # Wait for WORKFLOW_SUSPENDED event before recording STEP_FAILED
+    # This prevents race conditions where both events get the same sequence number
+    # Use has_event() for memory-efficient polling instead of loading all events
+    max_wait_attempts = 50  # 50 * 10ms = 500ms max wait
+    wait_interval = 0.01  # 10ms between checks
+
+    for _attempt in range(max_wait_attempts):
+        has_suspended = await storage.has_event(
+            run_id, EventType.WORKFLOW_SUSPENDED.value, step_id=step_id
+        )
+        if has_suspended:
+            break
+
+        # Wait and check again
+        await asyncio.sleep(wait_interval)
+
+        # Also check if step was already handled by another task during wait
+        already_completed = await storage.has_event(
+            run_id, EventType.STEP_COMPLETED.value, step_id=step_id
+        )
+        already_failed_terminal = await storage.has_event(
+            run_id, EventType.STEP_FAILED.value, step_id=step_id, is_retryable="False"
+        )
+        if already_completed or already_failed_terminal:
+            logger.info(
+                "Step already completed/failed by another task during wait, skipping",
+                run_id=run_id,
+                step_id=step_id,
+                step_name=step_name,
+            )
+            return
+    else:
+        # Timeout waiting for suspension - log warning but proceed anyway
+        logger.warning(
+            "Timeout waiting for WORKFLOW_SUSPENDED event, proceeding with failure",
+            run_id=run_id,
+            step_id=step_id,
+            step_name=step_name,
+        )
+
     # Record STEP_FAILED event
     failure_event = create_step_failed_event(
         run_id=run_id,
@@ -511,35 +571,17 @@ async def _record_step_failure_and_resume(
     )
     await storage.record_event(failure_event)
 
-    # Refresh events to include the one we just recorded
-    events = await storage.get_events(run_id)
-
-    # Check if workflow has suspended (WORKFLOW_SUSPENDED event exists)
-    # Only schedule resume if workflow has properly suspended
-    has_suspended = any(evt.type == EventType.WORKFLOW_SUSPENDED for evt in events)
-
-    if has_suspended:
-        # Workflow has suspended, safe to schedule resume
-        schedule_workflow_resumption(
-            run_id, datetime.now(UTC), storage_config, triggered_by="step_failed"
-        )
-        logger.info(
-            "Step failed and workflow resumption scheduled",
-            run_id=run_id,
-            step_id=step_id,
-            step_name=step_name,
-            error=error,
-        )
-    else:
-        # Workflow hasn't suspended yet - don't schedule resume
-        # The suspension handler will check for step failure and schedule resume
-        logger.info(
-            "Step failed but workflow not yet suspended, skipping resume scheduling",
-            run_id=run_id,
-            step_id=step_id,
-            step_name=step_name,
-            error=error,
-        )
+    # Schedule workflow resumption
+    schedule_workflow_resumption(
+        run_id, datetime.now(UTC), storage_config, triggered_by="step_failed"
+    )
+    logger.info(
+        "Step failed and workflow resumption scheduled",
+        run_id=run_id,
+        step_id=step_id,
+        step_name=step_name,
+        error=error,
+    )
 
 
 async def _get_workflow_run_safe(
@@ -839,13 +881,13 @@ async def _execute_child_workflow_on_worker(
 
         # For step dispatch suspensions, check if step already completed/failed
         if step_id and e.reason.startswith("step_dispatch:"):
-            events = await storage.get_events(child_run_id)
-            step_finished = any(
-                evt.type in (EventType.STEP_COMPLETED, EventType.STEP_FAILED)
-                and evt.data.get("step_id") == step_id
-                for evt in events
+            step_completed = await storage.has_event(
+                child_run_id, EventType.STEP_COMPLETED.value, step_id=step_id
+            )
+            step_failed = await storage.has_event(
+                child_run_id, EventType.STEP_FAILED.value, step_id=step_id
             )
-            if step_finished:
+            if step_completed or step_failed:
                 logger.info(
                     "Child step finished before suspension completed, scheduling resume",
                     child_run_id=child_run_id,
@@ -1092,8 +1134,8 @@ async def _handle_workflow_recovery(
         return False
 
     # Get last event sequence
-    events = await storage.get_events(run.run_id)
-    last_event_sequence = max((e.sequence or 0 for e in events), default=0) if events else None
+    latest_event = await storage.get_latest_event(run.run_id)
+    last_event_sequence = latest_event.sequence if latest_event else None
 
     # Record interruption event
     interrupted_event = create_workflow_interrupted_event(
@@ -1235,13 +1277,13 @@ async def _recover_workflow_on_worker(
 
         # For step dispatch suspensions, check if step already completed/failed
         if step_id and e.reason.startswith("step_dispatch:"):
-            events = await storage.get_events(run_id)
-            step_finished = any(
-                evt.type in (EventType.STEP_COMPLETED, EventType.STEP_FAILED)
-                and evt.data.get("step_id") == step_id
-                for evt in events
+            step_completed = await storage.has_event(
+                run_id, EventType.STEP_COMPLETED.value, step_id=step_id
             )
-            if step_finished:
+            step_failed = await storage.has_event(
+                run_id, EventType.STEP_FAILED.value, step_id=step_id
+            )
+            if step_completed or step_failed:
                 logger.info(
                     "Step finished before recovery suspension completed, scheduling resume",
                     run_id=run_id,
@@ -1627,13 +1669,13 @@ async def _start_workflow_on_worker(
         # For step dispatch suspensions, check if step already completed/failed (race condition)
         # If so, schedule resume immediately
         if step_id and e.reason.startswith("step_dispatch:"):
-            events = await storage.get_events(run_id)
-            step_finished = any(
-                evt.type in (EventType.STEP_COMPLETED, EventType.STEP_FAILED)
-                and evt.data.get("step_id") == step_id
-                for evt in events
+            step_completed = await storage.has_event(
+                run_id, EventType.STEP_COMPLETED.value, step_id=step_id
+            )
+            step_failed = await storage.has_event(
+                run_id, EventType.STEP_FAILED.value, step_id=step_id
             )
-            if step_finished:
+            if step_completed or step_failed:
                 logger.info(
                     "Step finished before suspension completed, scheduling resume",
                     run_id=run_id,
@@ -2217,13 +2259,13 @@ async def _resume_workflow_on_worker(
 
         # For step dispatch suspensions, check if step already completed/failed
         if step_id and e.reason.startswith("step_dispatch:"):
-            events = await storage.get_events(run_id)
-            step_finished = any(
-                evt.type in (EventType.STEP_COMPLETED, EventType.STEP_FAILED)
-                and evt.data.get("step_id") == step_id
-                for evt in events
+            step_completed = await storage.has_event(
+                run_id, EventType.STEP_COMPLETED.value, step_id=step_id
+            )
+            step_failed = await storage.has_event(
+                run_id, EventType.STEP_FAILED.value, step_id=step_id
             )
-            if step_finished:
+            if step_completed or step_failed:
                 logger.info(
                     "Step finished before resume suspension completed, scheduling resume",
                     run_id=run_id,

pyworkflow/storage/base.py

@@ -203,6 +203,42 @@ class StorageBackend(ABC):
         """
         pass
 
+    @abstractmethod
+    async def has_event(
+        self,
+        run_id: str,
+        event_type: str,
+        **filters: str,
+    ) -> bool:
+        """
+        Check if an event exists matching the criteria.
+
+        This is a memory-efficient alternative to get_events() when you only
+        need to check for existence. Uses SQL EXISTS queries in SQL backends
+        for O(1) memory usage instead of loading all events.
+
+        Args:
+            run_id: Workflow run identifier
+            event_type: Event type to check for (e.g., "step_completed")
+            **filters: Additional filters to match against event data fields.
+                       For example, step_id="abc" will check data->>'step_id' = 'abc'
+
+        Returns:
+            True if a matching event exists, False otherwise
+
+        Example:
+            # Check if step completed
+            exists = await storage.has_event(
+                run_id, "step_completed", step_id="step_123"
+            )
+
+            # Check if workflow suspended for a specific step
+            exists = await storage.has_event(
+                run_id, "workflow_suspended", step_id="step_123"
+            )
+        """
+        pass
+
     @abstractmethod
     async def get_latest_event(
         self,

pyworkflow/storage/cassandra.py

@@ -896,6 +896,40 @@ class CassandraStorageBackend(StorageBackend):
 
         return None
 
+    async def has_event(
+        self,
+        run_id: str,
+        event_type: str,
+        **filters: str,
+    ) -> bool:
+        """
+        Check if an event exists matching the criteria.
+
+        Loads events of the specified type and filters in Python for efficiency.
+
+        Args:
+            run_id: Workflow run identifier
+            event_type: Event type to check for
+            **filters: Additional filters for event data fields
+
+        Returns:
+            True if a matching event exists, False otherwise
+        """
+        # Load only events of the specific type
+        events = await self.get_events(run_id, event_types=[event_type])
+
+        # Filter in Python
+        for event in events:
+            match = True
+            for key, value in filters.items():
+                if str(event.data.get(key)) != str(value):
+                    match = False
+                    break
+            if match:
+                return True
+
+        return False
+
     # Step Operations
 
     async def create_step(self, step: StepExecution) -> None:

pyworkflow/storage/dynamodb.py

@@ -588,6 +588,40 @@ class DynamoDBStorageBackend(StorageBackend):
 
             return self._item_to_event(self._item_to_dict(items[0]))
 
+    async def has_event(
+        self,
+        run_id: str,
+        event_type: str,
+        **filters: str,
+    ) -> bool:
+        """
+        Check if an event exists matching the criteria.
+
+        Loads events of the specified type and filters in Python for efficiency.
+
+        Args:
+            run_id: Workflow run identifier
+            event_type: Event type to check for
+            **filters: Additional filters for event data fields
+
+        Returns:
+            True if a matching event exists, False otherwise
+        """
+        # Load only events of the specific type
+        events = await self.get_events(run_id, event_types=[event_type])
+
+        # Filter in Python
+        for event in events:
+            match = True
+            for key, value in filters.items():
+                if str(event.data.get(key)) != str(value):
+                    match = False
+                    break
+            if match:
+                return True
+
+        return False
+
     # Step Operations
 
     async def create_step(self, step: StepExecution) -> None:

pyworkflow/storage/file.py

@@ -373,6 +373,58 @@ class FileStorageBackend(StorageBackend):
         events = await self.get_events(run_id, event_types=[event_type] if event_type else None)
         return events[-1] if events else None
 
+    async def has_event(
+        self,
+        run_id: str,
+        event_type: str,
+        **filters: str,
+    ) -> bool:
+        """
+        Check if an event exists using file-based iteration with early termination.
+
+        Reads the events file line by line and returns as soon as a match is found,
+        avoiding loading the entire event log into memory.
+
+        Args:
+            run_id: Workflow run identifier
+            event_type: Event type to check for
+            **filters: Additional filters for event data fields
+
+        Returns:
+            True if a matching event exists, False otherwise
+        """
+        events_file = self.events_dir / f"{run_id}.jsonl"
+
+        if not events_file.exists():
+            return False
+
+        def _check() -> bool:
+            with events_file.open("r") as f:
+                for line in f:
+                    if not line.strip():
+                        continue
+
+                    data = json.loads(line)
+
+                    # Check event type
+                    if data["type"] != event_type:
+                        continue
+
+                    # Check all data filters
+                    match = True
+                    event_data = data.get("data", {})
+                    for key, value in filters.items():
+                        if str(event_data.get(key)) != str(value):
+                            match = False
+                            break
+
+                    if match:
+                        return True
+
+            return False
+
+        return await asyncio.to_thread(_check)
+
     # Step Operations
 
     async def create_step(self, step: StepExecution) -> None:

pyworkflow/storage/memory.py

@@ -250,6 +250,43 @@ class InMemoryStorageBackend(StorageBackend):
             # Return event with highest sequence
             return max(events, key=lambda e: e.sequence or 0)
 
+    async def has_event(
+        self,
+        run_id: str,
+        event_type: str,
+        **filters: str,
+    ) -> bool:
+        """
+        Check if an event exists by loading events of the specific type and filtering.
+
+        This approach:
+        1. Uses the event_types filter to load only events of the target type
+        2. Filters in Python on the loaded data
+        3. Significantly reduces memory vs loading ALL events
+
+        Args:
+            run_id: Workflow run identifier
+            event_type: Event type to check for
+            **filters: Additional filters for event data fields
+
+        Returns:
+            True if a matching event exists, False otherwise
+        """
+        # Load only events of the specific type
+        events = await self.get_events(run_id, event_types=[event_type])
+
+        # Filter in Python
+        for event in events:
+            match = True
+            for key, value in filters.items():
+                if str(event.data.get(key)) != str(value):
+                    match = False
+                    break
+            if match:
+                return True
+
+        return False
+
     # Step Operations
 
     async def create_step(self, step: StepExecution) -> None:

pyworkflow/storage/migrations/__init__.py

@@ -0,0 +1,15 @@
+"""
+Database schema migration framework for PyWorkflow storage backends.
+
+This module provides a migration framework that allows storage backends to
+evolve their schema over time while maintaining backward compatibility with
+existing databases.
+"""
+
+from pyworkflow.storage.migrations.base import (
+    Migration,
+    MigrationRegistry,
+    MigrationRunner,
+)
+
+__all__ = ["Migration", "MigrationRegistry", "MigrationRunner"]