Fixing multiple state creation while unites

state-manager/app/controller/enqueue_states.py

@@ -1,4 +1,4 @@
-from beanie.operators import In
+import asyncio
 
 from ..models.enqueue_request import EnqueueRequestModel
 from ..models.enqueue_response import EnqueueResponseModel, StateModel
@@ -10,27 +10,38 @@
 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 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,11 @@
+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 typing import Any, Optional
+import hashlib
+import json
 
 
 class State(BaseDatabaseModel):
@@ -15,4 +18,41 @@ 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()    
+
+    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)

state-manager/tests/unit/controller/test_enqueue_states.py

@@ -1,5 +1,5 @@
 import pytest
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch
 from beanie import PydanticObjectId
 from datetime import datetime
 
@@ -36,29 +36,19 @@ def mock_state(self):
         state.created_at = datetime.now()
         return state
 
-    @patch('app.controller.enqueue_states.State')
+    @patch('app.controller.enqueue_states.find_state')
     async def test_enqueue_states_success(
         self,
-        mock_state_class,
+        mock_find_state,
         mock_namespace,
         mock_enqueue_request,
         mock_state,
         mock_request_id
     ):
         """Test successful enqueuing of states"""
         # Arrange
-        # Mock State.find().limit().to_list() chain
-        mock_query = MagicMock()
-        mock_query.limit = MagicMock(return_value=mock_query)
-        mock_query.to_list = AsyncMock(return_value=[mock_state])
-
-        # Mock State.find().set() chain for updating states
-        mock_update_query = MagicMock()
-        mock_update_query.set = AsyncMock()
-
-        # Configure State.find to return different mocks based on call
-        mock_state_class.find = MagicMock()
-        mock_state_class.find.side_effect = [mock_query, mock_update_query]
+        # Mock find_state to return the mock_state for all calls
+        mock_find_state.return_value = mock_state
 
         # Act
         result = await enqueue_states(
@@ -68,36 +58,31 @@ async def test_enqueue_states_success(
         )
 
         # Assert
-        assert result.count == 1
+        assert result.count == 10  # batch_size=10, so 10 states should be returned
         assert result.namespace == mock_namespace
         assert result.status == StateStatusEnum.QUEUED
-        assert len(result.states) == 1
+        assert len(result.states) == 10
         assert result.states[0].state_id == str(mock_state.id)
         assert result.states[0].node_name == "node1"
         assert result.states[0].identifier == "test_identifier"
         assert result.states[0].inputs == {"key": "value"}
 
-        # Verify the find query was called correctly
-        assert mock_state_class.find.call_count == 2  # Called twice: once for finding, once for updating
-        mock_query.limit.assert_called_once_with(10)
-        mock_update_query.set.assert_called_once()
+        # Verify find_state was called correctly
+        assert mock_find_state.call_count == 10  # Called batch_size times
+        mock_find_state.assert_called_with(mock_namespace, ["node1", "node2"])
 
-    @patch('app.controller.enqueue_states.State')
+    @patch('app.controller.enqueue_states.find_state')
     async def test_enqueue_states_no_states_found(
         self,
-        mock_state_class,
+        mock_find_state,
         mock_namespace,
         mock_enqueue_request,
         mock_request_id
     ):
         """Test when no states are found to enqueue"""
         # Arrange
-        mock_query = MagicMock()
-        mock_query.limit = MagicMock(return_value=mock_query)
-        mock_query.to_list = AsyncMock(return_value=[])
-
-        # When no states are found, the second State.find() call won't happen
-        mock_state_class.find = MagicMock(return_value=mock_query)
+        # Mock find_state to return None for all calls
+        mock_find_state.return_value = None
 
         # Act
         result = await enqueue_states(
@@ -112,10 +97,10 @@ async def test_enqueue_states_no_states_found(
         assert result.status == StateStatusEnum.QUEUED
         assert len(result.states) == 0
 
-    @patch('app.controller.enqueue_states.State')
+    @patch('app.controller.enqueue_states.find_state')
     async def test_enqueue_states_multiple_states(
         self,
-        mock_state_class,
+        mock_find_state,
         mock_namespace,
         mock_enqueue_request,
         mock_request_id
@@ -136,17 +121,8 @@ async def test_enqueue_states_multiple_states(
         state2.inputs = {"input2": "value2"}
         state2.created_at = datetime.now()
 
-        mock_query = MagicMock()
-        mock_query.limit = MagicMock(return_value=mock_query)
-        mock_query.to_list = AsyncMock(return_value=[state1, state2])
-
-        # Mock State.find().set() chain for updating states
-        mock_update_query = MagicMock()
-        mock_update_query.set = AsyncMock()
-
-        # Configure State.find to return different mocks based on call
-        mock_state_class.find = MagicMock()
-        mock_state_class.find.side_effect = [mock_query, mock_update_query]
+        # Mock find_state to return different states
+        mock_find_state.side_effect = [state1, state2, None, None, None, None, None, None, None, None]
 
         # Act
         result = await enqueue_states(
@@ -161,32 +137,36 @@ async def test_enqueue_states_multiple_states(
         assert result.states[0].node_name == "node1"
         assert result.states[1].node_name == "node2"
 
-    @patch('app.controller.enqueue_states.State')
+    @patch('app.controller.enqueue_states.find_state')
     async def test_enqueue_states_database_error(
         self,
-        mock_state_class,
+        mock_find_state,
         mock_namespace,
         mock_enqueue_request,
         mock_request_id
     ):
         """Test handling of database errors"""
         # Arrange
-        mock_state_class.find = MagicMock(side_effect=Exception("Database error"))
-
-        # Act & Assert
-        with pytest.raises(Exception) as exc_info:
-            await enqueue_states(
-                mock_namespace,
-                mock_enqueue_request,
-                mock_request_id
-            )
-
-        assert str(exc_info.value) == "Database error"
-
-    @patch('app.controller.enqueue_states.State')
+        # Mock find_state to raise an exception
+        mock_find_state.side_effect = Exception("Database error")
+
+        # Act
+        result = await enqueue_states(
+            mock_namespace,
+            mock_enqueue_request,
+            mock_request_id
+        )
+
+        # Assert - the function should handle exceptions gracefully and return empty result
+        assert result.count == 0
+        assert result.namespace == mock_namespace
+        assert result.status == StateStatusEnum.QUEUED
+        assert len(result.states) == 0
+
+    @patch('app.controller.enqueue_states.find_state')
     async def test_enqueue_states_with_different_batch_size(
         self,
-        mock_state_class,
+        mock_find_state,
         mock_namespace,
         mock_request_id
     ):
@@ -197,12 +177,8 @@ async def test_enqueue_states_with_different_batch_size(
             batch_size=5
         )
 
-        mock_query = MagicMock()
-        mock_query.limit = MagicMock(return_value=mock_query)
-        mock_query.to_list = AsyncMock(return_value=[])
-
-        # When no states are found, the second State.find() call won't happen
-        mock_state_class.find = MagicMock(return_value=mock_query)
+        # Mock find_state to return None
+        mock_find_state.return_value = None
 
         # Act
         result = await enqueue_states(
@@ -213,4 +189,4 @@ async def test_enqueue_states_with_different_batch_size(
 
         # Assert
         assert result.count == 0
-        mock_query.limit.assert_called_once_with(5)
+        assert mock_find_state.call_count == 5  # Called batch_size times