Add error handling for remote function.

Author: Ao Guo

requirements/extras/test_requirements.txt

@@ -22,3 +22,4 @@ Jinja2==3.0.3
 pandas>=1.3.5,<1.5
 scikit-learn==1.0.2
 cloudpickle==2.2.0
+tblib==1.7.0

src/sagemaker/remote_function/client.py

@@ -22,8 +22,10 @@
 import inspect
 
 from botocore.exceptions import ClientError
+from sagemaker.exceptions import UnexpectedStatusException
 
 import sagemaker.remote_function.core.serialization as serialization
+from sagemaker.remote_function.errors import RemoteFunctionError, ServiceError
 
 from sagemaker.session import Session
 from sagemaker.s3 import s3_path_join
@@ -136,18 +138,47 @@ def wrapper(*args, **kwargs):
                 volume_size=volume_size,
             )
             job = _Job.start(job_settings, func, args, kwargs)
-            job.wait()
+
+            try:
+                job.wait()
+            except UnexpectedStatusException as usex:
+                if usex.actual_status == "Failed":
+                    try:
+                        exception = serialization.deserialize_obj_from_s3(
+                            sagemaker_session=job_settings.sagemaker_session,
+                            s3_uri=s3_path_join(
+                                job_settings.s3_root_uri, job.job_name, "exception.pkl"
+                            ),
+                        )
+                    except ServiceError as serr:
+                        chained_e = serr.__cause__
+                        if (
+                            isinstance(chained_e, ClientError)
+                            and chained_e.response["Error"]["Code"]  # pylint: disable=no-member
+                            == "404"
+                            and chained_e.response["Error"]["Message"]  # pylint: disable=no-member
+                            == "Not Found"
+                        ):
+                            raise RemoteFunctionError(
+                                "Failed to execute remote function. "
+                                + "Check corresponding job for details."
+                            )
+                        raise serr
+
+                    raise exception
+
+                raise TimeoutError(
+                    "Job for remote function timed out before reaching a termination status."
+                )
 
             if job.describe()["TrainingJobStatus"] == "Completed":
                 return serialization.deserialize_obj_from_s3(
                     sagemaker_session=job_settings.sagemaker_session,
                     s3_uri=s3_path_join(job_settings.s3_root_uri, job.job_name, "results.pkl"),
                 )
 
-            if job.describe()["TrainingJobStatus"] == "Failed":
-                # TODO: the exception should be constructed based on the job failure reason
-                #  and function output
-                raise RuntimeError("Job failed")
+            if job.describe()["TrainingJobStatus"] == "Stopped":
+                raise RemoteFunctionError("Job for remote function has been aborted.")
 
             return None
 
@@ -522,7 +553,10 @@ def result(self, timeout: int = None) -> Any:
         Returns:
             The Python object returned by the function
         """
-        self.wait(timeout)
+        try:
+            self.wait(timeout)
+        except UnexpectedStatusException:
+            pass
 
         with self._condition:
             if self._state == _PENDING:
@@ -539,13 +573,36 @@ def result(self, timeout: int = None) -> Any:
                     self._state = _FINISHED
                     return self._return
                 if self._job.describe()["TrainingJobStatus"] == "Failed":
-                    # TODO: the exception should be constructed based on the job failure reason
-                    #  and function output
-                    self._exception = RuntimeError()
+                    try:
+                        self._exception = serialization.deserialize_obj_from_s3(
+                            sagemaker_session=self._job_settings.sagemaker_session,
+                            s3_uri=s3_path_join(
+                                self._job_settings.s3_root_uri, self._job.job_name, "exception.pkl"
+                            ),
+                        )
+                    except ServiceError as serr:
+                        chained_e = serr.__cause__
+                        if (
+                            isinstance(chained_e, ClientError)
+                            and chained_e.response["Error"]["Code"]  # pylint: disable=no-member
+                            == "404"
+                            and chained_e.response["Error"]["Message"]  # pylint: disable=no-member
+                            == "Not Found"
+                        ):
+                            self._exception = RemoteFunctionError(
+                                "Failed to execute remote function. "
+                                + "Check corresponding job for details."
+                            )
+                        else:
+                            self._exception = serr
                     self._state = _FINISHED
-                    raise self._exception
-
-                raise RuntimeError()
+                elif self._job.describe()["TrainingJobStatus"] == "Stopped":
+                    self._state = _CANCELLED
+                    raise RemoteFunctionError("Job for remote function has been aborted.")
+                else:
+                    raise TimeoutError(
+                        "Job for remote function timed out before reaching a termination status."
+                    )
 
             if self._state == _FINISHED:
                 if self._exception:

src/sagemaker/remote_function/core/serialization.py

@@ -16,6 +16,7 @@
 import cloudpickle
 
 from typing import Any, Callable
+from sagemaker.remote_function.errors import ServiceError, SerializationError, DeserializationError
 from sagemaker.s3 import S3Downloader, S3Uploader
 
 
@@ -32,10 +33,19 @@ def serialize_func_to_s3(func: Callable, sagemaker_session, s3_uri, s3_kms_key=N
         s3_uri (str): S3 root uri to which resulting serialized artifacts will be uploaded.
         s3_kms_key (str): KMS key used to encrypt artifacts uploaded to S3.
         func: function to be serialized and persisted
+    Raises:
+        SerializationError: when fail to serialize function to bytes.
     """
-    S3Uploader.upload_bytes(
-        cloudpickle.dumps(func), s3_uri, kms_key=s3_kms_key, sagemaker_session=sagemaker_session
-    )
+    try:
+        bytes_to_upload = cloudpickle.dumps(func)
+    except Exception as e:
+        raise SerializationError(
+            "Error when serializing function [{}]: {}".format(
+                getattr(func, "__name__", repr(func)), e
+            )
+        ) from e
+
+    _upload_bytes_to_s3(bytes_to_upload, s3_uri, s3_kms_key, sagemaker_session)
 
 
 def deserialize_func_from_s3(sagemaker_session, s3_uri) -> Callable:
@@ -50,8 +60,17 @@ def deserialize_func_from_s3(sagemaker_session, s3_uri) -> Callable:
         s3_uri (str): S3 root uri to which resulting serialized artifacts will be uploaded.
     Returns :
         The deserialized function.
+    Raises:
+        DeserializationError: when fail to serialize function to bytes.
     """
-    return cloudpickle.loads(S3Downloader.read_bytes(s3_uri, sagemaker_session=sagemaker_session))
+    bytes_to_deserialize = _read_bytes_from_s3(s3_uri, sagemaker_session)
+
+    try:
+        return cloudpickle.loads(bytes_to_deserialize)
+    except Exception as e:
+        raise DeserializationError(
+            "Error when deserializing bytes downloaded from {}: {}".format(s3_uri, e)
+        ) from e
 
 
 def serialize_obj_to_s3(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: str = None):
@@ -63,10 +82,17 @@ def serialize_obj_to_s3(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: st
         s3_uri (str): S3 root uri to which resulting serialized artifacts will be uploaded.
         s3_kms_key (str): KMS key used to encrypt artifacts uploaded to S3.
         obj: object to be serialized and persisted
+    Raises:
+        SerializationError: when fail to serialize object to bytes.
     """
-    S3Uploader.upload_bytes(
-        cloudpickle.dumps(obj), s3_uri, kms_key=s3_kms_key, sagemaker_session=sagemaker_session
-    )
+    try:
+        bytes_to_upload = cloudpickle.dumps(obj)
+    except Exception as e:
+        raise SerializationError(
+            "Error when serializing object of type [{}]: {}".format(type(obj).__name__, e)
+        ) from e
+
+    _upload_bytes_to_s3(bytes_to_upload, s3_uri, s3_kms_key, sagemaker_session)
 
 
 def deserialize_obj_from_s3(sagemaker_session, s3_uri) -> Any:
@@ -78,5 +104,32 @@ def deserialize_obj_from_s3(sagemaker_session, s3_uri) -> Any:
         s3_uri (str): S3 root uri to which resulting serialized artifacts will be uploaded.
     Returns :
         List of deserialized python objects.
+    Raises:
+        DeserializationError: when fail to serialize object to bytes.
     """
-    return cloudpickle.loads(S3Downloader.read_bytes(s3_uri, sagemaker_session=sagemaker_session))
+    bytes_to_deserialize = _read_bytes_from_s3(s3_uri, sagemaker_session)
+
+    try:
+        return cloudpickle.loads(bytes_to_deserialize)
+    except Exception as e:
+        raise DeserializationError(
+            "Error when deserializing bytes downloaded from {}: {}".format(s3_uri, e)
+        ) from e
+
+
+def _upload_bytes_to_s3(bytes, s3_uri, s3_kms_key, sagemaker_session):
+    """Wrapping s3 uploading with exception translation for remote function."""
+    try:
+        S3Uploader.upload_bytes(
+            bytes, s3_uri, kms_key=s3_kms_key, sagemaker_session=sagemaker_session
+        )
+    except Exception as e:
+        raise ServiceError("Failed to upload serialized bytes to {}: {}".format(s3_uri, e)) from e
+
+
+def _read_bytes_from_s3(s3_uri, sagemaker_session):
+    """Wrapping s3 downloading with exception translation for remote function."""
+    try:
+        return S3Downloader.read_bytes(s3_uri, sagemaker_session=sagemaker_session)
+    except Exception as e:
+        raise ServiceError("Failed to read serialized bytes from {}: {}".format(s3_uri, e)) from e

src/sagemaker/remote_function/errors.py

@@ -24,6 +24,7 @@
 FAILURE_REASON_PATH = "/opt/ml/output/failure"
 
 
+@pickling_support.install
 class RemoteFunctionError(Exception):
     """The base exception class for remote function excepitons"""
 
@@ -32,18 +33,22 @@ def __init__(self, message):
         super().__init__(self.message)
 
 
+@pickling_support.install
 class ServiceError(RemoteFunctionError):
     """Raised when errors encountered during interaction with SageMaker, S3 service APIs"""
 
 
+@pickling_support.install
 class RuntimeEnvironmentError(RemoteFunctionError):
     """Raised when errors encountered during remote function runtime environment setup"""
 
 
+@pickling_support.install
 class SerializationError(RemoteFunctionError):
     """Raised when errors encountered during serialization of remote function objects"""
 
 
+@pickling_support.install
 class DeserializationError(RemoteFunctionError):
     """Raised when errors encountered during deserialization of remote function objects"""
 
@@ -80,7 +85,7 @@ def handle_error(error, sagemaker_session, s3_base_uri, s3_kms_key) -> int:
         s3_base_uri (str): S3 root uri to which resulting serialized exception will be uploaded.
         s3_kms_key (str): KMS key used to encrypt artifacts uploaded to S3.
     Returns :
-        List of deserialized python objects.
+        exit_code (int): Exit code to terminate current job.
     """
     pickling_support.install()
     failure_reason = repr(error)

src/sagemaker/remote_function/job_driver.py

@@ -25,6 +25,7 @@
 
 
 SUCCESS_EXIT_CODE = 0
+DEFAULT_FAILURE_CODE = 1
 
 
 def _parse_agrs():
@@ -85,7 +86,7 @@ def main():
     logging_config.basic_config()
     logger = logging_config.get_logger()
 
-    exit_code = SUCCESS_EXIT_CODE
+    exit_code = DEFAULT_FAILURE_CODE
     try:
         args = _parse_agrs()
         region = args.region
@@ -109,7 +110,7 @@ def main():
         )
 
         _execute_remote_function(sagemaker_session, s3_base_uri, s3_kms_key)
-
+        exit_code = SUCCESS_EXIT_CODE
     except Exception as e:  # pylint: disable=broad-except
         logger.exception("Error encountered when invoking the remote function.")
         exit_code = handle_error(e, sagemaker_session, s3_base_uri, s3_kms_key)