added port argument for ssh

deepspeed/launcher/multinode_runner.py

@@ -65,6 +65,8 @@ def parse_user_args(self):
 
     def get_cmd(self, environment, active_resources):
         environment['PDSH_RCMD_TYPE'] = 'ssh'
+        if self.args.ssh_port is not None:  # only specify ssh port if it is specified
+            environment["PDSH_SSH_ARGS_APPEND"] = f" -p {self.args.ssh_port}"
 
         active_workers = ",".join(active_resources.keys())
         logger.info("Running on the following workers: %s" % active_workers)
@@ -101,7 +103,7 @@ def get_cmd(self, environment, active_resources):
         cmd_to_search = [i + "\\" for i in deepspeed_launch[2:6]]
 
         kill_command = pdsh_cmd_args + ["pkill -f ", " ".join(cmd_to_search)[:-2]]
-        return pdsh_cmd_args + deepspeed_launch + [self.user_script] + self.user_arguments, kill_command
+        return pdsh_cmd_args + deepspeed_launch + [self.user_script] + self.user_arguments, kill_command, environment
 
 
 class OpenMPIRunner(MultiNodeRunner):

deepspeed/launcher/runner.py

@@ -193,6 +193,8 @@ def parse_args(args=None):
                         "numbers and range. i.e. 1,3-5,7 => [1,3,4,5,7].  When not "
                         "specified, all cores on system would be used rank binding")
 
+    parser.add_argument("--ssh_port", type=int, default=None, help="SSH port to use for remote connections")
+
     return parser.parse_args(args=args)
 
 
@@ -432,10 +434,11 @@ def main(args=None):
     if multi_node_exec and not args.no_ssh_check:
         first_host = list(active_resources.keys())[0]
         try:
-            subprocess.check_call(f'ssh -o PasswordAuthentication=no {first_host} hostname',
-                                  stderr=subprocess.DEVNULL,
-                                  stdout=subprocess.DEVNULL,
-                                  shell=True)
+            ssh_check_cmd = "ssh -o PasswordAuthentication=no "
+            if args.ssh_port is not None:
+                ssh_check_cmd += f"-p {args.ssh_port} "
+            ssh_check_cmd += f"{first_host} hostname"
+            subprocess.check_call(ssh_check_cmd, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL, shell=True)
         except subprocess.CalledProcessError:
             raise RuntimeError(
                 f"Using hostfile at {args.hostfile} but host={first_host} was not reachable via ssh. If you are running with a single node please remove {args.hostfile} or setup passwordless ssh."
@@ -560,7 +563,7 @@ def main(args=None):
                         runner.add_export(key, val)
 
         if args.launcher == PDSH_LAUNCHER:
-            cmd, kill_cmd = runner.get_cmd(env, active_resources)
+            cmd, kill_cmd, env = runner.get_cmd(env, active_resources)
         else:
             cmd = runner.get_cmd(env, active_resources)
 

tests/unit/launcher/test_multinode_runner.py

@@ -22,7 +22,7 @@ def runner_info():
 def test_pdsh_runner(runner_info):
     env, resource_pool, world_info, args = runner_info
     runner = mnrunner.PDSHRunner(args, world_info)
-    cmd, kill_cmd = runner.get_cmd(env, resource_pool)
+    cmd, kill_cmd, env = runner.get_cmd(env, resource_pool)
     assert cmd[0] == 'pdsh'
     assert env['PDSH_RCMD_TYPE'] == 'ssh'