ruff format, durability

Author: luke-beamcloud

fastloop/__init__.py

@@ -2,12 +2,14 @@
 from .context import LoopContext
 from .fastloop import FastLoop
 from .loop import Loop, LoopEvent, Workflow, WorkflowBlock
+from .types import RetryPolicy
 
 __all__ = [
     "FastLoop",
     "Loop",
     "LoopContext",
     "LoopEvent",
+    "RetryPolicy",
     "Workflow",
     "WorkflowBlock",
     "integrations",

fastloop/constants.py

@@ -5,3 +5,5 @@
 CLAIM_LOCK_SLEEP_S = 0.1
 CLAIM_LOCK_BLOCKING_TIMEOUT_S = 5
 DEFAULT_SYSTEM_CONFIG_DIR = "/etc/fastloop.d/"
+LEASE_TTL_S = 30
+LEASE_HEARTBEAT_INTERVAL_S = 10

fastloop/exceptions.py

@@ -52,3 +52,20 @@ class WorkflowRepeatError(Exception):
 
 class WorkflowNotFoundError(Exception):
     pass
+
+
+class WorkflowMaxRetriesError(Exception):
+    def __init__(
+        self,
+        workflow_id: str,
+        block_index: int,
+        attempts: int,
+        last_error: str | None = None,
+    ):
+        self.workflow_id = workflow_id
+        self.block_index = block_index
+        self.attempts = attempts
+        self.last_error = last_error
+        super().__init__(
+            f"Workflow {workflow_id} block {block_index} failed after {attempts} attempts"
+        )

fastloop/fastloop.py

@@ -27,7 +27,7 @@
 from .logging import configure_logging, setup_logger
 from .loop import Loop, LoopEvent, LoopManager, Workflow, WorkflowManager
 from .state.state import StateManager, create_state_manager
-from .types import BaseConfig, LoopStatus
+from .types import BaseConfig, LoopStatus, RetryPolicy
 from .utils import get_func_import_path, import_func_from_path, infer_application_path
 
 logger = setup_logger()
@@ -460,6 +460,7 @@ def workflow(
         on_stop: Callable[..., Any] | None = None,
         on_block_complete: Callable[..., Any] | None = None,
         on_error: Callable[..., Any] | None = None,
+        retry: RetryPolicy | None = None,
     ) -> Callable[
         [Callable[..., Any] | type[Workflow]], Callable[..., Any] | type[Workflow]
     ]:
@@ -497,6 +498,7 @@ def _decorator(
                 "on_block_complete": workflow_on_block_complete,
                 "on_error": workflow_on_error,
                 "workflow_instance": workflow_instance,
+                "retry_policy": retry,
             }
 
             async def _start_handler(request: dict[str, Any]):
@@ -559,6 +561,7 @@ async def _start_handler(request: dict[str, Any]):
                     on_stop=workflow_on_stop,
                     on_block_complete=workflow_on_block_complete,
                     on_error=workflow_on_error,
+                    retry_policy=retry,
                 )
                 return (
                     await self.state_manager.get_workflow(workflow.workflow_id)
@@ -713,6 +716,13 @@ async def restart_workflow(self, workflow_id: str) -> bool:
             if not workflow.workflow_name:
                 return False
 
+            if workflow.status == LoopStatus.FAILED:
+                logger.info(
+                    "Workflow is failed, not restarting",
+                    extra={"workflow_id": workflow_id},
+                )
+                return False
+
             metadata = self._workflow_metadata.get(workflow.workflow_name)
             if not metadata:
                 logger.warning(
@@ -735,6 +745,7 @@ async def restart_workflow(self, workflow_id: str) -> bool:
                 on_stop=metadata.get("on_stop"),
                 on_block_complete=metadata.get("on_block_complete"),
                 on_error=metadata.get("on_error"),
+                retry_policy=metadata.get("retry_policy"),
             )
 
             if started:

fastloop/loop.py

@@ -20,15 +20,18 @@
     LoopNotFoundError,
     LoopPausedError,
     LoopStoppedError,
+    WorkflowMaxRetriesError,
     WorkflowNextError,
     WorkflowNotFoundError,
     WorkflowRepeatError,
 )
 from .logging import setup_logger
 from .state.state import LoopState, StateManager
-from .types import BaseConfig, LoopEventSender, LoopStatus
+from .types import BaseConfig, LoopEventSender, LoopStatus, RetryPolicy
 from .utils import get_func_import_path
 
+DEFAULT_RETRY_POLICY = RetryPolicy()
+
 if TYPE_CHECKING:
     from .context import LoopContext
 
@@ -135,16 +138,25 @@ class WorkflowState:
     blocks: list[dict[str, Any]] = field(default_factory=list)
     current_block_index: int = 0
     next_payload: dict[str, Any] | None = None
+    completed_blocks: list[int] = field(default_factory=list)
+    block_attempts: dict[int, int] = field(default_factory=dict)
+    last_error: str | None = None
 
     def to_dict(self) -> dict[str, Any]:
-        return self.__dict__.copy()
+        d = self.__dict__.copy()
+        d["block_attempts"] = {str(k): v for k, v in self.block_attempts.items()}
+        return d
 
     def to_string(self) -> str:
-        return json.dumps(self.__dict__, default=str)
+        return json.dumps(self.to_dict(), default=str)
 
     @classmethod
     def from_json(cls, json_str: str) -> "WorkflowState":
         data = json.loads(json_str)
+        if "block_attempts" in data and isinstance(data["block_attempts"], dict):
+            data["block_attempts"] = {
+                int(k): v for k, v in data["block_attempts"].items()
+            }
         return cls(**data)
 
 
@@ -434,6 +446,27 @@ def __init__(self, state_manager: StateManager):
         self.tasks: dict[str, asyncio.Task[None]] = {}
         self.state_manager = state_manager
 
+    async def _persist_block_attempt(
+        self, workflow_id: str, idx: int, error: str | None = None
+    ) -> int:
+        workflow = await self.state_manager.get_workflow(workflow_id)
+        attempts = workflow.block_attempts.get(idx, 0) + 1
+        workflow.block_attempts[idx] = attempts
+        workflow.last_error = error
+        await self.state_manager.update_workflow(workflow_id, workflow)
+        return attempts
+
+    async def _mark_block_completed(
+        self, workflow_id: str, idx: int, next_idx: int, payload: dict | None
+    ) -> None:
+        workflow = await self.state_manager.get_workflow(workflow_id)
+        if idx not in workflow.completed_blocks:
+            workflow.completed_blocks.append(idx)
+        workflow.current_block_index = next_idx
+        workflow.next_payload = payload
+        workflow.block_attempts.pop(idx, None)
+        await self.state_manager.update_workflow(workflow_id, workflow)
+
     async def _run(
         self,
         func: Callable[..., Any],
@@ -442,12 +475,13 @@ async def _run(
         on_stop: Callable[..., Any] | None,
         on_block_complete: Callable[..., Any] | None,
         on_error: Callable[..., Any] | None,
+        retry_policy: RetryPolicy,
     ) -> None:
         try:
             async with self.state_manager.with_workflow_claim(workflow_id):
                 while True:
                     workflow = await self.state_manager.get_workflow(workflow_id)
-                    if workflow.status == LoopStatus.STOPPED:
+                    if workflow.status in (LoopStatus.STOPPED, LoopStatus.FAILED):
                         break
 
                     blocks = [WorkflowBlock(**b) for b in workflow.blocks]
@@ -456,6 +490,12 @@ async def _run(
                     if idx >= len(blocks):
                         raise LoopStoppedError()
 
+                    if idx in workflow.completed_blocks:
+                        await self._mark_block_completed(
+                            workflow_id, idx, idx + 1, workflow.next_payload
+                        )
+                        continue
+
                     current_block = blocks[idx]
                     context.block_index = idx
                     context.block_count = len(blocks)
@@ -467,12 +507,12 @@ async def _run(
                         raise LoopStoppedError()
 
                     except WorkflowNextError as e:
+                        await self._mark_block_completed(
+                            workflow_id, idx, idx + 1, e.payload
+                        )
                         await _call(
                             on_block_complete, context, current_block, e.payload
                         )
-                        await self.state_manager.update_workflow_block_index(
-                            workflow_id, idx + 1, e.payload
-                        )
 
                     except WorkflowRepeatError:
                         pass
@@ -481,20 +521,61 @@ async def _run(
                         raise
 
                     except BaseException as e:
+                        error_str = str(e)
                         logger.error(
-                            "Workflow error",
+                            "Workflow block error",
                             extra={
                                 "workflow_id": workflow_id,
-                                "error": str(e),
+                                "block_index": idx,
+                                "error": error_str,
                                 "traceback": traceback.format_exc(),
                             },
                         )
+
+                        attempts = await self._persist_block_attempt(
+                            workflow_id, idx, error_str
+                        )
+
+                        should_retry = False
                         if on_error:
                             try:
                                 await _call(on_error, context, current_block, e)
                             except WorkflowRepeatError:
-                                continue
-                        raise LoopStoppedError() from e
+                                should_retry = True
+
+                        if not should_retry and attempts < retry_policy.max_attempts:
+                            should_retry = True
+
+                        if should_retry and attempts < retry_policy.max_attempts:
+                            delay = retry_policy.compute_delay(attempts)
+                            logger.info(
+                                "Retrying workflow block",
+                                extra={
+                                    "workflow_id": workflow_id,
+                                    "block_index": idx,
+                                    "attempt": attempts,
+                                    "delay": delay,
+                                },
+                            )
+                            await asyncio.sleep(delay)
+                            continue
+
+                        max_retries_error = WorkflowMaxRetriesError(
+                            workflow_id, idx, attempts, error_str
+                        )
+                        logger.error(
+                            "Workflow block failed after max retries",
+                            extra={
+                                "workflow_id": workflow_id,
+                                "block_index": idx,
+                                "attempts": attempts,
+                            },
+                        )
+                        await _call(on_error, context, current_block, max_retries_error)
+                        await self.state_manager.update_workflow_status(
+                            workflow_id, LoopStatus.FAILED
+                        )
+                        return
 
         except asyncio.CancelledError:
             pass
@@ -521,6 +602,7 @@ async def start(
         on_stop: Callable[..., Any] | None = None,
         on_block_complete: Callable[..., Any] | None = None,
         on_error: Callable[..., Any] | None = None,
+        retry_policy: RetryPolicy | None = None,
     ) -> bool:
         if workflow.workflow_id in self.tasks:
             return False
@@ -535,6 +617,7 @@ async def start(
                 on_stop,
                 on_block_complete,
                 on_error,
+                retry_policy or DEFAULT_RETRY_POLICY,
             )
         )
         return True