runnable 0.28.7__py3-none-any.whl → 0.29.0__py3-none-any.whl

extensions/pipeline_executor/__init__.py

@@ -198,7 +198,6 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         node: BaseNode,
         map_variable: TypeMapVariable = None,
         mock: bool = False,
-        **kwargs,
     ):
         """
         This is the entry point when we do the actual execution of the function.
@@ -232,7 +231,6 @@ class GenericPipelineExecutor(BasePipelineExecutor):
             map_variable=map_variable,
             attempt_number=self.step_attempt_number,
             mock=mock,
-            **kwargs,
         )
 
         data_catalogs_put: Optional[List[DataCatalog]] = self._sync_catalog(stage="put")
@@ -248,7 +246,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
 
         self._context.run_log_store.add_step_log(step_log, self._context.run_id)
 
-    def add_code_identities(self, node: BaseNode, step_log: StepLog, **kwargs):
+    def add_code_identities(self, node: BaseNode, step_log: StepLog):
         """
         Add code identities specific to the implementation.
 
@@ -260,9 +258,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         """
         step_log.code_identities.append(utils.get_git_code_identity())
 
-    def execute_from_graph(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_from_graph(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         This is the entry point to from the graph execution.
 
@@ -303,12 +299,12 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         # Add the step log to the database as per the situation.
         # If its a terminal node, complete it now
         if node.node_type in ["success", "fail"]:
-            self._execute_node(node, map_variable=map_variable, **kwargs)
+            self._execute_node(node, map_variable=map_variable)
             return
 
         # We call an internal function to iterate the sub graphs and execute them
         if node.is_composite:
-            node.execute_as_graph(map_variable=map_variable, **kwargs)
+            node.execute_as_graph(map_variable=map_variable)
             return
 
         task_console.export_text(clear=True)
@@ -317,10 +313,10 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         console.print(
             f":runner: Executing the node {task_name} ... ", style="bold color(208)"
         )
-        self.trigger_node_execution(node=node, map_variable=map_variable, **kwargs)
+        self.trigger_node_execution(node=node, map_variable=map_variable)
 
     def trigger_node_execution(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
+        self, node: BaseNode, map_variable: TypeMapVariable = None
     ):
         """
         Call this method only if we are responsible for traversing the graph via
@@ -376,7 +372,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
 
         return step_log.status, next_node_name
 
-    def execute_graph(self, dag: Graph, map_variable: TypeMapVariable = None, **kwargs):
+    def execute_graph(self, dag: Graph, map_variable: TypeMapVariable = None):
         """
         The parallelization is controlled by the nodes and not by this function.
 
@@ -430,7 +426,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
             )
 
             try:
-                self.execute_from_graph(working_on, map_variable=map_variable, **kwargs)
+                self.execute_from_graph(working_on, map_variable=map_variable)
                 status, next_node_name = self._get_status_and_next_node_name(
                     current_node=working_on, dag=dag, map_variable=map_variable
                 )
@@ -593,7 +589,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         step_log.status = defaults.PROCESSING
         self._context.run_log_store.add_step_log(step_log, self._context.run_id)
 
-        node.fan_out(executor=self, map_variable=map_variable)
+        node.fan_out(map_variable=map_variable)
 
     def fan_in(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
@@ -614,7 +610,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
             map_variable (dict, optional): If the node if of a map state,.Defaults to None.
 
         """
-        node.fan_in(executor=self, map_variable=map_variable)
+        node.fan_in(map_variable=map_variable)
 
         step_log = self._context.run_log_store.get_step_log(
             node._get_step_log_name(map_variable=map_variable), self._context.run_id

extensions/pipeline_executor/argo.py

@@ -810,7 +810,6 @@ class ArgoExecutor(GenericPipelineExecutor):
         self,
         dag: Graph,
         map_variable: dict[str, str | int | float] | None = None,
-        **kwargs,
     ):
         # All the arguments set at the spec level can be referred as "{{workflow.parameters.*}}"
         # We want to use that functionality to override the parameters at the task level
@@ -886,7 +885,6 @@ class ArgoExecutor(GenericPipelineExecutor):
         self,
         node: BaseNode,
         map_variable: dict[str, str | int | float] | None = None,
-        **kwargs,
     ):
         error_on_existing_run_id = os.environ.get("error_on_existing_run_id", "false")
         exists_ok = error_on_existing_run_id == "false"
@@ -904,7 +902,7 @@ class ArgoExecutor(GenericPipelineExecutor):
         step_log.status = defaults.PROCESSING
         self._context.run_log_store.add_step_log(step_log, self._context.run_id)
 
-        self._execute_node(node=node, map_variable=map_variable, **kwargs)
+        self._execute_node(node=node, map_variable=map_variable)
 
         # Raise exception if the step failed
         step_log = self._context.run_log_store.get_step_log(

extensions/pipeline_executor/local.py

@@ -29,16 +29,14 @@ class LocalExecutor(GenericPipelineExecutor):
 
     _is_local: bool = PrivateAttr(default=True)
 
-    def execute_from_graph(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_from_graph(self, node: BaseNode, map_variable: TypeMapVariable = None):
         if not self.object_serialisation:
             self._context.object_serialisation = False
 
-        super().execute_from_graph(node=node, map_variable=map_variable, **kwargs)
+        super().execute_from_graph(node=node, map_variable=map_variable)
 
     def trigger_node_execution(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
+        self, node: BaseNode, map_variable: TypeMapVariable = None
     ):
         """
         In this mode of execution, we prepare for the node execution and execute the node
@@ -47,11 +45,9 @@ class LocalExecutor(GenericPipelineExecutor):
             node (BaseNode): [description]
             map_variable (str, optional): [description]. Defaults to ''.
         """
-        self.execute_node(node=node, map_variable=map_variable, **kwargs)
+        self.execute_node(node=node, map_variable=map_variable)
 
-    def execute_node(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         For local execution, we just execute the node.
 
@@ -59,4 +55,4 @@ class LocalExecutor(GenericPipelineExecutor):
             node (BaseNode): _description_
             map_variable (dict[str, str], optional): _description_. Defaults to None.
         """
-        self._execute_node(node=node, map_variable=map_variable, **kwargs)
+        self._execute_node(node=node, map_variable=map_variable)

extensions/pipeline_executor/local_container.py

@@ -59,7 +59,7 @@ class LocalContainerExecutor(GenericPipelineExecutor):
     _container_secrets_location = "/tmp/dotenv"
     _volumes: Dict[str, Dict[str, str]] = {}
 
-    def add_code_identities(self, node: BaseNode, step_log: StepLog, **kwargs):
+    def add_code_identities(self, node: BaseNode, step_log: StepLog):
         """
         Call the Base class to add the git code identity and add docker identity
 
@@ -86,18 +86,18 @@ class LocalContainerExecutor(GenericPipelineExecutor):
             code_id.code_identifier_url = "local docker host"
             step_log.code_identities.append(code_id)
 
-    def execute_node(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         We are already in the container, we just execute the node.
         The node is already prepared for execution.
         """
         self._use_volumes()
-        return self._execute_node(node, map_variable, **kwargs)
+        return self._execute_node(node, map_variable)
 
     def execute_from_graph(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
+        self,
+        node: BaseNode,
+        map_variable: TypeMapVariable = None,
     ):
         """
         This is the entry point to from the graph execution.
@@ -139,12 +139,12 @@ class LocalContainerExecutor(GenericPipelineExecutor):
         # Add the step log to the database as per the situation.
         # If its a terminal node, complete it now
         if node.node_type in ["success", "fail"]:
-            self._execute_node(node, map_variable=map_variable, **kwargs)
+            self._execute_node(node, map_variable=map_variable)
             return
 
         # We call an internal function to iterate the sub graphs and execute them
         if node.is_composite:
-            node.execute_as_graph(map_variable=map_variable, **kwargs)
+            node.execute_as_graph(map_variable=map_variable)
             return
 
         task_console.export_text(clear=True)
@@ -153,10 +153,10 @@ class LocalContainerExecutor(GenericPipelineExecutor):
         console.print(
             f":runner: Executing the node {task_name} ... ", style="bold color(208)"
         )
-        self.trigger_node_execution(node=node, map_variable=map_variable, **kwargs)
+        self.trigger_node_execution(node=node, map_variable=map_variable)
 
     def trigger_node_execution(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
+        self, node: BaseNode, map_variable: TypeMapVariable = None
     ):
         """
         We come into this step via execute from graph, use trigger job to spin up the container.
@@ -181,7 +181,6 @@ class LocalContainerExecutor(GenericPipelineExecutor):
             command=command,
             map_variable=map_variable,
             auto_remove_container=auto_remove_container,
-            **kwargs,
         )
 
         step_log = self._context.run_log_store.get_step_log(
@@ -203,7 +202,6 @@ class LocalContainerExecutor(GenericPipelineExecutor):
         command: str,
         map_variable: TypeMapVariable = None,
         auto_remove_container: bool = True,
-        **kwargs,
     ):
         """
         During the flow run, we have to spin up a container with the docker image mentioned

extensions/pipeline_executor/mocked.py

@@ -36,9 +36,7 @@ class MockedExecutor(GenericPipelineExecutor):
     def _context(self):
         return context.run_context
 
-    def execute_from_graph(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_from_graph(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         This is the entry point to from the graph execution.
 
@@ -80,18 +78,18 @@ class MockedExecutor(GenericPipelineExecutor):
         # If its a terminal node, complete it now
         if node.node_type in ["success", "fail"]:
             self._context.run_log_store.add_step_log(step_log, self._context.run_id)
-            self._execute_node(node, map_variable=map_variable, **kwargs)
+            self._execute_node(node, map_variable=map_variable)
             return
 
         # We call an internal function to iterate the sub graphs and execute them
         if node.is_composite:
             self._context.run_log_store.add_step_log(step_log, self._context.run_id)
-            node.execute_as_graph(map_variable=map_variable, **kwargs)
+            node.execute_as_graph(map_variable=map_variable)
             return
 
         if node.name not in self.patches:
             # node is not patched, so mock it
-            self._execute_node(node, map_variable=map_variable, mock=True, **kwargs)
+            self._execute_node(node, map_variable=map_variable, mock=True)
         else:
             # node is patched
             # command as the patch value
@@ -103,9 +101,7 @@ class MockedExecutor(GenericPipelineExecutor):
                 node_name=node.name,
             )
             node_to_send.executable = executable
-            self._execute_node(
-                node_to_send, map_variable=map_variable, mock=False, **kwargs
-            )
+            self._execute_node(node_to_send, map_variable=map_variable, mock=False)
 
     def _resolve_executor_config(self, node: BaseNode):
         """
@@ -144,9 +140,7 @@ class MockedExecutor(GenericPipelineExecutor):
 
         return effective_node_config
 
-    def execute_node(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         The entry point for all executors apart from local.
         We have already prepared for node execution.

extensions/pipeline_executor/retry.py

@@ -63,9 +63,7 @@ class RetryExecutor(GenericPipelineExecutor):
         # Should the parameters be copied from previous execution
         # self._set_up_for_re_run(params=params)
 
-    def execute_from_graph(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_from_graph(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         This is the entry point to from the graph execution.
 
@@ -103,7 +101,7 @@ class RetryExecutor(GenericPipelineExecutor):
         # If its a terminal node, complete it now
         if node.node_type in ["success", "fail"]:
             self._context.run_log_store.add_step_log(step_log, self._context.run_id)
-            self._execute_node(node, map_variable=map_variable, **kwargs)
+            self._execute_node(node, map_variable=map_variable)
             return
 
         # In retry step
@@ -118,12 +116,12 @@ class RetryExecutor(GenericPipelineExecutor):
         # We call an internal function to iterate the sub graphs and execute them
         if node.is_composite:
             self._context.run_log_store.add_step_log(step_log, self._context.run_id)
-            node.execute_as_graph(map_variable=map_variable, **kwargs)
+            node.execute_as_graph(map_variable=map_variable)
             return
 
         # Executor specific way to trigger a job
         self._context.run_log_store.add_step_log(step_log, self._context.run_id)
-        self.execute_node(node=node, map_variable=map_variable, **kwargs)
+        self.execute_node(node=node, map_variable=map_variable)
 
     def _is_step_eligible_for_rerun(
         self, node: BaseNode, map_variable: TypeMapVariable = None
@@ -174,7 +172,5 @@ class RetryExecutor(GenericPipelineExecutor):
         self._restart_initiated = True
         return True
 
-    def execute_node(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
-        self._execute_node(node, map_variable=map_variable, **kwargs)
+    def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None):
+        self._execute_node(node, map_variable=map_variable)

extensions/run_log_store/generic_chunked.py

@@ -318,7 +318,6 @@ class ChunkedRunLogStore(BaseRunLogStore):
         tag: str = "",
         original_run_id: str = "",
         status: str = defaults.CREATED,
-        **kwargs,
     ):
         """
         Creates a Run Log object by using the config
@@ -549,7 +548,7 @@ class ChunkedRunLogStore(BaseRunLogStore):
         )
 
     def get_branch_log(
-        self, internal_branch_name: str, run_id: str, **kwargs
+        self, internal_branch_name: str, run_id: str
     ) -> Union[BranchLog, RunLog]:
         """
         Returns the branch log by the internal branch name for the run id

extensions/secrets/dotenv.py

@@ -37,7 +37,7 @@ class DotEnvSecrets(BaseSecrets):
         """
         self.secrets = dotenv_values(self.secrets_location)
 
-    def get(self, name: str = "", **kwargs) -> str:
+    def get(self, name: str = "") -> str:
         """
         Get a secret of name from the secrets file.
 

extensions/tasks/torch.py

@@ -0,0 +1,52 @@
+from typing import List, Optional
+
+from pydantic import Field, field_validator
+
+from runnable import defaults
+from runnable.datastore import StepAttempt
+from runnable.defaults import TypeMapVariable
+from runnable.tasks import BaseTaskType
+
+
+def run_torch_task(
+    rank: int = 1,
+    world_size: int = 1,
+    entrypoint: str = "some function",
+    catalog: Optional[dict[str, List[str]]] = None,
+    task_returns: Optional[List[str]] = None,
+    secrets: Optional[list[str]] = None,
+):
+    # Entry point that creates a python job using simpler python types
+    # and and executes them. The run_id for the job is set to be run_id_rank
+    # Since the configuration file is passes as environmental variable,
+    # The job will use the configuration file to get the required information.
+
+    # In pseudocode, the following is done:
+    # Create the catalog object
+    # Create the secrets and other objects required for the PythonJob
+    # Init the process group using:
+    # https://github.com/pytorch/examples/blob/main/imagenet/main.py#L140
+    # Execute the job, the job is expected to use the environmental variables
+    # to identify the rank or can have them as variable in the signature.
+    # Once the job is executed, we destroy the process group
+    pass
+
+
+class TorchTaskType(BaseTaskType):
+    task_type: str = Field(default="torch", serialization_alias="command_type")
+    command: str
+    num_gpus: int = Field(default=1, description="Number of GPUs to use")
+
+    @field_validator("num_gpus")
+    @classmethod
+    def check_if_cuda_is_available(cls, num_gpus: int) -> int:
+        # Import torch and check if cuda is available
+        # validate if the number of gpus is less than or equal to available gpus
+        return num_gpus
+
+    def execute_command(
+        self,
+        map_variable: TypeMapVariable = None,
+    ) -> StepAttempt:
+        # We have to spawn here
+        return StepAttempt(attempt_number=1, status=defaults.SUCCESS)

runnable/__init__.py

@@ -31,6 +31,7 @@ from runnable.sdk import (  # noqa
     ShellTask,
     Stub,
     Success,
+    Torch,
     metric,
     pickled,
 )

runnable/entrypoints.py

@@ -174,7 +174,7 @@ def set_pipeline_spec_from_yaml(run_context: context.Context, pipeline_file: str
 def set_pipeline_spec_from_python(run_context: context.Context, python_module: str):
     # Call the SDK to get the dag
     # Import the module and call the function to get the dag
-    module_file = python_module.strip(".py")
+    module_file = python_module.rstrip(".py")
     module, func = utils.get_module_and_attr_names(module_file)
     sys.path.insert(0, os.getcwd())  # Need to add the current directory to path
     imported_module = importlib.import_module(module)
@@ -429,7 +429,7 @@ def set_job_spec_from_yaml(run_context: context.Context, job_definition_file: st
 
 def set_job_spec_from_python(run_context: context.Context, python_module: str):
     # Import the module and call the function to get the task
-    module_file = python_module.strip(".py")
+    module_file = python_module.rstrip(".py")
     module, func = utils.get_module_and_attr_names(module_file)
     sys.path.insert(0, os.getcwd())  # Need to add the current directory to path
     imported_module = importlib.import_module(module)

runnable/executor.py

@@ -107,7 +107,7 @@ class BaseJobExecutor(BaseExecutor):
         ...
 
     @abstractmethod
-    def add_code_identities(self, job_log: JobLog, **kwargs):
+    def add_code_identities(self, job_log: JobLog):
         """
         Add code identities specific to the implementation.
 
@@ -161,7 +161,7 @@ class BasePipelineExecutor(BaseExecutor):
     _context_node: Optional[BaseNode] = PrivateAttr(default=None)
 
     @abstractmethod
-    def add_code_identities(self, node: BaseNode, step_log: StepLog, **kwargs):
+    def add_code_identities(self, node: BaseNode, step_log: StepLog):
         """
         Add code identities specific to the implementation.
 
@@ -204,7 +204,6 @@ class BasePipelineExecutor(BaseExecutor):
         node: BaseNode,
         map_variable: TypeMapVariable = None,
         mock: bool = False,
-        **kwargs,
     ):
         """
         This is the entry point when we do the actual execution of the function.
@@ -227,9 +226,7 @@ class BasePipelineExecutor(BaseExecutor):
         ...
 
     @abstractmethod
-    def execute_node(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         The entry point for all executors apart from local.
         We have already prepared for node execution.
@@ -244,9 +241,7 @@ class BasePipelineExecutor(BaseExecutor):
         ...
 
     @abstractmethod
-    def execute_from_graph(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
-    ):
+    def execute_from_graph(self, node: BaseNode, map_variable: TypeMapVariable = None):
         """
         This is the entry point to from the graph execution.
 
@@ -294,7 +289,7 @@ class BasePipelineExecutor(BaseExecutor):
         ...
 
     @abstractmethod
-    def execute_graph(self, dag: Graph, map_variable: TypeMapVariable = None, **kwargs):
+    def execute_graph(self, dag: Graph, map_variable: TypeMapVariable = None):
         """
         The parallelization is controlled by the nodes and not by this function.
 
@@ -395,7 +390,7 @@ class BasePipelineExecutor(BaseExecutor):
 
     @abstractmethod
     def trigger_node_execution(
-        self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
+        self, node: BaseNode, map_variable: TypeMapVariable = None
     ):
         """
         Executor specific way of triggering jobs when runnable does both traversal and execution

runnable/nodes.py

@@ -37,6 +37,7 @@ class BaseNode(ABC, BaseModel):
     internal_name: str = Field(exclude=True)
     internal_branch_name: str = Field(default="", exclude=True)
     is_composite: bool = Field(default=False, exclude=True)
+    is_distributed: bool = Field(default=False, exclude=True)
 
     @property
     def _context(self):
@@ -280,7 +281,6 @@ class BaseNode(ABC, BaseModel):
         mock=False,
         map_variable: TypeMapVariable = None,
         attempt_number: int = 1,
-        **kwargs,
     ) -> StepLog:
         """
         The actual function that does the execution of the command in the config.
@@ -299,7 +299,7 @@ class BaseNode(ABC, BaseModel):
         """
 
     @abstractmethod
-    def execute_as_graph(self, map_variable: TypeMapVariable = None, **kwargs):
+    def execute_as_graph(self, map_variable: TypeMapVariable = None):
         """
         This function would be called to set up the execution of the individual
         branches of a composite node.
@@ -314,7 +314,7 @@ class BaseNode(ABC, BaseModel):
         """
 
     @abstractmethod
-    def fan_out(self, map_variable: TypeMapVariable = None, **kwargs):
+    def fan_out(self, map_variable: TypeMapVariable = None):
         """
         This function would be called to set up the execution of the individual
         branches of a composite node.
@@ -330,7 +330,7 @@ class BaseNode(ABC, BaseModel):
         """
 
     @abstractmethod
-    def fan_in(self, map_variable: TypeMapVariable = None, **kwargs):
+    def fan_in(self, map_variable: TypeMapVariable = None):
         """
         This function would be called to tear down the execution of the individual
         branches of a composite node.
@@ -439,33 +439,25 @@ class ExecutableNode(TraversalNode):
             "This is an executable node and does not have branches"
         )
 
-    def execute_as_graph(self, map_variable: TypeMapVariable = None, **kwargs):
+    def execute_as_graph(self, map_variable: TypeMapVariable = None):
         raise exceptions.NodeMethodCallError(
             "This is an executable node and does not have a graph"
         )
 
-    def fan_in(self, map_variable: TypeMapVariable = None, **kwargs):
+    def fan_in(self, map_variable: TypeMapVariable = None):
         raise exceptions.NodeMethodCallError(
             "This is an executable node and does not have a fan in"
         )
 
-    def fan_out(self, map_variable: TypeMapVariable = None, **kwargs):
+    def fan_out(self, map_variable: TypeMapVariable = None):
         raise exceptions.NodeMethodCallError(
             "This is an executable node and does not have a fan out"
         )
 
-    def prepare_for_job_execution(self):
-        raise exceptions.NodeMethodCallError(
-            "This is an executable node and does not have a prepare_for_job_execution"
-        )
-
-    def tear_down_after_job_execution(self):
-        raise exceptions.NodeMethodCallError(
-            "This is an executable node and does not have a tear_down_after_job_execution",
-        )
-
 
 class CompositeNode(TraversalNode):
+    is_composite: bool = True
+
     def _get_catalog_settings(self) -> Dict[str, Any]:
         """
         If the node defines a catalog settings, return it or None
@@ -485,20 +477,44 @@ class CompositeNode(TraversalNode):
         mock=False,
         map_variable: TypeMapVariable = None,
         attempt_number: int = 1,
-        **kwargs,
     ) -> StepLog:
         raise exceptions.NodeMethodCallError(
             "This is a composite node and does not have an execute function"
         )
 
-    def prepare_for_job_execution(self):
+
+class DistributedNode(TraversalNode):
+    """
+    Use this node for distributed execution of tasks.
+    eg: torch distributed, horovod, etc.
+    """
+
+    is_distributed: bool = True
+    catalog: Optional[CatalogStructure] = Field(default=None)
+    max_attempts: int = Field(default=1, ge=1)
+
+    def _get_catalog_settings(self) -> Dict[str, Any]:
+        """
+        If the node defines a catalog settings, return it or None
+
+        Returns:
+            dict: catalog settings defined as per the node or None
+        """
+        if self.catalog:
+            return self.catalog.model_dump()
+        return {}
+
+    def _get_max_attempts(self) -> int:
+        return self.max_attempts
+
+    def _get_branch_by_name(self, branch_name: str):
         raise exceptions.NodeMethodCallError(
-            "This is an executable node and does not have a prepare_for_job_execution"
+            "This is an distributed node and does not have branches"
         )
 
-    def tear_down_after_job_execution(self):
+    def execute_as_graph(self, map_variable: TypeMapVariable = None):
         raise exceptions.NodeMethodCallError(
-            "This is an executable node and does not have a tear_down_after_job_execution"
+            "This is an executable node and does not have a graph"
         )
 
 
@@ -524,13 +540,16 @@ class TerminalNode(BaseNode):
     def _get_max_attempts(self) -> int:
         return 1
 
-    def execute_as_graph(self, map_variable: TypeMapVariable = None, **kwargs):
+    def execute_as_graph(self, map_variable: TypeMapVariable = None):
         raise exceptions.TerminalNodeError()
 
-    def fan_in(self, map_variable: TypeMapVariable = None, **kwargs):
+    def fan_in(self, map_variable: TypeMapVariable = None):
         raise exceptions.TerminalNodeError()
 
-    def fan_out(self, map_variable: TypeMapVariable = None, **kwargs):
+    def fan_out(
+        self,
+        map_variable: TypeMapVariable = None,
+    ):
         raise exceptions.TerminalNodeError()
 
     @classmethod