ray-project / ray

Ray is a unified framework for scaling AI and Python applications. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.
https://ray.io
Apache License 2.0
33.2k stars 5.61k forks source link

[Serve] `Can't pickle <functools._lru_cache_wrapper object` when using generic types in type hint #43452

Open auderson opened 7 months ago

auderson commented 7 months ago

What happened + What you expected to happen

Using typing.Generic[T] causes PicklingError with ray serve. See reproduction script below. A workaround of this is to use string annotations or use from __future__ import annoations.

Versions / Dependencies

ray.__version__ '2.9.1'

Reproduction script

import typing

import ray
from fastapi import FastAPI
from ray import serve

app = FastAPI()

T = typing.TypeVar("T")

@ray.remote
class Worker(typing.Generic[T]):
    pass

@serve.deployment
@serve.ingress(app)
class Ingress:
    def __init__(self, worker: Worker):
        self.worker = worker
Traceback (most recent call last):
  File "/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3553, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-2-1663390c0cbe>", line 24, in <module>
    class Ingress:
  File "/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/ray/serve/api.py", line 429, in deployment
    return decorator(_func_or_class) if callable(_func_or_class) else decorator
  File "/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/ray/serve/api.py", line 394, in decorator
    replica_config = ReplicaConfig.create(
  File "/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/ray/serve/_private/config.py", line 426, in create
    pickle_dumps(
  File "/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/ray/_private/serialization.py", line 65, in pickle_dumps
    return pickle.dumps(obj)
  File "/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/ray/cloudpickle/cloudpickle_fast.py", line 88, in dumps
    cp.dump(obj)
  File "/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/ray/cloudpickle/cloudpickle_fast.py", line 733, in dump
    return Pickler.dump(self, obj)
_pickle.PicklingError: Can't pickle <functools._lru_cache_wrapper object at 0x7f315813bd70>: it's not the same object as typing.Generic.__class_getitem__
from ray.util import inspect_serializability

inspect_serializability(Worker)
================================================================================
Checking Serializability of <__main__.ActorClass(Worker) object at 0x7f3134ecc1f0>
================================================================================
!!! FAIL serialization: Can't pickle <functools._lru_cache_wrapper object at 0x7f315813bd70>: it's not the same object as typing.Generic.__class_getitem__
    Serializing '_annotated' ActorClass...
    Serializing '_default_options' {}...
    Serializing '_is_protocol' False...
    Serializing '_ray_from_function_descriptor' <bound method ActorClass._ray_from_function_descriptor of <class '__main__.ActorClass(Worker)'>>...
    Serializing '_ray_from_modified_class' <bound method ActorClass._ray_from_modified_class of <class '__main__.ActorClass(Worker)'>>...
    Serializing '_remote' <bound method ActorClass._remote of <__main__.ActorClass(Worker) object at 0x7f3134ecc1f0>>...
    !!! FAIL serialization: Can't pickle <functools._lru_cache_wrapper object at 0x7f315813bd70>: it's not the same object as typing.Generic.__class_getitem__
        Serializing '__func__' <function ActorClass._remote at 0x7f314ce37a30>...
        Serializing '__wrapped__' <function ActorClass._remote at 0x7f314ce379a0>...
    WARNING: Did not find non-serializable object in <bound method ActorClass._remote of <__main__.ActorClass(Worker) object at 0x7f3134ecc1f0>>. This may be an oversight.
================================================================================
Variable: 
    FailTuple(_remote [obj=<bound method ActorClass._remote of <__main__.ActorClass(Worker) object at 0x7f3134ecc1f0>>, parent=<__main__.ActorClass(Worker) object at 0x7f3134ecc1f0>])
was found to be non-serializable. There may be multiple other undetected variables that were non-serializable. 
Consider either removing the instantiation/imports of these variables or moving the instantiation into the scope of the function/class. 
================================================================================
Check https://docs.ray.io/en/master/ray-core/objects/serialization.html#troubleshooting for more information.
If you have any suggestions on how to improve this error message, please reach out to the Ray developers on github.com/ray-project/ray/issues/
================================================================================
Out[4]: 
(False,
 {FailTuple(_remote [obj=<bound method ActorClass._remote of <__main__.ActorClass(Worker) object at 0x7f3134ecc1f0>>, parent=<__main__.ActorClass(Worker) object at 0x7f3134ecc1f0>])})

Issue Severity

Medium: It is a significant difficulty but I can work around it.

auderson commented 7 months ago

cloudpickle.dumps raise PicklingError with Worker actor class, but if @ray.remote removed, the original class can be pickled by cloudpickle. image