def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse: object_refs = [] for rid in request.object_ids: if rid not in self.object_refs[request.client_id]: raise Exception( "Asking for a ref not associated with this client: %s" % str(rid)) object_refs.append(self.object_refs[request.client_id][rid]) num_returns = request.num_returns timeout = request.timeout try: with disable_client_hook(): ready_object_refs, remaining_object_refs = ray.wait( object_refs, num_returns=num_returns, timeout=timeout if timeout != -1 else None, ) except Exception as e: # TODO(ameer): improve exception messages. logger.error(f"Exception {e}") return ray_client_pb2.WaitResponse(valid=False) logger.debug("wait: %s %s" % (str(ready_object_refs), str(remaining_object_refs))) ready_object_ids = [ ready_object_ref.binary() for ready_object_ref in ready_object_refs ] remaining_object_ids = [ remaining_object_ref.binary() for remaining_object_ref in remaining_object_refs ] return ray_client_pb2.WaitResponse( valid=True, ready_object_ids=ready_object_ids, remaining_object_ids=remaining_object_ids, )
def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse: object_refs = [cloudpickle.loads(o) for o in request.object_handles] num_returns = request.num_returns timeout = request.timeout object_refs_ids = [] for object_ref in object_refs: if object_ref.binary() not in self.object_refs: return ray_client_pb2.WaitResponse(valid=False) object_refs_ids.append(self.object_refs[object_ref.binary()]) try: ready_object_refs, remaining_object_refs = ray.wait( object_refs_ids, num_returns=num_returns, timeout=timeout if timeout != -1 else None) except Exception: # TODO(ameer): improve exception messages. return ray_client_pb2.WaitResponse(valid=False) logger.info("wait: %s %s" % (str(ready_object_refs), str(remaining_object_refs))) ready_object_ids = [ make_remote_ref( id=ready_object_ref.binary(), handle=cloudpickle.dumps(ready_object_ref), ) for ready_object_ref in ready_object_refs ] remaining_object_ids = [ make_remote_ref( id=remaining_object_ref.binary(), handle=cloudpickle.dumps(remaining_object_ref), ) for remaining_object_ref in remaining_object_refs ] return ray_client_pb2.WaitResponse( valid=True, ready_object_ids=ready_object_ids, remaining_object_ids=remaining_object_ids)