wehub-resource-sync
1024 lines
38 KiB
Python
1024 lines
38 KiB
Python
"""Unit tests for ShuffleStrategy.GPU_SHUFFLE.
|
|
|
|
These tests do NOT require GPUs or the rapidsmpf/cudf/ucxx packages.
|
|
All Ray actor calls are mocked so the tests run on a standard CPU cluster.
|
|
"""
|
|
|
|
from typing import List
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
import ray.data._internal.gpu_shuffle.hash_shuffle as hash_shuffle
|
|
from ray.data._internal.execution.interfaces import (
|
|
BlockEntry,
|
|
ExecutionResources,
|
|
PhysicalOperator,
|
|
RefBundle,
|
|
)
|
|
from ray.data._internal.gpu_shuffle.hash_shuffle import (
|
|
GPURankPool,
|
|
GPUShuffleActor,
|
|
GPUShuffleOperator,
|
|
_derive_num_gpu_ranks,
|
|
)
|
|
from ray.data._internal.logical.interfaces import LogicalOperator
|
|
from ray.data._internal.logical.operators import Repartition
|
|
from ray.data._internal.planner.plan_all_to_all_op import plan_all_to_all_op
|
|
from ray.data._internal.util import explain_plan
|
|
from ray.data.block import BlockMetadata
|
|
from ray.data.context import DataContext, ShuffleStrategy
|
|
from ray.data.tests.conftest import noop_counter
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_input_op_mock(num_blocks=None, size_bytes=None):
|
|
"""Return a minimal PhysicalOperator mock compatible with HashShufflingOperatorBase."""
|
|
logical_mock = MagicMock(LogicalOperator)
|
|
logical_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=size_bytes,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
logical_mock.estimated_num_outputs.return_value = num_blocks
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
return op_mock
|
|
|
|
|
|
def _make_bundle(num_blocks: int = 1) -> RefBundle:
|
|
"""Return a RefBundle with *num_blocks* placeholder block refs."""
|
|
meta = BlockMetadata(num_rows=10, size_bytes=100, exec_stats=None, input_files=None)
|
|
blocks = [
|
|
BlockEntry(ray.ObjectRef(bytes([i % 256]) * 28), meta)
|
|
for i in range(num_blocks)
|
|
]
|
|
return RefBundle(blocks, schema=None, owns_blocks=False)
|
|
|
|
|
|
def _make_data_context(
|
|
*,
|
|
gpu_shuffle_num_actors: int = 4,
|
|
gpu_shuffle_rmm_pool_size="auto",
|
|
gpu_shuffle_spill_memory_limit="auto",
|
|
) -> DataContext:
|
|
ctx = DataContext()
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
ctx.gpu_shuffle_num_actors = gpu_shuffle_num_actors
|
|
ctx.gpu_shuffle_rmm_pool_size = gpu_shuffle_rmm_pool_size
|
|
ctx.gpu_shuffle_spill_memory_limit = gpu_shuffle_spill_memory_limit
|
|
return ctx
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Enum / DataContext field tests (no Ray required)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDataContextGpuFields:
|
|
def test_gpu_shuffle_default_values(self):
|
|
ctx = DataContext()
|
|
assert ctx.gpu_shuffle_num_actors is None
|
|
assert ctx.gpu_shuffle_rmm_pool_size is None
|
|
assert ctx.gpu_shuffle_spill_memory_limit == "auto"
|
|
|
|
def test_gpu_shuffle_fields_settable(self):
|
|
ctx = DataContext()
|
|
ctx.shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
ctx.gpu_shuffle_num_actors = 8
|
|
ctx.gpu_shuffle_rmm_pool_size = 4 * 1024**3
|
|
ctx.gpu_shuffle_spill_memory_limit = None
|
|
assert ctx.shuffle_strategy == ShuffleStrategy.GPU_SHUFFLE
|
|
assert ctx.gpu_shuffle_num_actors == 8
|
|
assert ctx.gpu_shuffle_rmm_pool_size == 4 * 1024**3
|
|
assert ctx.gpu_shuffle_spill_memory_limit is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Import isolation — gpu_shuffle.py must be importable without GPU packages
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestImportIsolation:
|
|
def test_module_importable_without_rapidsmpf(self):
|
|
"""The gpu_shuffle module must not import rapidsmpf at module level."""
|
|
|
|
import ray.data._internal.gpu_shuffle.hash_shuffle as mod
|
|
|
|
# If we got here the import succeeded on a CPU-only env.
|
|
assert hasattr(mod, "GPUShuffleOperator")
|
|
assert hasattr(mod, "GPURankPool")
|
|
assert hasattr(mod, "GPUShuffleActor")
|
|
|
|
def test_ray_data_importable_without_gpu_packages(self):
|
|
import ray.data # noqa: F401 — must not raise
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _derive_num_gpu_ranks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDeriveNumGpuRanks:
|
|
def test_explicit_count_used(self):
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 7
|
|
assert _derive_num_gpu_ranks(ctx) == 7
|
|
|
|
def test_auto_detect_from_cluster(self):
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = None
|
|
|
|
with patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle"
|
|
"._get_total_cluster_resources"
|
|
) as mock_res:
|
|
mock_res.return_value = ExecutionResources(cpu=8, gpu=4)
|
|
assert _derive_num_gpu_ranks(ctx) == 4
|
|
|
|
def test_zero_gpus_raises(self):
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = None
|
|
|
|
with patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle"
|
|
"._get_total_cluster_resources"
|
|
) as mock_res:
|
|
mock_res.return_value = ExecutionResources(cpu=8, gpu=0)
|
|
with pytest.raises(RuntimeError, match="GPU resources"):
|
|
_derive_num_gpu_ranks(ctx)
|
|
|
|
def test_fractional_gpu_count_truncated(self):
|
|
"""ExecutionResources.gpu may be fractional; int() truncates."""
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = None
|
|
|
|
with patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle"
|
|
"._get_total_cluster_resources"
|
|
) as mock_res:
|
|
mock_res.return_value = ExecutionResources(cpu=4, gpu=3.9)
|
|
assert _derive_num_gpu_ranks(ctx) == 3
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPURankPool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGPURankPool:
|
|
def _make_pool(self, nranks=4, total_nparts=8):
|
|
return GPURankPool(
|
|
nranks=nranks,
|
|
total_nparts=total_nparts,
|
|
setup_timeout_s=60.0,
|
|
actor_cls_factory=lambda: hash_shuffle.GPUShuffleActor,
|
|
actor_kwargs={
|
|
"key_columns": ["user_id"],
|
|
"columns": None,
|
|
"rmm_pool_size": "auto",
|
|
"spill_memory_limit": "auto",
|
|
},
|
|
log_label="GPUShufflePool",
|
|
)
|
|
|
|
def test_actors_empty_before_start(self):
|
|
pool = self._make_pool()
|
|
assert pool.actors == []
|
|
|
|
def test_start_creates_correct_number_of_actors(self):
|
|
pool = self._make_pool(nranks=3)
|
|
|
|
mock_actor_handles = [MagicMock() for _ in range(3)]
|
|
mock_root_address = b"ucxx://fake-address"
|
|
|
|
# setup_root.remote() on rank 0 returns a future for (rank, address)
|
|
mock_actor_handles[0].setup_root.remote.return_value = MagicMock()
|
|
for handle in mock_actor_handles:
|
|
handle.setup_worker.remote.return_value = MagicMock()
|
|
|
|
with patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle.GPUShuffleActor"
|
|
) as mock_actor_cls, patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle.ray.get"
|
|
) as mock_ray_get, patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle.ray.wait"
|
|
) as mock_ray_wait:
|
|
mock_actor_cls.options.return_value.remote.side_effect = mock_actor_handles
|
|
# First ray.get returns (rank, root_address); second returns None list (setup done)
|
|
mock_ray_get.side_effect = [(0, mock_root_address), [None, None, None]]
|
|
# ray.wait returns all refs as ready
|
|
worker_refs = [
|
|
h.setup_worker.remote.return_value for h in mock_actor_handles
|
|
]
|
|
mock_ray_wait.return_value = (worker_refs, [])
|
|
|
|
pool.start()
|
|
|
|
assert len(pool.actors) == 3
|
|
assert mock_actor_cls.options.call_count == 3
|
|
|
|
def test_get_actor_for_block_round_robin(self):
|
|
pool = self._make_pool(nranks=3)
|
|
mock_actors = [MagicMock(name=f"actor_{i}") for i in range(3)]
|
|
pool._actors = mock_actors
|
|
|
|
# Blocks 0,1,2,3,4 should map to actors 0,1,2,0,1
|
|
expected = [mock_actors[i % 3] for i in range(5)]
|
|
actual = [pool.get_actor_for_block(i) for i in range(5)]
|
|
assert actual == expected
|
|
|
|
def test_shutdown_force_kills_actors(self):
|
|
pool = self._make_pool(nranks=2)
|
|
mock_actors = [MagicMock(), MagicMock()]
|
|
pool._actors = mock_actors
|
|
|
|
with patch("ray.data._internal.gpu_shuffle.hash_shuffle.ray.kill") as mock_kill:
|
|
pool.shutdown(force=True)
|
|
|
|
assert mock_kill.call_count == 2
|
|
assert pool.actors == []
|
|
assert pool.is_shutdown
|
|
|
|
def test_shutdown_without_force_clears_actors(self):
|
|
pool = self._make_pool(nranks=2)
|
|
pool._actors = [MagicMock(), MagicMock()]
|
|
|
|
with patch("ray.data._internal.gpu_shuffle.hash_shuffle.ray.kill") as mock_kill:
|
|
pool.shutdown(force=False)
|
|
|
|
mock_kill.assert_not_called()
|
|
assert pool.actors == []
|
|
assert pool.is_shutdown
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUShuffleOperator constructor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGPUShuffleOperatorConstructor:
|
|
def _make_op(self, num_partitions=None, nranks=4, default_parallelism=200):
|
|
ctx = _make_data_context(gpu_shuffle_num_actors=nranks)
|
|
ctx.default_hash_shuffle_parallelism = default_parallelism
|
|
input_op = _make_input_op_mock()
|
|
|
|
return GPUShuffleOperator(
|
|
input_op,
|
|
ctx,
|
|
key_columns=("user_id",),
|
|
num_partitions=num_partitions,
|
|
)
|
|
|
|
def test_name_contains_key_columns(self):
|
|
op = self._make_op(num_partitions=8)
|
|
assert "user_id" in op.name
|
|
|
|
def test_name_contains_num_partitions(self):
|
|
op = self._make_op(num_partitions=8)
|
|
assert "8" in op.name
|
|
|
|
def test_num_partitions_clamped_to_nranks(self):
|
|
"""When requested partitions < nranks, partitions is raised to nranks."""
|
|
op = self._make_op(num_partitions=2, nranks=8)
|
|
assert op._num_partitions == 8
|
|
|
|
def test_num_partitions_exceeds_nranks_unchanged(self):
|
|
op = self._make_op(num_partitions=16, nranks=4)
|
|
assert op._num_partitions == 16
|
|
|
|
def test_num_partitions_defaults_to_context_parallelism(self):
|
|
op = self._make_op(num_partitions=None, nranks=4, default_parallelism=200)
|
|
assert op._num_partitions == 200
|
|
|
|
def test_base_resource_usage_is_nranks_gpus(self):
|
|
op = self._make_op(nranks=6, num_partitions=6)
|
|
usage = op.base_resource_usage
|
|
assert usage.gpu == 6
|
|
|
|
def test_current_logical_usage_reserves_nranks_before_pool_start(self):
|
|
"""Empty actor list before start must still reserve configured GPUs."""
|
|
op = self._make_op(nranks=5, num_partitions=5)
|
|
assert op.current_logical_usage().gpu == 5
|
|
|
|
def test_current_logical_usage_matches_len_actors_when_running(self):
|
|
op = self._make_op(nranks=4, num_partitions=4)
|
|
op._rank_pool._actors = [MagicMock() for _ in range(4)]
|
|
assert op.current_logical_usage().gpu == 4
|
|
|
|
def test_current_logical_usage_zero_after_pool_shutdown(self):
|
|
"""Early actor release must drop logical GPU usage for the scheduler."""
|
|
op = self._make_op(nranks=4, num_partitions=4)
|
|
op._rank_pool._actors = [MagicMock() for _ in range(4)]
|
|
op._rank_pool.shutdown(force=False)
|
|
assert op.current_logical_usage().gpu == 0
|
|
|
|
def test_incremental_resource_usage_is_one_gpu(self):
|
|
op = self._make_op()
|
|
usage = op.incremental_resource_usage()
|
|
assert usage.gpu == 1
|
|
|
|
def test_progress_bar_names(self):
|
|
op = self._make_op()
|
|
names = op.get_sub_progress_bar_names()
|
|
assert names == ["GPU Shuffle", "GPU Reduce"]
|
|
|
|
def test_set_sub_progress_bar_shuffle(self):
|
|
op = self._make_op()
|
|
mock_bar = MagicMock()
|
|
op.set_sub_progress_bar("GPU Shuffle", mock_bar)
|
|
assert op._shuffle_bar is mock_bar
|
|
|
|
def test_set_sub_progress_bar_reduce(self):
|
|
op = self._make_op()
|
|
mock_bar = MagicMock()
|
|
op.set_sub_progress_bar("GPU Reduce", mock_bar)
|
|
assert op._reduce_bar is mock_bar
|
|
|
|
def test_initial_state(self):
|
|
op = self._make_op()
|
|
assert op._next_block_idx == 0
|
|
assert op._insert_tasks == {}
|
|
assert op._extraction_tasks == {}
|
|
assert not op._finalization_started
|
|
assert len(op._output_queue) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUShuffleOperator: _add_input_inner block routing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGPUShuffleOperatorInputRouting:
|
|
def _make_op_with_mock_pool(self, nranks=3, num_partitions=6):
|
|
ctx = _make_data_context(gpu_shuffle_num_actors=nranks)
|
|
input_op = _make_input_op_mock()
|
|
op = GPUShuffleOperator(
|
|
input_op, ctx, key_columns=("k",), num_partitions=num_partitions
|
|
)
|
|
# Replace the real pool with a mock
|
|
mock_actors = [MagicMock(name=f"actor_{i}") for i in range(nranks)]
|
|
for actor in mock_actors:
|
|
actor.insert_batch.remote.return_value = MagicMock()
|
|
op._rank_pool._actors = mock_actors
|
|
op._rank_pool._nranks = nranks
|
|
return op, mock_actors
|
|
|
|
def test_single_block_routed_to_first_actor(self):
|
|
op, actors = self._make_op_with_mock_pool(nranks=3)
|
|
bundle = _make_bundle(num_blocks=1)
|
|
op._add_input_inner(bundle, input_index=0)
|
|
|
|
actors[0].insert_batch.remote.assert_called_once()
|
|
|
|
def test_round_robin_across_three_ranks(self):
|
|
op, actors = self._make_op_with_mock_pool(nranks=3)
|
|
|
|
# Submit 6 single-block bundles
|
|
for _ in range(6):
|
|
op._add_input_inner(_make_bundle(1), input_index=0)
|
|
|
|
# Each actor should have received exactly 2 blocks
|
|
for actor in actors:
|
|
assert actor.insert_batch.remote.call_count == 2
|
|
|
|
def test_block_idx_increments_per_block(self):
|
|
op, actors = self._make_op_with_mock_pool(nranks=3)
|
|
|
|
bundle_with_2 = _make_bundle(num_blocks=2)
|
|
op._add_input_inner(bundle_with_2, input_index=0)
|
|
|
|
assert op._next_block_idx == 2
|
|
|
|
def test_insert_tasks_tracked(self):
|
|
op, actors = self._make_op_with_mock_pool(nranks=2)
|
|
op._add_input_inner(_make_bundle(1), 0)
|
|
assert len(op._insert_tasks) == 1
|
|
|
|
def test_insert_task_callback_removes_task(self):
|
|
op, actors = self._make_op_with_mock_pool(nranks=2)
|
|
op._add_input_inner(_make_bundle(1), 0)
|
|
# Grab the callback and invoke it
|
|
task = list(op._insert_tasks.values())[0]
|
|
assert 0 in op._insert_tasks
|
|
task._task_done_callback()
|
|
assert 0 not in op._insert_tasks
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUShuffleOperator: finalization and completion
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGPUShuffleOperatorFinalization:
|
|
def _make_op(self, nranks=2, num_partitions=4):
|
|
ctx = _make_data_context(gpu_shuffle_num_actors=nranks)
|
|
input_op = _make_input_op_mock()
|
|
op = GPUShuffleOperator(
|
|
input_op, ctx, key_columns=("k",), num_partitions=num_partitions
|
|
)
|
|
mock_actors = [MagicMock(name=f"actor_{i}") for i in range(nranks)]
|
|
for actor in mock_actors:
|
|
actor.finish_and_extract.options.return_value.remote.return_value = (
|
|
MagicMock()
|
|
)
|
|
op._rank_pool._actors = mock_actors
|
|
op._rank_pool._nranks = nranks
|
|
op._block_ref_counter = noop_counter()
|
|
return op, mock_actors
|
|
|
|
def test_finalization_not_started_until_inputs_complete(self):
|
|
op, _ = self._make_op()
|
|
op._inputs_complete = False
|
|
op._try_finalize()
|
|
assert not op._finalization_started
|
|
|
|
def test_finalization_not_started_while_inserts_pending(self):
|
|
op, _ = self._make_op()
|
|
op._inputs_complete = True
|
|
op._insert_tasks[0] = MagicMock() # fake pending insert
|
|
op._try_finalize()
|
|
assert not op._finalization_started
|
|
|
|
def test_finalization_starts_after_all_inserts_done(self):
|
|
op, mock_actors = self._make_op(nranks=2)
|
|
op._inputs_complete = True
|
|
# No pending inserts
|
|
|
|
with patch.object(op._reduce_metrics, "on_task_submitted"):
|
|
op._try_finalize()
|
|
|
|
assert op._finalization_started
|
|
|
|
def test_finish_and_extract_called_on_all_ranks(self):
|
|
op, mock_actors = self._make_op(nranks=2)
|
|
op._inputs_complete = True
|
|
|
|
with patch.object(op._reduce_metrics, "on_task_submitted"):
|
|
op._try_finalize()
|
|
|
|
for actor in mock_actors:
|
|
actor.finish_and_extract.options.assert_called_once()
|
|
|
|
def test_try_finalize_idempotent(self):
|
|
op, mock_actors = self._make_op(nranks=2)
|
|
op._inputs_complete = True
|
|
|
|
with patch.object(op._reduce_metrics, "on_task_submitted"):
|
|
op._try_finalize()
|
|
op._try_finalize() # second call should be no-op
|
|
|
|
# finish_and_extract should only be called once per actor
|
|
for actor in mock_actors:
|
|
assert actor.finish_and_extract.options.call_count == 1
|
|
|
|
def test_has_next_false_initially(self):
|
|
op, _ = self._make_op()
|
|
op._inputs_complete = False
|
|
assert not op.has_next()
|
|
|
|
def test_has_next_true_when_output_queued(self):
|
|
op, _ = self._make_op()
|
|
bundle = _make_bundle(1)
|
|
op._output_queue.add(bundle, key=0)
|
|
op._output_queue.finalize(key=0)
|
|
op._finalization_started = True
|
|
assert op.has_next()
|
|
|
|
def test_get_next_inner_dequeues(self):
|
|
op, _ = self._make_op()
|
|
b1 = _make_bundle(1)
|
|
b2 = _make_bundle(1)
|
|
op._output_queue.add(b1, key=0)
|
|
op._output_queue.finalize(key=0)
|
|
op._output_queue.add(b2, key=1)
|
|
op._output_queue.finalize(key=1)
|
|
|
|
with patch.object(op._reduce_metrics, "on_output_dequeued"), patch.object(
|
|
op._reduce_metrics, "on_output_taken"
|
|
):
|
|
result = op._get_next_inner()
|
|
|
|
assert result is b1
|
|
assert op._output_queue.has_next()
|
|
|
|
def test_has_completed_false_while_extracting(self):
|
|
op, _ = self._make_op()
|
|
op._finalization_started = True
|
|
op._extraction_tasks[0] = MagicMock() # still running
|
|
assert not op.has_completed()
|
|
|
|
def test_output_order_is_partition_order_regardless_of_arrival(self):
|
|
"""Bundles arriving out of order must be emitted in ascending partition order."""
|
|
op, _ = self._make_op(nranks=2, num_partitions=4)
|
|
op._finalization_started = True
|
|
|
|
# Build 4 bundles, insert in reverse order (3, 2, 1, 0)
|
|
bundles = {}
|
|
for partition_id in reversed(range(4)):
|
|
meta = BlockMetadata(
|
|
num_rows=1, size_bytes=8, exec_stats=None, input_files=None
|
|
)
|
|
bundle = RefBundle(
|
|
[BlockEntry(ray.ObjectRef(bytes([partition_id]) * 28), meta)],
|
|
schema=None,
|
|
owns_blocks=False,
|
|
)
|
|
bundles[partition_id] = bundle
|
|
op._output_queue.add(bundle, key=partition_id)
|
|
op._output_queue.finalize(key=partition_id)
|
|
|
|
with patch.object(op._reduce_metrics, "on_output_dequeued"), patch.object(
|
|
op._reduce_metrics, "on_output_taken"
|
|
):
|
|
results = [op._get_next_inner() for _ in range(4)]
|
|
|
|
# Output must be in partition order 0, 1, 2, 3 — not insertion order 3, 2, 1, 0
|
|
assert results == [bundles[i] for i in range(4)]
|
|
|
|
def test_get_active_tasks_combines_both_phases(self):
|
|
op, _ = self._make_op()
|
|
insert_task = MagicMock()
|
|
extract_task = MagicMock()
|
|
op._insert_tasks[0] = insert_task
|
|
op._extraction_tasks[0] = extract_task
|
|
|
|
active = op.get_active_tasks()
|
|
assert insert_task in active
|
|
assert extract_task in active
|
|
assert len(active) == 2
|
|
|
|
def test_shutdown_clears_tasks_and_kills_actors(self):
|
|
op, mock_actors = self._make_op(nranks=2)
|
|
op._insert_tasks[0] = MagicMock()
|
|
op._extraction_tasks[0] = MagicMock()
|
|
expected_kill_count = len(mock_actors)
|
|
|
|
with patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle.ray.kill"
|
|
) as mock_kill, patch.object(
|
|
PhysicalOperator, "_do_shutdown", return_value=None
|
|
):
|
|
op._do_shutdown(force=True)
|
|
|
|
assert op._insert_tasks == {}
|
|
assert op._extraction_tasks == {}
|
|
assert mock_kill.call_count == expected_kill_count
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# plan_all_to_all_op routing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPlanAllToAllOpRouting:
|
|
"""Verify that plan_all_to_all_op routes GPU_SHUFFLE to GPUShuffleOperator."""
|
|
|
|
def _make_repartition_op(self, keys=("user_id",), num_outputs=8):
|
|
return Repartition(
|
|
num_outputs=num_outputs,
|
|
input_dependencies=[MagicMock(LogicalOperator)],
|
|
shuffle=True,
|
|
keys=list(keys),
|
|
)
|
|
|
|
def test_gpu_shuffle_routes_to_gpu_operator(self):
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
|
|
logical_op = self._make_repartition_op(keys=["user_id"], num_outputs=8)
|
|
input_physical_op = _make_input_op_mock()
|
|
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert isinstance(op, GPUShuffleOperator)
|
|
|
|
def test_hash_shuffle_routes_to_shuffle_reduce_op(self):
|
|
"""V2 hash shuffle is a two-op DAG; planner returns the ShuffleReduceOp
|
|
with the ShuffleMapOp as its upstream input dependency."""
|
|
from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import ( # noqa: E501
|
|
ShuffleMapOp,
|
|
)
|
|
from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import ( # noqa: E501
|
|
ShuffleReduceOp,
|
|
)
|
|
|
|
ctx = DataContext()
|
|
ctx._shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE
|
|
|
|
logical_op = self._make_repartition_op(keys=["user_id"], num_outputs=8)
|
|
input_physical_op = _make_input_op_mock()
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert isinstance(op, ShuffleReduceOp)
|
|
assert isinstance(op.input_dependencies[0], ShuffleMapOp)
|
|
|
|
def test_unsupported_strategy_with_keys_raises(self):
|
|
ctx = DataContext()
|
|
ctx._shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
|
|
|
|
logical_op = self._make_repartition_op(keys=["user_id"], num_outputs=8)
|
|
input_physical_op = _make_input_op_mock()
|
|
|
|
with pytest.raises(ValueError, match="HASH_SHUFFLE"):
|
|
plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
def test_gpu_shuffle_respects_num_outputs(self):
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
|
|
logical_op = self._make_repartition_op(keys=["id"], num_outputs=16)
|
|
input_physical_op = _make_input_op_mock()
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert op._num_partitions == 16
|
|
|
|
def test_gpu_shuffle_key_columns_normalised(self):
|
|
"""Key columns from SortKey.get_columns() should propagate correctly."""
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
|
|
logical_op = self._make_repartition_op(keys=["col_a", "col_b"], num_outputs=8)
|
|
input_physical_op = _make_input_op_mock()
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert "col_a" in op._key_columns
|
|
assert "col_b" in op._key_columns
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUShuffleActor: deferred import guard
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGPUShuffleActorImportGuard:
|
|
"""GPUShuffleActor.__init__ must raise ImportError with a helpful message
|
|
when rapidsmpf is not installed, not a generic ModuleNotFoundError."""
|
|
|
|
def test_missing_rapidsmpf_raises_import_error(self):
|
|
from ray.data._internal.gpu_shuffle.hash_shuffle import GPUShuffleActor
|
|
|
|
# Access the underlying class (bypass Ray actor wrapper)
|
|
ActorClass = GPUShuffleActor.__ray_actor_class__
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{"ray.data._internal.gpu_shuffle.rapidsmpf_backend": None},
|
|
):
|
|
with pytest.raises(ImportError, match="rapidsmpf"):
|
|
ActorClass(
|
|
nranks=2,
|
|
total_nparts=4,
|
|
key_columns=["k"],
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU fixtures — shared by all real-GPU test classes below
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _num_cluster_gpus() -> int:
|
|
"""Return the number of GPUs in the Ray cluster (0 if Ray not initialised)."""
|
|
if not ray.is_initialized():
|
|
return 0
|
|
return int(ray.cluster_resources().get("GPU", 0))
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ray_with_gpu():
|
|
"""Skip the test if GPU packages or GPU hardware are absent.
|
|
|
|
Imports ``cudf``, ``rapidsmpf``, and ``ucp`` (ucxx Python bindings) with
|
|
``pytest.importorskip`` so the skip message clearly names the missing
|
|
package. Also initialises Ray (if not already running) and checks that
|
|
at least one GPU is visible in the cluster.
|
|
"""
|
|
pytest.importorskip("cudf", reason="cudf (GPU DataFrame library) not installed")
|
|
pytest.importorskip("rapidsmpf", reason="rapidsmpf not installed")
|
|
|
|
if not ray.is_initialized():
|
|
ray.init()
|
|
|
|
num_gpus = _num_cluster_gpus()
|
|
if num_gpus < 1:
|
|
pytest.skip("No GPU resources found in the Ray cluster")
|
|
|
|
yield num_gpus
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUShuffleActor — real GPU paths (conditional)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.gpu
|
|
class TestGPUShuffleActorReal:
|
|
"""Exercises GPU actor methods on actual hardware.
|
|
|
|
All tests are skipped automatically when GPU packages or GPU resources are
|
|
absent. Run them explicitly with ``pytest -m gpu``.
|
|
"""
|
|
|
|
def _make_setup_actor(self, total_nparts: int = 2, key_columns=None):
|
|
"""Create, UCXX-initialise, and return a single-rank GPUShuffleActor."""
|
|
key_columns = key_columns or ["id"]
|
|
actor = GPUShuffleActor.options(num_gpus=1).remote(
|
|
nranks=1,
|
|
total_nparts=total_nparts,
|
|
key_columns=key_columns,
|
|
)
|
|
_, root_address = ray.get(actor.setup_root.remote())
|
|
ray.get(actor.setup_worker.remote(root_address))
|
|
return actor
|
|
|
|
def test_actor_init_succeeds(self, ray_with_gpu):
|
|
"""GPUShuffleActor.__init__ succeeds (rapidsmpf import guard passes)."""
|
|
actor = GPUShuffleActor.options(num_gpus=1).remote(
|
|
nranks=1,
|
|
total_nparts=2,
|
|
key_columns=["id"],
|
|
)
|
|
ray.kill(actor)
|
|
|
|
def test_setup_root_returns_rank_and_address(self, ray_with_gpu):
|
|
"""setup_root() returns a (rank, address_bytes) tuple for UCXX setup."""
|
|
actor = GPUShuffleActor.options(num_gpus=1).remote(
|
|
nranks=1, total_nparts=1, key_columns=["k"]
|
|
)
|
|
rank, addr = ray.get(actor.setup_root.remote())
|
|
assert isinstance(rank, int)
|
|
assert isinstance(addr, bytes)
|
|
assert len(addr) > 0
|
|
ray.kill(actor)
|
|
|
|
def test_insert_batch_returns_row_count(self, ray_with_gpu):
|
|
"""insert_batch() returns the number of rows in the Arrow batch."""
|
|
actor = self._make_setup_actor()
|
|
table = pa.table({"id": [1, 2, 3], "val": [0.1, 0.2, 0.3]})
|
|
count = ray.get(actor.insert_batch.remote(table))
|
|
assert count == 3
|
|
ray.kill(actor)
|
|
|
|
def test_insert_batch_large_table(self, ray_with_gpu):
|
|
"""insert_batch handles a larger Arrow Table without error."""
|
|
|
|
n = 5_000
|
|
actor = self._make_setup_actor(total_nparts=4)
|
|
table = pa.table(
|
|
{
|
|
"id": pa.array(np.arange(n, dtype=np.int64)),
|
|
"val": pa.array(np.random.rand(n)),
|
|
}
|
|
)
|
|
count = ray.get(actor.insert_batch.remote(table))
|
|
assert count == n
|
|
ray.kill(actor)
|
|
|
|
def test_insert_batch_multiple_batches(self, ray_with_gpu):
|
|
"""Multiple insert_batch calls each return the correct row count."""
|
|
actor = self._make_setup_actor(total_nparts=2)
|
|
sizes = [3, 7, 5]
|
|
offset = 0
|
|
for size in sizes:
|
|
table = pa.table(
|
|
{
|
|
"id": list(range(offset, offset + size)),
|
|
"label": ["x"] * size,
|
|
}
|
|
)
|
|
count = ray.get(actor.insert_batch.remote(table))
|
|
assert count == size
|
|
offset += size
|
|
ray.kill(actor)
|
|
|
|
def test_finish_and_extract_succeeds_after_inserts(self, ray_with_gpu):
|
|
"""finish_and_extract() completes without error after a batch insert."""
|
|
actor = self._make_setup_actor()
|
|
table = pa.table({"id": [0, 1, 2], "v": [10, 20, 30]})
|
|
ray.get(actor.insert_batch.remote(table))
|
|
gen = actor.finish_and_extract.options(num_returns="streaming").remote()
|
|
# Drain the generator to ensure it completes.
|
|
for ref in gen:
|
|
ray.get(ref)
|
|
ray.kill(actor)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Single-rank end-to-end roundtrip (conditional)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.gpu
|
|
class TestGPUSingleRankRoundtrip:
|
|
"""Full insert → finish_and_extract roundtrip (1 GPU)."""
|
|
|
|
@staticmethod
|
|
def _collect_partitions(actor) -> List[pa.Table]:
|
|
"""Drain a streaming finish_and_extract generator.
|
|
|
|
finish_and_extract follows the Ray Data streaming protocol: each
|
|
partition yields a block (pa.Table) followed by a BlockMetadataWithSchema.
|
|
Collect only the blocks.
|
|
"""
|
|
|
|
gen = actor.finish_and_extract.options(num_returns="streaming").remote()
|
|
return [
|
|
item for ref in gen for item in [ray.get(ref)] if isinstance(item, pa.Table)
|
|
]
|
|
|
|
def _actor_with_data(
|
|
self,
|
|
table: pa.Table,
|
|
key_columns: List[str],
|
|
total_nparts: int = 2,
|
|
):
|
|
"""Create a single-rank actor, feed *table* into it, return ready actor."""
|
|
actor = GPUShuffleActor.options(num_gpus=1).remote(
|
|
nranks=1,
|
|
total_nparts=total_nparts,
|
|
key_columns=key_columns,
|
|
)
|
|
_, root_address = ray.get(actor.setup_root.remote())
|
|
ray.get(actor.setup_worker.remote(root_address))
|
|
ray.get(actor.insert_batch.remote(table))
|
|
return actor
|
|
|
|
def test_roundtrip_preserves_row_count(self, ray_with_gpu):
|
|
"""All inserted rows appear in the extracted partitions."""
|
|
n_rows = 30
|
|
table = pa.table(
|
|
{
|
|
"key": list(range(n_rows)),
|
|
"data": [float(i) for i in range(n_rows)],
|
|
}
|
|
)
|
|
actor = self._actor_with_data(table, ["key"], total_nparts=3)
|
|
partitions = self._collect_partitions(actor)
|
|
assert sum(t.num_rows for t in partitions) == n_rows
|
|
ray.kill(actor)
|
|
|
|
def test_roundtrip_output_is_arrow_tables(self, ray_with_gpu):
|
|
"""finish_and_extract yields pyarrow.Table objects."""
|
|
table = pa.table({"id": [1, 2, 3, 4], "name": ["a", "b", "c", "d"]})
|
|
actor = self._actor_with_data(table, ["id"], total_nparts=2)
|
|
partitions = self._collect_partitions(actor)
|
|
for part in partitions:
|
|
assert isinstance(part, pa.Table)
|
|
ray.kill(actor)
|
|
|
|
def test_roundtrip_multiple_batches_no_rows_lost(self, ray_with_gpu):
|
|
"""Rows from multiple insert_batch calls are all recovered."""
|
|
actor = GPUShuffleActor.options(num_gpus=1).remote(
|
|
nranks=1, total_nparts=2, key_columns=["k"]
|
|
)
|
|
_, root_address = ray.get(actor.setup_root.remote())
|
|
ray.get(actor.setup_worker.remote(root_address))
|
|
|
|
batch_sizes = [4, 6, 10]
|
|
offset = 0
|
|
for size in batch_sizes:
|
|
table = pa.table({"k": list(range(offset, offset + size)), "v": [0] * size})
|
|
ray.get(actor.insert_batch.remote(table))
|
|
offset += size
|
|
|
|
partitions = self._collect_partitions(actor)
|
|
assert sum(t.num_rows for t in partitions) == sum(batch_sizes)
|
|
ray.kill(actor)
|
|
|
|
def test_roundtrip_column_names_preserved(self, ray_with_gpu):
|
|
"""Column names in extracted partitions match the inserted schema."""
|
|
col_names = ["alpha", "beta", "gamma"]
|
|
table = pa.table({"alpha": [1, 2], "beta": [3.0, 4.0], "gamma": ["x", "y"]})
|
|
actor = self._actor_with_data(table, ["alpha"], total_nparts=1)
|
|
partitions = self._collect_partitions(actor)
|
|
for part in partitions:
|
|
if part.num_rows > 0:
|
|
for name in col_names:
|
|
assert name in part.schema.names
|
|
ray.kill(actor)
|
|
|
|
def test_roundtrip_key_column_hash_partitions_consistently(self, ray_with_gpu):
|
|
"""Each key value is always routed to exactly one partition."""
|
|
# Hash partitioning guarantees that all rows sharing a key land in the
|
|
# same partition, but makes no promise that *distinct* keys go to
|
|
# *distinct* partitions (collisions are possible, especially with few
|
|
# partitions). Test the actual guarantee: no key is split across
|
|
# multiple partitions.
|
|
n_rows, n_keys = 100, 10
|
|
table = pa.table(
|
|
{
|
|
"group": [i % n_keys for i in range(n_rows)],
|
|
"val": list(range(n_rows)),
|
|
}
|
|
)
|
|
actor = self._actor_with_data(table, ["group"], total_nparts=2)
|
|
all_partitions = self._collect_partitions(actor)
|
|
ray.kill(actor)
|
|
|
|
# For each key, collect the set of partition indices it appears in.
|
|
key_to_part_indices: dict = {}
|
|
for idx, part in enumerate(all_partitions):
|
|
for key in part.column("group").unique().to_pylist():
|
|
key_to_part_indices.setdefault(key, set()).add(idx)
|
|
|
|
for key, part_indices in key_to_part_indices.items():
|
|
assert (
|
|
len(part_indices) == 1
|
|
), f"Key {key!r} was split across partitions {part_indices}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPURankPool — real GPU lifecycle (conditional)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.gpu
|
|
class TestGPURankPoolReal:
|
|
"""Tests that exercise GPURankPool with actual GPU actors."""
|
|
|
|
def _make_pool(self, nranks: int = 1, total_nparts: int = 2) -> GPURankPool:
|
|
return GPURankPool(
|
|
nranks=nranks,
|
|
total_nparts=total_nparts,
|
|
setup_timeout_s=60.0,
|
|
actor_cls_factory=lambda: hash_shuffle.GPUShuffleActor,
|
|
actor_kwargs={
|
|
"key_columns": ["id"],
|
|
"columns": None,
|
|
"rmm_pool_size": "auto",
|
|
"spill_memory_limit": "auto",
|
|
},
|
|
log_label="GPUShufflePool",
|
|
)
|
|
|
|
def test_pool_start_creates_actors(self, ray_with_gpu):
|
|
"""GPURankPool.start() creates the expected number of actors."""
|
|
pool = self._make_pool(nranks=1)
|
|
pool.start()
|
|
assert len(pool.actors) == 1
|
|
pool.shutdown(force=True)
|
|
|
|
def test_pool_shutdown_clears_actors(self, ray_with_gpu):
|
|
"""GPURankPool.shutdown(force=True) kills actors and empties the list."""
|
|
pool = self._make_pool(nranks=1)
|
|
pool.start()
|
|
pool.shutdown(force=True)
|
|
assert pool.actors == []
|
|
|
|
def test_pool_actors_respond_after_start(self, ray_with_gpu):
|
|
"""Actors returned by the pool respond to remote calls after start()."""
|
|
pool = self._make_pool(nranks=1, total_nparts=1)
|
|
pool.start()
|
|
actor = pool.actors[0]
|
|
# Actor is fully set up by pool.start(); insert_batch should work immediately
|
|
table = pa.table({"id": [1], "v": [2]})
|
|
ray.get(actor.insert_batch.remote(table))
|
|
pool.shutdown(force=True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU Hash Shuffle - end to end
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.gpu
|
|
class TestGPUHashShuffle:
|
|
def test_hash_shuffle(self, ray_with_gpu):
|
|
"""Test that hash shuffle works end to end."""
|
|
# ray.init(num_gpus=1)
|
|
num_gpus = ray_with_gpu
|
|
ray.data.context.DataContext.get_current().shuffle_strategy = (
|
|
ShuffleStrategy.GPU_SHUFFLE
|
|
)
|
|
|
|
num_rows = 10000
|
|
parallelism = 1000
|
|
num_blocks = int(parallelism / 10)
|
|
|
|
ds = ray.data.range(num_rows, parallelism=parallelism).materialize()
|
|
ds = ds.repartition(keys=["id"], num_blocks=num_blocks)
|
|
assert "GPUShuffle" in explain_plan(ds._logical_plan)
|
|
ds = ds.materialize()
|
|
assert ds.num_blocks() == max(num_blocks, num_gpus)
|
|
assert ds.count() == num_rows
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|