runnable 0.34.0a2__py3-none-any.whl → 0.34.0a3__py3-none-any.whl

runnable/sdk.py

@@ -283,14 +283,15 @@ class PythonTask(BaseTask):
 
 
 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"
-    )
+    # 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:
@@ -984,13 +985,14 @@ class PythonJob(BaseJob):
 
 
 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
+    # 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

@@ -5,7 +5,6 @@ import io
 import json
 import logging
 import os
-import runpy
 import subprocess
 import sys
 from datetime import datetime
@@ -357,16 +356,15 @@ class PythonTaskType(BaseTaskType):  # pylint: disable=too-few-public-methods
 
 class TorchTaskType(BaseTaskType):
     task_type: str = Field(default="torch", serialization_alias="command_type")
-
-    entrypoint: str = Field(default="torch.distributed.run", frozen=True)
-    args_to_torchrun: dict[str, str | bool] = Field(default_factory=dict)  # For example
-    # {"nproc_per_node": 2, "nnodes": 1,}
+    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 (
@@ -376,39 +374,37 @@ class TorchTaskType(BaseTaskType):
             self.expose_secrets() as _,
         ):
             try:
-                entry_point_args = [self.entrypoint]
-
-                for key, value in self.args_to_torchrun.items():
-                    entry_point_args.append(f"--{key}")
-                    if type(value) is not bool:
-                        entry_point_args.append(str(value))
-
-                entry_point_args.append(self.script_to_call)
+                script_args = []
                 for key, value in params.items():
-                    entry_point_args.append(f"--{key}")
-                    if type(value.value) is not bool:  # type: ignore
-                        entry_point_args.append(str(value.value))  # type: ignore
+                    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(entry_point_args)
+                logger.info(script_args)
                 out_file = TeeIO()
                 try:
                     with contextlib.redirect_stdout(out_file):
-                        sys.argv = entry_point_args
-                        runpy.run_module(self.entrypoint, run_name="__main__")
+                        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 entrypoint {self.entrypoint} with {self.script_to_call} did not succeed."
+                        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 entrypoint {self.entrypoint} with {self.script_to_call} did not succeed."
+                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)

{runnable-0.34.0a2.dist-info → runnable-0.34.0a3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: runnable
-Version: 0.34.0a2
+Version: 0.34.0a3
 Summary: Add your description here
 Author-email: "Vammi, Vijay" <vijay.vammi@astrazeneca.com>
 License-File: LICENSE
@@ -15,6 +15,7 @@ Requires-Dist: rich>=13.9.4
 Requires-Dist: ruamel-yaml>=0.18.6
 Requires-Dist: setuptools>=75.6.0
 Requires-Dist: stevedore>=5.4.0
+Requires-Dist: torchvision>=0.21.0
 Requires-Dist: typer>=0.15.1
 Provides-Extra: docker
 Requires-Dist: docker>=7.1.0; extra == 'docker'
@@ -27,8 +28,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.34.0a2.dist-info → runnable-0.34.0a3.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=Cl6wVJj_pBnHmcszf-kh4nVqbiQaIruGJn06cm9epm4,35097
+runnable/sdk.py,sha256=-hsoZctbGKsrfOQW3Z7RqWVGJI4GhbsOjqjMRb2OAUo,35181
 runnable/secrets.py,sha256=4L_dBFxTgr8r_hHUD6RlZEtqaOHDRsFG5PXO5wlvMI0,2324
-runnable/tasks.py,sha256=OW9pzjEKMRFpB256KJm__jWwsF37gs-tkIUcfnOTJwA,32382
+runnable/tasks.py,sha256=lOtCninvosGI2bNIzblrzNa-lN7TMwel1KQ1g23M85A,32088
 runnable/utils.py,sha256=hBr7oGwGL2VgfITlQCTz-a1iwvvf7Mfl-HY8UdENZac,19929
-runnable-0.34.0a2.dist-info/METADATA,sha256=DzGQTVqxRAN95MoyRc5TQXG_OC85uf6PH5NGtru3qSg,10170
-runnable-0.34.0a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-runnable-0.34.0a2.dist-info/entry_points.txt,sha256=wKfW6aIWMQFlwrwpPBVWlMQDcxQmOupDKNkKyXoPFV4,1917
-runnable-0.34.0a2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-runnable-0.34.0a2.dist-info/RECORD,,
+runnable-0.34.0a3.dist-info/METADATA,sha256=AYMw1jtTzhBN_Y2dMJiguAnYwc82LLxa-WHYApUYpCs,10203
+runnable-0.34.0a3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+runnable-0.34.0a3.dist-info/entry_points.txt,sha256=wKfW6aIWMQFlwrwpPBVWlMQDcxQmOupDKNkKyXoPFV4,1917
+runnable-0.34.0a3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+runnable-0.34.0a3.dist-info/RECORD,,