Separate shuffler's communication into new interface by pentschev · Pull Request #437 · rapidsai/rapidsmpf

pentschev · 2025-08-14T21:29:48Z

Move communication operations from the Shuffler::ProgressThread into a new CommunicationInterface abstract class to handle all communication, with the use of a class Message. This approach has the benefit of making a more distinct and easier to understand separation of the communication routine from the shuffler making the progress operator considerably smaller (will be done in follow-up PR), simultaneously allowing extensions to the communication interface disjoint from the shuffler, for example, to allow implementing an Active Message-based communicator for the shuffler. It also generalizes the concept of a "message", instead of transferring chunks all the CommunicationInterface knows about are messages represented by the Message class that combines metadata, payload and src/dst rank.

Currently implement only a TagCommunicationInterface that implements exactly the behavior that is implemented in the shuffler progress, additionally removing the need for the ack message. Extending to support Active Messages will also be handled in a separate PR.

Removes also ack messages, closes #475, closes #535 .

copy-pr-bot · 2025-08-14T21:29:51Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

pentschev · 2025-08-14T21:32:31Z

/ok to test

pentschev · 2025-08-14T21:41:45Z

/ok to test

pentschev · 2025-08-15T06:23:27Z

/ok to test

…face

wence- · 2025-10-31T10:54:44Z

+    std::vector<std::unique_ptr<MetadataPayloadExchange::Message>> received_messages;
+    for (int iter = 0; iter < 10 && received_messages.empty(); ++iter) {
+        auto messages = comm_interface->recv();
+        received_messages.insert(
+            received_messages.end(),
+            std::make_move_iterator(messages.begin()),
+            std::make_move_iterator(messages.end())
+        );
+
+        if (!received_messages.empty())
+            break;
+        std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    }


In this test, rank 0 sends to rank 1, and everyone else "does nothing".

But the test is not really following that flow.

I think we want something like:

if (comm->nranks() != 2) { GTEST_SKIP() << "Only for two ranks"; } if (comm->rank() == 0) { send(message); } else { auto received = ...; while (received.empty()) { std::ranges::move(comm_interface->recv(), std::back_inserter(received)); std::this_thread::yield(); } EXPECT_EQ(received.size(), 1); EXPECT_EQ(msg->peer_rank(), 0); ... }

But now I am worried about rank-0 finishing the test with stuff still in flight.

Ah, wait, rank-0 must also call recv repeatedly because that's the only thing that actually progresses the underlying tag communication.

This is very counter intuitive.

wence- · 2025-10-31T10:57:03Z

+    completed_messages.insert(
+        completed_messages.end(),
+        std::make_move_iterator(completed_data.begin()),
+        std::make_move_iterator(completed_data.end())
+    );


std::ranges::move(completed_data, std::back_inserter(completed_messages));

wence-

Sorry, reading through the usage in the tests gives me some pause for thought.

It seems like the sender must call recv repeatedly to progress sends (even if they're never receiving anything). Do I have this right?

pentschev · 2025-10-31T17:53:40Z

@wence- I have pushed multiple changes that do what you requested in your latest comments and what we've discussed offline. Here's a summary of the changes:

Add a progress() method to allow progressing the state of MetadataPayloadExchange independent of recv()
recv() now only returns messages that were previously received during progress()
Address test issues
1. Use a ring topology instead of the (also previously broken) root to all: send data to "next rank", receive data from "previous rank"
2. Remove timeouts, that was an overly eager optimization and could also lead to incomplete processing as you noticed. If tests hang this will cause a deadlock but it is how most of our tests are currently written we can always revisit and add timeouts later if they begin to show signs of being problematic
Fixed a bug with message ordering that was introduced with some of the changes made when addressing previous change requests in this PR

I hope this now addresses all the concerns, and more importantly does it all correctly.

One more thing: I just want to say I will start pushing against changes in the Tag communication flow in this PR, my initial idea was to start simple with porting over the Shuffler comms implementation, and while some of the changes requested previously were indeed important to make MetadataPayloadExchange better, others were deeply involved in changing the actual implementation details of the Tag code which led me to make mistakes that are not necessarily well-tested in this PR because the Shuffler is not currently using the changes here. This made me spend over an hour now to address item 4, since that change was not necessary from an interface perspective. So for any implementation changes that do not pertain to the interface but rather to the Tag API implementation let's do so in a follow-up PR where needed.

wence- · 2025-11-07T17:22:08Z

+        std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    }
+
+    if (comm->rank() == peer_rank) {


This condition is never true.

wence- · 2025-11-07T17:22:43Z

+        std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    }
+
+    if (comm->rank() == peer_rank) {


This condition is never true

…rface

wence-

Some small suggestions for fixing, but I think this now looks like a good interface

wence- · 2025-11-07T17:40:16Z

+
+    EXPECT_TRUE(comm_interface->is_idle());
+    comm_interface->send(std::move(messages));
+    EXPECT_FALSE(comm_interface->is_idle());


OK, we're idle here because no progress has been called yet?

Correct, in more specific terms it means there's still unfinished work, calling progress() only may not immediately change state.

wence- · 2025-11-07T17:45:04Z

+    // Store all completed messages
+    received_messages_.insert(
+        received_messages_.end(),
+        std::make_move_iterator(completed_metadata_only.begin()),
+        std::make_move_iterator(completed_metadata_only.end())
+    );
+    received_messages_.insert(
+        received_messages_.end(),
+        std::make_move_iterator(completed_data.begin()),
+        std::make_move_iterator(completed_data.end())
+    );


std::ranges::move(completed_metadata_only, std::back_inserter(received_messages_)); std::ranges::move(completed_data, std::back_inserter(received_messages_));

Done in f3530ef.

wence- · 2025-11-07T17:46:31Z

+    auto const t0 = Clock::now();
+
+    // Return all completed messages and clear the internal storage
+    auto messages = std::move(received_messages_);
+    received_messages_.clear();
+
+    statistics_->add_duration_stat("comms-interface-receive-messages", Clock::now() - t0);
+
+    return messages;


I don't think this makes sense to time now, since it's just a pointer move, so the timing will be orders of magnitude longer than move.

I think this whole function should be (no need for the clear() because moving from a vector clears it):

return std::move(received_messages_);

Good point, removed in bd3c42b and c20d981.

wence- · 2025-11-07T17:50:10Z

+        /**
+         * @brief Get the destination rank for outgoing or source rank for incoming
+         * messages.
+         *
+         * @return The rank of the destination or source.
+         */
+        [[nodiscard]] constexpr Rank peer_rank() const {
+            return peer_rank_;
+        }
+
+        /**
+         * @brief Get the serialized metadata for this message.
+         *
+         * This metadata is sent first to inform the receiver about the incoming message.
+         *
+         * @return The serialized metadata.
+         */
+        [[nodiscard]] constexpr std::vector<std::uint8_t> const& metadata() const {
+            return metadata_;
+        }
+
+        /**
+         * @brief Release ownership of the metadata.
+         *
+         * This is typically called when transferring metadata to the communication layer.
+         *
+         * @return Metadata with ownership transferred.
+         */
+        [[nodiscard]] std::vector<std::uint8_t> release_metadata() {
+            return std::move(metadata_);
+        }


Move implementation to core.cpp?

Moved release_metadata in ff13c0a, but the others would require removing the constexpr qualifier, did you really mean to suggest that?

wence- · 2025-11-07T17:50:48Z

+        /**
+         * @brief Construct a new Message.
+         *
+         * @param peer_rank Destination (outgoing) or source (incoming) rank.
+         * @param metadata Serialized metadata.
+         * @param data Data buffer (can be nullptr for metadata-only messages).
+         */
+        Message(
+            Rank peer_rank,
+            std::vector<std::uint8_t> metadata,
+            std::unique_ptr<Buffer> data = nullptr
+        );


Should it be valid to only send data?

I don't see why we can't allow it, it should be possible but the overhead will be the same as sending metadata, because we need to pack some "header" together with the metadata. I'll make the change to support it.

Actually, we can pass an empty vector, added a test for that in aa68e8c, does that suffice? I think it keeps the interface cleaner (instead of having something like std::optional) and delivers essentially the same result with minimal overhead. Note that there's always some sort of metadata/header required, so implementation-wise it's identical.

wence- · 2025-11-07T17:51:03Z

+         */
+        Message(
+            Rank peer_rank,
+            std::vector<std::uint8_t> metadata,


Accept by && reference?

Done in efa9d94.

pentschev · 2025-11-10T20:49:03Z

Thanks all for the reviews, I'll go ahead and merge this now and continue the integration in a follow-up PR. Let me know if you see other issues we need to address.

pentschev · 2025-11-10T20:49:07Z

/merge

Follow-up on #437 by switching `Shuffler` to use the new `MetadataPayloadExchange`. This helps simplify the code in `Shuffler`, and simultaneously remove the need for an "ack" message, opening also the possiblity for using an Active Message-based communicator (TBD). Authors: - Peter Andreas Entschev (https://github.com/pentschev) Approvers: - Mads R. B. Kristensen (https://github.com/madsbk) URL: #649

pentschev added 11 commits August 14, 2025 08:54

Add communication interface

06617ca

Make ShufflerCommunicationInterface stateful

f60e76c

Fix build errors

4cb58d4

Fix segfault

89bd2aa

Fix logic errors that caused tests failures

94a1c34

Add statistics back

a05498e

Add missing log

0cc2ea1

Add stats for total communication

9d2351a

Rename default to TagShufflerCommunication

c7007ff

Fix build errors

8f32166

Make class names more consistent

ff5db84

pentschev self-assigned this Aug 14, 2025

pentschev added feature request New feature or request non-breaking Introduces a non-breaking change labels Aug 14, 2025

Mark appropriate methods nodiscard

3e0c826

pentschev added 2 commits August 14, 2025 23:19

Use std::ranges::all_of

131d952

Use std::ranges::all_of

44e73f4

pentschev changed the title ~~Separate shuffler's communication into separate interface~~ Separate shuffler's communication into new interface Aug 15, 2025

pentschev mentioned this pull request Aug 29, 2025

Remove need for ack messages in shuffler implementation #475

Closed

pentschev added 5 commits September 24, 2025 13:13

Merge remote-tracking branch 'upstream/branch-25.10' into comms-inter…

088896a

…face

Add MessageInterface

fc765e0

Refactor CommunicationInterface for better generalization

ec1dfb8

Add ChunkMessageAdapter for MessageInterface

67deec8

Update Shuffler for the refactored CommunicationInterface

d926526

pentschev marked this pull request as ready for review September 24, 2025 20:47

pentschev requested a review from a team as a code owner September 24, 2025 20:47

wence- reviewed Oct 31, 2025

View reviewed changes

wence- requested changes Oct 31, 2025

View reviewed changes

pentschev added 6 commits October 31, 2025 08:10

Add progress() method

9622e6a

Block indefinitely until progress completes

0058a70

Use std::ranges::move

064eb23

Fix tests to send/recv in ring fashion

00d8999

Further testing improvements

b091863

Fix tag message ordering, previousyl broken due to 0d28a8c

74a8321

Merge branch 'main' into comms-interface

4793117

wence- reviewed Nov 7, 2025

View reviewed changes

pentschev added 2 commits November 7, 2025 09:27

Merge remote-tracking branch 'origin/comms-interface' into comms-inte…

f665d05

…rface

Merge remote-tracking branch 'upstream/main' into comms-interface

915cfe1

wence- approved these changes Nov 7, 2025

View reviewed changes

pentschev added 7 commits November 7, 2025 09:58

Use std::ranges::move

f3530ef

Remove irrelevant timing of recv()

bd3c42b

Move implementation to cpp file

ff13c0a

Move completed messages

c20d981

Accept metadata by constref

efa9d94

Add metadata-only test

aa68e8c

Add broadcast test

9befd54

wence- approved these changes Nov 10, 2025

View reviewed changes

rapids-bot Bot merged commit dc5850d into rapidsai:main Nov 10, 2025
91 checks passed

pentschev mentioned this pull request Nov 10, 2025

Integrate MetadataPayloadExchange into Shuffler #649

Merged

pentschev deleted the comms-interface branch December 12, 2025 20:02

Conversation

pentschev commented Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Aug 14, 2025

Uh oh!

pentschev commented Aug 14, 2025

Uh oh!

pentschev commented Aug 14, 2025

Uh oh!

pentschev commented Aug 15, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wence- left a comment

Choose a reason for hiding this comment

Uh oh!

pentschev commented Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wence- left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pentschev Nov 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pentschev commented Nov 10, 2025

Uh oh!

pentschev commented Nov 10, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pentschev commented Aug 14, 2025 •

edited

Loading

pentschev commented Oct 31, 2025 •

edited

Loading

pentschev Nov 7, 2025 •

edited

Loading