pyworkflow-engine 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.1.7"
+__version__ = "0.1.9"
 
 # Configuration
 from pyworkflow.config import (
@@ -44,11 +44,15 @@ from pyworkflow.config import (
 from pyworkflow.context import (
     LocalContext,
     MockContext,
+    StepContext,
     WorkflowContext,
     get_context,
+    get_step_context,
     has_context,
+    has_step_context,
     reset_context,
     set_context,
+    set_step_context,
 )
 
 # Exceptions
@@ -224,6 +228,11 @@ __all__ = [
     "has_context",
     "set_context",
     "reset_context",
+    # Step context for distributed execution
+    "StepContext",
+    "get_step_context",
+    "has_step_context",
+    "set_step_context",
     # Registry
     "list_workflows",
     "get_workflow",

pyworkflow/celery/tasks.py

@@ -13,7 +13,10 @@ import asyncio
 import uuid
 from collections.abc import Callable
 from datetime import UTC, datetime
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from pyworkflow.context.step_context import StepContext
 
 from celery import Task
 from celery.exceptions import WorkerLostError
@@ -103,11 +106,16 @@ def execute_step_task(
     step_id: str,
     max_retries: int = 3,
     storage_config: dict[str, Any] | None = None,
+    context_data: dict[str, Any] | None = None,
+    context_class_name: str | None = None,
 ) -> Any:
     """
-    Execute a workflow step in a Celery worker.
+    Execute a workflow step on a Celery worker.
 
-    This task runs a single step and handles retries automatically.
+    This task:
+    1. Executes the step function
+    2. Records STEP_COMPLETED/STEP_FAILED event in storage
+    3. Triggers workflow resumption via resume_workflow_task
 
     Args:
         step_name: Name of the step function
@@ -117,6 +125,8 @@ def execute_step_task(
         step_id: Step execution ID
         max_retries: Maximum retry attempts
         storage_config: Storage backend configuration
+        context_data: Optional step context data (from workflow)
+        context_class_name: Optional fully qualified context class name
 
     Returns:
         Step result (serialized)
@@ -128,7 +138,7 @@ def execute_step_task(
     from pyworkflow.core.registry import _registry
 
     logger.info(
-        f"Executing step: {step_name}",
+        f"Executing dispatched step: {step_name}",
         run_id=run_id,
         step_id=step_id,
         attempt=self.request.retries + 1,
@@ -137,12 +147,49 @@ def execute_step_task(
     # Get step metadata
     step_meta = _registry.get_step(step_name)
     if not step_meta:
+        # Record failure and resume workflow
+        asyncio.run(
+            _record_step_failure_and_resume(
+                storage_config=storage_config,
+                run_id=run_id,
+                step_id=step_id,
+                step_name=step_name,
+                error=f"Step '{step_name}' not found in registry",
+                error_type="FatalError",
+                is_retryable=False,
+            )
+        )
         raise FatalError(f"Step '{step_name}' not found in registry")
 
     # Deserialize arguments
     args = deserialize_args(args_json)
     kwargs = deserialize_kwargs(kwargs_json)
 
+    # Set up step context if provided (read-only mode)
+    step_context_token = None
+    readonly_token = None
+
+    if context_data and context_class_name:
+        try:
+            from pyworkflow.context.step_context import (
+                _set_step_context_internal,
+                _set_step_context_readonly,
+            )
+
+            # Import context class dynamically
+            context_class = _resolve_context_class(context_class_name)
+            if context_class is not None:
+                step_ctx = context_class.from_dict(context_data)
+                step_context_token = _set_step_context_internal(step_ctx)
+                # Set readonly mode to prevent mutation in steps
+                readonly_token = _set_step_context_readonly(True)
+        except Exception as e:
+            logger.warning(
+                f"Failed to load step context: {e}",
+                run_id=run_id,
+                step_id=step_id,
+            )
+
     # Execute step function
     try:
         # Get the original function (unwrapped from decorator)
@@ -160,32 +207,225 @@ def execute_step_task(
             step_id=step_id,
         )
 
+        # Record STEP_COMPLETED event and trigger workflow resumption
+        asyncio.run(
+            _record_step_completion_and_resume(
+                storage_config=storage_config,
+                run_id=run_id,
+                step_id=step_id,
+                step_name=step_name,
+                result=result,
+            )
+        )
+
         return result
 
-    except FatalError:
+    except FatalError as e:
         logger.error(f"Step failed (fatal): {step_name}", run_id=run_id, step_id=step_id)
+        # Record failure and resume workflow (workflow will fail on replay)
+        asyncio.run(
+            _record_step_failure_and_resume(
+                storage_config=storage_config,
+                run_id=run_id,
+                step_id=step_id,
+                step_name=step_name,
+                error=str(e),
+                error_type=type(e).__name__,
+                is_retryable=False,
+            )
+        )
         raise
 
     except RetryableError as e:
-        logger.warning(
-            f"Step failed (retriable): {step_name}",
-            run_id=run_id,
-            step_id=step_id,
-            retry_after=e.retry_after,
-        )
-        # Let Celery handle the retry
-        raise self.retry(exc=e, countdown=e.get_retry_delay_seconds() or 60)
+        # Check if we have retries left
+        if self.request.retries < max_retries:
+            logger.warning(
+                f"Step failed (retriable): {step_name}, retrying...",
+                run_id=run_id,
+                step_id=step_id,
+                retry_after=e.retry_after,
+                attempt=self.request.retries + 1,
+                max_retries=max_retries,
+            )
+            # Let Celery handle the retry - don't resume workflow yet
+            raise self.retry(exc=e, countdown=e.get_retry_delay_seconds() or 60)
+        else:
+            # Max retries exhausted - record failure and resume workflow
+            logger.error(
+                f"Step failed after {max_retries + 1} attempts: {step_name}",
+                run_id=run_id,
+                step_id=step_id,
+            )
+            asyncio.run(
+                _record_step_failure_and_resume(
+                    storage_config=storage_config,
+                    run_id=run_id,
+                    step_id=step_id,
+                    step_name=step_name,
+                    error=str(e),
+                    error_type=type(e).__name__,
+                    is_retryable=False,  # Mark as not retryable since we exhausted retries
+                )
+            )
+            raise
 
     except Exception as e:
-        logger.error(
-            f"Step failed (unexpected): {step_name}",
-            run_id=run_id,
-            step_id=step_id,
-            error=str(e),
-            exc_info=True,
-        )
-        # Treat unexpected errors as retriable
-        raise self.retry(exc=RetryableError(str(e)), countdown=60)
+        # Check if we have retries left
+        if self.request.retries < max_retries:
+            logger.warning(
+                f"Step failed (unexpected): {step_name}, retrying...",
+                run_id=run_id,
+                step_id=step_id,
+                error=str(e),
+                attempt=self.request.retries + 1,
+            )
+            # Treat unexpected errors as retriable
+            raise self.retry(exc=RetryableError(str(e)), countdown=60)
+        else:
+            # Max retries exhausted
+            logger.error(
+                f"Step failed after {max_retries + 1} attempts: {step_name}",
+                run_id=run_id,
+                step_id=step_id,
+                error=str(e),
+                exc_info=True,
+            )
+            asyncio.run(
+                _record_step_failure_and_resume(
+                    storage_config=storage_config,
+                    run_id=run_id,
+                    step_id=step_id,
+                    step_name=step_name,
+                    error=str(e),
+                    error_type=type(e).__name__,
+                    is_retryable=False,
+                )
+            )
+            raise
+
+    finally:
+        # Clean up step context
+        if readonly_token is not None:
+            from pyworkflow.context.step_context import _reset_step_context_readonly
+
+            _reset_step_context_readonly(readonly_token)
+        if step_context_token is not None:
+            from pyworkflow.context.step_context import _reset_step_context
+
+            _reset_step_context(step_context_token)
+
+
+async def _record_step_completion_and_resume(
+    storage_config: dict[str, Any] | None,
+    run_id: str,
+    step_id: str,
+    step_name: str,
+    result: Any,
+) -> None:
+    """
+    Record STEP_COMPLETED event and trigger workflow resumption.
+
+    Called by execute_step_task after successful step execution.
+    """
+    from pyworkflow.engine.events import create_step_completed_event
+    from pyworkflow.serialization.encoder import serialize
+
+    # Get storage backend
+    storage = _get_storage_backend(storage_config)
+
+    # Ensure storage is connected
+    if hasattr(storage, "connect"):
+        await storage.connect()
+
+    # 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 immediately
+    schedule_workflow_resumption(run_id, datetime.now(UTC), storage_config)
+
+    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,
+    run_id: str,
+    step_id: str,
+    step_name: str,
+    error: str,
+    error_type: str,
+    is_retryable: bool,
+) -> None:
+    """
+    Record STEP_FAILED event and trigger workflow resumption.
+
+    Called by execute_step_task after step failure (when retries are exhausted).
+    The workflow will fail when it replays and sees the failure event.
+    """
+    from pyworkflow.engine.events import create_step_failed_event
+
+    # Get storage backend
+    storage = _get_storage_backend(storage_config)
+
+    # Ensure storage is connected
+    if hasattr(storage, "connect"):
+        await storage.connect()
+
+    # Record STEP_FAILED event
+    failure_event = create_step_failed_event(
+        run_id=run_id,
+        step_id=step_id,
+        error=error,
+        error_type=error_type,
+        is_retryable=is_retryable,
+        attempt=1,  # Final attempt
+    )
+    await storage.record_event(failure_event)
+
+    # Schedule workflow resumption - workflow will fail on replay
+    schedule_workflow_resumption(run_id, datetime.now(UTC), storage_config)
+
+    logger.info(
+        "Step failed and workflow resumption scheduled",
+        run_id=run_id,
+        step_id=step_id,
+        step_name=step_name,
+        error=error,
+    )
+
+
+def _resolve_context_class(class_name: str) -> type["StepContext"] | None:
+    """
+    Resolve a context class from its fully qualified name.
+
+    Args:
+        class_name: Fully qualified class name (e.g., "myapp.contexts.OrderContext")
+
+    Returns:
+        The class type, or None if resolution fails
+    """
+    try:
+        import importlib
+
+        parts = class_name.rsplit(".", 1)
+        if len(parts) == 2:
+            module_name, cls_name = parts
+            module = importlib.import_module(module_name)
+            return getattr(module, cls_name, None)
+        # Simple class name - try to get from globals
+        return None
+    except Exception:
+        return None
 
 
 @celery_app.task(
@@ -365,6 +605,8 @@ async def _execute_child_workflow_on_worker(
             storage=storage,
             durable=True,
             event_log=None,  # Fresh execution
+            runtime="celery",
+            storage_config=storage_config,
         )
 
         # Update status to COMPLETED
@@ -714,6 +956,8 @@ async def _recover_workflow_on_worker(
             args=args,
             kwargs=kwargs,
             event_log=events,
+            runtime="celery",
+            storage_config=storage_config,
         )
 
         # Update run status to completed
@@ -949,7 +1193,7 @@ async def _start_workflow_on_worker(
         input_kwargs=serialize_kwargs(**kwargs),
         idempotency_key=idempotency_key,
         max_duration=workflow_meta.max_duration,
-        metadata={},  # Run-level metadata (not from decorator)
+        context={},  # Step context (not from decorator)
         recovery_attempts=0,
         max_recovery_attempts=max_recovery_attempts,
         recover_on_worker_loss=recover_on_worker_loss,
@@ -977,6 +1221,8 @@ async def _start_workflow_on_worker(
             storage=storage,
             args=args,
             kwargs=kwargs,
+            runtime="celery",
+            storage_config=storage_config,
         )
 
         # Update run status to completed
@@ -1231,7 +1477,7 @@ async def _execute_scheduled_workflow(
             kwargs_json=kwargs_json,
             run_id=run_id,
             storage_config=storage_config,
-            metadata={"schedule_id": schedule_id, "scheduled_time": scheduled_time.isoformat()},
+            # Note: context data is passed through for scheduled workflows to include schedule info
         )
 
         # Record trigger event - use schedule_id as run_id since workflow run may not exist yet
@@ -1381,6 +1627,8 @@ async def _resume_workflow_on_worker(
             kwargs=kwargs,
             event_log=events,
             cancellation_requested=cancellation_requested,
+            runtime="celery",
+            storage_config=storage_config,
         )
 
         # Update run status to completed

pyworkflow/cli/__init__.py

@@ -27,7 +27,10 @@ from pyworkflow.cli.utils.storage import create_storage
 )
 @click.option(
     "--storage",
-    type=click.Choice(["file", "memory", "sqlite", "dynamodb"], case_sensitive=False),
+    type=click.Choice(
+        ["file", "memory", "sqlite", "postgres", "mysql", "dynamodb", "cassandra"],
+        case_sensitive=False,
+    ),
     envvar="PYWORKFLOW_STORAGE_BACKEND",
     help="Storage backend type (default: file)",
 )

pyworkflow/cli/commands/runs.py

@@ -227,7 +227,7 @@ async def run_status(ctx: click.Context, run_id: str) -> None:
                 "input_kwargs": json.loads(run.input_kwargs) if run.input_kwargs else None,
                 "result": json.loads(run.result) if run.result else None,
                 "error": run.error,
-                "metadata": run.metadata,
+                "context": run.context,
             }
             format_json(data)
 
@@ -266,9 +266,9 @@ async def run_status(ctx: click.Context, run_id: str) -> None:
             if run.error:
                 data["Error"] = run.error
 
-            # Add metadata if present
-            if run.metadata:
-                data["Metadata"] = json.dumps(run.metadata, indent=2)
+            # Add context if present
+            if run.context:
+                data["Context"] = json.dumps(run.context, indent=2)
 
             format_key_value(data, title=f"Workflow Run: {run_id}")