def wait(
self,
object_refs: List[ClientObjectRef],
*,
num_returns: int = 1,
timeout: float = None
) -> Tuple[List[ClientObjectRef], List[ClientObjectRef]]:
assert isinstance(object_refs, list)
for ref in object_refs:
assert isinstance(ref, ClientObjectRef)
data = {
"object_handles":
[object_ref.handle for object_ref in object_refs],
"num_returns": num_returns,
"timeout": timeout if timeout else -1
}
req = ray_client_pb2.WaitRequest(**data)
resp = self.server.WaitObject(req, metadata=self.metadata)
if not resp.valid:
# TODO(ameer): improve error/exceptions messages.
raise Exception("Client Wait request failed. Reference invalid?")
client_ready_object_ids = [
ClientObjectRef.from_remote_ref(ref)
for ref in resp.ready_object_ids
]
client_remaining_object_ids = [
ClientObjectRef.from_remote_ref(ref)
for ref in resp.remaining_object_ids
]
return (client_ready_object_ids, client_remaining_object_ids)
def call_remote(self, instance, *args, **kwargs):
task = instance._prepare_client_task()
for arg in args:
pb_arg = convert_to_arg(arg)
task.args.append(pb_arg)
logging.debug("Scheduling %s" % task)
ticket = self.server.Schedule(task, metadata=self.metadata)
return ClientObjectRef.from_remote_ref(ticket.return_ref)
def _put(self, val):
data = cloudpickle.dumps(val)
req = ray_client_pb2.PutRequest(data=data)
resp = self.server.PutObject(req, metadata=self.metadata)
return ClientObjectRef.from_remote_ref(resp.ref)