import collections import time from contextlib import contextmanager, nullcontext from typing import Any, Callable, Dict, Iterator, List, Optional import ray from ray._common.utils import env_integer from ray.data._internal.block_batching.interfaces import ( Batch, BlockPrefetcher, ) from ray.data._internal.block_batching.util import ( ActorBlockPrefetcher, WaitBlockPrefetcher, blocks_to_batches, collate, finalize_batches, format_batches, iter_threaded, resolve_block_refs, ) from ray.data._internal.execution.interfaces.ref_bundle import RefBundle from ray.data._internal.memory_tracing import trace_deallocation from ray.data._internal.stats import DatasetStats, TimeSpan, _StatsManager from ray.data.block import Block, DataBatch from ray.data.context import DataContext from ray.types import ObjectRef DEFAULT_FORMAT_THREADPOOL_NUM_WORKERS = env_integer( "RAY_DATA_MAX_FORMAT_THREADPOOL_NUM_WORKERS", 4 ) def _merged_duration( spans: List["TimeSpan"], blocked_start_s: float, blocked_end_s: float ) -> float: """Total time ``spans`` overlap with ``[blocked_start_s, blocked_end_s]``, with overlapping spans merged so nothing is double-counted.""" intervals = [] for s in spans: lo = max(s.start_s, blocked_start_s) hi = min(s.end_s, blocked_end_s) if hi > lo: intervals.append((lo, hi)) if not intervals: return 0.0 intervals.sort() merged = [intervals[0]] for i in range(1, len(intervals)): lo, hi = intervals[i] if lo <= merged[-1][1]: merged[-1] = (merged[-1][0], max(merged[-1][1], hi)) else: merged.append((lo, hi)) return sum(hi - lo for lo, hi in merged) class BatchIterator: """Defines an iterator pipeline to convert a stream of block object references into a stream of formatted batches ready to be consumed by the user. This takes a block iterator and creates batch_size batches, slicing, unioning, shuffling, prefetching, and formatting blocks as needed. This involves both pipeline parallelism (e.g. prefetching) and data parallelism (e.g. threadpool operations): If prefetch_batches=2, these are all the batches in flight: [User thread] trains on Batch 0 - [Fetch thread] Batch 1 finalization + move to output queue - [Worker thread 1] Batch 2 formatting + collating - [Worker thread 2] Batch 3 formatting + collating - [Raylet] Batches 4 + 5 fetched to local object store memory At any point in time there are prefetch_batches+1 batches in local heap memory. And the next set of prefetch_batches in local object store memory. The actual steps are as follows: In a single async thread, do the following: 1. Trigger Ray local prefetching of `prefetch_batches` worth of block object references. 2. Resolve (i.e. call `ray.get()`) on the block references. 3. Perform the necessary batch slicing to construct full batches, possibly shuffling if necessary. 4. Then, in a threadpool consisting of `prefetch_batches` threads: a. Format the batches to the provided batch format. b. Apply the collate function. 5. If preserve_order, restore the original batch order from the threadpool output. 6. Finalize each of the (now ordered) collated batches. Args: ref_bundles: An iterator over RefBundles. stats: DatasetStats object to record timing and other statistics. dataset_tag: The tag of the dataset to record timing and other statistics. clear_block_after_read: Whether to clear the block from object store manually (i.e. without waiting for Python's automatic GC) after it is read. Doing so will reclaim memory faster and hence reduce the memory footprint. However, the caller has to ensure the safety, i.e. the block will never be accessed again. batch_size: Record batch size, or None to let the system pick. batch_format: The format in which to return each batch. Specify "default" to use the current block format (promoting Arrow to pandas automatically), "pandas" to select ``pandas.DataFrame`` or "pyarrow" to select ``pyarrow.Table``, or None to use entire blocks as batches. Default is "default". drop_last: Whether to drop the last batch if it's incomplete. collate_fn: A function to apply to each data batch before returning it. finalize_fn: A function to apply to each data batch after it has been collated. This function is not run in a threadpool so it can be used for memory-intensive operations such as GPU preloading. shuffle_buffer_min_size: If non-None, the data will be randomly shuffled using a local in-memory shuffle buffer, and this value will serve as the minimum number of rows that must be in the local in-memory shuffle buffer in order to yield a batch. shuffle_seed: The seed to use for the local random shuffle. ensure_copy: Whether batches are always copied from the underlying base blocks (not zero-copy views). prefetch_batches: The number of batches to fetch ahead of the current batch to process. If set to greater than 0, a separate thread will be used to fetch the specified amount of formatted batches from blocks. This improves performance for non-CPU bound UDFs, allowing batch fetching compute and formatting to be overlapped with the UDF. Defaults to 1. prefetch_bytes_callback: A callback to report prefetched bytes to the executor's resource manager. preserve_order: Whether to maintain the original order that the batches were formed from the blocks (e.g., the input block order). This only takes effect in the case that the format/collate threadpool has more than one thread and the output batches have non-deterministic order. """ UPDATE_METRICS_INTERVAL_S: float = 5.0 def __init__( self, ref_bundles: Iterator[RefBundle], *, stats: Optional[DatasetStats] = None, dataset_tag: Optional[str] = None, clear_block_after_read: bool = False, batch_size: Optional[int] = None, batch_format: Optional[str] = "default", drop_last: bool = False, collate_fn: Optional[Callable[[DataBatch], Any]] = None, finalize_fn: Optional[Callable[[Any], Any]] = None, shuffle_buffer_min_size: Optional[int] = None, shuffle_seed: Optional[int] = None, ensure_copy: bool = False, prefetch_batches: int = 1, prefetch_bytes_callback: Optional[Callable[[int], None]] = None, preserve_order: bool = False, ): self._ref_bundles = ref_bundles self._stats = stats self._dataset_tag = dataset_tag self._batch_size = batch_size self._batch_format = batch_format self._drop_last = drop_last self._collate_fn = collate_fn self._finalize_fn = finalize_fn self._shuffle_buffer_min_size = shuffle_buffer_min_size self._shuffle_seed = shuffle_seed self._ensure_copy = ensure_copy self._prefetch_batches = prefetch_batches self._prefetch_bytes_callback = prefetch_bytes_callback self._preserve_order = preserve_order self._eager_free = ( clear_block_after_read and DataContext.get_current().eager_free ) actor_prefetcher_enabled = ( prefetch_batches > 0 and DataContext.get_current().actor_prefetcher_enabled and not ray.util.client.ray.is_connected() ) self._prefetcher = ( ActorBlockPrefetcher() if actor_prefetcher_enabled else WaitBlockPrefetcher() ) self._yielded_first_batch = False # This stores the last time we updated the metrics. # This allows us to update metrics on some interval, # by comparing it with the current timestamp. self._metrics_last_updated: float = 0.0 def _prefetch_blocks( self, ref_bundles: Iterator[RefBundle] ) -> Iterator[ObjectRef[Block]]: return prefetch_batches_locally( ref_bundles=ref_bundles, prefetcher=self._prefetcher, num_batches_to_prefetch=self._prefetch_batches, batch_size=self._batch_size, eager_free=self._eager_free, stats=self._stats, ) def _resolve_block_refs( self, block_refs: Iterator[ObjectRef[Block]] ) -> Iterator[Any]: return resolve_block_refs(block_ref_iter=block_refs, stats=self._stats) def _blocks_to_batches(self, blocks: Iterator[Block]) -> Iterator[Batch]: return blocks_to_batches( block_iter=blocks, stats=self._stats, batch_size=self._batch_size, drop_last=self._drop_last, shuffle_buffer_min_size=self._shuffle_buffer_min_size, shuffle_seed=self._shuffle_seed, ensure_copy=self._ensure_copy, ) def _format_batches(self, batches: Iterator[Batch]) -> Iterator[Batch]: num_threadpool_workers = min( DEFAULT_FORMAT_THREADPOOL_NUM_WORKERS, self._prefetch_batches ) return _format_in_threadpool( batch_iter=batches, stats=self._stats, batch_format=self._batch_format, collate_fn=self._collate_fn, num_threadpool_workers=num_threadpool_workers, ensure_copy=self._ensure_copy, ) def _finalize_batches( self, batch_iter: Iterator[Batch], ) -> Iterator[Batch]: if self._finalize_fn is None: return batch_iter return finalize_batches( batch_iter, finalize_fn=self._finalize_fn, stats=self._stats ) def _restore_original_batch_order( self, batches: Iterator[Batch] ) -> Iterator[Batch]: return restore_original_order(batches) def _pipeline(self, ref_bundles: Iterator[RefBundle]) -> Iterator[Batch]: # Step 1: Prefetch logical batches locally. block_iter = self._prefetch_blocks(ref_bundles) # Step 2: Resolve the blocks. block_iter = self._resolve_block_refs(block_iter) # Step 3: Batch and shuffle the resolved blocks. batch_iter = self._blocks_to_batches(block_iter) # Step 4: Format and collate the batches in a threadpool. batch_iter = self._format_batches(batch_iter) # Step 5 (optional): Restore the original order of the batches # if preserve_order is True, in the case that the format/collate threadpool # shuffles around the batches non-deterministically. # NOTE: This should happen before finalize_fn so the reorder buffer # holds CPU batches rather than finalize_fn outputs (e.g., GPU tensors). if self._preserve_order: batch_iter = self._restore_original_batch_order(batch_iter) # Step 6: Finalize the batches (e.g., move to GPU). batch_iter = self._finalize_batches(batch_iter) yield from batch_iter def _iter_batches(self) -> Iterator[DataBatch]: """Pull batches from the pipeline and yield batch data. Captures the training thread's blocked window around each ``next()`` call and attributes it to pipeline stages via ``_attribute_blocked_time``. """ batch_iter = iter_threaded(self._ref_bundles, fn=self._pipeline) self.before_epoch_start() while True: with self.get_next_batch_context(): blocked_start_s = time.perf_counter() try: batch = next(batch_iter) except StopIteration: break blocked_end_s = time.perf_counter() self._attribute_blocked_time(batch, blocked_start_s, blocked_end_s) with self.yield_batch_context(batch): yield batch.data self.after_epoch_end() def _attribute_blocked_time( self, batch: Batch, blocked_start_s: float, blocked_end_s: float ) -> None: """Attribute per-stage blocked time via overlap with the training window. Each stage's spans on ``batch.metadata.stage_timings`` are intersected with the training thread's blocked window ``[blocked_start_s, blocked_end_s]``. Overlapping spans are merged first, so the result is the total time the stage was active during the stall (no double-counting). Limitation: only the yielded batch's spans are attributed. Other in-flight batches (being processed by background threads) may also overlap with the training stall window but are not counted. TODO: track in-flight batches and union their spans for complete attribution. The current implementation suffices for capturing data-loading bottlenecks. TODO: reorder buffer wait under ``preserve_order`` is unattributed (per-stage spans are recorded at format/collate completion, before the batch leaves ``restore_original_order``). Args: batch: Batch whose per-stage timings should be attributed. blocked_start_s: perf_counter() just before next(). blocked_end_s: perf_counter() just after next() returned. """ if self._stats is None: return timings = batch.metadata.stage_timings for stage, spans in timings.stages(): overlap_s = _merged_duration(spans, blocked_start_s, blocked_end_s) if overlap_s > 0: self._stats.get_blocked_timer(stage).add(overlap_s) self._stats.iter_batches_total += 1 self._stats.iter_rows_total += batch.metadata.num_rows def __iter__(self) -> Iterator[DataBatch]: return self._iter_batches() def before_epoch_start(self): self._yielded_first_batch = False def after_epoch_end(self): # Report 0 prefetched bytes at the end of iteration. if self._prefetch_bytes_callback is not None: self._prefetch_bytes_callback(0) if self._stats is None: return _StatsManager.update_iteration_metrics(self._stats, self._dataset_tag) @contextmanager def get_next_batch_context(self): """Context around ``next(batch_iter)``: tracks total blocked time and time-to-first-batch.""" try: if self._stats: # Always track total blocked time total_timer = self._stats.iter_total_blocked_s.timer() # Also track the time until the first batch is ready first_batch_ready_timer = ( self._stats.iter_time_to_first_batch_s.timer() if not self._yielded_first_batch else nullcontext() ) with total_timer, first_batch_ready_timer: yield else: yield finally: self._yielded_first_batch = True @contextmanager def yield_batch_context(self, batch: Batch): """Context around yielding a batch to the user: tracks user time and periodically flushes metrics.""" with self._stats.iter_user_s.timer() if self._stats else nullcontext(): yield # Report prefetched bytes to the executor's resource manager. if self._prefetch_bytes_callback is not None and self._stats is not None: self._prefetch_bytes_callback(self._stats.iter_prefetched_bytes) if self._stats is None: return now = time.time() if (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S: _StatsManager.update_iteration_metrics(self._stats, self._dataset_tag) self._metrics_last_updated = now def _format_in_threadpool( batch_iter: Iterator[Batch], stats: DatasetStats, batch_format: Optional[str], collate_fn: Optional[Callable[[DataBatch], Any]], num_threadpool_workers: int, ensure_copy: bool = False, ) -> Iterator[Batch]: """Executes the batching, formatting, and collation logic in a threadpool. Args: batch_iter: An iterator over logical batches. stats: DatasetStats object to record timing and other statistics. batch_format: The format in which to return each batch. Specify "default" to use the current block format (promoting Arrow to pandas automatically), "pandas" to select ``pandas.DataFrame`` or "pyarrow" to select ``pyarrow.Table``, or None to use entire blocks as batches. collate_fn: A function to apply to each data batch before returning it. num_threadpool_workers: The number of threads to use in the threadpool. ensure_copy: Whether batches are always copied from the underlying base blocks (not zero-copy views). Returns: An iterator over batches with formatting and collation applied. """ def threadpool_computations_format_collate( batch_iter: Iterator[Batch], ) -> Iterator[Batch]: # Step 4a: Format the batches. formatted_batch_iter = format_batches( batch_iter, batch_format=batch_format, stats=stats, ensure_copy=ensure_copy ) # Step 4b: Apply the collate function if applicable. if collate_fn is not None: formatted_batch_iter = collate( formatted_batch_iter, collate_fn=collate_fn, stats=stats ) return formatted_batch_iter if num_threadpool_workers > 0: # Output order is non-deterministic across workers and is restored # downstream by `restore_original_order`. collated_iter = iter_threaded( base_iterator=batch_iter, fn=threadpool_computations_format_collate, num_workers=num_threadpool_workers, output_buffer_size=num_threadpool_workers, ) else: collated_iter = threadpool_computations_format_collate(batch_iter) return collated_iter def prefetch_batches_locally( ref_bundles: Iterator[RefBundle], prefetcher: BlockPrefetcher, num_batches_to_prefetch: int, batch_size: Optional[int], eager_free: bool = False, stats: Optional[DatasetStats] = None, ) -> Iterator[ObjectRef[Block]]: """Given an iterator of batched RefBundles, returns an iterator over the corresponding block references while prefetching `num_batches_to_prefetch` batches in advance. Args: ref_bundles: An iterator over batched RefBundles. prefetcher: The prefetcher to use. num_batches_to_prefetch: The number of batches to prefetch ahead of the current batch during the scan. batch_size: User specified batch size, or None to let the system pick. eager_free: Whether to eagerly free the object reference from the object store. stats: Dataset stats object used to store ref bundle retrieval time. Yields: Block: Block references (as ObjectRefs), in order. """ def get_next_ref_bundle() -> RefBundle: with stats.iter_get_ref_bundles_s.timer() if stats else nullcontext(): return next(ref_bundles) sliding_window = collections.deque() current_window_size = 0 if num_batches_to_prefetch <= 0: if stats: stats.iter_prefetched_bytes = 0 for ref_bundle in ref_bundles: for block_ref in ref_bundle.block_refs: yield block_ref return if batch_size is not None: num_rows_to_prefetch = num_batches_to_prefetch * batch_size else: num_rows_to_prefetch = None # Create and fetch the initial window. # Stop adding if the number of rows in this window is greater than requested # batch size, or if the batch size is None and the number of blocks in this window # is greater than requested batches to prefetch. while (batch_size is not None and current_window_size < num_rows_to_prefetch) or ( batch_size is None and len(sliding_window) < num_batches_to_prefetch ): try: next_ref_bundle = get_next_ref_bundle() sliding_window.extend(next_ref_bundle.blocks) current_window_size += next_ref_bundle.num_rows() except StopIteration: break prefetcher.prefetch_blocks([entry.ref for entry in sliding_window]) if stats: stats.iter_prefetched_bytes = sum( entry.metadata.size_bytes or 0 for entry in sliding_window ) while sliding_window: entry = sliding_window.popleft() current_window_size -= entry.metadata.num_rows if batch_size is None or current_window_size < num_rows_to_prefetch: try: next_ref_bundle = get_next_ref_bundle() for next_entry in next_ref_bundle.blocks: sliding_window.append(next_entry) current_window_size += next_entry.metadata.num_rows prefetcher.prefetch_blocks([entry.ref for entry in sliding_window]) except StopIteration: pass if stats: stats.iter_prefetched_bytes = sum( entry.metadata.size_bytes or 0 for entry in sliding_window ) yield entry.ref trace_deallocation(entry.ref, loc="iter_batches", free=eager_free) prefetcher.stop() def restore_original_order(batch_iter: Iterator[Batch]) -> Iterator[Batch]: """Restores the original order of the provided `batch_iter` This function will yield items from `base_iterator` in the correct order based on each batch's batch_idx. All indexes are expected to be unique. `batch_iter` is expected to not have any missing indexes. All indexes from 0 to len (base_iterator) must be present. """ next_index_required = 0 buffer: Dict[int, Batch] = {} for batch in batch_iter: assert batch.metadata.batch_idx not in buffer buffer[batch.metadata.batch_idx] = batch while next_index_required in buffer: yield buffer.pop(next_index_required) next_index_required += 1 while next_index_required in buffer: yield buffer.pop(next_index_required) next_index_required += 1