wehub-resource-sync
482 lines
19 KiB
Python
482 lines
19 KiB
Python
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)
|