[MINOR] Reject invalid CUDA streams in buffer factories

cpp/include/rapidsmpf/memory/buffer.hpp

@@ -254,6 +254,7 @@ class Buffer {
      * @param new_stream The new CUDA stream.
      *
      * @throws std::logic_error If the buffer is locked.
+     * @throws std::invalid_argument If @p new_stream is `cuda::invalid_stream`.
      *
      * @code{.cpp}
      * // Example: merge buffers from different streams onto a single stream.

cpp/include/rapidsmpf/memory/buffer_resource.hpp

@@ -483,6 +483,7 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      *
      * @throws std::invalid_argument if the memory type does not match the reservation.
      * @throws rapidsmpf::reservation_error if `size` exceeds the size of the reservation.
+     * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`.
      */
     std::unique_ptr<Buffer> make_buffer(
         std::size_t size, cuda::stream_ref stream, MemoryReservation& reservation
@@ -497,6 +498,8 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      * @param stream CUDA stream to use for device allocations.
      * @param reservation The memory reservation to consume for the allocation.
      * @return A unique pointer to the allocated Buffer.
+     *
+     * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`.
      */
     std::unique_ptr<Buffer> make_buffer(
         cuda::stream_ref stream, MemoryReservation&& reservation
@@ -512,13 +515,10 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      * the Buffer is created with `MemoryType::PINNED_HOST`; otherwise it is
      * created with `MemoryType::DEVICE`.
      *
-     * If @p stream differs from the device buffer's current stream:
-     *   - @p stream is synchronized with the device buffer's current stream, and
-     *   - the device buffer's current stream is updated to @p stream.
-     *
-     * If @p data is empty and has no stream (e.g. a default-constructed
-     * `rmm::device_buffer`), no synchronization is performed and its stream is set to
-     * @p stream.
+     * @p stream is ordered after any work pending on the device buffer's current
+     * stream, and the device buffer's stream is updated to @p stream. If @p data is
+     * empty (e.g. a default-constructed `rmm::device_buffer`, which has no stream), it
+     * owns no memory, so no synchronization is performed.
      *
      * @param data Unique pointer to the device or pinned host buffer.
      * @param stream CUDA stream associated with the new Buffer. Use or synchronize with
@@ -528,7 +528,7 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      * @return Unique pointer to the resulting Buffer.
      *
      * @throws std::invalid_argument If @p spill_token is set and @p data is not
-     * host-accessible.
+     * host-accessible, or if @p stream is `cuda::invalid_stream`.
      */
     std::unique_ptr<Buffer> move(
         std::unique_ptr<rmm::device_buffer> data,

cpp/include/rapidsmpf/memory/host_buffer.hpp

@@ -42,6 +42,8 @@ class HostBuffer {
      * so the buffer shares ownership of the resource (e.g. bumps the refcount
      * when constructed from a shared-ownership resource); an implicit
      * conversion from `rmm::host_async_resource_ref` is also supported.
+     *
+     * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`.
      */
     HostBuffer(
         std::size_t size,
@@ -141,6 +143,8 @@ class HostBuffer {
      * or is otherwise properly synchronized before switching streams.
      *
      * @param stream The new CUDA stream.
+     *
+     * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`.
      */
     void set_stream(cuda::stream_ref stream);
 
@@ -155,6 +159,8 @@ class HostBuffer {
      * @param mr Host memory resource used to allocate the buffer.
      *
      * @return A new `HostBuffer` containing a copy of `data`.
+     *
+     * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`.
      */
     static HostBuffer from_uint8_vector(
         std::vector<std::uint8_t> const& data,
@@ -173,6 +179,8 @@ class HostBuffer {
      * @param stream CUDA stream to associate with this buffer.
      *
      * @return A new `HostBuffer` owning the vector's memory.
+     *
+     * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`.
      */
     static HostBuffer from_owned_vector(
         std::vector<std::uint8_t>&& data, cuda::stream_ref stream
@@ -185,6 +193,9 @@ class HostBuffer {
      * the device buffer contains host-accessible memory (e.g., pinned host memory
      * allocated via a managed or pinned memory resource).
      *
+     * @p stream is ordered after any work pending on the device buffer's stream, and
+     * the memory is deallocated on the `HostBuffer`'s stream.
+     *
      * @warning The caller is responsible for ensuring the device buffer's memory is
      * host-accessible. Using this with non-host-accessible device memory will result
      * in a std::invalid_argument exception.
@@ -194,7 +205,8 @@ class HostBuffer {
      *
      * @return A new `HostBuffer` owning the device buffer's memory.
      *
-     * @throws std::invalid_argument if @p pinned_host_buffer is null.
+     * @throws std::invalid_argument if @p pinned_host_buffer is null or @p stream is
+     * `cuda::invalid_stream`.
      * @throws std::logic_error if @p pinned_host_buffer is not of pinned memory host
      * memory type (see warning for details).
      *
@@ -214,6 +226,8 @@ class HostBuffer {
      * @param deallocate_fn Callable invoked with the current stream to release the
      * underlying memory. It captures all resources needed for deallocation (e.g.
      * memory resource, raw pointer, size).
+     *
+     * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`.
      */
     HostBuffer(
         std::span<std::byte> span,

cpp/src/memory/buffer.cpp

@@ -146,6 +146,11 @@ Buffer::DiskBufferT Buffer::release_disk_buffer() {
 
 void Buffer::rebind_stream(cuda::stream_ref new_stream) {
     throw_if_locked();
+    RAPIDSMPF_EXPECTS(
+        new_stream != cuda::invalid_stream,
+        "new_stream must be valid",
+        std::invalid_argument
+    );
     if (new_stream.get() == stream_.get()) {
         return;
     }

cpp/src/memory/buffer_resource.cpp

@@ -276,6 +276,9 @@ std::size_t BufferResource::release(MemoryReservation& reservation, std::size_t
 std::unique_ptr<Buffer> BufferResource::make_buffer(
     std::size_t size, cuda::stream_ref stream, MemoryReservation& reservation
 ) {
+    RAPIDSMPF_EXPECTS(
+        stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument
+    );
     auto const mem_type = reservation.mem_type_;
     // A disk buffer is created with mkstemp, which never touches the stream, so a
     // stream-ordered timing would measure stream backlog rather than the creation.
@@ -328,21 +331,20 @@ std::unique_ptr<Buffer> BufferResource::move(
     cuda::stream_ref stream,
     std::shared_ptr<SpillTrackToken> spill_token
 ) {
+    RAPIDSMPF_EXPECTS(
+        stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument
+    );
     RAPIDSMPF_EXPECTS(
         spill_token == nullptr || is_host_accessible(data->memory_resource()),
         "a spill token cannot be attached to device memory",
         std::invalid_argument
     );
-    cuda::stream_ref upstream = data->stream();
-    if (upstream.get() != stream.get()) {
-        // Default-constructed RMM buffers have no stream or work to join.
-        bool const empty_without_stream =
-            data->capacity() == 0 && upstream == cuda::invalid_stream;
-        if (!empty_without_stream) {
-            cuda_stream_join(stream, upstream);
-        }
-        data->set_stream(stream);
+    // An empty buffer owns no memory, so there is no work to join. This also covers
+    // default-constructed RMM buffers, which have no stream.
+    if (data->capacity() > 0) {
+        cuda_stream_join(stream, data->stream());
     }
+    data->set_stream(stream);
 
     if (is_host_accessible(data->memory_resource())) {
         auto pinned_host_buffer = std::make_unique<HostBuffer>(

cpp/src/memory/host_buffer.cpp

@@ -7,6 +7,7 @@
 
 #include <cuda/stream>
 
+#include <rapidsmpf/cuda_stream.hpp>
 #include <rapidsmpf/memory/cuda_memcpy_async.hpp>
 #include <rapidsmpf/memory/host_buffer.hpp>
 #include <rapidsmpf/memory/memory_type.hpp>
@@ -20,6 +21,9 @@ HostBuffer::HostBuffer(
     cuda::mr::any_resource<cuda::mr::host_accessible> mr
 )
     : stream_{stream} {
+    RAPIDSMPF_EXPECTS(
+        stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument
+    );
     if (size > 0) {
         auto* ptr = static_cast<std::byte*>(
             mr.allocate(stream_, size, alignof(::cuda::std::max_align_t))
@@ -36,7 +40,11 @@ HostBuffer::HostBuffer(
     cuda::stream_ref stream,
     std::function<void(cuda::stream_ref)> deallocate_fn
 )
-    : stream_{stream}, span_{span}, deallocate_fn_{std::move(deallocate_fn)} {}
+    : stream_{stream}, span_{span}, deallocate_fn_{std::move(deallocate_fn)} {
+    RAPIDSMPF_EXPECTS(
+        stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument
+    );
+}
 
 void HostBuffer::deallocate_async() noexcept {
     if (!span_.empty() && deallocate_fn_) {
@@ -90,6 +98,9 @@ std::byte const* HostBuffer::data() const noexcept {
 }
 
 void HostBuffer::set_stream(cuda::stream_ref new_stream) {
+    RAPIDSMPF_EXPECTS(
+        new_stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument
+    );
     stream_ = new_stream;
 }
 
@@ -136,6 +147,9 @@ HostBuffer HostBuffer::from_owned_vector(
 HostBuffer HostBuffer::from_rmm_device_buffer(
     std::unique_ptr<rmm::device_buffer> pinned_host_buffer, cuda::stream_ref stream
 ) {
+    RAPIDSMPF_EXPECTS(
+        stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument
+    );
     RAPIDSMPF_EXPECTS(
         pinned_host_buffer != nullptr,
         "pinned_host_buffer must not be null",
@@ -148,6 +162,14 @@ HostBuffer HostBuffer::from_rmm_device_buffer(
         std::logic_error
     );
 
+    // Order `stream` after pending work on the buffer's stream. An empty buffer owns
+    // no memory, so there is no work to join. This also covers default-constructed
+    // RMM buffers, which have no stream.
+    if (pinned_host_buffer->capacity() > 0) {
+        cuda_stream_join(stream, pinned_host_buffer->stream());
+    }
+    pinned_host_buffer->set_stream(stream);
+
     // Wrap in shared_ptr so the lambda is copyable (required by std::function).
     auto shared_db = std::make_shared<rmm::device_buffer>(std::move(*pinned_host_buffer));
     std::span<std::byte> span{
@@ -157,7 +179,10 @@ HostBuffer HostBuffer::from_rmm_device_buffer(
     return HostBuffer{
         std::move(span),
         stream,
-        [shared_db_ = std::move(shared_db)](cuda::stream_ref) mutable {
+        // Deallocate on the HostBuffer's current stream, which `set_stream()` may
+        // have changed since construction.
+        [shared_db_ = std::move(shared_db)](cuda::stream_ref s) mutable {
+            shared_db_->set_stream(s);
             shared_db_.reset();
         }
     };

cpp/tests/test_buffer.cpp

@@ -238,3 +238,18 @@ TEST_P(BufferRebindStreamTest, ThrowsWhenLocked) {
     EXPECT_NO_THROW(buffer->rebind_stream(stream2));
     EXPECT_EQ(buffer->stream().get(), stream2.get());
 }
+
+TEST(BufferRebindStream, ThrowsOnInvalidStream) {
+    StreamPool stream_pool{1};
+    auto stream = stream_pool.get_stream();
+    auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref());
+
+    auto [reserve, _] = br->reserve(MemoryType::DEVICE, 1_KiB, AllowOverbooking::YES);
+    auto buffer = br->make_buffer(1_KiB, stream, reserve);
+
+    EXPECT_THROW(
+        buffer->rebind_stream(cuda::stream_ref{cuda::invalid_stream}),
+        std::invalid_argument
+    );
+    EXPECT_EQ(buffer->stream().get(), stream.get());
+}

cpp/tests/test_buffer_resource.cpp

@@ -79,6 +79,20 @@ TEST(BufferResource, MoveDefaultConstructedDeviceBuffer) {
     EXPECT_NO_THROW(buffer->latest_write_event().host_wait());
 }
 
+TEST(BufferResource, InvalidStreamThrows) {
+    cuda::stream_ref const invalid{cuda::invalid_stream};
+    auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref());
+
+    auto [reservation, _] = br->reserve(MemoryType::DEVICE, 0, AllowOverbooking::YES);
+    EXPECT_THROW(
+        std::ignore = br->make_buffer(0, invalid, reservation), std::invalid_argument
+    );
+    EXPECT_THROW(
+        std::ignore = br->move(std::make_unique<rmm::device_buffer>(), invalid),
+        std::invalid_argument
+    );
+}
+
 TEST(BufferResource, ReservationOverbooking) {
     // Create a buffer resource that always reports 10 KiB of available device memory.
     auto br = BufferResource::create(

cpp/tests/test_host_buffer.cpp

@@ -15,12 +15,14 @@
 
 #include <cuda/stream>
 
+#include <rmm/cuda_stream.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 
 #include <rapidsmpf/config.hpp>
 #include <rapidsmpf/config_defaults.hpp>
 #include <rapidsmpf/cuda_stream.hpp>
+#include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/buffer_resource.hpp>
 #include <rapidsmpf/memory/pinned_memory_resource.hpp>
 #include <rapidsmpf/system_info.hpp>
@@ -269,6 +271,79 @@ TEST(PinnedResource, transient_mr) {
     EXPECT_NO_THROW(test_buffer(std::move(buffer), source_data));
 }
 
+TEST(PinnedResource, from_rmm_device_buffer_stream_ordering) {
+    if (!rapidsmpf::is_pinned_memory_resources_supported()) {
+        GTEST_SKIP() << "PinnedMemoryResource is not supported";
+    }
+    auto br = rapidsmpf::BufferResource::create(
+        rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{}
+    );
+    auto mr = br->try_pinned_mr();
+    rmm::cuda_stream stream_a;
+    rmm::cuda_stream stream_b;
+    rmm::cuda_stream stream_c;
+    constexpr std::size_t size = 1 << 20;
+
+    // Fill the pinned buffer asynchronously on stream A, without synchronizing.
+    auto pinned_host_buffer = std::make_unique<rmm::device_buffer>(size, stream_a, *mr);
+    RAPIDSMPF_CUDA_TRY(
+        cudaMemsetAsync(pinned_host_buffer->data(), 0xAB, size, stream_a.value())
+    );
+
+    // The new stream B must be ordered after the fill on A.
+    auto buffer = rapidsmpf::HostBuffer::from_rmm_device_buffer(
+        std::move(pinned_host_buffer), stream_b
+    );
+    EXPECT_EQ(buffer.stream().get(), stream_b.value());
+    EXPECT_EQ(buffer.copy_to_uint8_vector(), std::vector<std::uint8_t>(size, 0xAB));
+
+    buffer.set_stream(stream_c);  // Switch to stream C; the memory must be freed on C.
+    EXPECT_EQ(buffer.stream().get(), stream_c.value());
+    buffer.deallocate_async();
+    EXPECT_TRUE(buffer.empty());
+    stream_c.synchronize();
+}
+
+TEST(HostBuffer, invalid_stream_throws) {
+    cuda::stream_ref const valid{cudaStreamLegacy};
+    cuda::stream_ref const invalid{cuda::invalid_stream};
+    std::vector<std::uint8_t> const data(8);
+
+    EXPECT_THROW(
+        std::ignore = rapidsmpf::HostBuffer::from_owned_vector(
+            std::vector<std::uint8_t>(data), invalid
+        ),
+        std::invalid_argument
+    );
+    auto br = rapidsmpf::BufferResource::create(
+        rmm::mr::get_current_device_resource_ref(),
+        rapidsmpf::is_pinned_memory_resources_supported()
+            ? std::make_optional<rapidsmpf::PinnedPoolProperties>()
+            : std::nullopt
+    );
+    EXPECT_THROW(
+        std::ignore =
+            rapidsmpf::HostBuffer::from_uint8_vector(data, invalid, br->host_mr()),
+        std::invalid_argument
+    );
+    EXPECT_THROW(
+        rapidsmpf::HostBuffer(data.size(), invalid, br->host_mr()), std::invalid_argument
+    );
+
+    auto buffer =
+        rapidsmpf::HostBuffer::from_owned_vector(std::vector<std::uint8_t>(data), valid);
+    EXPECT_THROW(buffer.set_stream(invalid), std::invalid_argument);
+
+    if (auto pinned_mr = br->try_pinned_mr()) {
+        EXPECT_THROW(
+            std::ignore = rapidsmpf::HostBuffer::from_rmm_device_buffer(
+                std::make_unique<rmm::device_buffer>(8, valid, *pinned_mr), invalid
+            ),
+            std::invalid_argument
+        );
+    }
+}
+
 namespace {
 
 /// Discover the actual pool size the driver creates when a small max is requested.