runnable 0.28.8__py3-none-any.whl → 0.30.0__py3-none-any.whl

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

runnable/sdk.py

@@ -5,7 +5,7 @@ import os
 import re
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from pydantic import (
     BaseModel,
@@ -34,17 +34,20 @@ from extensions.nodes.nodes import (
     SuccessNode,
     TaskNode,
 )
+from extensions.nodes.torch_config import TorchConfig
 from runnable import console, defaults, entrypoints, exceptions, graph, utils
 from runnable.executor import BaseJobExecutor, BasePipelineExecutor
 from runnable.nodes import TraversalNode
 from runnable.tasks import BaseTaskType as RunnableTask
 from runnable.tasks import TaskReturns
 
-# TODO: This might have to be an extension
-
 logger = logging.getLogger(defaults.LOGGER_NAME)
 
-StepType = Union["Stub", "PythonTask", "NotebookTask", "ShellTask", "Parallel", "Map"]
+StepType = Union[
+    "Stub", "PythonTask", "NotebookTask", "ShellTask", "Parallel", "Map", "Torch"
+]
+if TYPE_CHECKING:
+    from extensions.nodes.torch import TorchNode
 
 
 def pickled(name: str) -> TaskReturns:
@@ -322,7 +325,7 @@ class NotebookTask(BaseTask):
         catalog Optional[Catalog]: The files sync data from/to, refer to Catalog.
 
         secrets List[str]: List of secrets to pass to the task. They are exposed as environment variables
-        and removed after execution.
+            and removed after execution.
 
         overrides (Dict[str, Any]): Any overrides to the command.
             Individual tasks can override the global configuration config by referring to the
@@ -388,7 +391,7 @@ class ShellTask(BaseTask):
         catalog Optional[Catalog]: The files sync data from/to, refer to Catalog.
 
         secrets List[str]: List of secrets to pass to the task. They are exposed as environment variables
-        and removed after execution.
+            and removed after execution.
 
         overrides (Dict[str, Any]): Any overrides to the command.
             Individual tasks can override the global configuration config by referring to the
@@ -456,6 +459,43 @@ class Stub(BaseTraversal):
         return StubNode.parse_from_config(self.model_dump(exclude_none=True))
 
 
+class Torch(BaseTraversal, TorchConfig):
+    function: Callable = Field(exclude=True)
+    catalog: Optional[Catalog] = Field(default=None, alias="catalog")
+    overrides: Dict[str, Any] = Field(default_factory=dict, alias="overrides")
+    returns: List[Union[str, TaskReturns]] = Field(
+        default_factory=list, alias="returns"
+    )
+    secrets: List[str] = Field(default_factory=list)
+
+    @computed_field
+    def command_type(self) -> str:
+        return "python"
+
+    @computed_field
+    def command(self) -> str:
+        module = self.function.__module__
+        name = self.function.__name__
+
+        return f"{module}.{name}"
+
+    def create_node(self) -> TorchNode:
+        if not self.next_node:
+            if not (self.terminate_with_failure or self.terminate_with_success):
+                raise AssertionError(
+                    "A node not being terminated must have a user defined next node"
+                )
+
+        if self.on_failure:
+            self.on_failure = self.on_failure.steps[0].name  # type: ignore
+
+        from extensions.nodes.torch import TorchNode
+
+        return TorchNode.parse_from_config(
+            self.model_dump(exclude_none=True, by_alias=True)
+        )
+
+
 class Parallel(BaseTraversal):
     """
     A node that executes multiple branches in parallel.

runnable/secrets.py

@@ -29,7 +29,7 @@ class BaseSecrets(ABC, BaseModel):
         return context.run_context
 
     @abstractmethod
-    def get(self, name: str, **kwargs) -> str:
+    def get(self, name: str) -> str:
         """
         Return the secret by name.
 
@@ -53,7 +53,7 @@ class DoNothingSecretManager(BaseSecrets):
 
     service_name: str = "do-nothing"
 
-    def get(self, name: str, **kwargs) -> str:
+    def get(self, name: str) -> str:
         """
         If a name is provided, return None else return empty dict.
 
@@ -76,7 +76,7 @@ class EnvSecretsManager(BaseSecrets):
 
     service_name: str = "env-secrets"
 
-    def get(self, name: str, **kwargs) -> str:
+    def get(self, name: str) -> str:
         """
         If a name is provided, return None else return empty dict.
 

runnable/tasks.py

@@ -93,7 +93,6 @@ class BaseTaskType(BaseModel):
     def execute_command(
         self,
         map_variable: TypeMapVariable = None,
-        **kwargs,
     ) -> StepAttempt:
         """The function to execute the command.
 
@@ -271,7 +270,6 @@ class PythonTaskType(BaseTaskType):  # pylint: disable=too-few-public-methods
     def execute_command(
         self,
         map_variable: TypeMapVariable = None,
-        **kwargs,
     ) -> StepAttempt:
         """Execute the notebook as defined by the command."""
         attempt_log = StepAttempt(status=defaults.FAIL, start_time=str(datetime.now()))
@@ -441,7 +439,6 @@ class NotebookTaskType(BaseTaskType):
     def execute_command(
         self,
         map_variable: TypeMapVariable = None,
-        **kwargs,
     ) -> StepAttempt:
         """Execute the python notebook as defined by the command.
 
@@ -620,7 +617,6 @@ class ShellTaskType(BaseTaskType):
     def execute_command(
         self,
         map_variable: TypeMapVariable = None,
-        **kwargs,
     ) -> StepAttempt:
         # Using shell=True as we want to have chained commands to be executed in the same shell.
         """Execute the shell command as defined by the command.

{runnable-0.28.8.dist-info → runnable-0.30.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: runnable
-Version: 0.28.8
+Version: 0.30.0
 Summary: Add your description here
 Author-email: "Vammi, Vijay" <vijay.vammi@astrazeneca.com>
 License-File: LICENSE
@@ -26,6 +26,8 @@ Provides-Extra: notebook
 Requires-Dist: ploomber-engine>=0.0.33; extra == 'notebook'
 Provides-Extra: s3
 Requires-Dist: cloudpathlib[s3]; extra == 's3'
+Provides-Extra: torch
+Requires-Dist: torch>=2.6.0; extra == 'torch'
 Description-Content-Type: text/markdown
 
 

{runnable-0.28.8.dist-info → runnable-0.30.0.dist-info}/RECORD

@@ -8,37 +8,39 @@ extensions/catalog/pyproject.toml,sha256=lLNxY6v04c8I5QK_zKw_E6sJTArSJRA_V-79kta
 extensions/catalog/s3.py,sha256=Sw5t8_kVRprn3uGGJCiHn7M9zw1CLaCOFj6YErtfG0o,287
 extensions/job_executor/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 extensions/job_executor/__init__.py,sha256=E2R6GV5cZTlZdqA5SVJ6ajZFh4oruM0k8AKHkpOZ3W8,5772
-extensions/job_executor/k8s.py,sha256=Mlmuqm1D70gfj471KKZrQF1PeDSCpOiW1bd4vD0pQ-w,16243
+extensions/job_executor/k8s.py,sha256=erzw4UOsOf2JSOiQio5stgW_rMryAsIQSBd8wiL6nBY,16214
 extensions/job_executor/k8s_job_spec.yaml,sha256=7aFpxHdO_p6Hkc3YxusUOuAQTD1Myu0yTPX9DrhxbOg,1158
-extensions/job_executor/local.py,sha256=FvxTk0vyxdrbLOAyNkLyjvmmowypabWOSITQBK_ffVE,1907
-extensions/job_executor/local_container.py,sha256=hyFnpicCp3_87mZsW64P6KSVbz7XMLjwJUWVjeCJ0_I,6627
+extensions/job_executor/local.py,sha256=3v6F8SOaPbCfPVVmU07RFr1wgs8iC8WoSn6Evfi8o3M,2033
+extensions/job_executor/local_container.py,sha256=8-dLhzY34pOVjJ_x0VmeTwVvYkESXBnp4j-XLsSsgBk,6688
 extensions/job_executor/pyproject.toml,sha256=UIEgiCYHTXcRWSByNMFuKJFKgxTBpQqTqyUecIsb_Vc,286
 extensions/nodes/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-extensions/nodes/nodes.py,sha256=WdOmep4uxmY2mTOtsuVZ5QhYl96jqJprkG6jkIg7BVg,34774
+extensions/nodes/nodes.py,sha256=s9ub1dqy4qHjRQG6YElCdL7rCOTYNs9RUIrStZ6tEB4,28256
 extensions/nodes/pyproject.toml,sha256=YTu-ETN3JNFSkMzzWeOwn4m-O2nbRH-PmiPBALDCUw4,278
+extensions/nodes/torch.py,sha256=kB4a72YMcrxDDzbR5LffODtrdA7vUo9dRJlaVr8KEEM,5570
+extensions/nodes/torch_config.py,sha256=yDvDADpnLhQsNtfH8qIztLHQ2LhYiOJEWljxpH9GZzs,1222
 extensions/pipeline_executor/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-extensions/pipeline_executor/__init__.py,sha256=lk_QmbfzXNrgpF_KvMPuPpzxp0B8SJobDHWrK_0Q5FE,24359
-extensions/pipeline_executor/argo.py,sha256=nnlR_D6arQMUSgAevnW1RXeN48SoB1wVcEfQ4TBireY,34543
-extensions/pipeline_executor/local.py,sha256=H8s6AdML_9_f-vdGG_6k0y9FbLqAqvA1S_7xMNyARzY,1946
-extensions/pipeline_executor/local_container.py,sha256=HOT9I-cPDCvgy6_bzNEtl4jPhTyeYSn1GK7lplH3vDA,12515
-extensions/pipeline_executor/mocked.py,sha256=SuObJ6Myt7p8duW8sylIp1cYIAnFutsJW1avWaOUY3c,5798
+extensions/pipeline_executor/__init__.py,sha256=9ZMHcieSYdTiYyjSkc8eT8yhOlKEUFnrbrdbqdOgvP0,24195
+extensions/pipeline_executor/argo.py,sha256=svnje6l1mYfLhrHZf54pE_ljwxvH8Gt0ZXSLWLDs4aQ,37869
+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
 extensions/pipeline_executor/pyproject.toml,sha256=ykTX7srR10PBYb8LsIwEj8vIPPIEZQ5V_R7VYbZ-ido,291
-extensions/pipeline_executor/retry.py,sha256=KGenhWrLLmOQgzMvqloXHDRJyoNs91t05rRW8aLW6FA,6969
+extensions/pipeline_executor/retry.py,sha256=6ClFXJYtr0M6nWIZiI-mbUGshobOtVH_KADN8JCfvH0,6881
 extensions/run_log_store/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 extensions/run_log_store/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 extensions/run_log_store/any_path.py,sha256=0nN_LHbm2W6AHkerQmsVHq3EoybFQF8lxpCicacHo8Y,2861
 extensions/run_log_store/chunked_fs.py,sha256=wHMKcAx6uFI4OOTp7QWCdGq9WvEFesbLp9VxHZU28l0,3341
 extensions/run_log_store/chunked_minio.py,sha256=Itfkw4Ycf0uLCqxH3Uk_itmVgT7ipJp05yKfD22WBiY,4007
 extensions/run_log_store/file_system.py,sha256=hhrbhSnuzv8yzBr6DAu45NT8-sawPP86WA2-LY70vjw,2781
-extensions/run_log_store/generic_chunked.py,sha256=bsGgChTDZN3dSbLmLJ9SIpcvArzVmzhTVAOYZytAUNc,20483
+extensions/run_log_store/generic_chunked.py,sha256=EnhRxlqm1jG-Tdxul4sY8OeCX5fK9FY2v8DZanX9-5o,20455
 extensions/run_log_store/minio.py,sha256=omrKDSdRzmnVBg9xXkkdQb-icBIgBDRdpmwGRlMyCGk,3453
 extensions/run_log_store/pyproject.toml,sha256=YnmXsFvFG9uv_c0spLYBsNI_1sbktqxtHsOuClyvZ3g,288
 extensions/run_log_store/db/implementation_FF.py,sha256=euTnh0xzNF0e_DyfHQ4W-kG1AwTr8u7OuO3_cZkR5bM,5237
 extensions/run_log_store/db/integration_FF.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 extensions/secrets/README.md,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-extensions/secrets/dotenv.py,sha256=FbYYd_pVuJuVuIDIvXbzKuSSQ9GPq7xJXTDbJMTQbhM,1583
+extensions/secrets/dotenv.py,sha256=nADHXI6KJ_LUYOIe5EbtYH-21OBebSNVr0Pjb1GlZ7w,1573
 extensions/secrets/pyproject.toml,sha256=mLJNImNcBlbLKHh-0ugVWT9V83R4RibyyYDtBCSqVF4,282
-runnable/__init__.py,sha256=n14AnTUUEYxXlTJ6-YLT0tMmeFb7Co_3kNldV6pgKSs,662
+runnable/__init__.py,sha256=swvqdCjeddn40o4zjsluyahdVcU0r1arSRrxmRsvFEQ,673
 runnable/catalog.py,sha256=W_erYbLZ-ffuA9RQuWVqz1DUJOuWayf32ne32IDbAbc,4358
 runnable/cli.py,sha256=3BiKSj95h2Drn__YlchMPZ5rBMafuRb2OGIsVpbsO5Y,8788
 runnable/context.py,sha256=by5uepmuCP0dmM9BmsliXihSes5QEFejwAsmekcqylE,1388
@@ -46,18 +48,18 @@ runnable/datastore.py,sha256=ZobM1aVkgeUJ2fZYt63IFDsoNzObwc93hdByegS5YKQ,32396
 runnable/defaults.py,sha256=3o9IVGryyCE6PoQTOoaIaHHTbJGEzmdXMcwzOhwAYoI,3518
 runnable/entrypoints.py,sha256=cDbhtmLUWdBh9K6hNusfQpSd5NadcX8V1K2JEDf_YAg,18984
 runnable/exceptions.py,sha256=LFbp0-Qxg2PAMLEVt7w2whhBxSG-5pzUEv5qN-Rc4_c,3003
-runnable/executor.py,sha256=UCBBtyD0khl9QjT4SRTFMQDHDLWfJUC2U4_b3KQzaBE,15127
+runnable/executor.py,sha256=J8-Ri9nBZCb-ao6okePb9FUVlhAaPc0ojQ2l48-FUqc,15031
 runnable/graph.py,sha256=poQz5zcvq89ju_u5sYlunQLPbHnXTaUmjcvstPwvT4U,16536
 runnable/names.py,sha256=vn92Kv9ANROYSZX6Z4z1v_WA3WiEdIYmG6KEStBFZug,8134
-runnable/nodes.py,sha256=YU9u7r1ESzui1uVtJ1dgwdv1ozyJnF2k-MCFieT8CLI,17519
+runnable/nodes.py,sha256=d1eLttMAcV7CTwTEqOuNwZqItANoLUkXJ73Xp-srlyI,17811
 runnable/parameters.py,sha256=sT3DNGczivP9z7r4Cp_brbudg1z4J-zjmvrq3ppIrVs,5089
 runnable/pickler.py,sha256=ydJ_eti_U1F4l-YacFp7BWm6g5vTn04UXye25S1HVok,2684
-runnable/sdk.py,sha256=y3fbK4KAP6FiTF37ickVF9gqaZXrJIgSnHGOSG5ZWpI,33803
-runnable/secrets.py,sha256=PXcEJw-4WPzeWRLfsatcPPyr1zkqgHzdRWRcS9vvpvM,2354
-runnable/tasks.py,sha256=X6xijut7ffwpfYDcXoN6y0AcRVd7fWHs676DJ00Kma4,29134
+runnable/sdk.py,sha256=NZVQGaL4Zm2hwloRmqEgp8UPbBg9hY1abQGYnOgniPI,35128
+runnable/secrets.py,sha256=4L_dBFxTgr8r_hHUD6RlZEtqaOHDRsFG5PXO5wlvMI0,2324
+runnable/tasks.py,sha256=Qb1IhVxHv68E7vf3M3YCf7MGRHyjmsEEYBpEpiZ4mRI,29062
 runnable/utils.py,sha256=hBr7oGwGL2VgfITlQCTz-a1iwvvf7Mfl-HY8UdENZac,19929
-runnable-0.28.8.dist-info/METADATA,sha256=zA7aiT_kZ5NikkA0SRFmCC-83ZcjN883gHIQqVJXntw,10047
-runnable-0.28.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-runnable-0.28.8.dist-info/entry_points.txt,sha256=ioMbWojILtdibYVgh1jXJ00SpK-tX3gy7oVGDq61cSk,1839
-runnable-0.28.8.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-runnable-0.28.8.dist-info/RECORD,,
+runnable-0.30.0.dist-info/METADATA,sha256=mFY61nh2mT8DGjewQKOJIlNVxcZisIWsLZ6FlcvHvZU,10115
+runnable-0.30.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+runnable-0.30.0.dist-info/entry_points.txt,sha256=PrjKrlfXPZaV_7hz8orGu4FDnatLqnhPOXljyllszdw,1880
+runnable-0.30.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+runnable-0.30.0.dist-info/RECORD,,

{runnable-0.28.8.dist-info → runnable-0.30.0.dist-info}/entry_points.txt

@@ -21,6 +21,7 @@ parallel = extensions.nodes.nodes:ParallelNode
 stub = extensions.nodes.nodes:StubNode
 success = extensions.nodes.nodes:SuccessNode
 task = extensions.nodes.nodes:TaskNode
+torch = extensions.nodes.torch:TorchNode
 
 [pickler]
 pickle = runnable.pickler:NativePickler