wehub-resource-sync
135 lines
4.1 KiB
Python
135 lines
4.1 KiB
Python
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
|