Unable to Pickle a Cython coroutine

[XavierGeerinck]

What happened + What you expected to happen

cannot pickle '_cython_3_0_8.coroutine' object

also posted on Ray GitHub, but it seems to be a Cloudpickle issue

Test output

_____________________________________________________________________________________________________________________________ test_termination _____________________________________________________________________________________________________________________________

    def test_termination():
        """
        It should allow us to terminate an actor
        """
>       a1 = DemoCounterActor.options(name="demo").remote()

tests/unit/actor/test_actor.py:143: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/actor.py:686: in remote
    return actor_cls._remote(args=args, kwargs=kwargs, **updated_options)
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/_private/auto_init_hook.py:24: in auto_init_wrapper
    return fn(*args, **kwargs)
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/util/tracing/tracing_helper.py:388: in _invocation_actor_class_remote_span
    return method(self, args, kwargs, *_args, **_kwargs)
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/actor.py:890: in _remote
    worker.function_actor_manager.export_actor_class(
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/_private/function_manager.py:517: in export_actor_class
    serialized_actor_class = pickle_dumps(
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/_private/serialization.py:67: in pickle_dumps
    return pickle.dumps(obj)
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/cloudpickle/cloudpickle_fast.py:88: in dumps
    cp.dump(obj)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <ray.cloudpickle.cloudpickle_fast.CloudPickler object at 0x7facd7d30880>, obj = <class 'demo.actor_counter._modify_class.<locals>.Class'>

    def dump(self, obj):
        try:
>           return Pickler.dump(self, obj)
E           _pickle.PicklingError: Can't pickle <cyfunction DemoCounterActor.__init__ at 0x7fac8ee3c110>: it's not the same object as demo.actor_counter.DemoCounterActor.__init__

../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/ray/cloudpickle/cloudpickle_fast.py:733: PicklingError
============================================================================================================================= warnings summary =============================================================================================================================
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/pydantic/fields.py:792
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/pydantic/fields.py:792
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/pydantic/fields.py:792
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/pydantic/fields.py:792
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/pydantic/fields.py:792
../../../.pyenv/versions/3.8.18/lib/python3.8/site-packages/pydantic/fields.py:792
  /home/xanrin/.pyenv/versions/3.8.18/lib/python3.8/site-packages/pydantic/fields.py:792: PydanticDeprecatedSince20: Using extra keyword arguments on `Field` is deprecated and will be removed. Use `json_schema_extra` instead. (Extra keys: 'env'). Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.3/migration/
    warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
========================================================================================================================= short test summary info ==========================================================================================================================
FAILED tests/unit/actor/test_actor.py::test_termination - _pickle.PicklingError: Can't pickle <cyfunction DemoCounterActor.__init__ at 0x7fac8ee3c110>: it's not the same object as demo.actor_counter.DemoCounterActor.__init__
================================================================================================================= 1 failed, 1 passed, 6 warnings in 5.59s ==================================================================================================================

Stacktrace

__________________________________________________________________________________________________________________________ test_call_async_method __________________________________________________________________________________________________________________________

    @pytest.mark.asyncio
    async def test_call_async_method():
        """
        It should allow us to call a synchronous method on an actor
        """
        cls = ray.remote(DemoCounterActor)
        demo_actor = cls.options(name="demo").remote(counter=12)
>       assert await demo_actor.get_counter_async.remote() == 12
E       ray.exceptions.RayTaskError(TypeError): ray::DemoCounterActor.get_counter_async() (pid=1375208, ip=172.20.248.27, actor_id=2a97be7da0f4ff9c8a967ed001000000, repr=<demo.actor_counter.DemoCounterActor object at 0x7fea2b858e50>)
E         File "/home/xanrin/.pyenv/versions/3.9.18/lib/python3.9/site-packages/ray/cloudpickle/cloudpickle_fast.py", line 88, in dumps
E           cp.dump(obj)
E         File "/home/xanrin/.pyenv/versions/3.9.18/lib/python3.9/site-packages/ray/cloudpickle/cloudpickle_fast.py", line 733, in dump
E           return Pickler.dump(self, obj)
E       TypeError: cannot pickle '_cython_3_0_8.coroutine' object

tests/unit/actor/test_actor.py:131: RayTaskError(TypeError)

Versions / Dependencies

Tested on Ray 2.8.0 and Ray 2.9.2 on Ubuntu and Mac OS

Reproduction script

Full code: https://github.com/XavierGeerinck/issue-async-cython-pickle

Define a Demo Actor:

class DemoCounterActor:
    def __init__(self, counter: int = 0):
        self.counter = counter
        self.counter_nested = DemoCounterActorNested.remote(counter)

    def get_counter(self):
        return self.counter

    def get_counter_nested(self):
        return ray.get(self.counter_nested.get_counter.remote())

    def set_counter_nested(self, value: int):
        self.counter_nested.set_counter.remote(value)

    async def get_counter_async(self):
        return self.counter

    def increment(self) -> None:
        self.counter += 1

    async def increment_async(self) -> None:
        self.counter += 1

    def throw(self):
        raise ValueError("This is a sync test error")

    async def throw_async(self):
        raise ValueError("This is an async test error")

Cythonize it

from setuptools import setup
from Cython.Build import cythonize

setup(
    ext_modules = cythonize("example.pyx")
)

pip install cython
python setup.py build_ext --inplace

Then use a test:

@pytest.mark.asyncio
async def test_call_async_method():
    """
    It should allow us to call a synchronous method on an actor
    """
    cls = ray.remote(DemoCounterActor)
    demo_actor = cls.options(name="demo").remote(counter=12)
    assert await demo_actor.get_counter_async.remote() == 12
    await demo_actor.increment_async.remote()
    assert await demo_actor.get_counter_async.remote() == 13

    ray.kill(demo_actor)

Issue Severity

High: It blocks me from completing my task.

[ogrisel]

Thanks for the detailed report. However I don't think there is any easy way to implement pickling support for Cython coroutines.

Even regular Python coroutines are not picklable, either by pickle from the standard library or by cloudpickle:

>>> import pickle, cloudpickle
>>> async def f(a):
...     return a + 1
... 
>>> coroutine = f(0)
>>> pickle.loads(pickle.dumps(coroutine))
Traceback (most recent call last):
  Cell In[30], line 1
    pickle.loads(pickle.dumps(coroutine))
TypeError: cannot pickle 'coroutine' object

>>> pickle.loads(cloudpickle.dumps(coroutine))
Traceback (most recent call last):
  Cell In[31], line 1
    pickle.loads(cloudpickle.dumps(coroutine))
  File ~/code/cloudpickle/cloudpickle/cloudpickle.py:1518 in dumps
    cp.dump(obj)
  File ~/code/cloudpickle/cloudpickle/cloudpickle.py:1284 in dump
    return super().dump(obj)
TypeError: cannot pickle 'coroutine' object

[ogrisel]

Note that there is an old unmaintained prototype by @llllllllll to add support for pickling instances of generators discussed in #146 and implemented at https://github.com/llllllllll/cloudpickle-generators. That could be a start but I am not sure if it would be easy to generalize to regular asyncio coroutines and their Cython counterparts.