Shuffler#

using rapidsmpf::shuffler::PartID = std::uint32_t#

Partition ID, which goes from 0 to the total number of partitions.

The PartID is always referring to a partition globally.

inline std::ostream &rapidsmpf::shuffler::operator<<( std::ostream &os, detail::FinishCounter const &obj )#

Overloads the stream insertion operator for the FinishCounter class.

This function allows a description of a FinishCounter to be written to an output stream.

Parameters:

os – The output stream to write to.
obj – The object to write.

Returns:

A reference to the modified output stream.

inline std::ostream &rapidsmpf::shuffler::operator<<( std::ostream &os, Shuffler const &obj )#

Overloads the stream insertion operator for the Shuffler class.

This function allows a description of a Shuffler to be written to an output stream.

Parameters:

os – The output stream to write to.
obj – The object to write.

Returns:

A reference to the modified output stream.

class Shuffler#

#include <shuffler.hpp>

Shuffle service for all-to-all style communication of partitioned data.

The Shuffler class provides an interface for performing a shuffle operation on distributed data, using a partitioning scheme to distribute and collect data chunks across different ranks.

Public Types

using PartitionOwner = std::function<Rank(std::shared_ptr<Communicator> const&, PartID, PartID)>#: Function that given a Communicator, PartID, and total partition count, returns the rapidsmpf::Rank of the owning node.

using FinishedCallback = std::function<void()>#: Callback function type called when all partitions are finished and data can be extracted.

Warning

A callback must be fast and non-blocking. Ideally it should be used to signal a separate thread to do the actual processing.

Public Functions

Shuffler( std::shared_ptr<Communicator> comm, OpID op_id, PartID total_num_partitions, BufferResource *br, FinishedCallback &&finished_callback, PartitionOwner partition_owner = round_robin, std::unique_ptr<communicator::MetadataPayloadExchange> mpe = nullptr )#

Construct a new shuffler for a single shuffle.

Note

It is safe to reuse the op_id as soon as wait has completed locally.

Note

The caller promises that inserted buffers are stream-ordered with respect to their own stream, and extracted buffers are likewise guaranteed to be stream- ordered with respect to their own stream.

Parameters:

comm – The communicator to use.
op_id – The operation ID of the shuffle.
total_num_partitions – Total number of partitions in the shuffle.
br – Buffer resource used to allocate temporary and the shuffle result.
finished_callback – Callback to notify when all partitions are finished.
partition_owner – Function to determine partition ownership.
mpe – Optional custom metadata payload exchange. If not provided, uses the default tag-based implementation.

inline Shuffler( std::shared_ptr<Communicator> comm, OpID op_id, PartID total_num_partitions, BufferResource *br, PartitionOwner partition_owner = round_robin, std::unique_ptr<communicator::MetadataPayloadExchange> mpe = nullptr )#

Construct a new shuffler for a single shuffle.

Note

The caller promises that inserted buffers are stream-ordered with respect to their own stream, and extracted buffers are likewise guaranteed to be stream- ordered with respect to their own stream.

Parameters:

comm – The communicator to use.
op_id – The operation ID of the shuffle. This ID is unique for this operation, and should not be reused until all nodes has called Shuffler::shutdown().
total_num_partitions – Total number of partitions in the shuffle.
br – Buffer resource used to allocate temporary and the shuffle result.
partition_owner – Function to determine partition ownership.
mpe – Optional custom metadata payload exchange. If not provided, uses the default tag-based implementation.

inline std::shared_ptr<Communicator> const &comm() const noexcept#

Gets the communicator associated with this Shuffler.

Returns:: Shared pointer to communicator.

void shutdown()#

Shutdown the shuffle, blocking until all inflight communication is done.

Throws:: std::logic_error – If the shuffler is already inactive.

void insert(std::unordered_map<PartID, PackedData> &&chunks)#

Insert a bunch of packed (serialized) chunks into the shuffle.

Note

Concurrent insertion by multiple threads is supported, the caller must ensure that insert_finished() is called after all insert() calls have completed.

Parameters:: chunks – A map of partition IDs and their packed chunks.

void insert_finished()#

Signal that no more data will be inserted into the shuffle.

This informs the shuffler that this rank has finished inserting data. Must be called exactly once.

Note

If multiple threads are insert()ing, you must establish a happens-before relationship between the completion of all insert()s and the final call to insert_finished().

std::vector<PackedData> extract(PartID pid)#

Extract all chunks belonging to the specified partition.

It is valid to extract a partition that has not yet been fully received. In such cases, only the chunks received so far are returned.

To ensure the partition is complete, use wait() or another appropriate synchronization mechanism beforehand.

Parameters:: pid – The ID of the partition to extract.
Returns:: A vector of PackedData chunks associated with the partition.

bool finished() const#

Check if all partitions are finished.

Returns:: True if all partitions are finished, otherwise False.

void wait(std::optional<std::chrono::milliseconds> timeout = {})#

Wait for all partitions to finish (blocking).

Parameters:: timeout – Optional timeout (ms) to wait.
Throws:: std::runtime_error – if the timeout is reached.

std::size_t spill(std::optional<std::size_t> amount = std::nullopt)#

Spills data to device if necessary.

This function has two modes:

If amount is specified, it tries to spill at least amount bytes of device memory.
If amount is not specified (the default case), it spills based on the current available device memory returned by the buffer resource.

Parameters:: amount – An optional amount of memory to spill. If not provided, the function will check the current available device memory.
Returns:: The amount of memory actually spilled.

std::string str() const#

Returns a description of this instance.

Returns:: The description.

std::span<PartID const> local_partitions() const#

Returns the local partition IDs owned by the shuffler`.

Returns:: A span of partition IDs owned by the shuffler.

Public Members

PartID const total_num_partitions#: Total number of partition in the shuffle.

PartitionOwner const partition_owner#: Function to determine partition ownership.

Public Static Functions

static inline Rank round_robin( std::shared_ptr<Communicator> const &comm, PartID pid, [[maybe_unused]] PartID total_num_partitions )#

A PartitionOwner that distributes partitions using round robin assignment.

Parameters:

comm – The communicator to use.
pid – The partition ID to query.
total_num_partitions – Total number of partitions (unused).

Returns:

The rank owning the partition.

static inline Rank contiguous( std::shared_ptr<Communicator> const &comm, PartID pid, PartID total_num_partitions )#

A PartitionOwner that assigns contiguous partition ID ranges to ranks.

Rank 0 gets [0, k), rank 1 gets [k, 2k), etc. Use for sort so that each rank’s local_partitions() are adjacent and in order.

Parameters:

comm – The communicator to use.
pid – The partition ID to query.
total_num_partitions – Total number of partitions (must match the shuffle).

Returns:

The rank owning the partition.

static std::vector<PartID> local_partitions( std::shared_ptr<Communicator> const &comm, PartID total_num_partitions, PartitionOwner partition_owner )#

Returns the local partition IDs owned by the current node.

Parameters:

comm – The communicator to use.
total_num_partitions – Total number of partitions in the shuffle.
partition_owner – Function that determines partition ownership.

Returns:

A vector of partition IDs owned by the current node.

static inline constexpr Rank extract_rank(detail::ChunkID cid)#

Extract the rank from a chunk ID.

Parameters:: cid – The chunk ID.
Returns:: The rank.

Public Static Attributes

static constexpr int chunk_id_counter_bits = 38#: The number of bits used to store the counter in a chunk ID.