chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .base import (
|
||||
BaseBundleQueue,
|
||||
QueueWithRemoval,
|
||||
)
|
||||
from .bundler import EstimateSize, ExactMultipleSize, RebundleQueue
|
||||
from .fifo import FIFOBundleQueue
|
||||
from .hash_link import HashLinkedQueue
|
||||
from .reordering import ReorderingBundleQueue
|
||||
from .thread_safe import ThreadSafeBundleQueue
|
||||
|
||||
|
||||
def create_bundle_queue() -> QueueWithRemoval:
|
||||
return HashLinkedQueue()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseBundleQueue",
|
||||
"create_bundle_queue",
|
||||
"HashLinkedQueue",
|
||||
"RebundleQueue",
|
||||
"EstimateSize",
|
||||
"ReorderingBundleQueue",
|
||||
"FIFOBundleQueue",
|
||||
"ExactMultipleSize",
|
||||
"QueueWithRemoval",
|
||||
"ThreadSafeBundleQueue",
|
||||
]
|
||||
@@ -0,0 +1,254 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Optional,
|
||||
)
|
||||
|
||||
from ray.data._internal.execution.util import memory_string
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.execution.interfaces import RefBundle
|
||||
|
||||
|
||||
class BundleQueue(abc.ABC):
|
||||
@abc.abstractmethod
|
||||
def estimate_size_bytes(self) -> int:
|
||||
"""Returns the estimated size in bytes of all bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def num_blocks(self) -> int:
|
||||
"""Returns the total # of blocks across all bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def num_bundles(self) -> int:
|
||||
"""Returns the total # of bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def num_rows(self) -> int:
|
||||
"""Return the total # of rows across all bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _on_enqueue_bundle(self, bundle: RefBundle):
|
||||
"""Hook called before a bundle is added to the queue."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _on_dequeue_bundle(self, bundle: RefBundle):
|
||||
"""Hook called after a bundle is removed from the queue."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _add_inner(self, bundle: RefBundle, **kwargs: Any):
|
||||
"""Add a bundle to the internal data structure."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
"""Remove and return the next bundle from the internal data structure."""
|
||||
...
|
||||
|
||||
def add(self, bundle: RefBundle, **kwargs: Any):
|
||||
"""Add a bundle to the tail(end) of the queue. Base classes should override
|
||||
the `_add_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Args:
|
||||
bundle: The bundle to add.
|
||||
**kwargs: Additional queue-specific parameters (e.g., `key` for ordered queues).
|
||||
This is used for `finalize`.
|
||||
"""
|
||||
self._on_enqueue_bundle(bundle)
|
||||
self._add_inner(bundle, **kwargs)
|
||||
|
||||
def get_next(self) -> RefBundle:
|
||||
"""Remove and return the head of the queue. Base classes should override
|
||||
the `_get_next_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Raises:
|
||||
IndexError: If the queue is empty.
|
||||
|
||||
Returns:
|
||||
The `RefBundle` at the head of the queue.
|
||||
"""
|
||||
bundle = self._get_next_inner()
|
||||
self._on_dequeue_bundle(bundle)
|
||||
return bundle
|
||||
|
||||
@abc.abstractmethod
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
"""Return the head of the queue. The only invariant is
|
||||
that the # of blocks, rows, and bytes must remain unchanged
|
||||
before and after this method call.
|
||||
|
||||
If queue.has_next() == False, return `None`.
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def has_next(self) -> bool:
|
||||
"""Check if the queue has a valid bundle."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def clear(self):
|
||||
"""Remove all bundles from the queue."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def finalize(self, **kwargs: Any):
|
||||
"""Signal that no additional bundles will be added to the bundler so
|
||||
the bundler can be finalized. The keys of kwargs provided should be the same
|
||||
as the ones passed into the `add()` method. This is important for ordered
|
||||
queues."""
|
||||
...
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Return the total # bundles."""
|
||||
return self.num_bundles()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return a string representation showing queue metrics."""
|
||||
nbytes = memory_string(self.estimate_size_bytes())
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"num_bundles={len(self)}, "
|
||||
f"num_blocks={self.num_blocks()}, "
|
||||
f"num_rows={self.num_rows()}, "
|
||||
f"nbytes={nbytes})"
|
||||
)
|
||||
|
||||
|
||||
class BaseBundleQueue(BundleQueue):
|
||||
"""Base class for storing bundles. Here and subclasses should adhere to the mental
|
||||
model that "first", "front", or "head" is the next bundle to be dequeued. Consequently,
|
||||
"last", "back", or "tail" is the last bundle to be dequeued.
|
||||
|
||||
Subclasses may choose to use the _on_dequeue_bundle and _on_enqueue_bundle methods to
|
||||
track num_blocks, nbytes, etc... If not, they should override those methods.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._nbytes: int = 0
|
||||
self._num_blocks: int = 0
|
||||
self._num_bundles: int = 0
|
||||
self._num_rows: int = 0
|
||||
|
||||
def _on_enqueue_bundle(self, bundle: RefBundle):
|
||||
self._nbytes += bundle.size_bytes()
|
||||
self._num_blocks += len(bundle.block_refs)
|
||||
self._num_bundles += 1
|
||||
self._num_rows += bundle.num_rows() or 0
|
||||
|
||||
def _on_dequeue_bundle(self, bundle: RefBundle):
|
||||
self._nbytes -= bundle.size_bytes()
|
||||
self._num_blocks -= len(bundle.block_refs)
|
||||
self._num_bundles -= 1
|
||||
self._num_rows -= bundle.num_rows() or 0
|
||||
|
||||
def estimate_size_bytes(self) -> int:
|
||||
"""Return the estimated size in bytes of all bundles."""
|
||||
return self._nbytes
|
||||
|
||||
def num_blocks(self) -> int:
|
||||
"""Return the total # of blocks across all bundles."""
|
||||
return self._num_blocks
|
||||
|
||||
def num_bundles(self) -> int:
|
||||
return self._num_bundles
|
||||
|
||||
def num_rows(self) -> int:
|
||||
"""Return the total # of rows across all bundles."""
|
||||
return self._num_rows
|
||||
|
||||
def _reset_metrics(self):
|
||||
self._num_rows = 0
|
||||
self._num_blocks = 0
|
||||
self._num_bundles = 0
|
||||
self._nbytes = 0
|
||||
|
||||
def add(self, bundle: RefBundle, **kwargs: Any):
|
||||
"""Add a bundle to the tail(end) of the queue. Base classes should override
|
||||
the `_add_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Args:
|
||||
bundle: The bundle to add.
|
||||
**kwargs: Additional queue-specific parameters (e.g., `key` for ordered queues).
|
||||
This is used for `finalize`.
|
||||
"""
|
||||
self._on_enqueue_bundle(bundle)
|
||||
self._add_inner(bundle, **kwargs)
|
||||
|
||||
def _add_inner(self, bundle: RefBundle, **kwargs: Any) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_next(self) -> RefBundle:
|
||||
"""Remove and return the head of the queue. Base classes should override
|
||||
the `_get_next_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Raises:
|
||||
IndexError: If the queue is empty.
|
||||
|
||||
Returns:
|
||||
The `RefBundle` at the head of the queue.
|
||||
"""
|
||||
bundle = self._get_next_inner()
|
||||
self._on_dequeue_bundle(bundle)
|
||||
return bundle
|
||||
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
"""Return the head of the queue. The only invariant is
|
||||
that the # of blocks, rows, and bytes must remain unchanged
|
||||
before and after this method call.
|
||||
|
||||
If queue.has_next() == False, return `None`.
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def has_next(self) -> bool:
|
||||
"""Check if the queue has a valid bundle."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def clear(self):
|
||||
"""Remove all bundles from the queue."""
|
||||
...
|
||||
|
||||
def finalize(self, **kwargs: Any):
|
||||
"""Signal that no additional bundles will be added to the bundler so
|
||||
the bundler can be finalized. The keys of kwargs provided should be the same
|
||||
as the ones passed into the `add()` method. This is important for ordered
|
||||
queues."""
|
||||
return None
|
||||
|
||||
|
||||
class QueueWithRemoval(BaseBundleQueue):
|
||||
"""Base class for storing bundles AND supporting remove(bundle)
|
||||
and contains(bundle) operations."""
|
||||
|
||||
def __contains__(self, bundle: RefBundle) -> bool:
|
||||
"""Return whether the key is in the queue."""
|
||||
...
|
||||
|
||||
def remove(self, bundle: RefBundle) -> RefBundle:
|
||||
"""Remove the specified bundle from the queue. If multiple instances exist, remove the first one."""
|
||||
bundle = self._remove_inner(bundle)
|
||||
self._on_dequeue_bundle(bundle)
|
||||
return bundle
|
||||
|
||||
def _remove_inner(self, bundle: RefBundle) -> RefBundle:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,286 @@
|
||||
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
|
||||
@@ -0,0 +1,66 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from typing import TYPE_CHECKING, Any, Deque, Iterator, List, Optional
|
||||
|
||||
from typing_extensions import override
|
||||
|
||||
from .base import BaseBundleQueue
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.execution.interfaces import RefBundle
|
||||
|
||||
|
||||
class FIFOBundleQueue(BaseBundleQueue):
|
||||
"""A bundle queue that follows fifo-policy. Conceptually
|
||||
[ ] <- [ ] <- [ ] ...
|
||||
^ where the leftmost is popped first
|
||||
|
||||
NOTE: Not thread-safe
|
||||
"""
|
||||
|
||||
def __init__(self, bundles: Optional[List[RefBundle]] = None):
|
||||
super().__init__()
|
||||
self._inner: Deque[RefBundle] = deque([])
|
||||
if bundles is not None:
|
||||
for bundle in bundles:
|
||||
self.add(bundle)
|
||||
|
||||
@override
|
||||
def _add_inner(self, bundle: RefBundle, **kwargs: Any):
|
||||
self._inner.append(bundle)
|
||||
|
||||
@override
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
if not self.has_next():
|
||||
raise ValueError(
|
||||
f"Popping from empty {self.__class__.__name__} is prohibited"
|
||||
)
|
||||
|
||||
bundle = self._inner.popleft()
|
||||
return bundle
|
||||
|
||||
@override
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
if not self.has_next():
|
||||
return None
|
||||
return self._inner[0]
|
||||
|
||||
@override
|
||||
def has_next(self) -> bool:
|
||||
return len(self) > 0
|
||||
|
||||
@override
|
||||
def finalize(self, **kwargs: Any):
|
||||
pass
|
||||
|
||||
def __iter__(self) -> Iterator[RefBundle]:
|
||||
yield from self._inner
|
||||
|
||||
def to_list(self) -> List[RefBundle]:
|
||||
return list(self._inner)
|
||||
|
||||
@override
|
||||
def clear(self):
|
||||
self._reset_metrics()
|
||||
self._inner.clear()
|
||||
@@ -0,0 +1,126 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict, deque
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Deque, Dict, Iterator, Optional
|
||||
|
||||
from typing_extensions import override
|
||||
|
||||
from .base import QueueWithRemoval
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.execution.interfaces import RefBundle
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Node:
|
||||
value: RefBundle
|
||||
next: Optional[_Node] = None
|
||||
prev: Optional[_Node] = None
|
||||
|
||||
|
||||
class HashLinkedQueue(QueueWithRemoval):
|
||||
"""A bundle queue that supports these operations quickly:
|
||||
- contains(bundle)
|
||||
- remove(bundle)
|
||||
|
||||
NOTE: Not thread-safe
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# We manually implement a linked list because we need to remove elements
|
||||
# efficiently, and Python's built-in data structures have O(n) removal time.
|
||||
self._head: Optional[_Node] = None
|
||||
self._tail: Optional[_Node] = None
|
||||
# We use a dictionary to keep track of the nodes corresponding to each bundle.
|
||||
# This allows us to remove a bundle from the queue in O(1) time. We need a list
|
||||
# because a bundle can be added to the queue multiple times. Nodes in each list
|
||||
# are insertion-ordered.
|
||||
self._bundle_to_nodes: Dict[RefBundle, Deque[_Node]] = defaultdict(deque)
|
||||
|
||||
@override
|
||||
def __contains__(self, bundle: RefBundle) -> bool:
|
||||
return bundle in self._bundle_to_nodes
|
||||
|
||||
@override
|
||||
def _add_inner(self, bundle: RefBundle, **kwargs: Any):
|
||||
new_node = _Node(value=bundle, next=None, prev=self._tail)
|
||||
# Case 1: The queue is empty.
|
||||
if self._head is None:
|
||||
assert self._tail is None
|
||||
self._head = new_node
|
||||
self._tail = new_node
|
||||
# Case 2: The queue has at least one element.
|
||||
else:
|
||||
self._tail.next = new_node
|
||||
self._tail = new_node
|
||||
|
||||
self._bundle_to_nodes[bundle].append(new_node)
|
||||
|
||||
@override
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
# Case 1: The queue is empty.
|
||||
if not self._head:
|
||||
raise IndexError("You can't pop from an empty queue")
|
||||
|
||||
bundle = self._head.value
|
||||
self._remove_inner(bundle)
|
||||
|
||||
return bundle
|
||||
|
||||
@override
|
||||
def has_next(self) -> bool:
|
||||
return self._num_bundles > 0
|
||||
|
||||
@override
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
if self._head is None:
|
||||
return None
|
||||
|
||||
return self._head.value
|
||||
|
||||
@override
|
||||
def _remove_inner(self, bundle: RefBundle) -> RefBundle:
|
||||
# Case 1: The queue is empty.
|
||||
if bundle not in self._bundle_to_nodes:
|
||||
raise ValueError(f"The bundle {bundle} is not in the queue.")
|
||||
|
||||
node = self._bundle_to_nodes[bundle].popleft()
|
||||
if not self._bundle_to_nodes[bundle]:
|
||||
del self._bundle_to_nodes[bundle]
|
||||
|
||||
node = self._remove_node(node)
|
||||
return node.value
|
||||
|
||||
def _remove_node(self, node: _Node) -> _Node:
|
||||
# Case 2: The bundle is the only element in the queue.
|
||||
if self._head is self._tail:
|
||||
self._head = None
|
||||
self._tail = None
|
||||
# Case 3: The bundle is the first element in the queue.
|
||||
elif node is self._head:
|
||||
self._head = node.next
|
||||
self._head.prev = None
|
||||
# Case 4: The bundle is the last element in the queue.
|
||||
elif node is self._tail:
|
||||
self._tail = node.prev
|
||||
self._tail.next = None
|
||||
# Case 5: The bundle is in the middle of the queue.
|
||||
else:
|
||||
node.prev.next = node.next
|
||||
node.next.prev = node.prev
|
||||
|
||||
return node
|
||||
|
||||
def __iter__(self) -> Iterator[RefBundle]:
|
||||
curr = self._head
|
||||
while curr:
|
||||
yield curr.value
|
||||
curr = curr.next
|
||||
|
||||
def clear(self):
|
||||
self._reset_metrics()
|
||||
self._bundle_to_nodes.clear()
|
||||
self._head = None
|
||||
self._tail = None
|
||||
@@ -0,0 +1,89 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict, deque
|
||||
from typing import TYPE_CHECKING, DefaultDict, Deque, Optional, Set
|
||||
|
||||
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 ReorderingBundleQueue(BaseBundleQueue):
|
||||
"""A queue that iterates over the bundles in the order of provided "keys" rather than
|
||||
insertion order (for bundles inserted with the same key, insertion order is used)
|
||||
|
||||
User of this queue has to adhere to following invariants of this queue:
|
||||
|
||||
1. (!) Used keys have to be a *contiguous* range of `[0, N]`
|
||||
|
||||
Failure to follow this requirement might result in this queue getting
|
||||
irreversibly stuck.
|
||||
|
||||
NOTE: Not thread-safe
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self._inner: DefaultDict[int, Deque[RefBundle]] = defaultdict(lambda: deque())
|
||||
self._current_key: int = 0
|
||||
self._finalized_keys: Set[int] = set()
|
||||
|
||||
def _move_to_next_key(self):
|
||||
"""Move the output index to the next task.
|
||||
|
||||
This method should only be called when the current task is complete and all
|
||||
outputs have been taken.
|
||||
"""
|
||||
assert len(self._inner[self._current_key]) == 0
|
||||
assert self._current_key in self._finalized_keys
|
||||
|
||||
self._current_key += 1
|
||||
|
||||
@override
|
||||
def _add_inner(self, bundle: RefBundle, key: int) -> None:
|
||||
assert key is not None
|
||||
self._inner[key].append(bundle)
|
||||
|
||||
@override
|
||||
def has_next(self) -> bool:
|
||||
while (
|
||||
self._current_key in self._finalized_keys
|
||||
and len(self._inner[self._current_key]) == 0
|
||||
):
|
||||
self._move_to_next_key()
|
||||
|
||||
return len(self._inner[self._current_key]) > 0
|
||||
|
||||
@override
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
# It's vital to invoke `has_next` here, to potentially advance the pointer
|
||||
# to the next key
|
||||
if not self.has_next():
|
||||
raise ValueError("Cannot pop from empty queue.")
|
||||
|
||||
return self._inner[self._current_key].popleft()
|
||||
|
||||
@override
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
# It's vital to invoke `has_next` here, to potentially advance the pointer
|
||||
# to the next key
|
||||
if not self.has_next():
|
||||
return None
|
||||
|
||||
return self._inner[self._current_key][0]
|
||||
|
||||
@override
|
||||
def finalize(self, key: int):
|
||||
assert key is not None and key >= self._current_key
|
||||
self._finalized_keys.add(key)
|
||||
|
||||
@override
|
||||
def clear(self):
|
||||
self._reset_metrics()
|
||||
self._inner.clear()
|
||||
self._finalized_keys.clear()
|
||||
self._current_key = 0
|
||||
@@ -0,0 +1,74 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from .base import BundleQueue
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.execution.interfaces import RefBundle
|
||||
|
||||
|
||||
class ThreadSafeBundleQueue(BundleQueue):
|
||||
"""A thread-safe wrapper for a ``BundleQueue``.
|
||||
|
||||
Delegates all operations to the wrapped queue while a lock.
|
||||
|
||||
NOTE: Not safe to use with ``__contains__``/``remove`` from ``QueueWithRemoval``.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: BundleQueue):
|
||||
self._inner = inner
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def estimate_size_bytes(self) -> int:
|
||||
with self._lock:
|
||||
return self._inner.estimate_size_bytes()
|
||||
|
||||
def num_blocks(self) -> int:
|
||||
with self._lock:
|
||||
return self._inner.num_blocks()
|
||||
|
||||
def num_bundles(self) -> int:
|
||||
with self._lock:
|
||||
return self._inner.num_bundles()
|
||||
|
||||
def num_rows(self) -> int:
|
||||
with self._lock:
|
||||
return self._inner.num_rows()
|
||||
|
||||
def _on_enqueue_bundle(self, bundle: RefBundle):
|
||||
raise NotImplementedError("Use add() for thread-safe access")
|
||||
|
||||
def _on_dequeue_bundle(self, bundle: RefBundle):
|
||||
raise NotImplementedError("Use get_next() for thread-safe access")
|
||||
|
||||
def _add_inner(self, bundle: RefBundle, **kwargs: Any):
|
||||
raise NotImplementedError("Use add() for thread-safe access")
|
||||
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
raise NotImplementedError("Use get_next() for thread-safe access")
|
||||
|
||||
def add(self, bundle: RefBundle, **kwargs: Any):
|
||||
with self._lock:
|
||||
self._inner.add(bundle, **kwargs)
|
||||
|
||||
def get_next(self) -> RefBundle:
|
||||
with self._lock:
|
||||
return self._inner.get_next()
|
||||
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
with self._lock:
|
||||
return self._inner.peek_next()
|
||||
|
||||
def has_next(self) -> bool:
|
||||
with self._lock:
|
||||
return self._inner.has_next()
|
||||
|
||||
def clear(self):
|
||||
with self._lock:
|
||||
self._inner.clear()
|
||||
|
||||
def finalize(self, **kwargs: Any):
|
||||
with self._lock:
|
||||
return self._inner.finalize(**kwargs)
|
||||
Reference in New Issue
Block a user