Implement retry policy and enhance errored state handling

- Added a retry policy model to manage state retries with configurable methods (fixed, linear, exponential).
- Updated the errored state function to create a retry state if the maximum retries have not been reached, improving error recovery.
- Enhanced the ErroredResponseModel to include a flag indicating whether a retry state was created.
- Modified the GraphTemplate and State models to incorporate retry policy attributes, ensuring better state management.
- Improved validation and error handling in the upsert_graph_template function to accommodate the new retry policy structure.

Author: NiveditJain

state-manager/app/controller/errored_state.py

@@ -5,9 +5,22 @@
 from app.models.db.state import State
 from app.models.state_status_enum import StateStatusEnum
 from app.singletons.logs_manager import LogsManager
+from app.models.retry_policy_model import RetryPolicyModel, RetryMethod
+from app.models.db.graph_template_model import GraphTemplate
 
 logger = LogsManager().get_logger()
 
+def _calculate_enqueue_after(retry_policy: RetryPolicyModel, retry_count: int) -> int:
+    # convert seconds to milliseconds
+    if retry_policy.method == RetryMethod.FIXED:
+        return (retry_policy.backoff_factor * 1000)
+    elif retry_policy.method == RetryMethod.LINEAR:
+        return (retry_policy.backoff_factor * retry_count) * 1000
+    elif retry_policy.method == RetryMethod.EXPONENTIAL:
+        return (retry_policy.backoff_factor ** retry_count) * 1000
+    else:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid retry method")
+
 async def errored_state(namespace_name: str, state_id: PydanticObjectId, body: ErroredRequestModel, x_exosphere_request_id: str) -> ErroredResponseModel:
 
     try:
@@ -23,11 +36,35 @@ async def errored_state(namespace_name: str, state_id: PydanticObjectId, body: E
         if state.status == StateStatusEnum.EXECUTED:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is already executed")
 
+        graph_template = await GraphTemplate.get(namespace_name, state.graph_name)
+
+        retry_created = False
+
+        if state.retry_count < graph_template.retry_policy.max_retries:
+            retry_state = State(
+                node_name=state.node_name,
+                namespace_name=state.namespace_name,
+                identifier=state.identifier,
+                graph_name=state.graph_name,
+                run_id=state.run_id,
+                status=StateStatusEnum.CREATED,
+                inputs=state.inputs,
+                outputs=state.outputs,
+                error=body.error,
+                parents=state.parents,
+                does_unites=state.does_unites,
+                enqueue_after=state.enqueue_after + _calculate_enqueue_after(graph_template.retry_policy, state.retry_count + 1),
+                retry_count=state.retry_count + 1
+            )
+            retry_state = await retry_state.insert()
+            logger.info(f"Retry state {retry_state.id} created for state {state_id}", x_exosphere_request_id=x_exosphere_request_id)
+            retry_created = True
+
         state.status = StateStatusEnum.ERRORED
         state.error = body.error
         await state.save()
 
-        return ErroredResponseModel(status=StateStatusEnum.ERRORED)
+        return ErroredResponseModel(status=StateStatusEnum.ERRORED, retry_created=retry_created)
 
     except Exception as e:
         logger.error(f"Error errored state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id, error=e)

state-manager/app/controller/upsert_graph_template.py

@@ -27,7 +27,8 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
                     Set({
                         GraphTemplate.nodes: body.nodes, # type: ignore
                         GraphTemplate.validation_status: GraphTemplateValidationStatus.PENDING, # type: ignore
-                        GraphTemplate.validation_errors: [] # type: ignore
+                        GraphTemplate.validation_errors: [], # type: ignore
+                        GraphTemplate.retry_policy: body.retry_policy # type: ignore
                     })
                 )
 
@@ -44,7 +45,8 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
                         namespace=namespace_name,
                         nodes=body.nodes,
                         validation_status=GraphTemplateValidationStatus.PENDING,
-                        validation_errors=[]
+                        validation_errors=[],
+                        retry_policy=body.retry_policy
                     ).set_secrets(body.secrets)
                 )
         except ValueError as e:
@@ -58,6 +60,7 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
             validation_status=graph_template.validation_status,
             validation_errors=graph_template.validation_errors,
             secrets={secret_name: True for secret_name in graph_template.get_secrets().keys()},
+            retry_policy=graph_template.retry_policy,
             created_at=graph_template.created_at,
             updated_at=graph_template.updated_at
         )

state-manager/app/models/db/graph_template_model.py

@@ -11,7 +11,7 @@
 from ..node_template_model import NodeTemplate
 from app.utils.encrypter import get_encrypter
 from app.models.dependent_string import DependentString
-
+from app.models.retry_policy_model import RetryPolicyModel
 
 class GraphTemplate(BaseDatabaseModel):
     name: str = Field(..., description="Name of the graph")
@@ -20,6 +20,7 @@ class GraphTemplate(BaseDatabaseModel):
     validation_status: GraphTemplateValidationStatus = Field(..., description="Validation status of the graph")
     validation_errors: List[str] = Field(default_factory=list, description="Validation errors of the graph")
     secrets: Dict[str, str] = Field(default_factory=dict, description="Secrets of the graph")
+    retry_policy: RetryPolicyModel = Field(default_factory=RetryPolicyModel, description="Retry policy of the graph")
 
     _node_by_identifier: Dict[str, NodeTemplate] | None = PrivateAttr(default=None)
     _parents_by_identifier: Dict[str, set[str]] | None = PrivateAttr(default=None) # type: ignore

state-manager/app/models/db/state.py

@@ -24,6 +24,7 @@ class State(BaseDatabaseModel):
     does_unites: bool = Field(default=False, description="Whether this state unites other states")
     state_fingerprint: str = Field(default="", description="Fingerprint of the state")
     enqueue_after: int = Field(default_factory=lambda: int(time.time() * 1000), gt=0, description="Unix time in milliseconds after which the state should be enqueued")
+    retry_count: int = Field(default=0, description="Number of times the state has been retried")
 
     @before_event([Insert, Replace, Save])
     def _generate_fingerprint(self):
@@ -37,6 +38,7 @@ def _generate_fingerprint(self):
             "identifier": self.identifier,
             "graph_name": self.graph_name,
             "run_id": self.run_id,
+            "retry_count": self.retry_count,
             "parents": {k: str(v) for k, v in self.parents.items()},
         }
         payload = json.dumps(

state-manager/app/models/errored_models.py

@@ -7,4 +7,5 @@ class ErroredRequestModel(BaseModel):
 
 
 class ErroredResponseModel(BaseModel):
-    status: StateStatusEnum = Field(..., description="Status of the state")
+    status: StateStatusEnum = Field(..., description="Status of the state")
+    retry_created: bool = Field(default=False, description="Whether a retry state was created")

state-manager/app/models/graph_models.py

@@ -3,16 +3,19 @@
 from typing import Dict, List, Optional
 from datetime import datetime
 from .graph_template_validation_status import GraphTemplateValidationStatus
+from .retry_policy_model import RetryPolicyModel
 
 
 class UpsertGraphTemplateRequest(BaseModel):
     secrets: Dict[str, str] = Field(..., description="Dictionary of secrets that are used while graph execution")
     nodes: List[NodeTemplate] = Field(..., description="List of node templates that define the graph structure")
+    retry_policy: RetryPolicyModel = Field(default_factory=RetryPolicyModel, description="Retry policy of the graph")
 
 
 class UpsertGraphTemplateResponse(BaseModel):
     nodes: List[NodeTemplate] = Field(..., description="List of node templates that define the graph structure")
     secrets: Dict[str, bool] = Field(..., description="Dictionary of secrets that are used while graph execution")
+    retry_policy: RetryPolicyModel = Field(default_factory=RetryPolicyModel, description="Retry policy of the graph")
     created_at: datetime = Field(..., description="Timestamp when the graph template was created")
     updated_at: datetime = Field(..., description="Timestamp when the graph template was last updated")
     validation_status: GraphTemplateValidationStatus = Field(..., description="Current validation status of the graph template")

state-manager/app/models/retry_policy_model.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel, Field
+from enum import Enum
+
+class RetryMethod(str, Enum):
+    EXPONENTIAL = "EXPONENTIAL"
+    LINEAR = "LINEAR"
+    FIXED = "FIXED"
+
+class RetryPolicyModel(BaseModel):
+    max_retries: int = Field(default=3, description="The maximum number of retries", ge=0)
+    method: RetryMethod = Field(default=RetryMethod.EXPONENTIAL, description="The method of retry")
+    backoff_factor: int = Field(default=2, description="The backoff factor in seconds (default: 2 = 2 seconds)", gt=0)