pyworkflow-engine 0.1.37__tar.gz → 0.2.0__tar.gz

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/PKG-INFO

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

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyproject.toml

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

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.1.37"
+__version__ = "0.2.0"
 
 # Configuration
 from pyworkflow.config import (
@@ -135,6 +135,12 @@ from pyworkflow.primitives.schedule import (
 )
 from pyworkflow.primitives.shield import shield
 from pyworkflow.primitives.sleep import sleep
+from pyworkflow.primitives.step_checkpoint import (
+    delete_step_checkpoint,
+    load_step_checkpoint,
+    save_step_checkpoint,
+)
+from pyworkflow.primitives.step_hook import step_hook
 
 # Runtime
 from pyworkflow.runtime import LocalRuntime, Runtime, get_runtime, register_runtime
@@ -156,6 +162,21 @@ from pyworkflow.storage.schemas import (
     WorkflowRun,
 )
 
+# Streams (pub/sub signal pattern)
+from pyworkflow.streams import (
+    CheckpointBackend,
+    Signal,
+    Stream,
+    StreamConsumer,
+    StreamStepContext,
+    emit,
+    get_checkpoint,
+    get_current_signal,
+    save_checkpoint,
+    stream_step,
+    stream_workflow,
+)
+
 __all__ = [
     # Version
     "__version__",
@@ -264,4 +285,21 @@ __all__ = [
     "get_logger",
     "bind_workflow_context",
     "bind_step_context",
+    # Step checkpoint + hooks
+    "save_step_checkpoint",
+    "load_step_checkpoint",
+    "delete_step_checkpoint",
+    "step_hook",
+    # Streams
+    "stream_workflow",
+    "stream_step",
+    "emit",
+    "Signal",
+    "Stream",
+    "StreamStepContext",
+    "StreamConsumer",
+    "CheckpointBackend",
+    "get_current_signal",
+    "get_checkpoint",
+    "save_checkpoint",
 ]

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/celery/scheduler.py

@@ -181,9 +181,11 @@ class PyWorkflowScheduler(Scheduler):
                     execute_scheduled_workflow_task.apply_async(
                         kwargs={
                             "schedule_id": schedule.schedule_id,
-                            "scheduled_time": schedule.next_run_time.isoformat()
-                            if schedule.next_run_time
-                            else now.isoformat(),
+                            "scheduled_time": (
+                                schedule.next_run_time.isoformat()
+                                if schedule.next_run_time
+                                else now.isoformat()
+                            ),
                             "storage_config": self._storage_config,
                         },
                         queue="pyworkflow.schedules",

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/celery/tasks.py

@@ -198,14 +198,36 @@ def execute_step_task(
             _exec_lock_key, self.request.id or "unknown", expiry=self._lock_expiry
         )
         if not _exec_lock_acquired:
-            logger.warning(
-                "Step already being executed by another worker, skipping duplicate",
-                run_id=run_id,
-                step_id=step_id,
-                step_name=step_name,
-                existing_worker_task=_exec_lock_backend.get(_exec_lock_key),
-            )
-            return None
+            # Check for worker-loss re-delivery: when a worker dies mid-step
+            # (e.g., spot instance termination), Celery re-delivers the same
+            # message with the same task_id (task_reject_on_worker_lost=True).
+            # The dead worker's lock still exists in Redis (up to 1h TTL).
+            # If the lock holder matches our own task_id, this IS a re-delivery
+            # of the same task — force-acquire and proceed.
+            _holding_task_id = _exec_lock_backend.get(_exec_lock_key)
+            _our_task_id = self.request.id or "unknown"
+            if _holding_task_id == _our_task_id:
+                logger.warning(
+                    "Step lock held by same task_id — worker-loss re-delivery detected, "
+                    "force-acquiring lock",
+                    run_id=run_id,
+                    step_id=step_id,
+                    step_name=step_name,
+                    task_id=_our_task_id,
+                )
+                _exec_lock_backend.unlock(_exec_lock_key)
+                _exec_lock_acquired = _exec_lock_backend.lock(
+                    _exec_lock_key, _our_task_id, expiry=self._lock_expiry
+                )
+            else:
+                logger.warning(
+                    "Step already being executed by another worker, skipping duplicate",
+                    run_id=run_id,
+                    step_id=step_id,
+                    step_name=step_name,
+                    existing_worker_task=_holding_task_id,
+                )
+                return None
     else:
         _exec_lock_acquired = True  # No Redis = no locking
 
@@ -279,6 +301,18 @@ def execute_step_task(
                 step_id=step_id,
             )
 
+    # Set up step execution context for checkpoint/hook primitives
+    step_exec_key = f"{run_id}:{step_id}"
+    step_exec_tokens = None
+    try:
+        from pyworkflow.primitives.step_checkpoint import (
+            set_step_execution_context,
+        )
+
+        step_exec_tokens = set_step_execution_context(step_exec_key, storage)
+    except Exception as e:
+        logger.warning(f"Failed to set up step execution context: {e}")
+
     # Execute step function
     try:
         # Get the original function (unwrapped from decorator)
@@ -339,10 +373,29 @@ def execute_step_task(
         raise
 
     except SuspensionSignal as e:
-        # A primitive (sleep, hook, child_workflow) raised SuspensionSignal.
-        # This should not normally happen because primitives detect step worker
-        # context and adjust behavior. Treat as a fatal error rather than
-        # silently retrying, which would re-execute the step from scratch.
+        # Check if this is a step_hook suspension (supported)
+        if e.reason and e.reason.startswith("step_hook:"):
+            hook_id = e.data.get("hook_id")
+            logger.info(
+                f"Step suspended via step_hook: {step_name}",
+                run_id=run_id,
+                step_id=step_id,
+                hook_id=hook_id,
+            )
+            # Record STEP_SUSPENDED event so the workflow knows this step
+            # needs re-dispatch (clears in-progress state during replay)
+            run_async(
+                _record_step_suspended(
+                    storage_config=storage_config,
+                    run_id=run_id,
+                    step_id=step_id,
+                    step_name=step_name,
+                    hook_id=hook_id or "",
+                )
+            )
+            return None
+
+        # Other SuspensionSignals are not supported from steps
         logger.error(
             f"Step raised SuspensionSignal (unsupported from steps): {step_name}",
             run_id=run_id,
@@ -449,6 +502,12 @@ def execute_step_task(
             with contextlib.suppress(Exception):
                 _exec_lock_backend.unlock(_exec_lock_key)
 
+        # Clean up step execution context (checkpoint/hook primitives)
+        if step_exec_tokens is not None:
+            from pyworkflow.primitives.step_checkpoint import reset_step_execution_context
+
+            reset_step_execution_context(step_exec_tokens)
+
         # Clean up workflow context (must be reset before step context)
         if workflow_context_token is not None:
             from pyworkflow.context.base import reset_context
@@ -567,6 +626,63 @@ async def _record_step_completion_and_resume(
     )
 
 
+async def _record_step_suspended(
+    storage_config: dict[str, Any] | None,
+    run_id: str,
+    step_id: str,
+    step_name: str,
+    hook_id: str,
+) -> None:
+    """
+    Record STEP_SUSPENDED event when a step suspends via step_hook.
+
+    This clears the step's in-progress state during replay so that the
+    workflow re-dispatches the step on resume (instead of thinking it's
+    still running).
+
+    Does NOT schedule workflow resumption — that happens when resume_hook()
+    is called externally.
+    """
+    from pyworkflow.engine.events import create_step_suspended_event
+
+    storage = _get_storage_backend(storage_config)
+    if hasattr(storage, "connect"):
+        await storage.connect()
+
+    # Wait for WORKFLOW_SUSPENDED event to avoid sequence number race
+    max_wait_attempts = 50
+    wait_interval = 0.01
+
+    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
+        await asyncio.sleep(wait_interval)
+    else:
+        logger.warning(
+            "Timeout waiting for WORKFLOW_SUSPENDED before recording STEP_SUSPENDED",
+            run_id=run_id,
+            step_id=step_id,
+        )
+
+    event = create_step_suspended_event(
+        run_id=run_id,
+        step_id=step_id,
+        step_name=step_name,
+        hook_id=hook_id,
+    )
+    await storage.record_event(event)
+    logger.info(
+        "Step suspended event recorded",
+        run_id=run_id,
+        step_id=step_id,
+        step_name=step_name,
+        hook_id=hook_id,
+    )
+
+
 async def _record_step_failure_and_resume(
     storage_config: dict[str, Any] | None,
     run_id: str,
@@ -782,13 +898,31 @@ def start_workflow_task(
             _exec_lock_key, self.request.id or "unknown", expiry=self._lock_expiry
         )
         if not _exec_lock_acquired:
-            logger.warning(
-                "Workflow start already being executed by another worker, skipping duplicate",
-                run_id=run_id,
-                workflow_name=workflow_name,
-                existing_worker_task=_exec_lock_backend.get(_exec_lock_key),
-            )
-            return run_id
+            # Check for worker-loss re-delivery (same logic as step tasks):
+            # if the lock holder matches our task_id, this is the same message
+            # re-delivered after the original worker died.
+            _holding_task_id = _exec_lock_backend.get(_exec_lock_key)
+            _our_task_id = self.request.id or "unknown"
+            if _holding_task_id == _our_task_id:
+                logger.warning(
+                    "Workflow start lock held by same task_id — worker-loss re-delivery "
+                    "detected, force-acquiring lock",
+                    run_id=run_id,
+                    workflow_name=workflow_name,
+                    task_id=_our_task_id,
+                )
+                _exec_lock_backend.unlock(_exec_lock_key)
+                _exec_lock_acquired = _exec_lock_backend.lock(
+                    _exec_lock_key, _our_task_id, expiry=self._lock_expiry
+                )
+            else:
+                logger.warning(
+                    "Workflow start already being executed by another worker, skipping duplicate",
+                    run_id=run_id,
+                    workflow_name=workflow_name,
+                    existing_worker_task=_holding_task_id,
+                )
+                return run_id
     else:
         _exec_lock_acquired = True  # No Redis = no locking
 

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/cli/commands/hooks.py

@@ -299,12 +299,12 @@ async def list_hooks_cmd(
                     "Name": hook.name or "-",
                     "Status": hook.status.value,
                     "Run ID": hook.run_id,
-                    "Created": hook.created_at.strftime("%Y-%m-%d %H:%M")
-                    if hook.created_at
-                    else "-",
-                    "Expires": hook.expires_at.strftime("%Y-%m-%d %H:%M")
-                    if hook.expires_at
-                    else "-",
+                    "Created": (
+                        hook.created_at.strftime("%Y-%m-%d %H:%M") if hook.created_at else "-"
+                    ),
+                    "Expires": (
+                        hook.expires_at.strftime("%Y-%m-%d %H:%M") if hook.expires_at else "-"
+                    ),
                 }
                 for hook in hooks_list
             ]
@@ -376,15 +376,15 @@ async def hook_info_cmd(ctx: click.Context, token: str) -> None:
                 "Run ID": hook.run_id,
                 "Name": hook.name or "-",
                 "Status": hook.status.value,
-                "Created": hook.created_at.strftime("%Y-%m-%d %H:%M:%S")
-                if hook.created_at
-                else "-",
-                "Expires": hook.expires_at.strftime("%Y-%m-%d %H:%M:%S")
-                if hook.expires_at
-                else "-",
-                "Received": hook.received_at.strftime("%Y-%m-%d %H:%M:%S")
-                if hook.received_at
-                else "-",
+                "Created": (
+                    hook.created_at.strftime("%Y-%m-%d %H:%M:%S") if hook.created_at else "-"
+                ),
+                "Expires": (
+                    hook.expires_at.strftime("%Y-%m-%d %H:%M:%S") if hook.expires_at else "-"
+                ),
+                "Received": (
+                    hook.received_at.strftime("%Y-%m-%d %H:%M:%S") if hook.received_at else "-"
+                ),
             }
 
             # Show payload if received

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/cli/commands/runs.py

@@ -151,9 +151,9 @@ async def list_runs(
                     "Run ID": run.run_id,
                     "Workflow": run.workflow_name,
                     "Status": run.status.value,
-                    "Started": run.started_at.strftime("%Y-%m-%d %H:%M:%S")
-                    if run.started_at
-                    else "-",
+                    "Started": (
+                        run.started_at.strftime("%Y-%m-%d %H:%M:%S") if run.started_at else "-"
+                    ),
                     "Duration": durations.get(run.run_id, "-"),
                 }
                 for run in runs_list
@@ -238,9 +238,9 @@ async def run_status(ctx: click.Context, run_id: str) -> None:
                 "Status": run.status.value,
                 "Created": run.created_at.strftime("%Y-%m-%d %H:%M:%S") if run.created_at else "-",
                 "Started": run.started_at.strftime("%Y-%m-%d %H:%M:%S") if run.started_at else "-",
-                "Completed": run.completed_at.strftime("%Y-%m-%d %H:%M:%S")
-                if run.completed_at
-                else "-",
+                "Completed": (
+                    run.completed_at.strftime("%Y-%m-%d %H:%M:%S") if run.completed_at else "-"
+                ),
                 "Duration": duration_str,
             }
 
@@ -622,9 +622,9 @@ async def list_children(
                     "Workflow": child.workflow_name,
                     "Status": child.status.value,
                     "Depth": child.nesting_depth,
-                    "Started": child.started_at.strftime("%Y-%m-%d %H:%M:%S")
-                    if child.started_at
-                    else "-",
+                    "Started": (
+                        child.started_at.strftime("%Y-%m-%d %H:%M:%S") if child.started_at else "-"
+                    ),
                     "Duration": _calc_duration(child),
                 }
                 for child in children

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/cli/commands/schedules.py

@@ -123,9 +123,9 @@ async def list_schedules_cmd(
                     "Workflow": s.workflow_name,
                     "Status": s.status.value,
                     "Schedule": describe_schedule(s.spec),
-                    "Next Run": s.next_run_time.strftime("%Y-%m-%d %H:%M:%S")
-                    if s.next_run_time
-                    else "-",
+                    "Next Run": (
+                        s.next_run_time.strftime("%Y-%m-%d %H:%M:%S") if s.next_run_time else "-"
+                    ),
                     "Runs": f"{s.successful_runs}/{s.total_runs}",
                 }
                 for s in schedules_list
@@ -241,9 +241,9 @@ async def create_schedule_cmd(
                 "status": schedule.status.value,
                 "spec": describe_schedule(schedule.spec),
                 "overlap_policy": schedule.overlap_policy.value,
-                "next_run_time": schedule.next_run_time.isoformat()
-                if schedule.next_run_time
-                else None,
+                "next_run_time": (
+                    schedule.next_run_time.isoformat() if schedule.next_run_time else None
+                ),
             }
             format_json(data)
         else:
@@ -308,9 +308,9 @@ async def show_schedule_cmd(
                     "timezone": schedule.spec.timezone,
                 },
                 "overlap_policy": schedule.overlap_policy.value,
-                "next_run_time": schedule.next_run_time.isoformat()
-                if schedule.next_run_time
-                else None,
+                "next_run_time": (
+                    schedule.next_run_time.isoformat() if schedule.next_run_time else None
+                ),
                 "last_run_at": schedule.last_run_at.isoformat() if schedule.last_run_at else None,
                 "total_runs": schedule.total_runs,
                 "successful_runs": schedule.successful_runs,
@@ -327,19 +327,25 @@ async def show_schedule_cmd(
                 "Status": schedule.status.value,
                 "Schedule": describe_schedule(schedule.spec),
                 "Overlap Policy": schedule.overlap_policy.value,
-                "Next Run": schedule.next_run_time.strftime("%Y-%m-%d %H:%M:%S")
-                if schedule.next_run_time
-                else "-",
-                "Last Run": schedule.last_run_at.strftime("%Y-%m-%d %H:%M:%S")
-                if schedule.last_run_at
-                else "-",
+                "Next Run": (
+                    schedule.next_run_time.strftime("%Y-%m-%d %H:%M:%S")
+                    if schedule.next_run_time
+                    else "-"
+                ),
+                "Last Run": (
+                    schedule.last_run_at.strftime("%Y-%m-%d %H:%M:%S")
+                    if schedule.last_run_at
+                    else "-"
+                ),
                 "Total Runs": schedule.total_runs,
                 "Successful": schedule.successful_runs,
                 "Failed": schedule.failed_runs,
                 "Skipped": schedule.skipped_runs,
-                "Created": schedule.created_at.strftime("%Y-%m-%d %H:%M:%S")
-                if schedule.created_at
-                else "-",
+                "Created": (
+                    schedule.created_at.strftime("%Y-%m-%d %H:%M:%S")
+                    if schedule.created_at
+                    else "-"
+                ),
             }
             format_key_value(data, title=f"Schedule: {schedule_id}")
 
@@ -440,9 +446,9 @@ async def resume_schedule_cmd(
             data = {
                 "schedule_id": schedule.schedule_id,
                 "status": schedule.status.value,
-                "next_run_time": schedule.next_run_time.isoformat()
-                if schedule.next_run_time
-                else None,
+                "next_run_time": (
+                    schedule.next_run_time.isoformat() if schedule.next_run_time else None
+                ),
             }
             format_json(data)
 
@@ -774,9 +780,9 @@ async def update_schedule_cmd(
                 "workflow_name": schedule.workflow_name,
                 "spec": describe_schedule(schedule.spec),
                 "overlap_policy": schedule.overlap_policy.value,
-                "next_run_time": schedule.next_run_time.isoformat()
-                if schedule.next_run_time
-                else None,
+                "next_run_time": (
+                    schedule.next_run_time.isoformat() if schedule.next_run_time else None
+                ),
             }
             format_json(data)
         else:

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/context/local.py

@@ -94,6 +94,10 @@ class LocalContext(WorkflowContext):
         self._cancellation_blocked: bool = False
         self._cancellation_reason: str | None = None
 
+        # Signal/stream state (used by EventReplayer for stream workflows)
+        self._signal_waits: dict[str, dict[str, Any]] = {}
+        self._received_signals: dict[str, dict[str, Any]] = {}
+
         # Child workflow state
         self._child_results: dict[str, dict[str, Any]] = {}
         self._pending_children: dict[str, str] = {}  # child_id -> child_run_id
@@ -161,6 +165,13 @@ class LocalContext(WorkflowContext):
                 # Step completed - no longer in progress
                 self._steps_in_progress.discard(step_id)
 
+            elif event.type == EventType.STEP_SUSPENDED:
+                step_id = event.data.get("step_id")
+                if step_id:
+                    # Step suspended via step_hook — no longer in progress,
+                    # needs re-dispatch on workflow resume
+                    self._steps_in_progress.discard(step_id)
+
             elif event.type == EventType.SLEEP_COMPLETED:
                 sleep_id = event.data.get("sleep_id")
                 self._completed_sleeps.add(sleep_id)

{pyworkflow_engine-0.1.37 → pyworkflow_engine-0.2.0}/pyworkflow/core/step.py

@@ -21,7 +21,7 @@ from typing import Any
 from loguru import logger
 
 from pyworkflow.context import get_context, has_context
-from pyworkflow.core.exceptions import FatalError, RetryableError
+from pyworkflow.core.exceptions import FatalError, RetryableError, SuspensionSignal
 from pyworkflow.core.registry import register_step
 from pyworkflow.core.validation import validate_step_parameters
 from pyworkflow.engine.events import (
@@ -29,6 +29,10 @@ from pyworkflow.engine.events import (
     create_step_failed_event,
     create_step_started_event,
 )
+from pyworkflow.primitives.step_checkpoint import (
+    reset_step_execution_context,
+    set_step_execution_context,
+)
 from pyworkflow.serialization.encoder import serialize, serialize_args, serialize_kwargs
 
 
@@ -177,8 +181,6 @@ def step(
                     step_id=step_id,
                 )
                 # Re-suspend and wait for existing task to complete
-                from pyworkflow.core.exceptions import SuspensionSignal
-
                 raise SuspensionSignal(
                     reason=f"step_dispatch:{step_id}",
                     step_id=step_id,
@@ -229,8 +231,6 @@ def step(
                             current_attempt=current_attempt,
                             resume_at=resume_at.isoformat(),
                         )
-                        from pyworkflow.core.exceptions import SuspensionSignal
-
                         raise SuspensionSignal(
                             reason=f"retry:{step_id}",
                             resume_at=resume_at,
@@ -268,6 +268,10 @@ def step(
             # Validate parameters before execution
             validate_step_parameters(func, args, kwargs, step_name)
 
+            # Set up step execution context for checkpoint/hook primitives
+            step_exec_key = f"{ctx.run_id}:{step_id}"
+            step_exec_tokens = set_step_execution_context(step_exec_key, ctx.storage)
+
             try:
                 # Execute step function
                 result = await func(*args, **kwargs)
@@ -295,6 +299,15 @@ def step(
 
                 return result
 
+            except SuspensionSignal:
+                # step_hook() raised SuspensionSignal — propagate to suspend workflow
+                logger.info(
+                    f"Step suspended via step_hook: {step_name}",
+                    run_id=ctx.run_id,
+                    step_id=step_id,
+                )
+                raise
+
             except FatalError as e:
                 # Fatal error - don't retry
                 logger.error(
@@ -393,8 +406,6 @@ def step(
 
                     # Raise suspension signal to pause workflow
                     # Note: The workflow-level exception handler will schedule automatic resumption
-                    from pyworkflow.core.exceptions import SuspensionSignal
-
                     raise SuspensionSignal(
                         reason=f"retry:{step_id}",
                         resume_at=resume_at,
@@ -434,6 +445,9 @@ def step(
                     else:
                         raise
 
+            finally:
+                reset_step_execution_context(step_exec_tokens)
+
         # Register step
         register_step(
             name=step_name,
@@ -645,7 +659,6 @@ async def _dispatch_step_to_celery(
         SuspensionSignal: To pause workflow while step executes on worker
     """
     from pyworkflow.celery.tasks import execute_step_task
-    from pyworkflow.core.exceptions import SuspensionSignal
     from pyworkflow.engine.events import EventType
 
     logger.info(
@@ -657,11 +670,17 @@ async def _dispatch_step_to_celery(
     # Defense-in-depth: check if STEP_STARTED was already recorded for this step.
     # This guards against duplicate dispatch when two resume tasks race and both
     # replay past the same step. If already started, re-suspend to wait.
+    # Exception: if STEP_SUSPENDED was recorded (step_hook suspension), the step
+    # needs re-dispatch so it can re-execute and find the HOOK_RECEIVED event.
     events = await ctx.storage.get_events(ctx.run_id)
     already_started = any(
         evt.type == EventType.STEP_STARTED and evt.data.get("step_id") == step_id for evt in events
     )
-    if already_started:
+    was_suspended = any(
+        evt.type == EventType.STEP_SUSPENDED and evt.data.get("step_id") == step_id
+        for evt in events
+    )
+    if already_started and not was_suspended:
         logger.info(
             f"Step {step_name} already has STEP_STARTED event, re-suspending",
             run_id=ctx.run_id,