diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index 52864f6065b2..9c0034b9f1df 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -11,6 +11,7 @@ For the most part, the public API of `cudf-polars` is the polars API. DynamicPlanningOptions, InMemoryExecutor, ParquetOptions, + TracingOptions, Cluster, ShuffleMethod, ShufflerInsertionMethod, diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 9d59650b9688..74ec91ade58e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -257,6 +257,7 @@ class RunConfig: collect_traces: bool = False stats_planning: bool dynamic_planning: bool | None = None + trace_output_path: str | None = None max_io_threads: int native_parquet: bool spill_to_pinned_memory: bool @@ -377,6 +378,7 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: collect_traces=args.collect_traces, stats_planning=args.stats_planning, dynamic_planning=args.dynamic_planning, + trace_output_path=args.trace_output_path, max_io_threads=args.max_io_threads, native_parquet=args.native_parquet, extra_info=args.extra_info, @@ -477,6 +479,8 @@ def get_executor_options( if run_config.dynamic_planning: # Pass empty dict to enable with defaults; None means disabled executor_options["dynamic_planning"] = {} + if run_config.trace_output_path is not None: + executor_options["tracing"] = {"output_path": run_config.trace_output_path} if ( benchmark @@ -974,6 +978,13 @@ def parse_args( default=False, help="Enable dynamic shuffle planning (not yet implemented). ", ) + parser.add_argument( + "--trace-output-path", + dest="trace_output_path", + type=str, + default=None, + help="Path to write tracing output (row counts per node).", + ) parser.add_argument( "--max-io-threads", default=2, diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 17cba1d17cb4..564039507ac6 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 @@ -26,11 +26,13 @@ if TYPE_CHECKING: from collections.abc import MutableMapping + from pathlib import Path import polars as pl from cudf_polars.dsl.ir import IR from cudf_polars.experimental.base import PartitionInfo, StatsCollector + from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer def explain_query( @@ -168,3 +170,67 @@ def _(ir: Sort, *, offset: str = "") -> str: def _(ir: Scan, *, offset: str = "") -> str: label = f"SCAN {ir.typ.upper()}" return _repr_header(offset, label, ir.schema) + + +def write_query_trace( + trace_output: str | Path, + ir: IR, + partition_info: MutableMapping[IR, PartitionInfo], + tracer: StreamingQueryTracer, +) -> None: + """ + Write a post-execution trace showing actual row counts and decisions. + + Parameters + ---------- + trace_output + Path to write the trace file. + ir + The lowered IR root node. + partition_info + Partition information for the IR nodes. + tracer + The tracer with actual row counts and decisions from execution. + """ + from pathlib import Path + + trace_repr = _repr_trace_tree(ir, partition_info, tracer) + Path(trace_output).write_text(trace_repr) + + +def _repr_trace_tree( + ir: IR, + partition_info: MutableMapping[IR, PartitionInfo], + tracer: StreamingQueryTracer, + *, + offset: str = "", +) -> str: + """Recursively build a tree representation with tracer data.""" + header = _repr_ir(ir, offset=offset) + header = header.rstrip("\n") + + # Get node tracer if it exists + if (node_tracer := tracer.node_tracers.get(ir)) is not None: + # Add actual row count if available + if node_tracer.row_count is not None: + header += f" rows={_fmt_row_count(node_tracer.row_count)}" + + # Add decision if present + if node_tracer.decision is not None: + header += f" decision={node_tracer.decision}" + + # Add actual chunk count + header += f" chunks={node_tracer.chunk_count}" + + children_strs = [ + _repr_trace_tree(child, partition_info, tracer, offset=offset + " ") + for child in ir.children + ] + + header += "\n" + return header + "".join( + f"{line}{offset} (repeated {count} times)\n" + if (count := sum(1 for _ in group)) > 1 + else line + for line, group in groupby(children_strs) + ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index e9408b8c520e..5d6093247770 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -32,6 +32,7 @@ import cudf_polars.experimental.rapidsmpf.union # noqa: F401 from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import DataFrameScan, IRExecutionContext, Join, Scan, Union +from cudf_polars.dsl.tracing import LOG_TRACES from cudf_polars.dsl.traversal import CachingVisitor, traversal from cudf_polars.experimental.rapidsmpf.collectives import ReserveOpIDs from cudf_polars.experimental.rapidsmpf.dispatch import FanoutInfo, lower_ir_node @@ -39,6 +40,7 @@ generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk from cudf_polars.experimental.statistics import collect_statistics from cudf_polars.experimental.utils import _concat @@ -107,7 +109,7 @@ def evaluate_logical_plan( # NOTE: Distributed execution requires Dask for now from cudf_polars.experimental.rapidsmpf.dask import evaluate_pipeline_dask - result, metadata_collector = evaluate_pipeline_dask( + result, metadata_collector, tracer = evaluate_pipeline_dask( evaluate_pipeline, ir, partition_info, @@ -118,7 +120,7 @@ def evaluate_logical_plan( ) else: # Single-process execution: Run locally - result, metadata_collector = evaluate_pipeline( + result, metadata_collector, tracer = evaluate_pipeline( ir, partition_info, config_options, @@ -127,6 +129,13 @@ def evaluate_logical_plan( collect_metadata=collect_metadata, ) + # Write tracer output if configured + tracing = config_options.executor.tracing + if tracing is not None and tracing.output_path is not None and tracer is not None: + from cudf_polars.experimental.explain import write_query_trace + + write_query_trace(tracing.output_path, ir, partition_info, tracer) + return result, metadata_collector @@ -139,7 +148,7 @@ def evaluate_pipeline( rmpf_context: Context | None = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: """ Build and evaluate a RapidsMPF streaming pipeline. @@ -162,7 +171,7 @@ def evaluate_pipeline( Returns ------- - The output DataFrame and metadata collector. + The output DataFrame, metadata collector, and tracer. """ assert config_options.executor.name == "streaming", "Executor must be streaming" assert config_options.executor.runtime == "rapidsmpf", "Runtime must be rapidsmpf" @@ -232,6 +241,11 @@ def evaluate_pipeline( metadata_collector: list[ChannelMetadata] | None = ( [] if collect_metadata else None ) + tracer: StreamingQueryTracer | None = ( + StreamingQueryTracer() + if config_options.executor.tracing is not None or LOG_TRACES + else None + ) nodes, output = generate_network( rmpf_context, ir, @@ -241,6 +255,7 @@ def evaluate_pipeline( ir_context=ir_context, collective_id_map=collective_id_map, metadata_collector=metadata_collector, + tracer=tracer, ) # Run the network @@ -294,7 +309,7 @@ def evaluate_pipeline( if _initial_mr is not None: rmm.mr.set_current_device_resource(_original_mr) - return result, metadata_collector + return result, metadata_collector, tracer def lower_ir_graph( @@ -416,6 +431,7 @@ def generate_network( ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], metadata_collector: list[ChannelMetadata] | None, + tracer: StreamingQueryTracer | None = None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. @@ -440,6 +456,8 @@ def generate_network( The list to collect the final metadata. This list will be mutated when the network is executed. If None, metadata will not be collected. + tracer + Profiler for collecting runtime statistics. Returns ------- @@ -472,6 +490,7 @@ def generate_network( "max_io_threads": max_io_threads_local, "stats": stats, "collective_id_map": collective_id_map, + "tracer": tracer, } mapper: SubNetGenerator = CachingVisitor( generate_ir_sub_network_wrapper, state=state diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py index f1e25b3296e0..565180d7752b 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py @@ -24,6 +24,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.experimental.base import PartitionInfo, StatsCollector from cudf_polars.experimental.parallel import ConfigOptions + from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer class EvaluatePipelineCallback(Protocol): @@ -39,8 +40,8 @@ def __call__( rmpf_context: Context | None = None, *, collect_metadata: bool = False, - ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: - """Evaluate a pipeline and return the result DataFrame and metadata.""" + ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: + """Evaluate a pipeline and return the result DataFrame, metadata, and tracer.""" ... @@ -61,7 +62,7 @@ def evaluate_pipeline_dask( collective_id_map: dict[IR, list[int]], *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: """ Evaluate a RapidsMPF streaming pipeline on a Dask cluster. @@ -84,7 +85,7 @@ def evaluate_pipeline_dask( Returns ------- - The output DataFrame and metadata collector. + The output DataFrame, metadata collector, and merged tracer. """ client = get_dask_client() result = client.run( @@ -99,12 +100,18 @@ def evaluate_pipeline_dask( ) dfs: list[pl.DataFrame] = [] metadata_collector: list[ChannelMetadata] = [] - for df, md in result.values(): + merged_tracer: StreamingQueryTracer | None = None + for df, md, tracer in result.values(): dfs.append(df) if md is not None: metadata_collector.extend(md) + if tracer is not None: + if merged_tracer is None: + merged_tracer = tracer + else: + merged_tracer.merge(tracer) - return pl.concat(dfs), metadata_collector or None + return pl.concat(dfs), metadata_collector or None, merged_tracer def _evaluate_pipeline_dask( @@ -117,7 +124,7 @@ def _evaluate_pipeline_dask( dask_worker: Any = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: """ Build and evaluate a RapidsMPF streaming pipeline. @@ -144,7 +151,7 @@ def _evaluate_pipeline_dask( Returns ------- - The output DataFrame and metadata collector. + The output DataFrame, metadata collector, and tracer. """ assert dask_worker is not None, "Dask worker must be provided" assert config_options.executor.name == "streaming", "Executor must be streaming" diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py index 0706869bd8f3..061a4cda89a7 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py @@ -15,10 +15,8 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import ( - PartitionInfo, - StatsCollector, - ) + from cudf_polars.experimental.base import PartitionInfo, StatsCollector + from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer from cudf_polars.experimental.rapidsmpf.utils import ChannelManager from cudf_polars.utils.config import ConfigOptions @@ -77,6 +75,8 @@ class GenState(TypedDict): Statistics collector. collective_id_map The mapping of IR nodes to lists of collective IDs. + tracer + Runtime tracer for collecting execution statistics. """ context: Context @@ -87,6 +87,7 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] + tracer: StreamingQueryTracer | None SubNetGenerator: TypeAlias = GenericTransformer[ diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py index 71c033a865f6..d79c3dccec42 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -54,6 +54,10 @@ 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.tracing import ( + StreamingNodeTracer, + StreamingQueryTracer, + ) from cudf_polars.utils.config import ParquetOptions @@ -144,6 +148,7 @@ async def dataframescan_node( num_producers: int, rows_per_partition: int, estimated_chunk_bytes: int, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ DataFrameScan node for rapidsmpf. @@ -165,8 +170,10 @@ async def dataframescan_node( estimated_chunk_bytes Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. + node_tracer + The node tracer for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_out): + async with shutdown_on_error(context, ch_out, node_tracer=node_tracer) as tracer: # Find local partition count. nrows = ir.df.shape()[0] global_count = math.ceil(nrows / rows_per_partition) if nrows > 0 else 0 @@ -216,6 +223,7 @@ async def dataframescan_node( ch_out, ir_context, estimated_chunk_bytes, + node_tracer=tracer, ) await ch_out.drain(context) return @@ -241,6 +249,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, + node_tracer=tracer, ) await ch_out.drain(context) @@ -266,6 +275,7 @@ def _( context = rec.state["context"] ir_context = rec.state["ir_context"] + tracer: StreamingQueryTracer | None = rec.state["tracer"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} nodes: dict[IR, list[Any]] = { ir: [ @@ -277,6 +287,7 @@ def _( num_producers=num_producers, rows_per_partition=rows_per_partition, estimated_chunk_bytes=estimated_chunk_bytes, + node_tracer=tracer.get_or_create(ir) if tracer else None, ) ] } @@ -322,6 +333,7 @@ async def read_chunk( ch_out: Channel[TableChunk], ir_context: IRExecutionContext, estimated_chunk_bytes: int, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Read a chunk from disk and send it to the output channel. @@ -341,6 +353,8 @@ async def read_chunk( estimated_chunk_bytes Estimated size of the chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. + node_tracer + The node tracer for collecting runtime statistics. """ with opaque_reservation(context, estimated_chunk_bytes): df = await asyncio.to_thread( @@ -348,6 +362,8 @@ async def read_chunk( *scan._non_child_args, context=ir_context, ) + if node_tracer is not None: + node_tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -372,6 +388,7 @@ async def scan_node( plan: IOPartitionPlan, parquet_options: ParquetOptions, estimated_chunk_bytes: int, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Scan node for rapidsmpf. @@ -395,8 +412,10 @@ async def scan_node( estimated_chunk_bytes Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. + node_tracer + The node tracer for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_out): + async with shutdown_on_error(context, ch_out, node_tracer=node_tracer) as tracer: # Build a list of local Scan operations scans: list[Scan | SplitScan] = [] if plan.flavor == IOPartitionFlavor.SPLIT_FILES: @@ -487,6 +506,7 @@ async def scan_node( ch_out, ir_context, estimated_chunk_bytes, + node_tracer=tracer, ) await ch_out.drain(context) return @@ -512,6 +532,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, + node_tracer=tracer, ) await ch_out.drain(context) @@ -642,6 +663,7 @@ def _( parquet_options = config_options.parquet_options partition_info = rec.state["partition_info"][ir] num_producers = rec.state["max_io_threads"] + tracer: StreamingQueryTracer | None = rec.state["tracer"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} assert partition_info.io_plan is not None, "Scan node must have a partition plan" @@ -693,6 +715,7 @@ def _( partition_info.count / rec.state["context"].comm().nranks ), ), + node_tracer=tracer.get_or_create(ir) if tracer else None, ) nodes[ir] = [native_node, metadata_node] else: @@ -709,6 +732,7 @@ def _( plan=plan, parquet_options=parquet_options, estimated_chunk_bytes=executor.target_partition_size, + node_tracer=tracer.get_or_create(ir) if tracer else None, ) ] return nodes, channels diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index 4330cdb176fe..8542633d9618 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -37,6 +37,10 @@ from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import ( + StreamingNodeTracer, + StreamingQueryTracer, + ) @define_py_node() @@ -48,6 +52,7 @@ async def default_node_single( ch_in: Channel[TableChunk], *, preserve_partitioning: bool = False, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Single-channel default node for rapidsmpf. @@ -66,12 +71,16 @@ async def default_node_single( The input Channel[TableChunk]. preserve_partitioning Whether to preserve the partitioning metadata of the input chunks. + node_tracer + Node tracer for collecting runtime statistics. Notes ----- Chunks are processed in the order they are received. """ - async with shutdown_on_error(context, ch_in, ch_out): + async with shutdown_on_error( + context, ch_in, ch_out, node_tracer=node_tracer + ) as tracer: # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) partitioning = None @@ -86,6 +95,8 @@ async def default_node_single( duplicated=metadata_in.duplicated, ) await send_metadata(ch_out, context, metadata_out) + if tracer is not None and metadata_in.duplicated: + tracer.set_duplicated() # Recv/send data. seq_num = 0 @@ -123,6 +134,8 @@ async def default_node_single( ), context=ir_context, ) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -146,6 +159,7 @@ async def default_node_multi( chs_in: tuple[Channel[TableChunk], ...], *, partitioning_index: int | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Pointwise node for rapidsmpf. @@ -165,8 +179,12 @@ async def default_node_multi( partitioning_index Index of the input channel to preserve partitioning information for. If None, no partitioning information is preserved. + node_tracer + Node tracer for collecting runtime statistics. """ - async with shutdown_on_error(context, *chs_in, ch_out): + async with shutdown_on_error( + context, *chs_in, ch_out, node_tracer=node_tracer + ) as tracer: # Merge and forward basic metadata. local_count = 1 duplicated = True @@ -189,6 +207,8 @@ async def default_node_multi( duplicated=duplicated, ) await send_metadata(ch_out, context, metadata) + if tracer is not None and duplicated: + tracer.set_duplicated() seq_num = 0 n_children = len(chs_in) @@ -253,6 +273,8 @@ async def default_node_multi( *dfs, context=ir_context, ) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -531,6 +553,9 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + tracer: StreamingQueryTracer | None = rec.state.get("tracer") + node_tracer = tracer.get_or_create(ir) if tracer is not None else None + if len(ir.children) == 1: # Single-channel default node preserve_partitioning = isinstance( @@ -549,6 +574,7 @@ def _( channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), preserve_partitioning=preserve_partitioning, + node_tracer=node_tracer, ) ] else: @@ -560,6 +586,7 @@ def _( rec.state["ir_context"], channels[ir].reserve_input_slot(), tuple(channels[c].reserve_output_slot() for c in ir.children), + node_tracer=node_tracer, ) ] @@ -673,6 +700,7 @@ async def metadata_feeder_node( ch_in: Channel[TableChunk], ch_out: Channel[TableChunk], metadata: ChannelMetadata, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Forward data with new metadata. @@ -687,11 +715,19 @@ async def metadata_feeder_node( The output channel to forward data to and add metadata to. metadata The metadata to add to the output channel. + node_tracer + Node tracer for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_in, ch_out): + async with shutdown_on_error( + context, ch_in, ch_out, node_tracer=node_tracer + ) as tracer: await send_metadata(ch_out, context, metadata) + if tracer is not None and metadata.duplicated: + tracer.set_duplicated() while (msg := await ch_in.recv(context)) is not None: await ch_out.send(context, msg) + if tracer is not None: + tracer.chunk_count += 1 await ch_out.drain(context) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py index 380b8a68e0ca..8854065f5231 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -33,6 +33,7 @@ from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import StreamingNodeTracer @define_py_node() @@ -45,6 +46,7 @@ async def concatenate_node( *, output_count: int, collective_id: int, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Concatenate node for rapidsmpf. @@ -75,8 +77,12 @@ async def concatenate_node( The expected global number of output chunks. collective_id Pre-allocated collective ID for this operation. + node_tracer + Node tracer for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_in, ch_out): + async with shutdown_on_error( + context, ch_in, ch_out, node_tracer=node_tracer + ) as tracer: # Receive metadata. input_metadata = await recv_metadata(ch_in, context) nranks = context.comm().nranks @@ -132,6 +138,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) + if tracer is not None and output_duplicated: + tracer.set_duplicated() allgather = AllGatherManager(context, collective_id) stream = context.get_stream_from_pool() @@ -144,6 +152,8 @@ async def concatenate_node( # Extract concatenated result result_table = await allgather.extract_concatenated(stream) + if tracer is not None: + tracer.add_chunk(table=result_table) # If no chunks were gathered, result_table has 0 columns. # We need to create an empty table with the correct schema. @@ -164,6 +174,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) + if tracer is not None and output_duplicated: + tracer.set_duplicated() # Local repartitioning seq_num = 0 @@ -201,6 +213,8 @@ async def concatenate_node( ), context=ir_context, ) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -243,6 +257,7 @@ def _( collective_id = rec.state["collective_id_map"][ir][0] # Add python node + tracer = rec.state["tracer"] nodes[ir] = [ concatenate_node( rec.state["context"], @@ -252,6 +267,7 @@ def _( channels[ir.children[0]].reserve_output_slot(), output_count=partition_info[ir].count, collective_id=collective_id, + node_tracer=tracer.get_or_create(ir) if tracer else None, ) ] return nodes, channels diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py new file mode 100644 index 000000000000..59b5e698f346 --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Tracing infrastructure for the RapidsMPF streaming runtime.""" + +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import pylibcudf as plc + + from cudf_polars.dsl.ir import IR + + +def _stable_ir_id(ir_node: IR) -> int: + """ + Compute a stable identifier for an IR node. + + Uses MD5 hash of the node's hashable representation for determinism + across process boundaries (Python's hash() uses PYTHONHASHSEED). + + Parameters + ---------- + ir_node + The IR node. + + Returns + ------- + int + A stable 32-bit identifier for this node. + """ + content = repr(ir_node.get_hashable()).encode("utf-8") + return int(hashlib.md5(content).hexdigest()[:8], 16) + + +class StreamingNodeTracer: + """ + Tracer for a single streaming IR node. + + Collects execution statistics and emits structured log events. + + Attributes + ---------- + ir_id + Stable identifier for the IR node (for tracing/logging). + ir_type + Type name of the IR node (e.g., "Sort", "Join"). + row_count + Total row count produced by this node during execution. + None if row counting is not available for this node. + chunk_count + Total chunk count produced by this node during execution. + decision + The algorithm decision made at runtime for this node + (e.g., "broadcast_left", "shuffle", "tree", etc.). + duplicated + Whether the output rows are duplicated across ranks + (e.g., after an allgather). Affects how rows are merged. + """ + + __slots__ = ( + "chunk_count", + "decision", + "duplicated", + "ir_id", + "ir_type", + "row_count", + ) + + def __init__(self, ir_id: int | None = None, ir_type: str | None = None) -> None: + self.ir_id = ir_id + self.ir_type = ir_type + self.row_count: int | None = None + self.chunk_count: int = 0 + self.decision: str | None = None + self.duplicated: bool = False + + def add_chunk(self, *, table: plc.Table | None = None) -> None: + """ + Record a chunk. + + If table is provided, both row_count and chunk_count are updated. + If table is None, only chunk_count is incremented. + """ + if table is not None: # pragma: no cover; Covered by rapidsmpf tests + self.row_count = (self.row_count or 0) + table.num_rows() + self.chunk_count += 1 + + def set_duplicated(self, *, duplicated: bool = True) -> None: + """ + Mark output rows as duplicated across ranks. + + Call this after sending metadata when the output is duplicated + (e.g., after an allgather). Affects how rows are merged across ranks. + """ + self.duplicated = duplicated + + def merge(self, other: StreamingNodeTracer) -> None: + """Merge another node tracer's stats into this one.""" + assert self.duplicated == other.duplicated, ( + "Nodes should have the same duplicated status" + ) + if self.duplicated: + assert self.row_count == other.row_count, ( + "Duplicated nodes should have the same row count" + ) + assert self.decision == other.decision, ( + "Duplicated nodes should have the same decision" + ) + else: + self.chunk_count += other.chunk_count + if other.row_count is not None: + self.row_count = (self.row_count or 0) + other.row_count + if other.decision is not None: + self.decision = other.decision + + +class StreamingQueryTracer: + """ + Tracer for collecting runtime statistics for an entire streaming query. + + Attributes + ---------- + node_tracers + Mapping from each IR node to its node tracer. + """ + + __slots__ = ("node_tracers",) + node_tracers: dict[IR, StreamingNodeTracer] + + def __init__(self) -> None: + self.node_tracers = {} + + def get_or_create(self, ir: IR) -> StreamingNodeTracer: + """ + Get or create a node tracer for the given IR. + + Use this when setting up tracing for a node. To check if a node + was traced without creating an entry, use `node_tracers.get(ir)`. + """ + if ir not in self.node_tracers: + ir_id = _stable_ir_id(ir) + ir_type = type(ir).__name__ + self.node_tracers[ir] = StreamingNodeTracer(ir_id, ir_type) + return self.node_tracers[ir] + + def merge(self, other: StreamingQueryTracer) -> None: + """Merge another query tracer's statistics into this one.""" + for ir, node_tracer in other.node_tracers.items(): + self.get_or_create(ir).merge(node_tracer) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 3a8d0522be53..8868b1b64472 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -10,6 +10,14 @@ from functools import reduce from typing import TYPE_CHECKING, Any +from cudf_polars.dsl.tracing import LOG_TRACES + +try: + import structlog + import structlog.contextvars +except ImportError: + pass + from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.cudf.channel_metadata import ( ChannelMetadata, @@ -34,29 +42,65 @@ from cudf_polars.dsl.ir import IR from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import StreamingNodeTracer from cudf_polars.typing import DataType @asynccontextmanager async def shutdown_on_error( - context: Context, *channels: Channel[Any] -) -> AsyncIterator[None]: + context: Context, + *channels: Channel[Any], + node_tracer: StreamingNodeTracer | None = None, +) -> AsyncIterator[StreamingNodeTracer | None]: """ Shutdown on error for rapidsmpf. + This context manager handles channel cleanup on errors and optionally + manages node tracing with structlog integration. + Parameters ---------- context The rapidsmpf context. channels - The channels to shutdown. + The channels to shutdown on error. + node_tracer + Optional node tracer for collecting runtime statistics. + If provided and tracing is enabled, structlog events are emitted. + + Yields + ------ + StreamingNodeTracer | None + The node tracer (if provided) for use within the context. """ - # TODO: This probably belongs in rapidsmpf. + # Setup tracing if enabled and tracer has ir_id + ir_id: int | None = None + if node_tracer is not None and node_tracer.ir_id is not None and LOG_TRACES: + ir_id = node_tracer.ir_id + structlog.contextvars.bind_contextvars(ir_id=ir_id) + try: - yield + yield node_tracer except BaseException: await asyncio.gather(*(ch.shutdown(context) for ch in channels)) raise + finally: + # Emit structlog event on exit if tracing is enabled + if ir_id is not None and LOG_TRACES: + assert node_tracer is not None # ir_id implies node_tracer exists + log = structlog.get_logger() + record: dict[str, Any] = { + "ir_id": ir_id, + "ir_type": node_tracer.ir_type, + "chunks": node_tracer.chunk_count, + "duplicated": node_tracer.duplicated, + } + if node_tracer.row_count is not None: + record["rows"] = node_tracer.row_count + if node_tracer.decision is not None: + record["decision"] = node_tracer.decision + log.info("Streaming Node", **record) + structlog.contextvars.unbind_contextvars("ir_id") def remap_partitioning( @@ -92,27 +136,19 @@ def remap_partitioning( new_name_to_idx = {name: i for i, name in enumerate(new_schema.keys())} def remap_hash_scheme(hs: HashScheme | None | str) -> HashScheme | None | str: - if hs is None or isinstance(hs, str): + if isinstance(hs, HashScheme): + try: + new_indices = tuple( + new_name_to_idx[old_names[i]] for i in hs.column_indices + ) + except (IndexError, KeyError): + return None # Column missing in old or new schema + return HashScheme(new_indices, hs.modulus) + else: return hs # None or "inherit" passes through unchanged - try: - new_indices = tuple( - new_name_to_idx[old_names[i]] for i in hs.column_indices - ) - except (IndexError, KeyError): - return None # Column missing in old or new schema - return HashScheme(new_indices, hs.modulus) new_inter_rank = remap_hash_scheme(partitioning.inter_rank) new_local = remap_hash_scheme(partitioning.local) - - # If inter_rank partitioning was invalidated, the whole partitioning is invalid - if isinstance(partitioning.inter_rank, HashScheme) and new_inter_rank is None: - return None - - # If only local partitioning was invalidated, we can still use inter_rank - if isinstance(partitioning.local, HashScheme) and new_local is None: - new_local = None - return Partitioning(inter_rank=new_inter_rank, local=new_local) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 31bbec83fc2a..4ff584cc0499 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -55,6 +55,7 @@ "StatsPlanningOptions", "StreamingExecutor", "StreamingFallbackMode", + "TracingOptions", ] @@ -493,6 +494,48 @@ def __post_init__(self) -> None: # noqa: D105 raise ValueError("sample_chunk_count must be at least 1") +@dataclasses.dataclass(frozen=True) +class TracingOptions: + """ + Configuration for coarse-grained streaming-node tracing (rapidsmpf only). + + This class controls tracing at the *streaming-node* level, collecting + aggregate metrics (row counts, chunk counts, algorithm decisions) for + each IR node processed by the rapidsmpf runtime. When ``output_path`` + is set, a summary is written after query execution. + + For fine-grained IR-execution tracing (timing, memory, dataframe shapes), + use the ``CUDF_POLARS_LOG_TRACES`` environment variable instead. See the + `Tracing section of the usage guide + `_ + for details on available environment variables. + + To also emit structlog events for each streaming node, set + ``CUDF_POLARS_LOG_TRACES=1``. + + Parameters + ---------- + output_path + Path to write the trace summary. The output format is similar to + :func:`~cudf_polars.experimental.explain.explain_query`, annotated + with actual row counts and algorithm decisions. + If ``None`` (the default), tracing data is collected in memory + but not written to a file. + + Notes + ----- + This option only applies to the ``"rapidsmpf"`` runtime. + """ + + output_path: str | None = None + + def __post_init__(self) -> None: # noqa: D105 + if self.output_path is not None and ( + not isinstance(self.output_path, str) or not self.output_path + ): + raise TypeError("output_path must be a non-empty str or None") + + @dataclasses.dataclass(frozen=True, eq=True) class MemoryResourceConfig: """ @@ -704,6 +747,18 @@ class StreamingExecutor: or use regular pageable host memory. Pinned host memory offers higher bandwidth and lower latency for device to host transfers compared to regular pageable host memory. + tracing + Options controlling query tracing. When set to a + :class:`~cudf_polars.utils.config.TracingOptions` instance, + per-node metrics (such as row counts and algorithm decisions) + are collected during execution and written to the specified + output file. When ``None`` (the default), tracing is disabled. + + To also emit structlog events for each streaming node, set the + environment variable ``CUDF_POLARS_LOG_TRACES=1``. + + .. note:: + This feature is only available for the "rapidsmpf" runtime. Notes ----- @@ -812,6 +867,7 @@ class StreamingExecutor: f"{_env_prefix}__SPILL_TO_PINNED_MEMORY", bool, default=False ) ) + tracing: TracingOptions | None = None def __post_init__(self) -> None: # noqa: D105 # Check for rapidsmpf runtime @@ -929,6 +985,15 @@ def __post_init__(self) -> None: # noqa: D105 DynamicPlanningOptions(**self.dynamic_planning), ) + # Handle tracing. + # Can be None, dict, or TracingOptions + if isinstance(self.tracing, dict): + object.__setattr__( + self, + "tracing", + TracingOptions(**self.tracing), + ) + if self.cluster == "distributed": if self.sink_to_directory is False: raise ValueError( diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py new file mode 100644 index 000000000000..2169427d1735 --- /dev/null +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Integration tests for runtime tracing with rapidsmpf.""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap + +import pytest + +import polars as pl + +from cudf_polars.testing.asserts import DEFAULT_CLUSTER, DEFAULT_RUNTIME +from cudf_polars.testing.io import make_partitioned_source + + +def get_engine(output_path: str, parquet_options: dict | None = None) -> pl.GPUEngine: + return pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + executor_options={ + "cluster": DEFAULT_CLUSTER, + "runtime": DEFAULT_RUNTIME, + "tracing": {"output_path": str(output_path)}, + "max_rows_per_partition": 10, + "target_partition_size": 500, + }, + parquet_options=parquet_options, + ) + + +@pytest.fixture +def df(): + return pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) + + +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_tracing_basic_query(tmp_path, df): + """Test tracing output with a DataFrameScan query.""" + output_path = tmp_path / "dataframe_scan_trace.txt" + engine = get_engine(output_path) + q = df.lazy().filter(pl.col("x") > 50).group_by("y").agg(pl.col("x").sum()) + q.collect(engine=engine) + content = output_path.read_text() + assert "GROUPBY ('y',) ('y', 'x') rows=2 chunks=1" in content + assert "REPARTITION ('y', 'x') rows=10 chunks=1" in content + assert "FILTER ('x', 'y') rows=49 chunks=10" in content + assert "DATAFRAMESCAN ('x', 'y') rows=100 chunks=10" in content + + +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_tracing_scan_parquet_python(tmp_path, df): + """Test tracing output with a ScanParquet query.""" + output_path = tmp_path / "scan_parquet_trace.txt" + engine = get_engine(output_path, parquet_options={"use_rapidsmpf_native": False}) + pq_path = tmp_path / "pq_python" + pq_path.mkdir() + make_partitioned_source(df, pq_path, "parquet", n_files=5) + q = pl.scan_parquet(pq_path).filter(pl.col("x") > 50).select(["x", "y"]) + q.collect(engine=engine) + content = output_path.read_text() + assert "SCAN PARQUET ('x', 'y') rows=49 chunks=5" in content + + +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_tracing_scan_parquet_native(tmp_path, df): + """Test tracing output with a ScanParquet query.""" + output_path = tmp_path / "scan_parquet_trace.txt" + engine = get_engine(output_path, parquet_options={"use_rapidsmpf_native": True}) + pq_path = tmp_path / "pq_native" + pq_path.mkdir() + make_partitioned_source(df, pq_path, "parquet", n_files=5) + q = pl.scan_parquet(pq_path).filter(pl.col("x") > 50).select(["x", "y"]) + q.collect(engine=engine) + content = output_path.read_text() + # We can count the chunks but not the rows for "native" parquet + # (row count is omitted when unavailable) + assert "SCAN PARQUET ('x', 'y') chunks=5" in content + + +@pytest.mark.skipif( + DEFAULT_RUNTIME != "rapidsmpf", reason="Requires 'rapidsmpf' runtime." +) +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_structlog_streaming_node_events(): + """Test that structlog emits 'Streaming Node' events when tracing is enabled.""" + # Run in subprocess to control CUDF_POLARS_LOG_TRACES environment variable + code = textwrap.dedent("""\ + import polars as pl + import rmm + + df = pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) + q = df.lazy().filter(pl.col("x") > 50).group_by("y").agg(pl.col("x").sum()) + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + executor_options={ + "cluster": "single", + "runtime": "rapidsmpf", + "max_rows_per_partition": 10, + }, + memory_resource=rmm.mr.ManagedMemoryResource(), + ) + q.collect(engine=engine) + """) + + env = { + "CUDF_POLARS__EXECUTOR": "streaming", + "CUDF_POLARS_LOG_TRACES": "1", + } + + result = subprocess.check_output([sys.executable, "-c", code], env=env) + + # Check for Streaming Node events emitted by shutdown_on_error + assert b"Streaming Node" in result + assert b"ir_id=" in result + assert b"ir_type=" in result + assert b"chunks=" in result diff --git a/python/cudf_polars/tests/experimental/test_tracing.py b/python/cudf_polars/tests/experimental/test_tracing.py new file mode 100644 index 000000000000..3d884c091e7b --- /dev/null +++ b/python/cudf_polars/tests/experimental/test_tracing.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for StreamingQueryTracer and related output functions.""" + +from __future__ import annotations + +import polars as pl + +from cudf_polars.experimental.base import PartitionInfo +from cudf_polars.experimental.explain import _repr_trace_tree, write_query_trace +from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer + + +def test_streaming_query_tracer_and_output(tmp_path): + """Test StreamingQueryTracer, merge, and output formatting.""" + + class MockIR: + children = () + + def __init__(self): + self.schema = {"x": pl.Int64} + + def get_hashable(self): + return (type(self), tuple(self.schema.items())) + + ir1, ir2 = MockIR(), MockIR() + + # Test node tracer accumulation + tracer1 = StreamingQueryTracer() + nt1 = tracer1.get_or_create(ir1) + nt1.row_count = 100 + nt1.chunk_count = 5 + nt1.decision = "shuffle" + + tracer2 = StreamingQueryTracer() + tracer2.get_or_create(ir1).row_count = 150 + tracer2.get_or_create(ir1).chunk_count = 3 + tracer2.get_or_create(ir1).decision = "shuffle" + tracer2.get_or_create(ir2).row_count = 200 + tracer2.get_or_create(ir2).chunk_count = 3 + tracer2.get_or_create(ir2).add_chunk() + + # Test merge + tracer1.merge(tracer2) + assert tracer1.node_tracers[ir1].row_count == 250 + assert tracer1.node_tracers[ir1].chunk_count == 8 + + # Test _repr_trace_tree output format + partition_info = {ir1: PartitionInfo(count=4)} + output = _repr_trace_tree(ir1, partition_info, tracer1) + assert "rows=250" in output + assert "chunks=8" in output + assert "decision=shuffle" in output + + # Test write_query_trace + output_path = tmp_path / "trace.txt" + write_query_trace(output_path, ir1, partition_info, tracer1) + assert output_path.exists() + content = output_path.read_text() + assert "rows=250" in content diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 8960633c3d23..c3e3f18f55d1 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -29,6 +29,7 @@ CUDAStreamPoolConfig, ConfigOptions, MemoryResourceConfig, + TracingOptions, ) from cudf_polars.utils.cuda_stream import get_cuda_stream, get_new_cuda_stream @@ -953,6 +954,39 @@ def test_dynamic_planning_from_instance() -> None: assert config.executor.dynamic_planning.sample_chunk_count == 2 # default +def test_tracing_options() -> None: + # Tracing is disabled (None) by default + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.name == "streaming" + assert config.executor.tracing is None + + # Can enable via dict + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"tracing": {"output_path": "/tmp/trace.txt"}}, + ) + ) + assert config.executor.name == "streaming" + assert config.executor.tracing is not None + assert config.executor.tracing.output_path == "/tmp/trace.txt" + + # Can enable via instance (output_path is optional) + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"tracing": TracingOptions()}, + ) + ) + assert config.executor.name == "streaming" + assert config.executor.tracing is not None + assert config.executor.tracing.output_path is None + + # Empty output_path is rejected + with pytest.raises(TypeError, match="output_path must be a non-empty str"): + TracingOptions(output_path="") + + def test_parse_memory_resource_config() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine(