feature: serialize objs using dask (aws#907)

Co-authored-by: Namrata Madan <[email protected]>

Author: nmadan

requirements/extras/remote_function_requirements.txt

@@ -1,2 +1,2 @@
-cloudpickle==2.2.0
+cloudpickle==2.2.1
 tblib==1.7.0

requirements/extras/test_requirements.txt

@@ -21,4 +21,5 @@ sagemaker-experiments==0.1.35
 Jinja2==3.0.3
 pandas>=1.3.5,<1.5
 scikit-learn==1.0.2
-cloudpickle==2.2.0
+cloudpickle==2.2.1
+distributed==2022.2.0

src/sagemaker/remote_function/core/serialization.py

@@ -17,26 +17,43 @@
 import json
 import os
 import sys
+import pickle
+from enum import Enum
 
 import cloudpickle
 
 from typing import Any, Callable
+
+from sagemaker.s3 import s3_path_join
+
 from sagemaker.remote_function.errors import ServiceError, SerializationError, DeserializationError
 from sagemaker.s3 import S3Downloader, S3Uploader
 from tblib import pickling_support
 
+METADATA_FILE = "metadata.json"
+PAYLOAD_FILE = "payload.pkl"
+HEADER_FILE = "headers.pkl"
+FRAME_FILE = "frame-{}.dat"
+
 
 def _get_python_version():
     return f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
 
 
+class SerializationModule(str, Enum):
+    """Represents various serialization modules used."""
+
+    CLOUDPICKLE = "cloudpickle"
+    DASK = "dask"
+
+
 @dataclasses.dataclass
 class _MetaData:
     """Metadata about the serialized data or functions."""
 
+    serialization_module: SerializationModule
     version: str = "2023-04-24"
     python_version: str = _get_python_version()
-    serialization_module: str = "cloudpickle"
 
     def to_json(self):
         return json.dumps(dataclasses.asdict(self)).encode()
@@ -45,16 +62,13 @@ def to_json(self):
     def from_json(s):
         try:
             obj = json.loads(s)
-        except json.decoder.JSONDecodeError:
+            metadata = _MetaData(**obj)
+        except (json.decoder.JSONDecodeError, TypeError):
             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"
+        if (
+            metadata.version != "2023-04-24"
+            or metadata.serialization_module not in SerializationModule.__members__.values()
         ):
             raise DeserializationError(
                 f"Corrupt metadata file. Serialization approach {s} is not supported."
@@ -79,6 +93,12 @@ def serialize(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: str = None):
         Raises:
             SerializationError: when fail to serialize object to bytes.
         """
+        _upload_bytes_to_s3(
+            _MetaData(SerializationModule.CLOUDPICKLE).to_json(),
+            os.path.join(s3_uri, METADATA_FILE),
+            s3_kms_key,
+            sagemaker_session,
+        )
         try:
             bytes_to_upload = cloudpickle.dumps(obj)
         except Exception as e:
@@ -95,7 +115,76 @@ def serialize(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: str = None):
                 "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(
+            bytes_to_upload, os.path.join(s3_uri, PAYLOAD_FILE), 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(
+            os.path.join(s3_uri, PAYLOAD_FILE), sagemaker_session
+        )
+
+        try:
+            return cloudpickle.loads(bytes_to_deserialize)
+        except Exception as e:
+            raise DeserializationError(
+                "Error when deserializing bytes downloaded from {}: {}".format(
+                    os.path.join(s3_uri, PAYLOAD_FILE), repr(e)
+                )
+            ) from e
+
+
+class DaskSerializer:
+    """Serializer using Dask."""
+
+    @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.
+        """
+        import distributed.protocol as dask
+
+        _upload_bytes_to_s3(
+            _MetaData(SerializationModule.DASK).to_json(),
+            os.path.join(s3_uri, METADATA_FILE),
+            s3_kms_key,
+            sagemaker_session,
+        )
+        try:
+            header, frames = dask.serialize(obj, on_error="raise")
+        except Exception as e:
+            raise SerializationError(
+                "Error when serializing object of type [{}]: {}".format(type(obj).__name__, repr(e))
+            ) from e
+
+        _upload_bytes_to_s3(
+            pickle.dumps(header), s3_path_join(s3_uri, HEADER_FILE), s3_kms_key, sagemaker_session
+        )
+        for idx, frame in enumerate(frames):
+            frame = bytes(frame) if isinstance(frame, memoryview) else frame
+            _upload_bytes_to_s3(
+                frame, s3_path_join(s3_uri, FRAME_FILE.format(idx)), s3_kms_key, sagemaker_session
+            )
 
     @staticmethod
     def deserialize(sagemaker_session, s3_uri) -> Any:
@@ -110,19 +199,29 @@ def deserialize(sagemaker_session, s3_uri) -> Any:
         Raises:
             DeserializationError: when fail to serialize object to bytes.
         """
-        bytes_to_deserialize = _read_bytes_from_s3(s3_uri, sagemaker_session)
+        import distributed.protocol as dask
 
+        header_to_deserialize = _read_bytes_from_s3(
+            s3_path_join(s3_uri, HEADER_FILE), sagemaker_session
+        )
+        headers = pickle.loads(header_to_deserialize)
+        num_frames = len(headers["frame-lengths"]) if "frame-lengths" in headers else 1
+        frames = []
+        for idx in range(num_frames):
+            frame = _read_bytes_from_s3(
+                s3_path_join(s3_uri, FRAME_FILE.format(idx)), sagemaker_session
+            )
+            frames.append(frame)
         try:
-            return cloudpickle.loads(bytes_to_deserialize)
+            return dask.deserialize(headers, frames)
         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.
+    """Serializes function using cloudpickle and uploads it to S3.
 
     Args:
         sagemaker_session (sagemaker.session.Session): The underlying Boto3 session which AWS service
@@ -133,13 +232,7 @@ def serialize_func_to_s3(func: Callable, sagemaker_session, s3_uri, s3_kms_key=N
     Raises:
         SerializationError: when fail to serialize function to bytes.
     """
-
-    _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
-    )
+    CloudpickleSerializer.serialize(func, sagemaker_session, s3_uri, s3_kms_key)
 
 
 def deserialize_func_from_s3(sagemaker_session, s3_uri) -> Callable:
@@ -157,16 +250,16 @@ def deserialize_func_from_s3(sagemaker_session, s3_uri) -> Callable:
     Raises:
         DeserializationError: when fail to serialize function 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"))
+    _MetaData.from_json(_read_bytes_from_s3(os.path.join(s3_uri, METADATA_FILE), sagemaker_session))
+    return CloudpickleSerializer.deserialize(sagemaker_session, s3_uri)
 
 
 def serialize_obj_to_s3(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: str = None):
     """Serializes data object and uploads it to S3.
 
+    This method uses the Dask library to perform serialization if its already installed, otherwise,
+    it uses cloudpickle.
+
     Args:
         sagemaker_session (sagemaker.session.Session): The underlying Boto3 session which AWS service
              calls are delegated to.
@@ -177,12 +270,12 @@ def serialize_obj_to_s3(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: st
         SerializationError: when fail to serialize object to bytes.
     """
 
-    _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
-    )
+    try:
+        import distributed.protocol as dask  # noqa: F401
+
+        DaskSerializer.serialize(obj, sagemaker_session, s3_uri, s3_kms_key)
+    except ModuleNotFoundError:
+        CloudpickleSerializer.serialize(obj, sagemaker_session, s3_uri, s3_kms_key)
 
 
 def deserialize_obj_from_s3(sagemaker_session, s3_uri) -> Any:
@@ -197,12 +290,12 @@ def deserialize_obj_from_s3(sagemaker_session, s3_uri) -> Any:
     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)
+    metadata = _MetaData.from_json(
+        _read_bytes_from_s3(os.path.join(s3_uri, METADATA_FILE), sagemaker_session)
     )
-
-    return CloudpickleSerializer.deserialize(sagemaker_session, os.path.join(s3_uri, "payload.pkl"))
+    if metadata.serialization_module == SerializationModule.DASK:
+        return DaskSerializer.deserialize(sagemaker_session, s3_uri)
+    return CloudpickleSerializer.deserialize(sagemaker_session, s3_uri)
 
 
 def serialize_exception_to_s3(
@@ -220,12 +313,7 @@ def serialize_exception_to_s3(
         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
-    )
+    CloudpickleSerializer.serialize(exc, sagemaker_session, s3_uri, s3_kms_key)
 
 
 def deserialize_exception_from_s3(sagemaker_session, s3_uri) -> Any:
@@ -240,12 +328,8 @@ def deserialize_exception_from_s3(sagemaker_session, s3_uri) -> Any:
     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"))
+    _MetaData.from_json(_read_bytes_from_s3(os.path.join(s3_uri, METADATA_FILE), sagemaker_session))
+    return CloudpickleSerializer.deserialize(sagemaker_session, s3_uri)
 
 
 def _upload_bytes_to_s3(bytes, s3_uri, s3_kms_key, sagemaker_session):

src/sagemaker/remote_function/core/stored_function.py

@@ -18,7 +18,6 @@
 
 import sagemaker.remote_function.core.serialization as serialization
 
-
 logger = logging_config.get_logger()
 
 

tests/data/remote_function/requirements_dask.txt

@@ -0,0 +1 @@
+distributed==2022.2.0