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

pyworkflow/cli/utils/docker_manager.py

@@ -195,6 +195,131 @@ volumes:
     return template
 
 
+def generate_cassandra_docker_compose_content(
+    cassandra_host: str = "cassandra",
+    cassandra_port: int = 9042,
+    cassandra_keyspace: str = "pyworkflow",
+    cassandra_user: str | None = None,
+    cassandra_password: str | None = None,
+) -> str:
+    """
+    Generate docker-compose.yml content for Cassandra-based PyWorkflow services.
+
+    Args:
+        cassandra_host: Cassandra host (container name for docker networking)
+        cassandra_port: Cassandra CQL native transport port
+        cassandra_keyspace: Cassandra keyspace name
+        cassandra_user: Optional Cassandra user (for authentication)
+        cassandra_password: Optional Cassandra password (for authentication)
+
+    Returns:
+        docker-compose.yml content as string
+
+    Example:
+        >>> compose_content = generate_cassandra_docker_compose_content()
+    """
+    # Build environment variables for dashboard
+    dashboard_env = f"""      - DASHBOARD_STORAGE_TYPE=cassandra
+      - DASHBOARD_CASSANDRA_HOST={cassandra_host}
+      - DASHBOARD_CASSANDRA_PORT=9042
+      - DASHBOARD_CASSANDRA_KEYSPACE={cassandra_keyspace}"""
+
+    if cassandra_user and cassandra_password:
+        dashboard_env += f"""
+      - DASHBOARD_CASSANDRA_USER={cassandra_user}
+      - DASHBOARD_CASSANDRA_PASSWORD={cassandra_password}"""
+
+    # Cassandra environment for authentication
+    cassandra_env = ""
+    if cassandra_user and cassandra_password:
+        cassandra_env = f"""
+      - CASSANDRA_AUTHENTICATOR=PasswordAuthenticator
+      - CASSANDRA_AUTHORIZER=CassandraAuthorizer
+      - CASSANDRA_USER={cassandra_user}
+      - CASSANDRA_PASSWORD={cassandra_password}"""
+
+    template = f"""services:
+  cassandra:
+    image: cassandra:4.1
+    container_name: pyworkflow-cassandra
+    ports:
+      - "{cassandra_port}:9042"
+    environment:
+      - CASSANDRA_CLUSTER_NAME=pyworkflow-cluster
+      - CASSANDRA_DC=dc1
+      - CASSANDRA_RACK=rack1
+      - CASSANDRA_ENDPOINT_SNITCH=GossipingPropertyFileSnitch
+      - MAX_HEAP_SIZE=512M
+      - HEAP_NEWSIZE=128M{cassandra_env}
+    volumes:
+      - cassandra_data:/var/lib/cassandra
+    healthcheck:
+      test: ["CMD-SHELL", "cqlsh -e 'describe cluster' || exit 1"]
+      interval: 15s
+      timeout: 10s
+      retries: 10
+      start_period: 60s
+    restart: unless-stopped
+
+  redis:
+    image: redis:7-alpine
+    container_name: pyworkflow-redis
+    ports:
+      - "6379:6379"
+    volumes:
+      - redis_data:/data
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 5s
+      timeout: 3s
+      retries: 5
+    restart: unless-stopped
+
+  dashboard-backend:
+    image: yashabro/pyworkflow-dashboard-backend:latest
+    container_name: pyworkflow-dashboard-backend
+    working_dir: /app/project
+    ports:
+      - "8585:8585"
+    environment:
+      - DASHBOARD_PYWORKFLOW_CONFIG_PATH=/app/project/pyworkflow.config.yaml
+{dashboard_env}
+      - DASHBOARD_HOST=0.0.0.0
+      - DASHBOARD_PORT=8585
+      - DASHBOARD_CORS_ORIGINS=["http://localhost:5173","http://localhost:3000"]
+      - PYWORKFLOW_CELERY_BROKER=redis://redis:6379/0
+      - PYWORKFLOW_CELERY_RESULT_BACKEND=redis://redis:6379/1
+      - PYTHONPATH=/app/project
+    volumes:
+      - .:/app/project:ro
+    depends_on:
+      cassandra:
+        condition: service_healthy
+      redis:
+        condition: service_healthy
+    restart: unless-stopped
+
+  dashboard-frontend:
+    image: yashabro/pyworkflow-dashboard-frontend:latest
+    container_name: pyworkflow-dashboard-frontend
+    ports:
+      - "5173:80"
+    environment:
+      - VITE_API_URL=http://localhost:8585
+    depends_on:
+      - dashboard-backend
+    restart: unless-stopped
+
+volumes:
+  cassandra_data:
+    driver: local
+  redis_data:
+    driver: local
+"""
+
+    return template
+
+
 def generate_postgres_docker_compose_content(
     postgres_host: str = "postgres",
     postgres_port: int = 5432,
@@ -301,6 +426,113 @@ volumes:
     return template
 
 
+def generate_mysql_docker_compose_content(
+    mysql_host: str = "mysql",
+    mysql_port: int = 3306,
+    mysql_user: str = "pyworkflow",
+    mysql_password: str = "pyworkflow",
+    mysql_database: str = "pyworkflow",
+) -> str:
+    """
+    Generate docker-compose.yml content for MySQL-based PyWorkflow services.
+
+    Args:
+        mysql_host: MySQL host (container name for docker networking)
+        mysql_port: MySQL port
+        mysql_user: MySQL user
+        mysql_password: MySQL password
+        mysql_database: MySQL database name
+
+    Returns:
+        docker-compose.yml content as string
+
+    Example:
+        >>> compose_content = generate_mysql_docker_compose_content()
+    """
+    template = f"""services:
+  mysql:
+    image: mysql:8.0
+    container_name: pyworkflow-mysql
+    ports:
+      - "{mysql_port}:3306"
+    environment:
+      - MYSQL_ROOT_PASSWORD={mysql_password}
+      - MYSQL_USER={mysql_user}
+      - MYSQL_PASSWORD={mysql_password}
+      - MYSQL_DATABASE={mysql_database}
+    volumes:
+      - mysql_data:/var/lib/mysql
+    healthcheck:
+      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-u", "{mysql_user}", "-p{mysql_password}"]
+      interval: 5s
+      timeout: 3s
+      retries: 5
+    restart: unless-stopped
+
+  redis:
+    image: redis:7-alpine
+    container_name: pyworkflow-redis
+    ports:
+      - "6379:6379"
+    volumes:
+      - redis_data:/data
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 5s
+      timeout: 3s
+      retries: 5
+    restart: unless-stopped
+
+  dashboard-backend:
+    image: yashabro/pyworkflow-dashboard-backend:latest
+    container_name: pyworkflow-dashboard-backend
+    working_dir: /app/project
+    ports:
+      - "8585:8585"
+    environment:
+      - DASHBOARD_PYWORKFLOW_CONFIG_PATH=/app/project/pyworkflow.config.yaml
+      - DASHBOARD_STORAGE_TYPE=mysql
+      - DASHBOARD_MYSQL_HOST={mysql_host}
+      - DASHBOARD_MYSQL_PORT=3306
+      - DASHBOARD_MYSQL_USER={mysql_user}
+      - DASHBOARD_MYSQL_PASSWORD={mysql_password}
+      - DASHBOARD_MYSQL_DATABASE={mysql_database}
+      - DASHBOARD_HOST=0.0.0.0
+      - DASHBOARD_PORT=8585
+      - DASHBOARD_CORS_ORIGINS=["http://localhost:5173","http://localhost:3000"]
+      - PYWORKFLOW_CELERY_BROKER=redis://redis:6379/0
+      - PYWORKFLOW_CELERY_RESULT_BACKEND=redis://redis:6379/1
+      - PYTHONPATH=/app/project
+    volumes:
+      - .:/app/project:ro
+    depends_on:
+      mysql:
+        condition: service_healthy
+      redis:
+        condition: service_healthy
+    restart: unless-stopped
+
+  dashboard-frontend:
+    image: yashabro/pyworkflow-dashboard-frontend:latest
+    container_name: pyworkflow-dashboard-frontend
+    ports:
+      - "5173:80"
+    environment:
+      - VITE_API_URL=http://localhost:8585
+    depends_on:
+      - dashboard-backend
+    restart: unless-stopped
+
+volumes:
+  mysql_data:
+    driver: local
+  redis_data:
+    driver: local
+"""
+
+    return template
+
+
 def write_docker_compose(content: str, path: Path) -> None:
     """
     Write docker-compose.yml to file.

pyworkflow/context/__init__.py

@@ -41,6 +41,13 @@ from pyworkflow.context.base import (
 )
 from pyworkflow.context.local import LocalContext
 from pyworkflow.context.mock import MockContext
+from pyworkflow.context.step_context import (
+    StepContext,
+    get_step_context,
+    get_step_context_class,
+    has_step_context,
+    set_step_context,
+)
 
 __all__ = [
     # Base context and helpers
@@ -49,6 +56,12 @@ __all__ = [
     "has_context",
     "set_context",
     "reset_context",
+    # Step context for distributed execution
+    "StepContext",
+    "get_step_context",
+    "has_step_context",
+    "set_step_context",
+    "get_step_context_class",
     # Context implementations
     "LocalContext",
     "MockContext",

pyworkflow/context/base.py

@@ -281,6 +281,32 @@ class WorkflowContext(ABC):
         """Get the storage backend."""
         return None
 
+    @property
+    def runtime(self) -> str | None:
+        """
+        Get the runtime environment slug.
+
+        Returns the runtime identifier (e.g., "celery", "temporal") or None
+        for local/inline execution. Used to determine step dispatch behavior.
+
+        Returns:
+            Runtime slug string or None for local execution
+        """
+        return None  # Default: local/inline execution
+
+    @property
+    def storage_config(self) -> dict[str, Any] | None:
+        """Get storage configuration for distributed workers."""
+        return None
+
+    def has_step_failed(self, step_id: str) -> bool:
+        """Check if a step has a recorded failure."""
+        return False  # Default: no failures
+
+    def get_step_failure(self, step_id: str) -> dict[str, Any] | None:
+        """Get failure info for a step."""
+        return None  # Default: no failure info
+
     def should_execute_step(self, step_id: str) -> bool:
         """Check if step should be executed (not already completed)."""
         return True  # Default: always execute

pyworkflow/context/local.py

@@ -97,6 +97,13 @@ class LocalContext(WorkflowContext):
         self._child_results: dict[str, dict[str, Any]] = {}
         self._pending_children: dict[str, str] = {}  # child_id -> child_run_id
 
+        # Runtime environment (e.g., "celery", "temporal", None for local)
+        self._runtime: str | None = None
+        self._storage_config: dict[str, Any] | None = None
+
+        # Step failure tracking (for handling failures during replay)
+        self._step_failures: dict[str, dict[str, Any]] = {}
+
         # Replay state if resuming
         if event_log:
             self._is_replaying = True
@@ -134,6 +141,22 @@ class LocalContext(WorkflowContext):
                     "last_error": event.data.get("error", ""),
                 }
 
+            elif event.type == EventType.STEP_FAILED:
+                # Track step failures for distributed step dispatch
+                # If a step failed (not retrying), record it for replay detection
+                step_id = event.data.get("step_id")
+                is_retryable = event.data.get("is_retryable", True)
+                # Only track as a final failure if not retryable or if this is
+                # the last failure before a STEP_COMPLETED event
+                # For now, track all failures - the step decorator will check
+                # if there's also a STEP_COMPLETED event
+                if step_id and not is_retryable:
+                    self._step_failures[step_id] = {
+                        "error": event.data.get("error", "Unknown error"),
+                        "error_type": event.data.get("error_type", "Exception"),
+                        "is_retryable": is_retryable,
+                    }
+
             elif event.type == EventType.CANCELLATION_REQUESTED:
                 self._cancellation_requested = True
                 self._cancellation_reason = event.data.get("reason")
@@ -208,6 +231,28 @@ class LocalContext(WorkflowContext):
         """Set replay mode."""
         self._is_replaying = value
 
+    @property
+    def runtime(self) -> str | None:
+        """
+        Get the runtime environment slug.
+
+        Returns the runtime identifier (e.g., "celery", "temporal") or None
+        for local/inline execution. Used to determine step dispatch behavior.
+
+        Returns:
+            Runtime slug string or None for local execution
+        """
+        return self._runtime
+
+    @property
+    def storage_config(self) -> dict[str, Any] | None:
+        """
+        Get storage configuration for distributed workers.
+
+        This is passed to step workers so they can connect to the same storage.
+        """
+        return self._storage_config
+
     # =========================================================================
     # Step result caching (for @step decorator compatibility)
     # =========================================================================
@@ -255,6 +300,41 @@ class LocalContext(WorkflowContext):
         """Clear retry state for a step."""
         self._retry_states.pop(step_id, None)
 
+    # =========================================================================
+    # Step failure tracking (for distributed step dispatch)
+    # =========================================================================
+
+    def has_step_failed(self, step_id: str) -> bool:
+        """
+        Check if a step has a recorded failure.
+
+        Used during replay to detect steps that failed on a remote worker.
+        """
+        return step_id in self._step_failures
+
+    def get_step_failure(self, step_id: str) -> dict[str, Any] | None:
+        """
+        Get step failure info.
+
+        Returns:
+            Dict with 'error', 'error_type', 'is_retryable' or None
+        """
+        return self._step_failures.get(step_id)
+
+    def record_step_failure(
+        self,
+        step_id: str,
+        error: str,
+        error_type: str,
+        is_retryable: bool,
+    ) -> None:
+        """Record a step failure for replay detection."""
+        self._step_failures[step_id] = {
+            "error": error,
+            "error_type": error_type,
+            "is_retryable": is_retryable,
+        }
+
     # =========================================================================
     # Sleep state management (for @step decorator and EventReplayer compatibility)
     # =========================================================================

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)