wehub-resource-sync
125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
"""Implements the client side of the client/server pickling protocol.
|
|
|
|
These picklers are aware of the server internals and can find the
|
|
references held for the client within the server.
|
|
|
|
More discussion about the client/server pickling protocol can be found in:
|
|
|
|
ray/util/client/client_pickler.py
|
|
|
|
ServerPickler dumps ray objects from the server into the appropriate stubs.
|
|
ClientUnpickler loads stubs from the client and finds their associated handle
|
|
in the server instance.
|
|
"""
|
|
import io
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import ray
|
|
import ray.cloudpickle as cloudpickle
|
|
from ray._private.client_mode_hook import disable_client_hook
|
|
from ray.util.client.client_pickler import PickleStub
|
|
from ray.util.client.server.server_stubs import (
|
|
ClientReferenceActor,
|
|
ClientReferenceFunction,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.util.client.server.server import RayletServicer
|
|
|
|
import pickle # noqa: F401
|
|
|
|
|
|
class ServerPickler(cloudpickle.CloudPickler):
|
|
def __init__(self, client_id: str, server: "RayletServicer", *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.client_id = client_id
|
|
self.server = server
|
|
|
|
def persistent_id(self, obj):
|
|
if isinstance(obj, ray.ObjectRef):
|
|
obj_id = obj.binary()
|
|
if obj_id not in self.server.object_refs[self.client_id]:
|
|
# We're passing back a reference, probably inside a reference.
|
|
# Let's hold onto it.
|
|
self.server.object_refs[self.client_id][obj_id] = obj
|
|
return PickleStub(
|
|
type="Object",
|
|
client_id=self.client_id,
|
|
ref_id=obj_id,
|
|
name=None,
|
|
baseline_options=None,
|
|
)
|
|
elif isinstance(obj, ray.actor.ActorHandle):
|
|
actor_id = obj._actor_id.binary()
|
|
if actor_id not in self.server.actor_refs:
|
|
# We're passing back a handle, probably inside a reference.
|
|
self.server.actor_refs[actor_id] = obj
|
|
if actor_id not in self.server.actor_owners[self.client_id]:
|
|
self.server.actor_owners[self.client_id].add(actor_id)
|
|
return PickleStub(
|
|
type="Actor",
|
|
client_id=self.client_id,
|
|
ref_id=obj._actor_id.binary(),
|
|
name=None,
|
|
baseline_options=None,
|
|
)
|
|
return None
|
|
|
|
|
|
class ClientUnpickler(pickle.Unpickler):
|
|
def __init__(self, server, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.server = server
|
|
|
|
def persistent_load(self, pid):
|
|
assert isinstance(pid, PickleStub)
|
|
if pid.type == "Ray":
|
|
return ray
|
|
elif pid.type == "Object":
|
|
return self.server.object_refs[pid.client_id][pid.ref_id]
|
|
elif pid.type == "Actor":
|
|
return self.server.actor_refs[pid.ref_id]
|
|
elif pid.type == "RemoteFuncSelfReference":
|
|
return ClientReferenceFunction(pid.client_id, pid.ref_id)
|
|
elif pid.type == "RemoteFunc":
|
|
return self.server.lookup_or_register_func(
|
|
pid.ref_id, pid.client_id, pid.baseline_options
|
|
)
|
|
elif pid.type == "RemoteActorSelfReference":
|
|
return ClientReferenceActor(pid.client_id, pid.ref_id)
|
|
elif pid.type == "RemoteActor":
|
|
return self.server.lookup_or_register_actor(
|
|
pid.ref_id, pid.client_id, pid.baseline_options
|
|
)
|
|
elif pid.type == "RemoteMethod":
|
|
actor = self.server.actor_refs[pid.ref_id]
|
|
return getattr(actor, pid.name)
|
|
else:
|
|
raise NotImplementedError("Uncovered client data type")
|
|
|
|
|
|
def dumps_from_server(
|
|
obj: Any, client_id: str, server_instance: "RayletServicer", protocol=None
|
|
) -> bytes:
|
|
with io.BytesIO() as file:
|
|
sp = ServerPickler(client_id, server_instance, file, protocol=protocol)
|
|
sp.dump(obj)
|
|
return file.getvalue()
|
|
|
|
|
|
def loads_from_client(
|
|
data: bytes,
|
|
server_instance: "RayletServicer",
|
|
*,
|
|
fix_imports=True,
|
|
encoding="ASCII",
|
|
errors="strict"
|
|
) -> Any:
|
|
with disable_client_hook():
|
|
if isinstance(data, str):
|
|
raise TypeError("Can't load pickle from unicode string")
|
|
file = io.BytesIO(data)
|
|
return ClientUnpickler(
|
|
server_instance, file, fix_imports=fix_imports, encoding=encoding
|
|
).load()
|