Merge pull request #222 from dotflow-io/feature/205

🪲 BUG-#205: Fix PR #202 review issues — duration, executor leak, checkpoint, docstring

Author: FernandoCelmer

dotflow/abc/flow.py

@@ -3,7 +3,6 @@
 from abc import ABC, abstractmethod
 from uuid import UUID
 
-from dotflow.core.context import Context
 from dotflow.core.task import Task
 from dotflow.core.types import TypeStatus
 
@@ -43,23 +42,20 @@ def _flow_callback(self, task: Task) -> None:
     def run(self) -> None:
         return None
 
-    def _has_checkpoint(self, task: Task) -> bool:
+    def _try_restore_checkpoint(self, task: Task):
+        """Attempts to restore a checkpoint. Returns the Context if found, None otherwise."""
         if not self.resume:
-            return False
+            return None
 
         context = task.config.storage.get(
             key=task.config.storage.key(task=task)
         )
 
-        return context.storage is not None
-
-    def _restore_checkpoint(self, task: Task) -> Context:
-        previous_context = task.config.storage.get(
-            key=task.config.storage.key(task=task)
-        )
+        if context.storage is None:
+            return None
 
         task.status = TypeStatus.COMPLETED
-        task.current_context = previous_context
+        task.current_context = context
         self._flow_callback(task=task)
 
-        return previous_context
+        return context

dotflow/core/action.py

@@ -19,7 +19,11 @@ def is_execution_with_class_internal_error(error: Exception) -> bool:
 
 
 class Action:
-    """
+    """Decorator for creating task steps.
+
+    Accepts configuration for retry, timeout, retry_delay, and backoff,
+    which are stored as attributes and consumed by TaskEngine during execution.
+
     Import:
         You can import the **action** decorator directly from dotflow:
 
@@ -40,7 +44,6 @@ def my_task():
             def my_task():
                 print("task")
 
-
         With Timeout
 
             @action(timeout=60)
@@ -59,6 +62,10 @@ def my_task():
             def my_task():
                 print("task")
 
+    Note:
+        Retry, timeout, and backoff are enforced by TaskEngine.execute_with_retry(),
+        not by Action itself. Action stores the configuration as attributes.
+
     Args:
         func (Callable):
 

dotflow/core/engine.py

@@ -78,15 +78,15 @@ def start(self):
             self.task.current_context = None
             self.task.status = TypeStatus.FAILED
         else:
-            self.task.duration = (
-                datetime.now() - self._start_time
-            ).total_seconds()
             if self.task.status in (
                 TypeStatus.IN_PROGRESS,
                 TypeStatus.RETRY,
             ):
                 self.task.status = TypeStatus.COMPLETED
         finally:
+            self.task.duration = (
+                datetime.now() - self._start_time
+            ).total_seconds()
             self.task.config.tracer.end_task(task=self.task)
 
     def execute(self):
@@ -172,15 +172,15 @@ def _execute_with_timeout(self, seconds: int):
             return future.result(timeout=seconds)
         except TimeoutError:
             future.cancel()
-            executor.shutdown(wait=False, cancel_futures=True)
-            raise
-        except Exception:
-            executor.shutdown(wait=False)
             raise
+        finally:
+            executor.shutdown(wait=False, cancel_futures=True)
 
     @staticmethod
     def _is_class_internal_error(error: Exception) -> bool:
         """Checks if an error is an internal class execution error."""
+        if isinstance(error, ExecutionWithClassError):
+            return True
         message = str(error)
         patterns = [
             "initial_context",

dotflow/core/workflow.py

@@ -218,8 +218,9 @@ def run(self) -> None:
         previous_context = Context(workflow_id=self.workflow_id)
 
         for task in self.tasks:
-            if self._has_checkpoint(task):
-                previous_context = self._restore_checkpoint(task)
+            restored = self._try_restore_checkpoint(task)
+            if restored:
+                previous_context = restored
                 continue
 
             engine = TaskEngine(
@@ -302,8 +303,9 @@ def _run_group(self, groups: list[Task]) -> None:
         previous_context = Context(workflow_id=self.workflow_id)
 
         for task in groups:
-            if self._has_checkpoint(task):
-                previous_context = self._restore_checkpoint(task)
+            restored = self._try_restore_checkpoint(task)
+            if restored:
+                previous_context = restored
                 continue
 
             engine = TaskEngine(
@@ -344,8 +346,9 @@ def _run_sequential(self) -> None:
         previous_context = Context(workflow_id=self.workflow_id)
 
         for task in self.tasks:
-            if self._has_checkpoint(task):
-                previous_context = self._restore_checkpoint(task)
+            restored = self._try_restore_checkpoint(task)
+            if restored:
+                previous_context = restored
                 continue
 
             engine = TaskEngine(