Add prune and re-enqueue signal functionality

- Introduced new routes for pruning states and re-enqueuing states after a specified time.
- Added corresponding controller functions to handle the logic for pruning and re-enqueuing states, including validation and error handling.
- Created new signal models for request and response structures related to pruning and re-enqueuing operations.
- Updated the State model to include a new field for enqueue_after, enhancing state management capabilities.
- Enhanced logging for better traceability of operations related to state management.

Author: NiveditJain

state-manager/app/controller/enqueue_states.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 
 from ..models.enqueue_request import EnqueueRequestModel
 from ..models.enqueue_response import EnqueueResponseModel, StateModel
@@ -18,7 +19,8 @@ async def find_state(namespace_name: str, nodes: list[str]) -> State | None:
             "status": StateStatusEnum.CREATED,
             "node_name": {
                 "$in": nodes
-            }
+            },
+            "enqueue_after": {"$lte": int(time.time() * 1000)}
         },
         {
             "$set": {"status": StateStatusEnum.QUEUED}

state-manager/app/controller/prune_signal.py

@@ -0,0 +1,32 @@
+from app.models.signal_models import PruneRequestModel, SignalResponseModel
+from fastapi import HTTPException, status
+from beanie import PydanticObjectId
+
+from app.models.db.state import State
+from app.models.state_status_enum import StateStatusEnum
+from app.singletons.logs_manager import LogsManager
+
+logger = LogsManager().get_logger()
+
+async def prune_signal(namespace_name: str, state_id: PydanticObjectId, body: PruneRequestModel, x_exosphere_request_id: str) -> SignalResponseModel:
+
+    try:
+        logger.info(f"Recieved prune signal for state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        state = await State.find_one(State.id == state_id)
+
+        if not state:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="State not found")
+        
+        if state.status != StateStatusEnum.CREATED:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not created")
+        
+        state.status = StateStatusEnum.PRUNED
+        state.data = body.data
+        await state.save()
+
+        return SignalResponseModel(status=StateStatusEnum.PRUNED, enqueue_after=state.enqueue_after)
+
+    except Exception as e:
+        logger.error(f"Error pruning state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id, error=e)
+        raise

state-manager/app/controller/re_queue_after_singal.py

@@ -0,0 +1,32 @@
+from app.models.signal_models import ReEnqueueAfterRequestModel, SignalResponseModel
+from fastapi import HTTPException, status
+from beanie import PydanticObjectId
+
+from app.models.db.state import State
+from app.models.state_status_enum import StateStatusEnum
+from app.singletons.logs_manager import LogsManager
+
+logger = LogsManager().get_logger()
+
+async def re_queue_after_signal(namespace_name: str, state_id: PydanticObjectId, body: ReEnqueueAfterRequestModel, x_exosphere_request_id: str) -> SignalResponseModel:
+
+    try:
+        logger.info(f"Recieved re-queue after signal for state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        state = await State.find_one(State.id == state_id)
+
+        if not state:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="State not found")
+        
+        if state.status != StateStatusEnum.CREATED:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not created")
+        
+        state.status = StateStatusEnum.CREATED
+        state.enqueue_after = state.enqueue_after + body.enqueue_after
+        await state.save()
+
+        return SignalResponseModel(status=StateStatusEnum.CREATED, enqueue_after=state.enqueue_after)
+
+    except Exception as e:
+        logger.error(f"Error re-queueing state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id, error=e)
+        raise

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

@@ -7,7 +7,7 @@
 from typing import Any, Optional
 import hashlib
 import json
-
+import time
 
 class State(BaseDatabaseModel):
     node_name: str = Field(..., description="Name of the node of the state")
@@ -18,10 +18,12 @@ class State(BaseDatabaseModel):
     status: StateStatusEnum = Field(..., description="Status of the state")
     inputs: dict[str, Any] = Field(..., description="Inputs of the state")
     outputs: dict[str, Any] = Field(..., description="Outputs of the state")
+    data: dict[str, Any] = Field(default_factory=dict, description="Data of the state")
     error: Optional[str] = Field(None, description="Error message")
     parents: dict[str, PydanticObjectId] = Field(default_factory=dict, description="Parents of the state")
     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), description="Unix time in milliseconds after which the state should be enqueued")
 
     @before_event([Insert, Replace, Save])
     def _generate_fingerprint(self):
@@ -65,5 +67,29 @@ class Settings:
                 partialFilterExpression={
                     "does_unites": True
                 }
+            ),
+            IndexModel(
+                [
+                    ("enqueue_after", 1)
+                ],
+                name="idx_enqueue_after"
+            ),
+            IndexModel(
+                [
+                    ("status", 1)
+                ],
+                name="idx_status"
+            ),
+            IndexModel(
+                [
+                    ("namespace_name", 1),
+                ],
+                name="idx_namespace_name"
+            ),
+            IndexModel(
+                [
+                    ("node_name", 1),
+                ],
+                name="idx_node_name"
             )
         ]

state-manager/app/models/signal_models.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, Field
+from .state_status_enum import StateStatusEnum
+from typing import Any
+
+
+class SignalResponseModel(BaseModel):
+    enqueue_after: int = Field(..., description="Unix time in milliseconds after which the state should be re-enqueued")
+    status: StateStatusEnum = Field(..., description="Status of the state")
+
+class PruneRequestModel(BaseModel):
+    data: dict[str, Any] = Field(..., description="Data of the state")
+
+class ReEnqueueAfterRequestModel(BaseModel):
+    enqueue_after: int = Field(..., description="Unix time in milliseconds after which the state should be re-enqueued")

state-manager/app/models/state_status_enum.py

@@ -6,10 +6,8 @@ class StateStatusEnum(str, Enum):
     CREATED = 'CREATED'
     QUEUED = 'QUEUED'
     EXECUTED = 'EXECUTED'
-    NEXT_CREATED = 'NEXT_CREATED'
-    RETRY_CREATED = 'RETRY_CREATED'
-    TIMEDOUT = 'TIMEDOUT'
     ERRORED = 'ERRORED'
     CANCELLED = 'CANCELLED'
     SUCCESS = 'SUCCESS'
     NEXT_CREATED_ERROR = 'NEXT_CREATED_ERROR'
+    PRUNED = 'PRUNED'

state-manager/app/routes.py

@@ -40,6 +40,14 @@
 from .models.graph_structure_models import GraphStructureResponse
 from .controller.get_graph_structure import get_graph_structure
 
+### singnals
+from .models.signal_models import SignalResponseModel
+from .models.signal_models import PruneRequestModel
+from .controller.prune_signal import prune_signal
+from .models.signal_models import ReEnqueueAfterRequestModel
+from .controller.re_queue_after_singal import re_queue_after_signal
+
+
 logger = LogsManager().get_logger()
 
 router = APIRouter(prefix="/v0/namespace/{namespace_name}")
@@ -145,6 +153,44 @@ async def errored_state_route(namespace_name: str, state_id: str, body: ErroredR
     return await errored_state(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id)
 
 
+@router.post(
+    "/states/{state_id}/prune",
+    response_model=SignalResponseModel,
+    status_code=status.HTTP_200_OK,
+    response_description="State skipped successfully",
+    tags=["state"]
+)
+async def prune_state_route(namespace_name: str, state_id: str, body: PruneRequestModel, request: Request, api_key: str = Depends(check_api_key)):
+    x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4()))
+
+    if api_key:
+        logger.info(f"API key is valid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+    else:
+        logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
+
+    return await prune_signal(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id)
+
+
+@router.post(
+    "/states/{state_id}/re-enqueue-after",
+    response_model=SignalResponseModel,
+    status_code=status.HTTP_200_OK,
+    response_description="State re-enqueued successfully",
+    tags=["state"]
+)
+async def re_enqueue_after_state_route(namespace_name: str, state_id: str, body: ReEnqueueAfterRequestModel, request: Request, api_key: str = Depends(check_api_key)):
+    x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4()))
+
+    if api_key:
+        logger.info(f"API key is valid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+    else:
+        logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
+    
+    return await re_queue_after_signal(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id)
+
+
 @router.put(
     "/graph/{graph_name}",
     response_model=UpsertGraphTemplateResponse,