def register_callback( self, ref: ClientObjectRef, callback: Callable[[ray_client_pb2.DataResponse], None], ) -> None: req = ray_client_pb2.GetRequest(ids=[ref.id], asynchronous=True) self.data_client.RegisterGetCallback(req, callback)
def _get(self, ref: List[ClientObjectRef], timeout: float): req = ray_client_pb2.GetRequest(ids=[r.id for r in ref], timeout=timeout) data = bytearray() try: resp = self._get_object_iterator(req, metadata=self.metadata) for chunk in resp: if not chunk.valid: try: err = cloudpickle.loads(chunk.error) except (pickle.UnpicklingError, TypeError): logger.exception("Failed to deserialize {}".format( chunk.error)) raise raise err if chunk.total_size > OBJECT_TRANSFER_WARNING_SIZE and log_once( "client_object_transfer_size_warning"): size_gb = chunk.total_size / 2**30 warnings.warn( "Ray Client is attempting to retrieve a " f"{size_gb:.2f} GiB object over the network, which may " "be slow. Consider serializing the object to a file " "and using S3 or rsync instead.", UserWarning, stacklevel=5, ) data.extend(chunk.data) except grpc.RpcError as e: raise decode_exception(e) return loads_from_server(data)
def _get(self, id: bytes): req = ray_client_pb2.GetRequest(id=id) data = self.server.GetObject(req) if not data.valid: raise Exception( "Client GetObject returned invalid data: id invalid?") return cloudpickle.loads(data.data)
def _get(self, handle: bytes, timeout: float): req = ray_client_pb2.GetRequest(handle=handle, timeout=timeout) try: data = self.server.GetObject(req, metadata=self.metadata) except grpc.RpcError as e: raise decode_exception(e.details()) if not data.valid: raise TaskCancelledError(handle) return cloudpickle.loads(data.data)
def _get(self, ref: ClientObjectRef, timeout: float): req = ray_client_pb2.GetRequest(id=ref.id, timeout=timeout) try: data = self.data_client.GetObject(req) except grpc.RpcError as e: raise e.details() if not data.valid: raise cloudpickle.loads(data.error) return loads_from_server(data.data)
def _get(self, ref: ClientObjectRef, timeout: float): req = ray_client_pb2.GetRequest(id=ref.id, timeout=timeout) try: data = self.data_client.GetObject(req) except grpc.RpcError as e: raise decode_exception(e.details()) if not data.valid: try: err = cloudpickle.loads(data.error) except pickle.UnpicklingError: logger.exception("Failed to deserialize {}".format(data.error)) raise raise err return loads_from_server(data.data)
def _get(self, ref: List[ClientObjectRef], timeout: float): req = ray_client_pb2.GetRequest(ids=[r.id for r in ref], timeout=timeout) try: resp = self._call_stub("GetObject", req, metadata=self.metadata) except grpc.RpcError as e: raise decode_exception(e) if not resp.valid: try: err = cloudpickle.loads(resp.error) except (pickle.UnpicklingError, TypeError): logger.exception("Failed to deserialize {}".format(resp.error)) raise raise err return loads_from_server(resp.data)
def _get(self, id: bytes): req = ray_client_pb2.GetRequest(id=id) data = self.server.GetObject(req) return cloudpickle.loads(data.data)