wehub-resource-sync
559 lines
22 KiB
Python
559 lines
22 KiB
Python
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
|