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()