Python Streaming API Reference#

Core features#

Submodule for streaming core operations.

class rapidsmpf.streaming.core.Channel#

A coroutine-based, bounded channel for asynchronously sending and receiving Message objects.

Methods

`drain`(self, Context ctx)	Drain pending messages and then shut down the channel.
`drain_metadata`(self, Context ctx)	Drain pending metadata messages and then shut down metadata processing.
`recv`(self, Context ctx)	Receive the next message from the channel.
`recv_metadata`(self, Context ctx)	Receive the next metadata message from the channel.
`send`(self, Context ctx, Message msg)	Send a message into the channel.
`send_metadata`(self, Context ctx, Message msg)	Send a metadata message into the channel.
`shutdown`(self, Context ctx)	Immediately shut down the channel.
`shutdown_metadata`(self, Context ctx)	Immediately shut down metadata handling.

async drain(self, Context ctx)#

Drain pending messages and then shut down the channel.

Parameters:

ctx: The current streaming context.

async drain_metadata(self, Context ctx)#

Drain pending metadata messages and then shut down metadata processing.

Parameters:

ctx: The current streaming context.

async recv(self, Context ctx)#

Receive the next message from the channel.

Parameters:

ctx: The current streaming context.

Returns:

A Message if a message is available, otherwise None if the channel is
shut down and empty.

async recv_metadata(self, Context ctx)#

Receive the next metadata message from the channel.

Parameters:

ctx: The current streaming context.

Returns:

A Message if a message is available, otherwise None if the
metadata queue is shut down and empty.

async send(self, Context ctx, Message msg)#

Send a message into the channel.

Parameters:

ctx: The current streaming context.
msg: Message to move into the channel.

Warning

msg is released and left empty after this call.

async send_metadata(self, Context ctx, Message msg)#

Send a metadata message into the channel.

Parameters:

ctx: The current streaming context.
msg: Metadata message to move into the channel.

Warning

msg is released and left empty after this call.

async shutdown(self, Context ctx)#

Immediately shut down the channel.

Completes when the shutdown has been processed.

Parameters:

ctx: The current streaming context.

Notes

Pending and future send/recv operations will complete with failure.

async shutdown_metadata(self, Context ctx)#

Immediately shut down metadata handling.

Completes when the shutdown has been processed.

Parameters:

ctx: The current streaming context.

class rapidsmpf.streaming.core.Context#

Context for actors (coroutines) in rapidsmpf.

The context owns shared resources used during execution, including the coroutine executor and memory reservation infrastructure.

A Context instance must be created and shut down on the same thread. Shutting down the context from a different thread results in program termination. This is particularly important in coroutine-based code, where execution and stack unwinding may occur on different threads if ownership is not carefully managed.

In Python, it is easy to accidentally keep dangling references to a Context instance, which may delay destruction and cause shutdown to occur on an unintended thread. For this reason, it is strongly recommended to use Context as a context manager (that is, via a with statement), which guarantees that shutdown() is invoked deterministically and on the same thread that created the context.

Parameters:

logger: The logger to use.
br: The buffer resource to use.
options: The configuration options to use. Missing options are read from environment variables.

Methods

`br`(self)	Get buffer resource.
`create_channel`(self)	Create a new channel associated with this context.
`from_options`(cls, Logger logger, ...[, ...])
`logger`(self)	Get the logger.
`memory`(self, MemoryType mem_type)	Get the memory reservation handle for a given memory type.
`options`(self)	Get options.
`shutdown`(self)	Shut down the context.
`spillable_messages`(self)	Get spillable messages.
`statistics`(self)	Get statistics.

Examples

>>> with streaming.Context(
...     logger=...,
...     br=BufferResource(...),
...     options=Options(...),
... ) as ctx:
...     ch = ctx.create_channel()

br(self)#

Get buffer resource.

Returns:

The buffer resource associated with this context.

create_channel(self)#

Create a new channel associated with this context.

Returns:

The newly created channel.

classmethod from_options( cls, Logger logger, DeviceMemoryResource mr, Options options, statistics=None, )#

logger(self)#

Get the logger.

Returns:

The logger associated with this context.

memory(self, MemoryType mem_type)#

Get the memory reservation handle for a given memory type.

Returns an object that coordinates asynchronous memory reservation requests for the specified memory type. The returned instance provides backpressure and global progress guarantees and should be used to reserve memory before performing operations that require memory.

A recommended usage pattern is to reserve all required memory up front as a single atomic reservation. This allows callers to await the reservation and only start executing the operation once all required memory is available.

Parameters:

mem_type: Memory type for which reservations are requested.

Returns:

Handle that coordinates memory reservation requests for the given memory type.

options(self)#

Get options.

Returns:

The options associated with this context.

shutdown(self)#

Shut down the context.

This method is idempotent and only performs shutdown once. Subsequent calls have no effect.

Warning

Shutdown must be initiated from the same thread that constructed the executor. Calling this method from a different thread results in program termination.

spillable_messages(self)#

Get spillable messages.

Returns:

The spillable messages associated with this context.

statistics(self)#

Get statistics.

Returns:

The statistics associated with this context.

class rapidsmpf.streaming.core.FanoutPolicy(*values)#

BOUNDED#

UNBOUNDED#

class rapidsmpf.streaming.core.MemoryReserveOrWait(Options options, MemoryType mem_type, Context ctx)#

Asynchronous coordinator for memory reservation requests.

MemoryReserveOrWait provides a coroutine-based mechanism for reserving memory with backpressure. Callers submit reservation requests via reserve_or_wait(), which suspends until sufficient memory becomes available or progress must be forced.

A background task is spawned on demand to periodically check available memory and fulfill pending requests. If no reservation request can be satisfied within the timeout specified by the "memory_reserve_timeout" option, the scheduler forces progress by selecting the smallest pending request and attempting to reserve memory for it. This attempt may result in an empty reservation if the request still cannot be satisfied.

The timeout provides a global progress guarantee and does not apply to a specific reservation request. Instead, it bounds how long the system may go without satisfying any pending request.

Parameters:

options: Configuration options. The option "memory_reserve_timeout" controls the global progress timeout.
mem_type: The memory type for which reservations are requested.
ctx: Actor context used during construction to read context properties. The context is not kept alive after initialization.

Methods

`reserve_or_wait`(self, size_t size, *, ...)	Attempt to reserve memory, or wait until progress can be made.
`reserve_or_wait_or_fail`(self, size_t size, ...)	Variant of `reserve_or_wait()` that fails if no progress is possible.
`reserve_or_wait_or_overbook`(self, ...)	Variant of `reserve_or_wait()` that allows overbooking on timeout.
`shutdown`(self)	Shut down all pending memory reservation requests.
`size`(self)	Return the number of pending memory reservation requests.

Raises:

RuntimeError: If shutdown occurs before a reservation request can be processed.

async reserve_or_wait( self, size_t size, *, int64_t net_memory_delta, )#

Attempt to reserve memory, or wait until progress can be made.

Submits a memory reservation request and suspends until either sufficient memory becomes available or no reservation request, including other pending requests, makes progress within the configured timeout.

While no pending request fits, spilling is attempted to free the memory the highest-priority one needs, and again just before progress is forced. Both are limited to DEVICE, since SpillManager measures headroom against device memory. Requests of other memory types wait without spilling.

The timeout does not apply specifically to this request. Instead, it bounds when a final recovery attempt begins, not when it completes. If no pending reservation request can be satisfied within the timeout, MemoryReserveOrWait forces progress by selecting the smallest pending request and attempting to reserve memory without spilling queued data. The forced reservation attempt may result in an empty MemoryReservation if the selected request still cannot be satisfied.

When multiple reservation requests are eligible, MemoryReserveOrWait uses net_memory_delta as a heuristic to prefer requests that are expected to reduce memory pressure sooner. The value represents the estimated net change in memory usage after the reservation has been granted and the dependent operation completes (that is, after both reserving size bytes and completing the work that consumes the reservation): - > 0: expected net increase in memory usage - = 0: memory-neutral - < 0: expected net decrease in memory usage

Smaller values have higher priority.

Parameters:

size: Number of bytes to reserve.
net_memory_delta: Estimated net change in memory usage after the reservation has been granted and the dependent operation completes. Smaller values have higher priority.

Returns:

A memory reservation representing the allocated memory. The reservation may be
empty if progress could not be made.

Raises:

RuntimeError: If shutdown occurs before the request can be processed.

Examples

Reading data from disk into memory typically has a positive net_memory_delta because memory usage increases.

A row-wise transformation that retains input and output typically has a net_memory_delta near zero.

Writing data to disk or a reduction that frees inputs typically has a negative net_memory_delta because memory usage decreases.

async reserve_or_wait_or_fail( self, size_t size, *, int64_t net_memory_delta, )#

Variant of reserve_or_wait() that fails if no progress is possible.

This coroutine behaves identically to reserve_or_wait() with respect to request submission, waiting, and progress guarantees until the progress timeout expires.

If no reservation request can be satisfied before the timeout, this method fails instead of forcing progress. Overbooking is not allowed, and no memory reservation is made.

Parameters:

size: Number of bytes to reserve.
net_memory_delta: Heuristic used to prioritize eligible requests. See reserve_or_wait() for details and semantics.

Returns:

The memory reservation representing the allocated memory.

Raises:

RuntimeError: If no progress is possible within the timeout, or shutdown occurs before the request can be processed.

See also

reserve_or_wait

async reserve_or_wait_or_overbook( self, size_t size, *, int64_t net_memory_delta, )#

Variant of reserve_or_wait() that allows overbooking on timeout.

This coroutine behaves identically to reserve_or_wait() with respect to request submission, waiting, and progress guarantees. The only difference is the behavior when the progress timeout expires.

If no reservation request can be satisfied before the timeout, this method attempts to reserve the requested memory by allowing overbooking. This guarantees forward progress, but may exceed the configured memory limits.

Parameters:

size: Number of bytes to reserve.
net_memory_delta: Heuristic used to prioritize eligible requests. See reserve_or_wait() for details and semantics.

Returns:

A pair consisting of:

A MemoryReservation representing the allocated memory.
The number of bytes by which the reservation overbooked the available memory. This value is zero if no overbooking occurred.

Raises:

RuntimeError: If shutdown occurs before the request can be processed.

async shutdown(self)#

Shut down all pending memory reservation requests.

Cancels all pending reservation requests and signals the background periodic memory check task to exit. The returned coroutine completes only after all pending requests have been cancelled and the periodic memory check task has fully exited.

Returns:

A coroutine that completes once shutdown is complete.

size(self)#

Return the number of pending memory reservation requests.

The returned value is a snapshot and may change concurrently as reservation requests are added or fulfilled.

Returns:

The number of outstanding reservation requests.

class rapidsmpf.streaming.core.Message(uint64_t sequence_number, payload)#

A message to be transferred between streaming actors.

Parameters:

sequence_number: Ordering identifier for the message.
payload: A payload object that implements the Payload protocol. The payload is moved into this message.

Attributes:

sequence_number: Return the sequence number of this message.

Methods

`copy`(self, MemoryReservation reservation)	Perform a deep copy of this message and its payload.
`copy_cost`(self)	Return the total memory size required for a deep copy of the payload.
`empty`(self)	Return whether this message is empty.
`get_content_description`(self)	Return a copy of the content description associated with the message.

Warning

payload is released by this call and must not be used afterwards.

copy(self, MemoryReservation reservation)#

Perform a deep copy of this message and its payload.

A new message is created by invoking the registered copy callback, allocating fresh buffers using the provided memory reservation. The resulting message contains a deep copy of the payload while preserving the same metadata and callbacks.

Parameters:

reservation: Memory reservation to use for allocations during the copy.

Returns:

A new message containing a deep copy of the original payload.

Raises:

ValueError: If the message does not support copying.

Examples

>>> res = br.reserve_device_memory_and_spill(
...    msg.copy_cost(), allow_overbooking=False
... )
>>> msg_copy = msg.copy(res)
>>> assert msg_copy.sequence_number == msg.sequence_number

copy_cost(self)#

Return the total memory size required for a deep copy of the payload.

The computed size represents the total amount of memory that must be reserved to duplicate all content buffers of the message, regardless of where they currently reside. For example, if the payload has content in both host and device memory, the returned size is the sum of both.

Returns:

The number of bytes required to perform a deep copy of the message.

empty(self)#

Return whether this message is empty.

Returns:

True if the message is empty; otherwise, False.

get_content_description(self)#

Return a copy of the content description associated with the message.

Returns:

A copy of the message’s content description.

sequence_number#

Return the sequence number of this message.

Returns:

The sequence number.

class rapidsmpf.streaming.core.SpillableMessages(BufferResource br)#

Container for individually spillable messages.

This class manages a collection of messages that can be spilled, extracted, or inspected independently. Each inserted message is assigned a unique identifier that can later be used to extract or spill it. The container is thread-safe for concurrent insertions, extractions, and spills.

Parameters:

br: A BufferResource to keep alive.

Methods

`extract`(self, *, uint64_t mid)	Extract a message by its identifier.
`get_content_descriptions`(self)	Retrieve content descriptions for all messages.
`insert`(self, Message message)	Insert a new message into the container.
`spill`(self, *, uint64_t mid, BufferResource br)	Spill a specific message to an external buffer resource.

Examples

>>> msgs = SpillableMessages(br)
>>> mid = msgs.insert(msg)
>>> msgs.spill(mid=mid, br=br)
>>> recovered = msgs.extract(mid=mid)

extract(self, *, uint64_t mid)#

Extract a message by its identifier.

Parameters:

mid: Identifier of the message to extract.

Returns:

The extracted message instance.

get_content_descriptions(self)#

Retrieve content descriptions for all messages.

Returns:

A dict from message identifiers to content descriptions.

insert(self, Message message)#

Insert a new message into the container.

Parameters:

message: The message to insert.

Returns:

The unique identifier assigned to the inserted message.

spill(self, *, uint64_t mid, BufferResource br)#

Spill a specific message to an external buffer resource.

Parameters:

mid: Identifier of the message to spill.
br: Buffer resource used for spill allocations.

Returns:

The number of bytes spilled.

rapidsmpf.streaming.core.define_actor(*, extra_channels=())#

Create a decorator for defining a Python streaming actor.

The decorated coroutine must take a Context as its first positional argument and return None. When the coroutine finishes (whether successfully or with an exception), the wrapper automatically shuts down: * any channels discovered from the coroutine’s arguments. * all channels listed in extra_channels.

Channels are discovered by recursively inspecting the coroutine’s bound arguments. Mapping values and general iterables are traversed but byte-like objects (str, bytes, bytearray, memoryview) are skipped.

Parameters:

extra_channels: Additional channels to shut down after the decorated coroutine completes.

Returns:

decorator: A decorator for an async function that defines a Python actor.

Raises:

TypeError: If the decorated function is not async.

Examples

In the following example, python_actor is defined as a Python actor. When it completes, ch1 is shut down automatically because it is passed as a coroutine argument, and ch2 is shut down because it is listed in extra_channels.

>>> ch1: Channel[TableChunk] = context.create_channel()
>>> ch2: Channel[TableChunk] = context.create_channel()
>>> @define_actor(extra_channels=(ch2,))
... async def python_actor(ctx: Context, /, ch_in: Channel) -> None:
...     msg = await ch_in.recv()
...     await ch2.send(msg)
... # Calling the coroutine doesn't run it but we can provide its arguments.
>>> actor = python_actor(context, ch_in=ch1)
... # Later we need to call run_actor_network() to actually run the actor.

rapidsmpf.streaming.core.fanout(Context ctx, Channel ch_in, chs_out, FanoutPolicy policy)#

Broadcast messages from one input channel to multiple output channels.

The actor continuously receives messages from the input channel and forwards them to all output channels according to the selected fanout policy.

Each output channel receives a shallow copy of the same message; no payload data is duplicated. All copies share the same underlying payload, ensuring zero-copy broadcast semantics.

Parameters:

ctx: The actor context to use.
ch_in: Input channel from which messages are received.
chs_out: Output channels to which messages are broadcast.
policy: The fanout policy to use. BOUNDED can be used if all output channels are being consumed by independent consumers in the downstream. UNBOUNDED can be used if the output channels are being consumed by a single/ shared consumer in the downstream.

Returns:

Streaming actor representing the fanout operation.

Raises:

ValueError: If an unknown fanout policy is specified.

Notes

Since messages are shallow-copied, releasing a payload (release<T>()) is only valid on messages that hold exclusive ownership of the payload.

>>> import rapidsmpf.streaming.core as streaming
>>> with streaming.Context(...) as ctx:
...     ch_in = ctx.create_channel()
...     ch_out1 = ctx.create_channel()
...     ch_out2 = ctx.create_channel()
...     actor = streaming.fanout(
...         ctx,
...         ch_in,
...         [ch_out1, ch_out2],
...         streaming.FanoutPolicy.BOUNDED,
...     )

async rapidsmpf.streaming.core.reserve_memory( Context ctx, size, *, net_memory_delta, mem_type=MemoryType.DEVICE, allow_overbooking=None, )#

Reserve memory using the context memory reservation mechanism.

Submits a memory reservation request for the specified memory type and suspends until the request is satisfied or no further progress can be made. The behavior when the progress timeout expires depends on whether overbooking is allowed.

This is a convenience helper that returns only the memory reservation. If more control is required, for example inspecting the amount of overbooking, callers should use the context memory reservation system directly, such as ctx.memory(MemoryType.DEVICE).reserve_or_wait_or_overbook(...).

Parameters:

ctx

Actor context used to obtain the memory reservation handle.

size

Number of bytes to reserve.

net_memory_delta

Heuristic used to prioritize eligible requests. See MemoryReserveOrWait.reserve_or_wait() for details and semantics.

mem_type

Memory type for which to reserve memory.

allow_overbooking

Whether to allow overbooking if no progress is possible.

If True, the reservation may overbook memory when no further progress can be made. If False, the call fails when no progress is possible.
If None (the default), the behavior is determined by the configuration option "allow_overbooking_by_default", which is read via ctx.options().

Returns:

The allocated memory reservation.

Raises:

RuntimeError: If shutdown occurs before the request can be processed, or no further progress is possible and overbooking is disabled.

Examples

Reserve device memory inside an actor: >>> res = await reserve_memory( … ctx, … size=1024, … net_memory_delta=0, … allow_overbooking=True, … ) >>> res.size 1024

Disable overbooking and fail if no progress is possible: >>> res = await reserve_memory( … ctx, … size=2048, … net_memory_delta=0, … allow_overbooking=False, … )

rapidsmpf.streaming.core.run_actor_network(Context ctx, *, actors)#

Run streaming actors to completion (blocking).

Accepts a collection of actors. Native C++ actors are moved into the C++ network and executed with minimal Python overhead, while Python actors are gathered and executed on a dedicated event loop.

Parameters:

ctx: Streaming context for execution.
actors: Iterable of actors. Each element is either a native C++ actor or a Python awaitable representing an actor.

Raises:

Exception: Any unhandled exception from an actor is re-raised after execution. If multiple actors raise exceptions, only one is re-raised, and it is unspecified which one.
TypeError: If actors contains an unknown actor type.

Warning

C++ actors are released and must not be used after this call.

Examples

>>> ch: Channel = context.create_channel()
>>> cpp_actor, output = pull_from_channel(context, ch_in=ch)
...
>>> @define_actor()
... async def python_actor(ctx: Context, ch_out: Channel) -> None:
...     # Send one message and close.
...     await ch_out.send(context, Message(42, payload))
...     await ch_out.drain(context)
...
>>> run_actor_network(
...     context,
...     actors=[cpp_actor, python_actor(context, ch_out=ch)]
... )
>>> results = output.release()
>>> results[0].sequence_number
42

async rapidsmpf.streaming.core.shutdown_channels(Context ctx, *chs)#

Shutdown channels, recording and then propagating any exceptions

Parameters:

ctx: Streaming context for channel shutdown
chs: Channels to shutdown

Raises:

Any exceptions that shutting down the channels produces.

Collectives#

Submodule for collective streaming operations.

class rapidsmpf.streaming.coll.AllGather(Context ctx, Communicator comm, int32_t op_id)#

An asynchronous AllGather.

Parameters:

ctx: Streaming context
op_id: Operation id identifying this allgather. Must not be reused while this object is still live.

Attributes:

comm: Get the communicator used by the allgather.

Methods

`extract_all`(self, Context ctx, *, bool ordered)	Suspend and extract all data from the AllGather.
`insert`(self, uint64_t sequence_number, ...)	Insert data into the AllGather.
`insert_finished`(self)	Insert a finished marker into the AllGather.

comm#

Get the communicator used by the allgather.

Returns:

The communicator.

async extract_all(self, Context ctx, *, bool ordered)#

Suspend and extract all data from the AllGather.

Parameters:

ctx: Streaming context.
ordered: Should the extraction be ordered?

Returns:

Awaitable that returns the gathered PackedData.

insert( self, uint64_t sequence_number, PackedData packed_data, )#

Insert data into the AllGather.

Parameters:

sequence_number: Sequence number of this piece of data, used to provide an ordering when extracting.
packed_data: The data to insert.

insert_finished(self)#: Insert a finished marker into the AllGather.

class rapidsmpf.streaming.coll.ShufflerAsync( Context ctx, Communicator comm, int32_t op_id, uint32_t total_num_partitions, PartitionAssignment partition_assignment=PartitionAssignment.ROUND_ROBIN, )#

An asynchronous shuffle.

Parameters:

ctx: Streaming context
comm: The communicator the shuffle is collective over.
op_id: Operation id identifying this shuffle. Must not be reused while this object is still live.
total_num_partitions: Global number of output partitions in the shuffle.
partition_assignment: How to assign partition IDs to ranks: ROUND_ROBIN (default) for load balance (e.g. hash shuffle), or CONTIGUOUS so each rank gets a contiguous range of partition IDs (e.g. for sort so concatenation order matches global order). A custom callable may be supported in the future.

Attributes:

comm: Get the communicator used by the shuffler.

Methods

`extract`(self, uint32_t pid)	Extract all chunks belonging to the specified partition.
`insert`(self, chunks)	Insert data into the shuffle.
`insert_finished`(self, Context ctx)	Insert a finished marker into the shuffle.
`local_partitions`(self)	Return the partition IDs owned by this rank.

comm#

Get the communicator used by the shuffler.

Returns:

The communicator.

extract(self, uint32_t pid)#

Extract all chunks belonging to the specified partition.

Must only be called after awaiting insert_finished().

Parameters:

pid: The partition to extract.

Returns:

list[PackedData]: The PackedData chunks associated with the partition.

insert(self, chunks)#

Insert data into the shuffle.

Parameters:

chunks: Map of partition ID to PackedData associated with that partition.

async insert_finished(self, Context ctx)#

Insert a finished marker into the shuffle.

Notes

This must be awaited before extraction can occur.

local_partitions(self)#

Return the partition IDs owned by this rank.

Returns:

Partition IDs owned by this shuffler.

class rapidsmpf.streaming.coll.SparseAlltoall( Context ctx, Communicator comm, int32_t op_id, srcs, dsts, )#

An asynchronous sparse all-to-all.

Parameters:

ctx: Streaming context.
comm: The communicator the collective is over.
op_id: Operation id identifying this sparse all-to-all. Must not be reused while this object is still live.
srcs: Source ranks this rank receives from.
dsts: Destination ranks this rank sends to.

Attributes:

comm: Get the communicator used by the sparse all-to-all.

Methods

`extract`(self, Rank src)	Extract all data received from the specified source rank.
`insert`(self, Rank dst, PackedData packed_data)	Insert data into the sparse all-to-all.
`insert_finished`(self, Context ctx)	Insert the finished marker and await local completion.

comm#

Get the communicator used by the sparse all-to-all.

Returns:

The communicator.

extract(self, Rank src)#

Extract all data received from the specified source rank.

Must only be called after awaiting insert_finished().

Parameters:

src: Source rank to extract from.

Returns:

list[PackedData]: The PackedData messages received from the source.

insert(self, Rank dst, PackedData packed_data)#

Insert data into the sparse all-to-all.

Parameters:

dst: Destination rank to send to.
packed_data: The data to insert.

async insert_finished(self, Context ctx)#

Insert the finished marker and await local completion.

Notes

This must be awaited before extraction can occur.

rapidsmpf.streaming.coll.allgather( Context ctx, Communicator comm, Channel ch_in, Channel ch_out, int32_t op_id, *, bool ordered, )#

Launch an allgather actor for a single allgather operation.

Streaming variant of the RapidsMPF allgather.

Parameters:

ctx: The actor context to use.
comm: The communicator the allgather is collective over.
ch_in: Input channel that supplies PackedDataChunks to be gathered.
ch_out: Output channel that receives gathered PackedDataChunks.
op_id: Unique (per-communicator) identifier for this allgather operation. Must not be reused until all actors participating in the allgather have shut down.
ordered: Should the output channel provide data in order of input sequence numbers?

Returns:

A streaming actor that finishes when the allgather is complete and ch_out has
been drained.

rapidsmpf.streaming.coll.shuffler( Context ctx, Communicator comm, Channel ch_in, Channel ch_out, int32_t op_id, uint32_t total_num_partitions, PartitionAssignment partition_assignment=PartitionAssignment.ROUND_ROBIN, )#

Launch a shuffler actor for a single shuffle operation.

Streaming variant of the RapidsMPF shuffler that reads packed, partitioned input chunks from an input channel and emits output chunks grouped by partition owner.

Parameters:

ctx: The actor context to use.
comm: The communicator the shuffle is collective over.
ch_in: Input channel that supplies partitioned map chunks to be shuffled.
ch_out: Output channel that receives the grouped (vector) chunks.
op_id: Unique identifier for this shuffle operation. Must not be reused until all actors participating in the shuffle have shut down.
total_num_partitions: Total number of logical partitions to shuffle the data into.
partition_assignment: How to assign partition IDs to ranks: ROUND_ROBIN (default) for load balance (e.g. hash shuffle), or CONTIGUOUS so each rank gets a contiguous range of partition IDs (e.g. for sort so concatenation order matches global order). A custom callable may be supported in the future.

Returns:

A streaming actor that finishes when shuffling is complete and ch_out has
been drained.

Data chunks#

Submodule for generic stream chunks.

class rapidsmpf.streaming.chunks.ArbitraryChunk(obj)#

A chunk containing an arbitrary Python object as a payload.

Parameters:

obj: The payload object.

Methods

`from_message`(Message message)	Construct an ArbitraryChunk by consuming a Message.
`into_message`(self, uint64_t sequence_number, ...)	Move this ArbitraryChunk into a Message.
`release`(self)

Notes

To extract the object from the chunk, use release(). The object is stored in a unique pointer with a custom deleter, so it is safe to drop the chunk in C++: deallocation will acquire the gil and decref the stored object.

static from_message(Message message)#

Construct an ArbitraryChunk by consuming a Message.

Parameters:

message: Message containing an ArbitraryChunk. The message is released and is empty after this call.

Returns:

A new ArbitraryChunk extracted from the given message.

into_message( self, uint64_t sequence_number, Message message, )#

Move this ArbitraryChunk into a Message.

This method is not typically called directly. Instead, it is invoked by the Message constructor when creating a new Message with this ArbitraryChunk as its payload.

Parameters:

sequence_number: Ordering identifier for the message.
message: Message object that will take ownership of this ArbitraryChunk.

Raises:

ValueError: If the provided message is not empty.

Warning

The ArbitraryChunk is released and must not be used after this call.

release(self)#

class rapidsmpf.streaming.chunks.PackedDataChunk#

Methods

`from_message`(Message message, BufferResource br)	Construct a PackedDataChunk by consuming a Message.
`from_packed_data`(PackedData obj, ...)	Construct a PackedDataChunk from an existing PackedData object.
`into_message`(self, uint64_t sequence_number, ...)	Move this PackedDataChunk into a `Message`.
`to_packed_data`(self)	Convert to a PackedData object.

static from_message(Message message, BufferResource br)#

Construct a PackedDataChunk by consuming a Message.

Parameters:

message: Message containing a PackedDataChunk. The message is released and is empty after this call.

Returns:

PackedDataChunk: A new PackedDataChunk extracted from the given message.

static from_packed_data(PackedData obj, BufferResource br)#

Construct a PackedDataChunk from an existing PackedData object.

Parameters:

obj: The PackedData to construct from. The packed data is empty after this call.

Returns:

PackedDataChunk: A new PackedDataChunk from the given object.

into_message( self, uint64_t sequence_number, Message message, )#

Move this PackedDataChunk into a Message.

This method is not typically called directly. Instead, it is invoked by the Message constructor when creating a new Message with this PackedDataChunk as its payload.

Parameters:

sequence_number: Ordering identifier for the message.
message: Message object that will take ownership of this PackedDataChunk.

Raises:

ValueError: If the provided message is not empty.

Warning

The PackedDataChunk is released and must not be used after this call.

to_packed_data(self)#

Convert to a PackedData object.

Returns:

PackedData: A new PackedData from this chunk. The chunk is left empty.

class rapidsmpf.streaming.chunks.PartitionMapChunk#

Methods

`from_message`(Message message, BufferResource br)	Construct a PartitionMapChunk by consuming a `Message`.
`from_packed_data_map`(data, BufferResource br)	Construct a PartitionMapChunk from a mapping of partition ID to `PackedData`.
`into_message`(self, uint64_t sequence_number, ...)	Move this PartitionMapChunk into a `Message`.
`to_packed_data_map`(self)	Extract the partition data as a mapping of partition ID to PackedData.

static from_message(Message message, BufferResource br)#

Construct a PartitionMapChunk by consuming a Message.

Parameters:

message: Message containing a PartitionMapChunk. The message is released and is empty after this call.

Returns:

PartitionMapChunk: A new PartitionMapChunk extracted from the given message.

static from_packed_data_map(data, BufferResource br)#

Construct a PartitionMapChunk from a mapping of partition ID to PackedData.

Parameters:

data: Mapping of partition ID to the PackedData it holds. Each PackedData is consumed and left empty after this call.
br: Buffer resource kept alive for the lifetime of the chunk.

Returns:

PartitionMapChunk: A new PartitionMapChunk owning the given packed data.

Raises:

ValueError: If any of the provided PackedData objects is empty.

into_message( self, uint64_t sequence_number, Message message, )#

Move this PartitionMapChunk into a Message.

This method is not typically called directly. Instead, it is invoked by the Message constructor when creating a new Message with this PartitionMapChunk as its payload.

Parameters:

sequence_number: Ordering identifier for the message.
message: Message object that will take ownership of this PartitionMapChunk.

Raises:

ValueError: If the provided message is not empty.

Warning

The PartitionMapChunk is released and must not be used after this call.

to_packed_data_map(self)#

Extract the partition data as a mapping of partition ID to PackedData.

The chunk is drained and left empty after this call.

Returns:

dict: A dict mapping partition ID to the PackedData it holds.

class rapidsmpf.streaming.chunks.PartitionVectorChunk#

Methods

`from_message`(Message message, BufferResource br)	Construct a PartitionVectorChunk by consuming a `Message`.
`from_packed_data_list`(data, BufferResource br)	Construct a PartitionVectorChunk from a sequence of `PackedData`.
`into_message`(self, uint64_t sequence_number, ...)	Move this PartitionVectorChunk into a `Message`.
`to_packed_data_list`(self)	Extract the partition data as a list of PackedData.

static from_message(Message message, BufferResource br)#

Construct a PartitionVectorChunk by consuming a Message.

Parameters:

message: Message containing a PartitionVectorChunk. The message is released and is empty after this call.

Returns:

PartitionVectorChunk: A new PartitionVectorChunk extracted from the given message.

static from_packed_data_list(data, BufferResource br)#

Construct a PartitionVectorChunk from a sequence of PackedData.

Parameters:

data: Sequence of PackedData objects, stored in order. Each PackedData is consumed and left empty after this call.
br: Buffer resource kept alive for the lifetime of the chunk.

Returns:

PartitionVectorChunk: A new PartitionVectorChunk owning the given packed data.

Raises:

ValueError: If any of the provided PackedData objects is empty.

into_message( self, uint64_t sequence_number, Message message, )#

Move this PartitionVectorChunk into a Message.

This method is not typically called directly. Instead, it is invoked by the Message constructor when creating a new Message with this PartitionVectorChunk as its payload.

Parameters:

sequence_number: Ordering identifier for the message.
message: Message object that will take ownership of this PartitionVectorChunk.

Raises:

ValueError: If the provided message is not empty.

Warning

The PartitionVectorChunk is released and must not be used after this call.

to_packed_data_list(self)#

Extract the partition data as a list of PackedData.

The chunk is drained and left empty after this call.

Returns:

list[PackedData]: A list of PackedData, in order.