check_reservation_size

Author: madsbk

python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx

@@ -14,6 +14,8 @@ from libcpp.vector cimport vector
 
 from rmm.pylibrmm import CudaStreamFlags
 
+from rapidsmpf.utils.memory import check_reservation_size
+
 from rmm.librmm.memory_resource cimport (any_resource, device_accessible,
                                          device_async_resource_ref,
                                          make_any_device_resource)
@@ -416,6 +418,7 @@ cdef class BufferResource:
             - On failure, the reservation's size equals zero (a zero-sized reservation
               never fails).
         """
+        check_reservation_size(size)
         cdef pair[unique_ptr[cpp_MemoryReservation], size_t] ret
         with nogil:
             ret = cpp_br_reserve(self._handle, mem_type, size, allow_overbooking)
@@ -451,6 +454,7 @@ cdef class BufferResource:
             If overbooking is disabled and the buffer resource cannot free enough
             device memory through spilling to satisfy the request.
         """
+        check_reservation_size(size)
         cdef unique_ptr[cpp_MemoryReservation] ret
         with nogil:
             ret = cpp_br_reserve_device_memory_and_spill(
@@ -485,6 +489,7 @@ cdef class BufferResource:
         RuntimeError
             If no memory type in ``mem_types`` could satisfy the reservation.
         """
+        check_reservation_size(size)
         cdef vector[MemoryType] cpp_mem_types
         for mt in mem_types:
             cpp_mem_types.push_back(<MemoryType?>mt)

python/rapidsmpf/rapidsmpf/streaming/core/memory_reserve_or_wait.pyx

@@ -22,6 +22,7 @@ from rapidsmpf.streaming.core.context cimport Context, cpp_Context
 import asyncio
 
 import rapidsmpf.utils.string
+from rapidsmpf.utils.memory import check_reservation_size
 
 # Sentinel indicating that net_memory_delta estimation has not yet been implemented.
 #
@@ -394,6 +395,7 @@ cdef class MemoryReserveOrWait:
         RuntimeError
             If shutdown occurs before the request can be processed.
         """
+        check_reservation_size(size)
         cdef shared_ptr[unique_ptr[cpp_MemoryReservation]] c_ret
         future = asyncio.get_running_loop().create_future()
         Py_INCREF(future)
@@ -444,6 +446,7 @@ cdef class MemoryReserveOrWait:
         RuntimeError
             If shutdown occurs before the request can be processed.
         """
+        check_reservation_size(size)
         cdef shared_ptr[pair[unique_ptr[cpp_MemoryReservation], size_t]] c_ret
         future = asyncio.get_running_loop().create_future()
         Py_INCREF(future)
@@ -497,6 +500,7 @@ cdef class MemoryReserveOrWait:
         --------
         reserve_or_wait
         """
+        check_reservation_size(size)
         cdef shared_ptr[unique_ptr[cpp_MemoryReservation]] c_ret
         future = asyncio.get_running_loop().create_future()
         Py_INCREF(future)
@@ -602,6 +606,8 @@ async def reserve_memory(
     ...     allow_overbooking=False,
     ... )
     """
+    check_reservation_size(size)
+
     if allow_overbooking is None:
         allow_overbooking = ctx.options().get(
             "allow_overbooking_by_default",

python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py

@@ -13,6 +13,7 @@
 from rapidsmpf.memory.buffer_resource import BufferResource, OwningDeviceMemoryResource
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.statistics import Statistics
+from rapidsmpf.utils.memory import _MAX_RESERVATION_BYTES
 
 
 def KiB(x: int) -> int:
@@ -76,6 +77,18 @@ def test_memory_reservation(mem_type: MemoryType) -> None:
         br.release(res1, KiB(10))
 
 
+@pytest.mark.parametrize("mem_type", [MemoryType.DEVICE, MemoryType.HOST])
+def test_reserve_rejects_size_underflow(mem_type: MemoryType) -> None:
+    mr = rmm.mr.CudaMemoryResource()
+    br = BufferResource(mr, memory_limits={mem_type: KiB(100)})
+
+    with pytest.raises(ValueError, match="underflow"):
+        br.reserve(mem_type, _MAX_RESERVATION_BYTES + 1, allow_overbooking=True)
+
+    res, _ = br.reserve(mem_type, _MAX_RESERVATION_BYTES, allow_overbooking=True)
+    assert res.size == _MAX_RESERVATION_BYTES
+
+
 def test_stream_pool() -> None:
     """Test that stream_pool parameter can be configured."""
     mr = rmm.mr.CudaMemoryResource()

python/rapidsmpf/rapidsmpf/utils/memory.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Memory utilities."""
+
+from __future__ import annotations
+
+# Largest reservation size addressable as a signed 64-bit byte count (INT64_MAX).
+_MAX_RESERVATION_BYTES = (1 << 63) - 1
+
+
+def check_reservation_size(size: int) -> None:
+    """
+    Reject a reservation size that cannot represent a real allocation.
+
+    A value above ``INT64_MAX`` almost always means the caller computed ``size``
+    using an unsigned subtraction that underflowed and wrapped around outside of
+    Python. Performing this check in Python, before the value is coerced into a
+    ``size_t``, lets us raise a traceback that points at the offending caller
+    rather than silently passing the wrapped value down to C++ where it only
+    surfaces later as an opaque "value out of range" cast error.
+
+    Parameters
+    ----------
+    size
+        Requested reservation size in bytes.
+
+    Raises
+    ------
+    ValueError
+        If ``size`` exceeds the addressable range.
+    """
+    if size > _MAX_RESERVATION_BYTES:
+        raise ValueError(
+            f"reservation size ({size}) exceeds the addressable range; this almost "
+            f"certainly indicates an unsigned (size_t) underflow in the caller's size "
+            "computation"
+        )