Add UnitesStrategyEnum and integrate into NodeTemplate validation (#328)

* Add UnitesStrategyEnum and integrate into NodeTemplate validation

- Introduced UnitesStrategyEnum to define strategies for unites: ALL_SUCCESS and ALL_DONE.
- Updated NodeTemplate to include a strategy field, enhancing the model's validation capabilities.
- Modified the check_unites_satisfied function in create_next_states.py to incorporate the new strategy logic for state satisfaction checks, improving the handling of unites conditions.

* Update .codespellignorewords to include 'Kusto' and 'NotIn'

- Added 'Kusto' and 'NotIn' to the list of ignored words for codespell, improving spell-checking accuracy in the project.

* uv run ruff check --fix

* Update strategy field in Unites model to have a default value

- Changed the `strategy` field in the `Unites` model to use a default value of `UnitesStrategyEnum.ALL_SUCCESS`, enhancing the model's usability by providing a predefined strategy option.

* Add tests for ALL_DONE strategy in check_unites_satisfied function

- Introduced two new test cases to validate the behavior of the check_unites_satisfied function when using the ALL_DONE strategy.
- The first test checks for pending states, asserting that the result is False.
- The second test verifies the absence of pending states, asserting that the result is True.
- Updated the NodeTemplate model to include UnitesStrategyEnum for enhanced strategy handling.

* Add Unites Strategy documentation for ALL_SUCCESS and ALL_DONE

* Enhance documentation and logic for unites strategies

- Updated the `ALL_DONE` strategy description to include additional terminal statuses: `CANCELLED`, `NEXT_CREATED_ERROR`, and `PRUNED`, clarifying its behavior.
- Added a caution note to the `ALL_SUCCESS` strategy, highlighting potential indefinite blocking scenarios and suggesting the use of timeouts or fallback strategies.
- Refactored the `check_unites_satisfied` function to improve clarity and maintainability in state satisfaction checks for both `ALL_SUCCESS` and `ALL_DONE` strategies.

* Refactor tests in check_unites_satisfied to use find_one method

- Updated test cases in `test_create_next_states.py` to replace the use of `find` with `find_one` for better clarity and accuracy in state checks.
- Adjusted mock return values to reflect pending and non-pending states, enhancing the reliability of the tests.

* Update docs/docs/exosphere/architecture.md

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: NiveditJain

.github/.codespellignorewords

@@ -13,4 +13,5 @@ SLA
 YML
 SDK
 S3
-Kusto
+Kusto
+NotIn

docs/docs/exosphere/architecture.md

@@ -80,6 +80,58 @@ When a node has a `unites` configuration:
 2. **State fingerprinting** ensures only one unites state is created per unique combination
 3. **Dependency validation** ensures the unites node depends on the specified identifier
 
+### Unites Strategy (Beta)
+
+The `unites` keyword supports different strategies to control when the uniting node should execute. This feature is currently in **beta**.
+
+#### Available Strategies
+
+- **`ALL_SUCCESS`** (default): The uniting node executes only when all states with the specified identifier have reached `SUCCESS` status. If any state fails or is still processing, the uniting node will wait.
+
+- **`ALL_DONE`**: The uniting node executes when all states with the specified identifier have reached any terminal status (`SUCCESS`, `ERRORED`, `CANCELLED`, `NEXT_CREATED_ERROR`, or `PRUNED`). This strategy allows the uniting node to proceed even if some states have failed.
+
+#### Strategy Configuration
+
+You can specify the strategy in your unites configuration:
+
+```json hl_lines="22-25"
+{
+  "nodes": [
+    {
+      "node_name": "DataSplitterNode",
+      "identifier": "data_splitter",
+      "next_nodes": ["processor_1"]
+    },
+    {
+      "node_name": "DataProcessorNode",
+      "identifier": "processor_1",
+      "inputs":{
+        "x":"${{data_splitter.outputs.data_chunk}}"
+      },
+      "next_nodes": ["result_merger"]
+    },    
+    {
+      "node_name": "ResultMergerNode",
+      "identifier": "result_merger",
+      "inputs":{
+        "x_processed":"${{processor_1.outputs.processed_data}}"
+      },
+      "unites": {
+        "identifier": "data_splitter",
+        "strategy": "ALL_SUCCESS"
+      },
+      "next_nodes": []
+    }
+  ]
+}
+```
+
+#### Use Cases
+
+- **`ALL_SUCCESS`**: Use when you need all parallel processes to complete successfully before proceeding. Ideal for data processing workflows where partial failures are not acceptable. **Caution**: This strategy can block indefinitely if any parallel branch never reaches a SUCCESS terminal state. Consider adding timeouts, explicit failure-to-success fallbacks, or using ALL_DONE when partial results are acceptable. Implement watchdogs or retry/timeout policies in workflows to prevent permanent blocking.
+
+- **`ALL_DONE`**: Use when you want to proceed with partial results or when you have error handling logic in the uniting node. Useful for scenarios where you want to aggregate results from successful processes while handling failures separately.
+
 ### Unites Example
 
 ```json hl_lines="22-24"
@@ -95,15 +147,15 @@ When a node has a `unites` configuration:
       "identifier": "processor_1",
       "inputs":{
         "x":"${{data_splitter.outputs.data_chunk}}"
-      }
+      },
       "next_nodes": ["result_merger"]
     },    
     {
       "node_name": "ResultMergerNode",
       "identifier": "result_merger",
       "inputs":{
         "x_processed":"${{processor_1.outputs.processed_data}}"
-      }
+      },
       "unites": {
         "identifier": "data_splitter"
       },

state-manager/app/models/node_template_model.py

@@ -1,10 +1,17 @@
 from pydantic import Field, BaseModel, field_validator
 from typing import Any, Optional, List
 from .dependent_string import DependentString
+from enum import Enum
+
+
+class UnitesStrategyEnum(str, Enum):
+    ALL_SUCCESS = "ALL_SUCCESS"
+    ALL_DONE = "ALL_DONE"
 
 
 class Unites(BaseModel):
     identifier: str = Field(..., description="Identifier of the node")
+    strategy: UnitesStrategyEnum = Field(default=UnitesStrategyEnum.ALL_SUCCESS, description="Strategy of the unites")
 
 
 class NodeTemplate(BaseModel):

state-manager/app/tasks/create_next_states.py

@@ -8,6 +8,7 @@
 from app.models.node_template_model import NodeTemplate
 from app.models.db.registered_node import RegisteredNode
 from app.models.dependent_string import DependentString
+from app.models.node_template_model import UnitesStrategyEnum
 from json_schema_to_pydantic import create_model
 from pydantic import BaseModel
 from typing import Type
@@ -30,15 +31,30 @@ async def check_unites_satisfied(namespace: str, graph_name: str, node_template:
     if not unites_id:
         raise ValueError(f"Unit identifier not found in parents: {node_template.unites.identifier}")
     else:
-        if await State.find(
+        if node_template.unites.strategy == UnitesStrategyEnum.ALL_SUCCESS:
+            any_one_pending = await State.find_one(
                 State.namespace_name == namespace,
                 State.graph_name == graph_name,
                 NE(State.status, StateStatusEnum.SUCCESS),
                 {
                     f"parents.{node_template.unites.identifier}": unites_id
                 }
-            ).count() > 0:
+            )
+            if any_one_pending:
+                return False
+        
+        if node_template.unites.strategy == UnitesStrategyEnum.ALL_DONE:
+            any_one_pending = await State.find_one(
+                State.namespace_name == namespace,
+                State.graph_name == graph_name,
+                In(State.status, [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED]),
+                {
+                    f"parents.{node_template.unites.identifier}": unites_id
+                }
+            )
+            if any_one_pending:
                 return False
+        
     return True
 
 

state-manager/tests/unit/tasks/test_create_next_states.py

@@ -9,7 +9,7 @@
 )
 from app.models.dependent_string import Dependent, DependentString
 from app.models.state_status_enum import StateStatusEnum
-from app.models.node_template_model import NodeTemplate, Unites
+from app.models.node_template_model import NodeTemplate, Unites, UnitesStrategyEnum
 from pydantic import BaseModel
 
 
@@ -233,9 +233,9 @@ async def test_check_unites_satisfied_pending_states(self):
         parents = {"parent1": PydanticObjectId()}
 
         with patch('app.tasks.create_next_states.State') as mock_state:
-            mock_find = AsyncMock()
-            mock_find.count.return_value = 1
-            mock_state.find.return_value = mock_find
+            mock_find_one = AsyncMock()
+            mock_find_one.return_value = {"some": "state"}  # Return a non-None value to indicate pending state
+            mock_state.find_one = mock_find_one
 
             result = await check_unites_satisfied("test_namespace", "test_graph", node_template, parents)
 
@@ -255,9 +255,53 @@ async def test_check_unites_satisfied_no_pending_states(self):
         parents = {"parent1": PydanticObjectId()}
 
         with patch('app.tasks.create_next_states.State') as mock_state:
-            mock_find = AsyncMock()
-            mock_find.count.return_value = 0
-            mock_state.find.return_value = mock_find
+            mock_find_one = AsyncMock()
+            mock_find_one.return_value = None  # Return None to indicate no pending state
+            mock_state.find_one = mock_find_one
+            
+            result = await check_unites_satisfied("test_namespace", "test_graph", node_template, parents)
+            
+            assert result is True
+
+    @pytest.mark.asyncio
+    async def test_check_unites_satisfied_all_done_strategy_pending_states(self):
+        """Test when there are pending states for ALL_DONE strategy"""
+        node_template = NodeTemplate(
+            node_name="test_node",
+            identifier="test_id",
+            namespace="test",
+            inputs={},
+            next_nodes=None,
+            unites=Unites(identifier="parent1", strategy=UnitesStrategyEnum.ALL_DONE)
+        )
+        parents = {"parent1": PydanticObjectId()}
+        
+        with patch('app.tasks.create_next_states.State') as mock_state:
+            mock_find_one = AsyncMock()
+            mock_find_one.return_value = {"some": "state"}  # Return a non-None value to indicate pending state
+            mock_state.find_one = mock_find_one
+            
+            result = await check_unites_satisfied("test_namespace", "test_graph", node_template, parents)
+            
+            assert result is False
+
+    @pytest.mark.asyncio
+    async def test_check_unites_satisfied_all_done_strategy_no_pending_states(self):
+        """Test when there are no pending states for ALL_DONE strategy"""
+        node_template = NodeTemplate(
+            node_name="test_node",
+            identifier="test_id",
+            namespace="test",
+            inputs={},
+            next_nodes=None,
+            unites=Unites(identifier="parent1", strategy=UnitesStrategyEnum.ALL_DONE)
+        )
+        parents = {"parent1": PydanticObjectId()}
+        
+        with patch('app.tasks.create_next_states.State') as mock_state:
+            mock_find_one = AsyncMock()
+            mock_find_one.return_value = None  # Return None to indicate no pending state
+            mock_state.find_one = mock_find_one
 
             result = await check_unites_satisfied("test_namespace", "test_graph", node_template, parents)