wehub-resource-sync
255 lines
8.1 KiB
Python
255 lines
8.1 KiB
Python
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
|