Simplify cudf-polars input/output bytes tracing

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py

@@ -43,9 +43,6 @@
     ir_context_for_node,
 )
 from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error
-from cudf_polars.streaming.actor_graph.tracing import (
-    trace_channel,
-)
 from cudf_polars.streaming.actor_graph.utils import (
     ChannelManager,
     _is_already_partitioned,
@@ -625,10 +622,12 @@ async def shuffle_actor(
         The collective ID.
     """
     async with shutdown_on_error(
-        context, ch_in, ch_out, trace_ir=ir, ir_context=ir_context
-    ) as tracer:
-        ch_in = trace_channel(ch_in, tracer)
-        ch_out = trace_channel(ch_out, tracer)
+        context,
+        chs_in=(ch_in,),
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
+    ):
         await _global_shuffle(
             context,
             comm,

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py

@@ -40,10 +40,7 @@
     default_node_single,
     shutdown_on_error,
 )
-from cudf_polars.streaming.actor_graph.tracing import (
-    send_chunk,
-    trace_channel,
-)
+from cudf_polars.streaming.actor_graph.tracing import send_chunk
 from cudf_polars.streaming.actor_graph.utils import (
     ChannelManager,
     ChunkStore,
@@ -819,15 +816,11 @@ async def sort_actor(
     ch_chunk_store = context.create_channel()
     async with shutdown_on_error(
         context,
-        ch_in,
-        ch_out,
-        ch_sample_replay,
-        ch_chunk_store,
+        chs_in=(ch_in, ch_sample_replay, ch_chunk_store),
+        chs_out=(ch_out,),
         trace_ir=ir,
         ir_context=ir_context,
     ) as tracer:
-        ch_in = trace_channel(ch_in, tracer)
-        ch_out = trace_channel(ch_out, tracer)
         # TODO: Skip sort if OrderScheme metadata is present and compatible.
         metadata_in = await recv_metadata(ch_in, context)
 

python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py

@@ -34,10 +34,7 @@
     generate_ir_sub_network,
     ir_context_for_node,
 )
-from cudf_polars.streaming.actor_graph.tracing import (
-    send_chunk,
-    trace_channel,
-)
+from cudf_polars.streaming.actor_graph.tracing import send_chunk
 from cudf_polars.streaming.actor_graph.utils import (
     MAX_ROWS_PER_PARTITION,
     ChannelManager,
@@ -876,10 +873,12 @@ async def groupby_actor(
         The collective IDs.
     """
     async with shutdown_on_error(
-        context, ch_in, ch_out, trace_ir=ir, ir_context=ir_context
+        context,
+        chs_in=(ch_in,),
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
     ) as tracer:
-        ch_in = trace_channel(ch_in, tracer)
-        ch_out = trace_channel(ch_out, tracer)
         metadata_in = await recv_metadata(ch_in, context)
 
         nranks = comm.nranks

python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py

@@ -20,9 +20,6 @@
 from cudf_polars.dsl.ir import IR, MapFunction
 from cudf_polars.dsl.utils.naming import names_to_indices
 from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network
-from cudf_polars.streaming.actor_graph.tracing import (
-    trace_channel,
-)
 from cudf_polars.streaming.actor_graph.utils import (
     ChannelManager,
     process_children,
@@ -164,10 +161,12 @@ async def hint_sorted_actor(
 ) -> None:
     """Forward data and attach safe ordering metadata for ``hint_sorted``."""
     async with shutdown_on_error(
-        context, ch_in, ch_replay, ch_out, trace_ir=ir, ir_context=ir_context
-    ) as tracer:
-        ch_in = trace_channel(ch_in, tracer)
-        ch_out = trace_channel(ch_out, tracer)
+        context,
+        chs_in=(ch_in, ch_replay),
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
+    ):
         metadata = await recv_metadata(ch_in, context)
         metadata, ch_forward = await extract_hint_sorted_metadata(
             context,

python/cudf_polars/cudf_polars/streaming/actor_graph/io.py

@@ -32,10 +32,7 @@
     ir_context_for_node,
 )
 from cudf_polars.streaming.actor_graph.nodes import define_actor, shutdown_on_error
-from cudf_polars.streaming.actor_graph.tracing import (
-    send_chunk,
-    trace_channel,
-)
+from cudf_polars.streaming.actor_graph.tracing import send_chunk
 from cudf_polars.streaming.actor_graph.utils import (
     ChannelManager,
     chunk_to_frame,
@@ -202,9 +199,11 @@ async def dataframescan_node(
         ``Cluster.SPMD`` mode.
     """
     async with shutdown_on_error(
-        context, ch_out, trace_ir=ir, ir_context=ir_context
+        context,
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
     ) as tracer:
-        ch_out = trace_channel(ch_out, tracer)
         # Find local partition count.
         nrows = ir.df.shape()[0]
         global_count = math.ceil(nrows / rows_per_partition) if nrows > 0 else 0
@@ -308,7 +307,10 @@ async def _producer(producer_id: int) -> None:
 
         async with (
             shutdown_on_error(
-                context, *lineariser.input_channels, trace_ir=ir, ir_context=ir_context
+                context,
+                chs_out=lineariser.input_channels,
+                trace_ir=ir,
+                ir_context=ir_context,
             ),
         ):
             await gather_in_task_group(
@@ -448,9 +450,11 @@ async def python_scan_node(
         The output Channel[TableChunk].
     """
     async with shutdown_on_error(
-        context, ch_out, trace_ir=ir, ir_context=ir_context
+        context,
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
     ) as tracer:
-        ch_out = trace_channel(ch_out, tracer)
         rank_aware_source = _find_rank_aware_source(ir.options[0])
         if rank_aware_source is None and comm.nranks > 1 and comm.rank != 0:
             # A plain (rank-unaware) source runs on rank 0 only; other ranks
@@ -645,9 +649,11 @@ async def scan_node(
     scans: Sequence[SplitScan] | Sequence[FusedScan] = ir.scans
 
     async with shutdown_on_error(
-        context, ch_out, trace_ir=ir, ir_context=ir_context
+        context,
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
     ) as tracer:
-        ch_out = trace_channel(ch_out, tracer)
         # Send basic metadata
         ir_context = dataclasses.replace(ir_context, tracer=tracer)
         await send_metadata(
@@ -706,7 +712,10 @@ async def _producer(producer_id: int) -> None:
 
         async with (
             shutdown_on_error(
-                context, *lineariser.input_channels, trace_ir=ir, ir_context=ir_context
+                context,
+                chs_out=lineariser.input_channels,
+                trace_ir=ir,
+                ir_context=ir_context,
             ),
         ):
             await gather_in_task_group(
@@ -791,10 +800,12 @@ async def sink_node(
     # with other files.
 
     async with shutdown_on_error(
-        context, ch_in, ch_out, ir_context=ir_context, trace_ir=ir
-    ) as tracer:
-        ch_in = trace_channel(ch_in, tracer)
-        ch_out = trace_channel(ch_out, tracer)
+        context,
+        chs_in=(ch_in,),
+        chs_out=(ch_out,),
+        ir_context=ir_context,
+        trace_ir=ir,
+    ):
         metadata = await recv_metadata(ch_in, context)
         await send_metadata(
             ch_out, context, ChannelMetadata(local_count=1, duplicated=True)

python/cudf_polars/cudf_polars/streaming/actor_graph/join.py

@@ -44,10 +44,7 @@
     ir_context_for_node,
 )
 from cudf_polars.streaming.actor_graph.nodes import default_node_multi
-from cudf_polars.streaming.actor_graph.tracing import (
-    send_chunk,
-    trace_channel,
-)
+from cudf_polars.streaming.actor_graph.tracing import send_chunk
 from cudf_polars.streaming.actor_graph.utils import (
     CUDF_ROW_LIMIT,
     MAX_ROWS_PER_PARTITION,
@@ -192,15 +189,11 @@ async def broadcast_join_actor(
     """
     async with shutdown_on_error(
         context,
-        ch_out,
-        ch_left,
-        ch_right,
+        chs_in=(ch_left, ch_right),
+        chs_out=(ch_out,),
         trace_ir=ir,
         ir_context=ir_context,
     ) as tracer:
-        ch_left = trace_channel(ch_left, tracer)
-        ch_right = trace_channel(ch_right, tracer)
-        ch_out = trace_channel(ch_out, tracer)
         ir_context = replace(ir_context, tracer=tracer)
         await _broadcast_join(
             context,
@@ -723,8 +716,7 @@ async def _shuffle_join(
     # note: this is an actor inside of an actor. How should we log that in our traces?
     async with shutdown_on_error(
         context,
-        ch_left_shuffle,
-        ch_right_shuffle,
+        chs_out=(ch_left_shuffle, ch_right_shuffle),
         trace_ir=ir,
         ir_context=ir_context,
     ):
@@ -885,8 +877,7 @@ async def _ordered_join(
     ch_right_adjusted = context.create_channel()
     async with shutdown_on_error(
         context,
-        ch_left_adjusted,
-        ch_right_adjusted,
+        chs_out=(ch_left_adjusted, ch_right_adjusted),
         trace_ir=ir,
         ir_context=ir_context,
     ):
@@ -1310,15 +1301,11 @@ async def join_actor(
     """
     async with shutdown_on_error(
         context,
-        ch_out,
-        ch_left,
-        ch_right,
+        chs_in=(ch_left, ch_right),
+        chs_out=(ch_out,),
         trace_ir=ir,
         ir_context=ir_context,
     ) as tracer:
-        ch_left = trace_channel(ch_left, tracer)
-        ch_right = trace_channel(ch_right, tracer)
-        ch_out = trace_channel(ch_out, tracer)
         ir_context = replace(ir_context, tracer=tracer)
         left_metadata, right_metadata = await gather_in_task_group(
             recv_metadata(ch_left, context),
@@ -1341,8 +1328,7 @@ async def join_actor(
         ch_right_replay = context.create_channel()
         async with shutdown_on_error(
             context,
-            ch_left_replay,
-            ch_right_replay,
+            chs_out=(ch_left_replay, ch_right_replay),
             trace_ir=ir,
             ir_context=ir_context,
         ):

python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py

@@ -23,10 +23,7 @@
     generate_ir_sub_network,
     ir_context_for_node,
 )
-from cudf_polars.streaming.actor_graph.tracing import (
-    send_chunk,
-    trace_channel,
-)
+from cudf_polars.streaming.actor_graph.tracing import send_chunk
 from cudf_polars.streaming.actor_graph.utils import (
     ChannelManager,
     _leading_order_keys,
@@ -94,10 +91,12 @@ async def default_node_single(
     Chunks are processed in the order they are received.
     """
     async with shutdown_on_error(
-        context, ch_in, ch_out, trace_ir=ir, ir_context=ir_context
+        context,
+        chs_in=(ch_in,),
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
     ) as tracer:
-        ch_in = trace_channel(ch_in, tracer)
-        ch_out = trace_channel(ch_out, tracer)
         # Recv metadata and prepare output metadata
         metadata_in = await recv_metadata(ch_in, context)
         partitioning = maybe_remap_partitioning(
@@ -158,10 +157,12 @@ async def default_node_multi(
         If None, no partitioning information is preserved.
     """
     async with shutdown_on_error(
-        context, *chs_in, ch_out, trace_ir=ir, ir_context=ir_context
+        context,
+        chs_in=chs_in,
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
     ) as tracer:
-        chs_in = tuple(trace_channel(ch, tracer) for ch in chs_in)
-        ch_out = trace_channel(ch_out, tracer)
         # Merge and forward basic metadata.
         local_count = 1
         duplicated = True
@@ -311,7 +312,11 @@ async def fanout_node_bounded(
     # See: https://github.com/rapidsai/rapidsmpf/issues/560
     # TODO: Use ir_context
     async with shutdown_on_error(
-        context, ch_in, *chs_out, trace_ir=trace_ir, ir_context=ir_context
+        context,
+        chs_in=(ch_in,),
+        chs_out=chs_out,
+        trace_ir=trace_ir,
+        ir_context=ir_context,
     ):
         # Forward metadata to all outputs.
         metadata = await recv_metadata(ch_in, context)
@@ -387,7 +392,11 @@ async def fanout_node_unbounded(
     # See: https://github.com/rapidsai/rapidsmpf/issues/560
     # TODO: Use ir_context
     async with shutdown_on_error(
-        context, ch_in, *chs_out, trace_ir=trace_ir, ir_context=ir_context
+        context,
+        chs_in=(ch_in,),
+        chs_out=chs_out,
+        trace_ir=trace_ir,
+        ir_context=ir_context,
     ):
         # Forward metadata to all outputs.
         metadata = await recv_metadata(ch_in, context)
@@ -623,9 +632,11 @@ async def empty_node(
         The output Channel[TableChunk].
     """
     async with shutdown_on_error(
-        context, ch_out, ir_context=ir_context, trace_ir=ir
+        context,
+        chs_out=(ch_out,),
+        ir_context=ir_context,
+        trace_ir=ir,
     ) as tracer:
-        ch_out = trace_channel(ch_out, tracer)
         # Send metadata indicating a single empty chunk
         await send_metadata(
             ch_out,
@@ -641,7 +652,7 @@ async def empty_node(
         chunk = TableChunk.from_pylibcudf_table(
             df.table, df.stream, exclusive_view=True, br=context.br()
         )
-        await ch_out.send(context, Message(0, chunk))
+        await send_chunk(context, ch_out, chunk, 0, tracer=tracer)
 
         await ch_out.drain(context)
 
@@ -739,7 +750,7 @@ async def metadata_feeder_node(
     """
     # TODO: Use ir_context
     async with shutdown_on_error(
-        context, ch_in, ch_out, trace_ir=ir, ir_context=ir_context
+        context, chs_in=(ch_in,), chs_out=(ch_out,), trace_ir=ir, ir_context=ir_context
     ):
         await send_metadata(ch_out, context, metadata)
         while (msg := await ch_in.recv(context)) is not None:
@@ -780,7 +791,7 @@ async def metadata_drain_node(
         If None, metadata will not be collected.
     """
     async with shutdown_on_error(
-        context, ch_in, ch_out, ir_context=ir_context, trace_ir=ir
+        context, chs_in=(ch_in,), chs_out=(ch_out,), ir_context=ir_context, trace_ir=ir
     ):
         # Drain metadata channel (we don't need it after this point)
         msg = await ch_in.recv_metadata(context)