danielgtaylor / danielgtaylor/python-betterproto

Add getnewargs_ex to Enum base class to fix Ray serialization error

Open
#658 5 comments 4 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.8k
Forks
234
PR merge metrics
No merged PRs in 30d

Description

When using betterproto enums in a Ray environment, I encounter a pickling error when a task attempts to deserialize an enum member. For example, given an enum defined as follows:

```python
class Event(betterproto.Enum):
START = 0
STOP = 1
```
and using this enum as part of a task’s arguments, I see the following error when deserializing:

```python
ray.exceptions.RaySystemError: System error: Enum.__new__() takes 1 positional argument but 2 were given
```

This occurs during the call to ray.util.get_objects() (or similar), where Ray’s unpickler passes an extra argument to the Enum’s new method.

Workaround:
I found that manually registering serializers and deserializes for each enum (e.g., using:

```python
ray.util.register_serializer(
Event,
serializer=lambda obj: obj.value,
deserializer=lambda value: Event(value),
)
```
solves the problem.

Proposed Solution:
The core issue seems to be that the default pickling mechanism for betterproto’s Enum (a subclass of IntEnum) is not handling the construction correctly in the Ray environment. The standard approach to helping pickle an enum instance correctly is implementing the __getnewargs_ex__ method.

For example, adding the following method to the Enum base class:

```python
def __getnewargs_ex__(self):
# Provide no positional arguments and the keyword arguments needed for __new__
return (), {"name": self.name, "value": self.value}
```

allows the default pickle protocol to serialize and deserialize the enum members properly. With this change, Ray’s deserialization will invoke the Enum’s new with the correct keyword arguments, avoiding the TypeError.

Additional Context:
- Environment: betterproto 2.0.0b7, Python 3.10 (with Ray)

Request:
- Would it be possible to incorporate a __getnewargs_ex__ method in the base Enum class of betterproto? This would allow the standard pickle mechanism (and hence Ray’s object deserialization) to work without additional registration steps.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.