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

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
 
 

pyworkflow/engine/replay.py

@@ -93,6 +93,9 @@ class EventReplayer:
         elif event.type == EventType.CANCELLATION_REQUESTED:
             await self._apply_cancellation_requested(ctx, event)
 
+        elif event.type == EventType.CONTEXT_UPDATED:
+            await self._apply_context_updated(ctx, event)
+
         # Other event types don't affect replay state
         # (workflow_started, step_started, step_failed, etc. are informational)
 
@@ -255,6 +258,42 @@ class EventReplayer:
             requested_by=requested_by,
         )
 
+    async def _apply_context_updated(self, ctx: LocalContext, event: Event) -> None:
+        """
+        Apply context_updated event - restore step context.
+
+        During replay, this restores the step context to its state at the time
+        the event was recorded. This ensures deterministic replay.
+        """
+        from pyworkflow.context.step_context import (
+            _set_step_context_internal,
+            get_step_context_class,
+        )
+
+        context_data = event.data.get("context", {})
+
+        if context_data:
+            # Get the registered context class
+            context_class = get_step_context_class()
+            if context_class is not None:
+                try:
+                    step_ctx = context_class.from_dict(context_data)
+                    _set_step_context_internal(step_ctx)
+                    logger.debug(
+                        "Restored step context from replay",
+                        run_id=ctx.run_id,
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to restore step context: {e}",
+                        run_id=ctx.run_id,
+                    )
+            else:
+                logger.debug(
+                    "No context class registered, skipping context restoration",
+                    run_id=ctx.run_id,
+                )
+
 
 # Singleton instance
 _replayer = EventReplayer()

pyworkflow/primitives/child_workflow.py

@@ -346,7 +346,7 @@ async def _start_child_on_worker(
         parent_run_id=ctx.run_id,
         nesting_depth=child_depth,
         max_duration=workflow_meta.max_duration,
-        metadata={},  # Run-level metadata
+        context={},  # Step context
     )
     await storage.create_run(child_run)
 

pyworkflow/runtime/local.py

@@ -167,7 +167,7 @@ class LocalRuntime(Runtime):
                     input_kwargs=serialize_kwargs(**kwargs),
                     idempotency_key=idempotency_key,
                     max_duration=max_duration,
-                    metadata=metadata or {},
+                    context=metadata or {},
                 )
                 await storage.create_run(workflow_run)
 

pyworkflow/storage/__init__.py

@@ -35,6 +35,18 @@ try:
 except ImportError:
     DynamoDBStorageBackend = None  # type: ignore
 
+# Cassandra backend - optional import (requires cassandra-driver)
+try:
+    from pyworkflow.storage.cassandra import CassandraStorageBackend
+except ImportError:
+    CassandraStorageBackend = None  # type: ignore
+
+# MySQL backend - optional import (requires aiomysql)
+try:
+    from pyworkflow.storage.mysql import MySQLStorageBackend
+except ImportError:
+    MySQLStorageBackend = None  # type: ignore
+
 __all__ = [
     "StorageBackend",
     "FileStorageBackend",
@@ -42,6 +54,8 @@ __all__ = [
     "SQLiteStorageBackend",
     "PostgresStorageBackend",
     "DynamoDBStorageBackend",
+    "CassandraStorageBackend",
+    "MySQLStorageBackend",
     "WorkflowRun",
     "StepExecution",
     "Hook",

pyworkflow/storage/base.py

@@ -109,6 +109,41 @@ class StorageBackend(ABC):
         """
         pass
 
+    @abstractmethod
+    async def update_run_context(
+        self,
+        run_id: str,
+        context: dict,
+    ) -> None:
+        """
+        Update the step context for a workflow run.
+
+        Called when set_step_context() is invoked in workflow code.
+        The context is stored and can be loaded by steps running on
+        remote workers.
+
+        Args:
+            run_id: Workflow run identifier
+            context: Context data as a dictionary (serialized StepContext)
+        """
+        pass
+
+    @abstractmethod
+    async def get_run_context(self, run_id: str) -> dict:
+        """
+        Get the current step context for a workflow run.
+
+        Called when a step starts execution on a remote worker to
+        load the context that was set by the workflow.
+
+        Args:
+            run_id: Workflow run identifier
+
+        Returns:
+            Context data as a dictionary, or empty dict if not set
+        """
+        pass
+
     @abstractmethod
     async def list_runs(
         self,