wehub-resource-sync
102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
from typing import Callable, Iterator, Optional, TypeVar
|
|
|
|
from ray.data._internal.block_batching.interfaces import ResolvedBlock
|
|
from ray.data._internal.block_batching.util import (
|
|
_MappingIterator,
|
|
blocks_to_batches,
|
|
collate,
|
|
format_batches,
|
|
)
|
|
from ray.data._internal.stats import DatasetStats
|
|
from ray.data.block import Block, DataBatch
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def batch_blocks(
|
|
blocks: Iterator[Block],
|
|
*,
|
|
stats: Optional[DatasetStats] = None,
|
|
batch_size: Optional[int] = None,
|
|
batch_format: str = "default",
|
|
drop_last: bool = False,
|
|
collate_fn: Optional[Callable[[DataBatch], DataBatch]] = None,
|
|
shuffle_buffer_min_size: Optional[int] = None,
|
|
shuffle_seed: Optional[int] = None,
|
|
ensure_copy: bool = False,
|
|
) -> Iterator[DataBatch]:
|
|
"""Create formatted batches of data from 1 or more blocks.
|
|
|
|
This function takes in an iterator of already fetched blocks. Consequently, this
|
|
function doesn't support block prefetching.
|
|
"""
|
|
# TODO: make stage timings optional at _BatchingIterator so this
|
|
# shim can be removed. map() avoids holding block references.
|
|
wrapped_blocks = map(lambda b: ResolvedBlock(block=b), blocks)
|
|
|
|
# Build the processing pipeline
|
|
batch_iter = format_batches(
|
|
blocks_to_batches(
|
|
block_iter=wrapped_blocks,
|
|
stats=stats,
|
|
batch_size=batch_size,
|
|
drop_last=drop_last,
|
|
shuffle_buffer_min_size=shuffle_buffer_min_size,
|
|
shuffle_seed=shuffle_seed,
|
|
ensure_copy=ensure_copy,
|
|
),
|
|
batch_format=batch_format,
|
|
stats=stats,
|
|
ensure_copy=ensure_copy,
|
|
)
|
|
|
|
if collate_fn is not None:
|
|
batch_iter = collate(batch_iter, collate_fn=collate_fn, stats=stats)
|
|
|
|
return _UserTimingIterator(
|
|
_MappingIterator(batch_iter, lambda batch: batch.data), stats
|
|
)
|
|
|
|
|
|
class _UserTimingIterator(Iterator[DataBatch]):
|
|
def __init__(self, iter: Iterator[DataBatch], stats: Optional[DatasetStats]):
|
|
self._iter = iter
|
|
self._stats = stats
|
|
self._active_timer = None
|
|
|
|
def __iter__(self) -> Iterator[DataBatch]:
|
|
return self
|
|
|
|
def __next__(self) -> DataBatch:
|
|
# Since we're tracking time spent in user-code, we stop
|
|
# the timer immediately when `__next__` is called
|
|
self._stop_timer()
|
|
|
|
try:
|
|
res = next(self._iter)
|
|
# Reset timer and return
|
|
#
|
|
# NOTE: It's crucial that we reset the timer only after we
|
|
# retrieved the result to avoid starting the timer before
|
|
# we retrieve the next value
|
|
self._reset_timer()
|
|
return res
|
|
except StopIteration:
|
|
self._stop_timer()
|
|
raise
|
|
|
|
def _stop_timer(self):
|
|
if not self._stats:
|
|
return
|
|
|
|
if self._active_timer:
|
|
self._active_timer.__exit__(None, None, None)
|
|
self._active_timer = None
|
|
|
|
def _reset_timer(self):
|
|
if not self._stats:
|
|
return
|
|
|
|
self._active_timer = self._stats.iter_user_s.timer()
|
|
self._active_timer.__enter__()
|