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

pyworkflow/context/step_context.py

@@ -0,0 +1,295 @@
+"""
+Step Context - User-defined context accessible from steps during distributed execution.
+
+StepContext provides a type-safe, immutable context that can be accessed from steps
+running on remote Celery workers. Unlike WorkflowContext which is process-local,
+StepContext is serialized and passed to workers.
+
+Key design decisions:
+- **Immutable in steps**: Steps can only read context, not mutate it. This prevents
+  race conditions when multiple steps execute in parallel.
+- **Mutable in workflow**: Workflow code can update context via set_step_context().
+  Updates are recorded as CONTEXT_UPDATED events for deterministic replay.
+- **User-extensible**: Users subclass StepContext to define their own typed fields.
+
+Usage:
+    from pyworkflow.context import StepContext, get_step_context, set_step_context
+
+    # Define custom context
+    class OrderContext(StepContext):
+        workspace_id: str = ""
+        user_id: str = ""
+        order_id: str = ""
+
+    @workflow(context_class=OrderContext)
+    async def process_order(order_id: str, user_id: str):
+        # Initialize context in workflow
+        ctx = OrderContext(order_id=order_id, user_id=user_id)
+        await set_step_context(ctx)  # Note: async call
+
+        # Update context (creates new immutable instance)
+        ctx = get_step_context()
+        ctx = ctx.with_updates(workspace_id="ws-123")
+        await set_step_context(ctx)
+
+        result = await validate_order()
+        return result
+
+    @step
+    async def validate_order():
+        # Read-only access in steps
+        ctx = get_step_context()
+        print(f"Validating order {ctx.order_id}")
+
+        # This would raise RuntimeError - context is read-only in steps:
+        # set_step_context(ctx.with_updates(workspace_id="new"))
+
+        return {"valid": True}
+"""
+
+from contextvars import ContextVar, Token
+from typing import Any, Self
+
+from pydantic import BaseModel, ConfigDict
+
+
+class StepContext(BaseModel):
+    """
+    Base class for user-defined step context.
+
+    StepContext is immutable (frozen) to prevent accidental mutation.
+    Use with_updates() to create a new context with modified values.
+
+    The context is automatically:
+    - Persisted to storage when set_step_context() is called in workflow code
+    - Loaded from storage when a step executes on a Celery worker
+    - Replayed from CONTEXT_UPDATED events during workflow resumption
+
+    Example:
+        class FlowContext(StepContext):
+            workspace_id: str = ""
+            user_id: str = ""
+            attachments: list[str] = []
+
+        @workflow(context_class=FlowContext)
+        async def my_workflow():
+            ctx = FlowContext(workspace_id="ws-123")
+            set_step_context(ctx)
+            ...
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    def with_updates(self: Self, **kwargs: Any) -> Self:
+        """
+        Create a new context with updated values.
+
+        Since StepContext is immutable, this creates a new instance
+        with the specified fields updated.
+
+        Args:
+            **kwargs: Fields to update
+
+        Returns:
+            New StepContext instance with updated values
+
+        Example:
+            ctx = ctx.with_updates(workspace_id="ws-456", user_id="user-789")
+        """
+        return self.model_copy(update=kwargs)
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize context to dictionary for storage.
+
+        Returns:
+            Dictionary representation of the context
+        """
+        return self.model_dump()
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Self:
+        """
+        Deserialize context from storage.
+
+        Args:
+            data: Dictionary representation of the context
+
+        Returns:
+            StepContext instance
+        """
+        return cls.model_validate(data)
+
+
+# ---------------------------------------------------------------------------
+# Context Variables
+# ---------------------------------------------------------------------------
+
+# Current step context (may be None if not set)
+_step_context: ContextVar[StepContext | None] = ContextVar("step_context", default=None)
+
+# Whether context is read-only (True when executing inside a step)
+_step_context_readonly: ContextVar[bool] = ContextVar("step_context_readonly", default=False)
+
+# The context class registered with the workflow (for deserialization)
+_step_context_class: ContextVar[type[StepContext] | None] = ContextVar(
+    "step_context_class", default=None
+)
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def get_step_context() -> StepContext:
+    """
+    Get the current step context.
+
+    This function can be called from both workflow code and step code.
+    In step code, the context is read-only.
+
+    Returns:
+        Current StepContext instance
+
+    Raises:
+        RuntimeError: If no step context is available
+
+    Example:
+        @step
+        async def my_step():
+            ctx = get_step_context()
+            print(f"Working in workspace: {ctx.workspace_id}")
+    """
+    ctx = _step_context.get()
+    if ctx is None:
+        raise RuntimeError(
+            "No step context available. "
+            "Ensure the workflow is decorated with @workflow(context_class=YourContext) "
+            "and set_step_context() was called."
+        )
+    return ctx
+
+
+async def set_step_context(ctx: StepContext) -> None:
+    """
+    Set the current step context and persist to storage.
+
+    This function can only be called from workflow code, not from within steps.
+    When called, the context is persisted to storage and a CONTEXT_UPDATED event
+    is recorded for deterministic replay.
+
+    Args:
+        ctx: The StepContext instance to set
+
+    Raises:
+        RuntimeError: If called from within a step (read-only mode)
+        TypeError: If ctx is not a StepContext instance
+
+    Example:
+        @workflow(context_class=OrderContext)
+        async def my_workflow():
+            ctx = OrderContext(order_id="123")
+            await set_step_context(ctx)  # OK - in workflow code
+
+            await my_step()  # Step cannot call set_step_context()
+    """
+    if _step_context_readonly.get():
+        raise RuntimeError(
+            "Cannot modify step context within a step. "
+            "Context is read-only during step execution to prevent race conditions. "
+            "Return data from the step and update context in workflow code instead."
+        )
+
+    if not isinstance(ctx, StepContext):
+        raise TypeError(f"Expected StepContext instance, got {type(ctx).__name__}")
+
+    # Set the context in the contextvar
+    _step_context.set(ctx)
+
+    # Persist to storage if we're in a durable workflow
+    from pyworkflow.context import get_context, has_context
+
+    if has_context():
+        workflow_ctx = get_context()
+        if workflow_ctx.is_durable and workflow_ctx.storage is not None:
+            from pyworkflow.engine.events import create_context_updated_event
+
+            # Record CONTEXT_UPDATED event for replay
+            event = create_context_updated_event(
+                run_id=workflow_ctx.run_id,
+                context_data=ctx.to_dict(),
+            )
+            await workflow_ctx.storage.record_event(event)
+
+            # Update the WorkflowRun.context field
+            await workflow_ctx.storage.update_run_context(workflow_ctx.run_id, ctx.to_dict())
+
+
+def has_step_context() -> bool:
+    """
+    Check if step context is available.
+
+    Returns:
+        True if step context is set, False otherwise
+    """
+    return _step_context.get() is not None
+
+
+def get_step_context_class() -> type[StepContext] | None:
+    """
+    Get the registered step context class for the current workflow.
+
+    Returns:
+        The StepContext subclass, or None if not registered
+    """
+    return _step_context_class.get()
+
+
+# ---------------------------------------------------------------------------
+# Internal API (for framework use)
+# ---------------------------------------------------------------------------
+
+
+def _set_step_context_internal(ctx: StepContext | None) -> Token[StepContext | None]:
+    """
+    Internal: Set step context without readonly check.
+
+    Used by the framework when loading context on workers.
+    """
+    return _step_context.set(ctx)
+
+
+def _reset_step_context(token: Token[StepContext | None]) -> None:
+    """
+    Internal: Reset step context to previous value.
+    """
+    _step_context.reset(token)
+
+
+def _set_step_context_readonly(readonly: bool) -> Token[bool]:
+    """
+    Internal: Set readonly mode for step execution.
+    """
+    return _step_context_readonly.set(readonly)
+
+
+def _reset_step_context_readonly(token: Token[bool]) -> None:
+    """
+    Internal: Reset readonly mode.
+    """
+    _step_context_readonly.reset(token)
+
+
+def _set_step_context_class(cls: type[StepContext] | None) -> Token[type[StepContext] | None]:
+    """
+    Internal: Set the context class for deserialization.
+    """
+    return _step_context_class.set(cls)
+
+
+def _reset_step_context_class(token: Token[type[StepContext] | None]) -> None:
+    """
+    Internal: Reset context class.
+    """
+    _step_context_class.reset(token)

pyworkflow/core/registry.py

@@ -22,6 +22,7 @@ class WorkflowMetadata:
     max_duration: str | None = None
     tags: list[str] | None = None
     description: str | None = None  # Docstring from the workflow function
+    context_class: type | None = None  # StepContext subclass for step context access
 
     def __post_init__(self) -> None:
         if self.tags is None:
@@ -70,6 +71,7 @@ class WorkflowRegistry:
         original_func: Callable[..., Any],
         max_duration: str | None = None,
         tags: list[str] | None = None,
+        context_class: type | None = None,
     ) -> None:
         """
         Register a workflow.
@@ -80,6 +82,7 @@ class WorkflowRegistry:
             original_func: Original unwrapped function
             max_duration: Optional maximum duration
             tags: Optional list of tags (max 3)
+            context_class: Optional StepContext subclass for step context access
         """
         if name in self._workflows:
             existing = self._workflows[name]
@@ -96,6 +99,7 @@ class WorkflowRegistry:
             original_func=original_func,
             max_duration=max_duration,
             tags=tags or [],
+            context_class=context_class,
         )
 
         self._workflows[name] = workflow_meta
@@ -250,9 +254,10 @@ def register_workflow(
     original_func: Callable[..., Any],
     max_duration: str | None = None,
     tags: list[str] | None = None,
+    context_class: type | None = None,
 ) -> None:
     """Register a workflow in the global registry."""
-    _registry.register_workflow(name, func, original_func, max_duration, tags)
+    _registry.register_workflow(name, func, original_func, max_duration, tags, context_class)
 
 
 def get_workflow(name: str) -> WorkflowMetadata | None:

pyworkflow/core/step.py

@@ -128,6 +128,20 @@ def step(
 
             # Check if step has already completed (replay)
             if not ctx.should_execute_step(step_id):
+                # Check if step failed (for distributed step dispatch)
+                if ctx.has_step_failed(step_id):
+                    error_info = ctx.get_step_failure(step_id)
+                    logger.error(
+                        f"Step {step_name} failed on remote worker",
+                        run_id=ctx.run_id,
+                        step_id=step_id,
+                        error=error_info.get("error") if error_info else "Unknown error",
+                    )
+                    raise FatalError(
+                        f"Step {step_name} failed: "
+                        f"{error_info.get('error') if error_info else 'Unknown error'}"
+                    )
+
                 logger.debug(
                     f"Step {step_name} already completed, using cached result",
                     run_id=ctx.run_id,
@@ -135,6 +149,22 @@ def step(
                 )
                 return ctx.get_step_result(step_id)
 
+            # ========== Distributed Step Dispatch ==========
+            # When running in a distributed runtime (e.g., Celery), dispatch steps
+            # to step workers instead of executing inline.
+            if ctx.runtime == "celery":
+                return await _dispatch_step_to_celery(
+                    ctx=ctx,
+                    func=func,
+                    args=args,
+                    kwargs=kwargs,
+                    step_name=step_name,
+                    step_id=step_id,
+                    max_retries=max_retries,
+                    retry_delay=retry_delay,
+                    timeout=timeout,
+                )
+
             # Check if we're resuming from a retry
             retry_state = ctx.get_retry_state(step_id)
             if retry_state:
@@ -492,3 +522,114 @@ def _generate_step_id(step_name: str, args: tuple, kwargs: dict) -> str:
     hash_hex = hashlib.sha256(content.encode()).hexdigest()[:16]
 
     return f"step_{step_name}_{hash_hex}"
+
+
+async def _dispatch_step_to_celery(
+    ctx: Any,  # WorkflowContext
+    func: Callable,
+    args: tuple,
+    kwargs: dict,
+    step_name: str,
+    step_id: str,
+    max_retries: int,
+    retry_delay: str | int | list[int],
+    timeout: int | None,
+) -> Any:
+    """
+    Dispatch step execution to Celery step worker.
+
+    Instead of executing the step inline, this function:
+    1. Records STEP_STARTED event
+    2. Dispatches the step to execute_step_task on the steps queue
+    3. Raises SuspensionSignal to pause the workflow
+
+    The step worker will:
+    1. Execute the step function
+    2. Record STEP_COMPLETED/STEP_FAILED event
+    3. Trigger workflow resumption
+
+    Args:
+        ctx: Workflow context
+        func: Step function to execute
+        args: Positional arguments
+        kwargs: Keyword arguments
+        step_name: Name of the step
+        step_id: Deterministic step ID
+        max_retries: Maximum retry attempts
+        retry_delay: Retry delay strategy
+        timeout: Optional timeout in seconds
+
+    Returns:
+        This function never returns normally - it always raises SuspensionSignal
+
+    Raises:
+        SuspensionSignal: To pause workflow while step executes on worker
+    """
+    from pyworkflow.celery.tasks import execute_step_task
+    from pyworkflow.core.exceptions import SuspensionSignal
+
+    logger.info(
+        f"Dispatching step to Celery worker: {step_name}",
+        run_id=ctx.run_id,
+        step_id=step_id,
+    )
+
+    # Validate event limits before recording step event
+    await ctx.validate_event_limits()
+
+    # Record STEP_STARTED event
+    start_event = create_step_started_event(
+        run_id=ctx.run_id,
+        step_id=step_id,
+        step_name=step_name,
+        args=serialize_args(*args),
+        kwargs=serialize_kwargs(**kwargs),
+        attempt=1,
+    )
+    await ctx.storage.record_event(start_event)
+
+    # Serialize arguments for Celery transport
+    args_json = serialize_args(*args)
+    kwargs_json = serialize_kwargs(**kwargs)
+
+    # Get step context data if available
+    context_data = None
+    context_class_name = None
+    try:
+        from pyworkflow.context.step_context import get_step_context, has_step_context
+
+        if has_step_context():
+            step_ctx = get_step_context()
+            context_data = step_ctx.to_dict()
+            context_class_name = f"{step_ctx.__class__.__module__}.{step_ctx.__class__.__name__}"
+    except Exception:
+        pass  # Step context not available
+
+    # Dispatch to Celery step queue
+    task_result = execute_step_task.delay(
+        step_name=step_name,
+        args_json=args_json,
+        kwargs_json=kwargs_json,
+        run_id=ctx.run_id,
+        step_id=step_id,
+        max_retries=max_retries,
+        storage_config=ctx.storage_config,
+        context_data=context_data,
+        context_class_name=context_class_name,
+    )
+
+    logger.info(
+        f"Step dispatched to Celery: {step_name}",
+        run_id=ctx.run_id,
+        step_id=step_id,
+        task_id=task_result.id,
+    )
+
+    # Raise suspension signal - workflow will pause until step completes
+    # The step worker will record STEP_COMPLETED and trigger resume
+    raise SuspensionSignal(
+        reason=f"step_dispatch:{step_id}",
+        step_id=step_id,
+        step_name=step_name,
+        task_id=task_result.id,
+    )

pyworkflow/core/workflow.py

@@ -33,6 +33,7 @@ def workflow(
     tags: list[str] | None = None,
     recover_on_worker_loss: bool | None = None,
     max_recovery_attempts: int | None = None,
+    context_class: type | None = None,
 ) -> Callable:
     """
     Decorator to mark async functions as workflows.
@@ -49,6 +50,9 @@ def workflow(
         recover_on_worker_loss: Whether to auto-recover on worker failure
             (None = True for durable, False for transient)
         max_recovery_attempts: Max recovery attempts on worker failure (default: 3)
+        context_class: Optional StepContext subclass for step context access.
+            When provided, enables get_step_context() and set_step_context()
+            for passing context data to steps running on remote workers.
 
     Returns:
         Decorated workflow function
@@ -85,6 +89,18 @@ def workflow(
         async def tagged_workflow():
             result = await my_step()
             return result
+
+    Example (with step context):
+        class OrderContext(StepContext):
+            workspace_id: str = ""
+            user_id: str = ""
+
+        @workflow(context_class=OrderContext)
+        async def workflow_with_context():
+            ctx = OrderContext(workspace_id="ws-123")
+            set_step_context(ctx)  # Persisted and available in steps
+            result = await my_step()  # Step can call get_step_context()
+            return result
     """
 
     # Validate tags
@@ -110,6 +126,7 @@ def workflow(
             original_func=func,
             max_duration=max_duration,
             tags=validated_tags,
+            context_class=context_class,
         )
 
         # Store metadata on wrapper
@@ -124,6 +141,7 @@ def workflow(
         wrapper.__workflow_max_recovery_attempts__ = (  # type: ignore[attr-defined]
             max_recovery_attempts  # None = use config default
         )
+        wrapper.__workflow_context_class__ = context_class  # type: ignore[attr-defined]
 
         return wrapper
 
@@ -140,6 +158,8 @@ async def execute_workflow_with_context(
     event_log: list | None = None,
     durable: bool = True,
     cancellation_requested: bool = False,
+    runtime: str | None = None,
+    storage_config: dict | None = None,
 ) -> Any:
     """
     Execute workflow function with proper context setup.
@@ -161,6 +181,8 @@ async def execute_workflow_with_context(
         event_log: Optional existing event log for replay
         durable: Whether this is a durable workflow
         cancellation_requested: Whether cancellation was requested before execution
+        runtime: Runtime environment slug (e.g., "celery") for distributed step dispatch
+        storage_config: Storage configuration dict for distributed step workers
 
     Returns:
         Workflow result
@@ -182,6 +204,10 @@ async def execute_workflow_with_context(
         durable=is_durable,
     )
 
+    # Set runtime environment for distributed step dispatch
+    ctx._runtime = runtime
+    ctx._storage_config = storage_config
+
     # Set cancellation state if requested before execution
     if cancellation_requested:
         ctx.request_cancellation(reason="Cancellation requested before execution")
@@ -189,6 +215,26 @@ async def execute_workflow_with_context(
     # Set as current context using new API
     token = set_context(ctx)
 
+    # Set up step context if workflow has a context_class
+    step_context_token = None
+    step_context_class_token = None
+    context_class = getattr(workflow_func, "__workflow_context_class__", None)
+    if context_class is not None:
+        from pyworkflow.context.step_context import (
+            _set_step_context_class,
+            _set_step_context_internal,
+        )
+
+        # Store context class for deserialization
+        step_context_class_token = _set_step_context_class(context_class)
+
+        # Load existing context from storage (for resumption)
+        if is_durable and storage is not None:
+            context_data = await storage.get_run_context(run_id)
+            if context_data:
+                step_ctx = context_class.from_dict(context_data)
+                step_context_token = _set_step_context_internal(step_ctx)
+
     try:
         # Note: Event replay is handled by LocalContext in its constructor
         # when event_log is provided
@@ -290,5 +336,15 @@ async def execute_workflow_with_context(
         raise
 
     finally:
+        # Clear step context if it was set
+        if step_context_token is not None:
+            from pyworkflow.context.step_context import _reset_step_context
+
+            _reset_step_context(step_context_token)
+        if step_context_class_token is not None:
+            from pyworkflow.context.step_context import _reset_step_context_class
+
+            _reset_step_context_class(step_context_class_token)
+
         # Clear context using new API
         reset_context(token)

pyworkflow/engine/events.py

@@ -45,6 +45,9 @@ class EventType(Enum):
     # Cancellation events
     CANCELLATION_REQUESTED = "cancellation.requested"
 
+    # Context events
+    CONTEXT_UPDATED = "context.updated"
+
     # Child workflow events
     CHILD_WORKFLOW_STARTED = "child_workflow.started"
     CHILD_WORKFLOW_COMPLETED = "child_workflow.completed"
@@ -480,6 +483,33 @@ def create_step_cancelled_event(
     )
 
 
+def create_context_updated_event(
+    run_id: str,
+    context_data: dict[str, Any],
+) -> Event:
+    """
+    Create a context updated event.
+
+    This event is recorded when set_step_context() is called in workflow code.
+    It captures the full context state for deterministic replay.
+
+    Args:
+        run_id: The workflow run ID
+        context_data: The serialized context data (from StepContext.to_dict())
+
+    Returns:
+        Event: The context updated event
+    """
+    return Event(
+        run_id=run_id,
+        type=EventType.CONTEXT_UPDATED,
+        data={
+            "context": context_data,
+            "updated_at": datetime.now(UTC).isoformat(),
+        },
+    )
+
+
 # Child workflow event creation helpers