proxystore · gpauloski · Apr 4, 2024 · Apr 4, 2024 · Apr 4, 2024 · Apr 4, 2024
diff --git a/proxystore/connectors/multi.py b/proxystore/connectors/multi.py
@@ -24,8 +24,8 @@
 
 from proxystore import utils
 from proxystore.connectors.protocols import Connector
-from proxystore.utils.imports import get_class_path
-from proxystore.utils.imports import import_class
+from proxystore.utils.imports import get_object_path
+from proxystore.utils.imports import import_from_path
 from proxystore.warnings import ExperimentalWarning
 
 warnings.warn(
@@ -305,7 +305,7 @@ def config(self) -> dict[str, ConnectorPolicyConfig]:
         configs.update(
             {
                 name: (
-                    get_class_path(type(connector)),
+                    get_object_path(type(connector)),
                     connector.config(),
                     policy.as_dict(),
                 )
@@ -329,7 +329,7 @@ def from_config(
         for name, (conn_path, conn_config, policy_dict) in config.items():
             policy = Policy(**policy_dict)
             if policy.is_valid_on_host():
-                connector_type = import_class(conn_path)
+                connector_type = import_from_path(conn_path)
                 connector = connector_type.from_config(conn_config)
                 connectors[name] = (connector, policy)
             else:

diff --git a/proxystore/serialize.py b/proxystore/serialize.py
@@ -30,6 +30,11 @@ def serialize(obj: Any) -> bytes:
     Returns:
         Bytes that can be passed to \
         [`deserialize()`][proxystore.serialize.deserialize].
+
+    Raises:
+        SerializationError: If serializing the object fails with all available
+            serializers. Cloudpickle is the last resort, so this error will
+            typically be raised from a cloudpickle error.
     """
     if isinstance(obj, bytes):
         identifier = b'01\n'
@@ -44,7 +49,12 @@ def serialize(obj: Any) -> bytes:
             obj = pickle.dumps(obj, protocol=5)
         except Exception:
             identifier = b'04\n'
-            obj = cloudpickle.dumps(obj)
+            try:
+                obj = cloudpickle.dumps(obj)
+            except Exception as e:
+                raise SerializationError(
+                    f'Object of type {type(obj)} is not serializable.',
+                ) from e
 
     assert isinstance(identifier, bytes)
     assert isinstance(obj, bytes)
@@ -69,6 +79,8 @@ def deserialize(data: bytes) -> Any:
             [`serialize()`][proxystore.serialize.serialize] to indicate which
             serialization method was used (e.g., no serialization, pickle,
             etc.).
+        SerializationError: If pickle or cloudpickle raise an exception
+            when deserializing the object.
     """
     if not isinstance(data, bytes):
         raise ValueError(
@@ -84,9 +96,19 @@ def deserialize(data: bytes) -> Any:
     elif identifier == b'02':
         return data.decode()
     elif identifier == b'03':
-        return pickle.loads(data)
+        try:
+            return pickle.loads(data)
+        except Exception as e:
+            raise SerializationError(
+                'Failed to deserialize object with pickle.',
+            ) from e
     elif identifier == b'04':
-        return cloudpickle.loads(data)
+        try:
+            return cloudpickle.loads(data)
+        except Exception as e:
+            raise SerializationError(
+                'Failed to deserialize object with cloudpickle.',
+            ) from e
     else:
         raise SerializationError(
             f'Unknown identifier {identifier!r} for deserialization,',

diff --git a/proxystore/store/base.py b/proxystore/store/base.py
@@ -26,6 +26,7 @@
 from proxystore.connectors.protocols import DeferrableConnector
 from proxystore.proxy import Proxy
 from proxystore.proxy import ProxyLocker
+from proxystore.serialize import SerializationError
 from proxystore.store.cache import LRUCache
 from proxystore.store.exceptions import NonProxiableTypeError
 from proxystore.store.factory import PollingStoreFactory
@@ -39,8 +40,8 @@
 from proxystore.store.types import ConnectorT
 from proxystore.store.types import DeserializerT
 from proxystore.store.types import SerializerT
-from proxystore.utils.imports import get_class_path
-from proxystore.utils.imports import import_class
+from proxystore.utils.imports import get_object_path
+from proxystore.utils.imports import import_from_path
 from proxystore.utils.timer import Timer
 from proxystore.warnings import ExperimentalWarning
 
@@ -190,7 +191,7 @@ def config(self) -> dict[str, Any]:
         """
         return {
             'name': self.name,
-            'connector_type': get_class_path(type(self.connector)),
+            'connector_type': get_object_path(type(self.connector)),
             'connector_config': self.connector.config(),
             'serializer': self._serializer,
             'deserializer': self._deserializer,
@@ -211,7 +212,7 @@ def from_config(cls, config: dict[str, Any]) -> Store[Any]:
         config = config.copy()  # Avoid messing with callers version
         connector_type = config.pop('connector_type')
         connector_config = config.pop('connector_config')
-        connector = import_class(connector_type)
+        connector = import_from_path(connector_type)
         config['connector'] = connector.from_config(connector_config)
         return cls(**config)
 
@@ -409,6 +410,10 @@ def get(
 
         Returns:
             Object or `None` if the object does not exist.
+
+        Raises:
+            SerializationError: If an exception is caught when deserializing
+                the object associated with the key.
         """
         timer = Timer()
         timer.start()
@@ -437,10 +442,19 @@ def get(
 
         if value is not None:
             with Timer() as deserializer_timer:
-                if deserializer is not None:
+                deserializer = (
+                    deserializer
+                    if deserializer is not None
+                    else self.deserializer
+                )
+                try:
                     result = deserializer(value)
-                else:
-                    result = self.deserializer(value)
+                except Exception as e:
+                    name = get_object_path(deserializer)
+                    raise SerializationError(
+                        'Failed to deserialize object '
+                        f'(deserializer={name}, key={key}).',
+                    ) from e
 
             if self.metrics is not None:
                 dtime = deserializer_timer.elapsed_ns

diff --git a/proxystore/stream/events.py b/proxystore/stream/events.py
@@ -17,8 +17,8 @@
 from typing import Any
 from typing import Union
 
-from proxystore.utils.imports import get_class_path
-from proxystore.utils.imports import import_class
+from proxystore.utils.imports import get_object_path
+from proxystore.utils.imports import import_from_path
 
 
 @dataclasses.dataclass
@@ -54,15 +54,15 @@ def from_key(
     ) -> NewObjectEvent:
         """Create a new event from a key and metadata."""
         return cls(
-            key_type=get_class_path(type(key)),
+            key_type=get_object_path(type(key)),
             raw_key=list(key),
             evict=evict,
             metadata=metadata,
         )
 
     def get_key(self) -> Any:
         """Get the object key associated with the event."""
-        key_type = import_class(self.key_type)
+        key_type = import_from_path(self.key_type)
         return key_type(*self.raw_key)
 
 

diff --git a/proxystore/utils/__init__.py b/proxystore/utils/__init__.py
@@ -7,5 +7,5 @@
 from proxystore.utils.data import readable_to_bytes
 from proxystore.utils.environment import home_dir
 from proxystore.utils.environment import hostname
-from proxystore.utils.imports import get_class_path
-from proxystore.utils.imports import import_class
+from proxystore.utils.imports import get_object_path
+from proxystore.utils.imports import import_from_path
diff --git a/proxystore/utils/imports.py b/proxystore/utils/imports.py
@@ -6,47 +6,47 @@
 from typing import Any
 
 
-def get_class_path(cls: type[Any]) -> str:
-    """Get the fully qualified path of a type.
+def get_object_path(obj: Any) -> str:
+    """Get the fully qualified path of an object.
 
     Example:
         ```python
         >>> from proxystore.connectors.protocols import Connector
-        >>> get_class_path(Connector)
+        >>> get_object_path(Connector)
         'proxystore.connectors.protocols.Connector'
         ```
 
     Args:
-        cls: Class type to get fully qualified path of.
+        obj: Object to get fully qualified path of.
 
     Returns:
-        Fully qualified path of `cls`.
+        Fully qualified path of `obj`.
     """
-    return f'{cls.__module__}.{cls.__qualname__}'
+    return f'{obj.__module__}.{obj.__qualname__}'
 
 
-def import_class(path: str) -> type[Any]:
-    """Import class via its fully qualified path.
+def import_from_path(path: str) -> type[Any]:
+    """Import object via its fully qualified path.
 
     Example:
         ```python
-        >>> import_class('proxystore.connectors.protocols.Connector')
+        >>> import_from_path('proxystore.connectors.protocols.Connector')
         <class 'proxystore.connectors.protocols.Connector'>
         ```
 
     Args:
-        path: Fully qualified path of class to import.
+        path: Fully qualified path of object to import.
 
     Returns:
-        Imported class.
+        Imported object.
 
     Raises:
-        ImportError: If a class at the `path` is not found.
+        ImportError: If an object at the `path` is not found.
     """
     module_path, _, name = path.rpartition('.')
     if len(module_path) == 0:
         raise ImportError(
-            f'Class path must contain at least one module. Got {path}',
+            f'Object path must contain at least one module. Got {path}',
         )
     module = importlib.import_module(module_path)
     return getattr(module, name)
diff --git a/tests/serialization_test.py b/tests/serialization_test.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from unittest import mock
+
 import pytest
 
 from proxystore.serialize import deserialize
@@ -35,3 +37,32 @@ def test_serialization() -> None:
     with pytest.raises(SerializationError):
         # Fake identifier 'xxx'
         deserialize(b'99\nxxx')
+
+
+def test_cloudpickle_dumps_error() -> None:
+    with mock.patch('cloudpickle.dumps', side_effect=Exception()):
+        with pytest.raises(
+            SerializationError,
+            match="Object of type <class 'function'> is not serializable.",
+        ):
+            serialize(lambda x: x + x)  # pragma: no cover
+
+
+def test_pickle_loads_error() -> None:
+    v = serialize([1, 2, 3])
+    with mock.patch('pickle.loads', side_effect=Exception()):
+        with pytest.raises(
+            SerializationError,
+            match='Failed to deserialize object with pickle.',
+        ):
+            deserialize(v)
+
+
+def test_cloudpickle_loads_error() -> None:
+    v = serialize(lambda x: x + x)  # pragma: no cover
+    with mock.patch('cloudpickle.loads', side_effect=Exception()):
+        with pytest.raises(
+            SerializationError,
+            match='Failed to deserialize object with cloudpickle.',
+        ):
+            deserialize(v)
diff --git a/tests/store/store_basics_test.py b/tests/store/store_basics_test.py
@@ -8,6 +8,7 @@
 
 from proxystore.connectors.local import LocalConnector
 from proxystore.proxy import Proxy
+from proxystore.serialize import SerializationError
 from proxystore.store import Store
 from proxystore.store.future import Future
 from proxystore.store.lifetimes import ContextLifetime
@@ -101,6 +102,19 @@ def test_custom_serializer(store: Store[LocalConnector]) -> None:
         store.put_batch([[1, 2, 3]], serializer=lambda s: s)
 
 
+def test_custom_deserializer_error(store: Store[LocalConnector]) -> None:
+    key = store.put('value')
+
+    def _deserialize(x: bytes) -> Any:
+        raise Exception()
+
+    with pytest.raises(
+        SerializationError,
+        match='Failed to deserialize object',
+    ):
+        store.get(key, deserializer=_deserialize)
+
+
 def test_put_batch(store: Store[LocalConnector]) -> None:
     values = ['test_value1', 'test_value2', 'test_value3']
 

diff --git a/tests/utils/imports_test.py b/tests/utils/imports_test.py
@@ -6,8 +6,8 @@
 
 from proxystore.connectors.file import FileConnector
 from proxystore.connectors.local import LocalConnector
-from proxystore.utils.imports import get_class_path
-from proxystore.utils.imports import import_class
+from proxystore.utils.imports import get_object_path
+from proxystore.utils.imports import import_from_path
 
 
 @pytest.mark.parametrize(
@@ -17,8 +17,8 @@
         (LocalConnector, 'proxystore.connectors.local.LocalConnector'),
     ),
 )
-def test_get_class_path(cls: type[Any], expected: str) -> None:
-    assert get_class_path(cls) == expected
+def test_get_object_path(cls: type[Any], expected: str) -> None:
+    assert get_object_path(cls) == expected
 
 
 @pytest.mark.parametrize(
@@ -29,10 +29,10 @@ def test_get_class_path(cls: type[Any], expected: str) -> None:
         ('typing.Any', Any),
     ),
 )
-def test_import_class(path: str, expected: type[Any]) -> None:
-    assert import_class(path) == expected
+def test_import_from_path(path: str, expected: type[Any]) -> None:
+    assert import_from_path(path) == expected
 
 
-def test_import_class_missing_path() -> None:
+def test_import_from_path_missing_path() -> None:
     with pytest.raises(ImportError):
-        import_class('FileConnector')
+        import_from_path('FileConnector')