fixing runtime and BaseNode

Author: NiveditJain

python-sdk/exospherehost/node/BaseNode.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Optional, List
+from typing import Optional, List
 from pydantic import BaseModel  
 
 
@@ -17,16 +17,15 @@ class BaseNode(ABC):
         state (dict[str, Any]): A dictionary for storing node state between executions.
     """
 
-    def __init__(self, unique_name: Optional[str] = None):
+    def __init__(self):
         """
         Initialize a BaseNode instance.
         
         Args:
             unique_name (Optional[str], optional): A unique identifier for this node.
                 If None, the class name will be used as the unique name. Defaults to None.
         """
-        self.unique_name: Optional[str] = unique_name
-        self.state: dict[str, Any] = {}
+        self.inputs: Optional[BaseNode.Inputs] = None
 
     class Inputs(BaseModel):
         """

python-sdk/exospherehost/runtime.py

@@ -54,16 +54,18 @@ def __init__(self, namespace: str, name: str, nodes: List[type[BaseNode]], state
         self._batch_size = batch_size
         self._state_queue = Queue(maxsize=2*batch_size)
         self._workers = workers
-        self._nodes = []
-        self._node_names = []
+        self._nodes = nodes
+        self._node_names = [node.__class__.__name__ for node in nodes]
         self._state_manager_uri = state_manager_uri
         self._state_manager_version = state_manage_version
         self._poll_interval = poll_interval
-        self._node_mapping = {}
+        self._node_mapping = {
+            node.__class__.__name__: node for node in nodes
+        }
 
         self._set_config_from_env()
         self._validate_runtime()
-        self._validate_nodes(nodes)
+        self._validate_nodes()
 
     def _set_config_from_env(self):
         """
@@ -115,7 +117,7 @@ def _get_register_endpoint(self):
         """
         return f"{self._state_manager_uri}/{str(self._state_manager_version)}/namespace/{self._namespace}/nodes/"
 
-    async def _register_nodes(self):
+    async def _register(self):
         """
         Register node schemas and runtime metadata with the state manager.
 
@@ -146,22 +148,6 @@ async def _register_nodes(self):
 
                 return res
 
-    async def _register(self, nodes: List[type[BaseNode]]):
-        """
-        Validate and register nodes with the runtime and state manager.
-
-        Args:
-            nodes (List[type[BaseNode]]): List of BaseNode subclasses to register.
-
-        Raises:
-            ValidationError: If any node is invalid.
-        """
-        self._nodes = self._validate_nodes(nodes)
-        self._node_names = [node.__class__.__name__ for node in nodes]
-        self._node_mapping = {node.__class__.__name__: node for node in self._nodes}
-
-        await self._register_nodes()
-
     async def _enqueue_call(self):
         """
         Request a batch of states to process from the state manager.
@@ -237,7 +223,7 @@ async def _notify_errored(self, state_id: str, error: str):
                 if response.status != 200:
                     logger.error(f"Failed to notify errored state {state_id}: {res}")
 
-    def _validate_nodes(self, nodes: List[type[BaseNode]]):
+    def _validate_nodes(self):
         """
         Validate that all provided nodes are valid BaseNode subclasses.
 
@@ -252,7 +238,7 @@ def _validate_nodes(self, nodes: List[type[BaseNode]]):
         """
         errors = []
 
-        for node in nodes:
+        for node in self._nodes:
             if not issubclass(node, BaseNode):
                 errors.append(f"{node.__class__.__name__} does not inherit from exospherehost.BaseNode")
             if not hasattr(node, "Inputs"):
@@ -265,16 +251,14 @@ def _validate_nodes(self, nodes: List[type[BaseNode]]):
                 errors.append(f"{node.__class__.__name__} does not have an Outputs class that inherits from pydantic.BaseModel")
 
         # Find nodes with the same __class__.__name__
-        class_names = [node.__class__.__name__ for node in nodes]
+        class_names = [node.__class__.__name__ for node in self._nodes]
         duplicate_class_names = [name for name in set(class_names) if class_names.count(name) > 1]
         if duplicate_class_names:
             errors.append(f"Duplicate node class names found: {duplicate_class_names}")
 
         if len(errors) > 0:
             raise ValidationError("Following errors while validating nodes: " + "\n".join(errors))
 
-        return nodes
-
     async def _worker(self):
         """
         Worker task that processes states from the queue.