Move RecordSerializer and RecordDeserializer to sagemaker.serializers and sagemaker.deserialzers (#5037)

pintaoz-aws · pintaoz · web-flow · commit 9e44f84ea4a9 · 2025-02-17T11:58:06.000-08:00
* Move RecordSerializer and RecordDeserializer to sagemaker.serializers and sagemaker.deserializers

* fix codestyle

* fix test

---------

Co-authored-by: pintaoz &lt;pintaoz@amazon.com&gt;
diff --git a/doc/v2.rst b/doc/v2.rst
@@ -324,9 +324,9 @@ The follow serializer/deserializer classes have been renamed and/or moved:
 +--------------------------------------------------------+-------------------------------------------------------+
 | ``sagemaker.predictor._NPYSerializer``                 | ``sagemaker.serializers.NumpySerializer``             |
 +--------------------------------------------------------+-------------------------------------------------------+
-| ``sagemaker.amazon.common.numpy_to_record_serializer`` | ``sagemaker.amazon.common.RecordSerializer``          |
+| ``sagemaker.amazon.common.numpy_to_record_serializer`` | ``sagemaker.serializers.RecordSerializer``            |
 +--------------------------------------------------------+-------------------------------------------------------+
-| ``sagemaker.amazon.common.record_deserializer``        | ``sagemaker.amazon.common.RecordDeserializer``        |
+| ``sagemaker.amazon.common.record_deserializer``        | ``sagemaker.deserializers.RecordDeserializer``        |
 +--------------------------------------------------------+-------------------------------------------------------+
 | ``sagemaker.predictor._JsonDeserializer``              | ``sagemaker.deserializers.JSONDeserializer``          |
 +--------------------------------------------------------+-------------------------------------------------------+
diff --git a/src/sagemaker/amazon/common.py b/src/sagemaker/amazon/common.py
@@ -13,84 +13,16 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-import io
 import logging
 import struct
 import sys
 
 import numpy as np
 
 from sagemaker.amazon.record_pb2 import Record
-from sagemaker.deprecations import deprecated_class
-from sagemaker.deserializers import SimpleBaseDeserializer
-from sagemaker.serializers import SimpleBaseSerializer
 from sagemaker.utils import DeferredError
 
 
-class RecordSerializer(SimpleBaseSerializer):
-    """Serialize a NumPy array for an inference request."""
-
-    def __init__(self, content_type="application/x-recordio-protobuf"):
-        """Initialize a ``RecordSerializer`` instance.
-
-        Args:
-            content_type (str): The MIME type to signal to the inference endpoint when sending
-                request data (default: "application/x-recordio-protobuf").
-        """
-        super(RecordSerializer, self).__init__(content_type=content_type)
-
-    def serialize(self, data):
-        """Serialize a NumPy array into a buffer containing RecordIO records.
-
-        Args:
-            data (numpy.ndarray): The data to serialize.
-
-        Returns:
-            io.BytesIO: A buffer containing the data serialized as records.
-        """
-        if len(data.shape) == 1:
-            data = data.reshape(1, data.shape[0])
-
-        if len(data.shape) != 2:
-            raise ValueError(
-                "Expected a 1D or 2D array, but got a %dD array instead." % len(data.shape)
-            )
-
-        buffer = io.BytesIO()
-        write_numpy_to_dense_tensor(buffer, data)
-        buffer.seek(0)
-
-        return buffer
-
-
-class RecordDeserializer(SimpleBaseDeserializer):
-    """Deserialize RecordIO Protobuf data from an inference endpoint."""
-
-    def __init__(self, accept="application/x-recordio-protobuf"):
-        """Initialize a ``RecordDeserializer`` instance.
-
-        Args:
-            accept (union[str, tuple[str]]): The MIME type (or tuple of allowable MIME types) that
-                is expected from the inference endpoint (default:
-                "application/x-recordio-protobuf").
-        """
-        super(RecordDeserializer, self).__init__(accept=accept)
-
-    def deserialize(self, data, content_type):
-        """Deserialize RecordIO Protobuf data from an inference endpoint.
-
-        Args:
-            data (object): The protobuf message to deserialize.
-            content_type (str): The MIME type of the data.
-        Returns:
-            list: A list of records.
-        """
-        try:
-            return read_records(data)
-        finally:
-            data.close()
-
-
 def _write_feature_tensor(resolved_type, record, vector):
     """Placeholder Docstring"""
     if resolved_type == "Int32":
@@ -288,7 +220,3 @@ def _resolve_type(dtype):
     if dtype == np.dtype("float32"):
         return "Float32"
     raise ValueError("Unsupported dtype {} on array".format(dtype))
-
-
-numpy_to_record_serializer = deprecated_class(RecordSerializer, "numpy_to_record_serializer")
-record_deserializer = deprecated_class(RecordDeserializer, "record_deserializer")
diff --git a/src/sagemaker/amazon/factorization_machines.py b/src/sagemaker/amazon/factorization_machines.py
@@ -17,11 +17,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/amazon/kmeans.py b/src/sagemaker/amazon/kmeans.py
@@ -17,11 +17,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge, le
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/amazon/knn.py b/src/sagemaker/amazon/knn.py
@@ -17,11 +17,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, isin
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/amazon/lda.py b/src/sagemaker/amazon/lda.py
@@ -18,11 +18,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/amazon/linear_learner.py b/src/sagemaker/amazon/linear_learner.py
@@ -18,11 +18,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import isin, gt, lt, ge, le
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/amazon/ntm.py b/src/sagemaker/amazon/ntm.py
@@ -17,11 +17,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le, isin
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/amazon/pca.py b/src/sagemaker/amazon/pca.py
@@ -17,11 +17,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/amazon/randomcutforest.py b/src/sagemaker/amazon/randomcutforest.py
@@ -17,11 +17,12 @@
 
 from sagemaker import image_uris
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
-from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le
+from sagemaker.deserializers import RecordDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.model import Model
+from sagemaker.serializers import RecordSerializer
 from sagemaker.session import Session
 from sagemaker.utils import pop_out_unused_kwarg
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
diff --git a/src/sagemaker/base_deserializers.py b/src/sagemaker/base_deserializers.py
@@ -23,6 +23,7 @@
 import numpy as np
 from six import with_metaclass
 
+from sagemaker.amazon.common import read_records
 from sagemaker.utils import DeferredError
 
 try:
@@ -388,3 +389,31 @@ def deserialize(self, stream, content_type="tensor/pt"):
                 "Unable to deserialize your data to torch.Tensor.\
                     Please provide custom deserializer in InferenceSpec."
             )
+
+
+class RecordDeserializer(SimpleBaseDeserializer):
+    """Deserialize RecordIO Protobuf data from an inference endpoint."""
+
+    def __init__(self, accept="application/x-recordio-protobuf"):
+        """Initialize a ``RecordDeserializer`` instance.
+
+        Args:
+            accept (union[str, tuple[str]]): The MIME type (or tuple of allowable MIME types) that
+                is expected from the inference endpoint (default:
+                "application/x-recordio-protobuf").
+        """
+        super(RecordDeserializer, self).__init__(accept=accept)
+
+    def deserialize(self, data, content_type):
+        """Deserialize RecordIO Protobuf data from an inference endpoint.
+
+        Args:
+            data (object): The protobuf message to deserialize.
+            content_type (str): The MIME type of the data.
+        Returns:
+            list: A list of records.
+        """
+        try:
+            return read_records(data)
+        finally:
+            data.close()
diff --git a/src/sagemaker/base_serializers.py b/src/sagemaker/base_serializers.py
@@ -22,6 +22,7 @@
 from pandas import DataFrame
 from six import with_metaclass
 
+from sagemaker.amazon.common import write_numpy_to_dense_tensor
 from sagemaker.utils import DeferredError
 
 try:
@@ -466,3 +467,39 @@ def serialize(self, data):
                 )
 
         raise ValueError("Object of type %s is not a torch.Tensor" % type(data))
+
+
+class RecordSerializer(SimpleBaseSerializer):
+    """Serialize a NumPy array for an inference request."""
+
+    def __init__(self, content_type="application/x-recordio-protobuf"):
+        """Initialize a ``RecordSerializer`` instance.
+
+        Args:
+            content_type (str): The MIME type to signal to the inference endpoint when sending
+                request data (default: "application/x-recordio-protobuf").
+        """
+        super(RecordSerializer, self).__init__(content_type=content_type)
+
+    def serialize(self, data):
+        """Serialize a NumPy array into a buffer containing RecordIO records.
+
+        Args:
+            data (numpy.ndarray): The data to serialize.
+
+        Returns:
+            io.BytesIO: A buffer containing the data serialized as records.
+        """
+        if len(data.shape) == 1:
+            data = data.reshape(1, data.shape[0])
+
+        if len(data.shape) != 2:
+            raise ValueError(
+                "Expected a 1D or 2D array, but got a %dD array instead." % len(data.shape)
+            )
+
+        buffer = io.BytesIO()
+        write_numpy_to_dense_tensor(buffer, data)
+        buffer.seek(0)
+
+        return buffer
diff --git a/src/sagemaker/cli/compatibility/v2/modifiers/serde.py b/src/sagemaker/cli/compatibility/v2/modifiers/serde.py
@@ -51,8 +51,8 @@
     "StreamDeserializer": ("sagemaker.deserializers",),
     "NumpyDeserializer": ("sagemaker.deserializers",),
     "JSONDeserializer": ("sagemaker.deserializers",),
-    "RecordSerializer ": ("sagemaker.amazon.common",),
-    "RecordDeserializer": ("sagemaker.amazon.common",),
+    "RecordSerializer ": ("sagemaker.serializers",),
+    "RecordDeserializer": ("sagemaker.deserializers",),
 }
 
 OLD_CLASS_NAME_TO_NEW_CLASS_NAME = {
@@ -101,8 +101,8 @@ def node_should_be_modified(self, node):
         - ``sagemaker.predictor.StreamDeserializer``
         - ``sagemaker.predictor._NumpyDeserializer``
         - ``sagemaker.predictor._JsonDeserializer``
-        - ``sagemaker.amazon.common.numpy_to_record_serializer``
-        - ``sagemaker.amazon.common.record_deserializer``
+        - ``sagemaker.serializers.numpy_to_record_serializer``
+        - ``sagemaker.deserializers.record_deserializer``
 
         Args:
             node (ast.Call): a node that represents a function call. For more,
@@ -128,8 +128,8 @@ def modify_node(self, node):
         - ``sagemaker.deserializers.StreamDeserializer``
         - ``sagemaker.deserializers.NumpyDeserializer``
         - ``sagemaker.deserializers._JsonDeserializer``
-        - ``sagemaker.amazon.common.RecordSerializer``
-        - ``sagemaker.amazon.common.RecordDeserializer``
+        - ``sagemaker.serializers.RecordSerializer``
+        - ``sagemaker.deserializers.RecordDeserializer``
 
         Args:
             node (ast.Call): a node that represents a SerDe constructor.
@@ -303,8 +303,8 @@ def node_should_be_modified(self, node):
         """Checks if the import statement imports a SerDe from the ``sagemaker.amazon.common``.
 
         This checks for:
-        - ``sagemaker.amazon.common.numpy_to_record_serializer``
-        - ``sagemaker.amazon.common.record_deserializer``
+        - ``sagemaker.serializers.numpy_to_record_serializer``
+        - ``sagemaker.deserializers.record_deserializer``
 
         Args:
             node (ast.ImportFrom): a node that represents a ``from ... import ... `` statement.
@@ -322,8 +322,8 @@ def modify_node(self, node):
         """Upgrades the ``numpy_to_record_serializer`` and ``record_deserializer`` imports.
 
         This upgrades the classes to (if applicable):
-        - ``sagemaker.amazon.common.RecordSerializer``
-        - ``sagemaker.amazon.common.RecordDeserializer``
+        - ``sagemaker.serializers.RecordSerializer``
+        - ``sagemaker.deserializers.RecordDeserializer``
 
         Args:
             node (ast.ImportFrom): a node that represents a ``from ... import ... `` statement.
diff --git a/src/sagemaker/deserializers.py b/src/sagemaker/deserializers.py
@@ -31,8 +31,10 @@
     StreamDeserializer,
     StringDeserializer,
     TorchTensorDeserializer,
+    RecordDeserializer,
 )
 
+from sagemaker.deprecations import deprecated_class
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 from sagemaker.jumpstart.enums import JumpStartModelType
@@ -150,3 +152,6 @@ def retrieve_default(
         model_type=model_type,
         config_name=config_name,
     )
+
+
+record_deserializer = deprecated_class(RecordDeserializer, "record_deserializer")
diff --git a/src/sagemaker/serializers.py b/src/sagemaker/serializers.py
@@ -30,8 +30,10 @@
     SparseMatrixSerializer,
     TorchTensorSerializer,
     StringSerializer,
+    RecordSerializer,
 )
 
+from sagemaker.deprecations import deprecated_class
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 from sagemaker.jumpstart.enums import JumpStartModelType
@@ -152,3 +154,6 @@ def retrieve_default(
         model_type=model_type,
         config_name=config_name,
     )
+
+
+numpy_to_record_serializer = deprecated_class(RecordSerializer, "numpy_to_record_serializer")
diff --git a/tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_serde.py b/tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_serde.py
diff --git a/tests/unit/test_common.py b/tests/unit/test_common.py
