Add a runner for MosaicML cloud (#44)

* add mosaic launcher string constant

* first attempt at mosaic multinode runner for gptneox

* typo

* actually add mosaic runner

* string cast

* debug

* fix

* actually fix?

* strip extra space

* debugging

* debugging

* correctly set env vars

* drop cd

* add env vars via env instead of export commands

* fix

* try using slurms arg parsing

* debug print

* debug print

* print debug

* cleanup

* try getting world info from the hostfile

* add missing init arg

* more cleanup

* remove more prints

Author: dakinggg

deepspeed/launcher/constants.py

@@ -5,6 +5,7 @@
 
 OPENMPI_LAUNCHER = 'openmpi'
 SLURM_LAUNCHER = 'slurm'
+MOSAICML_LAUNCHER = 'mosaicml'
 
 MVAPICH_LAUNCHER = 'mvapich'
 MVAPICH_TMP_HOSTFILE = '/tmp/deepspeed_mvapich_hostfile'

deepspeed/launcher/multinode_runner.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import os
 import sys
@@ -252,3 +253,40 @@ def get_cmd(self, environment, active_resources):
 
         return mpirun_cmd + export_cmd + python_exec + [self.user_script
                                                         ] + self.user_arguments
+class MosaicMLRunner(MultiNodeRunner):
+    def __init__(self, args, world_info_base64):
+        super().__init__(args, world_info_base64)
+
+    def backend_exists(self):
+        return True
+
+    def parse_user_args(self):
+        user_args = []
+        for arg in self.args.user_args:
+            if arg.startswith('{') and arg.endswith('}'):
+                try:
+                    arg_dict = json.loads(arg)
+                    if 'config_files' in arg_dict:
+                        config_files = {}
+                        for k, v in arg_dict.get('config_files', {}).items():
+                            config_files[k] = json.loads(v)
+                        arg_dict['config_files'] = config_files
+                except json.JSONDecodeError as jde:
+                    raise ValueError('Please use plain json for your configs. Check for comments and lowercase trues') from jde
+                arg = json.dumps(arg_dict, separators=(',', ':'))
+            user_args.append(arg)
+        return user_args
+
+    def get_cmd(self, environment, active_resources):
+        deepspeed_launch = [
+            sys.executable,
+            "-u",
+            "-m",
+            "deepspeed.launcher.launch",
+            '--world_info={}'.format(self.world_info_base64),
+            "--node_rank={}".format(os.environ['NODE_RANK']),
+            "--master_addr={}".format(os.environ['MASTER_ADDR']),
+            "--master_port={}".format(os.environ['MASTER_PORT']),
+        ]
+
+        return deepspeed_launch + [self.user_script] + self.user_arguments

deepspeed/launcher/runner.py

@@ -18,8 +18,8 @@
 
 import torch.cuda
 
-from .multinode_runner import PDSHRunner, OpenMPIRunner, MVAPICHRunner, SlurmRunner
-from .constants import PDSH_LAUNCHER, OPENMPI_LAUNCHER, MVAPICH_LAUNCHER, SLURM_LAUNCHER
+from .multinode_runner import PDSHRunner, OpenMPIRunner, MVAPICHRunner, SlurmRunner, MosaicMLRunner
+from .constants import PDSH_LAUNCHER, OPENMPI_LAUNCHER, MVAPICH_LAUNCHER, SLURM_LAUNCHER, MOSAICML_LAUNCHER
 from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT
 from ..utils import logger
 
@@ -343,6 +343,8 @@ def main(args=None):
             runner = MVAPICHRunner(args, world_info_base64, resource_pool)
         elif args.launcher == SLURM_LAUNCHER:
             runner = SlurmRunner(args, world_info_base64, resource_pool)
+        elif args.launcher == MOSAICML_LAUNCHER:
+            runner = MosaicMLRunner(args, world_info_base64)
         else:
             raise NotImplementedError(f"Unknown launcher {args.launcher}")
 
@@ -367,11 +369,10 @@ def main(args=None):
                     for var in fd.readlines():
                         key, val = var.split('=')
                         runner.add_export(key, val)
-
         cmd = runner.get_cmd(env, active_resources)
 
     logger.info("cmd = {}".format(' '.join(cmd)))
-    result = subprocess.Popen(cmd, env=env)
+    result = subprocess.Popen(cmd, env=dict(env, **runner.exports))
     result.wait()
 
     # In case of failure must propagate the error-condition back to the caller (usually shell). The