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)