chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,454 @@
|
||||
"""Manage, parse and validate options for Ray tasks, actors and actor methods."""
|
||||
import warnings
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Dict, Optional, Tuple, Union
|
||||
|
||||
import ray
|
||||
from ray._private import ray_constants
|
||||
from ray._private.label_utils import (
|
||||
validate_fallback_strategy,
|
||||
validate_label_selector,
|
||||
)
|
||||
from ray._private.utils import get_ray_doc_version
|
||||
from ray.util.placement_group import PlacementGroup
|
||||
from ray.util.scheduling_strategies import (
|
||||
NodeAffinitySchedulingStrategy,
|
||||
NodeLabelSchedulingStrategy,
|
||||
PlacementGroupSchedulingStrategy,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Option:
|
||||
# Type constraint of an option.
|
||||
type_constraint: Optional[Union[type, Tuple[type]]] = None
|
||||
# Value constraint of an option.
|
||||
# The callable should return None if there is no error.
|
||||
# Otherwise, return the error message.
|
||||
value_constraint: Optional[Callable[[Any], Optional[str]]] = None
|
||||
# Default value.
|
||||
default_value: Any = None
|
||||
|
||||
def validate(self, keyword: str, value: Any):
|
||||
"""Validate the option."""
|
||||
if self.type_constraint is not None:
|
||||
if not isinstance(value, self.type_constraint):
|
||||
raise TypeError(
|
||||
f"The type of keyword '{keyword}' must be {self.type_constraint}, "
|
||||
f"but received type {type(value)}"
|
||||
)
|
||||
if self.value_constraint is not None:
|
||||
possible_error_message = self.value_constraint(value)
|
||||
if possible_error_message:
|
||||
raise ValueError(possible_error_message)
|
||||
|
||||
|
||||
def _counting_option(name: str, infinite: bool = True, default_value: Any = None):
|
||||
"""This is used for positive and discrete options.
|
||||
|
||||
Args:
|
||||
name: The name of the option keyword.
|
||||
infinite: If True, user could use -1 to represent infinity.
|
||||
default_value: The default value for this option.
|
||||
|
||||
Returns:
|
||||
An Option object.
|
||||
"""
|
||||
if infinite:
|
||||
return Option(
|
||||
(int, type(None)),
|
||||
lambda x: None
|
||||
if (x is None or x >= -1)
|
||||
else f"The keyword '{name}' only accepts None, 0, -1"
|
||||
" or a positive integer, where -1 represents infinity.",
|
||||
default_value=default_value,
|
||||
)
|
||||
return Option(
|
||||
(int, type(None)),
|
||||
lambda x: None
|
||||
if (x is None or x >= 0)
|
||||
else f"The keyword '{name}' only accepts None, 0 or a positive integer.",
|
||||
default_value=default_value,
|
||||
)
|
||||
|
||||
|
||||
def _validate_resource_quantity(name, quantity):
|
||||
if quantity < 0:
|
||||
return f"The quantity of resource {name} cannot be negative"
|
||||
if (
|
||||
isinstance(quantity, float)
|
||||
and quantity != 0.0
|
||||
and int(quantity * ray._raylet.RESOURCE_UNIT_SCALING) == 0
|
||||
):
|
||||
return (
|
||||
f"The precision of the fractional quantity of resource {name}"
|
||||
" cannot go beyond 0.0001"
|
||||
)
|
||||
resource_name = "GPU" if name == "num_gpus" else name
|
||||
if resource_name in ray._private.accelerators.get_all_accelerator_resource_names():
|
||||
(
|
||||
valid,
|
||||
error_message,
|
||||
) = ray._private.accelerators.get_accelerator_manager_for_resource(
|
||||
resource_name
|
||||
).validate_resource_request_quantity(
|
||||
quantity
|
||||
)
|
||||
if not valid:
|
||||
return error_message
|
||||
return None
|
||||
|
||||
|
||||
def _resource_option(name: str, default_value: Any = None):
|
||||
"""This is used for resource related options."""
|
||||
return Option(
|
||||
(float, int, type(None)),
|
||||
lambda x: None if (x is None) else _validate_resource_quantity(name, x),
|
||||
default_value=default_value,
|
||||
)
|
||||
|
||||
|
||||
def _validate_resources(resources: Optional[Dict[str, float]]) -> Optional[str]:
|
||||
if resources is None:
|
||||
return None
|
||||
|
||||
if "CPU" in resources or "GPU" in resources:
|
||||
return (
|
||||
"Use the 'num_cpus' and 'num_gpus' keyword instead of 'CPU' and 'GPU' "
|
||||
"in 'resources' keyword"
|
||||
)
|
||||
|
||||
for name, quantity in resources.items():
|
||||
possible_error_message = _validate_resource_quantity(name, quantity)
|
||||
if possible_error_message:
|
||||
return possible_error_message
|
||||
|
||||
return None
|
||||
|
||||
|
||||
_common_options = {
|
||||
"label_selector": Option((dict, type(None)), lambda x: validate_label_selector(x)),
|
||||
"fallback_strategy": Option(
|
||||
(list, type(None)), lambda x: validate_fallback_strategy(x)
|
||||
),
|
||||
"accelerator_type": Option((str, type(None))),
|
||||
"memory": _resource_option("memory"),
|
||||
"name": Option((str, type(None))),
|
||||
"num_cpus": _resource_option("num_cpus"),
|
||||
"num_gpus": _resource_option("num_gpus"),
|
||||
"object_store_memory": _counting_option("object_store_memory", False),
|
||||
# TODO(suquark): "placement_group", "placement_group_bundle_index"
|
||||
# and "placement_group_capture_child_tasks" are deprecated,
|
||||
# use "scheduling_strategy" instead.
|
||||
"placement_group": Option(
|
||||
(type(None), str, PlacementGroup), default_value="default"
|
||||
),
|
||||
"placement_group_bundle_index": Option(int, default_value=-1),
|
||||
"placement_group_capture_child_tasks": Option((bool, type(None))),
|
||||
"resources": Option((dict, type(None)), lambda x: _validate_resources(x)),
|
||||
"runtime_env": Option((dict, type(None))),
|
||||
"scheduling_strategy": Option(
|
||||
(
|
||||
type(None),
|
||||
str,
|
||||
PlacementGroupSchedulingStrategy,
|
||||
NodeAffinitySchedulingStrategy,
|
||||
NodeLabelSchedulingStrategy,
|
||||
)
|
||||
),
|
||||
"enable_task_events": Option(bool, default_value=True),
|
||||
"_labels": Option((dict, type(None))),
|
||||
}
|
||||
|
||||
|
||||
def issubclass_safe(obj: Any, cls_: type) -> bool:
|
||||
try:
|
||||
return issubclass(obj, cls_)
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
|
||||
_task_only_options = {
|
||||
"max_calls": _counting_option("max_calls", False, default_value=0),
|
||||
# Normal tasks may be retried on failure this many times.
|
||||
# TODO(swang): Allow this to be set globally for an application.
|
||||
"max_retries": _counting_option(
|
||||
"max_retries", default_value=ray_constants.DEFAULT_TASK_MAX_RETRIES
|
||||
),
|
||||
# override "_common_options"
|
||||
"num_cpus": _resource_option("num_cpus", default_value=1),
|
||||
"num_returns": Option(
|
||||
(int, str, type(None)),
|
||||
lambda x: None
|
||||
if (x is None or x == "dynamic" or x == "streaming" or x >= 0)
|
||||
else "Default None. When None is passed, "
|
||||
"The default value is 1 for a task and actor task, and "
|
||||
"'streaming' for generator tasks and generator actor tasks. "
|
||||
"The keyword 'num_returns' only accepts None, "
|
||||
"a non-negative integer, "
|
||||
"'streaming' (for generators), or 'dynamic'. 'dynamic' flag "
|
||||
"will be deprecated in the future, and it is recommended to use "
|
||||
"'streaming' instead.",
|
||||
default_value=None,
|
||||
),
|
||||
"object_store_memory": Option( # override "_common_options"
|
||||
(int, type(None)),
|
||||
lambda x: None
|
||||
if (x is None)
|
||||
else "Setting 'object_store_memory' is not implemented for tasks",
|
||||
),
|
||||
"retry_exceptions": Option(
|
||||
(bool, list, tuple),
|
||||
lambda x: None
|
||||
if (
|
||||
isinstance(x, bool)
|
||||
or (
|
||||
isinstance(x, (list, tuple))
|
||||
and all(issubclass_safe(x_, Exception) for x_ in x)
|
||||
)
|
||||
)
|
||||
else "retry_exceptions must be either a boolean or a list of exceptions",
|
||||
default_value=False,
|
||||
),
|
||||
"_generator_backpressure_num_objects": Option(
|
||||
(int, type(None)),
|
||||
lambda x: None
|
||||
if x != 0
|
||||
else (
|
||||
"_generator_backpressure_num_objects=0 is not allowed. "
|
||||
"Use a value > 0. If the value is equal to 1, the behavior "
|
||||
"is identical to Python generator (generator 1 object "
|
||||
"whenever `next` is called). Use -1 to disable this feature. "
|
||||
),
|
||||
),
|
||||
"_num_objects_per_yield": Option(
|
||||
(int, type(None)),
|
||||
lambda x: None
|
||||
if (x is None or x > 0)
|
||||
else (
|
||||
"_num_objects_per_yield is a private streaming generator option "
|
||||
"that must be set to a positive integer."
|
||||
),
|
||||
default_value=1,
|
||||
),
|
||||
}
|
||||
|
||||
_actor_only_options = {
|
||||
"concurrency_groups": Option((list, dict, type(None))),
|
||||
"enable_tensor_transport": Option((bool, type(None)), default_value=None),
|
||||
"lifetime": Option(
|
||||
(str, type(None)),
|
||||
lambda x: None
|
||||
if x in (None, "detached", "non_detached")
|
||||
else "actor `lifetime` argument must be one of 'detached', "
|
||||
"'non_detached' and 'None'.",
|
||||
),
|
||||
"max_concurrency": _counting_option("max_concurrency", False),
|
||||
"max_restarts": _counting_option("max_restarts", default_value=0),
|
||||
"max_task_retries": _counting_option("max_task_retries", default_value=0),
|
||||
"max_pending_calls": _counting_option("max_pending_calls", default_value=-1),
|
||||
"namespace": Option((str, type(None))),
|
||||
"get_if_exists": Option(bool, default_value=False),
|
||||
"allow_out_of_order_execution": Option((bool, type(None))),
|
||||
# Actor-wide cap on the number of unconsumed streaming-generator
|
||||
# objects across all generator tasks running on the actor. Coexists
|
||||
# with the per-method `_generator_backpressure_num_objects`: both
|
||||
# apply, and the producer blocks on whichever is tighter. -1 (or
|
||||
# None / unset) disables the actor-wide cap.
|
||||
"_actor_generator_backpressure_num_objects": Option(
|
||||
(int, type(None)),
|
||||
lambda x: None
|
||||
if (x is None or x > 0 or x == -1)
|
||||
else (
|
||||
"_actor_generator_backpressure_num_objects must be > 0 to cap the "
|
||||
"actor's total unconsumed generator objects, or -1 to disable. "
|
||||
f"Got {x}."
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
# Priority is important here because during dictionary update, same key with higher
|
||||
# priority overrides the same key with lower priority. We make use of priority
|
||||
# to set the correct default value for tasks / actors.
|
||||
|
||||
# priority: _common_options > _actor_only_options > _task_only_options
|
||||
valid_options: Dict[str, Option] = {
|
||||
**_task_only_options,
|
||||
**_actor_only_options,
|
||||
**_common_options,
|
||||
}
|
||||
# priority: _task_only_options > _common_options
|
||||
task_options: Dict[str, Option] = {**_common_options, **_task_only_options}
|
||||
# priority: _actor_only_options > _common_options
|
||||
actor_options: Dict[str, Option] = {**_common_options, **_actor_only_options}
|
||||
|
||||
remote_args_error_string = (
|
||||
"The @ray.remote decorator must be applied either with no arguments and no "
|
||||
"parentheses, for example '@ray.remote', or it must be applied using some of "
|
||||
f"the arguments in the list {list(valid_options.keys())}, for example "
|
||||
"'@ray.remote(num_returns=2, resources={\"CustomResource\": 1})'."
|
||||
)
|
||||
|
||||
|
||||
def _check_deprecate_placement_group(options: Dict[str, Any]):
|
||||
"""Check if deprecated placement group option exists."""
|
||||
placement_group = options.get("placement_group", "default")
|
||||
scheduling_strategy = options.get("scheduling_strategy")
|
||||
# TODO(suquark): @ray.remote(placement_group=None) is used in
|
||||
# "python/ray.data._internal/remote_fn.py" and many other places,
|
||||
# while "ray.data.read_api.read_datasource" set "scheduling_strategy=SPREAD".
|
||||
# This might be a bug, but it is also ok to allow them co-exist.
|
||||
if (placement_group not in ("default", None)) and (scheduling_strategy is not None):
|
||||
raise ValueError(
|
||||
"Placement groups should be specified via the "
|
||||
"scheduling_strategy option. "
|
||||
"The placement_group option is deprecated."
|
||||
)
|
||||
|
||||
|
||||
def _warn_if_using_deprecated_placement_group(
|
||||
options: Dict[str, Any], caller_stacklevel: int
|
||||
):
|
||||
placement_group = options["placement_group"]
|
||||
placement_group_bundle_index = options["placement_group_bundle_index"]
|
||||
placement_group_capture_child_tasks = options["placement_group_capture_child_tasks"]
|
||||
if placement_group != "default":
|
||||
warnings.warn(
|
||||
"placement_group parameter is deprecated. Use "
|
||||
"scheduling_strategy=PlacementGroupSchedulingStrategy(...) "
|
||||
"instead, see the usage at "
|
||||
f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-remote.", # noqa: E501
|
||||
DeprecationWarning,
|
||||
stacklevel=caller_stacklevel + 1,
|
||||
)
|
||||
if placement_group_bundle_index != -1:
|
||||
warnings.warn(
|
||||
"placement_group_bundle_index parameter is deprecated. Use "
|
||||
"scheduling_strategy=PlacementGroupSchedulingStrategy(...) "
|
||||
"instead, see the usage at "
|
||||
f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-remote.", # noqa: E501
|
||||
DeprecationWarning,
|
||||
stacklevel=caller_stacklevel + 1,
|
||||
)
|
||||
if placement_group_capture_child_tasks:
|
||||
warnings.warn(
|
||||
"placement_group_capture_child_tasks parameter is deprecated. Use "
|
||||
"scheduling_strategy=PlacementGroupSchedulingStrategy(...) "
|
||||
"instead, see the usage at "
|
||||
f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-remote.", # noqa: E501
|
||||
DeprecationWarning,
|
||||
stacklevel=caller_stacklevel + 1,
|
||||
)
|
||||
|
||||
|
||||
def validate_task_options(
|
||||
options: Dict[str, Any],
|
||||
in_options: bool,
|
||||
is_generator_callable: Optional[bool] = None,
|
||||
):
|
||||
"""Options check for Ray tasks.
|
||||
|
||||
Args:
|
||||
options: Options for Ray tasks.
|
||||
in_options: If True, we are checking the options under the context of
|
||||
".options()".
|
||||
is_generator_callable: Optional bool indicating whether the callable is a
|
||||
generator function. If provided and num_returns is 'streaming' or
|
||||
'dynamic', validates that the callable is a generator.
|
||||
"""
|
||||
for k, v in options.items():
|
||||
if k not in task_options:
|
||||
raise ValueError(
|
||||
f"Invalid option keyword {k} for remote functions. "
|
||||
f"Valid ones are {list(task_options)}."
|
||||
)
|
||||
task_options[k].validate(k, v)
|
||||
if in_options and "max_calls" in options:
|
||||
raise ValueError("Setting 'max_calls' is not supported in '.options()'.")
|
||||
_check_deprecate_placement_group(options)
|
||||
|
||||
if is_generator_callable is not None:
|
||||
num_returns = options.get("num_returns")
|
||||
if num_returns is not None:
|
||||
validate_num_returns(is_generator_callable, num_returns)
|
||||
|
||||
|
||||
def validate_actor_options(options: Dict[str, Any], in_options: bool):
|
||||
"""Options check for Ray actors.
|
||||
|
||||
Args:
|
||||
options: Options for Ray actors.
|
||||
in_options: If True, we are checking the options under the context of
|
||||
".options()".
|
||||
"""
|
||||
for k, v in options.items():
|
||||
if k not in actor_options:
|
||||
raise ValueError(
|
||||
f"Invalid option keyword {k} for actors. "
|
||||
f"Valid ones are {list(actor_options)}."
|
||||
)
|
||||
actor_options[k].validate(k, v)
|
||||
|
||||
if in_options and "concurrency_groups" in options:
|
||||
raise ValueError(
|
||||
"Setting 'concurrency_groups' is not supported in '.options()'."
|
||||
)
|
||||
|
||||
if options.get("get_if_exists") and not options.get("name"):
|
||||
raise ValueError("The actor name must be specified to use `get_if_exists`.")
|
||||
|
||||
if "object_store_memory" in options:
|
||||
warnings.warn(
|
||||
"Setting 'object_store_memory'"
|
||||
" for actors is deprecated since it doesn't actually"
|
||||
" reserve the required object store memory."
|
||||
f" Use object spilling that's enabled by default (https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/objects/object-spilling.html) " # noqa: E501
|
||||
"instead to bypass the object store memory size limitation.",
|
||||
DeprecationWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
_check_deprecate_placement_group(options)
|
||||
|
||||
|
||||
def validate_num_returns(is_generator_callable: bool, num_returns: Any) -> None:
|
||||
"""Validate num_returns for @ray.remote and @ray.method decorators.
|
||||
|
||||
This function validates:
|
||||
1. If num_returns is an integer < 0, it should fail fast.
|
||||
2. If num_returns='streaming' or 'dynamic' is used with a non-generator
|
||||
function, it should fail fast.
|
||||
|
||||
Args:
|
||||
is_generator_callable: Whether the callable is a generator function or
|
||||
async generator function.
|
||||
num_returns: The num_returns value to validate.
|
||||
|
||||
Raises:
|
||||
ValueError: If num_returns < 0, or if num_returns is 'streaming' or 'dynamic'
|
||||
but the callable is not a generator function or async generator function.
|
||||
"""
|
||||
if num_returns is None:
|
||||
return
|
||||
|
||||
# Validate num_returns < 0
|
||||
if isinstance(num_returns, int) and num_returns < 0:
|
||||
raise ValueError(f"num_returns must be >= 0, but got {num_returns}.")
|
||||
|
||||
# Validate num_returns='streaming' or 'dynamic' for generator functions
|
||||
if num_returns in ("streaming", "dynamic") and not is_generator_callable:
|
||||
raise ValueError(
|
||||
f"num_returns='{num_returns}' can only be used with generator functions "
|
||||
f"(functions that use 'yield'). "
|
||||
f"The decorated function is not a generator function."
|
||||
)
|
||||
|
||||
|
||||
def update_options(
|
||||
original_options: Dict[str, Any], new_options: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Update original options with new options and return.
|
||||
The returned updated options contain shallow copy of original options.
|
||||
"""
|
||||
|
||||
return {**original_options, **new_options}
|
||||
Reference in New Issue
Block a user