wehub-resource-sync
562 lines
18 KiB
Python
562 lines
18 KiB
Python
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
|