Log all gpu rank stdout/err to MosaicML platform

composer/cli/launcher.py

@@ -21,6 +21,7 @@
 import torch
 
 import composer
+from composer.loggers.mosaicml_logger import MOSAICML_ACCESS_TOKEN_ENV_VAR, MOSAICML_LOG_DIR, MOSAICML_PLATFORM_ENV_VAR
 from composer.utils import get_free_tcp_port
 
 CLEANUP_TIMEOUT = datetime.timedelta(seconds=30)
@@ -259,8 +260,7 @@ def _launch_processes(
     module_mode: bool,
     command_mode: bool,
     training_script: str,
-    stdout_file_format: str,
-    stderr_file_format: str,
+    log_file_format: str,
     training_script_args: List[Any],
     processes: Dict[int, subprocess.Popen],
 ):
@@ -313,19 +313,17 @@ def _get_file(format: str):
                         local_world_size=nproc,
                         node_rank=node_rank,
                     )
-                    return open(filename, 'x+')
+                    return open(filename, 'a+')
 
-                stderr_file = _get_file(stderr_file_format)
-                stdout_file = _get_file(stdout_file_format)
+                log_file = _get_file(log_file_format)
 
                 process = subprocess.Popen(
                     cmd,
-                    stdout=stdout_file,
-                    stderr=stderr_file,
+                    stdout=log_file,
+                    stderr=subprocess.STDOUT,
                     text=True,
                 )
-                process.stderr = stderr_file
-                process.stdout = stdout_file
+                process.stdout = log_file
             processes[global_rank] = process
 
 
@@ -360,32 +358,16 @@ def _print_process_exit_status(global_rank: int, process: subprocess.Popen):
     if process.stdout is None:
         output = None
     else:
+        print(f'stdout found for rank {global_rank}')
         process.stdout.seek(0)
         output = process.stdout.read()
 
-    if process.stderr is None:
-        stderr = None
-    else:
-        process.stderr.seek(0)
-        stderr = process.stderr.read()
-    exc = subprocess.CalledProcessError(
-        process.returncode,
-        cmd=process.args,
-        output=output,
-        stderr=stderr,
-    )
     error_msg = [f'Global rank {global_rank} (PID {process.pid}) exited with code {process.returncode}']
     if output is not None:
         error_msg.extend([
-            f'----------Begin global rank {global_rank} STDOUT----------',
+            f'----------Begin global rank {global_rank} logs----------',
             output,
-            f'----------End global rank {global_rank} STDOUT----------',
-        ])
-    if stderr is not None:
-        error_msg.extend([
-            f'----------Begin global rank {global_rank} STDERR----------',
-            exc.stderr,
-            f'----------End global rank {global_rank} STDERR----------',
+            f'----------End global rank {global_rank} logs----------',
         ])
     print('\n'.join(error_msg))
 
@@ -466,12 +448,14 @@ def main():
     log.setLevel(logging.INFO if args.verbose else logging.WARN)
 
     processes = {}
-
     log_tmpdir = tempfile.TemporaryDirectory()
-    if args.stdout is None:
-        args.stdout = f'{log_tmpdir.name}/rank{{rank}}.stdout.txt'
-    if args.stderr is None:
-        args.stderr = f'{log_tmpdir.name}/rank{{rank}}.stderr.txt'
+    log_file_format = f'{log_tmpdir.name}/rank{{rank}}.txt'
+
+    # If running on the Mosaic platform, also log all gpu ranks' stderr and stdout to Mosaic platform
+    if os.environ.get(MOSAICML_PLATFORM_ENV_VAR, 'false').lower() == 'true' and os.environ.get(
+            MOSAICML_ACCESS_TOKEN_ENV_VAR) is not None and os.environ.get(MOSAICML_LOG_DIR) is not None:
+        log.info('Logging all gpu ranks to Mosaic Platform')
+        log_file_format = f'{os.environ.get(MOSAICML_LOG_DIR)}/gpu_{{rank}}.txt'
 
     try:
         _launch_processes(nproc=args.nproc,
@@ -482,8 +466,7 @@ def main():
                           master_port=args.master_port,
                           module_mode=args.module_mode,
                           command_mode=args.command_mode,
-                          stdout_file_format=args.stdout,
-                          stderr_file_format=args.stderr,
+                          log_file_format=log_file_format,
                           training_script=args.training_script,
                           training_script_args=args.training_script_args,
                           processes=processes)

composer/loggers/mosaicml_logger.py

@@ -36,6 +36,7 @@
 RUN_NAME_ENV_VAR = 'RUN_NAME'
 MOSAICML_PLATFORM_ENV_VAR = 'MOSAICML_PLATFORM'
 MOSAICML_ACCESS_TOKEN_ENV_VAR = 'MOSAICML_ACCESS_TOKEN_FILE'
+MOSAICML_LOG_DIR = 'MOSAICML_LOG_DIR'
 
 
 class MosaicMLLogger(LoggerDestination):