Add getnewargs_ex to Enum base class to fix Ray serialization error

[yerkoescalona]

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:

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:

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:

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:

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.

[AdrienVannson]

Hello, do you have a minimal working example?

I was not able to reproduce it simply. For example, this works well:

import ray
import betterproto

ray.init()

class Color(betterproto.Enum):
    RED = 1
    BLUE = 2

ray.get(ray.put(Color.RED))

I also tried using the enum as a parameter to a remote function.

As a note, I refactored the Enum in the new repo to make the implementation much simpler (see betterproto/python-betterproto2#72 ). I wouldn't recommend using this version for now as it is not documented and subject to changes, but I hope I can release something stable soon

[jbkoh]

I believe this is related to pickling as I had the same issue with my proto definitions. The following (anonymized) code produced the same error as above.

import cloudpickle
from myproto import MyEnum
cloudpickle.loads(cloudpickle.dumps(MyEnum.MY_VALUE))

From ChatGPT, I got the following response:

The error indicates that during unpickling, an Enum (or an Enum member) is being re-created in a way that passes an extra argument to its new method. In other words, cloudpickle is trying to reconstruct an Enum object from your ex_mapper, but it’s calling Enum.new with two arguments instead of one.

and was able to solve the problem by attaching this to betterproto.Enum. I'm not sure of its consequences to other places.

    def __reduce__(self):
        # Return a tuple that tells the unpickler how to rebuild the object.
        # Here we use the class itself and the value needed to reconstruct it.
        return (self.__class__, (self.value,))

Unfortunately, the Color example doesn't produce the same error but a different one

from enum import Enum
import betterproto
import cloudpickle

#class Color(Enum):
class Color(betterproto.Enum):
    RED = 1
    BLUE = 2

cloudpickle.loads(cloudpickle.dumps(Color.RED))

This produce the following error

File "......./site-packages/betterproto/enum.py", line 109, in setattr
raise AttributeError(f"{cls.name}: cannot reassign Enum members.")
AttributeError: Color: cannot reassign Enum members.

Tested this with master branch as of now: 6a65ca9

EDIT: If I comment out EnumType.setattr definition, the same error as Enum.__new__() takes 1 positional argument but 2 were given is produced.

[AdrienVannson]

Thank you, I was able to reproduce the second error ("cannot reassign Enum members.") on this version of betterproto.

Recently, I've been mostly working on a fork here: https://github.com/betterproto/python-betterproto2 In this fork, I redesigned the enumerations to use the standard IntEnum from Python. As far as I know, everything works already well there.

However, I don't recommend switching to this fork for now since I still need to make a few changes before it's ready. It shouldn't take too long

[jbkoh]

@AdrienVannson Thanks for the effort! What's the relationship between the fork and this repo? Will python-betterproto2 be the official repo for betterproto v2?

[AdrienVannson]

I think it will probably be the case, but I need to talk more about this with the original maintainers (they seem to agree but I've been doing a lot of things on my own so I need to make sure it's ok)