chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,375 @@
|
||||
import warnings
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
||||
from ray.data._internal.arrow_ops import transform_pyarrow
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import try_combine_chunked_columns
|
||||
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
||||
from ray.data._internal.execution.util import memory_string
|
||||
from ray.data._internal.util import get_total_obj_store_mem_on_node
|
||||
from ray.data.block import Block, BlockAccessor
|
||||
from ray.util import log_once
|
||||
|
||||
# Delay compaction until the shuffle buffer has reached this ratio over the min
|
||||
# shuffle buffer size. Setting this to 1 minimizes memory usage, at the cost of
|
||||
# frequent compactions. Setting this to higher values increases memory usage but
|
||||
# reduces compaction frequency.
|
||||
SHUFFLE_BUFFER_COMPACTION_RATIO = 1.5
|
||||
|
||||
# Ratio of remaining compacted rows to shuffle_buffer_min_size at which
|
||||
# compaction (and re-shuffling of indices) is triggered. Experiments show 0.5
|
||||
# is a good trade-off between throughput and randomness.
|
||||
SHUFFLE_BUFFER_COMPACTION_THRESHOLD = 0.5
|
||||
|
||||
|
||||
class BatcherInterface:
|
||||
def add(self, block: Block):
|
||||
"""Add a block to the block buffer.
|
||||
|
||||
Args:
|
||||
block: Block to add to the block buffer.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def done_adding(self) -> bool:
|
||||
"""Indicate to the batcher that no more blocks will be added to the buffer."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def has_batch(self) -> bool:
|
||||
"""Whether this Batcher has any full batches."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def has_any(self) -> bool:
|
||||
"""Whether this Batcher has any data."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def next_batch(self) -> Block:
|
||||
"""Get the next batch from the block buffer.
|
||||
|
||||
Returns:
|
||||
A batch represented as a Block.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class Batcher(BatcherInterface):
|
||||
"""Chunks blocks into batches."""
|
||||
|
||||
# Implementation Note: When there are multiple batches per block, this batcher will
|
||||
# slice off and return each batch and add the remaining block back to the buffer
|
||||
# instead of optimally slicing and returning all batches from the block at once.
|
||||
# This will result in extra (and nested) block slicing. However, since slices are
|
||||
# zero-copy views, we sacrifice what should be a small performance hit for better
|
||||
# readability.
|
||||
|
||||
def __init__(self, batch_size: Optional[int], ensure_copy: bool = False):
|
||||
"""
|
||||
Construct a batcher that yields batches of batch_sizes rows.
|
||||
|
||||
Args:
|
||||
batch_size: The size of batches to yield.
|
||||
ensure_copy: Whether batches are always copied from the underlying base
|
||||
blocks (not zero-copy views).
|
||||
"""
|
||||
self._batch_size = batch_size
|
||||
self._buffer = []
|
||||
self._buffer_size = 0
|
||||
self._done_adding = False
|
||||
self._ensure_copy = ensure_copy
|
||||
|
||||
def add(self, block: Block):
|
||||
"""Add a block to the block buffer.
|
||||
|
||||
Note empty block is not added to buffer.
|
||||
|
||||
Args:
|
||||
block: Block to add to the block buffer.
|
||||
"""
|
||||
if BlockAccessor.for_block(block).num_rows() > 0:
|
||||
self._buffer.append(block)
|
||||
self._buffer_size += BlockAccessor.for_block(block).num_rows()
|
||||
|
||||
def done_adding(self) -> bool:
|
||||
"""Indicate to the batcher that no more blocks will be added to the batcher."""
|
||||
self._done_adding = True
|
||||
|
||||
def has_batch(self) -> bool:
|
||||
"""Whether this Batcher has any full batches."""
|
||||
return self.has_any() and (
|
||||
self._batch_size is None or self._buffer_size >= self._batch_size
|
||||
)
|
||||
|
||||
def has_any(self) -> bool:
|
||||
"""Whether this Batcher has any data."""
|
||||
return self._buffer_size > 0
|
||||
|
||||
def next_batch(self) -> Block:
|
||||
"""Get the next batch from the block buffer.
|
||||
|
||||
Returns:
|
||||
A batch represented as a Block.
|
||||
"""
|
||||
assert self.has_batch() or (self._done_adding and self.has_any())
|
||||
needs_copy = self._ensure_copy
|
||||
# If no batch size, short-circuit.
|
||||
if self._batch_size is None:
|
||||
assert len(self._buffer) == 1
|
||||
block = self._buffer[0]
|
||||
if needs_copy:
|
||||
# Copy block if needing to ensure fresh batch copy.
|
||||
block = BlockAccessor.for_block(block)
|
||||
block = block.slice(0, block.num_rows(), copy=True)
|
||||
self._buffer = []
|
||||
self._buffer_size = 0
|
||||
return block
|
||||
output = DelegatingBlockBuilder()
|
||||
leftover = []
|
||||
needed = self._batch_size
|
||||
for block in self._buffer:
|
||||
accessor = BlockAccessor.for_block(block)
|
||||
if needed <= 0:
|
||||
# We already have a full batch, so add this block to
|
||||
# the leftovers.
|
||||
leftover.append(block)
|
||||
elif accessor.num_rows() <= needed:
|
||||
output.add_block(accessor.to_block())
|
||||
needed -= accessor.num_rows()
|
||||
else:
|
||||
# Try de-fragmenting table in case its columns
|
||||
# have too many chunks (potentially hindering performance of
|
||||
# subsequent slicing operation)
|
||||
if isinstance(accessor, ArrowBlockAccessor):
|
||||
accessor = BlockAccessor.for_block(
|
||||
transform_pyarrow.try_combine_chunked_columns(
|
||||
block, min_chunks_to_combine=1
|
||||
)
|
||||
)
|
||||
|
||||
# We only need part of the block to fill out a batch.
|
||||
output.add_block(accessor.slice(0, needed, copy=False))
|
||||
# Add the rest of the block to the leftovers.
|
||||
leftover.append(accessor.slice(needed, accessor.num_rows(), copy=False))
|
||||
needed = 0
|
||||
|
||||
# Move the leftovers into the block buffer so they're the first
|
||||
# blocks consumed on the next batch extraction.
|
||||
self._buffer = leftover
|
||||
self._buffer_size -= self._batch_size
|
||||
needs_copy = needs_copy and not output.will_build_yield_copy()
|
||||
batch = output.build()
|
||||
if needs_copy:
|
||||
# Need to ensure that the batch is a fresh copy.
|
||||
batch = BlockAccessor.for_block(batch)
|
||||
batch = batch.slice(0, batch.num_rows(), copy=True)
|
||||
return batch
|
||||
|
||||
|
||||
class ShufflingBatcher(BatcherInterface):
|
||||
"""Chunks blocks into shuffled batches, using a local in-memory shuffle buffer.
|
||||
|
||||
Uses an **incremental index** approach: on each compaction a permutation
|
||||
array is generated over the buffer rows, and each ``next_batch()`` call
|
||||
gathers a small slice of that permutation via ``take``.
|
||||
|
||||
Properties of this approach:
|
||||
|
||||
* **Memory-efficient** -- the data buffer is kept as-is; only a
|
||||
lightweight int64 index array is allocated on top.
|
||||
* **Smooth per-batch latency** -- each ``take`` operates on a small
|
||||
slice of indices, so per-batch work is short and uniform, making it
|
||||
easy to hide behind prefetch threads.
|
||||
|
||||
Example with ``batch_size=3`` and a 9-row buffer::
|
||||
|
||||
buffer: [A, B, C, D, E, F, G, H, I]
|
||||
indices: [4, 7, 1, 0, 8, 3, 6, 2, 5] # random permutation
|
||||
|
||||
next_batch() -> take([4, 7, 1]) -> [E, H, B] # batch_head 0 -> 3
|
||||
next_batch() -> take([0, 8, 3]) -> [A, I, D] # batch_head 3 -> 6
|
||||
next_batch() -> take([6, 2, 5]) -> [G, C, F] # batch_head 6 -> 9
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size: Optional[int],
|
||||
shuffle_buffer_min_size: int,
|
||||
shuffle_seed: Optional[int] = None,
|
||||
):
|
||||
"""Constructs a random-shuffling block batcher.
|
||||
|
||||
Args:
|
||||
batch_size: Record batch size.
|
||||
shuffle_buffer_min_size: Minimum number of rows that must be in the local
|
||||
in-memory shuffle buffer in order to yield a batch. When there are no
|
||||
more rows to be added to the buffer, the number of rows in the buffer
|
||||
*will* decrease below this value while yielding the remaining batches,
|
||||
and the final batch may have less than ``batch_size`` rows. Increasing
|
||||
this will improve the randomness of the shuffle but may increase the
|
||||
latency to the first batch.
|
||||
shuffle_seed: The seed to use for the local random shuffle.
|
||||
"""
|
||||
if batch_size is None:
|
||||
raise ValueError("Must specify a batch_size if using a local shuffle.")
|
||||
self._batch_size = batch_size
|
||||
self._rng = np.random.default_rng(shuffle_seed)
|
||||
if shuffle_buffer_min_size < batch_size:
|
||||
# Round it up internally to `batch_size` since our algorithm requires it.
|
||||
# This is harmless since it only offers extra randomization.
|
||||
shuffle_buffer_min_size = batch_size
|
||||
self._shuffle_buffer_min_size = shuffle_buffer_min_size
|
||||
|
||||
self._min_rows_to_yield_batch = max(
|
||||
1, int(shuffle_buffer_min_size * SHUFFLE_BUFFER_COMPACTION_THRESHOLD)
|
||||
)
|
||||
self._min_rows_to_trigger_compaction = int(
|
||||
shuffle_buffer_min_size * SHUFFLE_BUFFER_COMPACTION_RATIO
|
||||
)
|
||||
self._builder = DelegatingBlockBuilder()
|
||||
self._shuffle_buffer: Block = None
|
||||
self._shuffled_indices: Optional[np.ndarray] = None
|
||||
self._batch_head = 0
|
||||
self._done_adding = False
|
||||
|
||||
self._total_object_store_nbytes = get_total_obj_store_mem_on_node()
|
||||
self._total_num_rows_added = 0
|
||||
self._total_nbytes_added = 0
|
||||
|
||||
def add(self, block: Block):
|
||||
"""Add a block to the shuffle buffer.
|
||||
|
||||
Note empty block is not added to buffer.
|
||||
|
||||
Args:
|
||||
block: Block to add to the shuffle buffer.
|
||||
"""
|
||||
# Because Arrow tables are memory mapped, blocks in the builder reside in object
|
||||
# store memory and not local heap memory. So, if you specify a large buffer size
|
||||
# and there isn't enough object store memory on the node, you encounter
|
||||
# spilling.
|
||||
if (
|
||||
self._estimated_min_nbytes_in_buffers is not None
|
||||
and self._estimated_min_nbytes_in_buffers > self._total_object_store_nbytes
|
||||
and log_once("shuffle_buffer_mem_warning")
|
||||
):
|
||||
warnings.warn(
|
||||
"The node you're iterating on has "
|
||||
f"{memory_string(self._total_object_store_nbytes)} object "
|
||||
"store memory, but the shuffle buffer is estimated to use "
|
||||
f"{memory_string(self._estimated_min_nbytes_in_buffers)}. If you don't "
|
||||
f"decrease the shuffle buffer size from "
|
||||
f"{self._shuffle_buffer_min_size} rows, you might encounter spilling."
|
||||
)
|
||||
|
||||
block_accessor = BlockAccessor.for_block(block)
|
||||
if block_accessor.num_rows() > 0:
|
||||
self._builder.add_block(block)
|
||||
self._total_num_rows_added += block_accessor.num_rows()
|
||||
self._total_nbytes_added += block_accessor.size_bytes()
|
||||
|
||||
@property
|
||||
def _average_row_nbytes(self) -> Optional[int]:
|
||||
"""Return the average number of bytes per row added to this batcher."""
|
||||
return (
|
||||
self._total_nbytes_added // self._total_num_rows_added
|
||||
if self._total_num_rows_added > 0
|
||||
else None
|
||||
)
|
||||
|
||||
@property
|
||||
def _estimated_min_nbytes_in_buffers(self) -> Optional[int]:
|
||||
"""Return the estimated minimum number of bytes across all buffers.
|
||||
|
||||
This includes data in both the compacted and uncompacted buffers.
|
||||
"""
|
||||
if self._average_row_nbytes is None:
|
||||
return None
|
||||
|
||||
return self._average_row_nbytes * self._min_rows_to_trigger_compaction
|
||||
|
||||
def done_adding(self) -> bool:
|
||||
"""Indicate to the batcher that no more blocks will be added to the batcher.
|
||||
|
||||
No more blocks should be added to the batcher after calling this.
|
||||
"""
|
||||
self._done_adding = True
|
||||
|
||||
def has_any(self) -> bool:
|
||||
"""Whether this batcher has any data."""
|
||||
return self._num_rows() > 0
|
||||
|
||||
def has_batch(self) -> bool:
|
||||
"""Whether this batcher has any batches."""
|
||||
num_rows = self._num_rows()
|
||||
|
||||
if not self._done_adding:
|
||||
# Delay pulling of batches until the buffer is large enough in order to
|
||||
# amortize compaction overhead.
|
||||
return num_rows >= self._batch_size and (
|
||||
self._num_compacted_rows() >= self._min_rows_to_yield_batch
|
||||
or num_rows - self._batch_size >= self._min_rows_to_trigger_compaction
|
||||
)
|
||||
else:
|
||||
return num_rows >= self._batch_size
|
||||
|
||||
def _num_rows(self) -> int:
|
||||
"""Return the total number of rows that haven't been yielded yet.
|
||||
|
||||
This includes rows in both the compacted and uncompacted buffers.
|
||||
"""
|
||||
return self._num_compacted_rows() + self._num_uncompacted_rows()
|
||||
|
||||
def _num_compacted_rows(self) -> int:
|
||||
"""Return number of unyielded rows in the compacted buffer."""
|
||||
if self._shuffle_buffer is None:
|
||||
return 0
|
||||
return max(0, len(self._shuffled_indices) - self._batch_head)
|
||||
|
||||
def _num_uncompacted_rows(self) -> int:
|
||||
"""Return number of unyielded rows in the uncompacted buffer."""
|
||||
return self._builder.num_rows()
|
||||
|
||||
def next_batch(self) -> Block:
|
||||
"""Get the next shuffled batch from the shuffle buffer.
|
||||
|
||||
Returns:
|
||||
A batch represented as a Block.
|
||||
"""
|
||||
assert self.has_batch() or (self._done_adding and self.has_any())
|
||||
if self._num_uncompacted_rows() > 0 and (
|
||||
self._done_adding
|
||||
or self._num_compacted_rows() <= self._min_rows_to_yield_batch
|
||||
):
|
||||
if self._shuffle_buffer is not None and self._batch_head < len(
|
||||
self._shuffled_indices
|
||||
):
|
||||
remaining_indices = self._shuffled_indices[self._batch_head :]
|
||||
remaining_block = BlockAccessor.for_block(self._shuffle_buffer).take(
|
||||
remaining_indices
|
||||
)
|
||||
self._builder.add_block(remaining_block)
|
||||
self._shuffle_buffer = self._builder.build()
|
||||
|
||||
accessor = BlockAccessor.for_block(self._shuffle_buffer)
|
||||
if isinstance(accessor, ArrowBlockAccessor):
|
||||
self._shuffle_buffer = try_combine_chunked_columns(
|
||||
self._shuffle_buffer, min_chunks_to_combine=1
|
||||
)
|
||||
accessor = BlockAccessor.for_block(self._shuffle_buffer)
|
||||
|
||||
num_rows = accessor.num_rows()
|
||||
self._shuffled_indices = self._rng.permutation(num_rows)
|
||||
|
||||
self._builder = DelegatingBlockBuilder()
|
||||
self._batch_head = 0
|
||||
|
||||
assert self._shuffle_buffer is not None
|
||||
assert self._shuffled_indices is not None
|
||||
remaining = len(self._shuffled_indices) - self._batch_head
|
||||
batch_size = min(self._batch_size, remaining)
|
||||
batch_indices = self._shuffled_indices[
|
||||
self._batch_head : self._batch_head + batch_size
|
||||
]
|
||||
self._batch_head += batch_size
|
||||
return BlockAccessor.for_block(self._shuffle_buffer).take(batch_indices)
|
||||
Reference in New Issue
Block a user