Merge branch 'niki/dev/statemanagerfrontend' of https://github.com/nk-ag/exospherehost into niki/dev/statemanagerfrontend

Author: nk-ag

state-manager/app/controller/enqueue_states.py

@@ -1,36 +1,49 @@
-from beanie.operators import In
+import asyncio
 
 from ..models.enqueue_request import EnqueueRequestModel
 from ..models.enqueue_response import EnqueueResponseModel, StateModel
 from ..models.db.state import State
 from ..models.state_status_enum import StateStatusEnum
 
 from app.singletons.logs_manager import LogsManager
+from pymongo import ReturnDocument
 
 logger = LogsManager().get_logger()
 
 
+async def find_state(namespace_name: str, nodes: list[str]) -> State | None:
+    data = await State.get_pymongo_collection().find_one_and_update(
+        {
+            "namespace_name": namespace_name,
+            "status": StateStatusEnum.CREATED,
+            "node_name": {
+                "$in": nodes
+            }
+        },
+        {
+            "$set": {"status": StateStatusEnum.QUEUED}
+        },
+        return_document=ReturnDocument.AFTER
+    )
+    return State(**data) if data else None
+
 async def enqueue_states(namespace_name: str, body: EnqueueRequestModel, x_exosphere_request_id: str) -> EnqueueResponseModel:
 
     try:
         logger.info(f"Enqueuing states for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
 
-        states = await State.find(
-            State.namespace_name == namespace_name,
-            State.status == StateStatusEnum.CREATED,
-            In(State.node_name, body.nodes)
-        ).limit(
-            body.batch_size
-        ).to_list()
-
-        if states:
-            await State.find(
-                In(State.id, [state.id for state in states])
-            ).set(
-                {
-                    "status": StateStatusEnum.QUEUED,
-                }
-            ) # type: ignore
+        # Create tasks for parallel execution
+        tasks = [find_state(namespace_name, body.nodes) for _ in range(body.batch_size)]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        
+        # Filter out None results and exceptions
+        states = []
+        for result in results:
+            if isinstance(result, Exception):
+                logger.error(f"Error finding state: {result}", x_exosphere_request_id=x_exosphere_request_id)
+                continue
+            if result is not None:
+                states.append(result)
 
         response = EnqueueResponseModel(
             count=len(states),

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

@@ -1,8 +1,12 @@
+from pymongo import IndexModel
 from .base import BaseDatabaseModel
 from ..state_status_enum import StateStatusEnum
 from pydantic import Field
-from beanie import PydanticObjectId
+from beanie import Insert, PydanticObjectId, Replace, Save, before_event
+from pymongo.results import InsertManyResult
 from typing import Any, Optional
+import hashlib
+import json
 
 
 class State(BaseDatabaseModel):
@@ -15,4 +19,51 @@ class State(BaseDatabaseModel):
     inputs: dict[str, Any] = Field(..., description="Inputs of the state")
     outputs: dict[str, Any] = Field(..., description="Outputs of the state")
     error: Optional[str] = Field(None, description="Error message")
-    parents: dict[str, PydanticObjectId] = Field(default_factory=dict, description="Parents of the state")
+    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")
+    
+    @before_event([Insert, Replace, Save])
+    def _generate_fingerprint(self):
+        if not self.does_unites:
+            self.state_fingerprint = ""
+            return
+        
+        data = {
+            "node_name": self.node_name,
+            "namespace_name": self.namespace_name,
+            "identifier": self.identifier,
+            "graph_name": self.graph_name,
+            "run_id": self.run_id,
+            "parents": {k: str(v) for k, v in self.parents.items()},
+        }
+        payload = json.dumps(
+            data,
+            sort_keys=True,            # canonical key ordering at all levels
+            separators=(",", ":"),     # no whitespace variance
+            ensure_ascii=True,         # normalized non-ASCII escapes
+        ).encode("utf-8")
+        self.state_fingerprint = hashlib.sha256(payload).hexdigest()    
+    
+    @classmethod
+    async def insert_many(cls, documents: list["State"]) -> InsertManyResult:
+        """Override insert_many to ensure fingerprints are generated before insertion."""
+        # Generate fingerprints for states that need them
+        for state in documents:
+            state._generate_fingerprint()
+        
+        return await super().insert_many(documents) # type: ignore
+        
+    class Settings:
+        indexes = [
+            IndexModel(
+                [
+                    ("state_fingerprint", 1)
+                ],
+                unique=True,
+                name="uniq_state_fingerprint_unites",
+                partialFilterExpression={
+                    "does_unites": True
+                }
+            )
+        ]

state-manager/app/tasks/create_next_states.py

@@ -1,4 +1,5 @@
 from beanie import PydanticObjectId
+from pymongo.errors import DuplicateKeyError, BulkWriteError
 from beanie.operators import In, NE
 from app.singletons.logs_manager import LogsManager
 from app.models.db.graph_template_model import GraphTemplate
@@ -58,6 +59,7 @@ async def check_unites_satisfied(namespace: str, graph_name: str, node_template:
                 return False
     return True
 
+
 def get_dependents(syntax_string: str) -> DependentString:
     splits = syntax_string.split("${{")
     if len(splits) <= 1:
@@ -134,6 +136,7 @@ def generate_next_state(next_state_input_model: Type[BaseModel], next_state_node
         parents=new_parents,
         inputs=next_state_input_data,
         outputs={},
+        does_unites=next_state_node_template.unites is not None,
         run_id=current_state.run_id,
         error=None
     )
@@ -231,10 +234,17 @@ async def get_input_model(node_template: NodeTemplate) -> Type[BaseModel]:
             parent_state = parents[next_state_node_template.unites.identifier]
 
             new_unit_states.append(generate_next_state(next_state_input_model, next_state_node_template, parents, parent_state))
-
-        if len(new_unit_states) > 0:
-            await State.insert_many(new_unit_states)
-    
+        
+        try:
+            if len(new_unit_states) > 0:
+                await State.insert_many(new_unit_states)
+        except (DuplicateKeyError, BulkWriteError):
+            logger.warning(
+                f"Caught duplicate key error for new unit states in namespace={namespace}, "
+                f"graph={graph_name}, likely due to a race condition. "
+                f"Attempted to insert {len(new_unit_states)} states"
+            )
+            
     except Exception as e:
         await State.find(
             In(State.id, state_ids)