pysfi 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

sfi/workflowengine/workflowengine.py

@@ -3,6 +3,13 @@
 This module provides a comprehensive workflow engine for managing
 complex task dependencies with support for I/O tasks, CPU-intensive tasks,
 serial tasks, and parallel task execution.
+
+The engine supports:
+- Dependency management with cycle detection
+- Topological sorting for execution order
+- Concurrent execution with configurable limits
+- Error handling and timeout management
+- Execution monitoring and reporting
 """
 
 from __future__ import annotations
@@ -49,7 +56,13 @@ class TaskResult:
 class Task(ABC):
     """Task abstract base class"""
 
-    def __init__(self, task_id: str, task_type: TaskType, dependencies: list[str] | None = None, timeout: float = 30.0):
+    def __init__(
+        self,
+        task_id: str,
+        task_type: TaskType,
+        dependencies: list[str] | None = None,
+        timeout: float = 30.0,
+    ):
         self.task_id = task_id
         self.task_type = task_type
         self.dependencies = dependencies or []
@@ -82,18 +95,68 @@ class Task(ABC):
             return self.end_time - self.start_time
         return 0.0
 
+    async def _execute_with_error_handling(
+        self, execution_func, context: dict[str, TaskResult]
+    ):
+        """Common execution wrapper with error handling"""
+        self.start_time = time.time()
+        self.update_status(TaskStatus.RUNNING)
+
+        try:
+            data = await asyncio.wait_for(execution_func(context), timeout=self.timeout)
+            self.end_time = time.time()
+            self.result = TaskResult(
+                task_id=self.task_id,
+                success=True,
+                data=data,
+                execution_time=self.get_execution_time(),
+            )
+            self.update_status(TaskStatus.COMPLETED)
+            return self.result
+        except asyncio.TimeoutError as e:
+            self.end_time = time.time()
+            self.result = TaskResult(
+                task_id=self.task_id,
+                success=False,
+                data=None,
+                execution_time=self.get_execution_time(),
+                error=e,
+            )
+            self.update_status(TaskStatus.FAILED)
+            raise
+        except Exception as e:
+            self.end_time = time.time()
+            self.result = TaskResult(
+                task_id=self.task_id,
+                success=False,
+                data=None,
+                execution_time=self.get_execution_time(),
+                error=e,
+            )
+            self.update_status(TaskStatus.FAILED)
+            raise
+
 
 class IOTask(Task):
     """I/O-intensive task"""
 
-    def __init__(self, task_id: str, duration: float, dependencies: list[str] | None = None, timeout: float = 30.0):
+    def __init__(
+        self,
+        task_id: str,
+        duration: float,
+        dependencies: list[str] | None = None,
+        timeout: float = 30.0,
+    ):
         super().__init__(task_id, TaskType.ASYNC, dependencies, timeout)
         self.duration = duration
 
     async def execute(self, context: dict[str, TaskResult]) -> Any:
         """Simulate I/O operation"""
-        print(f"[IO] Starting task {self.task_id}, estimated duration: {self.duration}s")
-        return await asyncio.wait_for(self._execute_io(context), timeout=self.timeout)
+        print(
+            f"[IO] Starting task {self.task_id}, estimated duration: {self.duration}s"
+        )
+        result = await self._execute_with_error_handling(self._execute_io, context)
+        return result.data
 
     async def _execute_io(self, context: dict[str, TaskResult]) -> Any:
         """Internal I/O execution method"""
@@ -104,25 +167,34 @@ class IOTask(Task):
 class CPUTask(Task):
     """CPU-intensive task"""
 
-    def __init__(self, task_id: str, iterations: int, dependencies: list[str] | None = None, timeout: float = 30.0):
+    def __init__(
+        self,
+        task_id: str,
+        iterations: int,
+        dependencies: list[str] | None = None,
+        timeout: float = 30.0,
+    ):
         super().__init__(task_id, TaskType.CPU, dependencies, timeout)
         self.iterations = iterations
 
     async def execute(self, context: dict[str, TaskResult]) -> Any:
         """CPU-intensive computation task"""
         print(f"[CPU] Starting task {self.task_id}, iterations: {self.iterations}")
+        result = await self._execute_with_error_handling(self._execute_cpu, context)
+        return result.data
+
+    async def _execute_cpu(self, context: dict[str, TaskResult]) -> Any:
+        """Execute CPU-intensive work in thread pool"""
 
-        # Move CPU-intensive task to thread pool to avoid blocking event loop
         def cpu_intensive_work():
             result = 0
             for i in range(self.iterations):
                 result += i * i
             return result
 
-        # Use asyncio.wait_for with timeout to prevent infinite hangs
-        # Use run_in_executor for Python 3.8 compatibility (asyncio.to_thread is Python 3.9+)
+        # Use run_in_executor to avoid blocking event loop
         loop = asyncio.get_event_loop()
-        result = await asyncio.wait_for(loop.run_in_executor(None, cpu_intensive_work), timeout=self.timeout)
+        result = await loop.run_in_executor(None, cpu_intensive_work)
         return f"CPU task {self.task_id} completed, result: {result}"
 
 
@@ -130,7 +202,11 @@ class SerialTask(Task):
     """Serial task (stateful, must execute sequentially)"""
 
     def __init__(
-        self, task_id: str, process_func: Callable, dependencies: list[str] | None = None, timeout: float = 30.0
+        self,
+        task_id: str,
+        process_func: Callable,
+        dependencies: list[str] | None = None,
+        timeout: float = 30.0,
     ):
         super().__init__(task_id, TaskType.SERIAL, dependencies, timeout)
         self.process_func = process_func
@@ -139,7 +215,11 @@ class SerialTask(Task):
     async def execute(self, context: dict[str, TaskResult]) -> Any:
         """Execute serial task"""
         print(f"[Serial] Starting serial task {self.task_id}")
+        result = await self._execute_with_error_handling(self._execute_serial, context)
+        return result.data
 
+    async def _execute_serial(self, context: dict[str, TaskResult]) -> Any:
+        """Execute serial task logic"""
         # Collect results from dependent tasks
         inputs = {dep_id: context[dep_id].data for dep_id in self.dependencies}
 
@@ -172,61 +252,71 @@ class ParallelTask(Task):
 
     async def execute(self, context: dict[str, TaskResult]) -> Any:
         """Execute subtasks in parallel"""
-        print(f"[Parallel] Starting parallel task {self.task_id}, contains {len(self.subtasks)} subtasks")
-
+        print(
+            f"[Parallel] Starting parallel task {self.task_id}, contains {len(self.subtasks)} subtasks"
+        )
+        result = await self._execute_with_error_handling(
+            self._execute_parallel, context
+        )
+        return result.data
+
+    async def _execute_parallel(self, context: dict[str, TaskResult]) -> Any:
+        """Execute subtasks in parallel with controlled concurrency"""
         # Create semaphore to control concurrency
         semaphore = asyncio.Semaphore(self.max_concurrent)
 
-        async def execute_subtask(subtask: Task, sem: asyncio.Semaphore):
-            async with sem:
-                subtask.start_time = time.time()
-                subtask.update_status(TaskStatus.RUNNING)
-
+        async def execute_subtask_with_semaphore(subtask: Task):
+            async with semaphore:
+                # Execute subtask with its own context
                 try:
-                    data = await asyncio.wait_for(subtask.execute(context), timeout=subtask.timeout)
+                    data = await asyncio.wait_for(
+                        subtask.execute(context), timeout=subtask.timeout
+                    )
                     subtask.result = TaskResult(
                         task_id=subtask.task_id,
                         success=True,
                         data=data,
-                        execution_time=time.time() - subtask.start_time,
+                        execution_time=time.time() - subtask.start_time
+                        if subtask.start_time
+                        else 0,
                     )
                     subtask.update_status(TaskStatus.COMPLETED)
                     return subtask.result
-                except asyncio.TimeoutError as e:
-                    error = TimeoutError(f"Task {subtask.task_id} execution timeout")
-                    subtask.result = TaskResult(
-                        task_id=subtask.task_id,
-                        success=False,
-                        data=None,
-                        execution_time=time.time() - subtask.start_time,
-                        error=error,
-                    )
-                    subtask.update_status(TaskStatus.FAILED)
-                    raise error from e
                 except Exception as e:
+                    subtask.end_time = time.time()
                     subtask.result = TaskResult(
                         task_id=subtask.task_id,
                         success=False,
                         data=None,
-                        execution_time=time.time() - subtask.start_time,
+                        execution_time=time.time() - subtask.start_time
+                        if subtask.start_time
+                        else 0,
                         error=e,
                     )
                     subtask.update_status(TaskStatus.FAILED)
-                    raise e
+                    return subtask.result
 
         # Execute all subtasks in parallel
-        tasks = [execute_subtask(subtask, semaphore) for subtask in self.subtasks]
-        results = await asyncio.gather(*tasks, return_exceptions=True)
+        results = await asyncio.gather(
+            *[execute_subtask_with_semaphore(subtask) for subtask in self.subtasks],
+            return_exceptions=True,
+        )
 
-        # Process results using zip for better readability
+        # Process results
         successful_results = []
         failed_results = []
 
-        for subtask, result in zip(self.subtasks, results):
+        for i, result in enumerate(results):
+            subtask = self.subtasks[i]
             if isinstance(result, Exception):
                 failed_results.append(f"Subtask {subtask.task_id} failed: {result}")
             elif isinstance(result, TaskResult):
-                successful_results.append(result.data)
+                if result.success:
+                    successful_results.append(result.data)
+                else:
+                    failed_results.append(
+                        f"Subtask {subtask.task_id} failed: {result.error}"
+                    )
 
         if failed_results:
             return f"Parallel task {self.task_id} partially failed: {failed_results}"
@@ -336,86 +426,8 @@ class WorkflowEngine:
         for level_index, level in enumerate(execution_order, 1):
             print(f"\n{'=' * 20} Phase {level_index} ({len(level)} tasks) {'=' * 20}")
 
-            # Filter executable tasks in this level
-            ready_tasks = []
-            for task_id in level:
-                task = self.tasks[task_id]
-                if task.can_execute(completed_tasks):
-                    task.update_status(TaskStatus.READY)
-                    ready_tasks.append(task)
-
-            if not ready_tasks:
-                continue
-
-            # Create semaphore for this level to control concurrency
-            semaphore = asyncio.Semaphore(self.max_concurrent)
-
-            async def execute_single_task(task: Task, sem: asyncio.Semaphore):
-                async with sem:
-                    task.start_time = time.time()
-                    task.update_status(TaskStatus.RUNNING)
-
-                    try:
-                        # Collect results from dependent tasks
-                        dependency_results = {dep_id: self.results[dep_id] for dep_id in task.get_dependencies()}
-
-                        # Execute task
-                        data = await asyncio.wait_for(task.execute(dependency_results), timeout=task.timeout)
-
-                        task.end_time = time.time()
-                        task.result = TaskResult(
-                            task_id=task.task_id, success=True, data=data, execution_time=task.get_execution_time()
-                        )
-                        task.update_status(TaskStatus.COMPLETED)
-
-                        # Store result
-                        self.results[task.task_id] = task.result
-                        completed_tasks.add(task.task_id)
-
-                        print(f"[OK] Task {task.task_id} completed, duration: {task.get_execution_time():.2f}s")
-
-                        return task.result
-
-                    except asyncio.TimeoutError as e:
-                        task.end_time = time.time()
-                        task.result = TaskResult(
-                            task_id=task.task_id,
-                            success=False,
-                            data=None,
-                            execution_time=task.get_execution_time(),
-                            error=e,
-                        )
-                        task.update_status(TaskStatus.FAILED)
-
-                        # Store result and mark as completed (even if failed)
-                        self.results[task.task_id] = task.result
-                        completed_tasks.add(task.task_id)
-
-                        print(f"[FAIL] Task {task.task_id} timeout")
-                        raise e from e
-                    except Exception as e:
-                        task.end_time = time.time()
-                        task.result = TaskResult(
-                            task_id=task.task_id,
-                            success=False,
-                            data=None,
-                            execution_time=task.get_execution_time(),
-                            error=e,
-                        )
-                        task.update_status(TaskStatus.FAILED)
-
-                        # Store result and mark as completed (even if failed)
-                        self.results[task.task_id] = task.result
-                        completed_tasks.add(task.task_id)
-
-                        print(f"[FAIL] Task {task.task_id} failed: {e}")
-                        raise e
-
-            # Execute all ready tasks in this level in parallel
-            tasks_to_execute = [execute_single_task(task, semaphore) for task in ready_tasks]
-
-            # Use return_exceptions=True to ensure all tasks complete even if some fail
-            await asyncio.gather(*tasks_to_execute, return_exceptions=True)
+            # Execute the current level
+            await self._execute_level(level, completed_tasks)
 
         print(f"\n{'=' * 50}")
         print("Workflow execution completed")
@@ -423,22 +435,113 @@ class WorkflowEngine:
 
         return self.results
 
+    async def _execute_level(self, level: list[str], completed_tasks: set[str]):
+        """Execute a single level of tasks with controlled concurrency."""
+        # Filter executable tasks in this level
+        ready_tasks = []
+        for task_id in level:
+            task = self.tasks[task_id]
+            if task.can_execute(completed_tasks):
+                task.update_status(TaskStatus.READY)
+                ready_tasks.append(task)
+
+        if not ready_tasks:
+            return
+
+        # Use a shared semaphore for this level to control concurrency
+        semaphore = asyncio.Semaphore(self.max_concurrent)
+
+        # Execute all ready tasks in this level in parallel
+        tasks_to_execute = [
+            self._execute_single_task_with_semaphore(task, semaphore, completed_tasks)
+            for task in ready_tasks
+        ]
+
+        # Use return_exceptions=True to ensure all tasks complete even if some fail
+        await asyncio.gather(*tasks_to_execute, return_exceptions=True)
+
+    async def _execute_single_task_with_semaphore(
+        self, task: Task, semaphore: asyncio.Semaphore, completed_tasks: set[str]
+    ):
+        """Execute a single task with semaphore control for concurrency."""
+        async with semaphore:
+            return await self._execute_single_task(task, completed_tasks)
+
+    async def _execute_single_task(self, task: Task, completed_tasks: set[str]):
+        """Execute a single task with error handling."""
+        task.start_time = time.time()
+        task.update_status(TaskStatus.RUNNING)
+
+        # Collect results from dependent tasks
+        dependency_results = {
+            dep_id: self.results[dep_id] for dep_id in task.get_dependencies()
+        }
+
+        try:
+            # Execute task using the common error handling wrapper
+            result_data = await task._execute_with_error_handling(
+                lambda ctx: task.execute(ctx), dependency_results
+            )
+            # Update task with result
+            task.result = result_data
+        except Exception:
+            # Result is already stored in task.result by _execute_with_error_handling
+            pass
+
+        # Store result and update completed tasks (even if failed)
+        self.results[task.task_id] = task.result
+        completed_tasks.add(task.task_id)
+
+        # Print appropriate message based on result
+        if task.result.success:
+            print(
+                f"[OK] Task {task.task_id} completed, duration: {task.get_execution_time():.2f}s"
+            )
+        else:
+            error_msg = (
+                "timeout"
+                if isinstance(task.result.error, asyncio.TimeoutError)
+                else str(task.result.error)
+            )
+            print(f"[FAIL] Task {task.task_id} failed: {error_msg}")
+
+        return task.result
+
     def get_execution_summary(self) -> dict[str, Any]:
         """Get execution summary"""
         total_tasks = len(self.tasks)
-        completed = sum(1 for task in self.tasks.values() if task.status == TaskStatus.COMPLETED)
-        failed = sum(1 for task in self.tasks.values() if task.status == TaskStatus.FAILED)
-
+        if total_tasks == 0:
+            return {
+                "total_tasks": 0,
+                "completed": 0,
+                "failed": 0,
+                "pending": 0,
+                "total_execution_time": 0.0,
+                "success_rate": 0.0,
+            }
+
+        completed = 0
+        failed = 0
         total_time = 0.0
+
+        # Single pass through tasks to calculate all metrics
         for task in self.tasks.values():
+            if task.status == TaskStatus.COMPLETED:
+                completed += 1
+            elif task.status == TaskStatus.FAILED:
+                failed += 1
+
             if task.result:
                 total_time += task.result.execution_time
 
+        pending = total_tasks - completed - failed
+        success_rate = completed / total_tasks if total_tasks > 0 else 0
+
         return {
             "total_tasks": total_tasks,
             "completed": completed,
             "failed": failed,
-            "pending": total_tasks - completed - failed,
+            "pending": pending,
             "total_execution_time": total_time,
-            "success_rate": completed / total_tasks if total_tasks > 0 else 0,
+            "success_rate": success_rate,
         }