Support Running Any Command in PyTorchDistributed Runtime (#513)

Author: qiuosier

ads/jobs/templates/driver_pytorch.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2023 Oracle and/or its affiliates.
+# Copyright (c) 2023, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 """This module requires oracle-ads>=2.6.8
 """
@@ -341,13 +341,15 @@ def prepare_cmd(self, launch_args: list = None, prefix=""):
             launch_args = []
         # Append launch cmd args specified by the user.
         if self.launch_cmd:
-            if not self.launch_cmd.startswith(self.LAUNCHER):
-                raise ValueError(
-                    f"Command not supported: '{self.launch_cmd}'. "
-                    f"The command should start with '{self.LAUNCHER}'."
-                )
+            if self.LAUNCHER:
+                if not self.launch_cmd.startswith(self.LAUNCHER):
+                    raise ValueError(
+                        f"Command not supported: '{self.launch_cmd}'. "
+                    )
 
-            launch_args.append(self.launch_cmd[len(self.LAUNCHER) + 1 :])
+                launch_args.append(self.launch_cmd[len(self.LAUNCHER) + 1 :])
+            else:
+                launch_args.append(self.launch_cmd)
         else:
             launch_args.append(self.get_cmd_with_entrypoint_and_args())
 
@@ -673,7 +675,51 @@ def run(self):
             self.run_deepspeed_worker()
 
 
+class GenericRunner(TorchRunner, DeepSpeedRunner):
+    """Runner for running command other than ``torchrun``, ``deepspeed`` or ``accelerate``."""
+
+    LAUNCHER = ""
+
+    def use_deepspeed(self) -> bool:
+        """Indicate if DeepSpeed is used."""
+        if os.environ.get(CONST_ENV_DEEPSPEED):
+            return True
+        return False
+
+    def set_env_var(self):
+        """Set default environment variables."""
+        defaults = {
+            "WORLD_SIZE": self.node_count,
+            "MASTER_ADDR": self.host_ip,
+            "MASTER_PORT": self.RDZV_PORT,
+        }
+        for k, v in defaults.items():
+            if k not in os.environ:
+                os.environ[k] = str(v)
+
+    def run(self):
+        """Runs the user's command.
+        Note that for TorchRunner or DeepSpeedRunner,
+        we automatically add arguments for some settings,
+        like the number of nodes and the host node address.
+
+        This generic runner does not modify the command specified by the user.
+        User needs to make sure the command can work on all nodes.
+        User may use the environment variables in the command.
+        """
+        self.set_env_var()
+        if self.use_deepspeed():
+            if self.is_host:
+                self.run_deepspeed_host()
+            else:
+                self.run_deepspeed_worker()
+        else:
+            self.time_cmd(cmd=self.prepare_cmd(prefix=self.env_ld_preload()))
+
+
 class AccelerateRunner(TorchRunner, DeepSpeedRunner):
+    """Runner for HuggingFace Accelerate."""
+
     # accelerate launch will add main_process_port for deepspeed cmd even if it is not needed.
     # https://github.com/huggingface/accelerate/blob/70920895e80f78d96d8f91e0beeb3ebdb8e5e5d6/src/accelerate/utils/launch.py#L233
     DEFAULT_ARGS = [
@@ -704,11 +750,18 @@ def __init__(self, code_dir: str = driver_utils.DEFAULT_CODE_DIR) -> None:
         self.main_process_ip = None
 
     def use_deepspeed(self):
-        return os.environ.get(CONST_ENV_DEEPSPEED) or self.launch_cmd_contains(
+        """Indicate if DeepSpeed is used."""
+        # Accelerate support using DeepSpeed by adding the "--use_deepspeed" argument.
+        if os.environ.get(CONST_ENV_DEEPSPEED) or self.launch_cmd_contains(
             "use_deepspeed"
-        )
+        ):
+            return True
+        return False
 
     def accelerate_args(self):
+        """Gets the default arguments for the accelerate command.
+        The value of the default arguments are assigned in ``__init__()``.
+        """
         args = []
         for arg in self.DEFAULT_ARGS:
             arg_val = getattr(self, arg, None)
@@ -720,6 +773,7 @@ def accelerate_args(self):
         return args
 
     def run_with_torchrun(self):
+        """Runs the job with torchrun."""
         launch_args = self.accelerate_args()
         for arg in self.TORCHRUN_ARGS:
             if not self.launch_cmd_contains(arg):
@@ -728,6 +782,7 @@ def run_with_torchrun(self):
         self.time_cmd(cmd=cmd)
 
     def run_with_deepspeed(self):
+        """Runs the job with DeepSpeed."""
         if self.is_host:
             launch_args = self.accelerate_args()
             if self.num_machines > 1:
@@ -758,6 +813,8 @@ def main():
         runner_class = DeepSpeedRunner
     elif launch_cmd.startswith("accelerate "):
         runner_class = AccelerateRunner
+    else:
+        runner_class = GenericRunner
 
     runner = runner_class()
     runner: Runner

ads/jobs/templates/driver_utils.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2023 Oracle and/or its affiliates.
+# Copyright (c) 2023, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 import contextlib
 import importlib
@@ -397,6 +397,7 @@ def run_command(
             shell=True,
         )
         # Stream the outputs
+        logger.debug("Streaming command output from subprocess %s", process.pid)
         while True:
             output = process.stdout.readline()
             if process.poll() is not None and output == b"":
@@ -411,9 +412,15 @@ def run_command(
                     # logging will add line break
                     msg = msg.rstrip("\n")
                     logger.log(level=level, msg=msg)
+                if "pdsh@" in msg and "ssh exited with exit code 1" in msg:
+                    print("DeepSpeed Failed.")
+                    sys.exit(1)
             # Add a small delay so that
             # outputs from the subsequent code will have different timestamp for oci logging
             time.sleep(0.02)
+        logger.debug(
+            "subprocess %s returned exit code %s", process.pid, process.returncode
+        )
         if check and process.returncode != 0:
             # If there is an error, exit the main process with the same return code.
             sys.exit(process.returncode)