From 812cd97ac3dd888af365e66dfe5c6c4724b89426 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 10 Sep 2026 07:41:45 -0700
Subject: [PATCH 1/5] Simplify cudf-polars input/output bytes tracing

This updates how we track the number of bytes send / recv on rapidsmpf
channels, using the new metrics added in
https://github.com/rapidsai/rapidsmpf/pull/1207 rather than a proxy.
---
 .../actor_graph/collectives/shuffle.py        |  13 +--
 .../streaming/actor_graph/collectives/sort.py |  13 +--
 .../streaming/actor_graph/groupby.py          |  13 +--
 .../streaming/actor_graph/hint_sorted.py      |  13 +--
 .../cudf_polars/streaming/actor_graph/io.py   |  43 +++++---
 .../cudf_polars/streaming/actor_graph/join.py |  30 ++---
 .../streaming/actor_graph/nodes.py            |  45 +++++---
 .../cudf_polars/streaming/actor_graph/over.py |  13 +--
 .../streaming/actor_graph/repartition.py      |  13 +--
 .../streaming/actor_graph/tracing.py          | 104 +++++++-----------
 .../streaming/actor_graph/union.py            |  13 +--
 .../streaming/actor_graph/utils.py            |  52 +++++++--
 .../tests/streaming/test_tracing.py           |  35 +++++-
 13 files changed, 221 insertions(+), 179 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py
index d7bf5d9cf7f..d16d31c0e6c 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py
+++ b/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,
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
index 1d6932b9077..be6727d49bf 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
+++ b/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)
 
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py
index 894a358f968..95a97ccb7c4 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py
+++ b/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
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py
index a3fd0ff7f1c..cf0247708ed 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py
+++ b/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,
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
index f69c8757e80..49cf3da04ed 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/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)
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
index 75ad6e6cfd1..e7182533a89 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
+++ b/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,
         ):
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py
index 63a2904aa49..ddd303f69ab 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py
+++ b/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)
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
index 514b4e67d7f..9f4968e6549 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
@@ -66,10 +66,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,
     ChunkStore,
@@ -757,10 +754,12 @@ async def over_actor(
         time. ``None`` for non-scalar Over nodes.
     """
     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)
 
         partitioning = NormalizedPartitioning.from_keys(
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py
index 8c0daf1057d..8195f08d10f 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py
@@ -22,10 +22,7 @@
     ir_context_for_node,
 )
 from cudf_polars.streaming.actor_graph.nodes import 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,
     empty_table_chunk,
@@ -89,10 +86,12 @@ async def concatenate_node(
         Pre-allocated collective ID for this operation.
     """
     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)
         # Receive metadata.
         input_metadata = await recv_metadata(ch_in, context)
         nranks = comm.nranks
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py
index e41ffb7cccf..987bc2cbbdc 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py
@@ -5,14 +5,17 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
+from typing import TYPE_CHECKING, Any
 
+from rapidsmpf.memory.buffer import MemoryType
 from rapidsmpf.streaming.core.message import Message
 
 from cudf_polars.dsl.tracing import LOG_TRACES, Scope
 from cudf_polars.streaming.explain import SerializablePlan
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from cudf_streaming.table_chunk import TableChunk
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
@@ -20,7 +23,9 @@
     from cudf_polars.dsl.ir import IR
     from cudf_polars.utils.config import ConfigOptions
 
-T = TypeVar("T")
+
+def _zero_bytes_by_tier() -> dict[MemoryType, int]:
+    return dict.fromkeys(MemoryType, 0)
 
 
 @dataclasses.dataclass(slots=True)
@@ -41,6 +46,10 @@ class ActorTracer:
         None if row counting is not available for this node.
     chunk_count
         Total chunk count produced by this node during execution.
+    input_bytes
+        Bytes received on boundary input channels, stratified by memory tier.
+    output_bytes
+        Bytes sent on the boundary output channel, stratified by memory tier.
     decision
         The algorithm decision made at runtime for this node
         (e.g., "broadcast_left", "shuffle", "tree", etc.).
@@ -53,8 +62,12 @@ class ActorTracer:
     ir_type: str | None = None
     row_count: int | None = None
     chunk_count: int = 0
-    input_bytes: int = 0
-    output_bytes: int = 0
+    input_bytes: dict[MemoryType, int] = dataclasses.field(
+        default_factory=_zero_bytes_by_tier
+    )
+    output_bytes: dict[MemoryType, int] = dataclasses.field(
+        default_factory=_zero_bytes_by_tier
+    )
     decision: str | None = None
     duplicated: bool = False
     extra: dict[str, Any] = dataclasses.field(default_factory=dict)
@@ -89,69 +102,32 @@ def set_extra(self, key: str, value: Any) -> None:
         self.extra[key] = value
 
 
-class TracingChannel(Generic[T]):
+def record_channel_metrics(
+    tracer: ActorTracer | None,
+    *,
+    chs_in: Sequence[Channel[Any]] = (),
+    chs_out: Sequence[Channel[Any]] = (),
+) -> None:
     """
-    Channel proxy that records the bytes an actor reads and writes.
-
-    Wrap an actor's channels to attribute ``input_bytes`` (from ``recv``) and
-    ``output_bytes`` (from ``send``) to that actor.
+    Record boundary channel byte volumes on an actor tracer.
 
-    Internal channels an actor creates for its own sub-network are left
-    unwrapped so intermediate traffic is not counted.
+    Parameters
+    ----------
+    tracer
+        The actor tracer to update.
+    chs_in
+        Input boundary channels. ``recv_bytes`` are summed per memory tier.
+    chs_out
+        Output boundary channels. ``send_bytes`` are summed per memory tier.
     """
-
-    def __init__(self, channel: Channel[T], tracer: ActorTracer | None) -> None:
-        self._channel = channel
-        self._tracer = tracer
-
-    async def recv(self, context: Context) -> Message[T] | None:
-        """Wrapper around ``Channel.recv`` that records the input bytes."""
-        message = await self._channel.recv(context)
-        if message is not None and self._tracer is not None:
-            self._tracer.input_bytes += _message_size(message)
-        return message
-
-    async def send(self, context: Context, message: Message[T]) -> None:
-        """Wrapper around ``Channel.send`` that records the output bytes."""
-        if self._tracer is not None:
-            self._tracer.output_bytes += _message_size(message)
-        await self._channel.send(context, message)
-
-    # Implement the rest of the Channel interface.
-
-    async def drain(self, context: Context) -> None:
-        """Passthrough to ``Channel.drain``."""
-        await self._channel.drain(context)
-
-    async def drain_metadata(self, context: Context) -> None:
-        """Passthrough to ``Channel.drain_metadata``."""
-        await self._channel.drain_metadata(context)
-
-    async def shutdown(self, context: Context) -> None:
-        """Passthrough to ``Channel.shutdown``."""
-        await self._channel.shutdown(context)
-
-    async def shutdown_metadata(self, context: Context) -> None:
-        """Passthrough to ``Channel.shutdown_metadata``."""
-        await self._channel.shutdown_metadata(context)
-
-    async def recv_metadata(self, context: Context) -> Message[Any] | None:
-        """Passthrough to ``Channel.recv_metadata``."""
-        return await self._channel.recv_metadata(context)
-
-    async def send_metadata(self, context: Context, message: Message[Any]) -> None:
-        """Passthrough to ``Channel.send_metadata``."""
-        await self._channel.send_metadata(context, message)
-
-
-def _message_size(message: Message[Any]) -> int:
-    """Return the total data allocation size described by a message."""
-    return sum(message.get_content_description().content_sizes.values())
-
-
-def trace_channel(channel: Channel[T], tracer: ActorTracer | None) -> Channel[T]:
-    """Wrap one of an actor's boundary channels to record the bytes crossing it."""
-    return cast("Channel[T]", TracingChannel(channel, tracer))
+    if tracer is None:
+        return
+    for ch in chs_in:
+        for mem_type, nbytes in ch.metrics().recv_bytes.items():
+            tracer.input_bytes[mem_type] += nbytes
+    for ch in chs_out:
+        for mem_type, nbytes in ch.metrics().send_bytes.items():
+            tracer.output_bytes[mem_type] += nbytes
 
 
 async def send_chunk(
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py
index 07f670dce44..1b01b8a86e2 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py
@@ -19,9 +19,6 @@
     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 (
-    trace_channel,
-)
 from cudf_polars.streaming.actor_graph.utils import (
     ChannelManager,
     empty_table_chunk,
@@ -68,10 +65,12 @@ async def union_node(
         The input Channel[TableChunk]s.
     """
     async with shutdown_on_error(
-        context, *chs_in, 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)
+        context,
+        chs_in=chs_in,
+        chs_out=(ch_out,),
+        trace_ir=ir,
+        ir_context=ir_context,
+    ):
         # Merge and forward metadata.
         # Union loses partitioning/ordering info since sources may differ.
         # TODO: Warn users that Union does NOT preserve order?
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
index 6749cbdaf28..db3ede24240 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -31,6 +31,7 @@
     TableChunk,
     make_table_chunks_available_or_wait,
 )
+from rapidsmpf.memory.buffer import MemoryType
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.memory.packed_data import PackedData
 from rapidsmpf.streaming.coll.allgather import AllGather
@@ -52,7 +53,11 @@
 from cudf_polars.dsl.utils.column_domain import column_domain_bindings
 from cudf_polars.dsl.utils.naming import names_to_indices
 from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager
-from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk
+from cudf_polars.streaming.actor_graph.tracing import (
+    ActorTracer,
+    record_channel_metrics,
+    send_chunk,
+)
 from cudf_polars.streaming.utils import _concat
 from cudf_polars.utils.dtypes import make_empty_column
 
@@ -268,7 +273,9 @@ async def shutdown_channels_on_error(
 @asynccontextmanager
 async def shutdown_on_error(
     context: Context,
-    *channels: Channel[Any],
+    *,
+    chs_in: Sequence[Channel[Any]] = (),
+    chs_out: Sequence[Channel[Any]] = (),
     trace_ir: IR,
     ir_context: IRExecutionContext | None = None,
 ) -> AsyncIterator[ActorTracer | None]:
@@ -282,8 +289,12 @@ async def shutdown_on_error(
     ----------
     context
         The rapidsmpf context.
-    channels
-        The channels to shutdown on error.
+    chs_in
+        Boundary input channels. Shut down on error, and used to record
+        ``input_bytes`` from ``Channel.metrics().recv_bytes``.
+    chs_out
+        Boundary output channels. Shut down on error, and used to record
+        ``output_bytes`` from ``Channel.metrics().send_bytes``.
     trace_ir
         Optional IR node to enable tracing for this streaming actor.
         When provided and LOG_TRACES is enabled, an ActorTracer
@@ -298,6 +309,7 @@ async def shutdown_on_error(
     ActorTracer | None
         An actor tracer for collecting stats (if tracing enabled), else None.
     """
+    channels = (*chs_in, *chs_out)
     # Create tracer only if LOG_TRACES is enabled and IR is provided
     tracer: ActorTracer | None = None
     contextvars: dict[str, Any] = {}
@@ -320,6 +332,7 @@ async def shutdown_on_error(
             raise
         finally:
             stop = time.monotonic_ns()
+            record_channel_metrics(tracer, chs_in=chs_in, chs_out=chs_out)
             record: dict[str, Any] = {
                 "scope": Scope.ACTOR.value,
             }
@@ -370,16 +383,39 @@ async def shutdown_on_error(
                             value=tracer.decision,
                         )
                     )
+                if tracer is not None:
+                    for mem_type in MemoryType:
+                        tier = mem_type.name.lower()
+                        custom_attributes.append(
+                            cudf_polars.quent._types.StatisticsAttribute(
+                                key=f"input_bytes_{tier}",
+                                value_type="U64",
+                                value=tracer.input_bytes[mem_type],
+                            )
+                        )
+                        custom_attributes.append(
+                            cudf_polars.quent._types.StatisticsAttribute(
+                                key=f"output_bytes_{tier}",
+                                value_type="U64",
+                                value=tracer.output_bytes[mem_type],
+                            )
+                        )
                 if tracer is None or tracer.row_count is None:
                     # TODO: See if `output_rows` is nullable.
                     output_rows = 0
                 else:
                     output_rows = tracer.row_count
+                input_bytes = (
+                    sum(tracer.input_bytes.values()) if tracer is not None else 0
+                )
+                output_bytes = (
+                    sum(tracer.output_bytes.values()) if tracer is not None else 0
+                )
                 stats = quent_ir_execution_context.quent_operator.statistics(
                     statistics=cudf_polars.quent._types.Statistics(
                         output_rows=output_rows,
-                        input_bytes=tracer.input_bytes,
-                        output_bytes=tracer.output_bytes,
+                        input_bytes=input_bytes,
+                        output_bytes=output_bytes,
                         custom_attributes=custom_attributes,
                     )
                 )
@@ -1359,7 +1395,9 @@ async def replay_buffered_channel(
     trace_ir
         The IR node to trace. Passed through to shutdown_on_error.
     """
-    async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir):
+    async with shutdown_on_error(
+        context, chs_in=(ch_in,), chs_out=(ch_out,), trace_ir=trace_ir
+    ):
         await send_metadata(ch_out, context, metadata)
         for msg in buffered_chunks:
             await ch_out.send(context, msg)
diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py
index 6eb99da2b0b..f946085ebad 100644
--- a/python/cudf_polars/tests/streaming/test_tracing.py
+++ b/python/cudf_polars/tests/streaming/test_tracing.py
@@ -17,12 +17,17 @@
 import polars as pl
 
 from cudf_streaming.table_chunk import TableChunk
+from rapidsmpf.memory.buffer import MemoryType
 from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk
 from rapidsmpf.streaming.core.message import Message
 
 from cudf_polars.containers import DataFrame
 from cudf_polars.streaming.actor_graph.io import Lineariser
-from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk
+from cudf_polars.streaming.actor_graph.tracing import (
+    ActorTracer,
+    record_channel_metrics,
+    send_chunk,
+)
 
 if TYPE_CHECKING:
     import pathlib
@@ -48,6 +53,34 @@ def test_actor_tracer_counts_table_chunk_without_table_view(chunk: TableChunk) -
     assert tracer.row_count == 3
 
 
+@pytest.mark.spmd
+def test_record_channel_metrics_reads_send_and_recv_bytes(
+    spmd_engine: SPMDEngine, chunk: TableChunk
+) -> None:
+    context = spmd_engine.context
+    ch = context.create_channel()
+
+    async def send_and_recv() -> None:
+        async with asyncio.TaskGroup() as tg:
+            recv_task = tg.create_task(ch.recv(context))
+            tg.create_task(send_chunk(context, ch, chunk, 11, tracer=None))
+        recv_task.result()
+
+    asyncio.run(send_and_recv())
+
+    metrics = ch.metrics()
+    producer_tracer = ActorTracer()
+    consumer_tracer = ActorTracer()
+    record_channel_metrics(producer_tracer, chs_out=(ch,))
+    record_channel_metrics(consumer_tracer, chs_in=(ch,))
+
+    assert producer_tracer.output_bytes == metrics.send_bytes
+    assert consumer_tracer.input_bytes == metrics.recv_bytes
+    assert producer_tracer.output_bytes[MemoryType.DEVICE] > 0
+    assert producer_tracer.output_bytes[MemoryType.HOST] == 0
+    assert producer_tracer.output_bytes[MemoryType.PINNED_HOST] == 0
+
+
 @pytest.mark.spmd
 def test_send_chunk_traces_and_sends_message(
     spmd_engine: SPMDEngine, chunk: TableChunk

From 50ca871ceb78d720f378dc83901c42de3875ca9d Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 10 Sep 2026 11:49:46 -0700
Subject: [PATCH 2/5] Update test

---
 .../streaming/actor_graph/utils.py            |  4 +-
 .../tests/streaming/test_tracing.py           | 43 +++++++++++--------
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
index db3ede24240..a013089352e 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -278,7 +278,7 @@ async def shutdown_on_error(
     chs_out: Sequence[Channel[Any]] = (),
     trace_ir: IR,
     ir_context: IRExecutionContext | None = None,
-) -> AsyncIterator[ActorTracer | None]:
+) -> AsyncIterator[ActorTracer]:
     """
     Actor-level shutdown and tracing for rapidsmpf.
 
@@ -311,7 +311,7 @@ async def shutdown_on_error(
     """
     channels = (*chs_in, *chs_out)
     # Create tracer only if LOG_TRACES is enabled and IR is provided
-    tracer: ActorTracer | None = None
+    # tracer: ActorTracer | None = None
     contextvars: dict[str, Any] = {}
 
     ir_id = trace_ir.get_stable_id()
diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py
index f946085ebad..57fdaac0e8c 100644
--- a/python/cudf_polars/tests/streaming/test_tracing.py
+++ b/python/cudf_polars/tests/streaming/test_tracing.py
@@ -22,12 +22,10 @@
 from rapidsmpf.streaming.core.message import Message
 
 from cudf_polars.containers import DataFrame
+from cudf_polars.dsl.ir import Empty
 from cudf_polars.streaming.actor_graph.io import Lineariser
-from cudf_polars.streaming.actor_graph.tracing import (
-    ActorTracer,
-    record_channel_metrics,
-    send_chunk,
-)
+from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk
+from cudf_polars.streaming.actor_graph.utils import shutdown_on_error
 
 if TYPE_CHECKING:
     import pathlib
@@ -54,25 +52,36 @@ def test_actor_tracer_counts_table_chunk_without_table_view(chunk: TableChunk) -
 
 
 @pytest.mark.spmd
-def test_record_channel_metrics_reads_send_and_recv_bytes(
-    spmd_engine: SPMDEngine, chunk: TableChunk
-) -> None:
+def test_send_and_recv_bytes(spmd_engine: SPMDEngine, chunk: TableChunk) -> None:
     context = spmd_engine.context
     ch = context.create_channel()
+    ir = Empty({})
+
+    async def run() -> tuple[ActorTracer, ActorTracer]:
+
+        async def producer() -> ActorTracer:
+            async with shutdown_on_error(context, chs_out=(ch,), trace_ir=ir) as tracer:
+                await send_chunk(context, ch, chunk, 11, tracer=tracer)
+                await ch.drain(context)
+            return tracer
+
+        async def consumer() -> ActorTracer:
+            async with shutdown_on_error(context, chs_in=(ch,), trace_ir=ir) as tracer:
+                msg = await ch.recv(context)
+                assert msg is not None
+            return tracer
 
-    async def send_and_recv() -> None:
         async with asyncio.TaskGroup() as tg:
-            recv_task = tg.create_task(ch.recv(context))
-            tg.create_task(send_chunk(context, ch, chunk, 11, tracer=None))
-        recv_task.result()
+            producer_tracer_task = tg.create_task(producer())
+            consumer_tracer_task = tg.create_task(consumer())
+
+        producer_tracer = await producer_tracer_task
+        consumer_tracer = await consumer_tracer_task
 
-    asyncio.run(send_and_recv())
+        return producer_tracer, consumer_tracer
 
+    producer_tracer, consumer_tracer = asyncio.run(run())
     metrics = ch.metrics()
-    producer_tracer = ActorTracer()
-    consumer_tracer = ActorTracer()
-    record_channel_metrics(producer_tracer, chs_out=(ch,))
-    record_channel_metrics(consumer_tracer, chs_in=(ch,))
 
     assert producer_tracer.output_bytes == metrics.send_bytes
     assert consumer_tracer.input_bytes == metrics.recv_bytes

From 1f90c6e923c62f3f9969a5bd3e35c1b5ccd935e8 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 10 Sep 2026 13:00:24 -0700
Subject: [PATCH 3/5] cleanup

---
 python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
index f0fc7de13d8..ff575d71821 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -317,7 +317,6 @@ async def shutdown_on_error(
     """
     channels = (*chs_in, *chs_out)
     # Create tracer only if LOG_TRACES is enabled and IR is provided
-    # tracer: ActorTracer | None = None
     contextvars: dict[str, Any] = {}
 
     ir_id = trace_ir.get_stable_id()

From 08ddd6177ef26de35062a49f24fa6e431b68083b Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 11 Sep 2026 06:13:23 -0700
Subject: [PATCH 4/5] Aux channels

---
 .../cudf_polars/streaming/actor_graph/collectives/sort.py   | 3 ++-
 .../cudf_polars/streaming/actor_graph/hint_sorted.py        | 3 ++-
 .../cudf_polars/cudf_polars/streaming/actor_graph/join.py   | 3 ++-
 .../cudf_polars/streaming/actor_graph/prefilter_actor.py    | 3 ++-
 .../cudf_polars/cudf_polars/streaming/actor_graph/utils.py  | 6 +++++-
 5 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
index be6727d49bf..aa8629aa8b6 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
@@ -816,8 +816,9 @@ async def sort_actor(
     ch_chunk_store = context.create_channel()
     async with shutdown_on_error(
         context,
-        chs_in=(ch_in, ch_sample_replay, ch_chunk_store),
+        chs_in=(ch_in,),
         chs_out=(ch_out,),
+        auxiliary_channels=(ch_sample_replay, ch_chunk_store),
         trace_ir=ir,
         ir_context=ir_context,
     ) as tracer:
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py
index cf0247708ed..e0b03358f1b 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py
@@ -162,8 +162,9 @@ async def hint_sorted_actor(
     """Forward data and attach safe ordering metadata for ``hint_sorted``."""
     async with shutdown_on_error(
         context,
-        chs_in=(ch_in, ch_replay),
+        chs_in=(ch_in,),
         chs_out=(ch_out,),
+        auxiliary_channels=(ch_replay,),
         trace_ir=ir,
         ir_context=ir_context,
     ):
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
index 83bfbdeedc0..43f854eb147 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
@@ -1691,8 +1691,9 @@ async def join_actor(
     """
     async with shutdown_on_error(
         context,
-        chs_in=(ch_left, ch_right, *ch_prefilter_domains),
+        chs_in=(ch_left, ch_right),
         chs_out=(ch_out,),
+        auxiliary_channels=ch_prefilter_domains,
         trace_ir=ir,
         ir_context=ir_context,
     ) as tracer:
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py
index 9fee926d440..10adeb27533 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py
@@ -59,8 +59,9 @@ async def pushdown_filter_actor(
     collected_samples: Sequence[TableSizeStats] = []
     async with shutdown_on_error(
         context,
-        chs_in=(ch_target, ch_domain),
+        chs_in=(ch_target,),
         chs_out=(ch_out,),
+        auxiliary_channels=(ch_domain,),
         trace_ir=ir,
         ir_context=ir_context,
     ) as tracer:
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
index ff575d71821..3ec294aa794 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -282,6 +282,7 @@ async def shutdown_on_error(
     *,
     chs_in: Sequence[Channel[Any]] = (),
     chs_out: Sequence[Channel[Any]] = (),
+    auxiliary_channels: Sequence[Channel[Any]] = (),
     trace_ir: IR,
     ir_context: IRExecutionContext | None = None,
 ) -> AsyncIterator[ActorTracer]:
@@ -301,6 +302,9 @@ async def shutdown_on_error(
     chs_out
         Boundary output channels. Shut down on error, and used to record
         ``output_bytes`` from ``Channel.metrics().send_bytes``.
+    auxiliary_channels
+        Auxiliary channels. Shut down on error. Statistics from these channels
+        are not included in the actor tracing.
     trace_ir
         Optional IR node to enable tracing for this streaming actor.
         When provided and LOG_TRACES is enabled, an ActorTracer
@@ -315,7 +319,7 @@ async def shutdown_on_error(
     ActorTracer | None
         An actor tracer for collecting stats (if tracing enabled), else None.
     """
-    channels = (*chs_in, *chs_out)
+    channels = (*chs_in, *chs_out, *auxiliary_channels)
     # Create tracer only if LOG_TRACES is enabled and IR is provided
     contextvars: dict[str, Any] = {}
 

From e88ab4ce5771430e924413ba834611e18af151a0 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 11 Sep 2026 06:39:28 -0700
Subject: [PATCH 5/5] missed one

---
 .../cudf_polars/streaming/actor_graph/prefilter_actor.py        | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py
index 10adeb27533..1619cd915d5 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py
@@ -182,7 +182,7 @@ async def pushdown_filter_actor(
                 )
                 async with shutdown_on_error(
                     context,
-                    *execution.channels,
+                    auxiliary_channels=execution.channels,
                     trace_ir=ir,
                     ir_context=ir_context,
                 ):