Add shuffler disk residency for received chunks

cmake/thirdparty/get_kvikio.cmake

@@ -0,0 +1,12 @@
+# =============================================================================
+# cmake-format: off
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# cmake-format: on
+# =============================================================================
+
+rapids_find_package(
+  kvikio REQUIRED
+  BUILD_EXPORT_SET rapidsmpf-exports
+  INSTALL_EXPORT_SET rapidsmpf-exports
+)

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -26,6 +26,7 @@ dependencies:
 - gcc_linux-aarch64=14.*
 - gdb
 - ipython
+- libkvikio==26.10.*,>=0.0.0a0
 - libnuma
 - libpmix-devel >=5.0,<6.0
 - librmm==26.10.*,>=0.0.0a0

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -26,6 +26,7 @@ dependencies:
 - gcc_linux-64=14.*
 - gdb
 - ipython
+- libkvikio==26.10.*,>=0.0.0a0
 - libnuma
 - libpmix-devel >=5.0,<6.0
 - librmm==26.10.*,>=0.0.0a0

conda/environments/all_cuda-133_arch-aarch64.yaml

@@ -26,6 +26,7 @@ dependencies:
 - gcc_linux-aarch64=14.*
 - gdb
 - ipython
+- libkvikio==26.10.*,>=0.0.0a0
 - libnuma
 - libpmix-devel >=5.0,<6.0
 - librmm==26.10.*,>=0.0.0a0

conda/environments/all_cuda-133_arch-x86_64.yaml

@@ -26,6 +26,7 @@ dependencies:
 - gcc_linux-64=14.*
 - gdb
 - ipython
+- libkvikio==26.10.*,>=0.0.0a0
 - libnuma
 - libpmix-devel >=5.0,<6.0
 - librmm==26.10.*,>=0.0.0a0

conda/recipes/librapidsmpf/recipe.yaml

@@ -77,6 +77,7 @@ cache:
       - cuda-version =${{ cuda_version }}
       - cuda-cupti-dev
       - cuda-nvml-dev
+      - libkvikio =${{ minor_version }}
       - libnuma
       - libpmix-devel >=5.0,<6.0
       - librmm =${{ minor_version }}
@@ -111,13 +112,15 @@ outputs:
         - cuda-version =${{ cuda_version }}
         - cuda-cupti-dev
         - cuda-nvml-dev
+        - libkvikio =${{ minor_version }}
         - libpmix-devel >=5.0,<6.0
         - libucxx ${{ ucxx_version }}
         - openmpi >=5.0
         - libnuma
       run:
         - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
         - cuda-cupti
+        - libkvikio =${{ minor_version }}
         - librmm =${{ minor_version }}
         - libpmix >=5.0,<6.0
         - libucxx ${{ ucxx_version }}
@@ -130,6 +133,7 @@ outputs:
           - ${{ stdlib("c") }}
         by_name:
           - cuda-cupti
+          - libkvikio
           - librmm
           - openmpi
     tests:

cpp/CMakeLists.txt

@@ -132,6 +132,7 @@ include(../cmake/thirdparty/get_nvtx.cmake)
 include(../cmake/thirdparty/get_rmm.cmake)
 include(../cmake/thirdparty/get_cuco.cmake)
 include(../cmake/thirdparty/get_cucascade.cmake)
+include(../cmake/thirdparty/get_kvikio.cmake)
 
 if(RAPIDSMPF_HAVE_UCXX)
   rapids_find_package(
@@ -202,6 +203,8 @@ add_library(
   src/memory/scoped_memory_record.cpp
   src/memory/spill.cpp
   src/memory/spill_manager.cpp
+  src/disk/disk_resource.cpp
+  src/disk/disk_buffer.cpp
   src/pausable_thread_loop.cpp
   src/progress_thread.cpp
   src/rmm_resource_adaptor.cpp
@@ -306,8 +309,8 @@ endif()
 
 target_link_libraries(
   rapidsmpf
-  PUBLIC rmm::rmm CCCL::CCCL $<TARGET_NAME_IF_EXISTS:ucxx::ucxx> $<TARGET_NAME_IF_EXISTS:libcoro>
-         CUDA::cudart_static
+  PUBLIC rmm::rmm CCCL::CCCL kvikio::kvikio $<TARGET_NAME_IF_EXISTS:ucxx::ucxx>
+         $<TARGET_NAME_IF_EXISTS:libcoro> CUDA::cudart_static
   PRIVATE cuco::cuco
           cuCascade::cucascade_topology_discovery
           $<$<BOOL:${RAPIDSMPF_HAVE_NUMA}>:numa>

cpp/include/rapidsmpf/communicator/metadata_payload_exchange/tag.hpp

@@ -30,7 +30,9 @@ class TagMetadataPayloadExchange : public MetadataPayloadExchange {
      *
      * @param comm The communicator to use for operations.
      * @param op_id The operation ID for tagging messages.
-     * @param allocate_buffer_fn Function to allocate buffers for incoming data.
+     * @param allocate_buffer_fn Function to allocate buffers for incoming data. May
+     * return `nullptr` when allocation is temporarily unavailable; allocation is
+     * retried by subsequent calls to `progress()`.
      * @param statistics The statistics to use for tracking communication operations.
      */
     TagMetadataPayloadExchange(
@@ -124,6 +126,8 @@ class TagMetadataPayloadExchange : public MetadataPayloadExchange {
         peer_expected_;  ///< Expected application messages per peer (0 = unknown).
     std::vector<bool>
         peer_terminated_;  ///< Whether we received the termination marker from each peer.
+    std::vector<bool>
+        peer_allocation_deferred_;  ///< Whether payload allocation is deferred per peer.
 
     /// Sentinel message_id value used to identify protocol-level termination markers.
     static constexpr std::uint64_t termination_sentinel_ = UINT64_MAX;
@@ -158,8 +162,7 @@ class TagMetadataPayloadExchange : public MetadataPayloadExchange {
      *
      * @return A vector of completed metadata-only messages.
      *
-     * @throw std::runtime_error if an in-transit message or future is not found, or
-     * if a data buffer is not available.
+     * @throw std::runtime_error if an in-transit message or future is not found.
      */
     std::vector<std::unique_ptr<Message>> setup_data_receives();
 

cpp/include/rapidsmpf/config_defaults.hpp

@@ -30,13 +30,15 @@ inline const std::unordered_map<std::string, std::string> DEFAULTS{
     {"pinned_initial_pool_size", "0%"},
     {"pinned_max_pool_size", "80%"},
     {"spill_device_limit", "80%"},
+    {"spill_host_limit", "disabled"},
     {"periodic_spill_check", "1ms"},
     {"num_streams", "16"},
     {"num_streaming_threads", "1"},
     {"memory_reserve_timeout", "100ms"},
     {"allow_overbooking_by_default", "true"},
     {"log", "WARN"},
     {"ucxx_progress_mode", "thread-blocking"},
+    {"disk_spill_dir", ""},
 };
 
 }  // namespace rapidsmpf::config

cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp

@@ -36,12 +36,14 @@ namespace rapidsmpf::detail {
  * `cuda::mr::resource` concept and is held by `RmmResourceAdaptor` via
  * `cuda::mr::shared_resource` for reference-counted ownership.
  *
- * @tparam PrimaryMR The type of the primary memory resource. Use a concrete
- * resource type (e.g. `cuda::pinned_memory_pool`) to store the resource
- * directly inside the shared control block, avoiding an extra heap allocation.
+ * @tparam PrimaryMR The type of the primary memory resource. Its properties are
+ * forwarded via `cuda::forward_property`. Use a concrete resource type (e.g.
+ * `cuda::pinned_memory_pool`) to store the resource directly inside the shared
+ * control block, avoiding an extra heap allocation.
  */
-template <cuda::mr::resource_with<cuda::mr::device_accessible> PrimaryMR>
-class RmmResourceAdaptorImpl {
+template <cuda::mr::resource PrimaryMR>
+class RmmResourceAdaptorImpl
+    : public cuda::forward_property<RmmResourceAdaptorImpl<PrimaryMR>, PrimaryMR> {
   public:
     /**
      * @brief Construct with a primary memory resource.
@@ -91,12 +93,21 @@ class RmmResourceAdaptorImpl {
 
     /**
      * @brief Returns a reference to the primary upstream resource.
+     *
+     * Named `upstream_resource` so `cuda::forward_property` can forward stateful
+     * properties and `cuda::mr::dynamic_accessibility_property`.
+     *
      * @return Reference to the primary resource.
      */
-    [[nodiscard]] PrimaryMR const& get_upstream_resource() const noexcept {
+    [[nodiscard]] PrimaryMR const& upstream_resource() const noexcept {
         return primary_mr_;
     }
 
+    /// @copydoc upstream_resource
+    [[nodiscard]] PrimaryMR const& get_upstream_resource() const noexcept {
+        return upstream_resource();
+    }
+
     /// @copydoc RmmResourceAdaptor::get_main_record
     [[nodiscard]] ScopedMemoryRecord get_main_record() const {
         std::lock_guard<std::mutex> lock(mutex_);
@@ -224,11 +235,6 @@ class RmmResourceAdaptorImpl {
         sync_stream_.synchronize_no_throw();
     }
 
-    /// @brief Tag this resource as device-accessible for the CCCL concept.
-    friend void get_property(
-        RmmResourceAdaptorImpl const&, cuda::mr::device_accessible
-    ) noexcept {}
-
   private:
     mutable std::mutex mutex_;
     PrimaryMR primary_mr_;

cpp/include/rapidsmpf/disk/disk_buffer.hpp

@@ -0,0 +1,127 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <cstddef>
+#include <filesystem>
+#include <memory>
+
+#include <cuda/stream>
+
+#include <rapidsmpf/disk/disk_resource.hpp>
+#include <rapidsmpf/memory/buffer.hpp>
+#include <rapidsmpf/memory/memory_reservation.hpp>
+
+namespace rapidsmpf {
+
+class BufferResource;
+
+namespace disk {
+
+/**
+ * @brief File-backed handle to a byte buffer.
+ *
+ * `DiskBuffer` is intentionally separate from `Buffer` and outside the
+ * `MemoryType` taxonomy. It owns a backing file and deletes it on
+ * `deallocate()` or destruction.
+ *
+ * `from_buffer` and `restore` are synchronous: they block until the disk
+ * transfer completes.
+ */
+class DiskBuffer {
+  public:
+    /**
+     * @brief Move constructor.
+     *
+     * Transfers ownership of the backing file. The moved-from object is empty
+     * (`path()` is empty and `size()` is zero).
+     *
+     * @param other Buffer to move from.
+     */
+    DiskBuffer(DiskBuffer&& other) noexcept;
+    DiskBuffer& operator=(DiskBuffer&& other) = delete;  ///< Not move-assignable.
+
+    ~DiskBuffer();
+
+    DiskBuffer(DiskBuffer const&) = delete;  ///< Not copyable.
+    DiskBuffer& operator=(DiskBuffer const&) = delete;  ///< Not copy-assignable.
+
+    /**
+     * @brief Payload size in bytes stored in the backing file.
+     *
+     * @return Size in bytes.
+     */
+    [[nodiscard]] constexpr std::size_t size() const noexcept {
+        return size_;
+    }
+
+    /**
+     * @brief Path to the backing file.
+     *
+     * @return Filesystem path to the backing file.
+     */
+    [[nodiscard]] std::filesystem::path const& path() const noexcept {
+        return path_;
+    }
+
+    /**
+     * @brief Delete the backing file, if any.
+     *
+     * After deallocation the buffer is empty (`path()` is empty and `size()`
+     * is zero). The associated disk resource is retained until destruction,
+     * matching the resource-lifetime behavior of `rmm::device_buffer`. Safe to
+     * call multiple times.
+     */
+    void deallocate() noexcept;
+
+    /**
+     * @brief Write the contents of @p source to a new file on disk.
+     *
+     * Takes ownership of @p source, creates a unique file under @p br's
+     * configured disk directory, and blocks until the write completes. The
+     * in-memory buffer is released after the write finishes. @p source must
+     * already have no pending stream-ordered writes
+     * (`is_latest_write_done()`).
+     *
+     * @param source Buffer whose bytes are written to disk.
+     * @param br Buffer resource supplying disk I/O and directory configuration.
+     * @return A disk-resident handle to the written bytes.
+     */
+    [[nodiscard]] static std::unique_ptr<DiskBuffer> from_buffer(
+        std::unique_ptr<Buffer> source, BufferResource& br
+    );
+
+    /**
+     * @brief Restore bytes from a disk handle into a newly allocated `Buffer`.
+     *
+     * Takes ownership of @p source, allocates from @p reservation on @p stream,
+     * reads the backing file with blocking disk I/O, and records the fill as a
+     * stream-ordered write on the destination buffer. The backing file is
+     * deleted when @p source is destroyed.
+     *
+     * @param source Disk-resident handle whose bytes are restored.
+     * @param reservation Memory reservation covering at least `source->size()`
+     * bytes.
+     * @param stream CUDA stream associated with the destination buffer.
+     * @return In-memory buffer containing the file contents.
+     */
+    [[nodiscard]] static std::unique_ptr<Buffer> restore(
+        std::unique_ptr<DiskBuffer> source,
+        MemoryReservation& reservation,
+        cuda::stream_ref stream
+    );
+
+  private:
+    DiskBuffer(
+        std::shared_ptr<DiskResource> disk, std::filesystem::path path, std::size_t size
+    );
+
+    std::shared_ptr<DiskResource> disk_;
+    std::filesystem::path path_;
+    std::size_t size_{};
+};
+
+}  // namespace disk
+}  // namespace rapidsmpf