add test and make error class optional

Author: roywei

src/sagemaker_training/process.py

@@ -36,7 +36,7 @@
 _DEFAULT_BUF_SIZE = 1024 * 64
 
 
-async def watch(stream, error_classes, proc_per_host):
+async def watch(stream, proc_per_host, error_classes=None):
     """Process the stdout and stderr streams on the fly.
     Decode the output lines
     Remove new line characters (if any)
@@ -45,12 +45,16 @@ async def watch(stream, error_classes, proc_per_host):
 
     Args:
         stream: asyncio subprocess PIPE
-        error_classes (list): List of exception classes to watch and raise
         proc_per_host (int): Number of processes per each host
+        error_classes (list): List of exception classes to watch and raise
 
     Returns:
         output: Filtered stderr
     """
+    if not error_classes:
+        error_classes = []
+    if not isinstance(error_classes, list):
+        error_classes = [error_classes]
     output = []
     buf_size = _DEFAULT_BUF_SIZE
     start = False
@@ -105,17 +109,17 @@ async def watch(stream, error_classes, proc_per_host):
     return " ".join(output)
 
 
-async def run_async(cmd, error_classes, processes_per_host, env, cwd, stderr, **kwargs):
+async def run_async(cmd, processes_per_host, env, cwd, stderr, error_classes=None, **kwargs):
     """Method responsible for launching asyncio subprocess shell
-    Use asyncio gather to collect processed stdout and stderr
+    Usyncse asyncio gather to collect processed stdout and stderr
 
     Args:
         cmd (list): The command to be run
-        error_classes (list): List of exception classes to watch and raise
         processes_per_host (int): Number of processes per host
         env: os.environ
         cwd (str): The location from which to run the command (default: None).
             If None, this defaults to the ``code_dir`` of the environment.
+        error_classes (list): List of exception classes to watch and raise
         **kwargs: Extra arguments that are passed to the asyncio create subprocess constructor.
 
     Returns:
@@ -126,14 +130,18 @@ async def run_async(cmd, error_classes, processes_per_host, env, cwd, stderr, **
     Raises:
         ExecuteUserScriptError: If there is an exception raised when creating the process.
     """
+    if not error_classes:
+        error_classes = []
+    if not isinstance(error_classes, list):
+        error_classes = [error_classes]
     cmd = " ".join(cmd)
     proc = await asyncio.create_subprocess_shell(
         cmd, env=env, cwd=cwd, stdout=PIPE, stderr=stderr, **kwargs
     )
 
     output = await asyncio.gather(
-        watch(proc.stdout, error_classes, processes_per_host),
-        watch(proc.stderr, error_classes, processes_per_host),
+        watch(proc.stdout, processes_per_host, error_classes=error_classes),
+        watch(proc.stderr, processes_per_host, error_classes=error_classes),
     )
     return_code = proc.returncode
     return return_code, output, proc
@@ -166,16 +174,20 @@ def create(
     Raises:
         ExecuteUserScriptError: If there is an exception raised when creating the process.
     """
+    if not error_classes:
+        error_classes = []
+    if not isinstance(error_classes, list):
+        error_classes = [error_classes]
     try:
         stderr = PIPE if capture_error else None
         rc, output, proc = asyncio.run(
             run_async(
                 cmd,
-                error_classes,
                 processes_per_host,
                 env=env or os.environ,
                 cwd=cwd or environment.code_dir,
                 stderr=stderr,
+                error_classes=error_classes,
                 **kwargs,
             )
         )
@@ -206,7 +218,10 @@ def check_error(cmd, error_classes, processes_per_host, cwd=None, capture_error=
     Raises:
         ExecuteUserScriptError: If there is an exception raised when creating the process.
     """
-
+    if not error_classes:
+        error_classes = []
+    if not isinstance(error_classes, list):
+        error_classes = [error_classes]
     if capture_error:
         return_code, output, process = create(
             cmd,
@@ -239,10 +254,11 @@ def check_error(cmd, error_classes, processes_per_host, cwd=None, capture_error=
         # default error class will be user script error
         error_class = errors.ExecuteUserScriptError
         # use first found target error class if available
-        for error_name in error_classes:
-            if error_name in stderr:
-                error_class = type(error_name, (errors.ExecuteUserScriptError,), {})
-                break
+        if stderr:
+            for error_name in error_classes:
+                if str(error_name) in stderr:
+                    error_class = type(error_name, (errors.ExecuteUserScriptError,), {})
+                    break
 
         raise error_class(
             cmd=" ".join(cmd) if isinstance(cmd, list) else cmd,

test/unit/test_environment.py

@@ -212,6 +212,7 @@ def test_env_mapping_properties(training_env):
         "output_intermediate_dir",
         "is_master",
         "master_hostname",
+        "is_modelparallel_enabled",
     }
 
 

test/unit/test_process.py

@@ -105,7 +105,7 @@ async def test_watch(event_loop, capsys):
     expected_stream += (
         "[1,mpirank:0,algo-1]<stderr>:FileNotFoundError: [Errno 2] No such file or directory\n"
     )
-    expected_errmsg = ":FileNotFoundError: [Errno 2] No such file or directory\n"
+    expected_errmsg = "FileNotFoundError: [Errno 2] No such file or directory\n"
 
     stream = asyncio.StreamReader()
     stream.feed_data(b"[1,10]<stdout>:This is stdout\n")
@@ -119,6 +119,48 @@ async def test_watch(event_loop, capsys):
     assert output == expected_errmsg
 
 
+@pytest.mark.asyncio
+async def test_watch_custom_error(event_loop, capsys):
+    num_processes_per_host = 8
+    expected_stream = "[1,mpirank:10,algo-2]<stdout>:This is stdout\n"
+    expected_stream += "[1,mpirank:10,algo-2]<stderr>:This is stderr\n"
+    expected_stream += "[1,mpirank:0,algo-1]<stderr>:SMDDPNCCLError: unhandled cuda error\n"
+    expected_errmsg = "SMDDPNCCLError: unhandled cuda error\n"
+
+    stream = asyncio.StreamReader()
+    stream.feed_data(b"[1,10]<stdout>:This is stdout\n")
+    stream.feed_data(b"[1,10]<stderr>:This is stderr\n")
+    stream.feed_data(b"[1,0]<stderr>:SMDDPNCCLError: unhandled cuda error")
+    stream.feed_eof()
+
+    error_classes = ["SMDDPNCCLError"]
+    output = await process.watch(stream, num_processes_per_host, error_classes=error_classes)
+    captured_stream = capsys.readouterr()
+    assert captured_stream.out == expected_stream
+    assert output == expected_errmsg
+
+    # test errors piped in stdout
+    stream = asyncio.StreamReader()
+    stream.feed_data(b"[1,10]<stdout>:This is stdout\n")
+    stream.feed_data(b"[1,10]<stderr>:This is stderr\n")
+    stream.feed_data(b"[1,0]<stdout>:SMDDPNCCLError: unhandled cuda error")
+    stream.feed_eof()
+
+    error_classes = ["SMDDPNCCLError"]
+    output = await process.watch(stream, num_processes_per_host, error_classes=error_classes)
+    assert output == expected_errmsg
+
+    # test single item
+    stream = asyncio.StreamReader()
+    stream.feed_data(b"[1,10]<stdout>:This is stdout\n")
+    stream.feed_data(b"[1,10]<stderr>:This is stderr\n")
+    stream.feed_data(b"[1,0]<stdout>:SMDDPNCCLError: unhandled cuda error")
+    stream.feed_eof()
+    error_classes = "SMDDPNCCLError"
+    output = await process.watch(stream, num_processes_per_host, error_classes=error_classes)
+    assert output == expected_errmsg
+
+
 @patch("asyncio.run", AsyncMock(side_effect=ValueError("FAIL")))
 def test_create_error():
     with pytest.raises(errors.ExecuteUserScriptError):