wehub-resource-sync
287 lines
12 KiB
Python
287 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import abc
|
|
from collections import deque
|
|
from typing import TYPE_CHECKING, Any, Deque, List, Optional, Tuple
|
|
|
|
from typing_extensions import override
|
|
|
|
from ray.data._internal.execution.bundle_queue import (
|
|
BaseBundleQueue,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data._internal.execution.interfaces import RefBundle
|
|
|
|
|
|
class RebundlingStrategy(abc.ABC):
|
|
"""Base class for strategies describing how to rebundle queues."""
|
|
|
|
@abc.abstractmethod
|
|
def can_build_ready_bundle(self, num_pending_rows: int) -> bool:
|
|
"""Signifies whether we can build a ready bundle. A ready bundle is a bundle
|
|
that will be returned from `get_next()` calls. Pending bundles merge into Ready bundles."""
|
|
...
|
|
|
|
@abc.abstractmethod
|
|
def rows_needed_from_last_pending_bundle(
|
|
self, total_pending_rows: int, last_pending_bundle: RefBundle
|
|
) -> int:
|
|
"""Used to determine how to rebundle and slice an existing bundle.
|
|
|
|
Args:
|
|
total_pending_rows: The number of rows in a batch of pending bundles that will be merged to form
|
|
a ready bundle, including the last_pending_bundle.
|
|
last_pending_bundle: The last pending bundles in that batch ^. The term *last* means the bundle that caused
|
|
`can_build_ready_bundle(num_pending_rows)` to be `True` for the first time.
|
|
|
|
Returns:
|
|
The # of rows needed from the last pending bundle. This should be > 0, unless bundle.num_rows() is None.
|
|
"""
|
|
...
|
|
|
|
|
|
class EstimateSize(RebundlingStrategy):
|
|
"""Rebundles RefBundles to get them close to a particular number of rows."""
|
|
|
|
def __init__(self, min_rows_per_bundle: Optional[int]):
|
|
"""Creates a strategy for combining bundles close to a particular row count.
|
|
|
|
Args:
|
|
min_rows_per_bundle: The target number of rows per bundle. Note that we
|
|
bundle up to this target, but only exceed it if not doing so would
|
|
result in an empty bundle. If None, this behaves like a normal queue.
|
|
"""
|
|
|
|
assert (
|
|
min_rows_per_bundle is None or min_rows_per_bundle >= 0
|
|
), "Min rows per bundle has to be non-negative"
|
|
|
|
self._min_rows_per_bundle: Optional[int] = min_rows_per_bundle
|
|
|
|
@override
|
|
def can_build_ready_bundle(self, num_pending_rows: int) -> bool:
|
|
return num_pending_rows > 0 and (
|
|
self._min_rows_per_bundle is None
|
|
or num_pending_rows >= self._min_rows_per_bundle
|
|
)
|
|
|
|
@override
|
|
def rows_needed_from_last_pending_bundle(
|
|
self, total_pending_rows: int, last_pending_bundle: RefBundle
|
|
) -> int:
|
|
"""Returns all the rows in the pending bundle, since we only care about an estimate"""
|
|
return last_pending_bundle.num_rows() or 0
|
|
|
|
|
|
class ExactMultipleSize(RebundlingStrategy):
|
|
def __init__(self, target_num_rows_per_block: int):
|
|
assert (
|
|
target_num_rows_per_block > 0
|
|
), "target_num_rows_per_block must be positive for streaming repartition."
|
|
self._target_num_rows = target_num_rows_per_block
|
|
|
|
@override
|
|
def can_build_ready_bundle(self, num_pending_rows: int) -> bool:
|
|
return num_pending_rows >= self._target_num_rows
|
|
|
|
@override
|
|
def rows_needed_from_last_pending_bundle(
|
|
self, total_pending_rows: int, last_pending_bundle: RefBundle
|
|
) -> int:
|
|
"""Returns an exact MULTIPLE of target_num_rows from the last pending bundle."""
|
|
pending_rows = last_pending_bundle.num_rows() or 0
|
|
assert total_pending_rows - pending_rows < self._target_num_rows, (
|
|
f"Total pending rows={total_pending_rows} should be less than target_num_rows={self._target_num_rows}, "
|
|
"because last_pending_bundle should trigger building ready bundles"
|
|
)
|
|
extra_rows = total_pending_rows % self._target_num_rows
|
|
assert extra_rows < pending_rows
|
|
return pending_rows - extra_rows
|
|
|
|
|
|
"""**For `ExactMultipleSize` strategy ONLY**
|
|
|
|
Streaming repartition builds fixed-size outputs from a stream of inputs.
|
|
|
|
We construct batches here to produce exactly sized outputs from arbitrary [start, end) slices across input blocks.
|
|
The task builder submits a map task only after the total number of rows accumulated across pending blocks reaches
|
|
target num rows (except during the final flush, which may emit a smaller tail block). This allows us to create
|
|
target-sized batches without materializing entire large blocks on the driver.
|
|
|
|
Detailed Implementation:
|
|
1. When a new bundle arrives, buffer it in the pending list.
|
|
2. Whenever the total number of rows in the pending bundles reaches the target row count, try to build a ready bundle.
|
|
3. Determine the slice needed from the final bundle so the ready bundle holds an exact multiple of the target rows,
|
|
and add the remaining bundle to the pending bundles for the next iteration.
|
|
4. Submit that ready bundle to a remote map task; the task slices each block according to the slice metadata stored
|
|
in the RefBundle (the bundle now contains n * target rows for n ≥ 1).
|
|
5. We configured the `OutputBlockSizeOption.target_num_rows_per_block` to the target number of rows per block in
|
|
plan_streaming_repartition_op so the output buffer further splits the n * target rows into n blocks of exactly
|
|
the target size.
|
|
6. Once upstream input is exhausted, flush any leftover pending bundles and repeat steps 1-5 for the tail.
|
|
7. The resulting blocks have lengths `[target, …, target, (total_rows % target)]`; ordering isn't guaranteed, but the
|
|
remainder block should appear near the end.
|
|
|
|
"""
|
|
|
|
|
|
class RebundleQueue(BaseBundleQueue):
|
|
"""Incrementally builds task inputs to produce multiples of target-sized outputs."""
|
|
|
|
def __init__(self, strategy: RebundlingStrategy):
|
|
super().__init__()
|
|
|
|
self._strategy = strategy
|
|
self._pending_bundles: Deque[RefBundle] = deque()
|
|
self._ready_bundles: Deque[RefBundle] = deque()
|
|
|
|
self._curr_consumed_bundles: List[RefBundle] = []
|
|
# The original bundles that formed a ready bundle
|
|
self._consumed_bundles_list: Deque[List[RefBundle]] = deque()
|
|
self._total_pending_rows: int = 0
|
|
|
|
def _merge_bundles(self):
|
|
"""Combine *ALL* pending_bundles into a single, ready bundle."""
|
|
from ray.data._internal.execution.interfaces import RefBundle
|
|
|
|
merged_bundle = RefBundle.merge_ref_bundles(self._pending_bundles)
|
|
# Update the metrics
|
|
self._ready_bundles.append(merged_bundle)
|
|
self._on_enqueue_bundle(merged_bundle)
|
|
|
|
# Clear the pending queue since all bundles have been processed
|
|
for bundle in self._pending_bundles:
|
|
self._on_dequeue_bundle(bundle)
|
|
self._pending_bundles.clear()
|
|
self._total_pending_rows = 0
|
|
|
|
def _try_build_ready_bundle(self, flush_remaining: bool) -> int:
|
|
"""Attempts to build a ready bundle from a list of pending bundles by:
|
|
|
|
- Checking the threshold to build a ready bundle defined by `RebundlingStrategy`
|
|
- Appropiately keeping track of queue metrics
|
|
|
|
Returns `True` if ready bundle built, otherwise `False`
|
|
"""
|
|
|
|
ready_bundles_built: int = 0
|
|
if self._pending_bundles and self._strategy.can_build_ready_bundle(
|
|
self._total_pending_rows
|
|
):
|
|
last_pending_bundle = self._pending_bundles.pop()
|
|
|
|
# We now know `pending_bundle` is the bundle that enabled us to
|
|
# build a ready bundle. Therefore, we may need to slice the bundle.
|
|
rows_needed = self._strategy.rows_needed_from_last_pending_bundle(
|
|
total_pending_rows=self._total_pending_rows,
|
|
last_pending_bundle=last_pending_bundle,
|
|
)
|
|
assert rows_needed > 0, (
|
|
"A refbundle has zero row-count but triggered building a ready bundle"
|
|
"This is a bug in the Ray Data code."
|
|
)
|
|
remaining_bundle: Optional[RefBundle] = None
|
|
last_num_rows = last_pending_bundle.num_rows() or 0
|
|
if rows_needed < last_num_rows:
|
|
sliced_bundle, remaining_bundle = last_pending_bundle.slice(rows_needed)
|
|
# The original bundle was enqueued in add(). We need to dequeue it
|
|
# and enqueue the sliced portion, since _merge_bundles will dequeue
|
|
# sliced_bundle (which has different metrics than the original).
|
|
self._on_dequeue_bundle(last_pending_bundle)
|
|
self._on_enqueue_bundle(sliced_bundle)
|
|
self._pending_bundles.append(sliced_bundle)
|
|
else:
|
|
assert rows_needed == last_num_rows
|
|
self._pending_bundles.append(last_pending_bundle)
|
|
|
|
self._merge_bundles()
|
|
ready_bundles_built += 1
|
|
|
|
if remaining_bundle is not None:
|
|
# Add back remaining sliced bundle that was not included to build
|
|
# a ready bundle.
|
|
self._pending_bundles.appendleft(remaining_bundle)
|
|
self._total_pending_rows += remaining_bundle.num_rows() or 0
|
|
self._on_enqueue_bundle(remaining_bundle)
|
|
|
|
# If we're flushing and have leftover bundles, convert them to a ready bundle.
|
|
# Note: add() eagerly calls _try_build_ready_bundle after every insertion, so
|
|
# pending rows are always below the threshold when finalize() is called. This
|
|
# means at most one ready bundle is built per call (only the flush path fires).
|
|
if flush_remaining and self._pending_bundles:
|
|
self._merge_bundles()
|
|
ready_bundles_built += 1
|
|
|
|
return ready_bundles_built
|
|
|
|
@override
|
|
def add(self, bundle: RefBundle, **kwargs: Any):
|
|
from ray.data._internal.execution.interfaces import RefBundle
|
|
|
|
num_rows = bundle.num_rows() or 0
|
|
if num_rows == 0:
|
|
if self._pending_bundles:
|
|
last = self._pending_bundles.pop()
|
|
self._on_dequeue_bundle(last)
|
|
merged = RefBundle.merge_ref_bundles([last, bundle])
|
|
self._pending_bundles.append(merged)
|
|
self._on_enqueue_bundle(merged)
|
|
else:
|
|
self._pending_bundles.append(bundle)
|
|
self._on_enqueue_bundle(bundle)
|
|
return
|
|
self._total_pending_rows += num_rows
|
|
self._pending_bundles.append(bundle)
|
|
self._on_enqueue_bundle(bundle)
|
|
self._curr_consumed_bundles.append(bundle)
|
|
ready_bundles_built = self._try_build_ready_bundle(flush_remaining=False)
|
|
if ready_bundles_built > 0:
|
|
assert ready_bundles_built == 1
|
|
self._consumed_bundles_list.append(self._curr_consumed_bundles)
|
|
self._curr_consumed_bundles = []
|
|
|
|
@override
|
|
def has_next(self) -> bool:
|
|
return len(self._ready_bundles) > 0
|
|
|
|
@override
|
|
def _get_next_inner(self) -> RefBundle:
|
|
if not self.has_next():
|
|
raise ValueError("You can't pop from empty queue")
|
|
ready_bundle = self._ready_bundles.popleft()
|
|
# discard the original bundle
|
|
self._consumed_bundles_list.popleft()
|
|
return ready_bundle
|
|
|
|
def get_next_with_original(self) -> Tuple[RefBundle, List[RefBundle]]:
|
|
if not self.has_next():
|
|
raise ValueError("You can't pop from empty queue")
|
|
ready_bundle = self._ready_bundles.popleft()
|
|
self._on_dequeue_bundle(ready_bundle)
|
|
consumed_bundle = self._consumed_bundles_list.popleft()
|
|
return ready_bundle, consumed_bundle
|
|
|
|
@override
|
|
def peek_next(self) -> Optional[RefBundle]:
|
|
if not self.has_next():
|
|
return None
|
|
return self._ready_bundles[0]
|
|
|
|
@override
|
|
def finalize(self, **kwargs: Any):
|
|
if len(self._pending_bundles) > 0:
|
|
ready_bundles_built = self._try_build_ready_bundle(flush_remaining=True)
|
|
assert ready_bundles_built == 1
|
|
self._consumed_bundles_list.append(self._curr_consumed_bundles)
|
|
self._curr_consumed_bundles = []
|
|
|
|
@override
|
|
def clear(self):
|
|
self._reset_metrics()
|
|
self._pending_bundles.clear()
|
|
self._ready_bundles.clear()
|
|
self._curr_consumed_bundles.clear()
|
|
self._consumed_bundles_list.clear()
|
|
self._total_pending_rows = 0
|