feature: allow framework libraries to supply exceptions to track and report as failure reason.

Added support for SMDDP and SMMP custom exceptions.
Include custom exception as error class and de-duplicated stack trace as error message.
Added tests for wacthing single, list of exceptions and also support existing internal exceptions.

Author: roywei

src/sagemaker_training/environment.py

@@ -563,7 +563,7 @@ def __init__(self, resource_config=None, input_data_config=None, hyperparameters
         self._is_master = current_host == self._master_hostname
 
         mp_parameters = os.environ.get(params.SM_HP_MP_PARAMETERS)
-        self._is_modelparallel_enabled = mp_parameters and mp_parameters != '{}'
+        self._is_modelparallel_enabled = mp_parameters and mp_parameters != "{}"
 
     @property
     def model_dir(self):  # type: () -> str
@@ -921,6 +921,7 @@ def is_modelparallel_enabled(self):  # type: () -> bool
         """
         return self._is_modelparallel_enabled
 
+
 def write_env_vars(env_vars=None):  # type: (dict) -> None
     """Write the dictionary env_vars in the system, as environment variables.
 

src/sagemaker_training/mpi.py

@@ -13,7 +13,7 @@
 """This module contains functionality related to distributed training using
 MPI (Message Passing Interface)."""
 import argparse
-import inspect
+from inspect import getfile, isclass
 import logging
 import os
 import subprocess
@@ -24,23 +24,24 @@
 
 import gethostname
 from sagemaker_training import environment, errors, logging_config, process, timeout
-from inspect import isclass
 
 logger = logging_config.get_logger()
 logging.getLogger("paramiko").setLevel(logging.INFO)
 
 try:
     from smdistributed.modelparallel.backend import exceptions
+
     # list of exceptions SMMP wants training toolkit to catch and log
     exception_classes = [x for x in dir(exceptions) if isclass(getattr(exceptions, x))]
-except ImportError as e:
+except ImportError:
     logger.info("No exception classes found in smdistributed.modelparallel")
     exception_classes = []
 try:
     from smdistributed.modelparallel.torch import exceptions as torch_exceptions
+
     # list of torch exceptions SMMP wants training toolkit to catch and log
     exception_classes += [x for x in dir(torch_exceptions) if isclass(getattr(torch_exceptions, x))]
-except ImportError as e:
+except ImportError:
     logger.info("No torch exception classes found in smdistributed.modelparallel")
 
 
@@ -250,12 +251,16 @@ def _create_command(self):
             "-x",
             "PATH",
             "-x",
-            "LD_PRELOAD=%s" % inspect.getfile(gethostname),
+            "LD_PRELOAD=%s" % getfile(gethostname),
         ]
 
         command.extend(additional_options)
 
-        for credential in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"]:
+        for credential in [
+            "AWS_ACCESS_KEY_ID",
+            "AWS_SECRET_ACCESS_KEY",
+            "AWS_SESSION_TOKEN",
+        ]:
             if credential in os.environ:
                 command.extend(["-x", credential])
 
@@ -293,15 +298,19 @@ def run(self, wait=True, capture_error=False):
         if wait:
             process_spawned = process.check_error(
                 cmd,
-                exception_classes if training_env.is_modelparallel_enabled else errors.ExecuteUserScriptError,
+                exception_classes
+                if training_env.is_modelparallel_enabled
+                else errors.ExecuteUserScriptError,
                 self._processes_per_host,
                 capture_error=capture_error,
                 cwd=environment.code_dir,
             )
         else:
             _, _, process_spawned = process.create(
                 cmd,
-                exception_classes if training_env.is_modelparallel_enabled else errors.ExecuteUserScriptError,
+                exception_classes
+                if training_env.is_modelparallel_enabled
+                else errors.ExecuteUserScriptError,
                 self._processes_per_host,
                 capture_error=capture_error,
                 cwd=environment.code_dir,
@@ -311,7 +320,6 @@ def run(self, wait=True, capture_error=False):
         return process_spawned
 
 
-
 _SSH_DAEMON_NOT_FOUND_ERROR_MESSAGE = """
 SSH daemon not found, please install SSH to allow MPI to communicate different nodes in cluster.
 

src/sagemaker_training/params.py

@@ -59,4 +59,4 @@
 SMDATAPARALLEL_CUSTOM_MPI_OPTIONS = (
     "sagemaker_distributed_dataparallel_custom_mpi_options"
 )  # type: str
-SM_HP_MP_PARAMETERS = "SM_HP_MP_PARAMETERS" 
+SM_HP_MP_PARAMETERS = "SM_HP_MP_PARAMETERS"

src/sagemaker_training/process.py

@@ -17,6 +17,7 @@
 
 import asyncio
 from asyncio.subprocess import PIPE
+from inspect import isclass
 import os
 import re
 import subprocess
@@ -36,7 +37,24 @@
 _DEFAULT_BUF_SIZE = 1024 * 64
 
 
-async def watch(stream, error_classes, proc_per_host):
+def process_error_classes(error_classes):
+    """Process error classes and return a list of string.
+    Input could be class, string, or None
+
+    Args:
+        error_classes (list): List of error classes
+
+    Returns:
+        error_classes: processed classes
+    """
+    if not error_classes:
+        return []
+    if not isinstance(error_classes, list):
+        error_classes = [error_classes]
+    return [error.__name__ if isclass(error) else error for error in error_classes]
+
+
+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 +63,13 @@ 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
     """
+    error_classes = process_error_classes(error_classes)
     output = []
     buf_size = _DEFAULT_BUF_SIZE
     start = False
@@ -90,25 +109,32 @@ async def watch(stream, error_classes, proc_per_host):
                 if err_line not in output:
                     output.append(err_line.strip(" :\n") + "\n")
             else:
-                if any(str(err) in err_line for err in (_PYTHON_ERRORS_ + error_classes if type(error_classes) == list else [error_classes])):
+                if any(
+                    str(err) in err_line
+                    for err in (
+                        _PYTHON_ERRORS_ + error_classes
+                        if isinstance(error_classes, list)
+                        else [error_classes]
+                    )
+                ):
                     # start logging error message if target exceptions found
                     start = True
                     output.append(err_line.strip(" :\n") + "\n")
 
     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:
@@ -125,8 +151,8 @@ async def run_async(cmd, error_classes, processes_per_host, env, cwd, stderr, **
     )
 
     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
@@ -164,11 +190,11 @@ def create(
         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,
             )
         )
@@ -181,9 +207,7 @@ def create(
         )
 
 
-def check_error(
-    cmd, error_classes, processes_per_host, cwd=None, capture_error=True, **kwargs
-):
+def check_error(cmd, error_classes, processes_per_host, cwd=None, capture_error=True, **kwargs):
     """Run a commmand, raising an exception if there is an error.
 
     Args:
@@ -201,7 +225,7 @@ def check_error(
     Raises:
         ExecuteUserScriptError: If there is an exception raised when creating the process.
     """
-
+    error_classes = process_error_classes(error_classes)
     if capture_error:
         return_code, output, process = create(
             cmd,
@@ -231,13 +255,24 @@ def check_error(
         extra_info = None
         if return_code == 137:
             extra_info = "OutOfMemory: Process killed by SIGKILL (signal 9)"
-        # 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._CalledProcessError,), {})
-                break
+
+        # throw internal error classes first
+        internal_errors = [err for err in dir(errors) if isclass(getattr(errors, err))]
+        error_class = next(
+            (name for name in error_classes if name in internal_errors), "ExecuteUserScriptError"
+        )
+        error_class = getattr(errors, error_class)
+
+        # only replace ExecuteUserScriptError with custom library errors
+        if stderr and error_class == errors.ExecuteUserScriptError:
+            # find the first target error in stderr
+            error_name = next((str(name) for name in error_classes if str(name) in stderr), False)
+            if error_name:
+                error_class = type(
+                    error_name,
+                    (errors._CalledProcessError,),  # pylint: disable=protected-access
+                    {},
+                )
 
         raise error_class(
             cmd=" ".join(cmd) if isinstance(cmd, list) else cmd,
@@ -257,9 +292,7 @@ def python_executable():
         (str): The real path of the current Python executable.
     """
     if not sys.executable:
-        raise RuntimeError(
-            "Failed to retrieve the real path for the Python executable binary"
-        )
+        raise RuntimeError("Failed to retrieve the real path for the Python executable binary")
     return sys.executable
 
 
@@ -281,9 +314,7 @@ def __init__(self, user_entry_point, args, env_vars, processes_per_host):
         self._processes_per_host = processes_per_host
 
     def _create_command(self):
-        entrypoint_type = _entry_point_type.get(
-            environment.code_dir, self._user_entry_point
-        )
+        entrypoint_type = _entry_point_type.get(environment.code_dir, self._user_entry_point)
 
         if entrypoint_type is _entry_point_type.PYTHON_PACKAGE:
             entry_module = self._user_entry_point.replace(".py", "")

src/sagemaker_training/smdataparallel.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 """Contains functionality related to SM Distributed Data Parallel Training."""
 import argparse
-import inspect
+from inspect import getfile, isclass
 import json
 import logging
 import os
@@ -23,7 +23,7 @@
 
 import gethostname
 from sagemaker_training import environment, errors, logging_config, process, timeout
-from inspect import isclass
+
 
 logger = logging_config.get_logger()
 logging.getLogger("paramiko").setLevel(logging.INFO)
@@ -33,7 +33,7 @@
 
     # list of exceptions SMDDP wants training toolkit to catch and log
     exception_classes = [x for x in dir(exceptions) if isclass(getattr(exceptions, x))]
-except ImportError as e:
+except ImportError:
     logger.info("No exception classes found in smdistributed.dataparallel")
     exception_classes = [errors.ExecuteUserScriptError]
 
@@ -173,7 +173,7 @@ def _get_mpirun_command(
             "-x",
             "RDMAV_FORK_SAFE=1",
             "-x",
-            "LD_PRELOAD=%s" % inspect.getfile(gethostname),
+            "LD_PRELOAD=%s" % getfile(gethostname),
         ]
 
         mpirun_command.extend(additional_options)
@@ -231,9 +231,7 @@ def _create_command(self):
             # homogeneous mode uses 16 processes per host; 8 server; 8 worker
             smdataparallel_server_addr = self._master_hostname
             smdataparallel_server_port = 7592
-            host_list = [
-                "{}:{}".format(host, num_processes_per_host) for host in self._hosts
-            ]
+            host_list = ["{}:{}".format(host, num_processes_per_host) for host in self._hosts]
             smdataparallel_flag = "SMDATAPARALLEL_USE_HOMOGENEOUS=1"
             command = self._get_mpirun_command(
                 num_hosts,

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",
     }