change: Add allow_pickle parameter to NumpyDeserializer (#1755)

Author: bveeramani

src/sagemaker/deserializers.py

@@ -163,13 +163,15 @@ class NumpyDeserializer(BaseDeserializer):
 
     ACCEPT = "application/x-npy"
 
-    def __init__(self, dtype=None):
-        """Initialize the dtype.
+    def __init__(self, dtype=None, allow_pickle=True):
+        """Initialize the dtype and allow_pickle arguments.
 
         Args:
-            dtype (str): The dtype of the data.
+            dtype (str): The dtype of the data (default: None).
+            allow_pickle (bool): Allow loading pickled object arrays (default: True).
         """
         self.dtype = dtype
+        self.allow_pickle = allow_pickle
 
     def deserialize(self, stream, content_type):
         """Deserialize data from an inference endpoint into a NumPy array.
@@ -189,7 +191,7 @@ def deserialize(self, stream, content_type):
             if content_type == "application/json":
                 return np.array(json.load(codecs.getreader("utf-8")(stream)), dtype=self.dtype)
             if content_type == "application/x-npy":
-                return np.load(io.BytesIO(stream.read()))
+                return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
         finally:
             stream.close()
 

tests/unit/sagemaker/test_deserializers.py

@@ -141,7 +141,7 @@ def test_numpy_deserializer_from_npy(numpy_deserializer):
 
 
 def test_numpy_deserializer_from_npy_object_array(numpy_deserializer):
-    array = np.array(["one", "two"])
+    array = np.array([{"a": "", "b": ""}, {"c": "", "d": ""}])
     stream = io.BytesIO()
     np.save(stream, array)
     stream.seek(0)
@@ -151,6 +151,18 @@ def test_numpy_deserializer_from_npy_object_array(numpy_deserializer):
     assert np.array_equal(array, result)
 
 
+def test_numpy_deserializer_from_npy_object_array_with_allow_pickle_false():
+    numpy_deserializer = NumpyDeserializer(allow_pickle=False)
+
+    array = np.array([{"a": "", "b": ""}, {"c": "", "d": ""}])
+    stream = io.BytesIO()
+    np.save(stream, array)
+    stream.seek(0)
+
+    with pytest.raises(ValueError):
+        numpy_deserializer.deserialize(stream, "application/x-npy")
+
+
 @pytest.fixture
 def json_deserializer():
     return JSONDeserializer()