added bulk_create_next #236

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

NiveditJain merged 23 commits into exospherehost:main from NiveditJain:234-support-bulk-operation-in-createnextstate-taking-too-long

Aug 20, 2025

state-manager/app/controller/executed_state.py

-Original file line number
+Diff line change
@@ -1,13 +1,12 @@
     from beanie import PydanticObjectId
-    from beanie.operators import In
     from app.models.executed_models import ExecutedRequestModel, ExecutedResponseModel
     from fastapi import HTTPException, status, BackgroundTasks
     from app.models.db.state import State
     from app.models.state_status_enum import StateStatusEnum
     from app.singletons.logs_manager import LogsManager
-    from app.tasks.create_next_state import create_next_state
+    from app.tasks.create_next_states import create_next_states
     logger = LogsManager().get_logger()
@@ Expand All @@
             if state.status != StateStatusEnum.QUEUED:
                 raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not queued")
+            next_state_ids = []
             if len(body.outputs) == 0:
                 state.status = StateStatusEnum.EXECUTED
                 state.outputs = {}
                 await state.save()
-                background_tasks.add_task(create_next_state, state)
+                next_state_ids.append(state.id)
             else:
                 state.outputs = body.outputs[0]
                 state.status = StateStatusEnum.EXECUTED
                 await state.save()
-                background_tasks.add_task(create_next_state, state)
+                next_state_ids.append(state.id)
                 new_states = []
                 for output in body.outputs[1:]:
@@ Expand All @@
                 if len(new_states) > 0:
                     inserted_ids = (await State.insert_many(new_states)).inserted_ids
+                    next_state_ids.extend(inserted_ids)
-                    inserted_states = await State.find(
-                        In(State.id, inserted_ids)
-                    ).to_list()
-                    if len(inserted_states) != len(new_states):
-                        raise RuntimeError(f"Failed to insert all new states. Expected {len(new_states)} states, but only {len(inserted_states)} were inserted")
-                    for inserted_state in inserted_states:
-                        background_tasks.add_task(create_next_state, inserted_state)
+            background_tasks.add_task(create_next_states, next_state_ids, state.identifier, state.namespace_name, state.graph_name, state.parents)
             return ExecutedResponseModel(status=StateStatusEnum.EXECUTED)
@@ Expand Down @@

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

-Original file line number
+Diff line change
@@ -1,7 +1,9 @@
     import base64
+    import time
+    import asyncio
     from .base import BaseDatabaseModel
-    from pydantic import Field, field_validator
+    from pydantic import Field, field_validator, PrivateAttr
     from typing import Optional, List
     from ..graph_template_validation_status import GraphTemplateValidationStatus
     from ..node_template_model import NodeTemplate
@@ Expand All / @@ -17,6 +19,7 @@ class GraphTemplate(BaseDatabaseModel): @@
         validation_status: GraphTemplateValidationStatus = Field(..., description="Validation status of the graph")
         validation_errors: Optional[List[str]] = Field(None, description="Validation errors of the graph")
         secrets: Dict[str, str] = Field(default_factory=dict, description="Secrets of the graph")
+        _node_by_identifier: Dict[str, NodeTemplate] | None = PrivateAttr(default=None)
         class Settings:
             indexes = [
@@ Expand All / @@ -27,12 +30,18 @@ class Settings: @@
                 )
             ]
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+        def _build_node_by_identifier(self) -> None:
+            self._node_by_identifier = {node.identifier: node for node in self.nodes}
         def get_node_by_identifier(self, identifier: str) -> NodeTemplate | None:
             """Get a node by its identifier using O(1) dictionary lookup."""
-            for node in self.nodes:
-                if node.identifier == identifier:
-                    return node
-            return None
+            if self._node_by_identifier is None:
+                self._build_node_by_identifier()
+            return self._node_by_identifier.get(identifier) # type: ignore
         @field_validator('secrets')
         @classmethod
@@ Expand Down Expand Up / @@ -78,4 +87,40 @@ def get_secret(self, secret_name: str) -> str | None: @@
                 return None
             if secret_name not in self.secrets:
                 return None
-            return get_encrypter().decrypt(self.secrets[secret_name])
+            return get_encrypter().decrypt(self.secrets[secret_name])
+        def is_valid(self) -> bool:
+            return self.validation_status == GraphTemplateValidationStatus.VALID
+        def is_validating(self) -> bool:
+            return self.validation_status in (GraphTemplateValidationStatus.ONGOING, GraphTemplateValidationStatus.PENDING)
+        @staticmethod
+        async def get(namespace: str, graph_name: str) -> "GraphTemplate":
+            graph_template = await GraphTemplate.find_one(GraphTemplate.namespace == namespace, GraphTemplate.name == graph_name)
+            if not graph_template:
+                raise ValueError(f"Graph template not found for namespace: {namespace} and graph name: {graph_name}")
+            return graph_template
+        @staticmethod
+        async def get_valid(namespace: str, graph_name: str, polling_interval: float = 1.0, timeout: float = 300.0) -> "GraphTemplate":
+            # Validate polling_interval and timeout
+            if polling_interval <= 0:
+                raise ValueError("polling_interval must be positive")
+            if timeout <= 0:
+                raise ValueError("timeout must be positive")
+            # Coerce polling_interval to a sensible minimum
+            if polling_interval < 0.1:
+                polling_interval = 0.1
+            start_time = time.monotonic()
+            while time.monotonic() - start_time < timeout:
+                graph_template = await GraphTemplate.get(namespace, graph_name)
+                if graph_template.is_valid():
+                    return graph_template
+                if graph_template.is_validating():
+                    await asyncio.sleep(polling_interval)
+                else:
+                    raise ValueError(f"Graph template is in a non-validating state: {graph_template.validation_status.value} for namespace: {namespace} and graph name: {graph_name}")
+            raise ValueError(f"Graph template is not valid for namespace: {namespace} and graph name: {graph_name} after {timeout} seconds")

state-manager/app/models/state_status_enum.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -11,4 +11,5 @@ class StateStatusEnum(str, Enum): @@
         TIMEDOUT = 'TIMEDOUT'
         ERRORED = 'ERRORED'
         CANCELLED = 'CANCELLED'
-        SUCCESS = 'SUCCESS'
+        SUCCESS = 'SUCCESS'
+        NEXT_CREATED_ERROR = 'NEXT_CREATED_ERROR'

state-manager/app/tasks/create_next_state.py

This file was deleted.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

added bulk_create_next #236

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!