Version the serialization scheme (aws#891)

Co-authored-by: Zhankui Lu <[email protected]>

Author: 2 people

src/sagemaker/remote_function/client.py

@@ -28,6 +28,7 @@
 
 import sagemaker.remote_function.core.serialization as serialization
 from sagemaker.remote_function.errors import RemoteFunctionError, ServiceError, DeserializationError
+from sagemaker.remote_function.core.stored_function import RESULTS_FOLDER, EXCEPTION_FOLDER
 from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
     RuntimeEnvironmentError,
 )
@@ -162,10 +163,10 @@ def wrapper(*args, **kwargs):
             except UnexpectedStatusException as usex:
                 if usex.actual_status == "Failed":
                     try:
-                        exception = serialization.deserialize_obj_from_s3(
+                        exception = serialization.deserialize_exception_from_s3(
                             sagemaker_session=job_settings.sagemaker_session,
                             s3_uri=s3_path_join(
-                                job_settings.s3_root_uri, job.job_name, "exception.pkl"
+                                job_settings.s3_root_uri, job.job_name, EXCEPTION_FOLDER
                             ),
                         )
                     except ServiceError as serr:
@@ -202,7 +203,7 @@ def wrapper(*args, **kwargs):
             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"),
+                    s3_uri=s3_path_join(job_settings.s3_root_uri, job.job_name, RESULTS_FOLDER),
                 )
 
             if job.describe()["TrainingJobStatus"] == "Stopped":
@@ -596,7 +597,7 @@ def from_describe_response(describe_training_job_response, sagemaker_session):
             try:
                 job_return = serialization.deserialize_obj_from_s3(
                     sagemaker_session=sagemaker_session,
-                    s3_uri=s3_path_join(job.s3_uri, "results.pkl"),
+                    s3_uri=s3_path_join(job.s3_uri, RESULTS_FOLDER),
                 )
             except DeserializationError as e:
                 client_exception = e
@@ -605,9 +606,9 @@ def from_describe_response(describe_training_job_response, sagemaker_session):
         elif describe_training_job_response["TrainingJobStatus"] == "Failed":
             state = _FINISHED
             try:
-                job_exception = serialization.deserialize_obj_from_s3(
+                job_exception = serialization.deserialize_exception_from_s3(
                     sagemaker_session=sagemaker_session,
-                    s3_uri=s3_path_join(job.s3_uri, "exception.pkl"),
+                    s3_uri=s3_path_join(job.s3_uri, EXCEPTION_FOLDER),
                 )
             except ServiceError as serr:
                 chained_e = serr.__cause__
@@ -694,15 +695,15 @@ def result(self, timeout: float = None) -> Any:
                 if self._job.describe()["TrainingJobStatus"] == "Completed":
                     self._return = serialization.deserialize_obj_from_s3(
                         sagemaker_session=self._job.sagemaker_session,
-                        s3_uri=s3_path_join(self._job.s3_uri, "results.pkl"),
+                        s3_uri=s3_path_join(self._job.s3_uri, RESULTS_FOLDER),
                     )
                     self._state = _FINISHED
                     return self._return
                 if self._job.describe()["TrainingJobStatus"] == "Failed":
                     try:
-                        self._exception = serialization.deserialize_obj_from_s3(
+                        self._exception = serialization.deserialize_exception_from_s3(
                             sagemaker_session=self._job.sagemaker_session,
-                            s3_uri=s3_path_join(self._job.s3_uri, "exception.pkl"),
+                            s3_uri=s3_path_join(self._job.s3_uri, EXCEPTION_FOLDER),
                         )
                     except ServiceError as serr:
                         chained_e = serr.__cause__

src/sagemaker/remote_function/core/serialization.py

@@ -13,17 +13,114 @@
 """SageMaker remote function data serializer/deserializer."""
 from __future__ import absolute_import
 
+import dataclasses
+import json
+import os
+import sys
+
 import cloudpickle
 
 from typing import Any, Callable
 from sagemaker.remote_function.errors import ServiceError, SerializationError, DeserializationError
 from sagemaker.s3 import S3Downloader, S3Uploader
+from tblib import pickling_support
+
+
+def _get_python_version():
+    return f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+
+
+@dataclasses.dataclass
+class _MetaData:
+    """Metadata about the serialized data or functions."""
 
+    version: str = "2023-04-24"
+    python_version: str = _get_python_version()
+    serialization_module: str = "cloudpickle"
 
-# TODO: 1) use dask serializer instead of cloudpickle for data serialization.
-#  2) set the pickle protocol properly
-#  3) serialization/deserialization scheme needs to be explicitly versioned
-#  4) handle exceptions
+    def to_json(self):
+        return json.dumps(dataclasses.asdict(self)).encode()
+
+    @staticmethod
+    def from_json(s):
+        try:
+            obj = json.loads(s)
+        except json.decoder.JSONDecodeError:
+            raise DeserializationError("Corrupt metadata file. It is not a valid json file.")
+
+        metadata = _MetaData()
+        metadata.version = obj.get("version")
+        metadata.python_version = obj.get("python_version")
+        metadata.serialization_module = obj.get("serialization_module")
+
+        if not (
+            metadata.version == "2023-04-24" and metadata.serialization_module == "cloudpickle"
+        ):
+            raise DeserializationError(
+                f"Corrupt metadata file. Serialization approach {s} is not supported."
+            )
+
+        return metadata
+
+
+class CloudpickleSerializer:
+    """Serializer using cloudpickle."""
+
+    @staticmethod
+    def serialize(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: str = None):
+        """Serializes data object and uploads it to S3.
+
+        Args:
+            sagemaker_session (sagemaker.session.Session): The underlying Boto3 session which AWS service
+                 calls are delegated to.
+            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.
+        """
+        try:
+            bytes_to_upload = cloudpickle.dumps(obj)
+        except Exception as e:
+            if isinstance(
+                e, NotImplementedError
+            ) and "Instance of Run type is not allowed to be pickled." in str(e):
+                raise SerializationError(
+                    """You are trying to reference to a sagemaker.experiments.run.Run instance from within the function
+                       or passing it as a function argument.
+                       Instantiate a Run in the function or use load_run instead."""
+                ) from e
+
+            raise SerializationError(
+                "Error when serializing object of type [{}]: {}".format(type(obj).__name__, repr(e))
+            ) from e
+
+        _upload_bytes_to_s3(bytes_to_upload, s3_uri, s3_kms_key, sagemaker_session)
+
+    @staticmethod
+    def deserialize(sagemaker_session, s3_uri) -> Any:
+        """Downloads from S3 and then deserializes data objects.
+
+        Args:
+            sagemaker_session (sagemaker.session.Session): The underlying sagemaker session which AWS service
+                 calls are delegated to.
+            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.
+        """
+        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, repr(e))
+            ) from e
+
+
+# TODO: use dask serializer in case dask distributed is installed in users' environment.
 def serialize_func_to_s3(func: Callable, sagemaker_session, s3_uri, s3_kms_key=None):
     """Serializes function and uploads it to S3.
 
@@ -36,16 +133,13 @@ def serialize_func_to_s3(func: Callable, sagemaker_session, s3_uri, s3_kms_key=N
     Raises:
         SerializationError: when fail to serialize function to bytes.
     """
-    try:
-        bytes_to_upload = cloudpickle.dumps(func)
-    except Exception as e:
-        raise SerializationError(
-            "Error when serializing function [{}]: {}".format(
-                getattr(func, "__name__", repr(func)), repr(e)
-            )
-        ) from e
 
-    _upload_bytes_to_s3(bytes_to_upload, s3_uri, s3_kms_key, sagemaker_session)
+    _upload_bytes_to_s3(
+        _MetaData().to_json(), os.path.join(s3_uri, "metadata.json"), s3_kms_key, sagemaker_session
+    )
+    CloudpickleSerializer.serialize(
+        func, sagemaker_session, os.path.join(s3_uri, "payload.pkl"), s3_kms_key
+    )
 
 
 def deserialize_func_from_s3(sagemaker_session, s3_uri) -> Callable:
@@ -63,16 +157,11 @@ def deserialize_func_from_s3(sagemaker_session, s3_uri) -> Callable:
     Raises:
         DeserializationError: when fail to serialize function to bytes.
     """
-    bytes_to_deserialize = _read_bytes_from_s3(s3_uri, sagemaker_session)
+    _MetaData.from_json(
+        _read_bytes_from_s3(os.path.join(s3_uri, "metadata.json"), sagemaker_session)
+    )
 
-    try:
-        return cloudpickle.loads(bytes_to_deserialize)
-    except Exception as e:
-        raise DeserializationError(
-            "Error when deserializing bytes downloaded from {} to function: {}".format(
-                s3_uri, repr(e)
-            )
-        ) from e
+    return CloudpickleSerializer.deserialize(sagemaker_session, os.path.join(s3_uri, "payload.pkl"))
 
 
 def serialize_obj_to_s3(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: str = None):
@@ -87,21 +176,13 @@ def serialize_obj_to_s3(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: st
     Raises:
         SerializationError: when fail to serialize object to bytes.
     """
-    try:
-        bytes_to_upload = cloudpickle.dumps(obj)
-    except Exception as e:
-        if isinstance(
-            e, NotImplementedError
-        ) and "Instance of Run type is not allowed to be pickled." in str(e):
-            raise SerializationError(
-                "Remote function does not allow parameters of Run type."
-            ) from e
-
-        raise SerializationError(
-            "Error when serializing object of type [{}]: {}".format(type(obj).__name__, repr(e))
-        ) from e
 
-    _upload_bytes_to_s3(bytes_to_upload, s3_uri, s3_kms_key, sagemaker_session)
+    _upload_bytes_to_s3(
+        _MetaData().to_json(), os.path.join(s3_uri, "metadata.json"), s3_kms_key, sagemaker_session
+    )
+    CloudpickleSerializer.serialize(
+        obj, sagemaker_session, os.path.join(s3_uri, "payload.pkl"), s3_kms_key
+    )
 
 
 def deserialize_obj_from_s3(sagemaker_session, s3_uri) -> Any:
@@ -112,18 +193,59 @@ def deserialize_obj_from_s3(sagemaker_session, s3_uri) -> Any:
              calls are delegated to.
         s3_uri (str): S3 root uri to which resulting serialized artifacts will be uploaded.
     Returns :
-        List of deserialized python objects.
+        Deserialized python objects.
     Raises:
         DeserializationError: when fail to serialize object to bytes.
     """
-    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, repr(e))
-        ) from e
+    _MetaData.from_json(
+        _read_bytes_from_s3(os.path.join(s3_uri, "metadata.json"), sagemaker_session)
+    )
+
+    return CloudpickleSerializer.deserialize(sagemaker_session, os.path.join(s3_uri, "payload.pkl"))
+
+
+def serialize_exception_to_s3(
+    exc: Exception, sagemaker_session, s3_uri: str, s3_kms_key: str = None
+):
+    """Serializes exception with traceback and uploads it to S3.
+
+    Args:
+        sagemaker_session (sagemaker.session.Session): The underlying Boto3 session which AWS service
+             calls are delegated to.
+        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.
+        exc: Exception to be serialized and persisted
+    Raises:
+        SerializationError: when fail to serialize object to bytes.
+    """
+    pickling_support.install()
+    _upload_bytes_to_s3(
+        _MetaData().to_json(), os.path.join(s3_uri, "metadata.json"), s3_kms_key, sagemaker_session
+    )
+    CloudpickleSerializer.serialize(
+        exc, sagemaker_session, os.path.join(s3_uri, "payload.pkl"), s3_kms_key
+    )
+
+
+def deserialize_exception_from_s3(sagemaker_session, s3_uri) -> Any:
+    """Downloads from S3 and then deserializes exception.
+
+    Args:
+        sagemaker_session (sagemaker.session.Session): The underlying sagemaker session which AWS service
+             calls are delegated to.
+        s3_uri (str): S3 root uri to which resulting serialized artifacts will be uploaded.
+    Returns :
+        Deserialized exception with traceback.
+    Raises:
+        DeserializationError: when fail to serialize object to bytes.
+    """
+
+    _MetaData.from_json(
+        _read_bytes_from_s3(os.path.join(s3_uri, "metadata.json"), sagemaker_session)
+    )
+
+    return CloudpickleSerializer.deserialize(sagemaker_session, os.path.join(s3_uri, "payload.pkl"))
 
 
 def _upload_bytes_to_s3(bytes, s3_uri, s3_kms_key, sagemaker_session):