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 d7bf5d9cf7f2..d16d31c0e6cd 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 1d6932b90770..aa8629aa8b68 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,12 @@ 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,), + chs_out=(ch_out,), + auxiliary_channels=(ch_sample_replay, ch_chunk_store), 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 894a358f9682..95a97ccb7c4b 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 a3fd0ff7f1c1..e0b03358f1bf 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,13 @@ 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,), + chs_out=(ch_out,), + auxiliary_channels=(ch_replay,), + 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 f69c8757e809..49cf3da04ed5 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 256a505961b6..43f854eb1475 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -54,7 +54,6 @@ from cudf_polars.streaming.actor_graph.tracing import ( LOG_TRACES, send_chunk, - trace_channel, ) from cudf_polars.streaming.actor_graph.utils import ( CUDF_ROW_LIMIT, @@ -262,15 +261,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, @@ -950,8 +945,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, ): @@ -1112,8 +1106,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, ): @@ -1698,16 +1691,12 @@ async def join_actor( """ async with shutdown_on_error( context, - ch_out, - 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: - 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, @@ -1768,9 +1757,7 @@ async def join_actor( ) async with shutdown_on_error( context, - ch_left_replay, - ch_right_replay, - *prefilter_execution.channels, + chs_out=(ch_left_replay, ch_right_replay, *prefilter_execution.channels), 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 63a2904aa49a..ddd303f69ab7 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 514b4e67d7f8..9f4968e65494 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/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py index 76ecee442821..1619cd915d5e 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,9 +59,9 @@ async def pushdown_filter_actor( collected_samples: Sequence[TableSizeStats] = [] async with shutdown_on_error( context, - ch_out, - 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: @@ -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, ): 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 8c0daf1057d8..8195f08d10fb 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 e41ffb7cccf9..987bc2cbbdc5 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 07f670dce445..1b01b8a86e2c 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 209df4d99338..3ec294aa7946 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -32,6 +32,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 @@ -53,7 +54,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 @@ -274,10 +279,13 @@ 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]] = (), + auxiliary_channels: Sequence[Channel[Any]] = (), trace_ir: IR, ir_context: IRExecutionContext | None = None, -) -> AsyncIterator[ActorTracer | None]: +) -> AsyncIterator[ActorTracer]: """ Actor-level shutdown and tracing for rapidsmpf. @@ -288,8 +296,15 @@ 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``. + 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 @@ -304,8 +319,8 @@ async def shutdown_on_error( ActorTracer | None An actor tracer for collecting stats (if tracing enabled), else None. """ + channels = (*chs_in, *chs_out, *auxiliary_channels) # 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() @@ -326,6 +341,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, } @@ -377,16 +393,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, ) ) @@ -1439,7 +1478,9 @@ async def replay_buffered_channel( The IR node to trace. Passed through to shutdown_on_error. """ try: - 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 99ec3b52af55..d4c77d953d54 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -17,12 +17,15 @@ 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.dsl.ir import Empty 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.utils import shutdown_on_error from cudf_polars.utils.versions import POLARS_VERSION_LT_138 if TYPE_CHECKING: @@ -49,6 +52,45 @@ def test_actor_tracer_counts_table_chunk_without_table_view(chunk: TableChunk) - assert tracer.row_count == 3 +@pytest.mark.spmd +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 with asyncio.TaskGroup() as tg: + 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 + + return producer_tracer, consumer_tracer + + producer_tracer, consumer_tracer = asyncio.run(run()) + metrics = ch.metrics() + + 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