wehub-resource-sync
462 lines
21 KiB
Python
462 lines
21 KiB
Python
"""Metadata-fetch strategy for the streaming executor.
|
|
|
|
``DataOpTask.on_data_ready`` pulls ``(block_ref, meta_ref)`` pairs from a task's
|
|
streaming generator; a ``MetadataFetcher`` turns each pair into an emitted
|
|
``RefBundle``. Two modes, selected by ``RAY_DATA_METADATA_PREFETCH_ON_THREAD``
|
|
(default on):
|
|
|
|
- :class:`ThreadedMetadataFetcher` (default): defer every pair and fetch its
|
|
metadata on a dedicated background thread, so the scheduling loop never blocks
|
|
on ``ray.get(meta_ref)``. The output-budget size comes from the block's local
|
|
``object_size`` (no RPC); completion is postponed until the task's deferred
|
|
pairs have emitted, and the per-operator FIFO preserves emission order.
|
|
- :class:`InlineMetadataFetcher`: fetch each pair's metadata inline with
|
|
``ray.get`` and emit the ``RefBundle`` immediately, budgeting off
|
|
``meta.size_bytes``; completion and task-failure are handled inline by
|
|
``on_data_ready``.
|
|
"""
|
|
|
|
import logging
|
|
import pickle
|
|
import queue as queue_module
|
|
import threading
|
|
from abc import ABC, abstractmethod
|
|
from collections import defaultdict, deque
|
|
from collections.abc import Hashable
|
|
from enum import Enum
|
|
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
|
|
|
|
import ray
|
|
import ray.exceptions
|
|
from ray._common.utils import env_bool
|
|
from ray.data._internal.execution.interfaces.physical_operator import (
|
|
METADATA_GET_TIMEOUT_S,
|
|
METADATA_WAIT_TIMEOUT_S,
|
|
DataOpTask,
|
|
DeferredEmit,
|
|
)
|
|
from ray.data.block import BlockMetadataWithSchema
|
|
from ray.exceptions import GetTimeoutError
|
|
from ray.experimental.locations import get_local_object_locations
|
|
from ray.util.debug import log_once
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# How long ``ThreadedMetadataFetcher.stop`` waits for the fetch thread to exit.
|
|
_FETCH_THREAD_JOIN_TIMEOUT_S = 5.0
|
|
|
|
# How long the fetch thread's ``ray.wait`` blocks each pass — bounds the
|
|
# busy-wait when nothing is ready, and how long a straggler can delay a batch.
|
|
_FETCH_WAIT_TIMEOUT_S = 0.1
|
|
|
|
# Selects the mode (see module docstring). Threaded by default; set to 0/false to
|
|
# fall back to the synchronous inline path.
|
|
_PREFETCH_ON_THREAD = env_bool("RAY_DATA_METADATA_PREFETCH_ON_THREAD", True)
|
|
|
|
|
|
class MetadataFetcher(ABC):
|
|
def start(self) -> None:
|
|
"""Start any background machinery."""
|
|
|
|
def stop(self) -> None:
|
|
"""Stop any background machinery."""
|
|
|
|
@abstractmethod
|
|
def in_data_ready_get_object_size(self, task: DataOpTask) -> Optional[int]:
|
|
"""Handle one pulled pair inside the ``on_data_ready`` loop. The pair's
|
|
refs are read off the task (``task.pending_block_ref`` /
|
|
``task.pending_meta_ref``).
|
|
|
|
Returns the output-budget bytes for this pair (0 if the size is
|
|
unknown), or ``None`` to mean "the metadata isn't available yet — stop
|
|
and retry next iteration" (the caller breaks, leaving the refs set).
|
|
|
|
``None`` must be returned ONLY when the pair was NOT consumed: the
|
|
caller will hand the same pair back on the next call, so returning
|
|
``None`` after emitting/deferring it would emit the block twice.
|
|
"""
|
|
|
|
def in_data_ready_done(self, task: DataOpTask) -> None:
|
|
"""Called once a task is drained (generator exhausted/failed)."""
|
|
|
|
def submit(self, op_key: Hashable, tasks: List[DataOpTask]) -> None:
|
|
"""Hand the operator's deferred pairs off for processing, and record
|
|
any end-of-stream/failed tasks whose completion is postponed."""
|
|
|
|
def emit_ready_and_fire_done_callbacks(self) -> List[Tuple[str, BaseException]]:
|
|
"""Run once at the end of ``process_completed_tasks``. Returns
|
|
``(operator_name, exception)`` for each block-level fetch failure, for
|
|
the caller's ``max_errored_blocks`` accounting. Default: nothing to do."""
|
|
return []
|
|
|
|
|
|
class InlineMetadataFetcher(MetadataFetcher):
|
|
"""Synchronous mode: fetch metadata inline and emit immediately. Holds no
|
|
state and starts no thread."""
|
|
|
|
def in_data_ready_get_object_size(self, task: DataOpTask) -> Optional[int]:
|
|
# The ref resolves to pickled metadata bytes, not a BlockMetadata.
|
|
meta_ref: "ray.ObjectRef[Any]" = task.pending_meta_ref
|
|
try:
|
|
# The timeout includes the time to ship the metadata to this node,
|
|
# so a 0 timeout could cancel an in-flight download. Use a small
|
|
# non-zero value to avoid that.
|
|
meta_bytes: bytes = ray.get(meta_ref, timeout=METADATA_GET_TIMEOUT_S)
|
|
except ray.exceptions.GetTimeoutError:
|
|
# We have refs to the block and its metadata, but the metadata
|
|
# object isn't available. This can happen if the node dies. Leave
|
|
# the pair pending and retry next iteration.
|
|
logger.warning(
|
|
f"Timed out ({METADATA_GET_TIMEOUT_S}s) waiting for metadata from "
|
|
f"operator '{task.operator_name}' "
|
|
f"(metadata_ref={meta_ref.hex()}). "
|
|
f"Possible causes include a worker crash, node preemption, or an "
|
|
f"overloaded worker or head node. Will retry next iteration. "
|
|
f"If this repeats, check the Ray dashboard and logs for worker "
|
|
f"crashes, node preemption, or overload."
|
|
)
|
|
return None
|
|
return task.produce_block(task.pending_block_ref, meta_bytes)
|
|
|
|
def in_data_ready_done(self, task: DataOpTask) -> None:
|
|
# Inline mode fires the done-callback the moment the generator drains:
|
|
# all of the task's pairs have already emitted inline, so
|
|
# ``_pending_emit_count`` is 0. A task failure is re-raised after the
|
|
# callback.
|
|
task.mark_done()
|
|
if task.task_error is not None:
|
|
raise task.task_error from None
|
|
|
|
|
|
class _Signal(Enum):
|
|
"""Sentinels used by :class:`ThreadedMetadataFetcher`.
|
|
|
|
``STOP`` is enqueued on the request queue to tell the fetch thread to exit;
|
|
``NOT_READY`` marks "ref not fetched yet" in the result store. Members of a
|
|
single enum so identity checks narrow cleanly under type checkers.
|
|
"""
|
|
|
|
STOP = "stop"
|
|
NOT_READY = "not_ready"
|
|
|
|
|
|
# A request-queue item: a batch of meta_refs to fetch, or the stop sentinel.
|
|
_Request = Union[List["ray.ObjectRef"], _Signal]
|
|
|
|
|
|
class ThreadedMetadataFetcher(MetadataFetcher):
|
|
"""Asynchronous mode: defer every pulled pair and fetch its metadata on a
|
|
dedicated background thread, so the scheduling (executor) thread never blocks
|
|
on ``ray.get(meta_refs)``.
|
|
|
|
The two threads communicate through one thread-safe queue (``_request_q``);
|
|
fetched bytes come back via ``_results``. The background thread fetches the
|
|
refs ``ray.wait(fetch_local=True)`` reports ready; a ref stuck on a bad node
|
|
merely stays pending instead of wedging the thread.
|
|
|
|
Data flow (for a single operator)::
|
|
|
|
Executor thread Fetch thread
|
|
--------------- ------------
|
|
on_data_ready --defer--> _pending_deferred
|
|
submit(op) --meta_refs--> _request_q -----> ray.wait(ready)
|
|
+ ray.get
|
|
|
|
|
_results <----- fetched bytes --------+
|
|
emit_ready_and_fire_done_callbacks():
|
|
_fifos[op]: head [d0] -> [d1] -> [d2] tail (append = yield order)
|
|
|
|
|
`- emit front-first while its bytes are in
|
|
_results; stop at the first pair not back yet,
|
|
so this op's RefBundle order is preserved.
|
|
|
|
Operators each get their own FIFO and are independent, so one operator
|
|
waiting on metadata never blocks another.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
get_objects: Optional[Callable] = None,
|
|
wait_for_objects: Optional[Callable] = None,
|
|
get_object_locations: Optional[Callable] = None,
|
|
):
|
|
"""Create a ThreadedMetadataFetcher.
|
|
|
|
Args:
|
|
get_objects: Fetches object values by ref, like ``ray.get``
|
|
(called as ``get_objects(refs, timeout=...)``). Injectable so
|
|
tests can drive the fetch path without a real cluster.
|
|
wait_for_objects: Reports which refs are locally available, like
|
|
``ray.wait`` (called with ``fetch_local=True``).
|
|
get_object_locations: Returns per-ref location info (including
|
|
``object_size``), like ``get_local_object_locations``.
|
|
"""
|
|
self._get_objects: Callable = get_objects or ray.get
|
|
self._wait_for_objects: Callable = wait_for_objects or ray.wait
|
|
self._get_object_locations: Callable = (
|
|
get_object_locations or get_local_object_locations
|
|
)
|
|
|
|
self._request_q: "queue_module.Queue[_Request]" = queue_module.Queue()
|
|
# fetch thread -> executor: meta_ref -> bytes (or captured Exception).
|
|
self._results: Dict["ray.ObjectRef", Any] = {}
|
|
self._results_lock = threading.Lock()
|
|
|
|
# Executor-thread-only state below.
|
|
# Pairs deferred by ``in_data_ready_get_object_size`` for the current operator,
|
|
# flushed into the FIFOs by ``submit``.
|
|
self._pending_deferred: List[DeferredEmit] = []
|
|
# Per-operator (keyed by the caller's op key) FIFO of pairs awaiting
|
|
# metadata, in append (= emission) order. Each op's deque is drained
|
|
# front-first so that op's RefBundle emission order is preserved.
|
|
self._fifos: "defaultdict[Hashable, deque[DeferredEmit]]" = defaultdict(deque)
|
|
# Drained (end-of-stream/failed) tasks whose done-callback is postponed
|
|
# until all of their deferred pairs have been emitted. A set so a task
|
|
# re-seen on a later iteration (still pending) isn't registered — or
|
|
# fired — twice.
|
|
self._drained_tasks: Set[DataOpTask] = set()
|
|
|
|
self._thread = threading.Thread(
|
|
target=self._run, name="ray-data-metadata-prefetch", daemon=True
|
|
)
|
|
self._started = False
|
|
self._stopped = False
|
|
|
|
def start(self) -> None:
|
|
if not self._started:
|
|
self._started = True
|
|
self._thread.start()
|
|
|
|
def stop(self) -> None:
|
|
if self._stopped:
|
|
return
|
|
self._stopped = True
|
|
self._request_q.put(_Signal.STOP)
|
|
if self._started:
|
|
try:
|
|
self._thread.join(timeout=_FETCH_THREAD_JOIN_TIMEOUT_S)
|
|
if self._thread.is_alive():
|
|
logger.warning(
|
|
"Metadata-fetch thread did not exit within "
|
|
f"{_FETCH_THREAD_JOIN_TIMEOUT_S}s; leaving the daemon "
|
|
"thread behind."
|
|
)
|
|
except Exception:
|
|
logger.warning(
|
|
"Failed to join the metadata-fetch thread.", exc_info=True
|
|
)
|
|
|
|
def in_data_ready_get_object_size(self, task: DataOpTask) -> Optional[int]:
|
|
block_ref = task.pending_block_ref
|
|
meta_ref = task.pending_meta_ref
|
|
# Output-budget size from the block's local object_size (no RPC).
|
|
# Normally known: the driver owns the just-yielded block ref, so the
|
|
# value (which matches ``meta.size_bytes``) is in the local object
|
|
# directory.
|
|
# TODO: ``object_size`` is the object-store size of the block, which can
|
|
# differ from ``meta.size_bytes`` (the in-memory/logical size). We should
|
|
# add an explicit ``object_size_bytes`` to ``BlockMetadata`` and use it
|
|
# directly, so the fallback below doesn't conflate the two.
|
|
info: Optional[Dict[str, Any]] = self._get_object_locations([block_ref]).get(
|
|
block_ref
|
|
)
|
|
object_size: Optional[int] = (
|
|
info.get("object_size") if info is not None else None
|
|
)
|
|
if object_size is None:
|
|
# Rare: no local size record. Fall back to a short metadata
|
|
# ``ray.get`` for the size. Log once to flag the path without
|
|
# spamming if it recurs.
|
|
if log_once(f"data_object_size_unavailable_{task.operator_name}"):
|
|
logger.warning(
|
|
"Local object_size unavailable for a block from operator "
|
|
"'%s'; falling back to its metadata for the output-budget "
|
|
"size.",
|
|
task.operator_name,
|
|
)
|
|
try:
|
|
meta_with_schema: BlockMetadataWithSchema = pickle.loads(
|
|
self._get_objects(meta_ref, timeout=METADATA_WAIT_TIMEOUT_S)
|
|
)
|
|
except ray.exceptions.GetTimeoutError:
|
|
# Metadata isn't local yet either. Leave this pair pending and
|
|
# retry next iteration.
|
|
return None
|
|
# Coalesce a missing size to 0: None is reserved for the "pair not
|
|
# consumed, retry" signal above, and this pair IS consumed
|
|
# (deferred) below — returning None here would defer it twice.
|
|
object_size = meta_with_schema.metadata.size_bytes or 0
|
|
|
|
self._pending_deferred.append(
|
|
DeferredEmit(task=task, block_ref=block_ref, meta_ref=meta_ref)
|
|
)
|
|
return object_size
|
|
|
|
def submit(self, op_key: Hashable, tasks: List[DataOpTask]) -> None:
|
|
"""Queue the current operator's deferred pairs for metadata fetch +
|
|
emission — append them to the op's FIFO (preserving emission order) and
|
|
hand their ``meta_ref``s to the fetch thread — and record any drained
|
|
(end-of-stream/failed) tasks so their done-callback fires once all of
|
|
their deferred pairs have emitted. Must run on the executor thread."""
|
|
deferred = self._pending_deferred
|
|
self._pending_deferred = []
|
|
if deferred:
|
|
fifo = self._fifos[op_key]
|
|
for d in deferred:
|
|
d.task.add_pending_metadata_ref()
|
|
fifo.append(d)
|
|
self._request_q.put([d.meta_ref for d in deferred])
|
|
for task in tasks:
|
|
if task.is_drained():
|
|
self._drained_tasks.add(task)
|
|
|
|
def emit_ready_and_fire_done_callbacks(self) -> List[Tuple[str, BaseException]]:
|
|
"""Emit whatever's ready (per-op order) then fire postponed done
|
|
callbacks. Returns ``(operator_name, exception)`` for each pair whose
|
|
metadata fetch failed, for the caller's ``max_errored_blocks``
|
|
accounting. Must run on the executor thread."""
|
|
return self._emit_ready() + self._fire_done_callbacks()
|
|
|
|
def _emit_ready(self) -> List[Tuple[str, BaseException]]:
|
|
# Emit every pair whose metadata is now available, in per-op append
|
|
# order. A failed fetch is accounted as emitted (so the task can still
|
|
# complete) but its block is dropped and the error is surfaced to the
|
|
# caller rather than raised.
|
|
failures: List[Tuple[str, BaseException]] = []
|
|
for fifo in self._fifos.values():
|
|
while fifo:
|
|
d = fifo[0]
|
|
result = self._pop_result(d.meta_ref)
|
|
if result is _Signal.NOT_READY:
|
|
# Preserve order: stop at the first pair still in flight;
|
|
# this operator is retried next call.
|
|
# TODO: order only needs to be preserved when
|
|
# ``DataContext.get_current().execution_options.preserve_order``
|
|
# is True; otherwise we could skip past in-flight pairs and
|
|
# emit any ready ones.
|
|
break
|
|
fifo.popleft()
|
|
d.task.mark_emitted()
|
|
if isinstance(result, BaseException):
|
|
failures.append((d.task.operator_name, result))
|
|
continue
|
|
try:
|
|
d.task.produce_block(d.block_ref, result)
|
|
except Exception as e:
|
|
# Deserializing/emitting the fetched metadata can also fail
|
|
# (e.g. ``pickle.loads`` raising on a corrupt object). Treat
|
|
# it as a block-level error and route it through the same
|
|
# accounting, rather than letting it escape.
|
|
failures.append((d.task.operator_name, e))
|
|
return failures
|
|
|
|
def _fire_done_callbacks(self) -> List[Tuple[str, BaseException]]:
|
|
# Fire postponed done-callbacks for drained tasks whose pairs have all
|
|
# emitted. A failed task fires with its error, which is also surfaced for
|
|
# ``max_errored_blocks`` accounting.
|
|
if not self._drained_tasks:
|
|
return []
|
|
failures: List[Tuple[str, BaseException]] = []
|
|
to_mark_done = [t for t in self._drained_tasks if not t.has_pending_emits()]
|
|
for task in to_mark_done:
|
|
if task.task_error is not None:
|
|
failures.append((task.operator_name, task.task_error))
|
|
task.mark_done()
|
|
self._drained_tasks.difference_update(to_mark_done)
|
|
return failures
|
|
|
|
def _pop_result(self, ref: "ray.ObjectRef") -> Any:
|
|
with self._results_lock:
|
|
return self._results.pop(ref, _Signal.NOT_READY)
|
|
|
|
def _run(self) -> None:
|
|
# Fetch-thread loop: accumulate requested meta_refs into a pending set
|
|
# and hand them to ``_fetch``, which fetches the locally-available ones
|
|
# and returns those still in flight.
|
|
pending: List["ray.ObjectRef"] = []
|
|
while True:
|
|
# Block on the queue only when idle; while refs are in flight,
|
|
# don't block here — get back to ``_fetch`` to keep them moving.
|
|
try:
|
|
item = self._request_q.get(block=not pending)
|
|
except queue_module.Empty:
|
|
item = None
|
|
# Drain whatever else is already queued into a single fetch batch.
|
|
while item is not None:
|
|
if isinstance(item, _Signal):
|
|
assert item is _Signal.STOP
|
|
# ``stop()`` enqueued the STOP sentinel: fast teardown —
|
|
# drop any in-flight refs and exit. ``stop`` runs after the
|
|
# scheduling loop (which feeds us) is joined, so there's
|
|
# nothing left to emit.
|
|
return
|
|
pending.extend(item)
|
|
try:
|
|
item = self._request_q.get_nowait()
|
|
except queue_module.Empty:
|
|
item = None
|
|
|
|
if pending:
|
|
pending = self._fetch(pending)
|
|
|
|
def _fetch(self, pending: List["ray.ObjectRef"]) -> List["ray.ObjectRef"]:
|
|
"""One fetch pass over ``pending``:
|
|
|
|
1. ``ray.wait(fetch_local=True)`` pulls the metadata objects to this
|
|
(driver) node and reports which are locally available.
|
|
2. ``ray.get`` the ready refs in one batch (``timeout=0`` — they're
|
|
local) and publish the bytes to ``_results``.
|
|
3. If the batched get raises (it hides which ref failed), fall back to
|
|
per-ref gets to isolate the failure and keep the rest.
|
|
4. Return the refs to retry next pass: the not-yet-local ones, plus any
|
|
that raced out of the local store. A ref that resolved to an error is
|
|
published as that exception for ``_emit_ready`` to surface.
|
|
"""
|
|
ready, not_ready = self._wait_for_objects(
|
|
pending,
|
|
num_returns=len(pending),
|
|
timeout=_FETCH_WAIT_TIMEOUT_S,
|
|
fetch_local=True,
|
|
)
|
|
if not ready:
|
|
return not_ready
|
|
|
|
retry: List["ray.ObjectRef"] = []
|
|
try:
|
|
values = self._get_objects(ready, timeout=0)
|
|
results: Dict["ray.ObjectRef", Any] = dict(zip(ready, values))
|
|
except Exception:
|
|
# A batched get raises on the first error and hides which ref
|
|
# failed; retry per-ref to isolate it and keep the rest.
|
|
results = {}
|
|
for ref in ready:
|
|
try:
|
|
results[ref] = self._get_objects(ref, timeout=0)
|
|
except GetTimeoutError:
|
|
# ray.wait reported it ready but it's no longer local (e.g. a
|
|
# raced eviction). Re-queue rather than treating it as a
|
|
# block-level error. Shouldn't be common — log once.
|
|
if log_once("ray_data_metadata_prefetch_not_local"):
|
|
logger.warning(
|
|
"A metadata object reported ready by ray.wait was "
|
|
"not locally available for ray.get; re-queuing it. "
|
|
"If this repeats, the object store may be under "
|
|
"memory pressure (objects evicted/spilled)."
|
|
)
|
|
retry.append(ref)
|
|
except Exception as e:
|
|
results[ref] = e
|
|
if results:
|
|
with self._results_lock:
|
|
self._results.update(results)
|
|
return not_ready + retry
|
|
|
|
|
|
def make_metadata_fetcher() -> MetadataFetcher:
|
|
"""Build the metadata fetcher for the configured mode (see module
|
|
docstring)."""
|
|
if _PREFETCH_ON_THREAD:
|
|
return ThreadedMetadataFetcher()
|
|
return InlineMetadataFetcher()
|