def send_get_response(result: Any) -> None:
"""Pushes GetResponses to the main DataPath loop to send
to the client. This is called when the object is ready
on the server side."""
try:
serialized = dumps_from_server(result, client_id, self)
total_size = len(serialized)
assert total_size > 0, "Serialized object cannot be zero bytes"
total_chunks = math.ceil(total_size /
OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(request.start_chunk_id,
total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) *
OBJECT_TRANSFER_CHUNK_SIZE)
get_resp = ray_client_pb2.GetResponse(
valid=True,
data=serialized[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)
chunk_resp = ray_client_pb2.DataResponse(
get=get_resp, req_id=req_id)
result_queue.put(chunk_resp)
except Exception as exc:
get_resp = ray_client_pb2.GetResponse(
valid=False, error=cloudpickle.dumps(exc))
resp = ray_client_pb2.DataResponse(get=get_resp,
req_id=req_id)
result_queue.put(resp)
def GetObject(self, request, context=None):
if request.id not in self.object_refs:
return ray_client_pb2.GetResponse(valid=False)
objectref = self.object_refs[request.id]
logger.info("get: %s" % objectref)
item = ray.get(objectref)
item_ser = cloudpickle.dumps(item)
return ray_client_pb2.GetResponse(valid=True, data=item_ser)
def GetObject(self, request, context=None):
objectref = self.realref[request.id]
print("get: %s" % objectref)
item = ray.get(objectref)
if item is None:
return ray_client_pb2.GetResponse(valid=False)
data = cloudpickle.loads(item)
return ray_client_pb2.GetResponse(valid=True, data=data)
def GetObject(self, request, context=None):
request_ref = cloudpickle.loads(request.handle)
if request_ref.binary() not in self.object_refs:
return ray_client_pb2.GetResponse(valid=False)
objectref = self.object_refs[request_ref.binary()]
logger.info("get: %s" % objectref)
try:
item = ray.get(objectref, timeout=request.timeout)
except Exception as e:
return_exception_in_context(e, context)
item_ser = cloudpickle.dumps(item)
return ray_client_pb2.GetResponse(valid=True, data=item_ser)
def _async_get_object(
self,
request: ray_client_pb2.GetRequest,
client_id: str,
req_id: int,
result_queue: queue.Queue,
context=None,
) -> Optional[ray_client_pb2.GetResponse]:
"""Attempts to schedule a callback to push the GetResponse to the
main loop when the desired object is ready. If there is some failure
in scheduling, a GetResponse will be immediately returned.
"""
if len(request.ids) != 1:
raise ValueError(
"Async get() must have exactly 1 Object ID. " f"Actual: {request}"
)
rid = request.ids[0]
ref = self.object_refs[client_id].get(rid, None)
if not ref:
return ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
f"ClientObjectRef with id {rid} not found for "
f"client {client_id}"
)
),
)
try:
logger.debug("async get: %s" % ref)
with disable_client_hook():
def send_get_response(result: Any) -> None:
"""Pushes a GetResponse to the main DataPath loop to send
to the client. This is called when the object is ready
on the server side."""
try:
serialized = dumps_from_server(result, client_id, self)
get_resp = ray_client_pb2.GetResponse(
valid=True, data=serialized
)
except Exception as exc:
get_resp = ray_client_pb2.GetResponse(
valid=False, error=cloudpickle.dumps(exc)
)
resp = ray_client_pb2.DataResponse(get=get_resp, req_id=req_id)
result_queue.put(resp)
ref._on_completed(send_get_response)
return None
except Exception as e:
return ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
def _get_object(self, request, client_id: str, context=None):
if request.id not in self.object_refs[client_id]:
return ray_client_pb2.GetResponse(valid=False)
objectref = self.object_refs[client_id][request.id]
logger.debug("get: %s" % objectref)
try:
item = ray.get(objectref, timeout=request.timeout)
except Exception as e:
return ray_client_pb2.GetResponse(valid=False,
error=cloudpickle.dumps(e))
item_ser = dumps_from_server(item, client_id, self)
return ray_client_pb2.GetResponse(valid=True, data=item_ser)
def send_get_response(result: Any) -> None:
"""Pushes a GetResponse to the main DataPath loop to send
to the client. This is called when the object is ready
on the server side."""
try:
serialized = dumps_from_server(result, client_id, self)
get_resp = ray_client_pb2.GetResponse(valid=True,
data=serialized)
except Exception as exc:
get_resp = ray_client_pb2.GetResponse(
valid=False, error=cloudpickle.dumps(exc))
resp = ray_client_pb2.DataResponse(get=get_resp,
req_id=req_id)
result_queue.put(resp)
def _async_get_object(
self,
request: ray_client_pb2.GetRequest,
client_id: str,
req_id: int,
result_queue: queue.Queue,
context=None) -> Optional[ray_client_pb2.GetResponse]:
"""Attempts to schedule a callback to push the GetResponse to the
main loop when the desired object is ready. If there is some failure
in scheduling, a GetResponse will be immediately returned.
"""
refs = []
for rid in request.ids:
ref = self.object_refs[client_id].get(rid, None)
if ref:
refs.append(ref)
else:
return ray_client_pb2.GetResponse(valid=False)
try:
logger.debug("async get: %s" % refs)
with disable_client_hook():
def send_get_response(result: Any) -> None:
"""Pushes a GetResponse to the main DataPath loop to send
to the client. This is called when the object is ready
on the server side."""
try:
serialized = dumps_from_server(result, client_id, self)
get_resp = ray_client_pb2.GetResponse(valid=True,
data=serialized)
except Exception as e:
get_resp = ray_client_pb2.GetResponse(
valid=False, error=cloudpickle.dumps(e))
resp = ray_client_pb2.DataResponse(get=get_resp,
req_id=req_id)
resp.req_id = req_id
result_queue.put(resp)
for ref in refs:
ref._on_completed(send_get_response)
return None
except Exception as e:
return ray_client_pb2.GetResponse(valid=False,
error=cloudpickle.dumps(e))
def GetObject(self, request: ray_client_pb2.GetRequest, context):
metadata = {k: v for k, v in context.invocation_metadata()}
client_id = metadata.get("client_id")
if client_id is None:
return ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
"client_id is not specified in request metadata")))
return self._get_object(request, client_id)
def _get_object(self,
request: ray_client_pb2.GetRequest,
client_id: str,
context=None):
objectrefs = []
for rid in request.ids:
ref = self.object_refs[client_id].get(rid, None)
if ref:
objectrefs.append(ref)
else:
return ray_client_pb2.GetResponse(valid=False)
try:
logger.debug("get: %s" % objectrefs)
with disable_client_hook():
items = ray.get(objectrefs, timeout=request.timeout)
except Exception as e:
return ray_client_pb2.GetResponse(valid=False,
error=cloudpickle.dumps(e))
items_ser = dumps_from_server(items, client_id, self)
return ray_client_pb2.GetResponse(valid=True, data=items_ser)
def _get_object(self, request: ray_client_pb2.GetRequest, client_id: str):
objectrefs = []
for rid in request.ids:
ref = self.object_refs[client_id].get(rid, None)
if ref:
objectrefs.append(ref)
else:
return ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
f"ClientObjectRef {rid} is not found for client "
f"{client_id}")))
try:
logger.debug("get: %s" % objectrefs)
with disable_client_hook():
items = ray.get(objectrefs, timeout=request.timeout)
except Exception as e:
return ray_client_pb2.GetResponse(valid=False,
error=cloudpickle.dumps(e))
serialized = dumps_from_server(items, client_id, self)
return ray_client_pb2.GetResponse(valid=True, data=serialized)
def _get_object(self, request: ray_client_pb2.GetRequest, client_id: str):
objectrefs = []
for rid in request.ids:
ref = self.object_refs[client_id].get(rid, None)
if ref:
objectrefs.append(ref)
else:
yield ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
f"ClientObjectRef {rid} is not found for client "
f"{client_id}"
)
),
)
return
try:
logger.debug("get: %s" % objectrefs)
with disable_client_hook():
items = ray.get(objectrefs, timeout=request.timeout)
except Exception as e:
yield ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
return
serialized = dumps_from_server(items, client_id, self)
total_size = len(serialized)
assert total_size > 0, "Serialized object cannot be zero bytes"
total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(request.start_chunk_id, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
yield ray_client_pb2.GetResponse(
valid=True,
data=serialized[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)