wehub-resource-sync
738 lines
27 KiB
Python
738 lines
27 KiB
Python
import functools
|
|
import logging
|
|
import pickle
|
|
import time
|
|
import typing
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Tuple,
|
|
Union,
|
|
)
|
|
|
|
import ray
|
|
import ray.exceptions
|
|
from ray.actor import ActorHandle
|
|
from ray.data import ExecutionOptions
|
|
from ray.data._internal.execution.bundle_queue import ReorderingBundleQueue
|
|
from ray.data._internal.execution.interfaces import (
|
|
BlockEntry,
|
|
ExecutionResources,
|
|
PhysicalOperator,
|
|
RefBundle,
|
|
)
|
|
from ray.data._internal.execution.interfaces.physical_operator import (
|
|
DataOpTask,
|
|
MetadataOpTask,
|
|
OpTask,
|
|
estimate_total_num_of_blocks,
|
|
)
|
|
from ray.data._internal.execution.operators.hash_shuffle import (
|
|
_get_total_cluster_resources,
|
|
)
|
|
from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
|
|
from ray.data._internal.stats import OpRuntimeMetrics
|
|
from ray.data.block import Block, BlockAccessor, BlockStats, to_stats
|
|
from ray.data.context import DataContext
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from ray.data._internal.execution.block_ref_counter import BlockRefCounter
|
|
from ray.data._internal.execution.interfaces.physical_operator import ActorPoolInfo
|
|
from ray.data._internal.progress.base_progress import BaseProgressBar
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Arrow schema metadata key for the rapidsmpf partition ID.
|
|
_GPU_PARTITION_ID_KEY = b"_gpu_partition_id"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU shuffle actor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@ray.remote(num_gpus=1)
|
|
class GPUShuffleActor:
|
|
"""One GPU rank in a RAPIDS MPF-based distributed shuffle.
|
|
|
|
Each instance wraps a ``BulkRapidsMPFShuffler`` via composition rather than
|
|
inheritance to keep CPU-only environments unaffected.
|
|
|
|
Actors are arranged in a virtual communicator ring coordinated
|
|
through UCXX; data never passes through the Ray object store or the CPU
|
|
after initial ingestion.
|
|
|
|
Constructor is intentionally lightweight — expensive UCXX setup happens in
|
|
:meth:`setup_worker`, which is called once from :class:`GPURankPool`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
nranks: int,
|
|
total_nparts: int,
|
|
key_columns: List[str],
|
|
columns: Optional[List[str]] = None,
|
|
rmm_pool_size: Union[int, str, None] = None,
|
|
spill_memory_limit: Union[int, str, None] = "auto",
|
|
should_sort: bool = False,
|
|
):
|
|
from ray.data._internal.gpu_shuffle.rapidsmpf_backend import (
|
|
BulkRapidsMPFShuffler,
|
|
)
|
|
|
|
self._shuffler = BulkRapidsMPFShuffler(
|
|
nranks=nranks,
|
|
total_nparts=total_nparts,
|
|
shuffle_on=key_columns,
|
|
rmm_pool_size=rmm_pool_size,
|
|
spill_memory_limit=spill_memory_limit,
|
|
)
|
|
self._columns = columns
|
|
self._key_columns = key_columns
|
|
self._should_sort = should_sort
|
|
self._arrow_schema = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# UCXX communicator setup
|
|
# ------------------------------------------------------------------
|
|
|
|
def setup_root(self) -> tuple[int, bytes]:
|
|
"""Initialize the root communicator and return ``(rank, root_address_bytes)``.
|
|
|
|
Only called on rank 0; the returned address is broadcast to all ranks
|
|
via :meth:`setup_worker`.
|
|
"""
|
|
logger.info("UCXX setup_root starting on rank 0.")
|
|
t0 = time.perf_counter()
|
|
result = self._shuffler.setup_root()
|
|
elapsed = time.perf_counter() - t0
|
|
logger.info("UCXX setup_root completed in %.2fs (rank=%d).", elapsed, result[0])
|
|
return result
|
|
|
|
def setup_worker(self, root_address: bytes) -> None:
|
|
"""Finish UCXX communicator setup and create the internal shuffler.
|
|
|
|
Must be called on *every* rank (including rank 0) after
|
|
:meth:`get_root_address` has been called on rank 0 and its result
|
|
broadcast to all ranks.
|
|
"""
|
|
logger.info(
|
|
"UCXX setup_worker starting (root_address=%d bytes).",
|
|
len(root_address),
|
|
)
|
|
t0 = time.perf_counter()
|
|
self._shuffler.setup_worker(root_address)
|
|
elapsed = time.perf_counter() - t0
|
|
logger.info("UCXX setup_worker completed in %.2fs.", elapsed)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Insert / extract interface (called by GPUShuffleOperator)
|
|
# ------------------------------------------------------------------
|
|
|
|
def insert_batch(self, block: Block) -> int:
|
|
"""Hash-partition *block* and route shards to peers.
|
|
|
|
Returns the number of rows in the incoming block so the driver can
|
|
track throughput without serialising the data back.
|
|
"""
|
|
import cudf
|
|
|
|
table = BlockAccessor.for_block(block).to_arrow()
|
|
df = cudf.DataFrame.from_arrow(table)
|
|
if self._columns is None:
|
|
# save columns from first batch, if not already set
|
|
self._columns = list(df.columns)
|
|
if self._arrow_schema is None:
|
|
# save arrow schema from first batch
|
|
self._arrow_schema = table.schema
|
|
self._shuffler.insert_chunk(table=df, column_names=self._columns)
|
|
return len(df)
|
|
|
|
def finish_and_extract(self) -> Iterator:
|
|
"""Signal insertion is done, then yield one Arrow Table per output partition.
|
|
|
|
Combines insert-finished and extraction into a single actor call so
|
|
correctness does not depend on actor task ordering.
|
|
|
|
Follows the Ray Data streaming generator protocol: yield block then
|
|
BlockMetadataWithSchema for each output partition.
|
|
|
|
The partition ID from ``rapidsmpf``'s ``extract()`` is embedded in each
|
|
block's Arrow schema metadata under ``_gpu_partition_id`` so the operator
|
|
can reorder bundles into correct partition order on the driver side,
|
|
regardless of GPU completion order.
|
|
"""
|
|
self._shuffler.insert_finished()
|
|
|
|
import pyarrow as pa
|
|
from rapidsmpf.utils.cudf import pylibcudf_to_cudf_dataframe
|
|
|
|
from ray.data.block import BlockExecStats, BlockMetadataWithSchema
|
|
|
|
for partition_id, partition in self._shuffler.extract():
|
|
exec_stats_builder = BlockExecStats.builder()
|
|
if partition.num_columns() == 0:
|
|
# rapidsmpf returns a zero-column table when no rows were
|
|
# routed to this partition. Emit an empty arrow table, so every
|
|
# partition id produces exactly one block, so downstream queues
|
|
# that require contiguous key ranges (e.g. ReorderingBundleQueue)
|
|
# don't stall.
|
|
block = pa.Table.from_pylist([], schema=self._arrow_schema)
|
|
else:
|
|
cdf = pylibcudf_to_cudf_dataframe(
|
|
partition, column_names=self._columns
|
|
).copy(deep=True)
|
|
# Caveat: The following operation copies the data to CPU memory, unless we use Arrow CUDA.
|
|
if self._should_sort and len(cdf) > 0:
|
|
cdf = cdf.sort_values(by=self._key_columns)
|
|
block = cdf.to_arrow(preserve_index=False)
|
|
|
|
existing_metadata = block.schema.metadata or {}
|
|
tagged_schema = block.schema.with_metadata(
|
|
{**existing_metadata, _GPU_PARTITION_ID_KEY: str(partition_id).encode()}
|
|
)
|
|
exec_stats = exec_stats_builder.build()
|
|
stats = yield block
|
|
if stats:
|
|
object.__setattr__(
|
|
exec_stats, "block_ser_time_s", stats.object_creation_dur_s
|
|
)
|
|
block_meta = BlockMetadataWithSchema.from_block(
|
|
block, block_exec_stats=exec_stats
|
|
)
|
|
bm = BlockMetadataWithSchema.from_metadata(
|
|
block_meta.metadata, schema=tagged_schema
|
|
)
|
|
yield pickle.dumps(bm)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPURankPool — lifecycle manager for a set of GPUShuffleActors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class GPURankPool:
|
|
"""Manages the lifecycle of ``GPUShuffleActor`` instances.
|
|
|
|
Analogous to ``AggregatorPool`` in the CPU hash-shuffle path, but for GPU
|
|
ranks coordinated through UCXX.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
nranks: int,
|
|
total_nparts: int,
|
|
setup_timeout_s: float,
|
|
actor_cls_factory: Callable[[], Any],
|
|
actor_kwargs: Dict[str, Any],
|
|
log_label: str,
|
|
label_selector: Optional[Dict[str, str]] = None,
|
|
) -> None:
|
|
self._nranks = nranks
|
|
self._total_nparts = total_nparts
|
|
self._setup_timeout_s = setup_timeout_s
|
|
self._actor_cls_factory = actor_cls_factory
|
|
self._actor_kwargs = actor_kwargs
|
|
self._log_label = log_label
|
|
self._label_selector = label_selector
|
|
self._actors: List[ActorHandle] = []
|
|
self._shutdown: bool = False
|
|
|
|
@property
|
|
def is_shutdown(self) -> bool:
|
|
return self._shutdown
|
|
|
|
@property
|
|
def nranks(self) -> int:
|
|
return self._nranks
|
|
|
|
@property
|
|
def actors(self) -> List[ActorHandle]:
|
|
return self._actors
|
|
|
|
def start(self) -> None:
|
|
timeout = self._setup_timeout_s
|
|
t_start = time.perf_counter()
|
|
|
|
logger.info(
|
|
"%s: creating %d actor(s) (total_nparts=%d).",
|
|
self._log_label,
|
|
self._nranks,
|
|
self._total_nparts,
|
|
)
|
|
actor_cls = self._actor_cls_factory()
|
|
actor_options: Dict[str, typing.Any] = {
|
|
"num_gpus": 1,
|
|
"scheduling_strategy": "SPREAD",
|
|
}
|
|
if self._label_selector:
|
|
actor_options["label_selector"] = self._label_selector
|
|
self._actors = [
|
|
actor_cls.options(**actor_options).remote(
|
|
nranks=self._nranks,
|
|
total_nparts=self._total_nparts,
|
|
**self._actor_kwargs,
|
|
)
|
|
for _ in range(self._nranks)
|
|
]
|
|
t_actors = time.perf_counter()
|
|
logger.info(
|
|
"%s: %d actor(s) created in %.2fs.",
|
|
self._log_label,
|
|
self._nranks,
|
|
t_actors - t_start,
|
|
)
|
|
|
|
remaining = max(0, timeout - (time.perf_counter() - t_start))
|
|
logger.info("%s: calling setup_root on rank 0.", self._log_label)
|
|
try:
|
|
_, root_address_bytes = ray.get(
|
|
self._actors[0].setup_root.remote(), timeout=remaining
|
|
)
|
|
except ray.exceptions.GetTimeoutError:
|
|
raise TimeoutError(
|
|
f"UCXX setup_root on {self._log_label} rank 0 did not complete "
|
|
f"within {timeout}s. Check GPU/network health."
|
|
)
|
|
t_root = time.perf_counter()
|
|
logger.info(
|
|
"%s: setup_root completed in %.2fs, "
|
|
"broadcasting root address (%d bytes) to %d worker(s).",
|
|
self._log_label,
|
|
t_root - t_actors,
|
|
len(root_address_bytes),
|
|
self._nranks,
|
|
)
|
|
|
|
remaining = max(0, timeout - (time.perf_counter() - t_start))
|
|
worker_refs = [
|
|
actor.setup_worker.remote(root_address_bytes) for actor in self._actors
|
|
]
|
|
self._wait_for_refs_with_timeout(worker_refs, remaining, "setup_worker")
|
|
t_done = time.perf_counter()
|
|
logger.info(
|
|
"%s: all %d worker(s) setup completed in %.2fs "
|
|
"(total UCXX init: %.2fs).",
|
|
self._log_label,
|
|
self._nranks,
|
|
t_done - t_root,
|
|
t_done - t_start,
|
|
)
|
|
|
|
def get_actor_for_block(self, block_idx: int) -> ActorHandle:
|
|
"""Round-robin distribution of input blocks across ranks."""
|
|
return self._actors[block_idx % self._nranks]
|
|
|
|
def shutdown(self, force: bool = False) -> None:
|
|
if force:
|
|
for actor in self._actors:
|
|
ray.kill(actor)
|
|
self._actors.clear()
|
|
self._shutdown = True
|
|
|
|
def _wait_for_refs_with_timeout(
|
|
self,
|
|
refs: List[ray.ObjectRef],
|
|
timeout_s: float,
|
|
task_name: str,
|
|
) -> None:
|
|
"""Poll ``refs`` in a loop, raising on timeout or task failure."""
|
|
total = len(refs)
|
|
pending = list(refs)
|
|
t_start = time.perf_counter()
|
|
|
|
while pending:
|
|
elapsed = time.perf_counter() - t_start
|
|
if elapsed >= timeout_s:
|
|
pending_indices = [i for i, ref in enumerate(refs) if ref in pending]
|
|
raise TimeoutError(
|
|
f"{task_name} did not complete on {len(pending)}/{total} "
|
|
f"rank(s) within {timeout_s}s "
|
|
f"(pending ranks: {pending_indices}). "
|
|
f"Check GPU/network health."
|
|
)
|
|
ready, pending = ray.wait(
|
|
pending, num_returns=len(pending), timeout=min(0.1, timeout_s - elapsed)
|
|
)
|
|
if ready:
|
|
ray.get(ready)
|
|
logger.info(
|
|
"%s: %d/%d rank(s) completed %s.",
|
|
self._log_label,
|
|
total - len(pending),
|
|
total,
|
|
task_name,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper: derive number of GPU ranks from the cluster
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _derive_num_gpu_ranks(data_context: DataContext) -> int:
|
|
"""Return the configured or auto-detected number of GPU ranks."""
|
|
if data_context.gpu_shuffle_num_actors is not None:
|
|
return data_context.gpu_shuffle_num_actors
|
|
|
|
total_resources = _get_total_cluster_resources()
|
|
num_gpus = int(total_resources.gpu or 0)
|
|
if num_gpus == 0:
|
|
raise RuntimeError(
|
|
"ShuffleStrategy.GPU_SHUFFLE requires GPU resources in the cluster. "
|
|
"Set DataContext.gpu_shuffle_num_actors to override the number of ranks."
|
|
)
|
|
return num_gpus
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUShuffleOperator
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class GPUShuffleOperator(PhysicalOperator, SubProgressBarMixin):
|
|
"""GPU-native shuffle operator using RAPIDS MPF + UCXX.
|
|
|
|
Unlike the CPU ``HashShuffleOperator``, this operator:
|
|
|
|
* Uses UCXX point-to-point communication instead of the Ray object store
|
|
for inter-rank data movement.
|
|
* Accepts Arrow Tables from upstream (converting to cuDF on the actor) so
|
|
it remains compatible with non-GPU upstream operators.
|
|
* Supports repartition-only (no reduce/aggregate phase on the driver side).
|
|
|
|
Lifecycle::
|
|
|
|
start() # creates actors, blocks for UCXX setup
|
|
_add_input_inner(bundle) # routes blocks to actors round-robin
|
|
[inputs_done()] # called by the executor
|
|
has_next() / _get_next_inner() # streams output bundles
|
|
|
|
The ``finish_and_extract`` actor task is submitted once all inserts
|
|
complete; it signals insertion done and streams output partitions in a
|
|
single call.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_op: PhysicalOperator,
|
|
data_context: DataContext,
|
|
*,
|
|
key_columns: Tuple[str, ...],
|
|
columns: Optional[List[str]] = None,
|
|
num_partitions: Optional[int] = None,
|
|
should_sort: bool = False,
|
|
name: Optional[str] = None,
|
|
nranks: Optional[int] = None,
|
|
rank_pool: Optional[GPURankPool] = None,
|
|
) -> None:
|
|
nranks = nranks or _derive_num_gpu_ranks(data_context)
|
|
target_num_partitions = (
|
|
num_partitions or data_context.default_hash_shuffle_parallelism
|
|
)
|
|
# rapidsmpf requires total_nparts >= nranks
|
|
target_num_partitions = max(target_num_partitions, nranks)
|
|
|
|
super().__init__(
|
|
name=(
|
|
name
|
|
or (
|
|
f"GPUShuffle("
|
|
f"key_columns={key_columns}, "
|
|
f"num_partitions={target_num_partitions})"
|
|
)
|
|
),
|
|
input_dependencies=[input_op],
|
|
data_context=data_context,
|
|
)
|
|
|
|
self._key_columns = key_columns
|
|
self._num_partitions = target_num_partitions
|
|
self._columns = columns
|
|
self._should_sort = should_sort
|
|
self._rank_pool = rank_pool or GPURankPool(
|
|
nranks=nranks,
|
|
total_nparts=target_num_partitions,
|
|
setup_timeout_s=data_context.gpu_shuffle_setup_timeout_s,
|
|
actor_cls_factory=lambda: GPUShuffleActor,
|
|
actor_kwargs={
|
|
"key_columns": list(key_columns),
|
|
"columns": columns,
|
|
"rmm_pool_size": data_context.gpu_shuffle_rmm_pool_size,
|
|
"spill_memory_limit": data_context.gpu_shuffle_spill_memory_limit,
|
|
"should_sort": should_sort,
|
|
},
|
|
log_label="GPUShufflePool",
|
|
label_selector=data_context.execution_options.label_selector,
|
|
)
|
|
|
|
self._next_block_idx: int = 0
|
|
self._insert_tasks: Dict[int, MetadataOpTask] = {}
|
|
self._extraction_tasks: Dict[int, DataOpTask] = {}
|
|
self._finalization_started: bool = False
|
|
self._output_queue: ReorderingBundleQueue = ReorderingBundleQueue()
|
|
self._shuffled_blocks_stats: List[BlockStats] = []
|
|
self._output_blocks_stats: List[BlockStats] = []
|
|
|
|
# Progress bars (populated by SubProgressBarMixin callbacks)
|
|
self._shuffle_bar = None
|
|
self._reduce_bar = None
|
|
|
|
# Metrics
|
|
self._shuffle_metrics = OpRuntimeMetrics(self)
|
|
self._reduce_metrics = OpRuntimeMetrics(self)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Lifecycle
|
|
# ------------------------------------------------------------------
|
|
|
|
def start(
|
|
self,
|
|
options: ExecutionOptions,
|
|
block_ref_counter: "BlockRefCounter",
|
|
) -> None:
|
|
super().start(options, block_ref_counter)
|
|
self._rank_pool.start()
|
|
|
|
def _add_input_inner(self, bundle: RefBundle, input_index: int) -> None:
|
|
self._shuffle_metrics.on_input_received(bundle)
|
|
self._shuffled_blocks_stats.extend(to_stats(bundle.metadata))
|
|
|
|
for block_ref, metadata in zip(bundle.block_refs, bundle.metadata):
|
|
actor = self._rank_pool.get_actor_for_block(self._next_block_idx)
|
|
insert_ref = actor.insert_batch.remote(block_ref)
|
|
task_idx = self._next_block_idx
|
|
self._next_block_idx += 1
|
|
|
|
def _on_insert_done(idx: int = task_idx) -> None:
|
|
self._insert_tasks.pop(idx, None)
|
|
|
|
task = MetadataOpTask(
|
|
task_index=task_idx,
|
|
object_ref=insert_ref,
|
|
task_done_callback=_on_insert_done,
|
|
task_resource_bundle=None,
|
|
)
|
|
self._insert_tasks[task_idx] = task
|
|
self._shuffle_metrics.on_task_submitted(
|
|
task_idx,
|
|
RefBundle(
|
|
[BlockEntry(block_ref, metadata)],
|
|
schema=None,
|
|
owns_blocks=False,
|
|
),
|
|
task_id=task.get_task_id(),
|
|
)
|
|
|
|
if self._shuffle_bar is not None:
|
|
self._shuffle_bar.update(total=self._next_block_idx)
|
|
|
|
def _is_inserting_done(self) -> bool:
|
|
return self._inputs_complete and len(self._insert_tasks) == 0
|
|
|
|
def _try_finalize(self) -> None:
|
|
"""Schedule extraction once all inserts have completed."""
|
|
if self._finalization_started or not self._is_inserting_done():
|
|
return
|
|
|
|
self._finalization_started = True
|
|
# Running count of partitions extracted, used for metrics only.
|
|
# Real partition_id is read from each block's Arrow schema metadata
|
|
# ("_gpu_partition_id"), embedded by the actor because rapidsmpf's
|
|
# extract() uses wait_any() and yields in completion order, not
|
|
# partition order.
|
|
self._num_partitions_reduced = 0
|
|
|
|
def _on_bundle_ready(bundle: RefBundle) -> None:
|
|
assert (
|
|
bundle.schema is not None
|
|
and _GPU_PARTITION_ID_KEY in bundle.schema.metadata
|
|
), (
|
|
"Bundle is missing _gpu_partition_id in schema metadata. "
|
|
"Was finish_and_extract modified to skip tagging?"
|
|
)
|
|
partition_id = int(bundle.schema.metadata[_GPU_PARTITION_ID_KEY].decode())
|
|
clean_meta = {
|
|
k: v
|
|
for k, v in bundle.schema.metadata.items()
|
|
if k != _GPU_PARTITION_ID_KEY
|
|
}
|
|
bundle = RefBundle(
|
|
bundle.blocks,
|
|
schema=bundle.schema.with_metadata(clean_meta or None),
|
|
owns_blocks=bundle.owns_blocks,
|
|
)
|
|
self._num_partitions_reduced += 1
|
|
|
|
# Register a logical reduce "task" for this partition, mirroring
|
|
# the per-partition task lifecycle in the CPU path.
|
|
empty_bundle = RefBundle([], schema=None, owns_blocks=False)
|
|
self._reduce_metrics.on_task_submitted(
|
|
partition_id, empty_bundle, task_id=None
|
|
)
|
|
|
|
# Add to the reordering queue keyed by partition_id so output is
|
|
# always emitted in partition order (0, 1, 2, ...) regardless of
|
|
# the order GPU actors finish.
|
|
self._output_queue.add(bundle, key=partition_id)
|
|
self._output_queue.finalize(key=partition_id)
|
|
|
|
# Update Finalize Metrics on task output generated
|
|
self._reduce_metrics.on_output_queued(bundle)
|
|
self._reduce_metrics.on_task_output_generated(
|
|
task_index=partition_id, output=bundle
|
|
)
|
|
|
|
# Mark the logical partition task as finished (each GPU
|
|
# partition produces exactly one block).
|
|
self._reduce_metrics.on_task_finished(
|
|
task_index=partition_id,
|
|
exception=None,
|
|
task_exec_stats=None,
|
|
task_exec_driver_stats=None,
|
|
)
|
|
|
|
_, num_outputs, num_rows = estimate_total_num_of_blocks(
|
|
self._num_partitions_reduced,
|
|
self.upstream_op_num_outputs(),
|
|
self._reduce_metrics,
|
|
total_num_tasks=self._num_partitions,
|
|
)
|
|
self._estimated_num_output_bundles = num_outputs
|
|
self._estimated_output_num_rows = num_rows
|
|
|
|
# Update Finalize progress bar
|
|
self._reduce_bar.update(
|
|
increment=bundle.num_rows() or 0, total=self.num_output_rows_total()
|
|
)
|
|
|
|
def _on_extraction_done(
|
|
exc: Optional[Exception],
|
|
worker_stats=None,
|
|
driver_stats=None,
|
|
rank: int = -1,
|
|
) -> None:
|
|
self._extraction_tasks.pop(rank, None)
|
|
if not self._extraction_tasks:
|
|
# release GPU actors so downstream operators can acquire those GPUs
|
|
self._rank_pool.shutdown()
|
|
|
|
for rank_idx, actor in enumerate(self._rank_pool.actors):
|
|
block_gen = actor.finish_and_extract.options(
|
|
num_returns="streaming"
|
|
).remote()
|
|
|
|
data_task = DataOpTask(
|
|
task_index=rank_idx,
|
|
streaming_gen=block_gen,
|
|
block_ref_counter=self._block_ref_counter,
|
|
producer_id=self.id,
|
|
output_ready_callback=_on_bundle_ready,
|
|
task_done_callback=functools.partial(
|
|
_on_extraction_done, rank=rank_idx
|
|
),
|
|
)
|
|
self._extraction_tasks[rank_idx] = data_task
|
|
|
|
# ------------------------------------------------------------------
|
|
# Output interface
|
|
# ------------------------------------------------------------------
|
|
|
|
def has_next(self) -> bool:
|
|
self._try_finalize()
|
|
return self._output_queue.has_next()
|
|
|
|
def _get_next_inner(self) -> RefBundle:
|
|
bundle = self._output_queue.get_next()
|
|
self._reduce_metrics.on_output_dequeued(bundle)
|
|
self._reduce_metrics.on_output_taken(bundle)
|
|
self._output_blocks_stats.extend(to_stats(bundle.metadata))
|
|
return bundle
|
|
|
|
# ------------------------------------------------------------------
|
|
# Task / completion tracking
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_active_tasks(self) -> List[OpTask]:
|
|
return list(self._insert_tasks.values()) + list(self._extraction_tasks.values())
|
|
|
|
def has_completed(self) -> bool:
|
|
return (
|
|
self._finalization_started
|
|
and len(self._extraction_tasks) == 0
|
|
and super().has_completed()
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Shutdown
|
|
# ------------------------------------------------------------------
|
|
|
|
def _do_shutdown(self, force: bool = False) -> None:
|
|
self._rank_pool.shutdown(force=True)
|
|
super()._do_shutdown(force)
|
|
self._insert_tasks.clear()
|
|
self._extraction_tasks.clear()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Resource accounting
|
|
# ------------------------------------------------------------------
|
|
|
|
def current_logical_usage(self) -> ExecutionResources:
|
|
pool = self._rank_pool
|
|
if pool.is_shutdown:
|
|
return ExecutionResources(gpu=0)
|
|
gpus = len(pool.actors) or pool.nranks
|
|
return ExecutionResources(gpu=gpus)
|
|
|
|
@property
|
|
def base_resource_usage(self) -> ExecutionResources:
|
|
return ExecutionResources(gpu=self._rank_pool.nranks)
|
|
|
|
def incremental_resource_usage(self) -> ExecutionResources:
|
|
return ExecutionResources(gpu=1)
|
|
|
|
def get_actor_info(self) -> "ActorPoolInfo":
|
|
from ray.data._internal.execution.interfaces.physical_operator import (
|
|
ActorPoolInfo,
|
|
)
|
|
|
|
n = len(self._rank_pool.actors)
|
|
return ActorPoolInfo(
|
|
running=n,
|
|
pending=0,
|
|
restarting=0,
|
|
active=n,
|
|
idle=0,
|
|
pool_utilization=0,
|
|
tasks_in_flight=0,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# SubProgressBarMixin
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_sub_progress_bar_names(self) -> List[str]:
|
|
return ["GPU Shuffle", "GPU Reduce"]
|
|
|
|
def set_sub_progress_bar(self, name: str, pg: "BaseProgressBar") -> None:
|
|
if name == "GPU Shuffle":
|
|
self._shuffle_bar = pg
|
|
elif name == "GPU Reduce":
|
|
self._reduce_bar = pg
|
|
|
|
# ------------------------------------------------------------------
|
|
# Stats
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_stats(self) -> Dict[str, List[BlockStats]]:
|
|
shuffle_name = f"{self._name}_shuffle"
|
|
reduce_name = f"{self._name}_finalize"
|
|
return {
|
|
shuffle_name: self._shuffled_blocks_stats,
|
|
reduce_name: self._output_blocks_stats,
|
|
}
|