Use the new make_table_chunks_available_or_wait API from RapidsMPF (rapidsai#21291)

Replacing `opaque_reservation`.

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)
  - Richard (Rick) Zamora (https://github.com/rjzamora)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)
  - Richard (Rick) Zamora (https://github.com/rjzamora)

URL: rapidsai#21291

Author: madsbk

python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py

@@ -9,6 +9,10 @@
 import math
 from typing import TYPE_CHECKING, Any
 
+from rapidsmpf.memory.memory_reservation import opaque_memory_usage
+from rapidsmpf.streaming.core.memory_reserve_or_wait import (
+    reserve_memory,
+)
 from rapidsmpf.streaming.core.message import Message
 from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
 from rapidsmpf.streaming.cudf.table_chunk import TableChunk
@@ -37,7 +41,6 @@
 )
 from cudf_polars.experimental.rapidsmpf.utils import (
     ChannelManager,
-    opaque_reservation,
     send_metadata,
 )
 
@@ -345,25 +348,29 @@ async def read_chunk(
     tracer
         The actor tracer for collecting runtime statistics.
     """
-    with opaque_reservation(context, estimated_chunk_bytes):
+    with opaque_memory_usage(
+        await reserve_memory(
+            context, size=estimated_chunk_bytes, net_memory_delta=estimated_chunk_bytes
+        )
+    ):
         df = await asyncio.to_thread(
             scan.do_evaluate,
             *scan._non_child_args,
             context=ir_context,
         )
-        if tracer is not None:
-            tracer.add_chunk(table=df.table)
-        await ch_out.send(
-            context,
-            Message(
-                seq_num,
-                TableChunk.from_pylibcudf_table(
-                    df.table,
-                    df.stream,
-                    exclusive_view=True,
-                ),
+    if tracer is not None:
+        tracer.add_chunk(table=df.table)
+    await ch_out.send(
+        context,
+        Message(
+            seq_num,
+            TableChunk.from_pylibcudf_table(
+                df.table,
+                df.stream,
+                exclusive_view=True,
             ),
-        )
+        ),
+    )
 
 
 @define_py_node()

python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py

@@ -8,6 +8,11 @@
 from typing import TYPE_CHECKING, Any, Literal
 
 from rapidsmpf.memory.buffer import MemoryType
+from rapidsmpf.memory.memory_reservation import opaque_memory_usage
+from rapidsmpf.streaming.core.memory_reserve_or_wait import (
+    missing_net_memory_delta,
+    reserve_memory,
+)
 from rapidsmpf.streaming.core.message import Message
 from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
 from rapidsmpf.streaming.cudf.table_chunk import TableChunk
@@ -27,7 +32,6 @@
     ChannelManager,
     chunk_to_frame,
     empty_table_chunk,
-    opaque_reservation,
     process_children,
     recv_metadata,
     send_metadata,
@@ -191,28 +195,31 @@ async def broadcast_join_node(
                     return
             else:
                 large_chunk_processed = True
-                large_chunk = TableChunk.from_message(msg).make_available_and_spill(
-                    context.br(), allow_overbooking=True
+                large_chunk = await TableChunk.from_message(msg).make_available_or_wait(
+                    context,
+                    net_memory_delta=missing_net_memory_delta,
                 )
                 seq_num = msg.sequence_number
-                del msg
 
             large_df = DataFrame.from_table(
                 large_chunk.table_view(),
                 list(large_child.schema.keys()),
                 list(large_child.schema.values()),
                 large_chunk.stream,
             )
+            large_chunk_size = large_chunk.data_alloc_size(MemoryType.DEVICE)
+            del large_chunk  # `large_df` keeps `large_chunk` alive.
 
             # Lazily create empty small table if small_dfs is empty
             if not small_dfs:
                 stream = ir_context.get_cuda_stream()
                 empty_small_chunk = empty_table_chunk(small_child, context, stream)
                 small_dfs = [chunk_to_frame(empty_small_chunk, small_child)]
 
-            large_chunk_size = large_chunk.data_alloc_size(MemoryType.DEVICE)
             input_bytes = large_chunk_size + small_size
-            with opaque_reservation(context, input_bytes):
+            with opaque_memory_usage(
+                await reserve_memory(context, size=input_bytes, net_memory_delta=0)
+            ):
                 df = _concat(
                     *[
                         await asyncio.to_thread(
@@ -229,18 +236,18 @@ async def broadcast_join_node(
                     ],
                     context=ir_context,
                 )
+                del large_df
 
-                # Send output chunk
-                await ch_out.send(
-                    context,
-                    Message(
-                        seq_num,
-                        TableChunk.from_pylibcudf_table(
-                            df.table, df.stream, exclusive_view=True
-                        ),
+            # Send output chunk
+            await ch_out.send(
+                context,
+                Message(
+                    seq_num,
+                    TableChunk.from_pylibcudf_table(
+                        df.table, df.stream, exclusive_view=True
                     ),
-                )
-                del df, large_df, large_chunk
+                ),
+            )
 
         del small_dfs, small_chunks
         await ch_out.drain(context)

python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py

@@ -8,11 +8,15 @@
 from typing import TYPE_CHECKING, Any, cast
 
 from rapidsmpf.memory.buffer import MemoryType
+from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.streaming.core.message import Message
 from rapidsmpf.streaming.core.node import define_py_node
 from rapidsmpf.streaming.core.spillable_messages import SpillableMessages
 from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
+from rapidsmpf.streaming.cudf.table_chunk import (
+    TableChunk,
+    make_table_chunks_available_or_wait,
+)
 
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import IR, Cache, Empty, Filter, Projection
@@ -23,7 +27,6 @@
     ChannelManager,
     empty_table_chunk,
     make_spill_function,
-    opaque_reservation,
     process_children,
     recv_metadata,
     remap_partitioning,
@@ -102,18 +105,21 @@ async def default_node_single(
                 else:
                     # Make sure we have an empty chunk in case do_evaluate
                     # always produces rows (e.g. aggregation)
-                    stream = ir_context.get_cuda_stream()
-                    chunk = empty_table_chunk(ir.children[0], context, stream)
+                    chunk = empty_table_chunk(
+                        ir.children[0], context, ir_context.get_cuda_stream()
+                    )
             else:
                 received_any = True
-                chunk = TableChunk.from_message(msg).make_available_and_spill(
-                    context.br(), allow_overbooking=True
-                )
+                chunk = TableChunk.from_message(msg)
                 seq_num = msg.sequence_number
-            del msg
 
-            input_bytes = chunk.data_alloc_size(MemoryType.DEVICE)
-            with opaque_reservation(context, input_bytes):
+            chunk, extra = await make_table_chunks_available_or_wait(
+                context,
+                chunk,
+                reserve_extra=chunk.data_alloc_size(),
+                net_memory_delta=0,
+            )
+            with opaque_memory_usage(extra):
                 df = await asyncio.to_thread(
                     ir.do_evaluate,
                     *ir._non_child_args,
@@ -125,18 +131,18 @@ async def default_node_single(
                     ),
                     context=ir_context,
                 )
-                if tracer is not None:
-                    tracer.add_chunk(table=df.table)
-                await ch_out.send(
-                    context,
-                    Message(
-                        seq_num,
-                        TableChunk.from_pylibcudf_table(
-                            df.table, chunk.stream, exclusive_view=True
-                        ),
+            if tracer is not None:
+                tracer.add_chunk(table=df.table)
+            await ch_out.send(
+                context,
+                Message(
+                    seq_num,
+                    TableChunk.from_pylibcudf_table(
+                        df.table, chunk.stream, exclusive_view=True
                     ),
-                )
-                del df, chunk
+                ),
+            )
+            del df, chunk
 
         await ch_out.drain(context)
 
@@ -234,10 +240,15 @@ async def default_node_multi(
                     ready_chunks[ch_idx] = empty_table_chunk(child, context, stream)
 
             # Ensure all table chunks are unspilled and available.
-            ready_chunks = [
-                chunk.make_available_and_spill(context.br(), allow_overbooking=True)
-                for chunk in cast(list[TableChunk], ready_chunks)
-            ]
+            ready_chunks, extra = await make_table_chunks_available_or_wait(
+                context,
+                ready_chunks,
+                reserve_extra=sum(
+                    chunk.data_alloc_size()
+                    for chunk in cast(list[TableChunk], ready_chunks)
+                ),
+                net_memory_delta=0,
+            )
             dfs = [
                 DataFrame.from_table(
                     chunk.table_view(),  # type: ignore[union-attr]
@@ -247,33 +258,29 @@ async def default_node_multi(
                 )
                 for chunk, child in zip(ready_chunks, ir.children, strict=True)
             ]
-
-            input_bytes = sum(
-                chunk.data_alloc_size(MemoryType.DEVICE)
-                for chunk in cast(list[TableChunk], ready_chunks)
-            )
-            with opaque_reservation(context, input_bytes):
+            with opaque_memory_usage(extra):
                 df = await asyncio.to_thread(
                     ir.do_evaluate,
                     *ir._non_child_args,
                     *dfs,
                     context=ir_context,
                 )
-                if tracer is not None:
-                    tracer.add_chunk(table=df.table)
-                await ch_out.send(
-                    context,
-                    Message(
-                        seq_num,
-                        TableChunk.from_pylibcudf_table(
-                            df.table,
-                            df.stream,
-                            exclusive_view=True,
-                        ),
+                del dfs
+            if tracer is not None:
+                tracer.add_chunk(table=df.table)
+            await ch_out.send(
+                context,
+                Message(
+                    seq_num,
+                    TableChunk.from_pylibcudf_table(
+                        df.table,
+                        df.stream,
+                        exclusive_view=True,
                     ),
-                )
-                seq_num += 1
-                del df, dfs
+                ),
+            )
+            seq_num += 1
+            del df
 
         # Drain the output channel
         del ready_chunks

python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py

@@ -7,11 +7,14 @@
 import math
 from typing import TYPE_CHECKING, Any
 
-from rapidsmpf.memory.buffer import MemoryType
+from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.streaming.core.message import Message
 from rapidsmpf.streaming.core.node import define_py_node
 from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
+from rapidsmpf.streaming.cudf.table_chunk import (
+    TableChunk,
+    make_table_chunks_available_or_wait,
+)
 
 from cudf_polars.containers import DataFrame
 from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager
@@ -20,7 +23,6 @@
 from cudf_polars.experimental.rapidsmpf.utils import (
     ChannelManager,
     empty_table_chunk,
-    opaque_reservation,
     recv_metadata,
     send_metadata,
 )
@@ -183,18 +185,16 @@ async def concatenate_node(
                     if msg is None:
                         done_receiving = True
                         break
-                    chunks.append(
-                        TableChunk.from_message(msg).make_available_and_spill(
-                            context.br(), allow_overbooking=True
-                        )
-                    )
-                    del msg
+                    chunks.append(TableChunk.from_message(msg))
 
                 if chunks:
-                    input_bytes = sum(
-                        chunk.data_alloc_size(MemoryType.DEVICE) for chunk in chunks
+                    chunks, extra = await make_table_chunks_available_or_wait(
+                        context,
+                        chunks,
+                        reserve_extra=sum(chunk.data_alloc_size() for chunk in chunks),
+                        net_memory_delta=0,
                     )
-                    with opaque_reservation(context, input_bytes):
+                    with opaque_memory_usage(extra):
                         df = _concat(
                             *(
                                 DataFrame.from_table(
@@ -207,19 +207,20 @@ async def concatenate_node(
                             ),
                             context=ir_context,
                         )
-                        if tracer is not None:
-                            tracer.add_chunk(table=df.table)
-                        await ch_out.send(
-                            context,
-                            Message(
-                                seq_num,
-                                TableChunk.from_pylibcudf_table(
-                                    df.table, df.stream, exclusive_view=True
-                                ),
+                        del chunks
+                    if tracer is not None:
+                        tracer.add_chunk(table=df.table)
+                    await ch_out.send(
+                        context,
+                        Message(
+                            seq_num,
+                            TableChunk.from_pylibcudf_table(
+                                df.table, df.stream, exclusive_view=True
                             ),
-                        )
-                        seq_num += 1
-                        del df, chunks
+                        ),
+                    )
+                    seq_num += 1
+                    del df
 
                 # Break if we reached end of stream
                 if done_receiving: