runnable 0.33.1__py3-none-any.whl → 0.34.0__py3-none-any.whl

runnable/sdk.py

@@ -34,7 +34,6 @@ from extensions.nodes.nodes import (
     SuccessNode,
     TaskNode,
 )
-from extensions.tasks.torch_config import TorchConfig
 from runnable import console, defaults, entrypoints, exceptions, graph, utils
 from runnable.executor import BaseJobExecutor, BasePipelineExecutor
 from runnable.nodes import TraversalNode
@@ -44,7 +43,13 @@ from runnable.tasks import TaskReturns
 logger = logging.getLogger(defaults.LOGGER_NAME)
 
 StepType = Union[
-    "Stub", "PythonTask", "NotebookTask", "ShellTask", "Parallel", "Map", "TorchTask"
+    "Stub",
+    "PythonTask",
+    "NotebookTask",
+    "ShellTask",
+    "Parallel",
+    "Map",
+    "TorchTask",
 ]
 
 
@@ -189,36 +194,6 @@ class BaseTask(BaseTraversal):
         )
 
 
-class TorchTask(BaseTask, TorchConfig):
-    # The user will not know the rnnz variables for multi node
-    # They should be overridden in the environment
-    function: Callable = Field(exclude=True)
-
-    @field_validator("returns", mode="before")
-    @classmethod
-    def serialize_returns(
-        cls, returns: List[Union[str, TaskReturns]]
-    ) -> List[TaskReturns]:
-        assert len(returns) == 0, "Torch tasks cannot return any variables"
-        return []
-
-    @computed_field
-    def command_type(self) -> str:
-        return "torch"
-
-    @computed_field
-    def command(self) -> str:
-        module = self.function.__module__
-        name = self.function.__name__
-
-        return f"{module}.{name}"
-
-    def create_job(self) -> RunnableTask:
-        self.terminate_with_success = True
-        node = self.create_node()
-        return node.executable
-
-
 class PythonTask(BaseTask):
     """
     An execution node of the pipeline of python functions.
@@ -307,6 +282,27 @@ class PythonTask(BaseTask):
         return node.executable
 
 
+class TorchTask(BaseTask):
+    # entrypoint: str = Field(
+    #     alias="entrypoint", default="torch.distributed.run", frozen=True
+    # )
+    # args_to_torchrun: Dict[str, Any] = Field(
+    #     default_factory=dict, alias="args_to_torchrun"
+    # )
+
+    script_to_call: str
+    accelerate_config_file: str
+
+    @computed_field
+    def command_type(self) -> str:
+        return "torch"
+
+    def create_job(self) -> RunnableTask:
+        self.terminate_with_success = True
+        node = self.create_node()
+        return node.executable
+
+
 class NotebookTask(BaseTask):
     """
     An execution node of the pipeline of notebook.
@@ -978,7 +974,6 @@ class PythonJob(BaseJob):
 
         return f"{module}.{name}"
 
-    # TODO: can this be simplified to just self.model_dump(exclude_none=True)?
     def get_task(self) -> RunnableTask:
         # Piggy bank on existing tasks as a hack
         task = PythonTask(
@@ -989,9 +984,15 @@ class PythonJob(BaseJob):
         return task.create_node().executable
 
 
-class TorchJob(BaseJob, TorchConfig):
-    function: Callable = Field()
-    # min and max should always be 1
+class TorchJob(BaseJob):
+    # entrypoint: str = Field(default="torch.distributed.run", frozen=True)
+    # args_to_torchrun: dict[str, str | bool | int | float] = Field(
+    #     default_factory=dict
+    # )  # For example
+    # {"nproc_per_node": 2, "nnodes": 1,}
+
+    script_to_call: str  # For example train/script.py
+    accelerate_config_file: str
 
     def get_task(self) -> RunnableTask:
         # Piggy bank on existing tasks as a hack

runnable/tasks.py

@@ -354,6 +354,66 @@ class PythonTaskType(BaseTaskType):  # pylint: disable=too-few-public-methods
         return attempt_log
 
 
+class TorchTaskType(BaseTaskType):
+    task_type: str = Field(default="torch", serialization_alias="command_type")
+    accelerate_config_file: str
+
+    script_to_call: str  # For example train/script.py
+
+    def execute_command(
+        self, map_variable: Dict[str, str | int | float] | None = None
+    ) -> StepAttempt:
+        from accelerate.commands import launch
+
+        attempt_log = StepAttempt(status=defaults.FAIL, start_time=str(datetime.now()))
+
+        with (
+            self.execution_context(
+                map_variable=map_variable, allow_complex=False
+            ) as params,
+            self.expose_secrets() as _,
+        ):
+            try:
+                script_args = []
+                for key, value in params.items():
+                    script_args.append(f"--{key}")
+                    if type(value.value) is not bool:
+                        script_args.append(str(value.value))
+
+                # TODO: Check the typing here
+
+                logger.info("Calling the user script with the following parameters:")
+                logger.info(script_args)
+                out_file = TeeIO()
+                try:
+                    with contextlib.redirect_stdout(out_file):
+                        parser = launch.launch_command_parser()
+                        args = parser.parse_args(self.script_to_call)
+                        args.training_script = self.script_to_call
+                        args.config_file = self.accelerate_config_file
+                        args.training_script_args = script_args
+
+                        launch.launch_command(args)
+                    task_console.print(out_file.getvalue())
+                except Exception as e:
+                    raise exceptions.CommandCallError(
+                        f"Call to script{self.script_to_call} did not succeed."
+                    ) from e
+                finally:
+                    sys.argv = sys.argv[:1]
+
+                attempt_log.status = defaults.SUCCESS
+            except Exception as _e:
+                msg = f"Call to script: {self.script_to_call} did not succeed."
+                attempt_log.message = msg
+                task_console.print_exception(show_locals=False)
+                task_console.log(_e, style=defaults.error_style)
+
+        attempt_log.end_time = str(datetime.now())
+
+        return attempt_log
+
+
 class NotebookTaskType(BaseTaskType):
     """
     --8<-- [start:notebook_reference]
@@ -747,6 +807,31 @@ class ShellTaskType(BaseTaskType):
         return attempt_log
 
 
+def convert_binary_to_string(data):
+    """
+    Recursively converts 1 and 0 values in a nested dictionary to "1" and "0".
+
+    Args:
+        data (dict or any): The input data (dictionary, list, or other).
+
+    Returns:
+        dict or any: The modified data with binary values converted to strings.
+    """
+
+    if isinstance(data, dict):
+        for key, value in data.items():
+            data[key] = convert_binary_to_string(value)
+        return data
+    elif isinstance(data, list):
+        return [convert_binary_to_string(item) for item in data]
+    elif data == 1:
+        return "1"
+    elif data == 0:
+        return "0"
+    else:
+        return data  # Return other values unchanged
+
+
 def create_task(kwargs_for_init) -> BaseTaskType:
     """
     Creates a task object from the command configuration.
@@ -763,6 +848,8 @@ def create_task(kwargs_for_init) -> BaseTaskType:
     kwargs = kwargs_for_init.copy()
     command_type = kwargs.pop("command_type", defaults.COMMAND_TYPE)
 
+    kwargs = convert_binary_to_string(kwargs)
+
     try:
         task_mgr = driver.DriverManager(
             namespace="tasks",

{runnable-0.33.1.dist-info → runnable-0.34.0.dist-info}/METADATA

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

{runnable-0.33.1.dist-info → runnable-0.34.0.dist-info}/RECORD

@@ -56,12 +56,12 @@ runnable/names.py,sha256=vn92Kv9ANROYSZX6Z4z1v_WA3WiEdIYmG6KEStBFZug,8134
 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=xjMEDg1Uovrp0Kw7RR4GAQIJgG9zcIgueqOzQAaO7Bs,35363
+runnable/sdk.py,sha256=-hsoZctbGKsrfOQW3Z7RqWVGJI4GhbsOjqjMRb2OAUo,35181
 runnable/secrets.py,sha256=4L_dBFxTgr8r_hHUD6RlZEtqaOHDRsFG5PXO5wlvMI0,2324
-runnable/tasks.py,sha256=ABRhgiTY8F62pNlqJmVTDjwJwuzp8DqciUEOq1fpt1U,28989
+runnable/tasks.py,sha256=lOtCninvosGI2bNIzblrzNa-lN7TMwel1KQ1g23M85A,32088
 runnable/utils.py,sha256=hBr7oGwGL2VgfITlQCTz-a1iwvvf7Mfl-HY8UdENZac,19929
-runnable-0.33.1.dist-info/METADATA,sha256=Cgy39IH3KEP3b2eEVF8u934b_AQ4BWzzQqjb2BV5aKw,10168
-runnable-0.33.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-runnable-0.33.1.dist-info/entry_points.txt,sha256=uWHbbOSj0jlG54tFHw377xKkfVbjWvb_1Y9L_LgjJ0Q,1925
-runnable-0.33.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-runnable-0.33.1.dist-info/RECORD,,
+runnable-0.34.0.dist-info/METADATA,sha256=E_O8YUEotnppM6EG006CSRPs3XAuwspYAwrgNe5UgRc,10166
+runnable-0.34.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+runnable-0.34.0.dist-info/entry_points.txt,sha256=wKfW6aIWMQFlwrwpPBVWlMQDcxQmOupDKNkKyXoPFV4,1917
+runnable-0.34.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+runnable-0.34.0.dist-info/RECORD,,

{runnable-0.33.1.dist-info → runnable-0.34.0.dist-info}/entry_points.txt

@@ -49,4 +49,4 @@ env-secrets = runnable.secrets:EnvSecretsManager
 notebook = runnable.tasks:NotebookTaskType
 python = runnable.tasks:PythonTaskType
 shell = runnable.tasks:ShellTaskType
-torch = extensions.tasks.torch:TorchTaskType
+torch = runnable.tasks:TorchTaskType