chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,234 @@
|
||||
from typing import TYPE_CHECKING, Dict, Optional, Union
|
||||
|
||||
from ray._raylet import NodeID
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.util.placement_group import PlacementGroup
|
||||
|
||||
# "DEFAULT": The default hybrid scheduling strategy
|
||||
# based on config scheduler_spread_threshold.
|
||||
# This disables any potential placement group capture.
|
||||
|
||||
# "SPREAD": Spread scheduling on a best effort basis.
|
||||
|
||||
|
||||
@PublicAPI
|
||||
class PlacementGroupSchedulingStrategy:
|
||||
"""Placement group based scheduling strategy."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
placement_group: "PlacementGroup",
|
||||
placement_group_bundle_index: int = -1,
|
||||
placement_group_capture_child_tasks: Optional[bool] = None,
|
||||
):
|
||||
"""Initialize a ``PlacementGroupSchedulingStrategy``.
|
||||
|
||||
Args:
|
||||
placement_group: the placement group this actor belongs to,
|
||||
or None if it doesn't belong to any group.
|
||||
placement_group_bundle_index: the index of the bundle
|
||||
if the actor belongs to a placement group, which may be -1 to
|
||||
specify any available bundle.
|
||||
placement_group_capture_child_tasks: Whether or not children tasks
|
||||
of this actor should implicitly use the same placement group
|
||||
as its parent. It is False by default.
|
||||
"""
|
||||
self.placement_group = placement_group
|
||||
self.placement_group_bundle_index = placement_group_bundle_index
|
||||
self.placement_group_capture_child_tasks = placement_group_capture_child_tasks
|
||||
|
||||
|
||||
@PublicAPI
|
||||
class NodeAffinitySchedulingStrategy:
|
||||
"""Static scheduling strategy used to run a task or actor on a particular node."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
node_id: str,
|
||||
soft: bool,
|
||||
_spill_on_unavailable: bool = False,
|
||||
_fail_on_unavailable: bool = False,
|
||||
):
|
||||
"""Initialize a ``NodeAffinitySchedulingStrategy``.
|
||||
|
||||
Args:
|
||||
node_id: the hex id of the node where the task or actor should run.
|
||||
soft: whether the scheduler should run the task or actor somewhere
|
||||
else if the target node doesn't exist (e.g. the node dies) or
|
||||
is infeasible during scheduling. If the node exists and is
|
||||
feasible, the task or actor will only be scheduled there. This
|
||||
means if the node doesn't have the available resources, the
|
||||
task or actor will wait indefinitely until resources become
|
||||
available. If the node doesn't exist or is infeasible, the
|
||||
task or actor will fail if ``soft`` is False or be scheduled
|
||||
somewhere else if ``soft`` is True.
|
||||
_spill_on_unavailable: (Private) When ``soft`` is True, allow
|
||||
spilling to another node if the target node is unavailable.
|
||||
_fail_on_unavailable: (Private) When ``soft`` is False, fail
|
||||
immediately if the target node is unavailable rather than
|
||||
waiting.
|
||||
"""
|
||||
# This will be removed once we standardize on node id being hex string.
|
||||
if not isinstance(node_id, str):
|
||||
node_id = node_id.hex()
|
||||
|
||||
self.node_id = node_id
|
||||
self.soft = soft
|
||||
self._spill_on_unavailable = _spill_on_unavailable
|
||||
self._fail_on_unavailable = _fail_on_unavailable
|
||||
|
||||
self._validate_attributes()
|
||||
|
||||
def _validate_attributes(self):
|
||||
invalid_node_id_error = ValueError(
|
||||
f"Invalid node_id '{self.node_id}'. Node ID must be a valid "
|
||||
"hex string. To get a list of all nodes and their IDs in your cluster, "
|
||||
"use ray.nodes(). See https://docs.ray.io/en/latest/ray-core/miscellaneous.html#node-information for more details."
|
||||
)
|
||||
try:
|
||||
node_id = NodeID.from_hex(self.node_id)
|
||||
except Exception as e:
|
||||
raise invalid_node_id_error from e
|
||||
|
||||
if node_id.is_nil():
|
||||
raise invalid_node_id_error
|
||||
|
||||
if self._spill_on_unavailable and not self.soft:
|
||||
raise ValueError(
|
||||
"_spill_on_unavailable cannot be set when soft is "
|
||||
"False. Please set soft to True to use _spill_on_unavailable."
|
||||
)
|
||||
if self._fail_on_unavailable and self.soft:
|
||||
raise ValueError(
|
||||
"_fail_on_unavailable cannot be set when soft is "
|
||||
"True. Please set soft to False to use _fail_on_unavailable."
|
||||
)
|
||||
|
||||
|
||||
def _validate_label_match_operator_values(values, operator):
|
||||
if not values:
|
||||
raise ValueError(
|
||||
f"The variadic parameter of the {operator} operator"
|
||||
f' must be a non-empty tuple: e.g. {operator}("value1", "value2").'
|
||||
)
|
||||
|
||||
index = 0
|
||||
for value in values:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(
|
||||
f"Type of value in position {index} for the {operator} operator "
|
||||
f'must be str (e.g. {operator}("value1", "value2")) '
|
||||
f"but got {str(value)} of type {type(value)}."
|
||||
)
|
||||
index = index + 1
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class In:
|
||||
def __init__(self, *values):
|
||||
_validate_label_match_operator_values(values, "In")
|
||||
self.values = list(values)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class NotIn:
|
||||
def __init__(self, *values):
|
||||
_validate_label_match_operator_values(values, "NotIn")
|
||||
self.values = list(values)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class Exists:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class DoesNotExist:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
class _LabelMatchExpression:
|
||||
"""An expression used to select node by node's labels."""
|
||||
|
||||
def __init__(self, key: str, operator: Union[In, NotIn, Exists, DoesNotExist]):
|
||||
"""Initialize a ``_LabelMatchExpression``.
|
||||
|
||||
Args:
|
||||
key: the key of label.
|
||||
operator: one of ``In``, ``NotIn``, ``Exists``, ``DoesNotExist``.
|
||||
"""
|
||||
self.key = key
|
||||
self.operator = operator
|
||||
|
||||
|
||||
LabelMatchExpressionsT = Dict[str, Union[In, NotIn, Exists, DoesNotExist]]
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class NodeLabelSchedulingStrategy:
|
||||
"""
|
||||
Label based node affinity scheduling strategy
|
||||
|
||||
scheduling_strategy=NodeLabelSchedulingStrategy({
|
||||
"region": In("us"),
|
||||
"gpu_type": Exists(),
|
||||
})
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, hard: LabelMatchExpressionsT, *, soft: LabelMatchExpressionsT = None
|
||||
):
|
||||
self.hard = _convert_map_to_expressions(hard, "hard")
|
||||
self.soft = _convert_map_to_expressions(soft, "soft")
|
||||
self._check_usage()
|
||||
|
||||
def _check_usage(self):
|
||||
if not (self.hard or self.soft):
|
||||
raise ValueError(
|
||||
"The `hard` and `soft` parameter "
|
||||
"of NodeLabelSchedulingStrategy cannot both be empty."
|
||||
)
|
||||
|
||||
|
||||
def _convert_map_to_expressions(map_expressions: LabelMatchExpressionsT, param: str):
|
||||
expressions = []
|
||||
if map_expressions is None:
|
||||
return expressions
|
||||
|
||||
if not isinstance(map_expressions, Dict):
|
||||
raise ValueError(
|
||||
f'The {param} parameter must be a map (e.g. {{"key1": In("value1")}}) '
|
||||
f"but got type {type(map_expressions)}."
|
||||
)
|
||||
|
||||
for key, value in map_expressions.items():
|
||||
if not isinstance(key, str):
|
||||
raise ValueError(
|
||||
f"The map key of the {param} parameter must "
|
||||
f'be of type str (e.g. {{"key1": In("value1")}}) '
|
||||
f"but got {str(key)} of type {type(key)}."
|
||||
)
|
||||
|
||||
if not isinstance(value, (In, NotIn, Exists, DoesNotExist)):
|
||||
raise ValueError(
|
||||
f"The map value for key {key} of the {param} parameter "
|
||||
f"must be one of the `In`, `NotIn`, `Exists` or `DoesNotExist` "
|
||||
f'operator (e.g. {{"key1": In("value1")}}) '
|
||||
f"but got {str(value)} of type {type(value)}."
|
||||
)
|
||||
|
||||
expressions.append(_LabelMatchExpression(key, value))
|
||||
return expressions
|
||||
|
||||
|
||||
SchedulingStrategyT = Union[
|
||||
None,
|
||||
str, # Literal["DEFAULT", "SPREAD"]
|
||||
PlacementGroupSchedulingStrategy,
|
||||
NodeAffinitySchedulingStrategy,
|
||||
NodeLabelSchedulingStrategy,
|
||||
]
|
||||
Reference in New Issue
Block a user