def options( self, args=None, kwargs=None, **task_options, ): """Configures and overrides the task invocation parameters. The arguments are the same as those that can be passed to :obj:`ray.remote`. Overriding `max_calls` is not supported. Examples: .. code-block:: python @ray.remote(num_gpus=1, max_calls=1, num_returns=2) def f(): return 1, 2 # Task f will require 2 gpus instead of 1. g = f.options(num_gpus=2) """ func_cls = self # override original options default_options = self._default_options.copy() # max_calls could not be used in ".options()", we should remove it before # merging options from '@ray.remote'. default_options.pop("max_calls", None) updated_options = {**default_options, **task_options} ray_option_utils.validate_task_options(updated_options, in_options=True) # only update runtime_env when ".options()" specifies new runtime_env if "runtime_env" in task_options: updated_options["runtime_env"] = parse_runtime_env( updated_options["runtime_env"]) class FuncWrapper: def remote(self, *args, **kwargs): return func_cls._remote(args=args, kwargs=kwargs, **updated_options) def bind(self, *args, **kwargs): """ **Experimental** For ray DAG building. Implementation and interface subject to changes. """ from ray.experimental.dag.function_node import FunctionNode return FunctionNode(func_cls._function, args, kwargs, updated_options) return FuncWrapper()
def __init__( self, language, function, function_descriptor, task_options, ): if inspect.iscoroutinefunction(function): raise ValueError( "'async def' should not be used for remote tasks. You can wrap the " "async function with `asyncio.get_event_loop.run_until(f())`. " "See more at https://docs.ray.io/en/latest/ray-core/async_api.html#asyncio-for-remote-tasks" # noqa ) self._default_options = task_options # TODO(suquark): This is a workaround for class attributes of options. # They are being used in some other places, mostly tests. Need cleanup later. # E.g., actors uses "__ray_metadata__" to collect options, we can so something # similar for remote functions. for k, v in ray_option_utils.task_options.items(): setattr(self, "_" + k, task_options.get(k, v.default_value)) self._runtime_env = parse_runtime_env(self._runtime_env) if "runtime_env" in self._default_options: self._default_options["runtime_env"] = self._runtime_env self._language = language self._function = _inject_tracing_into_function(function) self._function_name = function.__module__ + "." + function.__name__ self._function_descriptor = function_descriptor self._is_cross_language = language != Language.PYTHON self._decorator = getattr(function, "__ray_invocation_decorator__", None) self._function_signature = ray._private.signature.extract_signature( self._function) self._last_export_session_and_job = None self._uuid = uuid.uuid4() # Override task.remote's signature and docstring @wraps(function) def _remote_proxy(*args, **kwargs): return self._remote(args=args, kwargs=kwargs, **self._default_options) self.remote = _remote_proxy
def _remote(self, args=None, kwargs=None, **task_options): """Submit the remote function for execution.""" # We pop the "max_calls" coming from "@ray.remote" here. We no longer need # it in "_remote()". task_options.pop("max_calls", None) if client_mode_should_convert(auto_init=True): return client_mode_convert_function(self, args, kwargs, **task_options) worker = ray.worker.global_worker worker.check_connected() # If this function was not exported in this session and job, we need to # export this function again, because the current GCS doesn't have it. if (not self._is_cross_language and self._last_export_session_and_job != worker.current_session_and_job): self._function_descriptor = PythonFunctionDescriptor.from_function( self._function, self._uuid) # There is an interesting question here. If the remote function is # used by a subsequent driver (in the same script), should the # second driver pickle the function again? If yes, then the remote # function definition can differ in the second driver (e.g., if # variables in its closure have changed). We probably want the # behavior of the remote function in the second driver to be # independent of whether or not the function was invoked by the # first driver. This is an argument for repickling the function, # which we do here. try: self._pickled_function = pickle.dumps(self._function) except TypeError as e: msg = ( "Could not serialize the function " f"{self._function_descriptor.repr}. Check " "https://docs.ray.io/en/master/ray-core/objects/serialization.html#troubleshooting " # noqa "for more information.") raise TypeError(msg) from e self._last_export_session_and_job = worker.current_session_and_job worker.function_actor_manager.export(self) kwargs = {} if kwargs is None else kwargs args = [] if args is None else args # fill task required options for k, v in ray_option_utils.task_options.items(): task_options[k] = task_options.get(k, v.default_value) # "max_calls" already takes effects and should not apply again. # Remove the default value here. task_options.pop("max_calls", None) # TODO(suquark): cleanup these fields name = task_options["name"] runtime_env = parse_runtime_env(task_options["runtime_env"]) placement_group = task_options["placement_group"] placement_group_bundle_index = task_options[ "placement_group_bundle_index"] placement_group_capture_child_tasks = task_options[ "placement_group_capture_child_tasks"] scheduling_strategy = task_options["scheduling_strategy"] num_returns = task_options["num_returns"] max_retries = task_options["max_retries"] retry_exceptions = task_options["retry_exceptions"] resources = ray._private.utils.resources_from_ray_options(task_options) if scheduling_strategy is None or isinstance( scheduling_strategy, PlacementGroupSchedulingStrategy): if isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy): placement_group = scheduling_strategy.placement_group placement_group_bundle_index = ( scheduling_strategy.placement_group_bundle_index) placement_group_capture_child_tasks = ( scheduling_strategy.placement_group_capture_child_tasks) if placement_group_capture_child_tasks is None: placement_group_capture_child_tasks = ( worker.should_capture_child_tasks_in_placement_group) placement_group = configure_placement_group_based_on_context( placement_group_capture_child_tasks, placement_group_bundle_index, resources, {}, # no placement_resources for tasks self._function_descriptor.function_name, placement_group=placement_group, ) if not placement_group.is_empty: scheduling_strategy = PlacementGroupSchedulingStrategy( placement_group, placement_group_bundle_index, placement_group_capture_child_tasks, ) else: scheduling_strategy = "DEFAULT" serialized_runtime_env_info = None if runtime_env is not None: serialized_runtime_env_info = get_runtime_env_info( runtime_env, is_job_runtime_env=False, serialize=True, ) def invocation(args, kwargs): if self._is_cross_language: list_args = cross_language.format_args(worker, args, kwargs) elif not args and not kwargs and not self._function_signature: list_args = [] else: list_args = ray._private.signature.flatten_args( self._function_signature, args, kwargs) if worker.mode == ray.worker.LOCAL_MODE: assert ( not self._is_cross_language ), "Cross language remote function cannot be executed locally." object_refs = worker.core_worker.submit_task( self._language, self._function_descriptor, list_args, name if name is not None else "", num_returns, resources, max_retries, retry_exceptions, scheduling_strategy, worker.debugger_breakpoint, serialized_runtime_env_info or "{}", ) # Reset worker's debug context from the last "remote" command # (which applies only to this .remote call). worker.debugger_breakpoint = b"" if len(object_refs) == 1: return object_refs[0] elif len(object_refs) > 1: return object_refs if self._decorator is not None: invocation = self._decorator(invocation) return invocation(args, kwargs)
def __init__( self, language, function, function_descriptor, num_cpus, num_gpus, memory, object_store_memory, resources, accelerator_type, num_returns, max_calls, max_retries, retry_exceptions, runtime_env, placement_group, scheduling_strategy: SchedulingStrategyT, ): if inspect.iscoroutinefunction(function): raise ValueError("'async def' should not be used for remote " "tasks. You can wrap the async function with " "`asyncio.get_event_loop.run_until(f())`. " "See more at docs.ray.io/async_api.html") self._language = language self._function = _inject_tracing_into_function(function) self._function_name = function.__module__ + "." + function.__name__ self._function_descriptor = function_descriptor self._is_cross_language = language != Language.PYTHON self._num_cpus = DEFAULT_REMOTE_FUNCTION_CPUS if num_cpus is None else num_cpus self._num_gpus = num_gpus self._memory = memory if object_store_memory is not None: raise NotImplementedError( "setting object_store_memory is not implemented for tasks") self._object_store_memory = None self._resources = resources self._accelerator_type = accelerator_type self._num_returns = (DEFAULT_REMOTE_FUNCTION_NUM_RETURN_VALS if num_returns is None else num_returns) self._max_calls = (DEFAULT_REMOTE_FUNCTION_MAX_CALLS if max_calls is None else max_calls) self._max_retries = (DEFAULT_REMOTE_FUNCTION_NUM_TASK_RETRIES if max_retries is None else max_retries) self._retry_exceptions = (DEFAULT_REMOTE_FUNCTION_RETRY_EXCEPTIONS if retry_exceptions is None else retry_exceptions) self._runtime_env = parse_runtime_env(runtime_env) self._placement_group = placement_group self._decorator = getattr(function, "__ray_invocation_decorator__", None) self._function_signature = ray._private.signature.extract_signature( self._function) self._scheduling_strategy = scheduling_strategy self._last_export_session_and_job = None self._uuid = uuid.uuid4() # Override task.remote's signature and docstring @wraps(function) def _remote_proxy(*args, **kwargs): return self._remote(args=args, kwargs=kwargs) self.remote = _remote_proxy
def options( self, args=None, kwargs=None, num_returns=None, num_cpus=None, num_gpus=None, memory=None, object_store_memory=None, accelerator_type=None, resources=None, max_retries=None, retry_exceptions=None, placement_group="default", placement_group_bundle_index=-1, placement_group_capture_child_tasks=None, runtime_env=None, name="", scheduling_strategy: SchedulingStrategyT = None, ): """Configures and overrides the task invocation parameters. The arguments are the same as those that can be passed to :obj:`ray.remote`. Overriding `max_calls` is not supported. Examples: .. code-block:: python @ray.remote(num_gpus=1, max_calls=1, num_returns=2) def f(): return 1, 2 # Task f will require 2 gpus instead of 1. g = f.options(num_gpus=2) """ func_cls = self new_runtime_env = parse_runtime_env(runtime_env) options = dict( num_returns=num_returns, num_cpus=num_cpus, num_gpus=num_gpus, memory=memory, object_store_memory=object_store_memory, accelerator_type=accelerator_type, resources=resources, max_retries=max_retries, retry_exceptions=retry_exceptions, placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_index, placement_group_capture_child_tasks=( placement_group_capture_child_tasks), runtime_env=new_runtime_env, name=name, scheduling_strategy=scheduling_strategy, ) class FuncWrapper: def remote(self, *args, **kwargs): return func_cls._remote(args=args, kwargs=kwargs, **options) def bind(self, *args, **kwargs): """ **Experimental** For ray DAG building. Implementation and interface subject to changes. """ from ray.experimental.dag.function_node import FunctionNode return FunctionNode( func_cls._function, args, kwargs, options, ) return FuncWrapper()