chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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