def wait( self, object_refs: List[ClientObjectRef], *, num_returns: int = 1, timeout: float = None ) -> (List[ClientObjectRef], List[ClientObjectRef]): assert isinstance(object_refs, list) for ref in object_refs: assert isinstance(ref, ClientObjectRef) data = { "object_refs": [cloudpickle.dumps(object_ref) 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) if not resp.valid: # TODO(ameer): improve error/exceptions messages. raise Exception("Client Wait request failed. Reference invalid?") client_ready_object_ids = [ ClientObjectRef(id) for id in resp.ready_object_ids ] client_remaining_object_ids = [ ClientObjectRef(id) for id in resp.remaining_object_ids ] return (client_ready_object_ids, client_remaining_object_ids)
def wait( self, object_refs: List[ClientObjectRef], *, num_returns: int = 1, timeout: float = None, fetch_local: bool = True ) -> Tuple[List[ClientObjectRef], List[ClientObjectRef]]: if not isinstance(object_refs, list): raise TypeError("wait() expected a list of ClientObjectRef, " f"got {type(object_refs)}") for ref in object_refs: if not isinstance(ref, ClientObjectRef): raise TypeError("wait() expected a list of ClientObjectRef, " f"got list containing {type(ref)}") data = { "object_ids": [object_ref.id for object_ref in object_refs], "num_returns": num_returns, "timeout": timeout if timeout else -1, "client_id": self._client_id, } 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(ref) for ref in resp.ready_object_ids ] client_remaining_object_ids = [ ClientObjectRef(ref) for ref in resp.remaining_object_ids ] return (client_ready_object_ids, client_remaining_object_ids)
def test_wait(ray_start_regular_shared): with ray_start_client_server() as ray: objectref = ray.put("hello world") ready, remaining = ray.wait([objectref]) assert remaining == [] retval = ray.get(ready[0]) assert retval == "hello world" objectref2 = ray.put(5) ready, remaining = ray.wait([objectref, objectref2]) assert (ready, remaining) == ([objectref], [objectref2]) or \ (ready, remaining) == ([objectref2], [objectref]) ready_retval = ray.get(ready[0]) remaining_retval = ray.get(remaining[0]) assert (ready_retval, remaining_retval) == ("hello world", 5) \ or (ready_retval, remaining_retval) == (5, "hello world") with pytest.raises(Exception): # Reference not in the object store. ray.wait([ClientObjectRef("blabla")]) with pytest.raises(AssertionError): ray.wait("blabla") with pytest.raises(AssertionError): ray.wait(ClientObjectRef("blabla")) with pytest.raises(AssertionError): ray.wait(["blabla"])
def test_wait(ray_start_regular_shared): server = ray_client_server.serve("localhost:50051", test_mode=True) ray.connect("localhost:50051") objectref = ray.put("hello world") ready, remaining = ray.wait([objectref]) assert remaining == [] retval = ray.get(ready[0]) assert retval == "hello world" objectref2 = ray.put(5) ready, remaining = ray.wait([objectref, objectref2]) assert (ready, remaining) == ([objectref], [objectref2]) or \ (ready, remaining) == ([objectref2], [objectref]) ready_retval = ray.get(ready[0]) remaining_retval = ray.get(remaining[0]) assert (ready_retval, remaining_retval) == ("hello world", 5) \ or (ready_retval, remaining_retval) == (5, "hello world") with pytest.raises(Exception): # Reference not in the object store. ray.wait([ClientObjectRef("blabla")]) with pytest.raises(AssertionError): ray.wait("blabla") with pytest.raises(AssertionError): ray.wait(ClientObjectRef("blabla")) with pytest.raises(AssertionError): ray.wait(["blabla"]) ray.disconnect() server.stop(0)
def persistent_load(self, pid): assert isinstance(pid, PickleStub) if pid.type == "Object": return ClientObjectRef(id=pid.ref_id) elif pid.type == "Actor": return ClientActorHandle(ClientActorRef(id=pid.ref_id)) else: raise NotImplementedError("Being passed back an unknown stub")
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 call_remote(self, func, *args, **kwargs): if not isinstance(func, ClientRemoteFunc): raise TypeError("Client not passing a ClientRemoteFunc stub") func_ref = self._put(func) task = ray_client_pb2.ClientTask() task.name = func._name task.payload_id = func_ref.id for arg in args: pb_arg = convert_to_arg(arg) task.args.append(pb_arg) ticket = self.server.Schedule(task) return ClientObjectRef(ticket.return_id)
def _put(self, val): if isinstance(val, ClientObjectRef): raise TypeError( "Calling 'put' on an ObjectRef is not allowed " "(similarly, returning an ObjectRef from a remote " "function is not allowed). If you really want to " "do this, you can wrap the ObjectRef in a list and " "call 'put' on it (or return it).") data = dumps_from_client(val, self._client_id) req = ray_client_pb2.PutRequest(data=data) resp = self.data_client.PutObject(req) return ClientObjectRef(resp.id)
def call_remote(self, instance, kind, *args, **kwargs): ticket = None if kind == ray_client_pb2.ClientTask.FUNCTION: ticket = self._put_and_schedule(instance, kind, *args, **kwargs) elif kind == ray_client_pb2.ClientTask.ACTOR: ticket = self._put_and_schedule(instance, kind, *args, **kwargs) return ClientActorRef(ticket.return_id) elif kind == ray_client_pb2.ClientTask.METHOD: ticket = self._call_method(instance, *args, **kwargs) if ticket is None: raise Exception("Couldn't call_remote on %s for type %s" % (instance, kind)) return ClientObjectRef(ticket.return_id)
def _put(self, val): data = cloudpickle.dumps(val) req = ray_client_pb2.PutRequest(data=data) resp = self.server.PutObject(req) return ClientObjectRef(resp.id)
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)
def _put(self, val): data = dumps_from_client(val, self._client_id) req = ray_client_pb2.PutRequest(data=data) resp = self.data_client.PutObject(req) return ClientObjectRef(resp.id)
def call_remote(self, instance: ClientStub, *args, **kwargs): task = instance._prepare_client_task() ticket = self.server.Schedule(task, prepared_args=args) return ClientObjectRef(ticket.return_id)
def _put(self, val: Any): resp = self.server._put_and_retain_obj(val) return ClientObjectRef(resp.id)