wehub-resource-sync
513 lines
19 KiB
Python
513 lines
19 KiB
Python
import functools
|
|
import math
|
|
import operator
|
|
import os
|
|
import warnings
|
|
from typing import Any, Callable, Dict, Iterable, List, Optional, Union
|
|
|
|
from .common import NodeIdStr
|
|
from ray.data._internal.execution.util import memory_string
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
class ExecutionResources:
|
|
"""Specifies resources usage or resource limits for execution.
|
|
|
|
By default this class represents resource usage. Use `for_limits` or
|
|
set `default_to_inf` to True to create an object that represents resource limits.
|
|
"""
|
|
|
|
# ``__slots__`` keeps instances small and makes attribute access go through
|
|
# slot descriptors instead of a per-instance ``__dict__``. The scheduler
|
|
# constructs many of these per iteration (every ``add``/``subtract``/
|
|
# ``max``/``copy`` returns a new object), so this is a hot-path win.
|
|
__slots__ = ("_cpu", "_gpu", "_object_store_memory", "_memory")
|
|
|
|
def __init__(
|
|
self,
|
|
cpu: Optional[float] = None,
|
|
gpu: Optional[float] = None,
|
|
object_store_memory: Optional[float] = None,
|
|
memory: Optional[float] = None,
|
|
):
|
|
"""Initializes ExecutionResources.
|
|
Args:
|
|
cpu: Amount of logical CPU slots.
|
|
gpu: Amount of logical GPU slots.
|
|
object_store_memory: Amount of object store memory.
|
|
memory: Amount of logical memory in bytes.
|
|
"""
|
|
|
|
# NOTE: Ray Core allocates fractional resources in up to 5th decimal
|
|
# digit, hence we round the values here up to it
|
|
self._cpu: Optional[float] = safe_round(cpu, 5)
|
|
self._gpu: Optional[float] = safe_round(gpu, 5)
|
|
self._object_store_memory: Optional[float] = safe_round(object_store_memory, 0)
|
|
self._memory: Optional[float] = safe_round(memory, 0)
|
|
|
|
@classmethod
|
|
def from_resource_dict(
|
|
cls,
|
|
resource_dict: Dict[str, float],
|
|
):
|
|
"""Create an ExecutionResources object from a resource dict."""
|
|
return ExecutionResources(
|
|
cpu=resource_dict.get("CPU", None) or resource_dict.get("num_cpus", None),
|
|
gpu=resource_dict.get("GPU", None) or resource_dict.get("num_gpus", None),
|
|
object_store_memory=resource_dict.get("object_store_memory", None),
|
|
memory=resource_dict.get("memory", None),
|
|
)
|
|
|
|
def to_resource_dict(self) -> Dict[str, float]:
|
|
"""Convert this ExecutionResources object to a resource dict."""
|
|
return {
|
|
"CPU": self.cpu,
|
|
"GPU": self.gpu,
|
|
"object_store_memory": self.object_store_memory,
|
|
"memory": self.memory,
|
|
}
|
|
|
|
@classmethod
|
|
def for_limits(
|
|
cls,
|
|
cpu: Optional[float] = None,
|
|
gpu: Optional[float] = None,
|
|
object_store_memory: Optional[float] = None,
|
|
memory: Optional[float] = None,
|
|
) -> "ExecutionResources":
|
|
"""Create an ExecutionResources object that represents resource limits.
|
|
|
|
Args:
|
|
cpu: Amount of logical CPU slots.
|
|
gpu: Amount of logical GPU slots.
|
|
object_store_memory: Amount of object store memory.
|
|
memory: Amount of logical memory in bytes.
|
|
|
|
Returns:
|
|
An ``ExecutionResources`` with the given limits (defaulting to
|
|
infinity for any unspecified field).
|
|
"""
|
|
return ExecutionResources(
|
|
cpu=safe_or(cpu, float("inf")),
|
|
gpu=safe_or(gpu, float("inf")),
|
|
object_store_memory=safe_or(object_store_memory, float("inf")),
|
|
memory=safe_or(memory, float("inf")),
|
|
)
|
|
|
|
@property
|
|
def cpu(self) -> float:
|
|
return self._cpu or 0.0
|
|
|
|
@property
|
|
def gpu(self) -> float:
|
|
return self._gpu or 0.0
|
|
|
|
@property
|
|
def object_store_memory(self) -> float:
|
|
return self._object_store_memory or 0
|
|
|
|
@property
|
|
def memory(self) -> float:
|
|
return self._memory or 0
|
|
|
|
def __repr__(self):
|
|
return (
|
|
f"ExecutionResources(cpu={self.cpu}, gpu={self.gpu}, "
|
|
f"object_store_memory={self.object_store_memory_str()}, "
|
|
f"memory={self.memory_str()})"
|
|
)
|
|
|
|
def __eq__(self, other: "ExecutionResources") -> bool:
|
|
return (
|
|
self.cpu == other.cpu
|
|
and self.gpu == other.gpu
|
|
and self.object_store_memory == other.object_store_memory
|
|
and self.memory == other.memory
|
|
)
|
|
|
|
def __hash__(self) -> int:
|
|
return hash(
|
|
(
|
|
self.cpu,
|
|
self.gpu,
|
|
self.object_store_memory,
|
|
self.memory,
|
|
)
|
|
)
|
|
|
|
@classmethod
|
|
@functools.cache
|
|
def zero(cls) -> "ExecutionResources":
|
|
"""Returns an ExecutionResources object with zero resources.
|
|
|
|
Returns a cached, shared singleton (``functools.cache`` keyed on ``cls``)
|
|
-- ``zero()`` is called all over the scheduler hot path (e.g.
|
|
``.max(zero())``) and instances are immutable in practice (every
|
|
arithmetic op returns a new object and there are no setters), so sharing
|
|
one instance is safe and avoids the per-call allocation.
|
|
"""
|
|
return ExecutionResources(0.0, 0.0, 0.0, 0.0)
|
|
|
|
@classmethod
|
|
@functools.cache
|
|
def inf(cls) -> "ExecutionResources":
|
|
"""Returns an ExecutionResources object with infinite resources.
|
|
|
|
Returns a cached, shared singleton (see :meth:`zero` for why this is
|
|
safe).
|
|
"""
|
|
return ExecutionResources.for_limits()
|
|
|
|
@classmethod
|
|
def combine(
|
|
cls,
|
|
resources: Iterable["ExecutionResources"],
|
|
fn: Callable[[float, float], float],
|
|
) -> Optional["ExecutionResources"]:
|
|
"""Fold an iterable of ``ExecutionResources`` per dimension with ``fn``.
|
|
|
|
``fn(acc, value)`` combines two per-dimension floats -- e.g.
|
|
``operator.add`` for a sum, or ``max``/``min`` for an element-wise
|
|
max/min. Accumulates raw floats in a single pass and allocates a single
|
|
result object, instead of one intermediate per element as
|
|
``reduce(lambda a, b: a.<op>(b), resources)`` would.
|
|
|
|
Seeds with the first element (so no per-``fn`` identity is needed) and
|
|
returns ``None`` for an empty iterable, which may be a one-shot
|
|
generator (so it's consumed exactly once).
|
|
"""
|
|
iterator = iter(resources)
|
|
first = next(iterator, None)
|
|
if first is None:
|
|
return None
|
|
cpu = first.cpu
|
|
gpu = first.gpu
|
|
object_store_memory = first.object_store_memory
|
|
memory = first.memory
|
|
for r in iterator:
|
|
cpu = fn(cpu, r.cpu)
|
|
gpu = fn(gpu, r.gpu)
|
|
object_store_memory = fn(object_store_memory, r.object_store_memory)
|
|
memory = fn(memory, r.memory)
|
|
return ExecutionResources(cpu, gpu, object_store_memory, memory)
|
|
|
|
@classmethod
|
|
def combine_sum(
|
|
cls, resources: Iterable["ExecutionResources"]
|
|
) -> "ExecutionResources":
|
|
"""Sum an iterable of ``ExecutionResources`` in a single pass.
|
|
|
|
Thin wrapper over :meth:`combine` with addition. Empty folds are common
|
|
(e.g. completed-ops / downstream-ineligible usage rollups on most
|
|
iterations), so an empty input reuses the shared ``zero()`` singleton
|
|
instead of allocating.
|
|
"""
|
|
result = cls.combine(resources, operator.add)
|
|
return result if result is not None else cls.zero()
|
|
|
|
def is_zero(self) -> bool:
|
|
"""Returns True if all resources are zero."""
|
|
return (
|
|
self.cpu == 0.0
|
|
and self.gpu == 0.0
|
|
and self.object_store_memory == 0.0
|
|
and self.memory == 0.0
|
|
)
|
|
|
|
def is_non_negative(self) -> bool:
|
|
"""Returns True if all resources are non-negative."""
|
|
return (
|
|
self.cpu >= 0
|
|
and self.gpu >= 0
|
|
and self.object_store_memory >= 0
|
|
and self.memory >= 0
|
|
)
|
|
|
|
def object_store_memory_str(self) -> str:
|
|
"""Returns a human-readable string for the object store memory field."""
|
|
if self.object_store_memory == float("inf"):
|
|
return "inf"
|
|
return memory_string(self.object_store_memory)
|
|
|
|
def memory_str(self) -> str:
|
|
"""Returns a human-readable string for the memory field."""
|
|
if self.memory == float("inf"):
|
|
return "inf"
|
|
return memory_string(self.memory)
|
|
|
|
def copy(
|
|
self,
|
|
cpu: Optional[float] = None,
|
|
gpu: Optional[float] = None,
|
|
memory: Optional[float] = None,
|
|
object_store_memory: Optional[float] = None,
|
|
) -> "ExecutionResources":
|
|
"""Returns a copy of this ExecutionResources object allowing to override
|
|
specific resources as necessary"""
|
|
return ExecutionResources(
|
|
cpu=safe_or(cpu, self.cpu),
|
|
gpu=safe_or(gpu, self.gpu),
|
|
object_store_memory=safe_or(object_store_memory, self.object_store_memory),
|
|
memory=safe_or(memory, self.memory),
|
|
)
|
|
|
|
def add(self, other: "ExecutionResources") -> "ExecutionResources":
|
|
"""Adds execution resources.
|
|
|
|
Args:
|
|
other: The other ``ExecutionResources`` to add to this one.
|
|
|
|
Returns:
|
|
A new ExecutionResource object with summed resources.
|
|
"""
|
|
return ExecutionResources(
|
|
cpu=self.cpu + other.cpu,
|
|
gpu=self.gpu + other.gpu,
|
|
object_store_memory=self.object_store_memory + other.object_store_memory,
|
|
memory=self.memory + other.memory,
|
|
)
|
|
|
|
def subtract(self, other: "ExecutionResources") -> "ExecutionResources":
|
|
"""Subtracts execution resources.
|
|
|
|
Args:
|
|
other: The other ``ExecutionResources`` to subtract from this one.
|
|
|
|
Returns:
|
|
A new ExecutionResource object with subtracted resources.
|
|
"""
|
|
return ExecutionResources(
|
|
cpu=self.cpu - other.cpu,
|
|
gpu=self.gpu - other.gpu,
|
|
object_store_memory=self.object_store_memory - other.object_store_memory,
|
|
memory=self.memory - other.memory,
|
|
)
|
|
|
|
def max(self, other: "ExecutionResources") -> "ExecutionResources":
|
|
"""Returns the maximum for each resource type."""
|
|
return ExecutionResources(
|
|
cpu=max(self.cpu, other.cpu),
|
|
gpu=max(self.gpu, other.gpu),
|
|
object_store_memory=max(
|
|
self.object_store_memory, other.object_store_memory
|
|
),
|
|
memory=max(self.memory, other.memory),
|
|
)
|
|
|
|
def min(self, other: "ExecutionResources") -> "ExecutionResources":
|
|
"""Returns the minimum for each resource type."""
|
|
return ExecutionResources(
|
|
cpu=min(self.cpu, other.cpu),
|
|
gpu=min(self.gpu, other.gpu),
|
|
object_store_memory=min(
|
|
self.object_store_memory, other.object_store_memory
|
|
),
|
|
memory=min(self.memory, other.memory),
|
|
)
|
|
|
|
def satisfies_limit(
|
|
self,
|
|
limit: "ExecutionResources",
|
|
*,
|
|
ignore_object_store_memory: bool = False,
|
|
) -> bool:
|
|
"""Return if this resource struct meets the specified limits.
|
|
|
|
Note that None for a field means no limit.
|
|
|
|
Args:
|
|
limit: The resource limits to check against.
|
|
ignore_object_store_memory: If True, ignore the object store memory
|
|
limit when checking if this resource struct meets the limits.
|
|
|
|
Returns:
|
|
``True`` if every resource is within the corresponding limit.
|
|
"""
|
|
return (
|
|
self.cpu <= limit.cpu
|
|
and self.gpu <= limit.gpu
|
|
and (
|
|
ignore_object_store_memory
|
|
or self.object_store_memory <= limit.object_store_memory
|
|
)
|
|
and self.memory <= limit.memory
|
|
)
|
|
|
|
def scale(self, f: float) -> "ExecutionResources":
|
|
"""Return copy with all set values scaled by `f`."""
|
|
if f < 0:
|
|
raise ValueError("Scaling factor must be non-negative.")
|
|
if f == 0:
|
|
# Explicitly handle the zero case, because `0 * inf` is undefined.
|
|
return ExecutionResources.zero()
|
|
|
|
return ExecutionResources(
|
|
cpu=self.cpu * f,
|
|
gpu=self.gpu * f,
|
|
object_store_memory=self.object_store_memory * f,
|
|
memory=self.memory * f,
|
|
)
|
|
|
|
def floordiv(self, other: "ExecutionResources") -> "ExecutionResources":
|
|
"""Returns the floor division of resources."""
|
|
|
|
def _div(a, b):
|
|
if b == 0:
|
|
return float("inf")
|
|
if a == float("inf"):
|
|
return float("inf")
|
|
return math.floor(a / b)
|
|
|
|
return ExecutionResources(
|
|
cpu=_div(self.cpu, other.cpu),
|
|
gpu=_div(self.gpu, other.gpu),
|
|
object_store_memory=_div(
|
|
self.object_store_memory, other.object_store_memory
|
|
),
|
|
memory=_div(self.memory, other.memory),
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class ExecutionOptions:
|
|
"""Common options for execution.
|
|
|
|
Some options may not be supported on all executors (e.g., resource limits).
|
|
|
|
Attributes:
|
|
resource_limits: Set a limit on the logical resources a Dataset can use.
|
|
Autodetected by default.
|
|
exclude_resources: Amount of resources to exclude from Ray Data.
|
|
Set this if you have other workloads running on the same cluster.
|
|
Note,
|
|
- If using Ray Data with Ray Train, training resources are
|
|
automatically reserved and you don't need to set exclude_resources
|
|
for them.
|
|
- For each resource type, resource_limits and exclude_resources can
|
|
not be both set.
|
|
preserve_order: Set this to preserve the ordering between blocks processed by
|
|
operators. Off by default.
|
|
actor_locality_enabled: Whether to enable locality-aware task dispatch to
|
|
actors (off by default). This parameter applies to both stateful map and
|
|
streaming_split operations.
|
|
verbose_progress: Whether to report progress individually per operator. By
|
|
default, only AllToAll operators and global progress is reported. This
|
|
option is useful for performance debugging. On by default.
|
|
label_selector: A mapping of label key to label value. When set, every task
|
|
and actor launched by this Dataset (including shuffle, sort, and
|
|
aggregator actors) carries this label selector in its remote args,
|
|
constraining placement to nodes whose labels satisfy the selector.
|
|
Used to scope a Dataset to a labeled subset of the cluster (e.g.
|
|
``{"ray-subcluster": "training"}``). Operator-level ``label_selector``
|
|
entries in ``ray_remote_args`` take precedence on key conflicts so
|
|
existing node-pin selectors are preserved.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
resource_limits: Optional[ExecutionResources] = None,
|
|
exclude_resources: Optional[ExecutionResources] = None,
|
|
preserve_order: bool = False,
|
|
actor_locality_enabled: bool = True,
|
|
verbose_progress: Optional[bool] = None,
|
|
label_selector: Optional[Dict[str, str]] = None,
|
|
):
|
|
"""Initialize execution options.
|
|
|
|
Args:
|
|
resource_limits: Limit on logical resources a Dataset can use.
|
|
Defaults to auto-detected limits.
|
|
exclude_resources: Resources to exclude from Ray Data.
|
|
preserve_order: Whether to preserve block processing order.
|
|
actor_locality_enabled: Whether to enable locality-aware dispatch for
|
|
stateful map and streaming split operations.
|
|
verbose_progress: Whether to report progress per operator. If None,
|
|
read from ``RAY_DATA_VERBOSE_PROGRESS``.
|
|
label_selector: Per-Dataset label selector applied to every task and
|
|
actor launched by Ray Data. ``None`` means no selector is added.
|
|
"""
|
|
if resource_limits is None:
|
|
resource_limits = ExecutionResources.for_limits()
|
|
self.resource_limits = resource_limits
|
|
if exclude_resources is None:
|
|
exclude_resources = ExecutionResources.zero()
|
|
self.exclude_resources = exclude_resources
|
|
self.preserve_order = preserve_order
|
|
self.actor_locality_enabled = actor_locality_enabled
|
|
if verbose_progress is None:
|
|
verbose_progress = bool(
|
|
int(os.environ.get("RAY_DATA_VERBOSE_PROGRESS", "1"))
|
|
)
|
|
self.verbose_progress = verbose_progress
|
|
self.label_selector = label_selector
|
|
|
|
def __repr__(self) -> str:
|
|
return (
|
|
f"ExecutionOptions(resource_limits={self.resource_limits}, "
|
|
f"exclude_resources={self.exclude_resources}, "
|
|
f"preserve_order={self.preserve_order}, "
|
|
f"actor_locality_enabled={self.actor_locality_enabled}, "
|
|
f"verbose_progress={self.verbose_progress}, "
|
|
f"label_selector={self.label_selector})"
|
|
)
|
|
|
|
@property
|
|
def resource_limits(self) -> ExecutionResources:
|
|
return self._resource_limits
|
|
|
|
@resource_limits.setter
|
|
def resource_limits(self, value: ExecutionResources) -> None:
|
|
self._resource_limits = ExecutionResources.for_limits(
|
|
cpu=value._cpu,
|
|
gpu=value._gpu,
|
|
object_store_memory=value._object_store_memory,
|
|
memory=value._memory,
|
|
)
|
|
|
|
def is_resource_limits_default(self):
|
|
"""Returns True if resource_limits is the default value."""
|
|
return self._resource_limits == ExecutionResources.for_limits()
|
|
|
|
def validate(self) -> None:
|
|
"""Validate the options."""
|
|
for attr in ["cpu", "gpu", "object_store_memory"]:
|
|
if (
|
|
getattr(self.resource_limits, attr) != float("inf")
|
|
and getattr(self.exclude_resources, attr, 0) > 0
|
|
):
|
|
raise ValueError(
|
|
"resource_limits and exclude_resources cannot "
|
|
f" both be set for {attr} resource."
|
|
)
|
|
|
|
@property
|
|
def locality_with_output(self) -> bool:
|
|
return False
|
|
|
|
@locality_with_output.setter
|
|
def locality_with_output(self, value: Union[bool, List[NodeIdStr]]) -> None:
|
|
if value:
|
|
warnings.warn(
|
|
"`ExecutionOptions.locality_with_output` has been removed and is now "
|
|
"a no-op. We don't recommend using it anymore, but if you still want "
|
|
"to replicate its behavior, follow the instructions in this gist: "
|
|
"https://gist.github.com/bveeramani/51e0383bb3680dd78fdfb92d76ea22a8.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
|
|
|
|
def safe_or(value: Optional[Any], alt: Any) -> Any:
|
|
return value if value is not None else alt
|
|
|
|
|
|
def safe_round(
|
|
value: Optional[float], ndigits: Optional[int] = None
|
|
) -> Optional[float]:
|
|
if value is None:
|
|
return None
|
|
elif ndigits is None or math.isinf(value):
|
|
return value
|
|
else:
|
|
return round(value, ndigits)
|