stabilize 0.9.2__py3-none-any.whl

stabilize/handlers/run_task.py

@@ -0,0 +1,369 @@
+"""
+RunTaskHandler - executes tasks.
+
+This is the handler that actually runs task implementations.
+It handles execution, retries, timeouts, and result processing.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from stabilize.handlers.base import StabilizeHandler
+from stabilize.models.status import WorkflowStatus
+from stabilize.queue.messages import (
+    CompleteTask,
+    PauseTask,
+    RunTask,
+)
+from stabilize.tasks.interface import RetryableTask, Task
+from stabilize.tasks.registry import TaskNotFoundError, TaskRegistry
+from stabilize.tasks.result import TaskResult
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+    from stabilize.models.task import TaskExecution
+    from stabilize.persistence.store import WorkflowStore
+    from stabilize.queue.queue import Queue
+
+logger = logging.getLogger(__name__)
+
+
+class TimeoutError(Exception):
+    """Raised when a task times out."""
+
+    pass
+
+
+class RunTaskHandler(StabilizeHandler[RunTask]):
+    """
+    Handler for RunTask messages.
+
+    This is where tasks are actually executed. The handler:
+    1. Resolves the task implementation
+    2. Checks for cancellation/pause
+    3. Checks for timeout
+    4. Executes the task
+    5. Processes the result
+    """
+
+    def __init__(
+        self,
+        queue: Queue,
+        repository: WorkflowStore,
+        task_registry: TaskRegistry,
+        retry_delay: timedelta = timedelta(seconds=15),
+    ) -> None:
+        super().__init__(queue, repository, retry_delay)
+        self.task_registry = task_registry
+
+    @property
+    def message_type(self) -> type[RunTask]:
+        return RunTask
+
+    def handle(self, message: RunTask) -> None:
+        """Handle the RunTask message."""
+
+        def on_task(stage: StageExecution, task_model: TaskExecution) -> None:
+            execution = stage.execution
+
+            # Resolve task implementation
+            try:
+                task = self._resolve_task(message.task_type, task_model)
+            except TaskNotFoundError as e:
+                logger.error("Task type not found: %s", message.task_type)
+                self._complete_with_error(stage, task_model, message, str(e))
+                return
+
+            # Check execution state
+            if execution.is_canceled:
+                self._handle_cancellation(stage, task_model, task, message)
+                return
+
+            if execution.status.is_complete:
+                self.queue.push(
+                    CompleteTask(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                        task_id=message.task_id,
+                        status=WorkflowStatus.CANCELED,
+                    )
+                )
+                return
+
+            if execution.status == WorkflowStatus.PAUSED:
+                self.queue.push(
+                    PauseTask(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                        task_id=message.task_id,
+                    )
+                )
+                return
+
+            # Check for manual skip
+            if stage.context.get("manualSkip"):
+                self.queue.push(
+                    CompleteTask(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                        task_id=message.task_id,
+                        status=WorkflowStatus.SKIPPED,
+                    )
+                )
+                return
+
+            # Check for timeout
+            try:
+                self._check_for_timeout(stage, task_model, task)
+            except TimeoutError as e:
+                logger.info("Task %s timed out: %s", task_model.name, e)
+                result = task.on_timeout(stage) if hasattr(task, "on_timeout") else None
+                if result is None:
+                    self._complete_with_error(stage, task_model, message, str(e))
+                    return
+                else:
+                    self._process_result(stage, task_model, result, message)
+                    return
+
+            # Execute the task
+            try:
+                result = task.execute(stage)
+                self._process_result(stage, task_model, result, message)
+            except Exception as e:
+                logger.error(
+                    "Error executing task %s: %s",
+                    task_model.name,
+                    e,
+                    exc_info=True,
+                )
+                self._handle_exception(stage, task_model, task, message, e)
+
+        self.with_task(message, on_task)
+
+    def _resolve_task(
+        self,
+        task_type: str,
+        task_model: TaskExecution,
+    ) -> Task:
+        """Resolve the task implementation."""
+        # Try by class name first
+        try:
+            return self.task_registry.get_by_class(task_model.implementing_class)
+        except TaskNotFoundError:
+            pass
+
+        # Try by type name
+        return self.task_registry.get(task_type)
+
+    def _check_for_timeout(
+        self,
+        stage: StageExecution,
+        task_model: TaskExecution,
+        task: Task,
+    ) -> None:
+        """Check if task has timed out."""
+        if not isinstance(task, RetryableTask):
+            return
+
+        if task_model.start_time is None:
+            return
+
+        start_time = task_model.start_time
+        current_time = self.current_time_millis()
+        elapsed = timedelta(milliseconds=current_time - start_time)
+
+        # Get timeout (potentially dynamic)
+        timeout = task.get_dynamic_timeout(stage)
+
+        # Account for paused time
+        paused_duration = timedelta(milliseconds=stage.execution.paused_duration_relative_to(start_time))
+        actual_elapsed = elapsed - paused_duration
+
+        if actual_elapsed > timeout:
+            raise TimeoutError(f"Task timed out after {actual_elapsed} (timeout: {timeout})")
+
+    def _process_result(
+        self,
+        stage: StageExecution,
+        task_model: TaskExecution,
+        result: TaskResult,
+        message: RunTask,
+    ) -> None:
+        """Process a task result."""
+        # Store outputs in stage
+        if result.context:
+            stage.context.update(result.context)
+        if result.outputs:
+            stage.outputs.update(result.outputs)
+
+        # Handle based on status
+        if result.status == WorkflowStatus.RUNNING:
+            # Task needs to be re-executed
+            self.repository.store_stage(stage)
+
+            # Calculate backoff
+            delay = self._get_backoff_period(stage, task_model, message)
+            self.queue.push(message, delay)
+
+            logger.debug(
+                "Task %s still running, re-queuing with %s delay",
+                task_model.name,
+                delay,
+            )
+
+        elif result.status in {
+            WorkflowStatus.SUCCEEDED,
+            WorkflowStatus.REDIRECT,
+            WorkflowStatus.SKIPPED,
+            WorkflowStatus.FAILED_CONTINUE,
+            WorkflowStatus.STOPPED,
+        }:
+            self.repository.store_stage(stage)
+            self.queue.push(
+                CompleteTask(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=message.stage_id,
+                    task_id=message.task_id,
+                    status=result.status,
+                )
+            )
+
+        elif result.status == WorkflowStatus.CANCELED:
+            self.repository.store_stage(stage)
+            status = stage.failure_status(default=result.status)
+            self.queue.push(
+                CompleteTask(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=message.stage_id,
+                    task_id=message.task_id,
+                    status=status,
+                    original_status=result.status,
+                )
+            )
+
+        elif result.status == WorkflowStatus.TERMINAL:
+            self.repository.store_stage(stage)
+            status = stage.failure_status(default=result.status)
+            self.queue.push(
+                CompleteTask(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=message.stage_id,
+                    task_id=message.task_id,
+                    status=status,
+                    original_status=result.status,
+                )
+            )
+
+        else:
+            logger.warning("Unhandled task status: %s", result.status)
+            self.repository.store_stage(stage)
+
+    def _get_backoff_period(
+        self,
+        stage: StageExecution,
+        task_model: TaskExecution,
+        message: RunTask,
+    ) -> timedelta:
+        """Calculate backoff period for retry."""
+        # Try to get the task and use its backoff
+        try:
+            task = self._resolve_task(message.task_type, task_model)
+            if isinstance(task, RetryableTask):
+                elapsed = timedelta(milliseconds=self.current_time_millis() - (task_model.start_time or 0))
+                return task.get_dynamic_backoff_period(stage, elapsed)
+        except TaskNotFoundError:
+            pass
+
+        # Default backoff
+        return timedelta(seconds=1)
+
+    def _handle_cancellation(
+        self,
+        stage: StageExecution,
+        task_model: TaskExecution,
+        task: Task,
+        message: RunTask,
+    ) -> None:
+        """Handle execution cancellation."""
+        result = task.on_cancel(stage) if hasattr(task, "on_cancel") else None
+        if result:
+            stage.context.update(result.context)
+            stage.outputs.update(result.outputs)
+            self.repository.store_stage(stage)
+
+        self.queue.push(
+            CompleteTask(
+                execution_type=message.execution_type,
+                execution_id=message.execution_id,
+                stage_id=message.stage_id,
+                task_id=message.task_id,
+                status=WorkflowStatus.CANCELED,
+            )
+        )
+
+    def _handle_exception(
+        self,
+        stage: StageExecution,
+        task_model: TaskExecution,
+        task: Task,
+        message: RunTask,
+        exception: Exception,
+    ) -> None:
+        """Handle task execution exception."""
+        # TODO: Check if exception is retryable
+        # For now, treat all exceptions as terminal
+
+        exception_details = {
+            "details": {
+                "error": str(exception),
+                "errors": [str(exception)],
+            }
+        }
+
+        stage.context["exception"] = exception_details
+        task_model.task_exception_details["exception"] = exception_details
+        self.repository.store_stage(stage)
+
+        status = stage.failure_status(default=WorkflowStatus.TERMINAL)
+        self.queue.push(
+            CompleteTask(
+                execution_type=message.execution_type,
+                execution_id=message.execution_id,
+                stage_id=message.stage_id,
+                task_id=message.task_id,
+                status=status,
+                original_status=WorkflowStatus.TERMINAL,
+            )
+        )
+
+    def _complete_with_error(
+        self,
+        stage: StageExecution,
+        task_model: TaskExecution,
+        message: RunTask,
+        error: str,
+    ) -> None:
+        """Complete task with an error."""
+        stage.context["exception"] = {
+            "details": {"error": error},
+        }
+        self.repository.store_stage(stage)
+
+        self.queue.push(
+            CompleteTask(
+                execution_type=message.execution_type,
+                execution_id=message.execution_id,
+                stage_id=message.stage_id,
+                task_id=message.task_id,
+                status=WorkflowStatus.TERMINAL,
+            )
+        )

stabilize/handlers/start_stage.py

@@ -0,0 +1,262 @@
+"""
+StartStageHandler - handles stage startup.
+
+This is one of the most critical handlers in the execution engine.
+It checks if upstream stages are complete, plans the stage's tasks
+and synthetic stages, and starts execution.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from stabilize.handlers.base import StabilizeHandler
+from stabilize.models.status import WorkflowStatus
+from stabilize.queue.messages import (
+    CompleteStage,
+    CompleteWorkflow,
+    SkipStage,
+    StartStage,
+    StartTask,
+)
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+
+logger = logging.getLogger(__name__)
+
+
+class StartStageHandler(StabilizeHandler[StartStage]):
+    """
+    Handler for StartStage messages.
+
+    Execution flow:
+    1. Check if any upstream stages failed -> CompleteWorkflow
+    2. Check if all upstream stages complete
+       - If not: Re-queue with retry delay
+       - If yes: Continue to step 3
+    3. Check if stage should be skipped -> SkipStage
+    4. Check if start time expired -> SkipStage
+    5. Plan the stage (build tasks and before stages)
+    6. Start the stage:
+       - If has before stages: StartStage for each
+       - Else if has tasks: StartTask for first task
+       - Else: CompleteStage
+    """
+
+    @property
+    def message_type(self) -> type[StartStage]:
+        return StartStage
+
+    def handle(self, message: StartStage) -> None:
+        """Handle the StartStage message."""
+
+        def on_stage(stage: StageExecution) -> None:
+            try:
+                # Check if upstream stages have failed
+                if stage.any_upstream_stages_failed():
+                    logger.warning(
+                        "Upstream stage failed for %s (%s), completing execution",
+                        stage.name,
+                        stage.id,
+                    )
+                    self.queue.push(
+                        CompleteWorkflow(
+                            execution_type=message.execution_type,
+                            execution_id=message.execution_id,
+                        )
+                    )
+                    return
+
+                # Check if all upstream stages are complete
+                if stage.all_upstream_stages_complete():
+                    self._start_if_ready(stage, message)
+                else:
+                    # Upstream not complete - re-queue
+                    logger.debug(
+                        "Re-queuing %s (%s) - upstream stages not complete",
+                        stage.name,
+                        stage.id,
+                    )
+                    self.queue.push(message, self.retry_delay)
+
+            except Exception as e:
+                logger.error(
+                    "Error starting stage %s (%s): %s",
+                    stage.name,
+                    stage.id,
+                    e,
+                    exc_info=True,
+                )
+                stage.context["exception"] = {
+                    "details": {"error": str(e)},
+                }
+                stage.context["beforeStagePlanningFailed"] = True
+                self.repository.store_stage(stage)
+                self.queue.push(
+                    CompleteStage(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=message.stage_id,
+                    )
+                )
+
+        self.with_stage(message, on_stage)
+
+    def _start_if_ready(
+        self,
+        stage: StageExecution,
+        message: StartStage,
+    ) -> None:
+        """Start the stage if it's ready to run."""
+        # Check if already processed
+        if stage.status != WorkflowStatus.NOT_STARTED:
+            logger.warning(
+                "Ignoring StartStage for %s - already %s",
+                stage.name,
+                stage.status,
+            )
+            return
+
+        # Check if should skip
+        if self._should_skip(stage):
+            logger.info("Skipping optional stage %s", stage.name)
+            self.queue.push(
+                SkipStage(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=message.stage_id,
+                )
+            )
+            return
+
+        # Check if start time expired
+        if self._is_after_start_time_expiry(stage):
+            logger.warning("Stage %s start time expired, skipping", stage.name)
+            self.queue.push(
+                SkipStage(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=message.stage_id,
+                )
+            )
+            return
+
+        # Plan and start the stage
+        stage.start_time = self.current_time_millis()
+        self._plan_stage(stage)
+        stage.status = WorkflowStatus.RUNNING
+        self.repository.store_stage(stage)
+
+        self._start_stage(stage, message)
+
+        logger.info("Started stage %s (%s)", stage.name, stage.id)
+
+    def _plan_stage(self, stage: StageExecution) -> None:
+        """
+        Plan the stage - build tasks and before stages.
+
+        This is where StageDefinitionBuilder.buildTasks() and
+        buildBeforeStages() would be called.
+
+        For now, we assume tasks are already defined on the stage.
+        """
+        # Mark first and last tasks
+        if stage.tasks:
+            stage.tasks[0].stage_start = True
+            stage.tasks[-1].stage_end = True
+
+        # TODO: Call stage definition builder to:
+        # 1. Build tasks
+        # 2. Build before stages
+        # 3. Add context flags
+
+    def _start_stage(
+        self,
+        stage: StageExecution,
+        message: StartStage,
+    ) -> None:
+        """
+        Start stage execution.
+
+        Order of execution:
+        1. Before stages (if any)
+        2. Tasks (if any)
+        3. After stages (if any)
+        4. Complete stage
+        """
+        # Check for before stages
+        before_stages = stage.first_before_stages()
+        if before_stages:
+            for before in before_stages:
+                self.queue.push(
+                    StartStage(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=before.id,
+                    )
+                )
+            return
+
+        # No before stages - start first task
+        first_task = stage.first_task()
+        if first_task:
+            self.queue.push(
+                StartTask(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=message.stage_id,
+                    task_id=first_task.id,
+                )
+            )
+            return
+
+        # No tasks - check for after stages
+        after_stages = stage.first_after_stages()
+        if after_stages:
+            for after in after_stages:
+                self.queue.push(
+                    StartStage(
+                        execution_type=message.execution_type,
+                        execution_id=message.execution_id,
+                        stage_id=after.id,
+                    )
+                )
+            return
+
+        # No tasks or synthetic stages - complete immediately
+        self.queue.push(
+            CompleteStage(
+                execution_type=message.execution_type,
+                execution_id=message.execution_id,
+                stage_id=message.stage_id,
+            )
+        )
+
+    def _should_skip(self, stage: StageExecution) -> bool:
+        """
+        Check if stage should be skipped.
+
+        Checks the 'stageEnabled' context for conditional execution.
+        """
+        stage_enabled = stage.context.get("stageEnabled")
+        if stage_enabled is None:
+            return False
+
+        if isinstance(stage_enabled, dict):
+            expr_type = stage_enabled.get("type")
+            if expr_type == "expression":
+                # TODO: Evaluate expression
+                expression = stage_enabled.get("expression", "true")
+                # For now, just check if it's explicitly "false"
+                if isinstance(expression, str):
+                    return expression.lower() == "false"
+
+        return False
+
+    def _is_after_start_time_expiry(self, stage: StageExecution) -> bool:
+        """Check if current time is past start time expiry."""
+        if stage.start_time_expiry is None:
+            return False
+        return self.current_time_millis() > stage.start_time_expiry

stabilize/handlers/start_task.py

@@ -0,0 +1,74 @@
+"""
+StartTaskHandler - handles task startup.
+
+This handler prepares a task for execution and triggers RunTask.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from stabilize.handlers.base import StabilizeHandler
+from stabilize.models.status import WorkflowStatus
+from stabilize.queue.messages import (
+    RunTask,
+    StartTask,
+)
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+    from stabilize.models.task import TaskExecution
+
+logger = logging.getLogger(__name__)
+
+
+class StartTaskHandler(StabilizeHandler[StartTask]):
+    """
+    Handler for StartTask messages.
+
+    Execution flow:
+    1. Check if task is enabled (SkippableTask)
+       - If not: Push CompleteTask(SKIPPED)
+    2. Set task status to RUNNING
+    3. Set task start time
+    4. Push RunTask
+    """
+
+    @property
+    def message_type(self) -> type[StartTask]:
+        return StartTask
+
+    def handle(self, message: StartTask) -> None:
+        """Handle the StartTask message."""
+
+        def on_task(stage: StageExecution, task: TaskExecution) -> None:
+            # Check if task should be skipped
+            # TODO: Check SkippableTask.isEnabled()
+
+            # Update task status
+            task.status = WorkflowStatus.RUNNING
+            task.start_time = self.current_time_millis()
+
+            # Save the stage (which includes the task)
+            self.repository.store_stage(stage)
+
+            # Push RunTask
+            self.queue.push(
+                RunTask(
+                    execution_type=message.execution_type,
+                    execution_id=message.execution_id,
+                    stage_id=message.stage_id,
+                    task_id=message.task_id,
+                    task_type=task.implementing_class,
+                )
+            )
+
+            logger.debug(
+                "Started task %s (%s) in stage %s",
+                task.name,
+                task.id,
+                stage.name,
+            )
+
+        self.with_task(message, on_task)