Python API Reference#

This page contains the API reference for rapidsmpf.

Integrations#

The subpackages under rapidsmpf.integrations contain integrations with other libraries.

Generic#

RapidsMPF Integrations.

Ray#

Integration for Ray clusters.

class rapidsmpf.integrations.ray.RapidsMPFActor( nranks: int, statistics: Statistics | None = None, )#

RapidsMPFActor is a base class that instantiates a UCXX communicator across all workers.

Parameters:

nranks: The number of workers in the cluster.
statistics: Optional statistics tracking object.

Attributes:

comm: The UCXX communicator object.
statistics: The statistics object used on this actor.

Methods

`is_initialized`()	Check if the communicator is initialized.
`nranks`()	Get the number of ranks in the UCXX communicator.
`rank`()	Get the rank of the worker, as inferred from the UCXX communicator.
`setup_root`()	Setup root communicator in the cluster.
`setup_worker`(root_address_bytes)	Setup the worker in the cluster once the root is initialized.
`to_string`()	Return a string representation of the actor.

Examples

>>> @ray.remote(num_cpus=1)
... class DummyActor(RapidsMPFActor): ...
>>> actors = setup_ray_ucx_cluster(DummyActor, 2)
>>> ray.get([actor.status_check.remote() for actor in actors]

property comm: Communicator#

The UCXX communicator object.

Returns:

The UCXX communicator object if initialized, otherwise None

Raises:

RuntimeError: If the communicator is not initialized.

Notes

This property is not meant to be called remotely from the client. Then Ray will attempt to serialize the Communicator object, which will fail. Instead, the subclasses can use the comm property to access the communicator. For example, to create a Shuffle operation

is_initialized() → bool#

Check if the communicator is initialized.

Returns:

True if the communicator is initialized, False otherwise.

nranks() → int#

Get the number of ranks in the UCXX communicator.

Returns:

The number of ranks in the UCXX communicator

rank() → int#

Get the rank of the worker, as inferred from the UCXX communicator.

Returns:

The rank of the worker

setup_root() → tuple[int, bytes]#

Setup root communicator in the cluster.

Returns:

rank: The rank of the root.
root_address_bytes: The address of the root.

setup_worker(root_address_bytes: bytes) → None#

Setup the worker in the cluster once the root is initialized.

This method needs to be called by every worker including the root.

Parameters:

root_address_bytes: The address of the root.

property statistics: Statistics#

The statistics object used on this actor.

Returns:

Statistics object.

to_string() → str#

Return a string representation of the actor.

Returns:

A string representation of the actor

Raises:

RuntimeError: If the communicator is not initialized.

rapidsmpf.integrations.ray.setup_ray_ucxx_cluster(

actor_cls: ActorClass,

num_workers: int,

*args: Any,

**kwargs: Any,

) → list[ActorHandle]#

A utility method to setup the UCXX communication using RapidsMPFActor actor objects.

Parameters:

actor_cls: The actor class to be instantiated in the cluster.
num_workers: The number of workers in the cluster.
*args: Additional arguments to be passed to the actor class.
**kwargs: Additional keyword arguments to be passed to the actor class.

Returns:

gpu_actors: A list of actors in the cluster.

Shuffler#

The Shuffler interface for RapidsMPF.

class rapidsmpf.shuffler.PartitionAssignment(*values)#

ROUND_ROBIN#

CONTIGUOUS#

class rapidsmpf.shuffler.Shuffler( Communicator comm, int32_t op_id, uint32_t total_num_partitions, BufferResource br, PartitionAssignment partition_assignment=PartitionAssignment.ROUND_ROBIN, )#

Shuffle service for partitioned data.

The rapidsmpf.shuffler.Shuffler class provides an interface for performing a shuffle operation on partitioned data. It uses a distribution scheme to distribute and collect data chunks across different ranks.

Parameters:

comm: The communicator to use for data exchange between ranks.
op_id: The operation ID of the shuffle. Must have a value between 0 and max_concurrent_shuffles-1.
total_num_partitions: Total number of partitions in the shuffle.
br: The buffer resource used to allocate temporary storage and shuffle results.
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:

max_concurrent_shuffles: Maximum number of concurrent shufflers.

Methods

`extract`(self, uint32_t pid)	Extract all chunks of the specified partition.
`finished`(self)	Check if all partitions are finished.
`insert_chunks`(self, chunks)	Insert a batch of packed (serialized) chunks into the shuffle.
`insert_finished`(self)	Signal that no more data will be inserted into the shuffle.
`local_partitions`(self)	Return the partition IDs owned by this rank.
`shutdown`(self)	Shutdown the shuffle, blocking until all inflight communication is completed.
`wait`(self)	Wait for all partitions to finish (blocking).

Notes

This class is designed to handle distributed operations by partitioning data and redistributing it across ranks in a cluster. It operates on caller-provided packed payloads and is independent of any particular dataframe implementation.

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.

comm#

Get the communicator used by the shuffler.

Returns:

The communicator.

extract(self, uint32_t pid)#

Extract all chunks of the specified partition.

Parameters:

pid: The partition ID to extract chunks for.

Returns:

A list of packed data belonging to the specified partition.

finished(self)#

Check if all partitions are finished.

This method verifies if all partitions have been completed, meaning all chunks have been inserted and no further data is expected from neither the local nor any remote nodes.

Returns:

True if all partitions are finished, otherwise False.

insert_chunks(self, chunks)#

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

Parameters:

chunks: A map where keys are partition IDs (int) and values are packed data (PackedData).

Notes

This method adds the given chunks to the shuffle, associating them with their respective partition IDs.

insert_finished(self)#

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.

local_partitions(self)#

Return the partition IDs owned by this rank.

Returns:

Partition IDs owned by this shuffler.

shutdown(self)#

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

Raises:

RuntimeError: If the shuffler is already inactive.

Notes

This method ensures that all pending shuffle operations and communications are completed before shutting down. It blocks until no inflight operations remain.

wait(self)#

Wait for all partitions to finish (blocking).

This method blocks until all partitions are finished and ready to be extracted.

Communicator#

Submodule for communication abstraction (e.g. UCXX and MPI).

rapidsmpf.communicator.COMMUNICATORS = ('single', 'ucxx', 'mpi')#

Tuple of available communicators.

RapidsMPF includes a collection of communicator backends, available as submodules under rapidsmpf.communicator.*. Typically, the Conda distribution includes both UCXX and MPI support, while the PIP installation generally supports only UCXX.

class rapidsmpf.communicator.communicator.Communicator#

Abstract base class for a communication mechanism between nodes.

Provides an interface for sending and receiving messages between nodes, supporting asynchronous operations, GPU data transfers, and custom logging. Concrete implementations must define the virtual methods to enable specific communication backends.

Attributes:

logger: Get the logger.
nranks: Get the total number of ranks.
progress_thread: Get the communicator’s progress thread.
rank: Get the rank of this communication node.

Methods

get_str(self)

Get a string representation of the communicator.

Notes

This class is designed as an abstract base class, meaning it cannot be instantiated directly. Subclasses are required to implement the necessary methods to support the desired communication backend and functionality.

get_str(self)#

Get a string representation of the communicator.

Returns:

A string describing the communicator

logger#

Get the logger.

Returns:

A logger instance.

nranks#

Get the total number of ranks.

Returns:

Total number of ranks.

progress_thread#

Get the communicator’s progress thread.

Returns:

The progress thread.

rank#

Get the rank of this communication node.

Returns:

The rank.

class rapidsmpf.communicator.communicator.LOG_LEVEL(*values)#

NONE#

PRINT#

WARN#

INFO#

DEBUG#

TRACE#

class rapidsmpf.communicator.communicator.Logger#

Logger.

To control the verbosity level, set the environment variable RAPIDSMPF_LOG:

NONE: No logging.
PRINT: General print messages.
WARN: Warning messages (default)
INFO: Informational messages.
DEBUG: Debug messages.
TRACE: Trace messages.

Attributes:

verbosity_level: Get the verbosity level of the logger.

Methods

`debug`(self, str msg)	Logs a debug message.
`info`(self, str msg)	Logs an informational message.
`print`(self, str msg)	Logs a print message.
`trace`(self, str msg)	Logs a trace message.
`warn`(self, str msg)	Logs a warning message.

debug(self, str msg)#

Logs a debug message.

Parameters:

msg: The message to log.

info(self, str msg)#

Logs an informational message.

Parameters:

msg: The message to log.

print(self, str msg)#

Logs a print message.

Parameters:

msg: The message to log.

trace(self, str msg)#

Logs a trace message.

Parameters:

msg: The message to log.

verbosity_level#

Get the verbosity level of the logger.

Returns:

The verbosity level.

warn(self, str msg)#

Logs a warning message.

Parameters:

msg: The message to log.

MPI Communicator#

rapidsmpf.communicator.mpi.new_communicator( Intracomm comm, Options options, ProgressThread progress_thread, )#

Create a new RapidsMPF-MPI communicator based on an existing mpi4py communicator.

Parameters:

comm: The existing mpi communicator from mpi4py.
options: Configuration options.
progress_thread: Progress thread for the communicator.

Returns:

A new RapidsMPF-MPI communicator.

UCXX Communicator#

ucxx-based implementation of a RapidsMPF Communicator.

rapidsmpf.communicator.ucxx.barrier(Communicator comm)#

Execute a barrier on the UCXX communicator.

Ensures all ranks connected to the root and all ranks reached the barrier before continuing.

Notes

Executing this barrier is required after the ranks are bootstrapped to ensure everyone is connected to the root. An alternative barrier, such as MPI_Barrier will not suffice for that purpose.

rapidsmpf.communicator.ucxx.get_root_ucxx_address(Communicator comm)#

Get the address of the communicator’s UCXX worker.

This function is intended to be called from the root rank to communicate to other processes how to reach the root, but it will return the address of UCXX worker of other ranks too.

Parameters:

comm: The RapidsMPF-UCXX communicator.

Returns:

A bytes sequence with the UCXX worker address.

Raises:

NotImplementedError: If the communicator was created with a HostPortPair, which is not yet supported.

rapidsmpf.communicator.ucxx.new_communicator( Rank nranks, UCXWorker ucx_worker, UCXAddress root_ucxx_address, Options options, ProgressThread progress_thread, )#

Create a new UCXX communicator with the given number of ranks.

An existing UCXWorker may be specified, otherwise one will be created. The root rank is created if no root_ucxx_address is specific, all other ranks must specify the the address of the root rank via that argument.

Parameters:

nranks: The number of ranks in the cluster.
ucx_worker: An existing UCXX worker to use if specified, otherwise one will be created.
root_ucxx_address: The UCXX address of the root rank (only specified for non-root ranks).
options: Configuration options.
progress_thread: Progress thread for the communicator.

Returns:

A new RapidsMPF-UCXX communicator.

Buffer#

Submodule for memory abstraction.

class rapidsmpf.memory.buffer.Buffer#

A stream-ordered host or device memory buffer managed by a BufferResource.

Buffers are not constructed directly; use make_buffer() to obtain one.

Attributes:

mem_type: Memory type of this buffer.
size: Size of the buffer in bytes.

Methods

host_view(self)

Context manager providing exclusive writable host access to the buffer.

host_view(self)#

Context manager providing exclusive writable host access to the buffer.

Acquires an exclusive lock on the buffer for the duration of the with block, preventing concurrent stream-ordered operations on the C++ side. The lock is released (and the returned memoryview must not be used) once the block exits.

Returns:

A context manager that yields a writable memoryview of the buffer.

Raises:

TypeError: If the buffer is not a host buffer (HOST or PINNED_HOST).
RuntimeError: If the buffer is already locked or a stream-ordered write is still in flight (is_latest_write_done() == False).

Examples

>>> with buf.host_view() as mv:
...     mv[:] = b"\x00" * buf.size

mem_type#: Memory type of this buffer.

size#: Size of the buffer in bytes.

class rapidsmpf.memory.buffer.BufferHostView#

Context manager providing exclusive writable host access to a Buffer.

Not constructed directly; use Buffer.host_view() to obtain one.

class rapidsmpf.memory.buffer.MemoryType(*values)#

DEVICE#

PINNED_HOST#

HOST#

class rapidsmpf.memory.memory_reservation.MemoryReservation#

Represents a reservation for future memory allocation.

A reservation is created by BufferResource.reserve and must be used when allocating buffers through the same BufferResource.

Attributes:

br: Get the buffer resource associated with this reservation.
mem_type: Get the type of memory associated with this reservation.
size: Get the remaining size of the reserved memory.

Methods

clear(self)

Clear the remaining size of the reservation.

br#

Get the buffer resource associated with this reservation.

Returns:

The buffer resource associated with this reservation.

clear(self)#

Clear the remaining size of the reservation.

Resets the reservation so that any remaining, unconsumed bytes are released back to the underlying memory resource. After this call, the reservation has a remaining size of zero and cannot be used to satisfy further allocations.

mem_type#

Get the type of memory associated with this reservation.

Returns:

The memory type associated with this reservation.

size#

Get the remaining size of the reserved memory.

Returns:

The size of the reserved memory in bytes.

rapidsmpf.memory.memory_reservation.opaque_memory_usage(MemoryReservation reservation)#

Associate untracked memory usage with an existing reservation.

This context manager is intended for code paths that use memory outside of RapidsMPF’s memory reservation system, for example internal allocations in third-party libraries. The memory may be of any type covered by a MemoryReservation, most commonly device memory.

While the context is active, the provided memory reservation is considered consumed by the enclosed code block. On exit, the reservation is cleared, releasing any remaining, unconsumed bytes back to the underlying memory resource.

Parameters:

reservation: Memory reservation that accounts for the untracked memory usage.

Yields:

The same reservation, which may be passed to APIs that require an explicit
reservation object.

Examples

Account for allocations outside RapidsMPF: >>> with opaque_memory_usage(ctx, reservation): … # library call that allocates memory unknown to RapidsMPF. … result = library_op(…)

class rapidsmpf.memory.buffer_resource.BufferResource#

Class managing buffer resources.

This class handles memory allocation and transfers between different memory types (e.g., host and device). All memory operations in RapidsMPF, such as those performed by the Shuffler, rely on a buffer resource for memory management.

Parameters:

device_mr: The RMM device memory resource used for device allocations. To ensure allocations are tracked for memory-limit accounting and statistics, use BufferResource.device_mr instead of the original device_mr after construction.
pinned_pool_properties: Configuration for the pinned host memory pool used for PINNED_HOST allocations, as a PinnedPoolProperties. When None (the default), pinned host allocations are disabled and any attempt to allocate pinned memory will fail regardless of any memory_limits entry for PINNED_HOST. When provided, pinned host memory must be supported on this system (see is_pinned_memory_resources_supported()); otherwise a RuntimeError is raised.
memory_limits: Optional mapping from MemoryType to an integer byte limit. Memory types not present in the mapping are treated as unlimited.
periodic_spill_check: Enable periodic spill checks. A dedicated thread continuously checks and performs spilling based on memory availability. The value of periodic_spill_check is used as the pause between checks (in seconds). If None, no periodic spill check is performed.
stream_pool: Optional CUDA stream pool to use. If None, a new pool with 16 streams will be created. Must be an instance of rmm.pylibrmm.cuda_stream_pool.CudaStreamPool.
statistics: The statistics instance to use. If None, a disabled statistics instance will be created.

Attributes:

device_mr: The tracked device memory resource.
pinned_mr: The memory resource used for pinned host memory allocations.
spill_manager
statistics: Gets the statistics instance associated with this buffer resource.
stream_pool: The stream pool associated with this buffer resource.

Methods

`device_mr_adaptor`(self)	The internal device memory resource adaptor with a back-reference installed.
`from_options`(cls, DeviceMemoryResource mr, ...)	Construct a BufferResource from configuration options.
`make_buffer`(self, size_t size, ...)	Allocate a buffer backed by the given memory reservation.
`memory_available`(self, MemoryType mem_type)	Get the current available memory of the specified memory type.
`memory_available_for_reservation`(self, ...)	Get the memory available to a new reservation, in bytes.
`release`(self, MemoryReservation reservation, ...)	Consume a portion of the reserved memory.
`reserve`(self, MemoryType mem_type, ...)	Reserve an amount of the specified memory type.
`reserve_device_memory_and_spill`(self, ...)	Reserve device memory and spill if necessary.
`reserve_or_fail`(self, size_t size, ...)	Make a memory reservation or fail based on the given order of memory types.
`set_memory_limit`(self, MemoryType mem_type, ...)	Set the byte limit for the specified memory type.

Notes

Allocation tracking only applies to allocations routed through this BufferResource.

To ensure allocations are included in memory-limit accounting and statistics, use BufferResource.device_mr for all CUDA allocations associated with this resource.

Allocations performed through other memory resources, including the original resource passed to the constructor or allocations made outside BufferResource, are not tracked by this class.

device_mr#

The tracked device memory resource.

Allocations made through this resource count against the BufferResource memory limits and appear in its statistics.

Returns:

The tracked device memory resource.

device_mr_adaptor(self) → RmmResourceAdaptor#

The internal device memory resource adaptor with a back-reference installed.

Returns a copyable RmmResourceAdaptor that holds shared ownership of this BufferResource, keeping it alive for as long as the returned adaptor (or any copies of it) lives.

This is the only way to obtain an RmmResourceAdaptor; use it when you need to pass one to APIs that copy the adaptor, such as memory_profiling() or report().

Returns:

A back-ref’d RmmResourceAdaptor whose copies keep this BufferResource
alive.

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

Construct a BufferResource from configuration options.

This factory method creates a BufferResource using configuration options to initialize all components. The supplied device memory resource is wrapped internally for allocation tracking — callers don’t need to pre-wrap it.

Parameters:

mr: A device-accessible RMM memory resource.
options: Configuration options.
statistics: The statistics instance to use. The caller is responsible for creating and owning this object. Defaults to Statistics.disabled().

Returns:

A BufferResource instance configured according to the options.

make_buffer( self, size_t size, Stream stream, MemoryReservation reservation, )#

Allocate a buffer backed by the given memory reservation.

Parameters:

size: Size of the buffer in bytes. Must not exceed the reservation size.
stream: CUDA stream to associate with the buffer.
reservation: Memory reservation that covers this allocation. The reservation’s memory type determines whether the buffer is device or host backed.

Returns:

A Buffer of the requested size.

Raises:

ValueError: If size exceeds the reservation size.

memory_available(self, MemoryType mem_type)#: Get the current available memory of the specified memory type.

memory_available_for_reservation( self, MemoryType mem_type, )#

Get the memory available to a new reservation, in bytes.

A snapshot of memory_available(mem_type) minus the outstanding reservations of that memory type. May be negative.

Parameters:

mem_type: The memory type to query.

Returns:

The memory available for reservation, in bytes.

pinned_mr#

The memory resource used for pinned host memory allocations.

The returned handle holds shared ownership of this BufferResource, keeping it alive for as long as the handle (or any copy of it) lives.

Returns:

The pinned host memory resource, or None if pinned host allocations
are disabled.

release( self, MemoryReservation reservation, size_t size, )#

Consume a portion of the reserved memory.

Reduces the remaining size of the reserved memory by the specified amount.

Parameters:

reservation: The memory reservation to consume from.
size: The number of bytes to consume.

Returns:

The remaining size of the reserved memory after consumption.

Raises:

ReservationError: If the released size exceeds the total reserved size.

reserve( self, MemoryType mem_type, size_t size, *, bool allow_overbooking, )#

Reserve an amount of the specified memory type.

Creates a new reservation of the specified size and memory type to inform the system about upcoming buffer allocations.

If overbooking is allowed, a reservation of the requested size is returned even if the memory is not currently available. In that case, the caller must guarantee that at least the overbooked amount of memory will be freed before the reservation is used.

If overbooking is not allowed, a reservation of size zero is returned on failure.

Parameters:

mem_type: The target memory type.
size: The number of bytes to reserve.
allow_overbooking: Whether overbooking is permitted.

Returns:

A tuple (reservation, overbooked_bytes):

On success, the reservation’s size equals size.
On failure, the reservation’s size equals zero (a zero-sized reservation never fails).

reserve_device_memory_and_spill( self, size_t size, *, bool allow_overbooking, )#

Reserve device memory and spill if necessary.

Attempts to reserve the requested amount of device memory. If insufficient memory is available, spilling is triggered to free space. When overbooking is allowed, the reservation may succeed even if spilling was not sufficient to fully satisfy the request.

Parameters:

size: The amount of memory to reserve.
allow_overbooking: Whether to allow overbooking. If false, ensures enough memory is freed to satisfy the reservation. If true, the reservation may succeed even if spilling was insufficient.

Returns:

The resulting memory reservation.

Raises:

ReservationError: If overbooking is disabled and the buffer resource cannot free enough device memory through spilling to satisfy the request.

reserve_or_fail(self, size_t size, list mem_types)#

Make a memory reservation or fail based on the given order of memory types.

Attempts to reserve memory by iterating over mem_types in the given order of preference. For each memory type, a reservation without overbooking is requested. If no memory type can satisfy the request, a RuntimeError is raised.

Parameters:

size: The number of bytes to reserve.
mem_types: List of MemoryType values specifying the order of preference in which memory types are tried.

Returns:

A MemoryReservation for the first memory type that could satisfy
the request.

Raises:

RuntimeError: If no memory type in mem_types could satisfy the reservation.

set_memory_limit( self, MemoryType mem_type, int64_t limit, )#

Set the byte limit for the specified memory type.

The store is atomic, but readers (e.g. memory_available() and reserve()) observe the limit and the allocation count independently. A concurrent set_memory_limit() call can change the limit between a caller’s read of memory_available() and a subsequent allocation decision; callers that need a coherent view must serialize updates with higher-level synchronization.

Parameters:

mem_type: The memory type whose limit is being updated.
limit: The new byte limit. Negative values are allowed; they make memory_available(mem_type) always negative and so trigger continuous spilling.

statistics#

Gets the statistics instance associated with this buffer resource.

Returns:

The Statistics instance.

stream_pool#

The stream pool associated with this buffer resource.

Returns:

An RMM CudaStreamPool.

class rapidsmpf.memory.buffer_resource.OwningDeviceMemoryResource#

Owning DeviceMemoryResource.

Useful for exposing device memory resources to Python in a form that is compatible with cuDF/RMM APIs while preserving ownership semantics.

Methods

`allocate`(self, size_t nbytes, ...)	Allocate `nbytes` bytes of memory.
`deallocate`(self, uintptr_t ptr, ...)	Deallocate memory pointed to by `ptr` of size `nbytes`.

Notes

RMM does not currently provide an equivalent owning wrapper. If one is added in the future, this class can likely be replaced by the RMM-provided implementation.

rapidsmpf.memory.buffer_resource.device_limit_from_options(Options options)#

Get the spill_device_limit parameter from configuration options.

Reads the spill_device_limit option, falling back to 80% of total device memory when unset.

Parameters:

options: Configuration options.

Returns:

int: The device memory limit in bytes.

rapidsmpf.memory.buffer_resource.periodic_spill_check_from_options(Options options)#

Get the periodic_spill_check parameter from configuration options.

Parameters:

options: Configuration options.

Returns:

The duration of the pause between spill checks in seconds, or None if
periodic spill checks are disabled.

rapidsmpf.memory.buffer_resource.stream_pool_from_options(Options options)#

Create a new CUDA stream pool from configuration options.

Parameters:

options: Configuration options.

Returns:

Pool of CUDA streams used throughout RapidsMPF for operations that do not take
an explicit CUDA stream.

class rapidsmpf.memory.pinned_memory_resource.PinnedMemoryResource(*args, **kwargs)#

Opaque handle to a pinned (page-locked) host memory resource.

The resource provides pinned host memory using a pool, enabling higher bandwidth and lower latency for device transfers compared to regular pageable host memory.

Construction

This class cannot be constructed directly. A pinned memory resource is owned by a BufferResource (which installs the back-reference that makes the handle copyable). Configure pinned memory on a BufferResource and obtain the handle via pinned_mr.

The returned handle holds shared ownership of its owning BufferResource, so it (and any copy of it) keeps the BufferResource alive.

Attributes:

enabled: PinnedMemoryResource.enabled: bool

Methods

`allocate`(self, size_t nbytes, Stream stream)	Allocate pinned host memory associated with a CUDA stream.
`deallocate`(self, size_t ptr, size_t nbytes, ...)	Deallocate pinned host memory associated with a CUDA stream.

allocate( self, size_t nbytes, Stream stream, ) → int#

Allocate pinned host memory associated with a CUDA stream.

Parameters:

nbytes: Number of bytes to allocate.
stream: CUDA stream to associate with the allocation.

Returns:

Integer address of the allocated memory.

deallocate( self, size_t ptr, size_t nbytes, Stream stream, ) → None#

Deallocate pinned host memory associated with a CUDA stream.

Parameters:

ptr: Integer address previously returned by allocate().
nbytes: Number of bytes originally allocated.
stream: CUDA stream associated with the allocation.

enabled#

PinnedMemoryResource.enabled: bool

Whether this handle wraps a valid pinned memory resource.

class rapidsmpf.memory.pinned_memory_resource.PinnedPoolProperties( initial_pool_size: int = 0, max_pool_size: object = None, numa_id: object = None, )#

Configuration for a pinned (page-locked) host memory pool.

Pass an instance to BufferResource to enable pinned host memory; passing None instead disables it. The pool is only created when pinned host memory is supported on this system (see is_pinned_memory_resources_supported()).

Attributes:

initial_pool_size: Initial size of the pinned host memory pool in bytes. The initial size is important for pinned-memory performance, especially for the first allocation. Defaults to 0.
max_pool_size: Maximum size of the pinned host memory pool in bytes, or None for no limit. Defaults to None.
numa_id: NUMA node from which pinned host memory should be allocated, or None to use the NUMA node of the calling thread. Defaults to None.

rapidsmpf.memory.pinned_memory_resource.is_pinned_memory_resources_supported() → bool#

Check whether pinned memory resources are supported for the current CUDA version.

RapidsMPF requires CUDA 12.6 or newer to support pinned memory resources.

class rapidsmpf.memory.packed_data.PackedData#

Methods

`from_device_buffer`(cls, ...)	Construct a PackedData from an rmm device buffer and host metadata.
`from_host_bytes`(cls, const uint8_t[, ...)	Construct a PackedData from raw host bytes.
`to_host_bytes`(self)	Extract the host bytes from this PackedData.

classmethod from_device_buffer( cls, DeviceBuffer gpu_data, const uint8_t[::1] metadata, Stream stream, BufferResource br, )#

Construct a PackedData from an rmm device buffer and host metadata.

Takes ownership of gpu_data; the input buffer is left empty after this call. The metadata bytes are copied into a host-side metadata buffer.

Parameters:

gpu_data: Device buffer holding the data payload. Consumed by this call.
metadata: Contiguous buffer of host bytes (bytes, bytearray, or buffer-protocol object). Must be non-empty.
stream: CUDA stream used to take ownership of the device buffer.
br: Buffer resource for memory management.

Returns:

A new PackedData instance owning the device buffer.

classmethod from_host_bytes( cls, const uint8_t[::1] data, BufferResource br, )#

Construct a PackedData from raw host bytes.

The bytes are stored in the data buffer (as host memory) with minimal metadata. This is useful for scalar allreduce operations.

Note: This makes a copy of the input data.

Parameters:

data: Contiguous buffer of bytes (bytes, bytearray, or buffer-protocol object).
br: Buffer resource for memory allocation.

Returns:

A new PackedData instance containing the bytes.

to_host_bytes(self) → bytes#

Extract the host bytes from this PackedData.

Returns the bytes stored in the data buffer. Works with both host and device memory buffers.

Returns a copy of the bytes stored in the data buffer. Works with both host and device memory buffers. The method synchronizes with the underlying buffer’s CUDA stream before returning.

Returns:

The raw bytes.

Raises:

ValueError: If the PackedData is empty.

class rapidsmpf.memory.scoped_memory_record.ScopedMemoryRecord#

Scoped memory record for tracking memory usage statistics.

Methods

`current`(self)	Current memory usage in bytes.
`num_current_allocs`(self)	Number of currently active (non-deallocated) allocations.
`num_total_allocs`(self)	Total number of allocations performed.
`peak`(self)	Peak memory usage in bytes.
`record_allocation`(self, uint64_t nbytes)	Record a memory allocation event.
`record_deallocation`(self, uint64_t nbytes)	Record a memory deallocation event.
`total`(self)	Total number of bytes allocated over the lifetime.

current(self)#

Current memory usage in bytes.

Returns:

Current memory usage in bytes.

num_current_allocs(self)#

Number of currently active (non-deallocated) allocations.

Returns:

Number of active allocations.

num_total_allocs(self)#

Total number of allocations performed.

Returns:

Number of total allocations.

peak(self)#

Peak memory usage in bytes.

Returns:

Peak memory usage in bytes.

record_allocation(self, uint64_t nbytes)#

Record a memory allocation event.

Updates internal statistics for memory allocation and adjusts peak usage if the current memory usage exceeds the previous peak.

Parameters:

nbytes: The number of bytes allocated.

record_deallocation(self, uint64_t nbytes)#

Record a memory deallocation event.

Updates internal statistics to reflect memory being freed.

Parameters:

nbytes: The number of bytes deallocated.

total(self)#

Total number of bytes allocated over the lifetime.

Returns:

Total allocated bytes.

Config Options#

class rapidsmpf.config.Optional(value)#

Represents an option value that can be explicitly disabled.

This class wraps an option value and interprets certain strings as indicators that the value is disabled (case-insensitive): {“false”, “no”, “off”, “disable”, “disabled”}.

This is typically used to simplify optional or Optional options with Options.get_or_default().

Parameters:

value: The input value to interpret.

Attributes:

value: The raw input value, unless it matched a disable keyword, in which case the value is None.

Examples

>>> from rapidsmpf.config import Optional, Options
>>> Optional("OFF").value
None

>>> Optional("no").value
None

>>> Optional("100").value
'100'

>>> Optional("").value
''

>>> opts = Options()
>>> opts.get_or_default(
...     "periodic_spill_check",
...     default_value=Optional(1e-3)
... ).value
0.001

class rapidsmpf.config.OptionalBytes(value)#

Represents a byte-sized option that can be explicitly disabled.

This class is a specialization of Optional that interprets the input as a human-readable byte size string (e.g., “100 MB”, “1KiB”, “1e6”). If the input is one of the disable keywords (e.g., “off”, “no”, “false”), the value is treated as disabled (None). Otherwise, it is parsed to an integer number of bytes using rapidsmpf.utils.string.parse_bytes().

This is useful for configuration options that may be set to a size limit or explicitly turned off.

Parameters:

value: A human-readable byte size (e.g., “1MiB”, “100 MB”) or a disable keyword (case-insensitive), or an integer number of bytes.

Attributes:

value: The size in bytes, or None if disabled.

Examples

>>> from rapidsmpf.config import OptionalBytes
>>> OptionalBytes("1KiB").value
1024

>>> OptionalBytes("OFF").value is None
True

>>> OptionalBytes(2048).value
2048

class rapidsmpf.config.Options#

Initialize an Options object with a dictionary of string options.

Parameters:

options_as_strings: A dictionary representing option names and their corresponding values.

Methods

`deserialize`(bytes serialized_buffer)	Deserialize a binary buffer into an Options object.
`get`(self, str key, *, return_type, factory)	Retrieves a configuration option by key.
`get_or_default`(self, str key, *, default_value)	Retrieve a configuration option by key, using a default value if not present.
`get_strings`(self)	Get all option key-value pairs as strings.
`insert_if_absent`(self, dict options_as_strings)	Insert multiple options if they are not already present.
`serialize`(self)	Serialize the Options object into a binary buffer.

static deserialize(bytes serialized_buffer)#

Deserialize a binary buffer into an Options object.

This method reconstructs an Options instance from a byte buffer produced by the Options.serialize() method.

See Options.serialize() for the binary format.

Parameters:

serialized_buffer: A buffer containing serialized options in the defined binary format.

Returns:

Options: A reconstructed Options instance containing the deserialized key-value pairs.

Raises:

ValueError: If the input buffer is malformed or inconsistent with the expected format.

get(self, str key, *, return_type, factory)#

Retrieves a configuration option by key.

If the option is not present, it is constructed using the provided factory function, which receives the string representation of the option (or an empty string if unset). The option is cached after the first access.

The option is cast to the specified return_type. To be accessible from C++, it must be one of: bool, int, float, str. Otherwise, it is stored as a PyObject*.

Once a key has been accessed with a particular return_type, subsequent calls to get with the same key must use the same return_type. Using a different type for the same key will result in a TypeError.

Parameters:

key: The option key. Should be in lowercase.
return_type: The return type. To be accessible from C++, it must be one of: bool, int, float, str. Use object to indicate any Python type.
factory: A factory function that constructs an instance of the desired type from a string representation.

Returns:

The value of the requested option, cast to the specified return_type.

Raises:

ValueError: If the return_type is unsupported, or if the stored option type does not match the expected type.
TypeError: If the option has already been accessed with a different return_type.

Warning

The factory must not access the Options instance, as this may lead to a deadlock due to internal locking.

get_or_default(self, str key, *, default_value)#

Retrieve a configuration option by key, using a default value if not present.

This is a convenience wrapper around get() that uses the type of the default_value as the return type and provides a default factory that parses a string into that type.

Parameters:

key: The name of the option to retrieve.
default_value: The default value to return if the option is not set. Its type is used to determine the expected return type.

Returns:

The value of the option if it exists and can be parsed to the type of
default_value, otherwise default_value.

Raises:

ValueError: If the stored option value cannot be parsed to the required type, or if key has a canonical default registered in rapidsmpf.config_defaults.DEFAULTS (call get() instead for those options and let the registered default apply).
TypeError: If the option has already been accessed with a different return type.

Notes

This method infers the return type from type(default_value).
If default_value is used, it will be cached and reused for subsequent accesses of the same key.

Examples

>>> opts = Options()
>>> opts.get_or_default("debug", default_value=False)
False
>>> opts.get_or_default("timeout", default_value=1.5)
1.5
>>> opts.get_or_default("level", default_value="info")
'info'

get_strings(self)#

Get all option key-value pairs as strings.

Options that do not have a string representation, such as options inserted as typed values in C++ are included with an empty string value.

Returns:

A dictionary containing all stored options, where the keys and values are
both strings.

insert_if_absent(self, dict options_as_strings)#

Insert multiple options if they are not already present.

Attempts to insert each key-value pair from the provided dictionary, skipping keys that already exist in the options.

Parameters:

options_as_strings: Dictionary of option keys mapped to their string representations. Keys are inserted only if they do not already exist. The keys are trimmed and converted to lower case before insertion.

Returns:

Number of newly inserted options (0 if none were added).

serialize(self) → bytes#

Serialize the Options object into a binary buffer.

This method produces a compact binary representation of the internal key-value options. The format is suitable for storage or transmission and can be later restored using Options.deserialize().

The binary format is:

[uint64_t count] — number of key-value pairs
[count * 2 * uint64_t] — offset pairs (key_offset, value_offset)
[raw bytes] — key and value strings stored contiguously

Returns:

bytes: A bytes object containing the serialized binary representation of the options.

Raises:

ValueError: If any option has already been accessed and cannot be serialized.

Notes

An Options instance can only be serialized if no options have been accessed. This is because serialization is based on the original string representations of the options. Once an option has been accessed and parsed, its string value may no longer accurately reflect its state, making serialization potentially inconsistent.

rapidsmpf.config.get_environment_variables(str key_regex='RAPIDSMPF_(.*)')#

Returns a dictionary of environment variables matching a given regular expression.

This function scans the current process’s environment variables and inserts those whose keys match the provided regular expression. The regular expression must contain exactly one capture group to extract the portion of the environment variable key to use as the dictionary key.

For example, to strip the RAPIDSMPF_ prefix, use r"RAPIDSMPF_(.*)" as the regex. The captured group will be used as the key in the output dictionary.

Example:

Environment variable: RAPIDSMPF_FOO=bar
key_regex: r”RAPIDSMPF_(.*)”
Resulting dictionary entry: { “FOO”: “bar” }

Parameters:

key_regex: A regular expression with a single capture group to match and extract the environment variable keys.

Returns:

A dictionary containing all matching environment variables, with keys as
extracted by the capture group.

Raises:

ValueError: If key_regex does not contain exactly one capture group.

Statistics#

class rapidsmpf.statistics.Formatter(*values)#

Default#

Bytes#

Duration#

HitRate#

MemoryThroughput#

class rapidsmpf.statistics.MemoryRecord( scoped: ScopedMemoryRecord, global_peak: int, num_calls: int, )#

Holds memory profiling statistics for a named scope.

Attributes:

scoped: Memory statistics collected while the scope was active, including number of allocations, peak bytes allocated, and total allocated bytes.
global_peak: The maximum global memory usage observed during the scope, including allocations from other threads or nested scopes.
num_calls: Number of times the profiling context with this name was entered.

class rapidsmpf.statistics.MemoryRecorder#

A context manager for recording memory allocation statistics within a code block.

This class is not intended to be used directly by end users. Instead, use Statistics.memory_profiling(), which creates and manages an instance of this class.

Parameters:

stats: The statistics object responsible for aggregating memory profiling data.
mr: The memory resource through which allocations are tracked.
name: The name of the profiling scope. Used as a key in the statistics record.

class rapidsmpf.statistics.Statistics(bool enable, *)#

Track statistics across RapidsMPF operations.

Parameters:

enable: Whether statistics tracking is enabled.

Attributes:

enabled: Checks if statistics is enabled.

Methods

`add_report_entry`(self, name, stat_names, ...)	Associate a predefined formatter with one or more stat names.
`add_stat`(self, name, double value)	Adds a value to a statistic.
`clear`(self)	Clears all statistics.
`copy`(self)	Creates a deep copy of this Statistics object.
`deserialize`(bytes buf)	Deserialize a binary buffer into a Statistics object.
`disable`(self)	Disable statistics tracking on this instance.
`disabled`(cls)	Returns a disabled (no-op) Statistics instance.
`enable`(self)	Enable statistics tracking on this instance.
`from_options`(cls, Options options)	Construct from configuration options.
`get_memory_records`(self)	Retrieves all memory profiling records stored by this instance.
`get_stat`(self, name)	Retrieves a statistic by name.
`list_stat_names`(self)	Returns a list of all statistic names.
`memory_profiling`(self, ...)	Create a scoped memory profiling context for a named code region.
`merge`(stats)	Merge a sequence of Statistics into a new one.
`report`(self, *, RmmResourceAdaptor mr=None, ...)	Generates a report of statistics in a formatted string.
`serialize`(self)	Serialize the stats and report entries to a binary buffer.
`to_dict`(self)	Return a plain dict snapshot of all statistics.
`write_json`(self, filepath)	Writes a JSON report of all collected statistics to a file.
`write_json_string`(self)	Returns a JSON representation of all collected statistics as a string.

add_report_entry( self, name, stat_names, Formatter formatter, )#

Associate a predefined formatter with one or more stat names.

Mirrors the C++ rapidsmpf::Statistics::add_report_entry. First-wins: if a report entry already exists under name, this call has no effect.

Parameters:

name: Report entry name. Becomes one line in report().
stat_names: Iterable of stat names this entry aggregates. The number of names must match the arity of formatter.
formatter: A Formatter selecting the predefined render function.

add_stat(self, name, double value)#

Adds a value to a statistic.

Parameters:

name: Name of the statistic.
value: Value to add.

clear(self) → None#

Clears all statistics.

Memory profiling records are not cleared.

copy(self)#

Creates a deep copy of this Statistics object.

Memory records are not copied.

Returns:

A new Statistics with the same stats and formatters.

static deserialize(bytes buf)#

Deserialize a binary buffer into a Statistics object.

Reconstructs a Statistics instance from a byte buffer produced by serialize(). The resulting object has no memory records and no associated memory-profiling resource.

Parameters:

buf: A buffer containing serialized statistics.

Returns:

A reconstructed Statistics instance.

Raises:

ValueError: If the input buffer is malformed or truncated.

disable(self)#: Disable statistics tracking on this instance.

classmethod disabled(cls)#

Returns a disabled (no-op) Statistics instance.

Useful when you need to pass a Statistics argument but do not want to collect any data.

Returns:

A Statistics instance with tracking disabled.

enable(self)#: Enable statistics tracking on this instance.

enabled#

Checks if statistics is enabled.

Operations on disabled statistics is no-ops.

Returns:

True if statistics is enabled, otherwise False.

classmethod from_options(cls, Options options)#

Construct from configuration options.

Parameters:

options: Configuration options.

Returns:

The constructed Statistics instance.

get_memory_records(self)#

Retrieves all memory profiling records stored by this instance.

Returns:

Dictionary mapping record names to memory usage data.

get_stat(self, name)#

Retrieves a statistic by name.

Parameters:

name: Name of the statistic to retrieve.

Returns:

A dict of the statistic.

Raises:

KeyError: If the statistic with the specified name does not exist.

list_stat_names(self)#: Returns a list of all statistic names.

memory_profiling(self, RmmResourceAdaptor mr, name)#

Create a scoped memory profiling context for a named code region.

Returns a context manager that tracks memory allocations and deallocations made through the associated memory resource while the context is active. The profiling data is aggregated under the provided name and made available via Statistics.get_memory_records().

The statistics include: - Total and peak memory allocated within the scope (scoped) - Global peak memory usage during the scope (global_peak) - Number of times the named scope was entered (num_calls)

Pass mr=None to get a no-op recorder.

Parameters:

mr: The memory resource through which allocations are tracked. Pass None to get a no-op recorder.
name: A unique identifier for the profiling scope. Used as a key when accessing profiling data via Statistics.get_memory_records().

Returns:

A context manager that collects memory profiling data.

Examples

>>> import rmm
>>> from rapidsmpf.memory.buffer_resource import BufferResource
>>> br = BufferResource(rmm.mr.CudaMemoryResource())
>>> mr = br.device_mr_adaptor()
>>> stats = Statistics(enable=True)
>>> with stats.memory_profiling(mr, "outer"):
...     b1 = rmm.DeviceBuffer(size=1024, mr=mr)
...     with stats.memory_profiling(mr, "inner"):
...         b2 = rmm.DeviceBuffer(size=1024, mr=mr)
>>> inner = stats.get_memory_records()["inner"]
>>> print(inner.scoped.peak())
1024
>>> outer = stats.get_memory_records()["outer"]
>>> print(outer.scoped.peak())
2048

static merge(stats)#

Merge a sequence of Statistics into a new one.

For each stat name present in any input, the result has the summed count, summed value, and the maximum of the maxes. Report entries with the same name must agree on formatter and stat-name list; otherwise the call raises ValueError. Memory records are not merged.

Parameters:

stats: A non-empty sequence of Statistics to merge.

Returns:

A new Statistics containing the merged data.

Raises:

ValueError: If stats is empty or two inputs have conflicting report entries.

report( self, *, RmmResourceAdaptor mr=None, PinnedMemoryResource pinned_mr=None, header=None, )#

Generates a report of statistics in a formatted string.

Operations on disabled statistics is no-ops.

Parameters:

mr: When provided, a memory profiling section is included in the report. When None, the memory profiling section shows “Disabled”.
pinned_mr: When provided, a pinned memory section is included in the report. Obtain the handle from rapidsmpf.memory.buffer_resource.BufferResource.pinned_mr.
header: Header line prepended to the report. When None, the C++ default is used.

Returns:

A string representing the formatted statistics report.

serialize(self) → bytes#

Serialize the stats and report entries to a binary buffer.

Memory records and the memory-profiling resource pointer are not included.

Returns:

A bytes object containing the serialized binary representation
of the Statistics.

to_dict(self)#

Return a plain dict snapshot of all statistics.

Each entry maps a stat name to a dict with count, value, and max keys, matching the shape returned by get_stat(). The snapshot is taken atomically and is detached thus mutating it does not affect the underlying Statistics.

Report-entry and formatter metadata is not included; use report() or write_json_string() for those.

Disabled statistics always return an empty dict.

Returns:

A dict mapping each stat name to its {"count", "value", "max"} dict.

write_json(self, filepath) → None#

Writes a JSON report of all collected statistics to a file.

Disabled statistics produce a JSON object with an empty statistics section.

Parameters:

filepath: Path to the output file. Created or overwritten.

Raises:

OSError: If the file cannot be opened or writing fails.

write_json_string(self) → str#

Returns a JSON representation of all collected statistics as a string.

Disabled statistics produce a JSON object with an empty statistics section.

Returns:

A JSON-formatted string.

RMM Resource Adaptor#

class rapidsmpf.rmm_resource_adaptor.RmmResourceAdaptor(*args, **kwargs)#

A RMM memory resource adaptor tailored to RapidsMPF.

Wraps a primary device memory resource and adds memory usage tracking (lifetime stats plus per-thread scoped records).

Construction

This class cannot be constructed directly. A usable RmmResourceAdaptor is always owned by a BufferResource (which installs the back-reference that makes the adaptor copyable). To obtain one, create a BufferResource from a device memory resource and call device_mr_adaptor():

>>> br = BufferResource(rmm.mr.CudaMemoryResource())
>>> mr = br.device_mr_adaptor()

The returned adaptor holds shared ownership of its owning BufferResource, so it (and any copy of it) keeps the BufferResource alive.

Attributes:

current_allocated: RmmResourceAdaptor.current_allocated: int

Methods

`allocate`(self, size_t nbytes, ...)	Allocate `nbytes` bytes of memory.
`deallocate`(self, uintptr_t ptr, ...)	Deallocate memory pointed to by `ptr` of size `nbytes`.
`get_main_record`(self)	Returns a copy of the main memory record.

current_allocated#

RmmResourceAdaptor.current_allocated: int

Get the total number of currently allocated bytes.

This includes both allocations on the primary and fallback memory resources.

Returns:

Total number of currently allocated bytes.

get_main_record(self)#

Returns a copy of the main memory record.

The main record tracks memory statistics for the lifetime of the resource.

Returns:

A copy of the current main memory record.