wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/data/_internal/execution/interfaces/__init__.py
+22
View File
@@ -0,0 +1,22 @@
from .common import NodeIdStr
from .execution_options import ExecutionOptions, ExecutionResources
from .executor import Executor, OutputIterator
from .physical_operator import PhysicalOperator, ReportsExtraResourceUsage
from .ref_bundle import BlockEntry, BlockSlice, RefBundle
from .task_context import TaskContext
from .transform_fn import AllToAllTransformFn
__all__ = [
"AllToAllTransformFn",
"BlockEntry",
"BlockSlice",
"ExecutionOptions",
"ExecutionResources",
"Executor",
"NodeIdStr",
"OutputIterator",
"PhysicalOperator",
"RefBundle",
"ReportsExtraResourceUsage",
"TaskContext",
]
python/ray/data/_internal/execution/interfaces/common.py
+186
View File
@@ -0,0 +1,186 @@
import bisect
import json
from typing import Dict, List, Tuple
from ray.data._internal.util import GiB, KiB, MiB
from ray.util.metrics import Histogram
# Node id string returned by `ray.get_runtime_context().get_node_id()`.
NodeIdStr = str
# Used for time-based histograms (e.g. task completion time, block completion time)
histogram_buckets_s = [
0.1,
0.25,
0.5,
1.0,
2.5,
5.0,
7.5,
10.0,
15.0,
20.0,
25.0,
50.0,
75.0,
100.0,
150.0,
500.0,
1000.0,
2500.0,
5000.0,
]
# Used for size-based histograms (e.g. block size in bytes)
histogram_buckets_bytes = [
KiB,
8 * KiB,
64 * KiB,
128 * KiB,
256 * KiB,
512 * KiB,
MiB,
8 * MiB,
64 * MiB,
128 * MiB,
256 * MiB,
512 * MiB,
GiB,
4 * GiB,
16 * GiB,
64 * GiB,
128 * GiB,
256 * GiB,
512 * GiB,
1024 * GiB,
4096 * GiB,
]
# Used for row count-based histograms (e.g. block size in rows)
histogram_bucket_rows = [
1,
5,
10,
25,
50,
100,
250,
500,
1_000,
2_500,
5_000,
10_000,
25_000,
50_000,
100_000,
250_000,
500_000,
1_000_000,
2_500_000,
5_000_000,
10_000_000,
]
class RuntimeMetricsHistogram:
"""
Class that tracks a histogram of values.
Contains helper methods to record the values and apply those values to a `ray.util.metrics.Histogram` metric.
"""
def __init__(self, boundaries: List[float]):
self._boundaries = boundaries
# Initialize bucket counts to 0 (+1 additional bucket to represent the +Inf bucket)
self._bucket_counts = [0 for _ in range(len(boundaries) + 1)]
self._memoized_avg = None
def observe(self, value: float, num_observations: int = 1):
self._bucket_counts[self._find_bucket_index(value)] += num_observations
self._memoized_avg = None
def export_to(
self,
metric: Histogram,
tags: Dict[str, str],
):
"""
This method calculates the difference between the current bucket counts and the previous bucket counts,
and applies those observations to the metric.
This method stores the previous_bucket_counts in the metric as `last_applied_bucket_counts_for_tags`.
"""
if getattr(metric, "last_applied_bucket_counts_for_tags", None) is None:
metric.last_applied_bucket_counts_for_tags = {}
tags_key = json.dumps(tags, sort_keys=True)
previous_bucket_counts = metric.last_applied_bucket_counts_for_tags.get(
tags_key
)
for i in range(len(self._bucket_counts)):
# Pick a value between the boundaries so the sample falls into the right bucket.
# We need to calculate the mid point because choosing the exact boundary value
# seems to have unreliable behavior on which bucket it ends up in.
boundary_upper_bound = (
self._boundaries[i]
if i < len(self._bucket_counts) - 1
# Since choosing an exact boundary value is unreliable to if it'll
# end up in the upper or lower bucket, we add a small buffer to the
# last boundary. The amount of the value doesn't matter much
# since it's the last bucket and should go to infinity.
else self._boundaries[-1] + 100
)
boundary_lower_bound = self._boundaries[i - 1] if i > 0 else 0
bucket_value = (boundary_upper_bound + boundary_lower_bound) / 2
# Calculate how many observations to add to the metric
diff = (
self._bucket_counts[i] - previous_bucket_counts[i]
if previous_bucket_counts is not None
else self._bucket_counts[i]
)
for _ in range(diff):
metric.observe(bucket_value, tags)
metric.last_applied_bucket_counts_for_tags[
tags_key
] = self._bucket_counts.copy()
def __repr__(self):
if self._memoized_avg is None:
self._memoized_avg = self._calculate_average_value()
total_samples, average = self._memoized_avg
return f"(samples: {total_samples}, avg: {average:.2f})"
def _calculate_average_value(self) -> Tuple[int, float]:
"""
Calculate the average value of all samples.
Used to show a representative value for the histogram when
printing the histogram as a string.
"""
total_samples = sum(self._bucket_counts)
if total_samples == 0:
return total_samples, 0
weighted_sum = 0.0
for i, count in enumerate(self._bucket_counts):
if count > 0:
# Calculate representative value for this bucket
if i == 0:
# First bucket: 0 to first boundary
bucket_value = self._boundaries[0] / 2
elif i == len(self._bucket_counts) - 1:
# Last bucket: last boundary to +inf
bucket_value = self._boundaries[-1] * 1.5
else:
# Middle buckets: between boundaries
bucket_value = (self._boundaries[i - 1] + self._boundaries[i]) / 2
weighted_sum += bucket_value * count
average = weighted_sum / total_samples
return total_samples, average
def _find_bucket_index(self, value: float):
return bisect.bisect_left(self._boundaries, value)
python/ray/data/_internal/execution/interfaces/distribution_tracker.py
+178
View File
@@ -0,0 +1,178 @@
import math
from typing import Dict, Optional, Union
try:
from datasketches import kll_doubles_sketch
_DATASKETCHES_AVAILABLE = True
except ImportError:
_DATASKETCHES_AVAILABLE = False
class DistributionTracker:
"""Tracks the running mean, variance, min, max, and approximate percentiles of a
stream of values using Welford's algorithm for moments and a KLL sketch for
quantiles.
More on Welford's algorithm:
https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
"""
def __init__(self):
self._count = 0
self._mean = 0.0
self._m2 = 0.0
self._min = float("inf")
self._max = float("-inf")
self._sketch = kll_doubles_sketch(200) if _DATASKETCHES_AVAILABLE else None
def add_sample(self, value: float) -> None:
self._count += 1
delta = value - self._mean
self._mean += delta / self._count
delta2 = value - self._mean
self._m2 += delta * delta2
if value < self._min:
self._min = value
if value > self._max:
self._max = value
if self._sketch is not None:
self._sketch.update(value)
def merge(self, other: "DistributionTracker") -> None:
"""Merge another tracker into this one (associative, commutative).
Uses Chan's parallel variant of Welford's algorithm for moments.
See: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford:~:text=Parallel%20algorithm%5Bedit%5D
"""
if other is self:
# Merging an accumulator into itself would double its samples
# (count, m2, and the sketch), so treat it as a no-op.
return
if other._count == 0:
return
if self._count == 0:
self._count = other._count
self._mean = other._mean
self._m2 = other._m2
self._min = other._min
self._max = other._max
else:
delta = other._mean - self._mean
total = self._count + other._count
self._m2 += other._m2 + (delta**2) * self._count * other._count / total
self._mean = (self._count * self._mean + other._count * other._mean) / total
self._count = total
self._min = min(self._min, other._min)
self._max = max(self._max, other._max)
if self._sketch is None or other._sketch is None:
# Moments above still merged; quantile detail is lost for the
# side(s) without a sketch.
self._sketch = None
else:
try:
self._sketch.merge(other._sketch)
except Exception:
self._sketch = None
@property
def num_samples(self) -> int:
return self._count
@property
def mean(self) -> float:
return self._mean
@property
def variance(self) -> float:
if self._count < 2:
return 0.0
return self._m2 / (self._count - 1)
@property
def stddev(self) -> float:
return math.sqrt(self.variance)
@property
def min(self) -> Optional[float]:
if self._count == 0:
return None
return self._min
@property
def max(self) -> Optional[float]:
if self._count == 0:
return None
return self._max
def _quantile(self, q: float) -> Optional[float]:
if self._sketch is None or self._count == 0:
return None
return self._sketch.get_quantiles([q])[0]
@property
def p25(self) -> Optional[float]:
return self._quantile(0.25)
@property
def p50(self) -> Optional[float]:
return self._quantile(0.5)
@property
def p75(self) -> Optional[float]:
return self._quantile(0.75)
@property
def p90(self) -> Optional[float]:
return self._quantile(0.9)
@property
def p95(self) -> Optional[float]:
return self._quantile(0.95)
@property
def p99(self) -> Optional[float]:
return self._quantile(0.99)
def as_dict(self) -> Dict[str, Optional[Union[int, float]]]:
return {
"num_samples": self.num_samples,
"mean": self.mean,
"variance": self.variance,
"min": self.min,
"max": self.max,
"p25": self.p25,
"p50": self.p50,
"p75": self.p75,
"p90": self.p90,
"p95": self.p95,
"p99": self.p99,
}
# ``kll_doubles_sketch`` is a C++-backed object that does not
# pickle natively. DistributionTracker rides on DatasetStats
# (via Timer), which is cloudpickled when Datasets cross actor /
# process boundaries — without these hooks any such transfer
# raises ``TypeError: cannot pickle 'kll_doubles_sketch' object``.
# The sketch exposes its own byte serialization, so we round-trip
# through that.
def __getstate__(self):
state = self.__dict__.copy()
if self._sketch is not None:
state["_sketch"] = self._sketch.serialize()
return state
def __setstate__(self, state):
self.__dict__.update(state)
# If the source had datasketches but this side doesn't, drop
# the sketch (percentiles will return None — same fallback as a
# default construction without datasketches installed).
if self._sketch is not None and not _DATASKETCHES_AVAILABLE:
self._sketch = None
elif self._sketch is not None and not isinstance(
self._sketch, kll_doubles_sketch
):
self._sketch = kll_doubles_sketch.deserialize(self._sketch)
python/ray/data/_internal/execution/interfaces/execution_options.py
+512
View File
@@ -0,0 +1,512 @@
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)
python/ray/data/_internal/execution/interfaces/executor.py
+106
View File
@@ -0,0 +1,106 @@
from abc import ABC, abstractmethod
from typing import ContextManager, Iterator, List, Optional
from .execution_options import ExecutionOptions
from .physical_operator import PhysicalOperator
from .ref_bundle import RefBundle
from ray.data._internal.stats import DatasetStats
class OutputIterator(Iterator[RefBundle], ABC):
"""Iterator used to access the output of an Executor execution.
This is a blocking iterator. Datasets guarantees that all its iterators are
thread-safe (i.e., multiple threads can block on them at the same time).
"""
@abstractmethod
def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
"""Can be used to pull outputs by a specified output index.
This is used to support the streaming_split() API, where the output of a
streaming execution is to be consumed by multiple processes.
Args:
output_split_idx: The output split index to get results for. This arg is
only allowed for iterators created by `Dataset.streaming_split()`.
Returns:
The next ``RefBundle`` of outputs for the given split index.
Raises:
StopIteration: If there are no more outputs to return.
"""
...
def __next__(self) -> RefBundle:
return self.get_next()
class Executor(ContextManager, ABC):
"""Abstract class for executors, which implement physical operator execution.
Subclasses:
StreamingExecutor
"""
def __init__(self, options: ExecutionOptions):
"""Create the executor."""
options.validate()
self._options = options
@abstractmethod
def execute(
self,
dag: PhysicalOperator,
initial_stats: Optional[DatasetStats] = None,
callbacks: Optional[List] = None,
) -> OutputIterator:
"""Start execution.
Args:
dag: The operator graph to execute.
initial_stats: The DatasetStats to prepend to the stats returned by the
executor. These stats represent actions done to compute inputs.
callbacks: A list of ExecutionCallbacks to run during execution.
This method keeps and uses the exact list you pass in, so do not
pass an empty list like ``[]`` directly. Create the list first,
then pass it.
Returns:
An ``OutputIterator`` over the execution's output ref bundles.
"""
...
def shutdown(self, force: bool, exception: Optional[Exception] = None):
"""Shutdown an executor, which may still be running.
This should interrupt execution and clean up any used resources.
Args:
force: Controls whether shutdown should forcefully terminate all execution
activity (making sure that upon returning from this method all
activities are stopped). When force=False, some activities could be
terminated asynchronously (ie this method won't provide guarantee
that they stop executing before returning from this method)
exception: The exception that causes the executor to shut down, or None if
the executor finishes successfully.
"""
pass
@abstractmethod
def get_stats(self) -> DatasetStats:
"""Return stats for the execution so far.
This is generally called after `execute` has completed, but may be called
while iterating over `execute` results for streaming execution.
"""
...
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback, /):
# NOTE: ``ContextManager`` semantic must guarantee that executor
# fully shutdown upon returning from this method
self.shutdown(force=True, exception=exc_value)
python/ray/data/_internal/execution/interfaces/op_runtime_metrics.py
+1231
View File
File diff suppressed because it is too large Load Diff
python/ray/data/_internal/execution/interfaces/physical_operator.py
+1286
View File
File diff suppressed because it is too large Load Diff
python/ray/data/_internal/execution/interfaces/ref_bundle.py
+481
View File
@@ -0,0 +1,481 @@
import itertools
from collections import defaultdict
from dataclasses import dataclass
from typing import Dict, Iterable, Iterator, List, Optional, Tuple
import ray
from .common import NodeIdStr
from ray.data._internal.memory_tracing import trace_deallocation
from ray.data.block import (
Block,
BlockAccessor,
BlockMetadata,
Schema,
_take_first_non_empty_schema,
)
from ray.data.context import DataContext
from ray.types import ObjectRef
@dataclass(frozen=True)
class BlockSlice:
"""A slice of a block."""
# Starting row offset (inclusive) within the block.
start_offset: int
# Ending row offset (exclusive) within the block.
end_offset: int
@property
def num_rows(self) -> int:
return self.end_offset - self.start_offset
@dataclass(frozen=True, slots=True)
class BlockEntry:
"""One block delivery: the ref + the block's measured metadata.
Used as the element type of ``RefBundle.blocks`` (replaces the legacy
``(ObjectRef, BlockMetadata)`` 2-tuple shape). Naming the fields makes
every call site self-describing and reserves room for the bundle entry
to grow without disturbing the surrounding shape.
"""
ref: ObjectRef[Block]
metadata: BlockMetadata
@dataclass(frozen=True)
class RefBundle:
"""A group of data block references and their metadata.
Operators take in and produce streams of RefBundles.
Most commonly a RefBundle consists of a single block object reference.
In some cases, e.g., due to block splitting, or for a reduce task, there may
be more than one block.
Block bundles have ownership semantics, i.e., shared ownership (similar to C++
shared_ptr, multiple operators share the same block bundle), or unique ownership
(similar to C++ unique_ptr, only one operator owns the block bundle). This
allows operators to know whether they can destroy blocks when they don't need
them. Destroying blocks eagerly is more efficient than waiting for Python GC /
Ray reference counting to kick in.
"""
# Per-block entries. The size_bytes must be known in the metadata,
# num_rows is optional. Legacy ``(ref, metadata)`` 2-tuples are no longer
# accepted at construction and must be explicitly wrapped in ``BlockEntry``
# (``__post_init__`` rejects anything else with an actionable assertion).
blocks: Tuple[BlockEntry, ...]
# The schema of the blocks in this bundle. This is optional, and may be None
# if blocks are empty.
schema: Optional["Schema"]
# Whether we own the blocks (can safely destroy them).
owns_blocks: bool
# The slices of the blocks in this bundle. After __post_init__, this is always
# a list with length equal to len(blocks). Individual entries can be None to
# represent a full block (equivalent to BlockSlice(0, num_rows)).
# Pass None during construction to initialize all slices as None (full blocks).
slices: Optional[Tuple[Optional[BlockSlice], ...]] = None
# This attribute is used by the split() operator to assign bundles to logical
# output splits. It is otherwise None.
output_split_idx: Optional[int] = None
# Object metadata (size, locations, spilling status)
_cached_object_meta: Optional[Dict[ObjectRef, "_ObjectMetadata"]] = None
# Preferred locations for this bundle determined based on the locations
# of individual objects and their corresponding size, ie location with the
# largest total number of bytes present there has the highest preference.
_cached_preferred_locations: Optional[Dict[NodeIdStr, int]] = None
def __post_init__(self):
if self.schema is not None:
import pyarrow as pa
from ray.data._internal.pandas_block import PandasBlockSchema
assert isinstance(
self.schema, (pa.lib.Schema, PandasBlockSchema)
), f"Schema must be a pyarrow or PandasBlockSchema, got {type(self.schema)}"
if not isinstance(self.blocks, tuple):
object.__setattr__(self, "blocks", tuple(self.blocks))
for entry in self.blocks:
assert isinstance(entry, BlockEntry), (
f"RefBundle.blocks must contain BlockEntry instances; got {type(entry).__name__}. "
"Construct entries with `BlockEntry(ref=..., metadata=...)`."
)
if self.slices is None:
object.__setattr__(self, "slices", (None,) * len(self.blocks))
else:
if not isinstance(self.slices, tuple):
object.__setattr__(self, "slices", tuple(self.slices))
assert len(self.blocks) == len(
self.slices
), "Number of blocks and slices must match"
# Validate slice ranges
for entry, block_slice in zip(self.blocks, self.slices):
if block_slice is not None:
assert (
block_slice.start_offset >= 0
), f"Slice start_offset must be non-negative: {block_slice.start_offset}"
assert (
block_slice.end_offset >= block_slice.start_offset
), f"Slice end_offset must be >= start_offset: [{block_slice.start_offset}, {block_slice.end_offset})"
if entry.metadata.num_rows is not None:
assert (
block_slice.end_offset <= entry.metadata.num_rows
), f"Slice range [{block_slice.start_offset}, {block_slice.end_offset}) exceeds block num_rows: {entry.metadata.num_rows}"
for entry in self.blocks:
if entry.metadata.size_bytes is None:
raise ValueError(
"The size in bytes of the block must be known: {}".format(entry)
)
@property
def block_refs(self) -> List[ObjectRef[Block]]:
"""List of block references in this bundle."""
return [entry.ref for entry in self.blocks]
@property
def metadata(self) -> List[BlockMetadata]:
"""List of block metadata in this bundle."""
return [entry.metadata for entry in self.blocks]
def num_rows(self) -> Optional[int]:
"""Number of rows present in this bundle, if known.
Iterates through blocks and their corresponding slices to calculate the total.
Note: Block metadata always refers to the full block, not the slice.
- If block_slice is None, uses the full block's metadata.num_rows
- If block_slice is present, uses the slice's num_rows (partial block portion)
- Returns None if any full block has unknown row count (metadata.num_rows is None)
"""
total = 0
for metadata, block_slice in zip(self.metadata, self.slices):
if block_slice is None:
if metadata.num_rows is None:
return None
total += metadata.num_rows
else:
total += block_slice.num_rows
return total
def size_bytes(self) -> int:
"""Size of the blocks of this bundle in bytes.
Iterates through blocks and their corresponding slices to calculate the total size.
Note: Block metadata always refers to the full block, not the slice.
- If block_slice is None, uses the full block's metadata.size_bytes
- If block_slice is present but num_rows is unknown or zero, uses full metadata.size_bytes
- If block_slice represents a partial block, estimates size proportionally based on
(metadata.size_bytes / metadata.num_rows) * block_slice.num_rows
- Otherwise, uses the full metadata.size_bytes
"""
total = 0
for entry, block_slice in zip(self.blocks, self.slices):
metadata = entry.metadata
if block_slice is None:
# Full block
total += metadata.size_bytes
elif metadata.num_rows is None or metadata.num_rows == 0:
# Unknown num_rows or empty block - use full metadata size
total += metadata.size_bytes
elif metadata.num_rows != block_slice.num_rows:
# Partial block - estimate size based on rows
per_row = metadata.size_bytes / metadata.num_rows
total += max(1, round(per_row * block_slice.num_rows))
else:
total += metadata.size_bytes
return total
def destroy_if_owned(self) -> int:
"""Clears the object store memory for these blocks if owned.
Returns:
The number of bytes freed.
"""
should_free = self.owns_blocks and DataContext.get_current().eager_free
for block_ref in self.block_refs:
trace_deallocation(
block_ref, "RefBundle.destroy_if_owned", free=should_free
)
return self.size_bytes() if should_free else 0
def get_preferred_object_locations(self) -> Dict[NodeIdStr, int]:
"""Returns a mapping of node IDs to total bytes stored on each node.
Returns:
Dict mapping node ID to total bytes stored on that node
"""
meta = self._get_cached_metadata()
if self._cached_preferred_locations is None:
preferred_locs: Dict[NodeIdStr, int] = defaultdict(int)
for ref, obj_meta in meta.items():
for loc in obj_meta.locs:
preferred_locs[loc] += obj_meta.size
# NOTE: We're working around object being immutable to update cached
# values (safe)
object.__setattr__(self, "_cached_preferred_locations", preferred_locs)
return self._cached_preferred_locations
def _get_cached_metadata(self) -> Dict[ObjectRef, "_ObjectMetadata"]:
if self._cached_object_meta is None:
# This call is pretty fast for owned objects (~5k/s), so we don't need to
# batch it for now.
meta = ray.experimental.get_local_object_locations(self.block_refs)
# Extract locations
object_metas: Dict[ObjectRef, _ObjectMetadata] = {
ref: _ObjectMetadata(
size=meta[ref]["object_size"],
spilled=meta[ref]["did_spill"],
locs=meta[ref]["node_ids"],
)
for ref in self.block_refs
}
# NOTE: We're working around object being immutable to update cached
# values (safe)
object.__setattr__(self, "_cached_object_meta", object_metas)
return self._cached_object_meta
def slice(self, needed_rows: int) -> Tuple["RefBundle", "RefBundle"]:
"""Slice a Ref Bundle into the first bundle containing the first `needed_rows` rows and the remaining bundle containing the remaining rows.
Args:
needed_rows: Number of rows to take from the head of the bundle.
Returns:
A tuple of (sliced_bundle, remaining_bundle). The needed rows must be less than the number of rows in the bundle.
"""
assert needed_rows > 0, "needed_rows must be positive."
assert (
self.num_rows() is not None
), "Cannot slice a RefBundle with unknown number of rows."
assert (
needed_rows < self.num_rows()
), f"To slice a RefBundle, the number of requested rows must be less than the number of rows in the bundle. Requested {needed_rows} rows but bundle only has {self.num_rows()} rows."
block_slices = []
for metadata, block_slice in zip(self.metadata, self.slices):
if block_slice is None:
# None represents a full block, convert to explicit BlockSlice
assert (
metadata.num_rows is not None
), "Cannot derive block slice for a RefBundle with unknown block row counts."
block_slices.append(
BlockSlice(start_offset=0, end_offset=metadata.num_rows)
)
else:
block_slices.append(block_slice)
consumed_blocks: List[BlockEntry] = []
consumed_slices: List[BlockSlice] = []
remaining_blocks: List[BlockEntry] = []
remaining_slices: List[BlockSlice] = []
rows_to_take = needed_rows
for entry, block_slice in zip(self.blocks, block_slices):
block_rows = block_slice.num_rows
if rows_to_take >= block_rows:
consumed_blocks.append(entry)
consumed_slices.append(block_slice)
rows_to_take -= block_rows
else:
if rows_to_take == 0:
remaining_blocks.append(entry)
remaining_slices.append(block_slice)
continue
consume_slice = BlockSlice(
start_offset=block_slice.start_offset,
end_offset=block_slice.start_offset + rows_to_take,
)
consumed_blocks.append(entry)
consumed_slices.append(consume_slice)
leftover_rows = block_rows - rows_to_take
if leftover_rows > 0:
remainder_slice = BlockSlice(
start_offset=consume_slice.end_offset,
end_offset=block_slice.end_offset,
)
remaining_blocks.append(entry)
remaining_slices.append(remainder_slice)
rows_to_take = 0
sliced_bundle = RefBundle(
blocks=tuple(consumed_blocks),
schema=self.schema,
owns_blocks=False,
slices=tuple(consumed_slices) if consumed_slices else None,
)
remaining_bundle = RefBundle(
blocks=tuple(remaining_blocks),
schema=self.schema,
owns_blocks=False,
slices=tuple(remaining_slices) if remaining_slices else None,
)
return sliced_bundle, remaining_bundle
@classmethod
def merge_ref_bundles(cls, bundles: Iterable["RefBundle"]) -> "RefBundle":
"""Merge multiple RefBundles into a single RefBundle.
Args:
bundles: An iterable of RefBundles to merge.
Returns:
A single RefBundle containing all blocks from the input bundles.
owns_blocks is True only if all input bundles own their blocks.
schema is the first non-empty schema found.
"""
bundles = list(bundles)
if not bundles:
return cls(blocks=(), owns_blocks=True, schema=None)
merged_blocks = list(
itertools.chain.from_iterable(bundle.blocks for bundle in bundles)
)
merged_slices = list(
itertools.chain.from_iterable(bundle.slices for bundle in bundles)
)
# Ray Data uses the `owns_blocks` flag to determine if the system can eagerly
# destroy blocks when they're no longer needed. To be safe, we only set this
# to True if all input bundles own their blocks.
owns_blocks = all(bundle.owns_blocks for bundle in bundles)
# TODO: Reconcile the schemas rather than taking the first non-empty schema.
schema = _take_first_non_empty_schema(bundle.schema for bundle in bundles)
return cls(
blocks=tuple(merged_blocks),
schema=schema,
owns_blocks=owns_blocks,
slices=merged_slices,
)
def __eq__(self, other: "RefBundle"):
if self is other:
return True
elif not isinstance(other, RefBundle):
return False
return (
self.blocks == other.blocks
and self.slices == other.slices
# NOTE: We're establishing a requirement of schemas for `RefBundle`
# to be exactly the same object for it to be considered equal.
#
# This is necessary to avoid a full schema equality check that
# is computationally intensive.
and self.schema is other.schema
and self.owns_blocks == other.owns_blocks
and self.output_split_idx == other.output_split_idx
)
def __hash__(self) -> int:
return hash(
(
# Only hash block refs
*[entry.ref for entry in self.blocks],
*self.slices,
# Check out comment in ``__eq__``
id(self.schema),
self.owns_blocks,
self.output_split_idx,
)
)
def __len__(self) -> int:
return len(self.blocks)
def __str__(self) -> str:
lines = [
f"RefBundle({len(self.blocks)} blocks,",
f" {self.num_rows()} rows,",
f" schema={self.schema},",
f" owns_blocks={self.owns_blocks},",
" blocks=(",
]
# Loop through each block and show details
for i, (entry, block_slice) in enumerate(zip(self.blocks, self.slices)):
metadata = entry.metadata
row_str = (
f"{metadata.num_rows} rows"
if metadata.num_rows is not None
else "unknown rows"
)
bytes_str = f"{metadata.size_bytes} bytes"
slice_str = (
f"slice={block_slice}"
if block_slice is not None
else "slice=None (full block)"
)
lines.append(f" {i}: {row_str}, {bytes_str}, {slice_str}")
lines.append(" )")
lines.append(")")
return "\n".join(lines)
@dataclass
class _ObjectMetadata:
# Object size in bytes
size: int
# Flag whether object has been spilled
spilled: bool
# List of nodes object exists on
locs: List[NodeIdStr] = None
def _ref_bundles_iterator_to_block_refs_list(
ref_bundles: Iterator[RefBundle],
) -> List[ObjectRef[Block]]:
"""Convert an iterator of RefBundles to a list of Block object references."""
return [
block_ref for ref_bundle in ref_bundles for block_ref in ref_bundle.block_refs
]
def _iter_sliced_blocks(
blocks: Iterable[Block],
slices: List[Optional[BlockSlice]],
) -> Iterator[Block]:
blocks_list = list(blocks)
for block, block_slice in zip(blocks_list, slices):
if block_slice is None:
# None represents a full block - yield it as is
yield block
else:
accessor = BlockAccessor.for_block(block)
start = block_slice.start_offset
end = block_slice.end_offset
assert start <= end, "start must be less than end"
assert start >= 0, "start must be non-negative"
assert (
end <= accessor.num_rows()
), "end must be less than or equal to the number of rows in the block"
yield accessor.slice(start, end, copy=False)
python/ray/data/_internal/execution/interfaces/task_context.py
+88
View File
@@ -0,0 +1,88 @@
import contextlib
import threading
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional
if TYPE_CHECKING:
from ray.data._internal.execution.operators.map_transformer import MapTransformer
from ray.data._internal.progress.base_progress import BaseProgressBar
_thread_local = threading.local()
@dataclass
class TaskContext:
"""This describes the information of a task running block transform."""
# The index of task. Each task has a unique task index within the same
# operator.
task_idx: int
# Name of the operator that this task belongs to.
op_name: str
# The dictionary of sub progress bar to update. The key is name of sub progress
# bar. Note this is only used on driver side.
# TODO(chengsu): clean it up from TaskContext with new optimizer framework.
sub_progress_bar_dict: Optional[Dict[str, "BaseProgressBar"]] = None
# NOTE(hchen): `upstream_map_transformer` and `upstream_map_ray_remote_args`
# are only used for `RandomShuffle`. DO NOT use them for other operators.
# Ideally, they should be handled by the optimizer, and should be transparent
# to the specific operators.
# But for `RandomShuffle`, the AllToAllOperator doesn't do the shuffle itself.
# It uses `ExchangeTaskScheduler` to launch new tasks to do the shuffle.
# That's why we need to pass them to `ExchangeTaskScheduler`.
# TODO(hchen): Use a physical operator to do the shuffle directly.
# The underlying function called in a MapOperator; this is used when fusing
# an AllToAllOperator with an upstream MapOperator.
upstream_map_transformer: Optional["MapTransformer"] = None
# The Ray remote arguments of the fused upstream MapOperator.
# This should be set if upstream_map_transformer is set.
upstream_map_ray_remote_args: Optional[Dict[str, Any]] = None
# Override of the target max-block-size for the task
target_max_block_size_override: Optional[int] = None
# Additional keyword arguments passed to the task.
kwargs: Dict[str, Any] = field(default_factory=dict)
@classmethod
def get_current(cls) -> Optional["TaskContext"]:
"""Get the TaskContext for the current thread.
Returns None if no TaskContext has been set.
"""
return getattr(_thread_local, "task_context", None)
@classmethod
def set_current(cls, context: "TaskContext") -> None:
"""Set the TaskContext for the current thread.
Args:
context: The TaskContext instance to set for this thread
"""
_thread_local.task_context = context
@classmethod
def reset_current(cls):
"""Clear the current thread's TaskContext."""
if hasattr(_thread_local, "task_context"):
delattr(_thread_local, "task_context")
@classmethod
@contextlib.contextmanager
def current(cls, context: "TaskContext") -> Iterator["TaskContext"]:
"""Sets this TaskContext as current for the scope
of the context block and resets it on exit.
"""
cls.set_current(context)
try:
yield context
finally:
cls.reset_current()
python/ray/data/_internal/execution/interfaces/transform_fn.py
+14
View File
@@ -0,0 +1,14 @@
from typing import Callable, List, Tuple
from .ref_bundle import RefBundle
from .task_context import TaskContext
from ray.data._internal.stats import StatsDict
# Result type of AllToAllTransformFn.
AllToAllTransformFnResult = Tuple[List[RefBundle], StatsDict]
# Block transform function applied in AllToAllOperator.
AllToAllTransformFn = Callable[
[List[RefBundle], TaskContext],
AllToAllTransformFnResult,
]