Skip to content

Commit 1cf46b0

Browse files
gpauloskigpauloski
authored
Merge pull request #531 from proxystore/issue-530
Improve serialization error handling when resolving a proxy
2 parents 75c59b1 + 6c73249 commit 1cf46b0

9 files changed

Lines changed: 122 additions & 41 deletions

File tree

‎proxystore/connectors/multi.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424

2525
from proxystore import utils
2626
from proxystore.connectors.protocols import Connector
27-
from proxystore.utils.imports import get_class_path
28-
from proxystore.utils.imports import import_class
27+
from proxystore.utils.imports import get_object_path
28+
from proxystore.utils.imports import import_from_path
2929
from proxystore.warnings import ExperimentalWarning
3030

3131
warnings.warn(
@@ -305,7 +305,7 @@ def config(self) -> dict[str, ConnectorPolicyConfig]:
305305
configs.update(
306306
{
307307
name: (
308-
get_class_path(type(connector)),
308+
get_object_path(type(connector)),
309309
connector.config(),
310310
policy.as_dict(),
311311
)
@@ -329,7 +329,7 @@ def from_config(
329329
for name, (conn_path, conn_config, policy_dict) in config.items():
330330
policy = Policy(**policy_dict)
331331
if policy.is_valid_on_host():
332-
connector_type = import_class(conn_path)
332+
connector_type = import_from_path(conn_path)
333333
connector = connector_type.from_config(conn_config)
334334
connectors[name] = (connector, policy)
335335
else:

‎proxystore/serialize.py‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,11 @@ def serialize(obj: Any) -> bytes:
3030
Returns:
3131
Bytes that can be passed to \
3232
[`deserialize()`][proxystore.serialize.deserialize].
33+
34+
Raises:
35+
SerializationError: If serializing the object fails with all available
36+
serializers. Cloudpickle is the last resort, so this error will
37+
typically be raised from a cloudpickle error.
3338
"""
3439
if isinstance(obj, bytes):
3540
identifier = b'01\n'
@@ -44,7 +49,12 @@ def serialize(obj: Any) -> bytes:
4449
obj = pickle.dumps(obj, protocol=5)
4550
except Exception:
4651
identifier = b'04\n'
47-
obj = cloudpickle.dumps(obj)
52+
try:
53+
obj = cloudpickle.dumps(obj)
54+
except Exception as e:
55+
raise SerializationError(
56+
f'Object of type {type(obj)} is not serializable.',
57+
) from e
4858

4959
assert isinstance(identifier, bytes)
5060
assert isinstance(obj, bytes)
@@ -69,6 +79,8 @@ def deserialize(data: bytes) -> Any:
6979
[`serialize()`][proxystore.serialize.serialize] to indicate which
7080
serialization method was used (e.g., no serialization, pickle,
7181
etc.).
82+
SerializationError: If pickle or cloudpickle raise an exception
83+
when deserializing the object.
7284
"""
7385
if not isinstance(data, bytes):
7486
raise ValueError(
@@ -84,9 +96,19 @@ def deserialize(data: bytes) -> Any:
8496
elif identifier == b'02':
8597
return data.decode()
8698
elif identifier == b'03':
87-
return pickle.loads(data)
99+
try:
100+
return pickle.loads(data)
101+
except Exception as e:
102+
raise SerializationError(
103+
'Failed to deserialize object with pickle.',
104+
) from e
88105
elif identifier == b'04':
89-
return cloudpickle.loads(data)
106+
try:
107+
return cloudpickle.loads(data)
108+
except Exception as e:
109+
raise SerializationError(
110+
'Failed to deserialize object with cloudpickle.',
111+
) from e
90112
else:
91113
raise SerializationError(
92114
f'Unknown identifier {identifier!r} for deserialization,',

‎proxystore/store/base.py‎

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from proxystore.connectors.protocols import DeferrableConnector
2727
from proxystore.proxy import Proxy
2828
from proxystore.proxy import ProxyLocker
29+
from proxystore.serialize import SerializationError
2930
from proxystore.store.cache import LRUCache
3031
from proxystore.store.exceptions import NonProxiableTypeError
3132
from proxystore.store.factory import PollingStoreFactory
@@ -39,8 +40,8 @@
3940
from proxystore.store.types import ConnectorT
4041
from proxystore.store.types import DeserializerT
4142
from proxystore.store.types import SerializerT
42-
from proxystore.utils.imports import get_class_path
43-
from proxystore.utils.imports import import_class
43+
from proxystore.utils.imports import get_object_path
44+
from proxystore.utils.imports import import_from_path
4445
from proxystore.utils.timer import Timer
4546
from proxystore.warnings import ExperimentalWarning
4647

@@ -190,7 +191,7 @@ def config(self) -> dict[str, Any]:
190191
"""
191192
return {
192193
'name': self.name,
193-
'connector_type': get_class_path(type(self.connector)),
194+
'connector_type': get_object_path(type(self.connector)),
194195
'connector_config': self.connector.config(),
195196
'serializer': self._serializer,
196197
'deserializer': self._deserializer,
@@ -211,7 +212,7 @@ def from_config(cls, config: dict[str, Any]) -> Store[Any]:
211212
config = config.copy() # Avoid messing with callers version
212213
connector_type = config.pop('connector_type')
213214
connector_config = config.pop('connector_config')
214-
connector = import_class(connector_type)
215+
connector = import_from_path(connector_type)
215216
config['connector'] = connector.from_config(connector_config)
216217
return cls(**config)
217218

@@ -409,6 +410,10 @@ def get(
409410
410411
Returns:
411412
Object or `None` if the object does not exist.
413+
414+
Raises:
415+
SerializationError: If an exception is caught when deserializing
416+
the object associated with the key.
412417
"""
413418
timer = Timer()
414419
timer.start()
@@ -437,10 +442,19 @@ def get(
437442

438443
if value is not None:
439444
with Timer() as deserializer_timer:
440-
if deserializer is not None:
445+
deserializer = (
446+
deserializer
447+
if deserializer is not None
448+
else self.deserializer
449+
)
450+
try:
441451
result = deserializer(value)
442-
else:
443-
result = self.deserializer(value)
452+
except Exception as e:
453+
name = get_object_path(deserializer)
454+
raise SerializationError(
455+
'Failed to deserialize object '
456+
f'(deserializer={name}, key={key}).',
457+
) from e
444458

445459
if self.metrics is not None:
446460
dtime = deserializer_timer.elapsed_ns

‎proxystore/stream/events.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@
1717
from typing import Any
1818
from typing import Union
1919

20-
from proxystore.utils.imports import get_class_path
21-
from proxystore.utils.imports import import_class
20+
from proxystore.utils.imports import get_object_path
21+
from proxystore.utils.imports import import_from_path
2222

2323

2424
@dataclasses.dataclass
@@ -54,15 +54,15 @@ def from_key(
5454
) -> NewObjectEvent:
5555
"""Create a new event from a key and metadata."""
5656
return cls(
57-
key_type=get_class_path(type(key)),
57+
key_type=get_object_path(type(key)),
5858
raw_key=list(key),
5959
evict=evict,
6060
metadata=metadata,
6161
)
6262

6363
def get_key(self) -> Any:
6464
"""Get the object key associated with the event."""
65-
key_type = import_class(self.key_type)
65+
key_type = import_from_path(self.key_type)
6666
return key_type(*self.raw_key)
6767

6868

‎proxystore/utils/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,5 @@
77
from proxystore.utils.data import readable_to_bytes
88
from proxystore.utils.environment import home_dir
99
from proxystore.utils.environment import hostname
10-
from proxystore.utils.imports import get_class_path
11-
from proxystore.utils.imports import import_class
10+
from proxystore.utils.imports import get_object_path
11+
from proxystore.utils.imports import import_from_path

‎proxystore/utils/imports.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,47 +6,47 @@
66
from typing import Any
77

88

9-
def get_class_path(cls: type[Any]) -> str:
10-
"""Get the fully qualified path of a type.
9+
def get_object_path(obj: Any) -> str:
10+
"""Get the fully qualified path of an object.
1111
1212
Example:
1313
```python
1414
>>> from proxystore.connectors.protocols import Connector
15-
>>> get_class_path(Connector)
15+
>>> get_object_path(Connector)
1616
'proxystore.connectors.protocols.Connector'
1717
```
1818
1919
Args:
20-
cls: Class type to get fully qualified path of.
20+
obj: Object to get fully qualified path of.
2121
2222
Returns:
23-
Fully qualified path of `cls`.
23+
Fully qualified path of `obj`.
2424
"""
25-
return f'{cls.__module__}.{cls.__qualname__}'
25+
return f'{obj.__module__}.{obj.__qualname__}'
2626

2727

28-
def import_class(path: str) -> type[Any]:
29-
"""Import class via its fully qualified path.
28+
def import_from_path(path: str) -> type[Any]:
29+
"""Import object via its fully qualified path.
3030
3131
Example:
3232
```python
33-
>>> import_class('proxystore.connectors.protocols.Connector')
33+
>>> import_from_path('proxystore.connectors.protocols.Connector')
3434
<class 'proxystore.connectors.protocols.Connector'>
3535
```
3636
3737
Args:
38-
path: Fully qualified path of class to import.
38+
path: Fully qualified path of object to import.
3939
4040
Returns:
41-
Imported class.
41+
Imported object.
4242
4343
Raises:
44-
ImportError: If a class at the `path` is not found.
44+
ImportError: If an object at the `path` is not found.
4545
"""
4646
module_path, _, name = path.rpartition('.')
4747
if len(module_path) == 0:
4848
raise ImportError(
49-
f'Class path must contain at least one module. Got {path}',
49+
f'Object path must contain at least one module. Got {path}',
5050
)
5151
module = importlib.import_module(module_path)
5252
return getattr(module, name)

‎tests/serialization_test.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
from __future__ import annotations
22

3+
from unittest import mock
4+
35
import pytest
46

57
from proxystore.serialize import deserialize
@@ -35,3 +37,32 @@ def test_serialization() -> None:
3537
with pytest.raises(SerializationError):
3638
# Fake identifier 'xxx'
3739
deserialize(b'99\nxxx')
40+
41+
42+
def test_cloudpickle_dumps_error() -> None:
43+
with mock.patch('cloudpickle.dumps', side_effect=Exception()):
44+
with pytest.raises(
45+
SerializationError,
46+
match="Object of type <class 'function'> is not serializable.",
47+
):
48+
serialize(lambda x: x + x) # pragma: no cover
49+
50+
51+
def test_pickle_loads_error() -> None:
52+
v = serialize([1, 2, 3])
53+
with mock.patch('pickle.loads', side_effect=Exception()):
54+
with pytest.raises(
55+
SerializationError,
56+
match='Failed to deserialize object with pickle.',
57+
):
58+
deserialize(v)
59+
60+
61+
def test_cloudpickle_loads_error() -> None:
62+
v = serialize(lambda x: x + x) # pragma: no cover
63+
with mock.patch('cloudpickle.loads', side_effect=Exception()):
64+
with pytest.raises(
65+
SerializationError,
66+
match='Failed to deserialize object with cloudpickle.',
67+
):
68+
deserialize(v)

‎tests/store/store_basics_test.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
from proxystore.connectors.local import LocalConnector
1010
from proxystore.proxy import Proxy
11+
from proxystore.serialize import SerializationError
1112
from proxystore.store import Store
1213
from proxystore.store.future import Future
1314
from proxystore.store.lifetimes import ContextLifetime
@@ -101,6 +102,19 @@ def test_custom_serializer(store: Store[LocalConnector]) -> None:
101102
store.put_batch([[1, 2, 3]], serializer=lambda s: s)
102103

103104

105+
def test_custom_deserializer_error(store: Store[LocalConnector]) -> None:
106+
key = store.put('value')
107+
108+
def _deserialize(x: bytes) -> Any:
109+
raise Exception()
110+
111+
with pytest.raises(
112+
SerializationError,
113+
match='Failed to deserialize object',
114+
):
115+
store.get(key, deserializer=_deserialize)
116+
117+
104118
def test_put_batch(store: Store[LocalConnector]) -> None:
105119
values = ['test_value1', 'test_value2', 'test_value3']
106120

‎tests/utils/imports_test.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66

77
from proxystore.connectors.file import FileConnector
88
from proxystore.connectors.local import LocalConnector
9-
from proxystore.utils.imports import get_class_path
10-
from proxystore.utils.imports import import_class
9+
from proxystore.utils.imports import get_object_path
10+
from proxystore.utils.imports import import_from_path
1111

1212

1313
@pytest.mark.parametrize(
@@ -17,8 +17,8 @@
1717
(LocalConnector, 'proxystore.connectors.local.LocalConnector'),
1818
),
1919
)
20-
def test_get_class_path(cls: type[Any], expected: str) -> None:
21-
assert get_class_path(cls) == expected
20+
def test_get_object_path(cls: type[Any], expected: str) -> None:
21+
assert get_object_path(cls) == expected
2222

2323

2424
@pytest.mark.parametrize(
@@ -29,10 +29,10 @@ def test_get_class_path(cls: type[Any], expected: str) -> None:
2929
('typing.Any', Any),
3030
),
3131
)
32-
def test_import_class(path: str, expected: type[Any]) -> None:
33-
assert import_class(path) == expected
32+
def test_import_from_path(path: str, expected: type[Any]) -> None:
33+
assert import_from_path(path) == expected
3434

3535

36-
def test_import_class_missing_path() -> None:
36+
def test_import_from_path_missing_path() -> None:
3737
with pytest.raises(ImportError):
38-
import_class('FileConnector')
38+
import_from_path('FileConnector')

0 commit comments

Comments
 (0)