pyworkflow-engine 0.1.7__py3-none-any.whl

pyworkflow/celery/app.py

@@ -0,0 +1,198 @@
+"""
+Celery application for distributed workflow execution.
+
+This module configures Celery for:
+- Distributed step execution across workers
+- Automatic retry with exponential backoff
+- Scheduled task execution (sleep resumption)
+- Result persistence
+
+Note: With Python 3.13, you may see "BufferError: Existing exports of data"
+warnings in Celery logs. This is a known compatibility issue between Python 3.13's
+garbage collector and Celery's saferepr module. It does not affect functionality.
+"""
+
+import os
+
+from celery import Celery
+from kombu import Exchange, Queue
+
+from pyworkflow.observability.logging import configure_logging
+
+
+def discover_workflows(modules: list[str] | None = None) -> None:
+    """
+    Discover and import workflow modules to register workflows with Celery workers.
+
+    This function imports Python modules containing workflow definitions so that
+    Celery workers can find and execute them.
+
+    Args:
+        modules: List of module paths to import (e.g., ["myapp.workflows", "myapp.tasks"])
+                If None, reads from PYWORKFLOW_DISCOVER environment variable
+
+    Environment Variables:
+        PYWORKFLOW_DISCOVER: Comma-separated list of modules to import
+                            Example: "myapp.workflows,myapp.tasks,examples.functional.basic_workflow"
+
+    Examples:
+        # Discover from environment variable
+        discover_workflows()
+
+        # Discover specific modules
+        discover_workflows(["myapp.workflows", "myapp.tasks"])
+    """
+    if modules is None:
+        # Read from environment variable
+        discover_env = os.getenv("PYWORKFLOW_DISCOVER", "")
+        if not discover_env:
+            return
+        modules = [m.strip() for m in discover_env.split(",") if m.strip()]
+
+    for module_path in modules:
+        try:
+            __import__(module_path)
+            print(f"✓ Discovered workflows from: {module_path}")
+        except ImportError as e:
+            print(f"✗ Failed to import {module_path}: {e}")
+
+
+def create_celery_app(
+    broker_url: str | None = None,
+    result_backend: str | None = None,
+    app_name: str = "pyworkflow",
+) -> Celery:
+    """
+    Create and configure a Celery application for PyWorkflow.
+
+    Args:
+        broker_url: Celery broker URL. Priority: parameter > PYWORKFLOW_CELERY_BROKER env var > redis://localhost:6379/0
+        result_backend: Result backend URL. Priority: parameter > PYWORKFLOW_CELERY_RESULT_BACKEND env var > redis://localhost:6379/1
+        app_name: Application name
+
+    Returns:
+        Configured Celery application
+
+    Environment Variables:
+        PYWORKFLOW_CELERY_BROKER: Celery broker URL (used if broker_url param not provided)
+        PYWORKFLOW_CELERY_RESULT_BACKEND: Result backend URL (used if result_backend param not provided)
+
+    Examples:
+        # Default configuration (uses env vars if set, otherwise localhost Redis)
+        app = create_celery_app()
+
+        # Custom Redis
+        app = create_celery_app(
+            broker_url="redis://redis-host:6379/0",
+            result_backend="redis://redis-host:6379/1"
+        )
+
+        # RabbitMQ with Redis backend
+        app = create_celery_app(
+            broker_url="amqp://guest:guest@rabbitmq:5672//",
+            result_backend="redis://localhost:6379/1"
+        )
+    """
+    # Priority: parameter > environment variable > hardcoded default
+    broker_url = broker_url or os.getenv("PYWORKFLOW_CELERY_BROKER") or "redis://localhost:6379/0"
+    result_backend = (
+        result_backend
+        or os.getenv("PYWORKFLOW_CELERY_RESULT_BACKEND")
+        or "redis://localhost:6379/1"
+    )
+
+    app = Celery(
+        app_name,
+        broker=broker_url,
+        backend=result_backend,
+        include=[
+            "pyworkflow.celery.tasks",
+        ],
+    )
+
+    # Configure Celery
+    app.conf.update(
+        # Task execution settings
+        task_serializer="json",
+        result_serializer="json",
+        accept_content=["json"],
+        timezone="UTC",
+        enable_utc=True,
+        # Task routing
+        task_default_queue="pyworkflow.default",
+        task_default_exchange="pyworkflow",
+        task_default_exchange_type="topic",
+        task_default_routing_key="workflow.default",
+        # Task queues
+        task_queues=(
+            Queue(
+                "pyworkflow.default",
+                Exchange("pyworkflow", type="topic"),
+                routing_key="workflow.#",
+            ),
+            Queue(
+                "pyworkflow.steps",
+                Exchange("pyworkflow", type="topic"),
+                routing_key="workflow.step.#",
+            ),
+            Queue(
+                "pyworkflow.workflows",
+                Exchange("pyworkflow", type="topic"),
+                routing_key="workflow.workflow.#",
+            ),
+            Queue(
+                "pyworkflow.schedules",
+                Exchange("pyworkflow", type="topic"),
+                routing_key="workflow.schedule.#",
+            ),
+        ),
+        # Result backend settings
+        result_expires=3600,  # 1 hour
+        result_persistent=True,
+        # Task execution
+        task_acks_late=True,
+        task_reject_on_worker_lost=True,
+        worker_prefetch_multiplier=1,  # Fair task distribution
+        # Retry settings
+        task_autoretry_for=(Exception,),
+        task_retry_backoff=True,
+        task_retry_backoff_max=600,  # 10 minutes max
+        task_retry_jitter=True,
+        # Monitoring
+        worker_send_task_events=True,
+        task_send_sent_event=True,
+        # Beat scheduler (for sleep resumption)
+        beat_schedule={},
+        # Logging
+        worker_log_format="[%(asctime)s: %(levelname)s/%(processName)s] %(message)s",
+        worker_task_log_format="[%(asctime)s: %(levelname)s/%(processName)s] [%(task_name)s(%(task_id)s)] %(message)s",
+    )
+
+    # Configure logging
+    configure_logging(level="INFO")
+
+    # Auto-discover workflows from environment variable or configured modules
+    discover_workflows()
+
+    return app
+
+
+# Global Celery app instance
+# Can be customized by calling create_celery_app() with custom config
+celery_app = create_celery_app()
+
+
+def get_celery_app() -> Celery:
+    """
+    Get the global Celery application instance.
+
+    Returns:
+        Celery application
+
+    Example:
+        from pyworkflow.celery.app import get_celery_app
+
+        app = get_celery_app()
+        app.conf.update(broker_url="redis://custom:6379/0")
+    """
+    return celery_app

pyworkflow/celery/scheduler.py

@@ -0,0 +1,315 @@
+"""
+Custom Celery Beat scheduler for PyWorkflow schedules.
+
+This scheduler integrates with PyWorkflow's storage backend to dynamically
+load and execute scheduled workflows without requiring Beat to be restarted
+when schedules change.
+
+Usage:
+    celery -A pyworkflow.celery.app beat \\
+        --scheduler pyworkflow.celery.scheduler:PyWorkflowScheduler \\
+        --loglevel INFO
+
+The scheduler:
+1. Polls storage for due schedules every sync_interval seconds
+2. Creates Celery tasks for each due schedule
+3. Handles overlap policies
+4. Updates schedule metadata after runs
+"""
+
+import asyncio
+from datetime import UTC, datetime
+from typing import Any
+
+from celery.beat import Scheduler
+from loguru import logger
+
+from pyworkflow.storage.base import StorageBackend
+from pyworkflow.storage.config import config_to_storage
+from pyworkflow.storage.schemas import OverlapPolicy, Schedule
+
+
+class PyWorkflowScheduler(Scheduler):
+    """
+    Custom Celery Beat scheduler that reads schedules from PyWorkflow storage.
+
+    This scheduler:
+    1. Polls storage for due schedules every sync_interval
+    2. Creates Celery tasks for each due schedule
+    3. Handles overlap policies (skip, buffer, cancel, allow)
+    4. Updates schedule metadata after runs
+
+    Configuration:
+        The scheduler reads configuration from environment variables:
+        - PYWORKFLOW_STORAGE_BACKEND: Storage backend type (file, memory)
+        - PYWORKFLOW_STORAGE_PATH: Path for file storage backend
+
+    Example:
+        celery -A pyworkflow.celery.app beat \\
+            --scheduler pyworkflow.celery.scheduler:PyWorkflowScheduler
+    """
+
+    #: How often to check for due schedules (seconds)
+    sync_interval = 5.0
+
+    def __init__(self, *args: Any, storage_config: dict[str, Any] | None = None, **kwargs: Any):
+        """
+        Initialize the scheduler.
+
+        Args:
+            storage_config: Storage backend configuration dict
+        """
+        super().__init__(*args, **kwargs)
+        self._storage_config = storage_config
+        self._storage: StorageBackend | None = None
+        self._last_schedule_check: datetime | None = None
+        self._initialized = False
+
+    def setup_schedule(self) -> None:
+        """Initialize the scheduler."""
+        super().setup_schedule()
+
+        # Activate any @scheduled_workflow decorated functions
+        self._activate_decorated_schedules()
+
+        logger.info("PyWorkflow scheduler initialized")
+        self._initialized = True
+
+    def _activate_decorated_schedules(self) -> None:
+        """Activate all @scheduled_workflow decorated functions."""
+        try:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                loop.run_until_complete(self._do_activate_schedules())
+            finally:
+                loop.close()
+        except Exception as e:
+            logger.error(f"Error activating decorated schedules: {e}")
+
+    async def _do_activate_schedules(self) -> None:
+        """Actually activate the decorated schedules."""
+        from pyworkflow.core.scheduled import activate_scheduled_workflows
+
+        storage = self._get_storage()
+        if storage is None:
+            logger.warning("Storage not configured, cannot activate decorated schedules")
+            return
+
+        try:
+            schedule_ids = await activate_scheduled_workflows(storage=storage)
+            if schedule_ids:
+                logger.info(f"Activated {len(schedule_ids)} decorated schedule(s): {schedule_ids}")
+            else:
+                logger.debug("No decorated schedules to activate")
+        except Exception as e:
+            logger.error(f"Failed to activate decorated schedules: {e}")
+
+    def tick(self) -> float:
+        """
+        Called by Celery Beat on each tick.
+
+        Checks for due schedules and triggers them.
+
+        Returns:
+            Seconds until next tick
+        """
+        # Call parent tick to handle existing celery beat entries
+        remaining = super().tick()
+
+        # Check for due schedules
+        now = datetime.now(UTC)
+        if (
+            self._last_schedule_check is None
+            or (now - self._last_schedule_check).total_seconds() >= self.sync_interval
+        ):
+            self._sync_schedules()
+            self._last_schedule_check = now
+
+        # Return the smaller of the two intervals
+        return min(remaining, self.sync_interval)
+
+    def _sync_schedules(self) -> None:
+        """Sync schedules from storage and trigger due ones."""
+        try:
+            # Run async code in sync context
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                loop.run_until_complete(self._process_due_schedules())
+            finally:
+                loop.close()
+        except Exception as e:
+            logger.error(f"Error syncing schedules: {e}")
+
+    async def _process_due_schedules(self) -> None:
+        """Process all schedules that are due to run."""
+        from pyworkflow.celery.tasks import execute_scheduled_workflow_task
+        from pyworkflow.utils.schedule import calculate_next_run_time
+
+        storage = self._get_storage()
+        if storage is None:
+            logger.warning("Storage not configured, skipping schedule sync")
+            return
+
+        now = datetime.now(UTC)
+
+        try:
+            due_schedules = await storage.get_due_schedules(now)
+        except Exception as e:
+            logger.error(f"Failed to get due schedules: {e}")
+            return
+
+        for schedule in due_schedules:
+            try:
+                should_run, reason = await self._check_overlap_policy(schedule, storage)
+
+                if should_run:
+                    # Calculate next run time before triggering
+                    next_run_time = calculate_next_run_time(
+                        schedule.spec,
+                        last_run=schedule.next_run_time,
+                        now=now,
+                    )
+
+                    # Update next_run_time immediately to prevent duplicate triggers
+                    schedule.next_run_time = next_run_time
+                    schedule.updated_at = datetime.now(UTC)
+                    await storage.update_schedule(schedule)
+
+                    # Trigger the scheduled workflow task
+                    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(),
+                            "storage_config": self._storage_config,
+                        },
+                        queue="pyworkflow.schedules",
+                    )
+
+                    logger.info(
+                        f"Triggered scheduled workflow: {schedule.workflow_name}",
+                        schedule_id=schedule.schedule_id,
+                        next_run=next_run_time.isoformat() if next_run_time else None,
+                    )
+                else:
+                    # Record skip and update next run time
+                    schedule.skipped_runs += 1
+                    next_run_time = calculate_next_run_time(
+                        schedule.spec,
+                        last_run=schedule.next_run_time,
+                        now=now,
+                    )
+                    schedule.next_run_time = next_run_time
+                    schedule.updated_at = datetime.now(UTC)
+                    await storage.update_schedule(schedule)
+
+                    logger.info(
+                        f"Skipped scheduled workflow: {schedule.workflow_name} ({reason})",
+                        schedule_id=schedule.schedule_id,
+                    )
+
+            except Exception as e:
+                logger.error(
+                    f"Error processing schedule {schedule.schedule_id}: {e}",
+                )
+
+    async def _check_overlap_policy(
+        self,
+        schedule: Schedule,
+        storage: StorageBackend,
+    ) -> tuple[bool, str | None]:
+        """
+        Check if schedule should run based on overlap policy.
+
+        Args:
+            schedule: The schedule to check
+            storage: Storage backend
+
+        Returns:
+            Tuple of (should_run, reason_if_not)
+        """
+        # No running runs means we can always run
+        if not schedule.running_run_ids:
+            return True, None
+
+        policy = schedule.overlap_policy
+
+        if policy == OverlapPolicy.ALLOW_ALL:
+            return True, None
+
+        elif policy == OverlapPolicy.SKIP:
+            return False, "Previous run still active (SKIP policy)"
+
+        elif policy == OverlapPolicy.BUFFER_ONE:
+            if schedule.buffered_count >= 1:
+                return False, "Buffer full (BUFFER_ONE policy)"
+            # Increment buffer count
+            schedule.buffered_count += 1
+            await storage.update_schedule(schedule)
+            return True, None
+
+        elif policy == OverlapPolicy.BUFFER_ALL:
+            # Always allow, buffered_count tracks pending runs
+            schedule.buffered_count += 1
+            await storage.update_schedule(schedule)
+            return True, None
+
+        elif policy == OverlapPolicy.CANCEL_OTHER:
+            # Cancel running runs
+            from pyworkflow.primitives.cancel import cancel_workflow
+
+            for run_id in schedule.running_run_ids:
+                try:
+                    await cancel_workflow(
+                        run_id,
+                        reason="Cancelled by CANCEL_OTHER overlap policy",
+                        storage=storage,
+                    )
+                except Exception as e:
+                    logger.warning(f"Failed to cancel run {run_id}: {e}")
+
+            # Clear the running runs list
+            schedule.running_run_ids = []
+            await storage.update_schedule(schedule)
+            return True, None
+
+        # Default: allow
+        return True, None
+
+    def _get_storage(self) -> StorageBackend | None:
+        """Get or create the storage backend."""
+        if self._storage is not None:
+            return self._storage
+
+        # Try to get storage config from environment or defaults
+        if self._storage_config:
+            self._storage = config_to_storage(self._storage_config)
+            return self._storage
+
+        # Try default file storage
+        import os
+
+        storage_type = os.getenv("PYWORKFLOW_STORAGE_BACKEND", "file")
+        storage_path = os.getenv("PYWORKFLOW_STORAGE_PATH", "./pyworkflow_data")
+
+        if storage_type == "file":
+            from pyworkflow.storage.file import FileStorageBackend
+
+            self._storage = FileStorageBackend(storage_path)
+        elif storage_type == "memory":
+            from pyworkflow.storage.memory import InMemoryStorageBackend
+
+            self._storage = InMemoryStorageBackend()
+        else:
+            logger.warning(f"Unknown storage type: {storage_type}")
+            return None
+
+        return self._storage
+
+    @property
+    def info(self) -> str:
+        """Return scheduler info string."""
+        return f"PyWorkflowScheduler (sync_interval={self.sync_interval}s)"