runnable 0.32.2__py3-none-any.whl → 0.32.3__py3-none-any.whl

extensions/nodes/torch.py

@@ -5,14 +5,14 @@ import random
 import string
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import Any, Callable, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
 from extensions.nodes.torch_config import EasyTorchConfig, TorchConfig
 from runnable import PythonJob, datastore, defaults
 from runnable.datastore import StepLog
-from runnable.nodes import DistributedNode
+from runnable.nodes import ExecutableNode
 from runnable.tasks import PythonTaskType, create_task
 from runnable.utils import TypeMapVariable
 
@@ -23,10 +23,7 @@ try:
     from torch.distributed.launcher.api import LaunchConfig, elastic_launch
 except ImportError:
     logger.exception("Torch is not installed. Please install torch first.")
-
-if TYPE_CHECKING:
-    from torch.distributed.elastic.multiprocessing.api import DefaultLogsSpecs, Std
-    from torch.distributed.launcher.api import LaunchConfig, elastic_launch
+    raise Exception("Torch is not installed. Please install torch first.")
 
 
 def training_subprocess():
@@ -120,7 +117,7 @@ def delete_env_vars_with_prefix(prefix):
 
 
 # TODO: The design of this class is not final
-class TorchNode(DistributedNode, TorchConfig):
+class TorchNode(ExecutableNode, TorchConfig):
     node_type: str = Field(default="torch", serialization_alias="type")
     executable: PythonTaskType = Field(exclude=True)
 

extensions/pipeline_executor/argo.py

@@ -20,14 +20,10 @@ from pydantic import (
 from pydantic.alias_generators import to_camel
 from ruamel.yaml import YAML
 
-from extensions.nodes.nodes import (
-    MapNode,
-    ParallelNode,
-    StubNode,
-    SuccessNode,
-    TaskNode,
-)
-from extensions.nodes.torch import TorchNode
+from extensions.nodes.nodes import MapNode, ParallelNode, TaskNode
+
+# TODO: Should be part of a wider refactor
+# from extensions.nodes.torch import TorchNode
 from extensions.pipeline_executor import GenericPipelineExecutor
 from runnable import defaults, utils
 from runnable.defaults import TypeMapVariable
@@ -590,12 +586,7 @@ class ArgoExecutor(GenericPipelineExecutor):
         task_name: str,
         inputs: Optional[Inputs] = None,
     ) -> ContainerTemplate:
-        assert (
-            isinstance(node, TaskNode)
-            or isinstance(node, StubNode)
-            or isinstance(node, SuccessNode)
-            or isinstance(node, TorchNode)
-        )
+        assert node.node_type in ["task", "torch", "success", "stub", "fail"]
 
         node_override = None
         if hasattr(node, "overrides"):
@@ -658,7 +649,7 @@ class ArgoExecutor(GenericPipelineExecutor):
     def _set_env_vars_to_task(
         self, working_on: BaseNode, container_template: CoreContainerTemplate
     ):
-        if not (isinstance(working_on, TaskNode) or isinstance(working_on, TorchNode)):
+        if working_on.node_type not in ["task", "torch"]:
             return
 
         global_envs: dict[str, str] = {}
@@ -667,7 +658,7 @@ class ArgoExecutor(GenericPipelineExecutor):
             env_var = cast(EnvVar, env_var)
             global_envs[env_var.name] = env_var.value
 
-        override_key = working_on.overrides.get(self.service_name, "")
+        override_key = working_on.overrides.get(self.service_name, "")  # type: ignore
         node_override = self.overrides.get(override_key, None)
 
         # Update the global envs with the node overrides
@@ -878,6 +869,8 @@ class ArgoExecutor(GenericPipelineExecutor):
                     self._templates.append(composite_template)
 
                 case "torch":
+                    from extensions.nodes.torch import TorchNode
+
                     assert isinstance(working_on, TorchNode)
                     # TODO: Need to add multi-node functionality
                     # Check notes on the torch node

extensions/tasks/torch.py

@@ -5,7 +5,7 @@ import random
 import string
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, field_serializer, model_validator
 from ruamel.yaml import YAML
@@ -19,16 +19,15 @@ from runnable.utils import get_module_and_attr_names
 
 logger = logging.getLogger(defaults.LOGGER_NAME)
 
+logger = logging.getLogger(defaults.LOGGER_NAME)
+
 try:
     from torch.distributed.elastic.multiprocessing.api import DefaultLogsSpecs, Std
     from torch.distributed.launcher.api import LaunchConfig, elastic_launch
 
-except ImportError:
+except ImportError as e:
     logger.exception("torch is not installed")
-
-if TYPE_CHECKING:
-    from torch.distributed.elastic.multiprocessing.api import DefaultLogsSpecs, Std
-    from torch.distributed.launcher.api import LaunchConfig, elastic_launch
+    raise Exception("torch is not installed") from e
 
 
 def get_min_max_nodes(nnodes: str) -> tuple[int, int]:
@@ -69,6 +68,7 @@ class TorchTaskType(BaseTaskType, TorchConfig):
             )
         )
         print("###", easy_torch_config)
+        print("###", easy_torch_config)
         launch_config = LaunchConfig(
             **easy_torch_config.model_dump(
                 exclude_none=True,
@@ -96,6 +96,28 @@ class TorchTaskType(BaseTaskType, TorchConfig):
 
         _, max_nodes = get_min_max_nodes(self.nnodes)
 
+        if max_nodes > 1 and not is_execute:
+            executor = self._context.executor
+            executor.scale_up(self)
+            return StepAttempt(
+                status=defaults.SUCCESS,
+                start_time=str(datetime.now()),
+                end_time=str(datetime.now()),
+                attempt_number=1,
+                message="Triggered a scale up",
+            )
+
+        # The below should happen only if we are in the node that we want to execute
+        # For a single node, multi worker setup, this should be the entry point
+        # For a multi-node, we need to:
+        # - create a service config
+        # - Create a stateful set with number of nodes
+        # - Create a job to run the torch.distributed.launcher.api.elastic_launch on every node
+        # - the entry point to runnnable could be a way to trigger execution instead of scaling
+        is_execute = os.environ.get("RUNNABLE_TORCH_EXECUTE", "true") == "true"
+
+        _, max_nodes = get_min_max_nodes(self.nnodes)
+
         if max_nodes > 1 and not is_execute:
             executor = self._context.executor
             executor.scale_up(self)
@@ -109,6 +131,7 @@ class TorchTaskType(BaseTaskType, TorchConfig):
 
         launch_config = self._get_launch_config()
         print("###****", launch_config)
+        print("###****", launch_config)
         logger.info(f"launch_config: {launch_config}")
 
         # ENV variables are shared with the subprocess, use that as communication

runnable/nodes.py

@@ -36,8 +36,8 @@ class BaseNode(ABC, BaseModel):
     name: str
     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):
@@ -483,41 +483,6 @@ class CompositeNode(TraversalNode):
         )
 
 
-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 distributed node and does not have branches"
-        )
-
-    def execute_as_graph(self, map_variable: TypeMapVariable = None):
-        raise exceptions.NodeMethodCallError(
-            "This is an executable node and does not have a graph"
-        )
-
-
 class TerminalNode(BaseNode):
     def _get_on_failure_node(self) -> str:
         return ""

{runnable-0.32.2.dist-info → runnable-0.32.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: runnable
-Version: 0.32.2
+Version: 0.32.3
 Summary: Add your description here
 Author-email: "Vammi, Vijay" <vijay.vammi@astrazeneca.com>
 License-File: LICENSE

{runnable-0.32.2.dist-info → runnable-0.32.3.dist-info}/RECORD

@@ -16,11 +16,11 @@ extensions/job_executor/pyproject.toml,sha256=UIEgiCYHTXcRWSByNMFuKJFKgxTBpQqTqy
 extensions/nodes/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 extensions/nodes/nodes.py,sha256=s9ub1dqy4qHjRQG6YElCdL7rCOTYNs9RUIrStZ6tEB4,28256
 extensions/nodes/pyproject.toml,sha256=YTu-ETN3JNFSkMzzWeOwn4m-O2nbRH-PmiPBALDCUw4,278
-extensions/nodes/torch.py,sha256=gydcRX5C7jEdPnxLsAQkpRD_by_0Lp4dFg96xDkRVW0,9510
+extensions/nodes/torch.py,sha256=64DTjdPNSJ8vfMwUN9h9Ly5g9qj-Bga7LSGrfCAO0BY,9389
 extensions/nodes/torch_config.py,sha256=tO3sG2_fj8a6FmPZZllwKVx3WaRr4QmQYcACseg8YXM,2839
 extensions/pipeline_executor/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 extensions/pipeline_executor/__init__.py,sha256=wfigTL2T9OHrmE8b2Ydmb8h6hr-oF--Yc2FectC7WaY,24623
-extensions/pipeline_executor/argo.py,sha256=AEGSWVZulBL6EsvbVCaeBeTl2m_t5ymc6RFpMKhivis,37946
+extensions/pipeline_executor/argo.py,sha256=lHM3TM_UnQc4I1ghkuYdeBLpyr4pBLg-Ubnaf55Zw54,37878
 extensions/pipeline_executor/local.py,sha256=6oWUJ6b6NvIkpeQJBoCT1hbfX4_6WCB4HzMgHZ4ik1A,1887
 extensions/pipeline_executor/local_container.py,sha256=3kZ2QCsrq_YjH9dcAz8v05knKShQ_JtbIU-IA_-G538,12724
 extensions/pipeline_executor/mocked.py,sha256=0sMmypuvstBIv9uQg-WAcPrF3oOFpeEXNi6N8Nzdnl0,5680
@@ -40,7 +40,7 @@ extensions/run_log_store/db/integration_FF.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeR
 extensions/secrets/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 extensions/secrets/dotenv.py,sha256=nADHXI6KJ_LUYOIe5EbtYH-21OBebSNVr0Pjb1GlZ7w,1573
 extensions/secrets/pyproject.toml,sha256=mLJNImNcBlbLKHh-0ugVWT9V83R4RibyyYDtBCSqVF4,282
-extensions/tasks/torch.py,sha256=At2eMpJas4sUUjzJfPrEBGamG-k3MsxXU6Bou0h9BEs,9274
+extensions/tasks/torch.py,sha256=oeXRkmuttFIAuBwH7-h4SOVXMDOZXX5mvqI2aFrR3Vo,10283
 extensions/tasks/torch_config.py,sha256=UjfMitT-TXASRDGR30I2vDRnyk7JQnR-5CsOVidjpSY,2833
 runnable/__init__.py,sha256=3ZKuvGEkY_zHVQlJtarXd4jkjICxjgnw-bbKN_5SiJI,691
 runnable/catalog.py,sha256=4msQxLhLKlsDDrHFnGauPYe-Or-q9g8_RYCn_4dpxaU,4466
@@ -53,15 +53,15 @@ runnable/exceptions.py,sha256=LFbp0-Qxg2PAMLEVt7w2whhBxSG-5pzUEv5qN-Rc4_c,3003
 runnable/executor.py,sha256=Jr9yJtSH7CzjXJLWx3VWIUAQblstuGqzpFtajv7d39M,15348
 runnable/graph.py,sha256=poQz5zcvq89ju_u5sYlunQLPbHnXTaUmjcvstPwvT4U,16536
 runnable/names.py,sha256=vn92Kv9ANROYSZX6Z4z1v_WA3WiEdIYmG6KEStBFZug,8134
-runnable/nodes.py,sha256=d1eLttMAcV7CTwTEqOuNwZqItANoLUkXJ73Xp-srlyI,17811
+runnable/nodes.py,sha256=QGHMznriEz4AcmntHICBZKrDT6zbc7WD1sV0MgwK10c,16691
 runnable/parameters.py,sha256=u77CdqqDAbVdzNeBFPNUfGnWPy9-SpBVmwEJ56xmDm8,5289
 runnable/pickler.py,sha256=ydJ_eti_U1F4l-YacFp7BWm6g5vTn04UXye25S1HVok,2684
 runnable/sdk.py,sha256=hwsEGCCFSijm0DZwDJGHmV8jdMuSU_3Pf-vYoomWYHw,35084
 runnable/secrets.py,sha256=4L_dBFxTgr8r_hHUD6RlZEtqaOHDRsFG5PXO5wlvMI0,2324
 runnable/tasks.py,sha256=ABRhgiTY8F62pNlqJmVTDjwJwuzp8DqciUEOq1fpt1U,28989
 runnable/utils.py,sha256=hBr7oGwGL2VgfITlQCTz-a1iwvvf7Mfl-HY8UdENZac,19929
-runnable-0.32.2.dist-info/METADATA,sha256=fcKKBj2v2AhRQFZ7ALqSdJrKF5r0Wg-QV6HVKqkBpRY,10168
-runnable-0.32.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-runnable-0.32.2.dist-info/entry_points.txt,sha256=uWHbbOSj0jlG54tFHw377xKkfVbjWvb_1Y9L_LgjJ0Q,1925
-runnable-0.32.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-runnable-0.32.2.dist-info/RECORD,,
+runnable-0.32.3.dist-info/METADATA,sha256=l0jxi_VKPXblTa_Kd-fnqKophmB2e8x1Dj1HDbJV570,10168
+runnable-0.32.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+runnable-0.32.3.dist-info/entry_points.txt,sha256=uWHbbOSj0jlG54tFHw377xKkfVbjWvb_1Y9L_LgjJ0Q,1925
+runnable-0.32.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+runnable-0.32.3.dist-info/RECORD,,