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 is not None) 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(b"blabla")]) with pytest.raises(TypeError): ray.wait("blabla") with pytest.raises(TypeError): ray.wait(ClientObjectRef("blabla")) with pytest.raises(TypeError): ray.wait(["blabla"])
def test_client_object_ref_basics(ray_start_regular): with ray_start_client_server_pair() as pair: ray, server = pair ref = ray.put("Hello World") # Make sure ClientObjectRef is a subclass of ObjectRef assert isinstance(ref, ClientObjectRef) assert isinstance(ref, ObjectRef) # Invalid ref format. with pytest.raises(Exception): ClientObjectRef(b"\0") obj_id = b"\0" * 28 fut = Future() fut.set_result(obj_id) server_ref = ObjectRef(obj_id) for client_ref in [ClientObjectRef(obj_id), ClientObjectRef(fut)]: client_members = set(client_ref.__dir__()) server_members = set(server_ref.__dir__()) assert client_members.difference(server_members) == {"id"} assert server_members.difference(client_members) == set() # Test __eq__() assert client_ref == ClientObjectRef(obj_id) assert client_ref != ref assert client_ref != server_ref # Test other methods assert client_ref.__repr__() == f"ClientObjectRef({obj_id.hex()})" assert client_ref.binary() == obj_id assert client_ref.hex() == obj_id.hex() assert not client_ref.is_nil() assert client_ref.task_id() == server_ref.task_id() assert client_ref.job_id() == server_ref.job_id()
def persistent_load(self, pid): assert isinstance(pid, PickleStub) if pid.type == "Object": return ClientObjectRef(pid.ref_id) elif pid.type == "Actor": return ClientActorHandle(ClientActorRef(pid.ref_id)) else: raise NotImplementedError("Being passed back an unknown stub")
def _put_pickled(self, data, client_ref_id: bytes): req = ray_client_pb2.PutRequest(data=data) if client_ref_id is not None: req.client_ref_id = client_ref_id resp = self.data_client.PutObject(req) if not resp.valid: try: raise cloudpickle.loads(resp.error) except (pickle.UnpicklingError, TypeError): logger.exception("Failed to deserialize {}".format(resp.error)) raise return ClientObjectRef(resp.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 test_put_get(ray_start_regular_shared): with ray_start_client_server() as ray: objectref = ray.put("hello world") print(objectref) retval = ray.get(objectref) assert retval == "hello world" # Make sure ray.put(1) == 1 is False and does not raise an exception. objectref = ray.put(1) assert not objectref == 1 # Make sure it returns True when necessary as well. assert objectref == ClientObjectRef(objectref.id)
def test_client_object_ref_basics(ray_start_regular): with ray_start_client_server_pair() as pair: ray, server = pair ref = ray.put("Hello World") # Make sure ClientObjectRef is a subclass of ObjectRef assert isinstance(ref, ClientObjectRef) assert isinstance(ref, ObjectRef) # Invalid ref format. with pytest.raises(Exception): ClientObjectRef(b"\0") # Test __eq__() id = b"\0" * 28 assert ClientObjectRef(id) == ClientObjectRef(id) assert ClientObjectRef(id) != ref assert ClientObjectRef(id) != ObjectRef(id) assert ClientObjectRef(id).__repr__() == f"ClientObjectRef({id.hex()})" assert ClientObjectRef(id).binary() == id assert ClientObjectRef(id).hex() == id.hex() assert not ClientObjectRef(id).is_nil()
def _put(self, val, *, client_ref_id: bytes = None): 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) if client_ref_id is not None: req.client_ref_id = client_ref_id resp = self.data_client.PutObject(req) if not resp.valid: try: raise cloudpickle.loads(resp.error) except pickle.UnpicklingError: logger.exception("Failed to deserialize {}".format(resp.error)) raise return ClientObjectRef(resp.id)