Revert "feature: serialize objs using dask (aws#907)"

This reverts commit 27cfeff29f961fd6bb28bd91545de6ec888b284a.

Author: Namrata Madan

requirements/extras/remote_function_requirements.txt

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

requirements/extras/test_requirements.txt

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

src/sagemaker/remote_function/core/serialization.py

@@ -17,43 +17,26 @@
 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()
@@ -62,13 +45,16 @@ def to_json(self):
     def from_json(s):
         try:
             obj = json.loads(s)
-            metadata = _MetaData(**obj)
-        except (json.decoder.JSONDecodeError, TypeError):
+        except json.decoder.JSONDecodeError:
             raise DeserializationError("Corrupt metadata file. It is not a valid json file.")
 
-        if (
-            metadata.version != "2023-04-24"
-            or metadata.serialization_module not in SerializationModule.__members__.values()
+        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."
@@ -93,12 +79,6 @@ 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:
@@ -116,76 +96,7 @@ 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, 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
-            )
+        _upload_bytes_to_s3(bytes_to_upload, s3_uri, s3_kms_key, sagemaker_session)
 
     @staticmethod
     def deserialize(sagemaker_session, s3_uri) -> Any:
@@ -200,29 +111,19 @@ def deserialize(sagemaker_session, s3_uri) -> Any:
         Raises:
             DeserializationError: when fail to serialize object to bytes.
         """
-        import distributed.protocol as dask
+        bytes_to_deserialize = _read_bytes_from_s3(s3_uri, sagemaker_session)
 
-        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 dask.deserialize(headers, frames)
+            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 using cloudpickle and uploads it to S3.
+    """Serializes function and uploads it to S3.
 
     Args:
         sagemaker_session (sagemaker.session.Session): The underlying Boto3 session which AWS service
@@ -233,7 +134,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.
     """
-    CloudpickleSerializer.serialize(func, sagemaker_session, s3_uri, s3_kms_key)
+
+    _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:
@@ -251,16 +158,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_FILE), sagemaker_session))
-    return CloudpickleSerializer.deserialize(sagemaker_session, s3_uri)
+    _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_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.
@@ -271,12 +178,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.
     """
 
-    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)
+    _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:
@@ -291,12 +198,12 @@ def deserialize_obj_from_s3(sagemaker_session, s3_uri) -> Any:
     Raises:
         DeserializationError: when fail to serialize object to bytes.
     """
-    metadata = _MetaData.from_json(
-        _read_bytes_from_s3(os.path.join(s3_uri, METADATA_FILE), sagemaker_session)
+
+    _MetaData.from_json(
+        _read_bytes_from_s3(os.path.join(s3_uri, "metadata.json"), sagemaker_session)
     )
-    if metadata.serialization_module == SerializationModule.DASK:
-        return DaskSerializer.deserialize(sagemaker_session, s3_uri)
-    return CloudpickleSerializer.deserialize(sagemaker_session, s3_uri)
+
+    return CloudpickleSerializer.deserialize(sagemaker_session, os.path.join(s3_uri, "payload.pkl"))
 
 
 def serialize_exception_to_s3(
@@ -314,7 +221,12 @@ def serialize_exception_to_s3(
         SerializationError: when fail to serialize object to bytes.
     """
     pickling_support.install()
-    CloudpickleSerializer.serialize(exc, sagemaker_session, s3_uri, s3_kms_key)
+    _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:
@@ -329,8 +241,12 @@ 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_FILE), sagemaker_session))
-    return CloudpickleSerializer.deserialize(sagemaker_session, s3_uri)
+
+    _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):

src/sagemaker/remote_function/core/stored_function.py

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