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

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:
@@ -456,6 +459,45 @@ class Stub(BaseTraversal):
         return StubNode.parse_from_config(self.model_dump(exclude_none=True))
 
 
+class Torch(BaseTraversal, TorchConfig):
+    # Its a wrapper of a python task
+    # TODO: Is there a way to not sync these with the torch node in extensions?
+    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.7.dist-info → runnable-0.29.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: runnable
-Version: 0.28.7
+Version: 0.29.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.7.dist-info → runnable-0.29.0.dist-info}/RECORD

@@ -8,56 +8,59 @@ 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=raobGxwoqZN8c-yCsAa0CDuPLWKuyEttB37U5wsqGF4,1968
+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=LlXtzcbJyOossTNd-gdC4xrkPe9qYxRc0czdNzNQzlY,34497
+extensions/pipeline_executor/local.py,sha256=orBIhG8QJesA3YqWhsuczjhhlKD_1s62MRNerxv9_Tg,1858
+extensions/pipeline_executor/local_container.py,sha256=PvMXy-zFTnT9hj7jjz1VkaPVJ9Dkrhnd7Hs6eqLMXZ8,12398
+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
+extensions/tasks/torch.py,sha256=uNO4qYMawNH5hPecANCiSUQZnUC8yqw3-rxtM526CeA,1955
+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
 runnable/datastore.py,sha256=ZobM1aVkgeUJ2fZYt63IFDsoNzObwc93hdByegS5YKQ,32396
 runnable/defaults.py,sha256=3o9IVGryyCE6PoQTOoaIaHHTbJGEzmdXMcwzOhwAYoI,3518
-runnable/entrypoints.py,sha256=xkUa568-7x9xALz13qW14DxS1nnLDKwLwdIBJZG-vM0,18982
+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=F0gQjJ10VQSlilqaxYm1gRdXpjmx5DqBP0KCPmz85zg,15021
 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=6OO_vsRuGSjVhME2AJEljs0cjobjIQKa2E3mGoILkZA,35237
+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.7.dist-info/METADATA,sha256=JktO2ExsbIta-x_wCW4WKO5GbosxMbLYf6r8cyKws70,10047
-runnable-0.28.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-runnable-0.28.7.dist-info/entry_points.txt,sha256=ioMbWojILtdibYVgh1jXJ00SpK-tX3gy7oVGDq61cSk,1839
-runnable-0.28.7.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-runnable-0.28.7.dist-info/RECORD,,
+runnable-0.29.0.dist-info/METADATA,sha256=1jW9CmQUxqDClGKzNPZGMlq1B7M4y8p9FwQf0Cz1bZg,10115
+runnable-0.29.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+runnable-0.29.0.dist-info/entry_points.txt,sha256=PrjKrlfXPZaV_7hz8orGu4FDnatLqnhPOXljyllszdw,1880
+runnable-0.29.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+runnable-0.29.0.dist-info/RECORD,,

{runnable-0.28.7.dist-info → runnable-0.29.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