Use the same channel for TableChunk and Metadata messages (rapidsai#21182)

Now that rapidsai/rapidsmpf#811 was merged in RapidsMPF, we can drop the confusing `ChannelPair` design in cudf-polars.

**NOTE**: After this and rapidsai/rapidsmpf#819 are both in, we can also drop the cudf-polasr `Metadata` definition in favor of the "standard" `ChannelMetadata` defined in RapidsMPF.

Authors:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

Approvers:
  - Matthew Murray (https://github.com/Matt711)
  - Tom Augspurger (https://github.com/TomAugspurger)

URL: rapidsai#21182

Author: rjzamora

python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py

@@ -24,18 +24,20 @@
     ChannelManager,
     HashPartitioned,
     Metadata,
+    recv_metadata,
+    send_metadata,
 )
 from cudf_polars.experimental.shuffle import Shuffle
 
 if TYPE_CHECKING:
+    from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
 
     import pylibcudf as plc
     from rmm.pylibrmm.stream import Stream
 
     from cudf_polars.dsl.ir import IR, IRExecutionContext
     from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator
-    from cudf_polars.experimental.rapidsmpf.utils import ChannelPair
 
 
 class ShuffleManager:
@@ -125,8 +127,8 @@ async def shuffle_node(
     context: Context,
     ir: Shuffle,
     ir_context: IRExecutionContext,
-    ch_in: ChannelPair,
-    ch_out: ChannelPair,
+    ch_in: Channel[TableChunk],
+    ch_out: Channel[TableChunk],
     columns_to_hash: tuple[int, ...],
     num_partitions: int,
     collective_id: int,
@@ -147,21 +149,19 @@ async def shuffle_node(
     ir_context
         The execution context for the IR node.
     ch_in
-        Input ChannelPair with metadata and data channels.
+        Input Channel[TableChunk] with metadata and data channels.
     ch_out
-        Output ChannelPair with metadata and data channels.
+        Output Channel[TableChunk] with metadata and data channels.
     columns_to_hash
         Tuple of column indices to use for hashing.
     num_partitions
         Number of partitions to shuffle into.
     collective_id
         The collective ID.
     """
-    async with shutdown_on_error(
-        context, ch_in.metadata, ch_in.data, ch_out.metadata, ch_out.data
-    ):
+    async with shutdown_on_error(context, ch_in, ch_out):
         # Receive and send updated metadata.
-        _ = await ch_in.recv_metadata(context)
+        _ = await recv_metadata(ch_in, context)
         column_names = list(ir.schema.keys())
         partitioned_on = tuple(column_names[i] for i in columns_to_hash)
         output_metadata = Metadata(
@@ -173,15 +173,15 @@ async def shuffle_node(
                 count=num_partitions,
             ),
         )
-        await ch_out.send_metadata(context, output_metadata)
+        await send_metadata(ch_out, context, output_metadata)
 
         # Create ShuffleManager instance
         shuffle = ShuffleManager(
             context, num_partitions, columns_to_hash, collective_id
         )
 
         # Process input chunks
-        while (msg := await ch_in.data.recv(context)) is not None:
+        while (msg := await ch_in.recv(context)) is not None:
             # Extract TableChunk from message and insert into shuffler
             shuffle.insert_chunk(
                 TableChunk.from_message(msg).make_available_and_spill(
@@ -202,7 +202,7 @@ async def shuffle_node(
             context.comm().nranks,
         ):
             # Extract and send the output chunk
-            await ch_out.data.send(
+            await ch_out.send(
                 context,
                 Message(
                     partition_id,
@@ -214,7 +214,7 @@ async def shuffle_node(
                 ),
             )
 
-        await ch_out.data.drain(context)
+        await ch_out.drain(context)
 
 
 @generate_ir_sub_network.register(Shuffle)

python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py

@@ -477,8 +477,8 @@ def generate_network(
     nodes_dict, channels = mapper(ir)
     ch_out = channels[ir].reserve_output_slot()
 
-    # Add node to drain metadata channel before pull_from_channel
-    # (since pull_from_channel doesn't accept a ChannelPair)
+    # Add node to drain metadata before pull_from_channel
+    # (since pull_from_channel doesn't handle metadata messages)
     ch_final_data: Channel[TableChunk] = context.create_channel()
     drain_node = metadata_drain_node(
         context,

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

@@ -41,6 +41,7 @@
     ChannelManager,
     Metadata,
     opaque_reservation,
+    send_metadata,
 )
 
 if TYPE_CHECKING:
@@ -53,7 +54,6 @@
     from cudf_polars.experimental.base import ColumnStat, StatsCollector
     from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator
     from cudf_polars.experimental.rapidsmpf.dispatch import LowerIRTransformer
-    from cudf_polars.experimental.rapidsmpf.utils import ChannelPair
     from cudf_polars.utils.config import ParquetOptions
 
 
@@ -139,7 +139,7 @@ async def dataframescan_node(
     context: Context,
     ir: DataFrameScan,
     ir_context: IRExecutionContext,
-    ch_out: ChannelPair,
+    ch_out: Channel[TableChunk],
     *,
     num_producers: int,
     rows_per_partition: int,
@@ -157,7 +157,7 @@ async def dataframescan_node(
     ir_context
         The execution context for the IR node.
     ch_out
-        The output ChannelPair.
+        The output Channel[TableChunk].
     num_producers
         The number of producers to use for the DataFrameScan node.
     rows_per_partition
@@ -166,7 +166,7 @@ async def dataframescan_node(
         Estimated size of each chunk in bytes. Used for memory reservation
         with block spilling to avoid thrashing.
     """
-    async with shutdown_on_error(context, ch_out.metadata, ch_out.data):
+    async with shutdown_on_error(context, ch_out):
         # Find local partition count.
         nrows = ir.df.shape()[0]
         global_count = math.ceil(nrows / rows_per_partition) if nrows > 0 else 0
@@ -180,7 +180,8 @@ async def dataframescan_node(
             local_offset = local_count * context.comm().rank
 
         # Send basic metadata
-        await ch_out.send_metadata(
+        await send_metadata(
+            ch_out,
             context,
             Metadata(local_count=local_count, global_count=global_count),
         )
@@ -201,7 +202,7 @@ async def dataframescan_node(
 
         # If there are no slices, drain the channel and return
         if len(ir_slices) == 0:
-            await ch_out.data.drain(context)
+            await ch_out.drain(context)
             return
 
         # If there is only one ir_slices or one producer, we can
@@ -212,16 +213,16 @@ async def dataframescan_node(
                     context,
                     ir_slice,
                     seq_num,
-                    ch_out.data,
+                    ch_out,
                     ir_context,
                     estimated_chunk_bytes,
                 )
-            await ch_out.data.drain(context)
+            await ch_out.drain(context)
             return
 
         # Use Lineariser to ensure ordered delivery
         num_producers = min(num_producers, len(ir_slices))
-        lineariser = Lineariser(context, ch_out.data, num_producers)
+        lineariser = Lineariser(context, ch_out, num_producers)
 
         # Assign tasks to producers using round-robin
         producer_tasks: list[list[tuple[int, DataFrameScan]]] = [
@@ -365,7 +366,7 @@ async def scan_node(
     context: Context,
     ir: Scan,
     ir_context: IRExecutionContext,
-    ch_out: ChannelPair,
+    ch_out: Channel[TableChunk],
     *,
     num_producers: int,
     plan: IOPartitionPlan,
@@ -384,7 +385,7 @@ async def scan_node(
     ir_context
         The execution context for the IR node.
     ch_out
-        The output ChannelPair.
+        The output Channel[TableChunk].
     num_producers
         The number of producers to use for the scan node.
     plan
@@ -395,7 +396,7 @@ async def scan_node(
         Estimated size of each chunk in bytes. Used for memory reservation
         with block spilling to avoid thrashing.
     """
-    async with shutdown_on_error(context, ch_out.metadata, ch_out.data):
+    async with shutdown_on_error(context, ch_out):
         # Build a list of local Scan operations
         scans: list[Scan | SplitScan] = []
         if plan.flavor == IOPartitionFlavor.SPLIT_FILES:
@@ -464,14 +465,15 @@ async def scan_node(
                     )
 
         # Send basic metadata
-        await ch_out.send_metadata(
+        await send_metadata(
+            ch_out,
             context,
             Metadata(local_count=len(scans), global_count=count),
         )
 
         # If there is nothing to scan, drain the channel and return
         if len(scans) == 0:
-            await ch_out.data.drain(context)
+            await ch_out.drain(context)
             return
 
         # If there is only one scan or one producer, we can
@@ -482,16 +484,16 @@ async def scan_node(
                     context,
                     scan,
                     seq_num,
-                    ch_out.data,
+                    ch_out,
                     ir_context,
                     estimated_chunk_bytes,
                 )
-            await ch_out.data.drain(context)
+            await ch_out.drain(context)
             return
 
         # Use Lineariser to ensure ordered delivery
         num_producers = min(num_producers, len(scans))
-        lineariser = Lineariser(context, ch_out.data, num_producers)
+        lineariser = Lineariser(context, ch_out, num_producers)
 
         # Assign tasks to producers using round-robin
         producer_tasks: list[list[tuple[int, Scan | SplitScan]]] = [
@@ -524,7 +526,7 @@ def make_rapidsmpf_read_parquet_node(
     context: Context,
     ir: Scan,
     num_producers: int,
-    ch_out: ChannelPair,
+    ch_out: Channel[TableChunk],
     stats: StatsCollector,
     partition_info: PartitionInfo,
 ) -> Any | None:
@@ -540,7 +542,7 @@ def make_rapidsmpf_read_parquet_node(
     num_producers
         The number of producers to use for the scan node.
     ch_out
-        The output ChannelPair.
+        The output Channel[TableChunk].
     stats
         The statistics collector.
     partition_info
@@ -611,7 +613,7 @@ def make_rapidsmpf_read_parquet_node(
     try:
         return read_parquet(
             context,
-            ch_out.data,
+            ch_out,
             num_producers,
             parquet_reader_options,
             num_rows_per_chunk,
@@ -652,7 +654,8 @@ def _(
     )
 
     # Use rapidsmpf native read_parquet node if possible
-    ch_pair = channels[ir].reserve_input_slot()
+    ch_in: Channel[TableChunk] | None = None
+    ch_out = channels[ir].reserve_input_slot()
     nodes: dict[IR, list[Any]] = {}
     native_node: Any = None
     if (
@@ -665,21 +668,24 @@ def _(
         and ir.skip_rows == 0
         and not distributed_split_files
     ):
+        # Create new channel to so ch_out can be used to add metadata
+        ch_in = rec.state["context"].create_channel()
         native_node = make_rapidsmpf_read_parquet_node(
             rec.state["context"],
             ir,
             num_producers,
-            ch_pair,
+            ch_in,
             rec.state["stats"],
             partition_info,
         )
 
-    if native_node is not None:
+    if native_node is not None and ch_in is not None:
         # Need metadata node, because the native read_parquet
         # node does not send metadata.
         metadata_node = metadata_feeder_node(
             rec.state["context"],
-            ch_pair,
+            ch_in,
+            ch_out,
             Metadata(
                 # partition_info.count is the estimated "global" count.
                 # Just estimate the local count as well.
@@ -689,7 +695,7 @@ def _(
                 global_count=partition_info.count,
             ),
         )
-        nodes[ir] = [metadata_node, native_node]
+        nodes[ir] = [native_node, metadata_node]
     else:
         # Fall back to scan_node (predicate not convertible, or other constraint)
         parquet_options = dataclasses.replace(parquet_options, chunked=False)
@@ -699,7 +705,7 @@ def _(
                 rec.state["context"],
                 ir,
                 rec.state["ir_context"],
-                ch_pair,
+                ch_out,
                 num_producers=num_producers,
                 plan=plan,
                 parquet_options=parquet_options,