def Schedule( self, task: ray_client_pb2.ClientTask, arglist: List[Any], kwargs: Dict[str, Any], context=None, ) -> ray_client_pb2.ClientTaskTicket: logger.debug( "schedule: %s %s" % (task.name, ray_client_pb2.ClientTask.RemoteExecType.Name( task.type))) try: with disable_client_hook(): if task.type == ray_client_pb2.ClientTask.FUNCTION: result = self._schedule_function(task, arglist, kwargs, context) elif task.type == ray_client_pb2.ClientTask.ACTOR: result = self._schedule_actor(task, arglist, kwargs, context) elif task.type == ray_client_pb2.ClientTask.METHOD: result = self._schedule_method(task, arglist, kwargs, context) elif task.type == ray_client_pb2.ClientTask.NAMED_ACTOR: result = self._schedule_named_actor(task, context) else: raise NotImplementedError( "Unimplemented Schedule task type: %s" % ray_client_pb2.ClientTask.RemoteExecType.Name( task.type)) result.valid = True return result except Exception as e: logger.debug("Caught schedule exception", exc_info=True) return ray_client_pb2.ClientTaskTicket(valid=False, error=cloudpickle.dumps(e))
def Schedule(self, task, context=None) -> ray_client_pb2.ClientTaskTicket: logger.debug( "schedule: %s %s" % (task.name, ray_client_pb2.ClientTask.RemoteExecType.Name( task.type))) with stash_api_for_tests(self._test_mode): try: if task.type == ray_client_pb2.ClientTask.FUNCTION: result = self._schedule_function(task, context) elif task.type == ray_client_pb2.ClientTask.ACTOR: result = self._schedule_actor(task, context) elif task.type == ray_client_pb2.ClientTask.METHOD: result = self._schedule_method(task, context) elif task.type == ray_client_pb2.ClientTask.NAMED_ACTOR: result = self._schedule_named_actor(task, context) else: raise NotImplementedError( "Unimplemented Schedule task type: %s" % ray_client_pb2.ClientTask.RemoteExecType.Name( task.type)) result.valid = True return result except Exception as e: logger.error(f"Caught schedule exception {e}") raise e return ray_client_pb2.ClientTaskTicket( valid=False, error=cloudpickle.dumps(e))
def _schedule_named_actor(self, task: ray_client_pb2.ClientTask, context=None) -> ray_client_pb2.ClientTaskTicket: assert len(task.payload_id) == 0 actor = ray.get_actor(task.name) self.actor_refs[actor._actor_id.binary()] = actor self.actor_owners[task.client_id].add(actor._actor_id.binary()) return ray_client_pb2.ClientTaskTicket( return_ids=[actor._actor_id.binary()])
def _schedule_named_actor( self, task: ray_client_pb2.ClientTask, context=None ) -> ray_client_pb2.ClientTaskTicket: assert len(task.payload_id) == 0 # Convert empty string back to None. actor = ray.get_actor(task.name, task.namespace or None) bin_actor_id = actor._actor_id.binary() self.actor_refs[bin_actor_id] = actor self.actor_owners[task.client_id].add(bin_actor_id) self.named_actors.add(bin_actor_id) return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()])
def _schedule_method(self, task: ray_client_pb2.ClientTask, context=None) -> ray_client_pb2.ClientTaskTicket: actor_handle = self.actor_refs.get(task.payload_id) if actor_handle is None: raise Exception( "Can't run an actor the server doesn't have a handle for") arglist = _convert_args(task.args) with stash_api_for_tests(self._test_mode): output = getattr(actor_handle, task.name).remote(*arglist) self.object_refs[output.binary()] = output return ray_client_pb2.ClientTaskTicket(return_id=output.binary())
def _schedule_function(self, task: ray_client_pb2.ClientTask, context=None) -> ray_client_pb2.ClientTaskTicket: remote_func = self.lookup_or_register_func( task.payload_id, task.client_id, decode_options(task.baseline_options)) arglist, kwargs = self._convert_args(task.args, task.kwargs) opts = decode_options(task.options) if opts is not None: remote_func = remote_func.options(**opts) with current_server(self): output = remote_func.remote(*arglist, **kwargs) ids = self.unify_and_track_outputs(output, task.client_id) return ray_client_pb2.ClientTaskTicket(return_ids=ids)
def _schedule_method(self, task: ray_client_pb2.ClientTask, context=None) -> ray_client_pb2.ClientTaskTicket: actor_handle = self.actor_refs.get(task.payload_id) if actor_handle is None: raise Exception( "Can't run an actor the server doesn't have a handle for") arglist, kwargs = self._convert_args(task.args, task.kwargs) method = getattr(actor_handle, task.name) opts = decode_options(task.options) if opts is not None: method = method.options(**opts) output = method.remote(*arglist, **kwargs) ids = self.unify_and_track_outputs(output, task.client_id) return ray_client_pb2.ClientTaskTicket(return_ids=ids)
def _schedule_function( self, task: ray_client_pb2.ClientTask, context=None, prepared_args=None) -> ray_client_pb2.ClientTaskTicket: remote_func = self.lookup_or_register_func(task.payload_id, task.client_id) arglist = self._convert_args(task.args, prepared_args) # Prepare call if we're in a test with current_func(remote_func): output = remote_func.remote(*arglist) if output.binary() in self.object_refs[task.client_id]: raise Exception("already found it") self.object_refs[task.client_id][output.binary()] = output return ray_client_pb2.ClientTaskTicket(return_id=output.binary())
def _schedule_actor(self, task: ray_client_pb2.ClientTask, context=None) -> ray_client_pb2.ClientTaskTicket: remote_class = self.lookup_or_register_actor( task.payload_id, task.client_id, decode_options(task.baseline_options)) arglist, kwargs = self._convert_args(task.args, task.kwargs) opts = decode_options(task.options) if opts is not None: remote_class = remote_class.options(**opts) with current_server(self): actor = remote_class.remote(*arglist, **kwargs) self.actor_refs[actor._actor_id.binary()] = actor self.actor_owners[task.client_id].add(actor._actor_id.binary()) return ray_client_pb2.ClientTaskTicket( return_ids=[actor._actor_id.binary()])
def Schedule(self, task, context=None): logger.info("schedule: %s" % task) if task.payload_id not in self.function_refs: funcref = self.object_refs[task.payload_id] func = ray.get(funcref) if not isinstance(func, ClientRemoteFunc): raise Exception("Attempting to schedule something that " "isn't a ClientRemoteFunc") ray_remote = ray.remote(func._func) func.set_remote_func(ray_remote) self.function_refs[task.payload_id] = func remote_func = self.function_refs[task.payload_id] arglist = _convert_args(task.args) output = remote_func.remote(*arglist) self.object_refs[output.binary()] = output return ray_client_pb2.ClientTaskTicket(return_id=output.binary())
def _schedule_function(self, task: ray_client_pb2.ClientTask, context=None) -> ray_client_pb2.ClientTaskTicket: if task.payload_id not in self.function_refs: funcref = self.object_refs[task.payload_id] func = ray.get(funcref) if not isinstance(func, ClientRemoteFunc): raise Exception("Attempting to schedule function that " "isn't a ClientRemoteFunc.") self.function_refs[task.payload_id] = func remote_func = self.function_refs[task.payload_id] arglist = _convert_args(task.args) # Prepare call if we're in a test with stash_api_for_tests(self._test_mode): output = remote_func.remote(*arglist) self.object_refs[output.binary()] = output return ray_client_pb2.ClientTaskTicket(return_id=output.binary())
def _schedule_actor(self, task: ray_client_pb2.ClientTask, context=None) -> ray_client_pb2.ClientTaskTicket: with stash_api_for_tests(self._test_mode): if task.payload_id not in self.registered_actor_classes: actor_class_ref = self.object_refs[task.payload_id] actor_class = ray.get(actor_class_ref) if not inspect.isclass(actor_class): raise Exception("Attempting to schedule actor that " "isn't a ClientActorClass.") reg_class = ray.remote(actor_class) self.registered_actor_classes[task.payload_id] = reg_class remote_class = self.registered_actor_classes[task.payload_id] arglist = _convert_args(task.args) actor = remote_class.remote(*arglist) actor_ref = actor._actor_id self.actor_refs[actor_ref.binary()] = actor return ray_client_pb2.ClientTaskTicket(return_id=actor_ref.binary())
def Schedule(self, task, context=None): logger.info("schedule: %s" % task) if task.payload_id not in self.function_refs: funcref = self.object_refs[task.payload_id] func = ray.get(funcref) if not isinstance(func, ClientRemoteFunc): raise Exception("Attempting to schedule something that " "isn't a ClientRemoteFunc") self.function_refs[task.payload_id] = func remote_func = self.function_refs[task.payload_id] arglist = _convert_args(task.args) # Prepare call if we're in a test api = None if self._test_mode: api = client_init.stash_api() output = remote_func.remote(*arglist) if self._test_mode: client_init.restore_api(api) self.object_refs[output.binary()] = output return ray_client_pb2.ClientTaskTicket(return_id=output.binary())
def _schedule_actor(self, task: ray_client_pb2.ClientTask, context=None, prepared_args=None) -> ray_client_pb2.ClientTaskTicket: if task.payload_id not in self.registered_actor_classes: actor_class_ref = \ self.object_refs[task.client_id][task.payload_id] actor_class = ray.get(actor_class_ref) if not inspect.isclass(actor_class): raise Exception("Attempting to schedule actor that " "isn't a class.") reg_class = ray.remote(actor_class) self.registered_actor_classes[task.payload_id] = reg_class remote_class = self.registered_actor_classes[task.payload_id] arglist = self._convert_args(task.args, prepared_args) actor = remote_class.remote(*arglist) self.actor_refs[actor._actor_id.binary()] = actor self.actor_owners[task.client_id].add(actor._actor_id.binary()) return ray_client_pb2.ClientTaskTicket( return_id=actor._actor_id.binary())
def _schedule_actor(self, task: ray_client_pb2.ClientTask, context=None, prepared_args=None) -> ray_client_pb2.ClientTaskTicket: with stash_api_for_tests(self._test_mode): payload_ref = cloudpickle.loads(task.payload_id) if payload_ref.binary() not in self.registered_actor_classes: actor_class_ref = self.object_refs[payload_ref.binary()] actor_class = ray.get(actor_class_ref) if not inspect.isclass(actor_class): raise Exception("Attempting to schedule actor that " "isn't a class.") reg_class = ray.remote(actor_class) self.registered_actor_classes[payload_ref.binary()] = reg_class remote_class = self.registered_actor_classes[payload_ref.binary()] arglist = _convert_args(task.args, prepared_args) actor = remote_class.remote(*arglist) actorhandle = cloudpickle.dumps(actor) self.actor_refs[actorhandle] = actor return ray_client_pb2.ClientTaskTicket( return_ref=make_remote_ref(actor._actor_id.binary(), actorhandle))
def _schedule_function( self, task: ray_client_pb2.ClientTask, context=None, prepared_args=None) -> ray_client_pb2.ClientTaskTicket: payload_ref = cloudpickle.loads(task.payload_id) if payload_ref.binary() not in self.function_refs: funcref = self.object_refs[payload_ref.binary()] func = ray.get(funcref) if not inspect.isfunction(func): raise Exception("Attempting to schedule function that " "isn't a function.") self.function_refs[payload_ref.binary()] = ray.remote(func) remote_func = self.function_refs[payload_ref.binary()] arglist = _convert_args(task.args, prepared_args) # Prepare call if we're in a test with stash_api_for_tests(self._test_mode): output = remote_func.remote(*arglist) if output.binary() in self.object_refs: raise Exception("already found it") self.object_refs[output.binary()] = output pickled_output = cloudpickle.dumps(output) return ray_client_pb2.ClientTaskTicket( return_ref=make_remote_ref(output.binary(), pickled_output))