Skip to content

Add generic streaming tests for cuDF-migrated coverage gaps #1093

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
pentschev wants to merge 1 commit into from
Closed

Add generic streaming tests for cuDF-migrated coverage gaps #1093

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions cpp/tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ if(RAPIDSMPF_HAVE_STREAMING)
streaming/test_channel.cpp
streaming/test_error_handling.cpp
streaming/test_fanout.cpp
streaming/test_leaf_actor.cpp
streaming/test_lineariser.cpp
streaming/test_memory_reserve_or_wait.cpp
streaming/test_message.cpp
Expand All @@ -123,8 +124,7 @@ if(BUILD_CUDF_TESTS)

if(RAPIDSMPF_HAVE_STREAMING)
target_sources(
test_sources PRIVATE streaming/test_allgather.cpp streaming/test_leaf_actor.cpp
streaming/test_shuffler.cpp
test_sources PRIVATE streaming/test_allgather.cpp streaming/test_shuffler.cpp
)
endif()
endif()
Expand Down
29 changes: 7 additions & 22 deletions cpp/tests/streaming/test_leaf_actor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,15 @@
*/

#include <atomic>
#include <cstdint>
#include <memory>
#include <stdexcept>
#include <tuple>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <cudf_streaming/streaming/table_chunk.hpp>
#include <cudf_test/table_utilities.hpp>

#include <rapidsmpf/communicator/single.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/content_description.hpp>
#include <rapidsmpf/streaming/core/actor.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
Expand All @@ -23,7 +21,6 @@
#include <rapidsmpf/streaming/core/leaf_actor.hpp>
#include <rapidsmpf/streaming/core/queue.hpp>

#include "../utils.hpp"
#include "base_streaming_fixture.hpp"

using namespace rapidsmpf;
Expand All @@ -33,12 +30,11 @@ namespace actor = rapidsmpf::streaming::actor;
using StreamingLeafTasks = BaseStreamingFixture;

TEST_F(StreamingLeafTasks, PushAndPullChunks) {
constexpr int num_rows = 100;
constexpr int num_chunks = 10;

std::vector<cudf::table> expects;
std::vector<int> expects;
for (int i = 0; i < num_chunks; ++i) {
expects.emplace_back(random_table_with_index(i, num_rows, 0, 10));
expects.push_back(i * 10);
}

std::vector<Actor> actors;
Expand All @@ -49,15 +45,7 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) {
std::vector<Message> inputs;
for (int i = 0; i < num_chunks; ++i) {
inputs.emplace_back(
cudf_streaming::streaming::to_message(
i,
std::make_unique<cudf_streaming::streaming::TableChunk>(
std::make_unique<cudf::table>(
expects[i], stream, ctx->br()->device_mr()
),
stream
)
)
i, std::make_unique<int>(expects[i]), ContentDescription{}
);
}

Expand All @@ -72,10 +60,7 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) {
EXPECT_EQ(expects.size(), outputs.size());
for (std::size_t i = 0; i < expects.size(); ++i) {
EXPECT_EQ(outputs[i].sequence_number(), i);
CUDF_TEST_EXPECT_TABLES_EQUIVALENT(
outputs[i].get<cudf_streaming::streaming::TableChunk>().table_view(),
expects[i].view()
);
EXPECT_EQ(outputs[i].release<int>(), expects[i]);
}
}

Expand Down
126 changes: 40 additions & 86 deletions python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,23 +6,14 @@
from contextlib import nullcontext
from typing import TYPE_CHECKING

import numpy as np
import pylibcudf as plc
import pytest

pytest.importorskip("cudf_streaming")
from cudf_streaming.integrations.partition import unpack_and_concat
from cudf_streaming.streaming.table_chunk import TableChunk

from rapidsmpf.memory.packed_data import PackedData
from rapidsmpf.streaming.chunks.packed_data import PackedDataChunk
from rapidsmpf.streaming.coll.allgather import AllGather, allgather
from rapidsmpf.streaming.core.actor import define_actor, run_actor_network
from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel
from rapidsmpf.streaming.core.message import Message
from rapidsmpf.testing import assert_eq

cudf = pytest.importorskip("cudf")
from rapidsmpf.testing import generate_packed_data, validate_packed_data

if TYPE_CHECKING:
from collections.abc import Awaitable
Expand All @@ -33,34 +24,28 @@
from rapidsmpf.streaming.core.context import Context


def _make_chunk(context: Context, num_rows: int, offset: int) -> PackedDataChunk:
stream = context.get_stream_from_pool()
return PackedDataChunk.from_packed_data(
generate_packed_data(num_rows, offset, stream, context.br()),
br=context.br(),
)


def _validate_message(
context: Context, msg: Message[PackedDataChunk], num_rows: int, offset: int
) -> None:
packed = PackedDataChunk.from_message(msg, br=context.br()).to_packed_data()
validate_packed_data(packed, num_rows, offset)


def test_allgather_actor(context: Context, comm: Communicator) -> None:
if comm.nranks != 1:
pytest.skip("Only support single-rank runs")

num_rows = 1000
op_id = 0
stream = context.get_stream_from_pool()
input_tables = [
plc.Table(
[
plc.Column.from_array(
np.arange(num_rows, dtype=np.int32) + i * num_rows, stream=stream
)
]
)
for i in range(3)
]
inputs = [
PackedDataChunk.from_packed_data(
PackedData.from_cudf_packed_columns( # type: ignore[attr-defined]
plc.contiguous_split.pack(table, stream=stream),
stream,
context.br(),
),
br=context.br(),
)
for table in input_tables
]
inputs = [_make_chunk(context, num_rows, i * num_rows) for i in range(3)]
actors = []

ch1: Channel[PackedDataChunk] = context.create_channel()
Expand All @@ -77,54 +62,28 @@ def test_allgather_actor(context: Context, comm: Communicator) -> None:
actors.append(actor)
run_actor_network(context, actors=actors)

result = unpack_and_concat(
(
PackedDataChunk.from_message(msg, br=context.br()).to_packed_data()
for msg in deferred.release()
),
stream,
context.br(),
)

expect = plc.concatenate.concatenate(input_tables, stream=stream)
stream.synchronize()
assert_eq(result, expect)
results = deferred.release()
assert len(results) == len(inputs)
for i, msg in enumerate(results):
assert msg.sequence_number == i
_validate_message(context, msg, num_rows, i * num_rows)


@define_actor()
async def generate_inputs(
context: Context, ch: Channel[PackedDataChunk], num_rows: int, num_chunks: int
) -> None:
for i in range(num_chunks):
stream = context.get_stream_from_pool()
table = plc.Table(
[
plc.Column.from_array(
np.arange(num_rows, dtype=np.int32) + i * num_rows, stream=stream
)
]
)
msg = Message(
i,
PackedDataChunk.from_packed_data(
PackedData.from_cudf_packed_columns( # type: ignore[attr-defined]
plc.contiguous_split.pack(table, stream=stream),
stream,
context.br(),
),
br=context.br(),
),
)
await ch.send(context, msg)
await ch.send(context, Message(i, _make_chunk(context, num_rows, i * num_rows)))
await ch.drain(context)


@define_actor()
async def allgather_and_concat(
async def allgather_and_forward(
context: Context,
comm: Communicator,
ch_in: Channel[PackedDataChunk],
ch_out: Channel[TableChunk],
ch_out: Channel[PackedDataChunk],
op_id: int,
use_context_manager: bool, # noqa: FBT001
) -> None:
Expand All @@ -137,12 +96,14 @@ async def allgather_and_concat(
if not use_context_manager:
gather.insert_finished()
gathered = await gather.extract_all(context, ordered=True)
stream = context.get_stream_from_pool()
table = unpack_and_concat(gathered, stream, context.br())
to_send = TableChunk.from_pylibcudf_table(
table, stream, exclusive_view=True, br=context.br()
)
await ch_out.send(context, Message(0, to_send))
for sequence, packed in enumerate(gathered):
await ch_out.send(
context,
Message(
sequence,
PackedDataChunk.from_packed_data(packed, br=context.br()),
),
)
await ch_out.drain(context)


Expand All @@ -158,29 +119,22 @@ def test_allgather_object_interface(
pytest.skip("Only support single-rank runs")

ch_in: Channel[PackedDataChunk] = context.create_channel()
ch_out: Channel[TableChunk] = context.create_channel()
ch_out: Channel[PackedDataChunk] = context.create_channel()
actors: list[CppActor | Awaitable[None]] = []
num_rows = 100
num_chunks = 10
op_id = 0
actors.append(generate_inputs(context, ch_in, num_rows, num_chunks))
actors.append(
allgather_and_concat(context, comm, ch_in, ch_out, op_id, use_context_manager)
allgather_and_forward(context, comm, ch_in, ch_out, op_id, use_context_manager)
)

actor, deferred = pull_from_channel(context, ch_out)
actors.append(actor)

run_actor_network(context, actors=actors)
(result_msg,) = deferred.release()
result = TableChunk.from_message(result_msg, br=context.br())
expect = plc.Table(
[
plc.Column.from_array(
np.arange(num_rows * num_chunks, dtype=np.int32), stream=result.stream
)
]
)
got = result.table_view()
result.stream.synchronize()
assert_eq(expect, got)
results = deferred.release()
assert len(results) == num_chunks
for i, msg in enumerate(results):
assert msg.sequence_number == i
_validate_message(context, msg, num_rows, i * num_rows)
Loading