chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,3 @@
|
||||
from ray.data._internal.block_batching.block_batching import batch_blocks
|
||||
|
||||
__all__ = ["batch_blocks"]
|
||||
@@ -0,0 +1,101 @@
|
||||
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__()
|
||||
@@ -0,0 +1,134 @@
|
||||
import abc
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Iterable, List, Optional, Tuple
|
||||
|
||||
from ray.data._internal.stats import IterationStage, TimeSpan
|
||||
from ray.data.block import Block, DataBatch
|
||||
from ray.types import ObjectRef
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockStageTimings:
|
||||
"""Per-block timing for production_wait + data_transfer.
|
||||
|
||||
Both fields are always populated when ``stage_timings`` is set on a
|
||||
``ResolvedBlock``; the outer ``ResolvedBlock.stage_timings`` Optional
|
||||
encodes "no timing recorded" (e.g. blocks already resolved before
|
||||
entering the pipeline).
|
||||
"""
|
||||
|
||||
production_wait: TimeSpan
|
||||
data_transfer: TimeSpan
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResolvedBlock:
|
||||
"""A resolved block paired with its per-block stage timings.
|
||||
|
||||
``stage_timings`` is None when no timing was recorded (e.g. blocks
|
||||
already resolved before entering the pipeline).
|
||||
"""
|
||||
|
||||
block: Block
|
||||
stage_timings: Optional[BlockStageTimings] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchStageTimings:
|
||||
"""Per-batch timing windows for each iteration stage.
|
||||
|
||||
Fetch stages (production_wait, data_transfer) accumulate one span per
|
||||
block, so they are ``List[TimeSpan]``. Other stages run at most once
|
||||
per batch, so they are ``Optional[TimeSpan]``. ``stages()`` yields
|
||||
``List[TimeSpan]`` for all stages (single spans wrapped in a 1-element
|
||||
list) so ``_attribute_blocked_time`` can use uniform overlap logic.
|
||||
"""
|
||||
|
||||
production_wait: List[TimeSpan] = field(default_factory=list)
|
||||
data_transfer: List[TimeSpan] = field(default_factory=list)
|
||||
batching: Optional[TimeSpan] = None
|
||||
format: Optional[TimeSpan] = None
|
||||
collate: Optional[TimeSpan] = None
|
||||
finalize: Optional[TimeSpan] = None
|
||||
|
||||
def stages(self) -> Iterable[Tuple[IterationStage, List[TimeSpan]]]:
|
||||
"""Yield (stage, spans) pairs, wrapping single spans in a list."""
|
||||
return (
|
||||
(IterationStage.PRODUCTION_WAIT, self.production_wait),
|
||||
(IterationStage.DATA_TRANSFER, self.data_transfer),
|
||||
(
|
||||
IterationStage.BATCHING,
|
||||
[self.batching] if self.batching is not None else [],
|
||||
),
|
||||
(IterationStage.FORMAT, [self.format] if self.format is not None else []),
|
||||
(
|
||||
IterationStage.COLLATE,
|
||||
[self.collate] if self.collate is not None else [],
|
||||
),
|
||||
(
|
||||
IterationStage.FINALIZE,
|
||||
[self.finalize] if self.finalize is not None else [],
|
||||
),
|
||||
)
|
||||
|
||||
def accumulate_block_timings(self, src: BlockStageTimings) -> None:
|
||||
"""Accumulate a block's fetch timings into this batch's lists.
|
||||
|
||||
A boundary block whose rows span multiple batches is attributed
|
||||
to the first batch it lands in.
|
||||
"""
|
||||
self.production_wait.append(src.production_wait)
|
||||
self.data_transfer.append(src.data_transfer)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchMetadata:
|
||||
"""Metadata associated with a batch.
|
||||
|
||||
Attributes:
|
||||
batch_idx: The global index of this batch so that downstream operations can
|
||||
maintain ordering.
|
||||
num_rows: Number of rows in this batch (for ``iter_rows_total``).
|
||||
stage_timings: Per-stage timing windows.
|
||||
"""
|
||||
|
||||
batch_idx: int
|
||||
num_rows: int = 0
|
||||
stage_timings: BatchStageTimings = field(default_factory=BatchStageTimings)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Batch:
|
||||
"""A batch of data.
|
||||
|
||||
Attributes:
|
||||
metadata: Metadata associated with this batch.
|
||||
data: The batch of data.
|
||||
"""
|
||||
|
||||
metadata: BatchMetadata
|
||||
data: DataBatch
|
||||
|
||||
|
||||
class CollatedBatch(Batch):
|
||||
"""A batch of collated data.
|
||||
|
||||
Attributes:
|
||||
data: The batch of data which is the output of a user provided collate_fn
|
||||
Therefore, the type of this data can be Any.
|
||||
"""
|
||||
|
||||
data: Any
|
||||
|
||||
|
||||
class BlockPrefetcher(metaclass=abc.ABCMeta):
|
||||
"""Interface for prefetching blocks."""
|
||||
|
||||
@abc.abstractmethod
|
||||
def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
|
||||
"""Prefetch the provided blocks to this node."""
|
||||
pass
|
||||
|
||||
def stop(self):
|
||||
"""Stop prefetching and release resources."""
|
||||
pass
|
||||
@@ -0,0 +1,558 @@
|
||||
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
|
||||
@@ -0,0 +1,561 @@
|
||||
import dataclasses
|
||||
import functools
|
||||
import logging
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Generator,
|
||||
Generic,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
TypeVar,
|
||||
)
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.data._internal.batcher import Batcher, ShufflingBatcher
|
||||
from ray.data._internal.block_batching.interfaces import (
|
||||
Batch,
|
||||
BatchMetadata,
|
||||
BatchStageTimings,
|
||||
BlockPrefetcher,
|
||||
BlockStageTimings,
|
||||
CollatedBatch,
|
||||
ResolvedBlock,
|
||||
)
|
||||
from ray.data._internal.stats import DatasetStats, TimeSpan, _maybe_time
|
||||
from ray.data.block import Block, BlockAccessor, DataBatch
|
||||
from ray.types import ObjectRef
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
U = TypeVar("U")
|
||||
I = TypeVar("I")
|
||||
O = TypeVar("O")
|
||||
|
||||
_SENTINEL = object()
|
||||
|
||||
|
||||
def iter_threaded(
|
||||
base_iterator: Iterator[T],
|
||||
fn: Callable[[Iterator[T]], Iterator[U]],
|
||||
num_workers: int = 1,
|
||||
output_buffer_size: int = 1,
|
||||
) -> Generator[U, None, None]:
|
||||
"""Apply ``fn`` to ``base_iterator`` across ``num_workers`` background
|
||||
threads, yielding results through a bounded queue.
|
||||
|
||||
Workers share ``base_iterator`` under a lock (so it may be a stateful,
|
||||
non-thread-safe generator) and run ``fn`` concurrently. With
|
||||
``num_workers > 1`` the output order is not preserved and must be restored
|
||||
downstream by the consumer.
|
||||
|
||||
Invariant: the number of output-queue items + items in-flight in workers is
|
||||
bounded by ``output_buffer_size``.
|
||||
Workers reserve an output buffer slot before pulling from ``fn``, ensuring
|
||||
they don't run ``fn`` (and hold the result) while waiting for queue space.
|
||||
|
||||
When the consumer stops early (``break``, ``.close()``, or GC), workers
|
||||
are signaled via a stop event so they don't leak. Note: a hanging
|
||||
``fn`` cannot be interrupted, so ``fn`` must terminate or raise within
|
||||
bounded time per element. For example, the user function should have
|
||||
timeouts if doing blocking I/O.
|
||||
|
||||
Args:
|
||||
base_iterator: Iterator consumed (under a lock) by the workers.
|
||||
fn: Transform applied by each worker to its view of ``base_iterator``.
|
||||
num_workers: Number of background worker threads.
|
||||
output_buffer_size: Max number of items held by the output-queue
|
||||
+ in-flight in the workers.
|
||||
"""
|
||||
if num_workers < 1:
|
||||
raise ValueError("num_workers must be at least 1.")
|
||||
if output_buffer_size < 1:
|
||||
raise ValueError("output_buffer_size must be at least 1.")
|
||||
|
||||
stopped = threading.Event()
|
||||
result_queue: queue.Queue = queue.Queue()
|
||||
slots = threading.Semaphore(output_buffer_size)
|
||||
iter_lock = threading.Lock()
|
||||
|
||||
def _locked_iter() -> Iterator[T]:
|
||||
while True:
|
||||
with iter_lock:
|
||||
if stopped.is_set():
|
||||
return
|
||||
try:
|
||||
item = next(base_iterator)
|
||||
except StopIteration:
|
||||
return
|
||||
yield item
|
||||
|
||||
def _acquire_slot() -> bool:
|
||||
# Block until a slot is acquired or the consumer has stopped.
|
||||
while not stopped.is_set():
|
||||
if slots.acquire(timeout=0.1):
|
||||
return True
|
||||
return False
|
||||
|
||||
remaining_workers = num_workers
|
||||
remaining_lock = threading.Lock()
|
||||
|
||||
def _worker():
|
||||
nonlocal remaining_workers
|
||||
slot_acquired = False
|
||||
try:
|
||||
# Construct `fn_iter` inside the try so any exception during
|
||||
# construction propagates to the consumer via the outer except.
|
||||
fn_iter = fn(_locked_iter())
|
||||
while True:
|
||||
slot_acquired = _acquire_slot()
|
||||
if not slot_acquired:
|
||||
break
|
||||
item = next(fn_iter)
|
||||
result_queue.put(item)
|
||||
# The consumer pulling from the result_queue will release the slot.
|
||||
# Resetting here prevents the finally block from double-releasing.
|
||||
slot_acquired = False
|
||||
except StopIteration:
|
||||
pass
|
||||
except Exception as e:
|
||||
# Handle errors in `fn` by propagating them to the consumer.
|
||||
if not stopped.is_set():
|
||||
result_queue.put(e)
|
||||
finally:
|
||||
if slot_acquired:
|
||||
slots.release()
|
||||
with remaining_lock:
|
||||
remaining_workers -= 1
|
||||
is_last = remaining_workers == 0
|
||||
# Signal the consumer that all thread workers have exhausted their input.
|
||||
if is_last and not stopped.is_set():
|
||||
result_queue.put(_SENTINEL)
|
||||
|
||||
worker_threads = [
|
||||
threading.Thread(target=_worker, name="iter_threaded", daemon=True)
|
||||
for _ in range(num_workers)
|
||||
]
|
||||
for t in worker_threads:
|
||||
t.start()
|
||||
|
||||
try:
|
||||
while True:
|
||||
item = result_queue.get()
|
||||
if item is _SENTINEL:
|
||||
break
|
||||
if isinstance(item, Exception):
|
||||
raise item
|
||||
# Release one slot at yield time so a worker can run `fn` for the next item.
|
||||
slots.release()
|
||||
yield item
|
||||
finally:
|
||||
stopped.set()
|
||||
|
||||
|
||||
class _MappingIterator(Iterator[O], Generic[I, O]):
|
||||
"""Iterator that applies a transform function to each element.
|
||||
|
||||
Unlike a generator, local variables in __next__ go out of scope when the method
|
||||
returns, avoiding holding references to yielded values.
|
||||
"""
|
||||
|
||||
def __init__(self, input_iter: Iterator[I], transform_fn: Callable[[I], O]):
|
||||
self._input_iter = input_iter
|
||||
self._transform_fn = transform_fn
|
||||
|
||||
def __iter__(self) -> "_MappingIterator[I, O]":
|
||||
return self
|
||||
|
||||
def __next__(self) -> O:
|
||||
return self._transform_fn(next(self._input_iter))
|
||||
|
||||
|
||||
def _calculate_ref_hits(refs: List[ObjectRef[Any]]) -> Tuple[int, int, int]:
|
||||
"""Given a list of object references, returns how many are already on the local
|
||||
node, how many require fetching from another node, and how many have unknown
|
||||
locations. If `DataContext.get_current().enable_get_object_locations_for_metrics` is
|
||||
False, this will return `(0, 0, 0)` as getting object locations is disabled."""
|
||||
current_node_id = ray.get_runtime_context().get_node_id()
|
||||
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
if ctx.enable_get_object_locations_for_metrics:
|
||||
locs = ray.experimental.get_object_locations(refs)
|
||||
nodes: List[List[str]] = [loc["node_ids"] for loc in locs.values()]
|
||||
hits = sum(current_node_id in node_ids for node_ids in nodes)
|
||||
unknowns = sum(1 for node_ids in nodes if not node_ids)
|
||||
misses = len(nodes) - hits - unknowns
|
||||
return hits, misses, unknowns
|
||||
|
||||
return 0, 0, 0
|
||||
|
||||
|
||||
def resolve_block_refs(
|
||||
block_ref_iter: Iterator[ObjectRef[Block]],
|
||||
stats: Optional[DatasetStats] = None,
|
||||
) -> Iterator[ResolvedBlock]:
|
||||
"""Resolve block references via ``ray.get()`` and attach per-block
|
||||
stage timings.
|
||||
|
||||
production_wait is captured manually (no Timer accumulation) to avoid
|
||||
double-counting with ``prefetch_batches_locally``'s
|
||||
``iter_get_ref_bundles_s`` timer; data_transfer uses ``_maybe_time``
|
||||
normally (no overlap with other timers).
|
||||
|
||||
Args:
|
||||
block_ref_iter: An iterator over block object references.
|
||||
stats: An optional stats object to record block hits, misses, and
|
||||
cumulative ray.get() time.
|
||||
|
||||
Yields:
|
||||
ResolvedBlock: Each resolved block with its stage timings.
|
||||
"""
|
||||
hits = 0
|
||||
misses = 0
|
||||
unknowns = 0
|
||||
|
||||
while True:
|
||||
production_wait_start = time.perf_counter() if stats else 0.0
|
||||
try:
|
||||
block_ref = next(block_ref_iter)
|
||||
except StopIteration:
|
||||
break
|
||||
production_wait_span = (
|
||||
TimeSpan(start_s=production_wait_start, end_s=time.perf_counter())
|
||||
if stats
|
||||
else None
|
||||
)
|
||||
|
||||
current_hit, current_miss, current_unknown = _calculate_ref_hits([block_ref])
|
||||
hits += current_hit
|
||||
misses += current_miss
|
||||
unknowns += current_unknown
|
||||
|
||||
# data_transfer: cross-node transfer via ray.get().
|
||||
# TODO(amogkam): batch multiple references in one ray.get() call.
|
||||
with _maybe_time(stats.iter_get_s if stats else None) as data_transfer_span:
|
||||
block = ray.get(block_ref)
|
||||
|
||||
if stats:
|
||||
assert production_wait_span is not None
|
||||
assert data_transfer_span is not None
|
||||
stage_timings = BlockStageTimings(
|
||||
production_wait=production_wait_span,
|
||||
data_transfer=data_transfer_span,
|
||||
)
|
||||
else:
|
||||
stage_timings = None
|
||||
yield ResolvedBlock(block=block, stage_timings=stage_timings)
|
||||
|
||||
if stats:
|
||||
stats.iter_blocks_local = hits
|
||||
stats.iter_blocks_remote = misses
|
||||
stats.iter_unknown_location = unknowns
|
||||
|
||||
|
||||
def blocks_to_batches(
|
||||
block_iter: Iterator[ResolvedBlock],
|
||||
stats: Optional[DatasetStats] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
drop_last: bool = False,
|
||||
shuffle_buffer_min_size: Optional[int] = None,
|
||||
shuffle_seed: Optional[int] = None,
|
||||
ensure_copy: bool = False,
|
||||
) -> Iterator[Batch]:
|
||||
"""Given an iterator over blocks, returns an iterator over batches."""
|
||||
return _BatchingIterator(
|
||||
block_iter,
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class _BatchingIterator(Iterator[Batch]):
|
||||
"""Iterator that converts blocks to batches.
|
||||
|
||||
Unlike a generator, local variables in __next__ go out of scope when the method
|
||||
returns, avoiding holding references to yielded values.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
block_iter: Iterator[ResolvedBlock],
|
||||
stats: Optional[DatasetStats] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
drop_last: bool = False,
|
||||
shuffle_buffer_min_size: Optional[int] = None,
|
||||
shuffle_seed: Optional[int] = None,
|
||||
ensure_copy: bool = False,
|
||||
):
|
||||
self._block_iter = block_iter
|
||||
self._stats = stats
|
||||
self._drop_last = drop_last
|
||||
self._global_counter = 0
|
||||
self._done_adding = False
|
||||
# Accumulates per-block stage timings until a batch is yielded.
|
||||
self._pending_timings = BatchStageTimings()
|
||||
|
||||
if shuffle_buffer_min_size is not None:
|
||||
self._batcher = ShufflingBatcher(
|
||||
batch_size=batch_size,
|
||||
shuffle_buffer_min_size=shuffle_buffer_min_size,
|
||||
shuffle_seed=shuffle_seed,
|
||||
)
|
||||
else:
|
||||
self._batcher = Batcher(batch_size=batch_size, ensure_copy=ensure_copy)
|
||||
|
||||
def __iter__(self) -> "_BatchingIterator":
|
||||
return self
|
||||
|
||||
def __next__(self) -> Batch:
|
||||
# Try to get a batch from current batcher state
|
||||
while True:
|
||||
can_yield = self._batcher.has_batch() or (
|
||||
self._batcher.has_any() and self._done_adding and not self._drop_last
|
||||
)
|
||||
|
||||
if can_yield:
|
||||
with _maybe_time(
|
||||
self._stats.iter_next_batch_s if self._stats else None
|
||||
) as span:
|
||||
next_batch = self._batcher.next_batch()
|
||||
self._pending_timings.batching = span
|
||||
|
||||
res = Batch(
|
||||
metadata=BatchMetadata(
|
||||
batch_idx=self._global_counter,
|
||||
num_rows=BlockAccessor.for_block(next_batch).num_rows(),
|
||||
stage_timings=self._pending_timings,
|
||||
),
|
||||
data=next_batch,
|
||||
)
|
||||
self._pending_timings = BatchStageTimings()
|
||||
|
||||
self._global_counter += 1
|
||||
return res
|
||||
|
||||
elif not self._done_adding:
|
||||
# If can't yield try adding more blocks
|
||||
try:
|
||||
# NOTE: Block ref is released immediately
|
||||
block_result = next(self._block_iter)
|
||||
if block_result.stage_timings is not None:
|
||||
self._pending_timings.accumulate_block_timings(
|
||||
block_result.stage_timings
|
||||
)
|
||||
self._batcher.add(block_result.block)
|
||||
except StopIteration:
|
||||
self._batcher.done_adding()
|
||||
self._done_adding = True
|
||||
else:
|
||||
# In case when
|
||||
# - We've exhausted input AND
|
||||
# - There's nothing to yield anymore
|
||||
#
|
||||
# We stop the iteration
|
||||
raise StopIteration
|
||||
|
||||
|
||||
def _format_batch(
|
||||
batch: Batch,
|
||||
batch_format: Optional[str],
|
||||
stats: Optional[DatasetStats],
|
||||
ensure_copy: bool = False,
|
||||
) -> Batch:
|
||||
with _maybe_time(stats.iter_format_batch_s if stats else None) as span:
|
||||
formatted_data = BlockAccessor.for_block(batch.data).to_batch_format(
|
||||
batch_format
|
||||
)
|
||||
if ensure_copy:
|
||||
formatted_data = _copy_batch(formatted_data)
|
||||
batch.metadata.stage_timings.format = span
|
||||
return dataclasses.replace(batch, data=formatted_data)
|
||||
|
||||
|
||||
def _copy_batch(batch: "DataBatch") -> "DataBatch":
|
||||
"""Return a copy of a batch, making it writable.
|
||||
|
||||
``pa.Array.to_numpy()`` returns read-only arrays by default, so when
|
||||
a caller passes ``ensure_copy=True`` (i.e. ``zero_copy_batch=False``) and the
|
||||
block is Arrow, the numpy-format batch must be explicitly copied to give the UDF
|
||||
writable arrays.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if isinstance(batch, dict):
|
||||
# Return a dictionary with the same keys (column names) and values (column numpy arrays),
|
||||
# with the values copied
|
||||
return {
|
||||
k: v.copy() if isinstance(v, np.ndarray) else v for k, v in batch.items()
|
||||
}
|
||||
elif isinstance(batch, np.ndarray):
|
||||
return batch.copy()
|
||||
return batch
|
||||
|
||||
|
||||
def format_batches(
|
||||
batch_iter: Iterator[Batch],
|
||||
batch_format: Optional[str],
|
||||
stats: Optional[DatasetStats] = None,
|
||||
ensure_copy: bool = False,
|
||||
) -> Iterator[Batch]:
|
||||
"""Given an iterator of batches, returns an iterator of formatted batches."""
|
||||
return _MappingIterator(
|
||||
batch_iter,
|
||||
functools.partial(
|
||||
_format_batch,
|
||||
batch_format=batch_format,
|
||||
stats=stats,
|
||||
ensure_copy=ensure_copy,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _collate_batch(
|
||||
batch: Batch,
|
||||
collate_fn: Callable[[DataBatch], Any],
|
||||
stats: Optional[DatasetStats],
|
||||
) -> CollatedBatch:
|
||||
with _maybe_time(stats.iter_collate_batch_s if stats else None) as span:
|
||||
collated_data = collate_fn(batch.data)
|
||||
batch.metadata.stage_timings.collate = span
|
||||
return CollatedBatch(metadata=batch.metadata, data=collated_data)
|
||||
|
||||
|
||||
def collate(
|
||||
batch_iter: Iterator[Batch],
|
||||
collate_fn: Optional[Callable[[DataBatch], Any]],
|
||||
stats: Optional[DatasetStats] = None,
|
||||
) -> Iterator[CollatedBatch]:
|
||||
"""Returns an iterator with the provided collate_fn applied to batches."""
|
||||
if not isinstance(batch_iter, Iterator):
|
||||
batch_iter = iter(batch_iter)
|
||||
|
||||
return _MappingIterator(
|
||||
batch_iter,
|
||||
functools.partial(_collate_batch, collate_fn=collate_fn, stats=stats),
|
||||
)
|
||||
|
||||
|
||||
def _finalize_batch(
|
||||
batch: CollatedBatch,
|
||||
finalize_fn: Callable[[Any], Any],
|
||||
stats: Optional[DatasetStats],
|
||||
) -> CollatedBatch:
|
||||
with _maybe_time(stats.iter_finalize_batch_s if stats else None) as span:
|
||||
finalized_data = finalize_fn(batch.data)
|
||||
batch.metadata.stage_timings.finalize = span
|
||||
return dataclasses.replace(batch, data=finalized_data)
|
||||
|
||||
|
||||
def finalize_batches(
|
||||
batch_iter: Iterator[CollatedBatch],
|
||||
finalize_fn: Callable[[Any], Any],
|
||||
stats: Optional[DatasetStats] = None,
|
||||
) -> Iterator[CollatedBatch]:
|
||||
"""Returns an iterator with finalize_fn applied to batches."""
|
||||
if not isinstance(batch_iter, Iterator):
|
||||
batch_iter = iter(batch_iter)
|
||||
|
||||
return _MappingIterator(
|
||||
batch_iter,
|
||||
functools.partial(_finalize_batch, finalize_fn=finalize_fn, stats=stats),
|
||||
)
|
||||
|
||||
|
||||
PREFETCHER_ACTOR_NAMESPACE = "ray.dataset"
|
||||
|
||||
|
||||
class WaitBlockPrefetcher(BlockPrefetcher):
|
||||
"""Block prefetcher using ray.wait."""
|
||||
|
||||
def __init__(self):
|
||||
self._blocks = []
|
||||
self._stopped = False
|
||||
self._condition = threading.Condition()
|
||||
self._thread = threading.Thread(
|
||||
target=self._run,
|
||||
name="Prefetcher",
|
||||
daemon=True,
|
||||
)
|
||||
self._thread.start()
|
||||
|
||||
def _run(self):
|
||||
while not self._stopped:
|
||||
try:
|
||||
with self._condition:
|
||||
if len(self._blocks) == 0:
|
||||
# Park, waiting for notification that prefetching
|
||||
# should resume
|
||||
self._condition.wait()
|
||||
|
||||
blocks_to_fetch, self._blocks = self._blocks[:], []
|
||||
|
||||
if len(blocks_to_fetch) > 0:
|
||||
ray.wait(
|
||||
blocks_to_fetch,
|
||||
num_returns=1,
|
||||
# NOTE: We deliberately setting timeout to 0 to avoid
|
||||
# blocking the fetching thread unnecessarily
|
||||
timeout=0,
|
||||
fetch_local=True,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Error in prefetcher thread.")
|
||||
|
||||
logger.debug("Exiting prefetcher's background thread")
|
||||
|
||||
def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
|
||||
with self._condition:
|
||||
if self._stopped:
|
||||
raise RuntimeError("Prefetcher is stopped.")
|
||||
self._blocks = blocks
|
||||
self._condition.notify()
|
||||
|
||||
def stop(self):
|
||||
with self._condition:
|
||||
if self._stopped:
|
||||
return
|
||||
self._stopped = True
|
||||
self._condition.notify()
|
||||
|
||||
def __del__(self):
|
||||
self.stop()
|
||||
|
||||
|
||||
class ActorBlockPrefetcher(BlockPrefetcher):
|
||||
"""Block prefetcher using a local actor."""
|
||||
|
||||
def __init__(self):
|
||||
self.prefetch_actor = self._get_or_create_actor_prefetcher()
|
||||
|
||||
@staticmethod
|
||||
def _get_or_create_actor_prefetcher() -> "ActorHandle":
|
||||
node_id = ray.get_runtime_context().get_node_id()
|
||||
actor_name = f"dataset-block-prefetcher-{node_id}"
|
||||
return _BlockPretcher.options(
|
||||
label_selector={ray._raylet.RAY_NODE_ID_KEY: node_id},
|
||||
name=actor_name,
|
||||
namespace=PREFETCHER_ACTOR_NAMESPACE,
|
||||
get_if_exists=True,
|
||||
).remote()
|
||||
|
||||
def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
|
||||
self.prefetch_actor.prefetch.remote(*blocks)
|
||||
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
class _BlockPretcher:
|
||||
"""Helper actor that prefetches blocks asynchronously."""
|
||||
|
||||
def prefetch(self, *blocks) -> None:
|
||||
pass
|
||||
Reference in New Issue
Block a user