From 64e6cb3ac9dc478211b9bbbae54002f7d05822dc Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 06:21:27 -0800 Subject: [PATCH 01/25] add ProfilingOptions --- docs/cudf/source/cudf_polars/api.md | 1 + .../cudf_polars/cudf_polars/utils/config.py | 50 +++++++++++++++++++ python/cudf_polars/tests/test_config.py | 49 ++++++++++++++++++ 3 files changed, 100 insertions(+) diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index 52864f6065b2..4a779c941bd6 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, + ProfilingOptions, Cluster, ShuffleMethod, ShufflerInsertionMethod, diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 31bbec83fc2a..6f5d5bf096d9 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -48,6 +48,7 @@ "DynamicPlanningOptions", "InMemoryExecutor", "ParquetOptions", + "ProfilingOptions", "Runtime", "Scheduler", # Deprecated, kept for backward compatibility "ShuffleMethod", @@ -493,6 +494,36 @@ def __post_init__(self) -> None: # noqa: D105 raise ValueError("sample_chunk_count must be at least 1") +@dataclasses.dataclass(frozen=True) +class ProfilingOptions: + """ + Configuration for query profiling. + + When enabled, the streaming executor collects per-node metrics + (such as row counts) and writes them to a file after execution. + This feature is only available for the "rapidsmpf" runtime. + + To enable profiling, pass a ``ProfilingOptions`` instance + to ``StreamingExecutor(profiling=...)``. To disable it, pass + ``None`` (the default). + + Parameters + ---------- + output_file + Path to write the profiling results. The file will contain + a JSON representation of per-node metrics collected during + query execution. Required when profiling is enabled. + """ + + output_file: str + + def __post_init__(self) -> None: # noqa: D105 + if not isinstance(self.output_file, str): + raise TypeError("output_file must be a str") + if not self.output_file: + raise ValueError("output_file must not be empty") + + @dataclasses.dataclass(frozen=True, eq=True) class MemoryResourceConfig: """ @@ -704,6 +735,15 @@ 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. + profiling + Options controlling query profiling. When set to a + :class:`~cudf_polars.utils.config.ProfilingOptions` instance, + per-node metrics (such as row counts) are collected during execution + and written to the specified output file. When ``None`` (the default), + profiling is disabled. + + .. note:: + This feature is only available for the "rapidsmpf" runtime. Notes ----- @@ -812,6 +852,7 @@ class StreamingExecutor: f"{_env_prefix}__SPILL_TO_PINNED_MEMORY", bool, default=False ) ) + profiling: ProfilingOptions | None = None def __post_init__(self) -> None: # noqa: D105 # Check for rapidsmpf runtime @@ -929,6 +970,15 @@ def __post_init__(self) -> None: # noqa: D105 DynamicPlanningOptions(**self.dynamic_planning), ) + # Handle profiling. + # Can be None, dict, or ProfilingOptions + if isinstance(self.profiling, dict): + object.__setattr__( + self, + "profiling", + ProfilingOptions(**self.profiling), + ) + if self.cluster == "distributed": if self.sink_to_directory is False: raise ValueError( diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 8960633c3d23..39b697abe62d 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -953,6 +953,55 @@ def test_dynamic_planning_from_instance() -> None: assert config.executor.dynamic_planning.sample_chunk_count == 2 # default +def test_profiling_defaults() -> None: + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.name == "streaming" + # Profiling is disabled (None) by default + assert config.executor.profiling is None + + +def test_profiling_from_dict() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"profiling": {"output_file": "/tmp/profile.json"}}, + ) + ) + assert config.executor.name == "streaming" + assert config.executor.profiling is not None + assert config.executor.profiling.output_file == "/tmp/profile.json" + + +def test_profiling_from_instance() -> None: + from cudf_polars.utils.config import ProfilingOptions + + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "profiling": ProfilingOptions(output_file="/tmp/out.json") + }, + ) + ) + assert config.executor.name == "streaming" + assert config.executor.profiling is not None + assert config.executor.profiling.output_file == "/tmp/out.json" + + +def test_profiling_output_file_required() -> None: + from cudf_polars.utils.config import ProfilingOptions + + with pytest.raises(TypeError): + ProfilingOptions() # type: ignore[call-arg] + + +def test_profiling_output_file_not_empty() -> None: + from cudf_polars.utils.config import ProfilingOptions + + with pytest.raises(ValueError, match="output_file must not be empty"): + ProfilingOptions(output_file="") + + def test_parse_memory_resource_config() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( From 27cdf6e22efa3a3cbfe94b2910b0d62f87f1ddc2 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 07:11:10 -0800 Subject: [PATCH 02/25] revise --- .../cudf_polars/cudf_polars/utils/config.py | 20 +++++----- python/cudf_polars/tests/test_config.py | 40 ++++++------------- 2 files changed, 24 insertions(+), 36 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 6f5d5bf096d9..f751364bed27 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -509,19 +509,21 @@ class ProfilingOptions: Parameters ---------- - output_file - Path to write the profiling results. The file will contain - a JSON representation of per-node metrics collected during - query execution. Required when profiling is enabled. + output_path + Path to write the profiling results. The output will be in a + human-readable text format similar to + :func:`~cudf_polars.experimental.explain.explain_query`. + If ``None`` (the default), profiling data is collected but not + written to a file. """ - output_file: str + output_path: str | None = None def __post_init__(self) -> None: # noqa: D105 - if not isinstance(self.output_file, str): - raise TypeError("output_file must be a str") - if not self.output_file: - raise ValueError("output_file must not be empty") + 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) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 39b697abe62d..82098476ebc7 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -953,53 +953,39 @@ def test_dynamic_planning_from_instance() -> None: assert config.executor.dynamic_planning.sample_chunk_count == 2 # default -def test_profiling_defaults() -> None: +def test_profiling_options() -> None: + from cudf_polars.utils.config import ProfilingOptions + + # Profiling is disabled (None) by default config = ConfigOptions.from_polars_engine(pl.GPUEngine()) assert config.executor.name == "streaming" - # Profiling is disabled (None) by default assert config.executor.profiling is None - -def test_profiling_from_dict() -> None: + # Can enable via dict config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"profiling": {"output_file": "/tmp/profile.json"}}, + executor_options={"profiling": {"output_path": "/tmp/profile.txt"}}, ) ) assert config.executor.name == "streaming" assert config.executor.profiling is not None - assert config.executor.profiling.output_file == "/tmp/profile.json" - - -def test_profiling_from_instance() -> None: - from cudf_polars.utils.config import ProfilingOptions + assert config.executor.profiling.output_path == "/tmp/profile.txt" + # Can enable via instance (output_path is optional) config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={ - "profiling": ProfilingOptions(output_file="/tmp/out.json") - }, + executor_options={"profiling": ProfilingOptions()}, ) ) assert config.executor.name == "streaming" assert config.executor.profiling is not None - assert config.executor.profiling.output_file == "/tmp/out.json" - - -def test_profiling_output_file_required() -> None: - from cudf_polars.utils.config import ProfilingOptions - - with pytest.raises(TypeError): - ProfilingOptions() # type: ignore[call-arg] - - -def test_profiling_output_file_not_empty() -> None: - from cudf_polars.utils.config import ProfilingOptions + assert config.executor.profiling.output_path is None - with pytest.raises(ValueError, match="output_file must not be empty"): - ProfilingOptions(output_file="") + # Empty output_path is rejected + with pytest.raises(TypeError, match="output_path must be a non-empty str"): + ProfilingOptions(output_path="") def test_parse_memory_resource_config() -> None: From 49c498c2a247ff2a891a19cb6aef5bdcb13c1454 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 09:38:21 -0800 Subject: [PATCH 03/25] cleanup and tests --- .../cudf_polars/experimental/base.py | 36 +++++++- .../experimental/benchmarks/utils.py | 13 +++ .../cudf_polars/experimental/explain.py | 74 ++++++++++++++++- .../experimental/rapidsmpf/core.py | 30 +++++-- .../experimental/rapidsmpf/dask.py | 28 +++++-- .../experimental/rapidsmpf/dispatch.py | 4 + .../experimental/rapidsmpf/nodes.py | 23 ++++++ .../experimental/test_runtime_profiler.py | 82 +++++++++++++++++++ 8 files changed, 273 insertions(+), 17 deletions(-) create mode 100644 python/cudf_polars/tests/experimental/test_runtime_profiler.py diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py index 35490c8ce7b9..6a8e24aeec63 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/experimental/base.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Multi-partition base classes.""" @@ -406,6 +406,40 @@ def __init__(self) -> None: self.join_info = JoinInfo() +class RuntimeProfiler: + """ + Profiler for collecting runtime statistics during execution. + + Attributes + ---------- + row_count + Mapping from IR node to actual row count produced during execution. + chunk_count + Mapping from IR node to actual chunk count produced during execution. + decisions + Mapping from IR node to the algorithm decision made at runtime + (e.g., "broadcast_left", "shuffle", "tree", etc.). + """ + + __slots__ = ("chunk_count", "decisions", "row_count") + row_count: defaultdict[IR, int] + chunk_count: defaultdict[IR, int] + decisions: dict[IR, str] + + def __init__(self) -> None: + self.row_count = defaultdict(int) + self.chunk_count = defaultdict(int) + self.decisions = {} + + def merge(self, other: RuntimeProfiler) -> None: + """Merge another profiler's statistics into this one.""" + for ir, n_rows in other.row_count.items(): + self.row_count[ir] += n_rows + for ir, n_chunks in other.chunk_count.items(): + self.chunk_count[ir] += n_chunks + self.decisions.update(other.decisions) + + class IOPartitionFlavor(IntEnum): """Flavor of IO partitioning.""" diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 9d59650b9688..61b1d0881c1e 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 + profile_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, + profile_output_path=args.profile_output_path, max_io_threads=args.max_io_threads, native_parquet=args.native_parquet, extra_info=args.extra_info, @@ -477,6 +479,10 @@ 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.profile_output_path is not None: + executor_options["profiling"] = { + "output_path": run_config.profile_output_path + } if ( benchmark @@ -974,6 +980,13 @@ def parse_args( default=False, help="Enable dynamic shuffle planning (not yet implemented). ", ) + parser.add_argument( + "--profile-output-path", + dest="profile_output_path", + type=str, + default=None, + help="Path to write profiling 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..ff5d59b4bd56 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,16 @@ 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.base import ( + PartitionInfo, + RuntimeProfiler, + StatsCollector, + ) def explain_query( @@ -168,3 +173,68 @@ 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_profile_output( + profile_output: str | Path, + ir: IR, + partition_info: MutableMapping[IR, PartitionInfo], + profiler: RuntimeProfiler, +) -> None: + """ + Write a post-execution profile showing actual row counts and decisions. + + Parameters + ---------- + profile_output + Path to write the profile file. + ir + The lowered IR root node. + partition_info + Partition information for the IR nodes. + profiler + The profiler with actual row counts and decisions from execution. + """ + from pathlib import Path + + profile_repr = _repr_profile_tree(ir, partition_info, profiler) + Path(profile_output).write_text(profile_repr) + + +def _repr_profile_tree( + ir: IR, + partition_info: MutableMapping[IR, PartitionInfo], + profiler: RuntimeProfiler, + *, + offset: str = "", +) -> str: + """Recursively build a tree representation with profiler data.""" + header = _repr_ir(ir, offset=offset) + + # Add actual row count + actual_rows = profiler.row_count.get(ir) + actual_str = _fmt_row_count(actual_rows) if actual_rows is not None else "?" + header = header.rstrip("\n") + f" rows={actual_str}" + + # Add decision if present + if ir in profiler.decisions: + header += f" decision={profiler.decisions[ir]}" + + # Add chunk count if available + actual_chunks = profiler.chunk_count.get(ir) + if actual_chunks is not None: + header += f" chunks={actual_chunks}" + + header += "\n" + + children_strs = [ + _repr_profile_tree(child, partition_info, profiler, offset=offset + " ") + for child in ir.children + ] + + 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 f308d3c11fbb..8863e59422df 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -33,6 +33,7 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import DataFrameScan, IRExecutionContext, Join, Scan, Union from cudf_polars.dsl.traversal import CachingVisitor, traversal +from cudf_polars.experimental.base import RuntimeProfiler from cudf_polars.experimental.rapidsmpf.collectives import ReserveOpIDs from cudf_polars.experimental.rapidsmpf.dispatch import FanoutInfo, lower_ir_node from cudf_polars.experimental.rapidsmpf.nodes import ( @@ -107,7 +108,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, profiler = evaluate_pipeline_dask( evaluate_pipeline, ir, partition_info, @@ -118,7 +119,7 @@ def evaluate_logical_plan( ) else: # Single-process execution: Run locally - result, metadata_collector = evaluate_pipeline( + result, metadata_collector, profiler = evaluate_pipeline( ir, partition_info, config_options, @@ -127,6 +128,17 @@ def evaluate_logical_plan( collect_metadata=collect_metadata, ) + # Write profiler output if configured + profiling = config_options.executor.profiling + if ( + profiling is not None + and profiling.output_path is not None + and profiler is not None + ): + from cudf_polars.experimental.explain import write_profile_output + + write_profile_output(profiling.output_path, ir, partition_info, profiler) + return result, metadata_collector @@ -139,7 +151,7 @@ def evaluate_pipeline( rmpf_context: Context | None = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None]: +) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: """ Build and evaluate a RapidsMPF streaming pipeline. @@ -162,7 +174,7 @@ def evaluate_pipeline( Returns ------- - The output DataFrame and metadata collector. + The output DataFrame, metadata collector, and profiler. """ assert config_options.executor.name == "streaming", "Executor must be streaming" assert config_options.executor.runtime == "rapidsmpf", "Runtime must be rapidsmpf" @@ -230,6 +242,9 @@ def evaluate_pipeline( # Generate network nodes assert rmpf_context is not None, "RapidsMPF context must defined." metadata_collector: list[Metadata] | None = [] if collect_metadata else None + profiler: RuntimeProfiler | None = ( + RuntimeProfiler() if config_options.executor.profiling is not None else None + ) nodes, output = generate_network( rmpf_context, ir, @@ -239,6 +254,7 @@ def evaluate_pipeline( ir_context=ir_context, collective_id_map=collective_id_map, metadata_collector=metadata_collector, + profiler=profiler, ) # Run the network @@ -292,7 +308,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, profiler def lower_ir_graph( @@ -414,6 +430,7 @@ def generate_network( ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], metadata_collector: list[Metadata] | None, + profiler: RuntimeProfiler | None = None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. @@ -438,6 +455,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. + profiler + Profiler for collecting runtime statistics. Returns ------- @@ -470,6 +489,7 @@ def generate_network( "max_io_threads": max_io_threads_local, "stats": stats, "collective_id_map": collective_id_map, + "profiler": profiler, } 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 f9527d5f963a..671a7db220b3 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py @@ -21,7 +21,11 @@ from distributed import Client from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import PartitionInfo, StatsCollector + from cudf_polars.experimental.base import ( + PartitionInfo, + RuntimeProfiler, + StatsCollector, + ) from cudf_polars.experimental.parallel import ConfigOptions from cudf_polars.experimental.rapidsmpf.utils import Metadata @@ -39,8 +43,8 @@ def __call__( rmpf_context: Context | None = None, *, collect_metadata: bool = False, - ) -> tuple[pl.DataFrame, list[Metadata] | None]: - """Evaluate a pipeline and return the result DataFrame and metadata.""" + ) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: + """Evaluate a pipeline and return the result DataFrame, metadata, and profiler.""" ... @@ -61,7 +65,7 @@ def evaluate_pipeline_dask( collective_id_map: dict[IR, list[int]], *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None]: +) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: """ Evaluate a RapidsMPF streaming pipeline on a Dask cluster. @@ -84,7 +88,7 @@ def evaluate_pipeline_dask( Returns ------- - The output DataFrame and metadata collector. + The output DataFrame, metadata collector, and merged profiler. """ client = get_dask_client() result = client.run( @@ -99,12 +103,18 @@ def evaluate_pipeline_dask( ) dfs: list[pl.DataFrame] = [] metadata_collector: list[Metadata] = [] - for df, md in result.values(): + merged_profiler: RuntimeProfiler | None = None + for df, md, profiler in result.values(): dfs.append(df) if md is not None: metadata_collector.extend(md) + if profiler is not None: + if merged_profiler is None: + merged_profiler = profiler + else: + merged_profiler.merge(profiler) - return pl.concat(dfs), metadata_collector or None + return pl.concat(dfs), metadata_collector or None, merged_profiler def _evaluate_pipeline_dask( @@ -117,7 +127,7 @@ def _evaluate_pipeline_dask( dask_worker: Any = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None]: +) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: """ Build and evaluate a RapidsMPF streaming pipeline. @@ -144,7 +154,7 @@ def _evaluate_pipeline_dask( Returns ------- - The output DataFrame and metadata collector. + The output DataFrame, metadata collector, and profiler. """ 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..6fe569a96ab6 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py @@ -17,6 +17,7 @@ from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.experimental.base import ( PartitionInfo, + RuntimeProfiler, StatsCollector, ) from cudf_polars.experimental.rapidsmpf.utils import ChannelManager @@ -77,6 +78,8 @@ class GenState(TypedDict): Statistics collector. collective_id_map The mapping of IR nodes to lists of collective IDs. + profiler + Runtime profiler for collecting execution statistics. """ context: Context @@ -87,6 +90,7 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] + profiler: RuntimeProfiler | None SubNetGenerator: TypeAlias = GenericTransformer[ diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index 86b7bfde5072..91da42d89d47 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -35,6 +35,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.experimental.base import RuntimeProfiler from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator @@ -47,6 +48,7 @@ async def default_node_single( ch_in: Channel[TableChunk], *, preserve_partitioning: bool = False, + profiler: RuntimeProfiler | None = None, ) -> None: """ Single-channel default node for rapidsmpf. @@ -65,6 +67,8 @@ async def default_node_single( The input Channel[TableChunk]. preserve_partitioning Whether to preserve the partitioning metadata of the input chunks. + profiler + Profiler for collecting runtime statistics. Notes ----- @@ -82,6 +86,8 @@ async def default_node_single( await send_metadata(ch_out, context, metadata_out) # Recv/send data. + n_rows_out = 0 + n_chunks_out = 0 seq_num = 0 receiving = True received_any = False @@ -117,6 +123,8 @@ async def default_node_single( ), context=ir_context, ) + n_rows_out += df.table.num_rows() + n_chunks_out += 1 await ch_out.send( context, Message( @@ -129,6 +137,9 @@ async def default_node_single( del df, chunk await ch_out.drain(context) + if profiler is not None: + profiler.row_count[ir] += n_rows_out + profiler.chunk_count[ir] += n_chunks_out @define_py_node() @@ -140,6 +151,7 @@ async def default_node_multi( chs_in: tuple[Channel[TableChunk], ...], *, partitioning_index: int | None = None, + profiler: RuntimeProfiler | None = None, ) -> None: """ Pointwise node for rapidsmpf. @@ -159,6 +171,8 @@ async def default_node_multi( partitioning_index Index of the input channel to preserve partitioning information for. If None, no partitioning information is preserved. + profiler + Profiler for collecting runtime statistics. """ async with shutdown_on_error(context, *chs_in, ch_out): # Merge and forward basic metadata. @@ -178,6 +192,8 @@ async def default_node_multi( metadata.partitioning = md_child.partitioning await send_metadata(ch_out, context, metadata) + n_rows_out = 0 + n_chunks_out = 0 seq_num = 0 n_children = len(chs_in) finished_channels: set[int] = set() @@ -241,6 +257,8 @@ async def default_node_multi( *dfs, context=ir_context, ) + n_rows_out += df.table.num_rows() + n_chunks_out += 1 await ch_out.send( context, Message( @@ -258,6 +276,9 @@ async def default_node_multi( # Drain the output channel del ready_chunks await ch_out.drain(context) + if profiler is not None: + profiler.row_count[ir] += n_rows_out + profiler.chunk_count[ir] += n_chunks_out @define_py_node() @@ -537,6 +558,7 @@ def _( channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), preserve_partitioning=preserve_partitioning, + profiler=rec.state.get("profiler"), ) ] else: @@ -548,6 +570,7 @@ def _( rec.state["ir_context"], channels[ir].reserve_input_slot(), tuple(channels[c].reserve_output_slot() for c in ir.children), + profiler=rec.state.get("profiler"), ) ] diff --git a/python/cudf_polars/tests/experimental/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/test_runtime_profiler.py new file mode 100644 index 000000000000..da7765d67a41 --- /dev/null +++ b/python/cudf_polars/tests/experimental/test_runtime_profiler.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +import polars as pl + +from cudf_polars.experimental.base import PartitionInfo, RuntimeProfiler +from cudf_polars.experimental.explain import _repr_profile_tree, write_profile_output +from cudf_polars.testing.asserts import DEFAULT_CLUSTER, DEFAULT_RUNTIME + + +def test_runtime_profiler_and_output(tmp_path): + """Test RuntimeProfiler, merge, and output formatting.""" + + class MockIR: + children = () + + def __init__(self): + self.schema = {"x": pl.Int64} + + ir1, ir2 = MockIR(), MockIR() + + # Test accumulation and merge + profiler1 = RuntimeProfiler() + profiler1.row_count[ir1] = 100 + profiler1.chunk_count[ir1] = 5 + profiler1.decisions[ir1] = "shuffle" + + profiler2 = RuntimeProfiler() + profiler2.row_count[ir1] = 150 + profiler2.row_count[ir2] = 200 + profiler2.chunk_count[ir1] = 3 + profiler2.chunk_count[ir2] = 4 + + profiler1.merge(profiler2) + assert profiler1.row_count[ir1] == 250 + assert profiler1.chunk_count[ir1] == 8 + + # Test _repr_profile_tree output format + partition_info = {ir1: PartitionInfo(count=1)} + output = _repr_profile_tree(ir1, partition_info, profiler1) + assert "rows=250" in output + assert "chunks=8" in output + assert "decision=shuffle" in output + + # Test write_profile_output + output_path = tmp_path / "profile.txt" + write_profile_output(output_path, ir1, partition_info, profiler1) + assert output_path.exists() + content = output_path.read_text() + assert "rows=250" in content + + +@pytest.mark.skipif( + DEFAULT_RUNTIME != "rapidsmpf", reason="Requires 'rapidsmpf' runtime." +) +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_profiling_with_real_query(tmp_path): + """Test profiling output with a real query execution.""" + output_path = tmp_path / "profile.txt" + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + executor_options={ + "cluster": DEFAULT_CLUSTER, + "runtime": DEFAULT_RUNTIME, + "profiling": {"output_path": str(output_path)}, + }, + ) + + df = pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) + q = df.lazy().filter(pl.col("x") > 50).select(["x", "y"]) + q.collect(engine=engine) + + # Verify profile was written + assert output_path.exists() + content = output_path.read_text() + assert "rows=" in content + assert "chunks=" in content From 08350a52ce632bb454b9fd13d5f2f7b2a43d2356 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 11:54:13 -0800 Subject: [PATCH 04/25] revise design --- .../cudf_polars/experimental/base.py | 85 +++++++++---- .../cudf_polars/experimental/explain.py | 27 ++-- .../experimental/rapidsmpf/core.py | 12 +- .../experimental/rapidsmpf/dask.py | 10 +- .../experimental/rapidsmpf/dispatch.py | 4 +- .../cudf_polars/experimental/rapidsmpf/io.py | 27 +++- .../experimental/rapidsmpf/nodes.py | 42 +++---- .../rapidsmpf/test_runtime_profiler.py | 118 ++++++++++++++++++ .../experimental/test_runtime_profiler.py | 82 ------------ 9 files changed, 258 insertions(+), 149 deletions(-) create mode 100644 python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py delete mode 100644 python/cudf_polars/tests/experimental/test_runtime_profiler.py diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py index 6a8e24aeec63..6a325059f56e 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/experimental/base.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from collections.abc import Generator, Iterator, MutableMapping + from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR from cudf_polars.dsl.nodebase import Node @@ -406,38 +407,80 @@ def __init__(self) -> None: self.join_info = JoinInfo() -class RuntimeProfiler: +class RuntimeNodeProfiler: """ - Profiler for collecting runtime statistics during execution. + Profiler for a single IR node. Attributes ---------- row_count - Mapping from IR node to actual row count produced during execution. + Total row count produced by this node during execution. + None if row counting is not available for this node. chunk_count - Mapping from IR node to actual chunk count produced during execution. - decisions - Mapping from IR node to the algorithm decision made at runtime + 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.). """ - __slots__ = ("chunk_count", "decisions", "row_count") - row_count: defaultdict[IR, int] - chunk_count: defaultdict[IR, int] - decisions: dict[IR, str] + __slots__ = ("chunk_count", "decision", "row_count") def __init__(self) -> None: - self.row_count = defaultdict(int) - self.chunk_count = defaultdict(int) - self.decisions = {} - - def merge(self, other: RuntimeProfiler) -> None: - """Merge another profiler's statistics into this one.""" - for ir, n_rows in other.row_count.items(): - self.row_count[ir] += n_rows - for ir, n_chunks in other.chunk_count.items(): - self.chunk_count[ir] += n_chunks - self.decisions.update(other.decisions) + self.row_count: int | None = None + self.chunk_count: int = 0 + self.decision: str | None = None + + def add_chunk(self, *, df: DataFrame | None = None) -> None: + """ + Record a chunk. + + If df is provided, both row_count and chunk_count are updated. + If df is None, only chunk_count is incremented. + """ + if df is not None: + self.row_count = (self.row_count or 0) + df.table.num_rows() + self.chunk_count += 1 + + def merge(self, other: RuntimeNodeProfiler) -> None: + """Merge another node profiler's stats into this one.""" + if other.row_count is not None: + self.row_count = (self.row_count or 0) + other.row_count + self.chunk_count += other.chunk_count + if other.decision is not None: + self.decision = other.decision + + +class RuntimeQueryProfiler: + """ + Profiler for collecting runtime statistics for an entire query. + + Attributes + ---------- + node_profilers + Mapping from each IR node to its node profiler. + """ + + __slots__ = ("node_profilers",) + node_profilers: dict[IR, RuntimeNodeProfiler] + + def __init__(self) -> None: + self.node_profilers = {} + + def get_or_create(self, ir: IR) -> RuntimeNodeProfiler: + """ + Get or create a node profiler for the given IR. + + Use this when setting up profiling for a node. To check if a node + was profiled without creating an entry, use `node_profilers.get(ir)`. + """ + if ir not in self.node_profilers: + self.node_profilers[ir] = RuntimeNodeProfiler() + return self.node_profilers[ir] + + def merge(self, other: RuntimeQueryProfiler) -> None: + """Merge another query profiler's statistics into this one.""" + for ir, node_profiler in other.node_profilers.items(): + self.get_or_create(ir).merge(node_profiler) class IOPartitionFlavor(IntEnum): diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index ff5d59b4bd56..d3e808680845 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -33,7 +33,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.experimental.base import ( PartitionInfo, - RuntimeProfiler, + RuntimeQueryProfiler, StatsCollector, ) @@ -179,7 +179,7 @@ def write_profile_output( profile_output: str | Path, ir: IR, partition_info: MutableMapping[IR, PartitionInfo], - profiler: RuntimeProfiler, + profiler: RuntimeQueryProfiler, ) -> None: """ Write a post-execution profile showing actual row counts and decisions. @@ -204,26 +204,33 @@ def write_profile_output( def _repr_profile_tree( ir: IR, partition_info: MutableMapping[IR, PartitionInfo], - profiler: RuntimeProfiler, + profiler: RuntimeQueryProfiler, *, offset: str = "", ) -> str: """Recursively build a tree representation with profiler data.""" header = _repr_ir(ir, offset=offset) + static_count = partition_info[ir].count if partition_info else None + + # Get node profiler if it exists + node_profiler = profiler.node_profilers.get(ir) # Add actual row count - actual_rows = profiler.row_count.get(ir) + actual_rows = node_profiler.row_count if node_profiler is not None else None actual_str = _fmt_row_count(actual_rows) if actual_rows is not None else "?" header = header.rstrip("\n") + f" rows={actual_str}" # Add decision if present - if ir in profiler.decisions: - header += f" decision={profiler.decisions[ir]}" + if node_profiler is not None and node_profiler.decision is not None: + header += f" decision={node_profiler.decision}" + + # Add actual chunk count if available + if node_profiler is not None and node_profiler.chunk_count > 0: + header += f" chunks={node_profiler.chunk_count}" - # Add chunk count if available - actual_chunks = profiler.chunk_count.get(ir) - if actual_chunks is not None: - header += f" chunks={actual_chunks}" + # Add expected partition count from PartitionInfo + if static_count is not None: + header += f" [{static_count}]" header += "\n" diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 8863e59422df..7685bb39ac4f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -33,7 +33,7 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import DataFrameScan, IRExecutionContext, Join, Scan, Union from cudf_polars.dsl.traversal import CachingVisitor, traversal -from cudf_polars.experimental.base import RuntimeProfiler +from cudf_polars.experimental.base import RuntimeQueryProfiler from cudf_polars.experimental.rapidsmpf.collectives import ReserveOpIDs from cudf_polars.experimental.rapidsmpf.dispatch import FanoutInfo, lower_ir_node from cudf_polars.experimental.rapidsmpf.nodes import ( @@ -151,7 +151,7 @@ def evaluate_pipeline( rmpf_context: Context | None = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: +) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeQueryProfiler | None]: """ Build and evaluate a RapidsMPF streaming pipeline. @@ -242,8 +242,10 @@ def evaluate_pipeline( # Generate network nodes assert rmpf_context is not None, "RapidsMPF context must defined." metadata_collector: list[Metadata] | None = [] if collect_metadata else None - profiler: RuntimeProfiler | None = ( - RuntimeProfiler() if config_options.executor.profiling is not None else None + profiler: RuntimeQueryProfiler | None = ( + RuntimeQueryProfiler() + if config_options.executor.profiling is not None + else None ) nodes, output = generate_network( rmpf_context, @@ -430,7 +432,7 @@ def generate_network( ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], metadata_collector: list[Metadata] | None, - profiler: RuntimeProfiler | None = None, + profiler: RuntimeQueryProfiler | None = None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py index 671a7db220b3..eeb8a2560782 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py @@ -23,7 +23,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.experimental.base import ( PartitionInfo, - RuntimeProfiler, + RuntimeQueryProfiler, StatsCollector, ) from cudf_polars.experimental.parallel import ConfigOptions @@ -43,7 +43,7 @@ def __call__( rmpf_context: Context | None = None, *, collect_metadata: bool = False, - ) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: + ) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeQueryProfiler | None]: """Evaluate a pipeline and return the result DataFrame, metadata, and profiler.""" ... @@ -65,7 +65,7 @@ def evaluate_pipeline_dask( collective_id_map: dict[IR, list[int]], *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: +) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeQueryProfiler | None]: """ Evaluate a RapidsMPF streaming pipeline on a Dask cluster. @@ -103,7 +103,7 @@ def evaluate_pipeline_dask( ) dfs: list[pl.DataFrame] = [] metadata_collector: list[Metadata] = [] - merged_profiler: RuntimeProfiler | None = None + merged_profiler: RuntimeQueryProfiler | None = None for df, md, profiler in result.values(): dfs.append(df) if md is not None: @@ -127,7 +127,7 @@ def _evaluate_pipeline_dask( dask_worker: Any = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeProfiler | None]: +) -> tuple[pl.DataFrame, list[Metadata] | None, RuntimeQueryProfiler | None]: """ Build and evaluate a RapidsMPF streaming pipeline. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py index 6fe569a96ab6..9403293e1785 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py @@ -17,7 +17,7 @@ from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.experimental.base import ( PartitionInfo, - RuntimeProfiler, + RuntimeQueryProfiler, StatsCollector, ) from cudf_polars.experimental.rapidsmpf.utils import ChannelManager @@ -90,7 +90,7 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] - profiler: RuntimeProfiler | None + profiler: RuntimeQueryProfiler | 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 d1ee3fcfd70f..b5da449330c3 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -51,7 +51,12 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import ColumnStat, StatsCollector + from cudf_polars.experimental.base import ( + ColumnStat, + RuntimeNodeProfiler, + RuntimeQueryProfiler, + StatsCollector, + ) from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator from cudf_polars.experimental.rapidsmpf.dispatch import LowerIRTransformer from cudf_polars.utils.config import ParquetOptions @@ -144,6 +149,7 @@ async def dataframescan_node( num_producers: int, rows_per_partition: int, estimated_chunk_bytes: int, + node_profiler: RuntimeNodeProfiler | None = None, ) -> None: """ DataFrameScan node for rapidsmpf. @@ -165,6 +171,8 @@ 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_profiler + The node profiler for collecting runtime statistics. """ async with shutdown_on_error(context, ch_out): # Find local partition count. @@ -216,6 +224,7 @@ async def dataframescan_node( ch_out, ir_context, estimated_chunk_bytes, + node_profiler=node_profiler, ) await ch_out.drain(context) return @@ -241,6 +250,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, + node_profiler=node_profiler, ) await ch_out.drain(context) @@ -266,6 +276,7 @@ def _( context = rec.state["context"] ir_context = rec.state["ir_context"] + profiler: RuntimeQueryProfiler | None = rec.state["profiler"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} nodes: dict[IR, list[Any]] = { ir: [ @@ -277,6 +288,7 @@ def _( num_producers=num_producers, rows_per_partition=rows_per_partition, estimated_chunk_bytes=estimated_chunk_bytes, + node_profiler=profiler.get_or_create(ir) if profiler else None, ) ] } @@ -322,6 +334,7 @@ async def read_chunk( ch_out: Channel[TableChunk], ir_context: IRExecutionContext, estimated_chunk_bytes: int, + node_profiler: RuntimeNodeProfiler | None = None, ) -> None: """ Read a chunk from disk and send it to the output channel. @@ -341,6 +354,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_profiler + The node profiler for collecting runtime statistics. """ with opaque_reservation(context, estimated_chunk_bytes): df = await asyncio.to_thread( @@ -348,6 +363,8 @@ async def read_chunk( *scan._non_child_args, context=ir_context, ) + if node_profiler is not None: + node_profiler.add_chunk(df=df) await ch_out.send( context, Message( @@ -372,6 +389,7 @@ async def scan_node( plan: IOPartitionPlan, parquet_options: ParquetOptions, estimated_chunk_bytes: int, + node_profiler: RuntimeNodeProfiler | None = None, ) -> None: """ Scan node for rapidsmpf. @@ -395,6 +413,8 @@ 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_profiler + The node profiler for collecting runtime statistics. """ async with shutdown_on_error(context, ch_out): # Build a list of local Scan operations @@ -487,6 +507,7 @@ async def scan_node( ch_out, ir_context, estimated_chunk_bytes, + node_profiler=node_profiler, ) await ch_out.drain(context) return @@ -512,6 +533,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, + node_profiler=node_profiler, ) await ch_out.drain(context) @@ -642,6 +664,7 @@ def _( parquet_options = config_options.parquet_options partition_info = rec.state["partition_info"][ir] num_producers = rec.state["max_io_threads"] + profiler: RuntimeQueryProfiler | None = rec.state["profiler"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} assert partition_info.io_plan is not None, "Scan node must have a partition plan" @@ -694,6 +717,7 @@ def _( ), global_count=partition_info.count, ), + node_profiler=profiler.get_or_create(ir) if profiler else None, ) nodes[ir] = [native_node, metadata_node] else: @@ -710,6 +734,7 @@ def _( plan=plan, parquet_options=parquet_options, estimated_chunk_bytes=executor.target_partition_size, + node_profiler=profiler.get_or_create(ir) if profiler 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 91da42d89d47..2a301c1f80b3 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -35,7 +35,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.base import RuntimeProfiler + from cudf_polars.experimental.base import RuntimeNodeProfiler, RuntimeQueryProfiler from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator @@ -48,7 +48,7 @@ async def default_node_single( ch_in: Channel[TableChunk], *, preserve_partitioning: bool = False, - profiler: RuntimeProfiler | None = None, + node_profiler: RuntimeNodeProfiler | None = None, ) -> None: """ Single-channel default node for rapidsmpf. @@ -67,8 +67,8 @@ async def default_node_single( The input Channel[TableChunk]. preserve_partitioning Whether to preserve the partitioning metadata of the input chunks. - profiler - Profiler for collecting runtime statistics. + node_profiler + Node profiler for collecting runtime statistics. Notes ----- @@ -86,8 +86,6 @@ async def default_node_single( await send_metadata(ch_out, context, metadata_out) # Recv/send data. - n_rows_out = 0 - n_chunks_out = 0 seq_num = 0 receiving = True received_any = False @@ -123,8 +121,8 @@ async def default_node_single( ), context=ir_context, ) - n_rows_out += df.table.num_rows() - n_chunks_out += 1 + if node_profiler is not None: + node_profiler.add_chunk(df=df) await ch_out.send( context, Message( @@ -137,9 +135,6 @@ async def default_node_single( del df, chunk await ch_out.drain(context) - if profiler is not None: - profiler.row_count[ir] += n_rows_out - profiler.chunk_count[ir] += n_chunks_out @define_py_node() @@ -151,7 +146,7 @@ async def default_node_multi( chs_in: tuple[Channel[TableChunk], ...], *, partitioning_index: int | None = None, - profiler: RuntimeProfiler | None = None, + node_profiler: RuntimeNodeProfiler | None = None, ) -> None: """ Pointwise node for rapidsmpf. @@ -171,8 +166,8 @@ async def default_node_multi( partitioning_index Index of the input channel to preserve partitioning information for. If None, no partitioning information is preserved. - profiler - Profiler for collecting runtime statistics. + node_profiler + Node profiler for collecting runtime statistics. """ async with shutdown_on_error(context, *chs_in, ch_out): # Merge and forward basic metadata. @@ -192,8 +187,6 @@ async def default_node_multi( metadata.partitioning = md_child.partitioning await send_metadata(ch_out, context, metadata) - n_rows_out = 0 - n_chunks_out = 0 seq_num = 0 n_children = len(chs_in) finished_channels: set[int] = set() @@ -257,8 +250,8 @@ async def default_node_multi( *dfs, context=ir_context, ) - n_rows_out += df.table.num_rows() - n_chunks_out += 1 + if node_profiler is not None: + node_profiler.add_chunk(df=df) await ch_out.send( context, Message( @@ -276,9 +269,6 @@ async def default_node_multi( # Drain the output channel del ready_chunks await ch_out.drain(context) - if profiler is not None: - profiler.row_count[ir] += n_rows_out - profiler.chunk_count[ir] += n_chunks_out @define_py_node() @@ -539,6 +529,7 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + profiler: RuntimeQueryProfiler | None = rec.state.get("profiler") if len(ir.children) == 1: # Single-channel default node @@ -558,7 +549,7 @@ def _( channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), preserve_partitioning=preserve_partitioning, - profiler=rec.state.get("profiler"), + node_profiler=profiler.get_or_create(ir) if profiler else None, ) ] else: @@ -570,7 +561,7 @@ def _( rec.state["ir_context"], channels[ir].reserve_input_slot(), tuple(channels[c].reserve_output_slot() for c in ir.children), - profiler=rec.state.get("profiler"), + node_profiler=profiler.get_or_create(ir) if profiler else None, ) ] @@ -684,6 +675,7 @@ async def metadata_feeder_node( ch_in: Channel[TableChunk], ch_out: Channel[TableChunk], metadata: Metadata, + node_profiler: RuntimeNodeProfiler | None = None, ) -> None: """ Forward data with new metadata. @@ -698,11 +690,15 @@ 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_profiler + Node profiler for collecting runtime statistics. """ async with shutdown_on_error(context, ch_in, ch_out): await send_metadata(ch_out, context, metadata) while (msg := await ch_in.recv(context)) is not None: await ch_out.send(context, msg) + if node_profiler is not None: + node_profiler.chunk_count += 1 await ch_out.drain(context) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py new file mode 100644 index 000000000000..ba951e5d0394 --- /dev/null +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +import polars as pl + +from cudf_polars.experimental.base import PartitionInfo, RuntimeQueryProfiler +from cudf_polars.experimental.explain import _repr_profile_tree, write_profile_output +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, + "profiling": {"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}) + + +def test_runtime_profiler_and_output(tmp_path): + """Test RuntimeQueryProfiler, merge, and output formatting.""" + + class MockIR: + children = () + + def __init__(self): + self.schema = {"x": pl.Int64} + + ir1, ir2 = MockIR(), MockIR() + + # Test node profiler accumulation + profiler1 = RuntimeQueryProfiler() + np1 = profiler1.get_or_create(ir1) + np1.row_count = 100 + np1.chunk_count = 5 + np1.decision = "shuffle" + + profiler2 = RuntimeQueryProfiler() + profiler2.get_or_create(ir1).row_count = 150 + profiler2.get_or_create(ir1).chunk_count = 3 + profiler2.get_or_create(ir2).row_count = 200 + profiler2.get_or_create(ir2).chunk_count = 4 + + # Test merge + profiler1.merge(profiler2) + assert profiler1.node_profilers[ir1].row_count == 250 + assert profiler1.node_profilers[ir1].chunk_count == 8 + + # Test _repr_profile_tree output format + partition_info = {ir1: PartitionInfo(count=4)} + output = _repr_profile_tree(ir1, partition_info, profiler1) + assert "rows=250" in output + assert "chunks=8" in output + assert "decision=shuffle" in output + assert "[4]" in output # static partition count + + # Test write_profile_output + output_path = tmp_path / "profile.txt" + write_profile_output(output_path, ir1, partition_info, profiler1) + assert output_path.exists() + content = output_path.read_text() + assert "rows=250" in content + + +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_profiling_dataframe_scan(tmp_path, df): + """Test profiling output with a DataFrameScan query.""" + output_path = tmp_path / "dataframe_scan_profile.txt" + engine = get_engine(output_path) + q = df.lazy().filter(pl.col("x") > 50).select(["x", "y"]) + q.collect(engine=engine) + content = output_path.read_text() + assert "FILTER ('x', 'y') rows=49 chunks=10 [10]" in content + assert "DATAFRAMESCAN ('x', 'y') rows=100 chunks=10 [10]" in content + + +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_profiling_scan_parquet_python(tmp_path, df): + """Test profiling output with a ScanParquet query.""" + output_path = tmp_path / "scan_parquet_profile.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 [5]" in content + + +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_profiling_scan_parquet_native(tmp_path, df): + """Test profiling output with a ScanParquet query.""" + output_path = tmp_path / "scan_parquet_profile.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() + assert "SCAN PARQUET ('x', 'y') rows=? chunks=5 [5]" in content diff --git a/python/cudf_polars/tests/experimental/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/test_runtime_profiler.py deleted file mode 100644 index da7765d67a41..000000000000 --- a/python/cudf_polars/tests/experimental/test_runtime_profiler.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import pytest - -import polars as pl - -from cudf_polars.experimental.base import PartitionInfo, RuntimeProfiler -from cudf_polars.experimental.explain import _repr_profile_tree, write_profile_output -from cudf_polars.testing.asserts import DEFAULT_CLUSTER, DEFAULT_RUNTIME - - -def test_runtime_profiler_and_output(tmp_path): - """Test RuntimeProfiler, merge, and output formatting.""" - - class MockIR: - children = () - - def __init__(self): - self.schema = {"x": pl.Int64} - - ir1, ir2 = MockIR(), MockIR() - - # Test accumulation and merge - profiler1 = RuntimeProfiler() - profiler1.row_count[ir1] = 100 - profiler1.chunk_count[ir1] = 5 - profiler1.decisions[ir1] = "shuffle" - - profiler2 = RuntimeProfiler() - profiler2.row_count[ir1] = 150 - profiler2.row_count[ir2] = 200 - profiler2.chunk_count[ir1] = 3 - profiler2.chunk_count[ir2] = 4 - - profiler1.merge(profiler2) - assert profiler1.row_count[ir1] == 250 - assert profiler1.chunk_count[ir1] == 8 - - # Test _repr_profile_tree output format - partition_info = {ir1: PartitionInfo(count=1)} - output = _repr_profile_tree(ir1, partition_info, profiler1) - assert "rows=250" in output - assert "chunks=8" in output - assert "decision=shuffle" in output - - # Test write_profile_output - output_path = tmp_path / "profile.txt" - write_profile_output(output_path, ir1, partition_info, profiler1) - assert output_path.exists() - content = output_path.read_text() - assert "rows=250" in content - - -@pytest.mark.skipif( - DEFAULT_RUNTIME != "rapidsmpf", reason="Requires 'rapidsmpf' runtime." -) -@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") -def test_profiling_with_real_query(tmp_path): - """Test profiling output with a real query execution.""" - output_path = tmp_path / "profile.txt" - engine = pl.GPUEngine( - raise_on_fail=True, - executor="streaming", - executor_options={ - "cluster": DEFAULT_CLUSTER, - "runtime": DEFAULT_RUNTIME, - "profiling": {"output_path": str(output_path)}, - }, - ) - - df = pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) - q = df.lazy().filter(pl.col("x") > 50).select(["x", "y"]) - q.collect(engine=engine) - - # Verify profile was written - assert output_path.exists() - content = output_path.read_text() - assert "rows=" in content - assert "chunks=" in content From 126c916d655dc525dcff101ea6e8b77da34280ed Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 12:39:20 -0800 Subject: [PATCH 05/25] tedious migration --- .../rapidsmpf/collectives/shuffle.py | 19 +- .../experimental/rapidsmpf/core.py | 12 +- .../experimental/rapidsmpf/dask.py | 10 +- .../cudf_polars/experimental/rapidsmpf/io.py | 9 +- .../experimental/rapidsmpf/join.py | 11 +- .../experimental/rapidsmpf/nodes.py | 42 ++-- .../experimental/rapidsmpf/repartition.py | 8 +- .../experimental/rapidsmpf/union.py | 8 +- .../experimental/rapidsmpf/utils.py | 206 +++++++++++------- .../experimental/rapidsmpf/test_metadata.py | 13 +- 10 files changed, 199 insertions(+), 139 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py index 31bae53839fd..7efacdfdc6c9 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py @@ -13,6 +13,11 @@ from rapidsmpf.streaming.coll.shuffler import ShufflerAsync from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.core.node import define_py_node +from rapidsmpf.streaming.cudf.channel_metadata import ( + ChannelMetadata, + HashScheme, + Partitioning, +) from rapidsmpf.streaming.cudf.table_chunk import TableChunk from cudf_polars.dsl.expr import Col @@ -22,8 +27,6 @@ from cudf_polars.experimental.rapidsmpf.nodes import shutdown_on_error from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, - HashPartitioned, - Metadata, recv_metadata, send_metadata, ) @@ -162,15 +165,11 @@ async def shuffle_node( async with shutdown_on_error(context, ch_in, ch_out): # Receive and send updated metadata. _ = await recv_metadata(ch_in, context) - column_names = list(ir.schema.keys()) - partitioned_on = tuple(column_names[i] for i in columns_to_hash) - output_metadata = Metadata( + output_metadata = ChannelMetadata( local_count=max(1, num_partitions // context.comm().nranks), - global_count=num_partitions, - partitioning=HashPartitioned( - columns=partitioned_on, - scope="global", - count=num_partitions, + partitioning=Partitioning( + inter_rank=HashScheme(columns_to_hash, num_partitions), + local="inherit", ), ) await send_metadata(ch_out, context, output_metadata) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index f308d3c11fbb..e9408b8c520e 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -49,6 +49,7 @@ from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.leaf_node import DeferredMessages + from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata import polars as pl @@ -63,7 +64,6 @@ LowerState, SubNetGenerator, ) - from cudf_polars.experimental.rapidsmpf.utils import Metadata def evaluate_logical_plan( @@ -71,7 +71,7 @@ def evaluate_logical_plan( config_options: ConfigOptions, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: """ Evaluate a logical plan with the RapidsMPF streaming runtime. @@ -139,7 +139,7 @@ def evaluate_pipeline( rmpf_context: Context | None = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: """ Build and evaluate a RapidsMPF streaming pipeline. @@ -229,7 +229,9 @@ def evaluate_pipeline( # Generate network nodes assert rmpf_context is not None, "RapidsMPF context must defined." - metadata_collector: list[Metadata] | None = [] if collect_metadata else None + metadata_collector: list[ChannelMetadata] | None = ( + [] if collect_metadata else None + ) nodes, output = generate_network( rmpf_context, ir, @@ -413,7 +415,7 @@ def generate_network( *, ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], - metadata_collector: list[Metadata] | None, + metadata_collector: list[ChannelMetadata] | None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py index f9527d5f963a..f1e25b3296e0 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py @@ -19,11 +19,11 @@ from collections.abc import MutableMapping from distributed import Client + from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata 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.utils import Metadata class EvaluatePipelineCallback(Protocol): @@ -39,7 +39,7 @@ def __call__( rmpf_context: Context | None = None, *, collect_metadata: bool = False, - ) -> tuple[pl.DataFrame, list[Metadata] | None]: + ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: """Evaluate a pipeline and return the result DataFrame and metadata.""" ... @@ -61,7 +61,7 @@ def evaluate_pipeline_dask( collective_id_map: dict[IR, list[int]], *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: """ Evaluate a RapidsMPF streaming pipeline on a Dask cluster. @@ -98,7 +98,7 @@ def evaluate_pipeline_dask( collect_metadata=collect_metadata, ) dfs: list[pl.DataFrame] = [] - metadata_collector: list[Metadata] = [] + metadata_collector: list[ChannelMetadata] = [] for df, md in result.values(): dfs.append(df) if md is not None: @@ -117,7 +117,7 @@ def _evaluate_pipeline_dask( dask_worker: Any = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[Metadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: """ Build and evaluate a RapidsMPF streaming pipeline. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py index d1ee3fcfd70f..f0764d7b693b 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from rapidsmpf.streaming.cudf.table_chunk import TableChunk import pylibcudf as plc @@ -39,7 +40,6 @@ ) from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, - Metadata, opaque_reservation, send_metadata, ) @@ -183,7 +183,7 @@ async def dataframescan_node( await send_metadata( ch_out, context, - Metadata(local_count=local_count, global_count=global_count), + ChannelMetadata(local_count=local_count), ) # Build list of IR slices to read @@ -468,7 +468,7 @@ async def scan_node( await send_metadata( ch_out, context, - Metadata(local_count=len(scans), global_count=count), + ChannelMetadata(local_count=len(scans)), ) # If there is nothing to scan, drain the channel and return @@ -686,13 +686,12 @@ def _( rec.state["context"], ch_in, ch_out, - Metadata( + ChannelMetadata( # partition_info.count is the estimated "global" count. # Just estimate the local count as well. local_count=math.ceil( partition_info.count / rec.state["context"].comm().nranks ), - global_count=partition_info.count, ), ) nodes[ir] = [native_node, metadata_node] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py index 6abe796172f7..d83e6f261533 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py @@ -9,6 +9,7 @@ from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from rapidsmpf.streaming.cudf.table_chunk import TableChunk from cudf_polars.containers import DataFrame @@ -24,7 +25,6 @@ ) from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, - Metadata, chunk_to_frame, empty_table_chunk, opaque_reservation, @@ -37,10 +37,10 @@ if TYPE_CHECKING: from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context + from rapidsmpf.streaming.cudf.channel_metadata import Partitioning from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator - from cudf_polars.experimental.rapidsmpf.utils import HashPartitioned @define_py_node() @@ -86,7 +86,7 @@ async def broadcast_join_node( recv_metadata(ch_right, context), ) - partitioning: HashPartitioned | None = None + partitioning: Partitioning | None = None if broadcast_side == "right": # Broadcast right, stream left small_ch = ch_right @@ -95,7 +95,6 @@ async def broadcast_join_node( large_child = ir.children[0] # Preserve left-side partitioning metadata local_count = left_metadata.local_count - global_count = left_metadata.global_count partitioning = left_metadata.partitioning # Check if the right-side is already broadcasted small_duplicated = right_metadata.duplicated @@ -107,16 +106,14 @@ async def broadcast_join_node( large_child = ir.children[1] # Preserve right-side partitioning metadata local_count = right_metadata.local_count - global_count = right_metadata.global_count if ir.options[0] == "Right": partitioning = right_metadata.partitioning # Check if the right-side is already broadcasted small_duplicated = left_metadata.duplicated # Send metadata. - output_metadata = Metadata( + output_metadata = ChannelMetadata( local_count=local_count, - global_count=global_count, partitioning=partitioning, # The result is only "duplicated" if both sides are duplicated duplicated=left_metadata.duplicated and right_metadata.duplicated, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index 86b7bfde5072..d49b3d947f7f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -11,6 +11,7 @@ from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.core.node import define_py_node from rapidsmpf.streaming.core.spillable_messages import SpillableMessages +from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from rapidsmpf.streaming.cudf.table_chunk import TableChunk from cudf_polars.containers import DataFrame @@ -20,12 +21,12 @@ ) from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, - Metadata, empty_table_chunk, make_spill_function, opaque_reservation, process_children, recv_metadata, + remap_partitioning, send_metadata, shutdown_on_error, ) @@ -73,10 +74,15 @@ async def default_node_single( async with shutdown_on_error(context, ch_in, ch_out): # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) - metadata_out = Metadata( + # Remap partitioning if preserving and schema might have changed + partitioning = None + if preserve_partitioning: + partitioning = remap_partitioning( + metadata_in.partitioning, ir.children[0].schema, ir.schema + ) + metadata_out = ChannelMetadata( local_count=metadata_in.local_count, - global_count=metadata_in.global_count, - partitioning=metadata_in.partitioning if preserve_partitioning else None, + partitioning=partitioning, duplicated=metadata_in.duplicated, ) await send_metadata(ch_out, context, metadata_out) @@ -162,20 +168,26 @@ async def default_node_multi( """ async with shutdown_on_error(context, *chs_in, ch_out): # Merge and forward basic metadata. - metadata = Metadata(local_count=1, duplicated=True) + local_count = 1 + duplicated = True + partitioning = None for idx, ch_in in enumerate(chs_in): md_child = await recv_metadata(ch_in, context) # Use simple "max" rule to determine counts. - metadata.local_count = max(md_child.local_count, metadata.local_count) - if md_child.global_count is not None: - metadata.global_count = max( - md_child.global_count, metadata.global_count or 0 - ) + local_count = max(md_child.local_count, local_count) # Set "duplicated" to False as soon as we # find a non-duplicated child. - metadata.duplicated = metadata.duplicated and md_child.duplicated + duplicated = duplicated and md_child.duplicated if idx == partitioning_index: - metadata.partitioning = md_child.partitioning + # Remap partitioning from child schema to output schema + partitioning = remap_partitioning( + md_child.partitioning, ir.children[idx].schema, ir.schema + ) + metadata = ChannelMetadata( + local_count=local_count, + partitioning=partitioning, + duplicated=duplicated, + ) await send_metadata(ch_out, context, metadata) seq_num = 0 @@ -581,7 +593,7 @@ async def empty_node( ch_out, context, # All ranks generate the same "empty" data. - Metadata(local_count=1, global_count=1, duplicated=True), + ChannelMetadata(local_count=1, duplicated=True), ) # Evaluate the IR node to create an empty DataFrame @@ -660,7 +672,7 @@ async def metadata_feeder_node( context: Context, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk], - metadata: Metadata, + metadata: ChannelMetadata, ) -> None: """ Forward data with new metadata. @@ -690,7 +702,7 @@ async def metadata_drain_node( ir_context: IRExecutionContext, ch_in: Channel[TableChunk], ch_out: Any, - metadata_collector: list[Metadata] | None, + metadata_collector: list[ChannelMetadata] | None, ) -> None: """ Drain metadata and forward data to a single channel. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py index e3e32d6217a7..380b8a68e0ca 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -10,6 +10,7 @@ from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.core.node import define_py_node +from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from rapidsmpf.streaming.cudf.table_chunk import TableChunk from cudf_polars.containers import DataFrame @@ -18,7 +19,6 @@ from cudf_polars.experimental.rapidsmpf.nodes import shutdown_on_error from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, - Metadata, empty_table_chunk, opaque_reservation, recv_metadata, @@ -127,9 +127,8 @@ async def concatenate_node( # Global repartitioning via AllGather to single duplicated chunk. # Send metadata. - metadata = Metadata( + metadata = ChannelMetadata( local_count=local_output_count, - global_count=output_count, duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) @@ -160,9 +159,8 @@ async def concatenate_node( # Local repartitioning (tree reduction). # Send metadata. - metadata = Metadata( + metadata = ChannelMetadata( local_count=local_output_count, - global_count=output_count, duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py index 627c8000efc8..4e7067f46093 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from rapidsmpf.streaming.cudf.table_chunk import TableChunk from cudf_polars.dsl.ir import Union @@ -16,7 +17,6 @@ from cudf_polars.experimental.rapidsmpf.nodes import define_py_node, shutdown_on_error from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, - Metadata, process_children, recv_metadata, send_metadata, @@ -59,20 +59,16 @@ async def union_node( # Union loses partitioning/ordering info since sources may differ. # TODO: Warn users that Union does NOT preserve order? total_local_count = 0 - total_global_count: int | None = None duplicated = True for ch_in in chs_in: metadata = await recv_metadata(ch_in, context) total_local_count += metadata.local_count - if metadata.global_count is not None: - total_global_count = (total_global_count or 0) + metadata.global_count duplicated = duplicated and metadata.duplicated await send_metadata( ch_out, context, - Metadata( + ChannelMetadata( local_count=total_local_count, - global_count=total_global_count, duplicated=duplicated, ), ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 63d657291650..75576861060a 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -6,12 +6,19 @@ import asyncio import operator +import struct from contextlib import asynccontextmanager, contextmanager from functools import reduce -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any -from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk +from rapidsmpf.memory.packed_data import PackedData +from rapidsmpf.streaming.coll.allgather import AllGather from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.cudf.channel_metadata import ( + ChannelMetadata, + HashScheme, + Partitioning, +) from rapidsmpf.streaming.cudf.table_chunk import TableChunk import pylibcudf as plc @@ -19,7 +26,7 @@ from cudf_polars.containers import DataFrame if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Iterator + from collections.abc import AsyncIterator, Callable, Iterator, Mapping from rapidsmpf.memory.memory_reservation import MemoryReservation from rapidsmpf.streaming.core.channel import Channel @@ -30,6 +37,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.typing import DataType @asynccontextmanager @@ -54,82 +62,71 @@ async def shutdown_on_error( raise -class HashPartitioned: +def remap_partitioning( + partitioning: Partitioning | None, + old_schema: Mapping[str, DataType], + new_schema: Mapping[str, DataType], +) -> Partitioning | None: """ - Hash-partitioned metadata. + Remap partitioning column indices from old schema to new schema. - Attributes - ---------- - columns - Columns the data is hash-partitioned on. - scope - Whether data is partitioned locally (within a rank) or - globally (across all ranks). - count - The modulus used for hash partitioning (number of partitions). - """ - - __slots__ = ("columns", "count", "scope") - - columns: tuple[str, ...] - scope: Literal["local", "global"] - count: int - - def __init__( - self, - columns: tuple[str, ...], - scope: Literal["local", "global"], - count: int, - ): - self.columns = columns - self.scope = scope - self.count = count + Since HashScheme uses column indices rather than names, we need to + remap indices when propagating partitioning through operations that + may change the schema (column order or presence). + Parameters + ---------- + partitioning + The partitioning to remap. + old_schema + The schema where the partitioning was established. + new_schema + The new schema to remap to. -class Metadata: - """Metadata payload for a channel.""" - - __slots__ = ( - "duplicated", - "global_count", - "local_count", - "partitioning", - ) - - # Chunk counts - local_count: int - """Local chunk-count estimate for the current rank.""" - global_count: int | None - """Global chunk-count estimate across all ranks.""" - - # Partitioning - partitioning: HashPartitioned | None - """How the data is hash-partitioned, or None if not partitioned.""" - - # Duplication - duplicated: bool - """Whether the data is duplicated (identical) on all workers.""" - - def __init__( - self, - local_count: int, - *, - global_count: int | None = None, - partitioning: HashPartitioned | None = None, - duplicated: bool = False, - ): - if local_count < 0: # pragma: no cover - raise ValueError(f"Local count must be non-negative. Got: {local_count}") - self.local_count = local_count - if global_count is not None and global_count < 0: # pragma: no cover - raise ValueError(f"Global count must be non-negative. Got: {global_count}") - self.global_count = global_count - self.partitioning = partitioning - self.duplicated = duplicated + Returns + ------- + The remapped partitioning, or None if any partitioning column + is not present in the new schema. + """ + if partitioning is None: + return None + + old_names = list(old_schema.keys()) + new_names = list(new_schema.keys()) + + def remap_hash_scheme(hs: HashScheme | None | str) -> HashScheme | None | str: + if hs is None or isinstance(hs, str): + # None or "inherit" - inherits parent partitioning unchanged + return hs + # Get column names from old indices + try: + column_names = [old_names[i] for i in hs.column_indices] + except IndexError: + return None # Invalid index in old schema + # Check all exist in new schema and map to new indices + new_indices = [] + for name in column_names: + if name not in new_names: + return None # Column not in new schema - partitioning invalidated + new_indices.append(new_names.index(name)) + return HashScheme(tuple(new_indices), hs.modulus) + + new_inter_rank = remap_hash_scheme(partitioning.inter_rank) + new_local = remap_hash_scheme(partitioning.local) + + # If inter_rank was a HashScheme and got invalidated, whole partitioning is invalid + if isinstance(partitioning.inter_rank, HashScheme) and new_inter_rank is None: + return None + + # If local was a HashScheme and got invalidated, set it to None + if isinstance(partitioning.local, HashScheme) and new_local is None: + new_local = None + + return Partitioning(inter_rank=new_inter_rank, local=new_local) async def send_metadata( - ch: Channel[TableChunk], ctx: Context, metadata: Metadata + ch: Channel[TableChunk], ctx: Context, metadata: ChannelMetadata ) -> None: """ Send metadata and drain the metadata queue. @@ -142,13 +139,25 @@ async def send_metadata( The streaming context. metadata : The metadata to send. + + Notes + ----- + This function copies the metadata before sending, so the caller + retains ownership of the original metadata object. """ - msg = Message(0, ArbitraryChunk(metadata)) + # Copy metadata before sending since Message consumes the handle. + # Metadata is small, so copying is cheap. + metadata_copy = ChannelMetadata( + local_count=metadata.local_count, + partitioning=metadata.partitioning, + duplicated=metadata.duplicated, + ) + msg = Message(0, metadata_copy) await ch.send_metadata(ctx, msg) await ch.drain_metadata(ctx) -async def recv_metadata(ch: Channel[TableChunk], ctx: Context) -> Metadata: +async def recv_metadata(ch: Channel[TableChunk], ctx: Context) -> ChannelMetadata: """ Receive metadata from a channel's metadata queue. @@ -161,12 +170,12 @@ async def recv_metadata(ch: Channel[TableChunk], ctx: Context) -> Metadata: Returns ------- - Metadata + ChannelMetadata The received metadata. """ msg = await ch.recv_metadata(ctx) - assert msg is not None, f"Expected Metadata message, got {msg}." - return ArbitraryChunk.from_message(msg).release() + assert msg is not None, f"Expected ChannelMetadata message, got {msg}." + return ChannelMetadata.from_message(msg) class ChannelManager: @@ -396,3 +405,46 @@ def opaque_reservation( yield context.br().reserve_device_memory_and_spill( estimated_bytes, allow_overbooking=True ) + + +async def allgather_reduce( + context: Context, + op_id: int, + *local_values: int, +) -> tuple[int, ...]: + """ + Allgather local scalar values and sum each across all ranks. + + Parameters + ---------- + context + The rapidsmpf context. + op_id + The collective operation ID for this allgather. + *local_values + One or more local scalar values to contribute. + + Returns + ------- + tuple[int, ...] + The sum of each local_value across all ranks. + """ + n = len(local_values) + fmt = "q" * n + data = struct.pack(fmt, *local_values) + packed = PackedData.from_host_bytes(data, context.br()) + + allgather = AllGather(context, op_id) + allgather.insert(0, packed) + allgather.insert_finished() + + results = await allgather.extract_all(context, ordered=False) + + totals = [0] * n + for packed_result in results: + result_bytes = packed_result.to_host_bytes() + values = struct.unpack(fmt, result_bytes) + for i, v in enumerate(values): + totals[i] += v + + return tuple(totals) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_metadata.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_metadata.py index 1265611bda83..52ec750b9e56 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_metadata.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_metadata.py @@ -6,6 +6,7 @@ from __future__ import annotations import pytest +from rapidsmpf.streaming.cudf.channel_metadata import HashScheme import polars as pl @@ -79,8 +80,12 @@ def test_rapidsmpf_join_metadata( assert metadata.local_count == left_count assert metadata.duplicated is False if right_count > broadcast_join_limit: - assert metadata.partitioning is not None - assert metadata.partitioning.columns == ("y",) - assert metadata.partitioning.scope == "global" + # After shuffle, partitioning has inter_rank=HashScheme, local="inherit" + assert isinstance(metadata.partitioning.inter_rank, HashScheme) + # "y" is at index 1 in the output schema: ["x", "y", "z", "xx", "zz"] + assert metadata.partitioning.inter_rank.column_indices == (1,) + assert metadata.partitioning.local == "inherit" else: - assert metadata.partitioning is None + # No partitioning (broadcast join preserves no partitioning from IO) + assert metadata.partitioning.inter_rank is None + assert metadata.partitioning.local is None From 1c5151e033ff604a216ca4373cf457e9aa188e15 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 12:44:22 -0800 Subject: [PATCH 06/25] remove unused allreduce def --- .../experimental/rapidsmpf/utils.py | 46 ------------------- 1 file changed, 46 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 75576861060a..60144c40f42a 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -6,13 +6,10 @@ import asyncio import operator -import struct from contextlib import asynccontextmanager, contextmanager from functools import reduce from typing import TYPE_CHECKING, Any -from rapidsmpf.memory.packed_data import PackedData -from rapidsmpf.streaming.coll.allgather import AllGather from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.cudf.channel_metadata import ( ChannelMetadata, @@ -405,46 +402,3 @@ def opaque_reservation( yield context.br().reserve_device_memory_and_spill( estimated_bytes, allow_overbooking=True ) - - -async def allgather_reduce( - context: Context, - op_id: int, - *local_values: int, -) -> tuple[int, ...]: - """ - Allgather local scalar values and sum each across all ranks. - - Parameters - ---------- - context - The rapidsmpf context. - op_id - The collective operation ID for this allgather. - *local_values - One or more local scalar values to contribute. - - Returns - ------- - tuple[int, ...] - The sum of each local_value across all ranks. - """ - n = len(local_values) - fmt = "q" * n - data = struct.pack(fmt, *local_values) - packed = PackedData.from_host_bytes(data, context.br()) - - allgather = AllGather(context, op_id) - allgather.insert(0, packed) - allgather.insert_finished() - - results = await allgather.extract_all(context, ordered=False) - - totals = [0] * n - for packed_result in results: - result_bytes = packed_result.to_host_bytes() - values = struct.unpack(fmt, result_bytes) - for i, v in enumerate(values): - totals[i] += v - - return tuple(totals) From d3b7583d5ed2583f397665650d6c9f84ad146c86 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 12:57:21 -0800 Subject: [PATCH 07/25] update helper function --- .../experimental/rapidsmpf/nodes.py | 2 +- .../experimental/rapidsmpf/utils.py | 46 +++++++++---------- 2 files changed, 22 insertions(+), 26 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index d49b3d947f7f..4330cdb176fe 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -74,9 +74,9 @@ async def default_node_single( async with shutdown_on_error(context, ch_in, ch_out): # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) - # Remap partitioning if preserving and schema might have changed partitioning = None if preserve_partitioning: + # Remap partitioning if schema has changed partitioning = remap_partitioning( metadata_in.partitioning, ir.children[0].schema, ir.schema ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 60144c40f42a..3a8d0522be53 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -82,40 +82,34 @@ def remap_partitioning( Returns ------- - The remapped partitioning, or None if any partitioning column - is not present in the new schema. + The remapped partitioning, or None if the inter-rank partitioning + columns are not present in the new schema. """ if partitioning is None: return None old_names = list(old_schema.keys()) - new_names = list(new_schema.keys()) + 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): - # None or "inherit" - inherits parent partitioning unchanged - return hs - # Get column names from old indices + return hs # None or "inherit" passes through unchanged try: - column_names = [old_names[i] for i in hs.column_indices] - except IndexError: - return None # Invalid index in old schema - # Check all exist in new schema and map to new indices - new_indices = [] - for name in column_names: - if name not in new_names: - return None # Column not in new schema - partitioning invalidated - new_indices.append(new_names.index(name)) - return HashScheme(tuple(new_indices), hs.modulus) + 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 was a HashScheme and got invalidated, whole partitioning is invalid + # 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 local was a HashScheme and got invalidated, set it to 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 @@ -142,14 +136,16 @@ async def send_metadata( This function copies the metadata before sending, so the caller retains ownership of the original metadata object. """ - # Copy metadata before sending since Message consumes the handle. - # Metadata is small, so copying is cheap. - metadata_copy = ChannelMetadata( - local_count=metadata.local_count, - partitioning=metadata.partitioning, - duplicated=metadata.duplicated, + msg = Message( + 0, + # Copy metadata before sending since Message consumes the handle. + # Metadata is small, so copying is cheap. + ChannelMetadata( + local_count=metadata.local_count, + partitioning=metadata.partitioning, + duplicated=metadata.duplicated, + ), ) - msg = Message(0, metadata_copy) await ch.send_metadata(ctx, msg) await ch.drain_metadata(ctx) From 07443e57bdfa9172d5e0d1179f1332a1c03a795c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 13:35:56 -0800 Subject: [PATCH 08/25] cleanup and repartition --- .../cudf_polars/cudf_polars/experimental/base.py | 13 +++++++------ .../cudf_polars/experimental/rapidsmpf/io.py | 2 +- .../cudf_polars/experimental/rapidsmpf/nodes.py | 4 ++-- .../experimental/rapidsmpf/repartition.py | 10 ++++++++++ .../rapidsmpf/test_runtime_profiler.py | 15 +++++++++------ 5 files changed, 29 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py index 6a325059f56e..06251fc22385 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/experimental/base.py @@ -14,7 +14,8 @@ if TYPE_CHECKING: from collections.abc import Generator, Iterator, MutableMapping - from cudf_polars.containers import DataFrame + import pylibcudf as plc + from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR from cudf_polars.dsl.nodebase import Node @@ -430,15 +431,15 @@ def __init__(self) -> None: self.chunk_count: int = 0 self.decision: str | None = None - def add_chunk(self, *, df: DataFrame | None = None) -> None: + def add_chunk(self, *, table: plc.Table | None = None) -> None: """ Record a chunk. - If df is provided, both row_count and chunk_count are updated. - If df is None, only chunk_count is incremented. + If table is provided, both row_count and chunk_count are updated. + If table is None, only chunk_count is incremented. """ - if df is not None: - self.row_count = (self.row_count or 0) + df.table.num_rows() + if table is not None: + self.row_count = (self.row_count or 0) + table.num_rows() self.chunk_count += 1 def merge(self, other: RuntimeNodeProfiler) -> None: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py index b5da449330c3..3d6ba74a68ad 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -364,7 +364,7 @@ async def read_chunk( context=ir_context, ) if node_profiler is not None: - node_profiler.add_chunk(df=df) + node_profiler.add_chunk(table=df.table) await ch_out.send( context, Message( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index 2a301c1f80b3..8f3128d7991e 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -122,7 +122,7 @@ async def default_node_single( context=ir_context, ) if node_profiler is not None: - node_profiler.add_chunk(df=df) + node_profiler.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -251,7 +251,7 @@ async def default_node_multi( context=ir_context, ) if node_profiler is not None: - node_profiler.add_chunk(df=df) + node_profiler.add_chunk(table=df.table) await ch_out.send( context, Message( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py index e3e32d6217a7..3cc91b80ac1f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -32,6 +32,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.experimental.base import RuntimeNodeProfiler from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator @@ -45,6 +46,7 @@ async def concatenate_node( *, output_count: int, collective_id: int, + node_profiler: RuntimeNodeProfiler | None = None, ) -> None: """ Concatenate node for rapidsmpf. @@ -75,6 +77,8 @@ async def concatenate_node( The expected global number of output chunks. collective_id Pre-allocated collective ID for this operation. + node_profiler + Node profiler for collecting runtime statistics. """ async with shutdown_on_error(context, ch_in, ch_out): # Receive metadata. @@ -145,6 +149,8 @@ async def concatenate_node( # Extract concatenated result result_table = await allgather.extract_concatenated(stream) + if node_profiler is not None: + node_profiler.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. @@ -203,6 +209,8 @@ async def concatenate_node( ), context=ir_context, ) + if node_profiler is not None: + node_profiler.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -245,6 +253,7 @@ def _( collective_id = rec.state["collective_id_map"][ir][0] # Add python node + profiler = rec.state["profiler"] nodes[ir] = [ concatenate_node( rec.state["context"], @@ -254,6 +263,7 @@ def _( channels[ir.children[0]].reserve_output_slot(), output_count=partition_info[ir].count, collective_id=collective_id, + node_profiler=profiler.get_or_create(ir) if profiler else None, ) ] return nodes, channels diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py index ba951e5d0394..24a339b11a01 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py @@ -79,15 +79,17 @@ def __init__(self): @pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") -def test_profiling_dataframe_scan(tmp_path, df): +def test_profiling_basic_query(tmp_path, df): """Test profiling output with a DataFrameScan query.""" output_path = tmp_path / "dataframe_scan_profile.txt" engine = get_engine(output_path) - q = df.lazy().filter(pl.col("x") > 50).select(["x", "y"]) + 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 "FILTER ('x', 'y') rows=49 chunks=10 [10]" in content - assert "DATAFRAMESCAN ('x', 'y') rows=100 chunks=10 [10]" in content + 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.") @@ -101,7 +103,7 @@ def test_profiling_scan_parquet_python(tmp_path, df): 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 [5]" in content + assert "SCAN PARQUET ('x', 'y') rows=49 chunks=5" in content @pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") @@ -115,4 +117,5 @@ def test_profiling_scan_parquet_native(tmp_path, df): 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=? chunks=5 [5]" in content + # We can count the chunks but not the rows for "native" parquet + assert "SCAN PARQUET ('x', 'y') rows=? chunks=5" in content From c8be0984b868be88995239f53904c35f85c13c2c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 13:51:45 -0800 Subject: [PATCH 09/25] leave out rows when unknown --- python/cudf_polars/cudf_polars/experimental/explain.py | 8 ++++---- .../tests/experimental/rapidsmpf/test_runtime_profiler.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index d3e808680845..771776ae42c1 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -214,11 +214,11 @@ def _repr_profile_tree( # Get node profiler if it exists node_profiler = profiler.node_profilers.get(ir) + header = header.rstrip("\n") - # Add actual row count - actual_rows = node_profiler.row_count if node_profiler is not None else None - actual_str = _fmt_row_count(actual_rows) if actual_rows is not None else "?" - header = header.rstrip("\n") + f" rows={actual_str}" + # Add actual row count if available + if node_profiler is not None and node_profiler.row_count is not None: + header += f" rows={_fmt_row_count(node_profiler.row_count)}" # Add decision if present if node_profiler is not None and node_profiler.decision is not None: diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py index 24a339b11a01..4cc6c3af26a6 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py @@ -118,4 +118,5 @@ def test_profiling_scan_parquet_native(tmp_path, df): q.collect(engine=engine) content = output_path.read_text() # We can count the chunks but not the rows for "native" parquet - assert "SCAN PARQUET ('x', 'y') rows=? chunks=5" in content + # (row count is omitted when unavailable) + assert "SCAN PARQUET ('x', 'y') chunks=5" in content From 05df3735f3b4cc7eaa6104f916bef2f91c72d3d8 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 13:56:20 -0800 Subject: [PATCH 10/25] leave off static partition count --- python/cudf_polars/cudf_polars/experimental/explain.py | 5 ----- .../tests/experimental/rapidsmpf/test_runtime_profiler.py | 1 - 2 files changed, 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 771776ae42c1..96683338a105 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -210,7 +210,6 @@ def _repr_profile_tree( ) -> str: """Recursively build a tree representation with profiler data.""" header = _repr_ir(ir, offset=offset) - static_count = partition_info[ir].count if partition_info else None # Get node profiler if it exists node_profiler = profiler.node_profilers.get(ir) @@ -228,10 +227,6 @@ def _repr_profile_tree( if node_profiler is not None and node_profiler.chunk_count > 0: header += f" chunks={node_profiler.chunk_count}" - # Add expected partition count from PartitionInfo - if static_count is not None: - header += f" [{static_count}]" - header += "\n" children_strs = [ diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py index 4cc6c3af26a6..7fd24cd62f85 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py @@ -68,7 +68,6 @@ def __init__(self): assert "rows=250" in output assert "chunks=8" in output assert "decision=shuffle" in output - assert "[4]" in output # static partition count # Test write_profile_output output_path = tmp_path / "profile.txt" From 032d13d1b0fd1bb33ab83cd8305e85c1e0fd0a92 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 15:12:43 -0800 Subject: [PATCH 11/25] improve cov --- .../rapidsmpf/test_runtime_profiler.py | 47 +--------------- .../experimental/test_runtime_profiler.py | 54 +++++++++++++++++++ 2 files changed, 55 insertions(+), 46 deletions(-) create mode 100644 python/cudf_polars/tests/experimental/test_runtime_profiler.py diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py index 7fd24cd62f85..57a1f614e035 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 +"""Integration tests for runtime profiling with rapidsmpf.""" from __future__ import annotations @@ -7,8 +8,6 @@ import polars as pl -from cudf_polars.experimental.base import PartitionInfo, RuntimeQueryProfiler -from cudf_polars.experimental.explain import _repr_profile_tree, write_profile_output from cudf_polars.testing.asserts import DEFAULT_CLUSTER, DEFAULT_RUNTIME from cudf_polars.testing.io import make_partitioned_source @@ -33,50 +32,6 @@ def df(): return pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) -def test_runtime_profiler_and_output(tmp_path): - """Test RuntimeQueryProfiler, merge, and output formatting.""" - - class MockIR: - children = () - - def __init__(self): - self.schema = {"x": pl.Int64} - - ir1, ir2 = MockIR(), MockIR() - - # Test node profiler accumulation - profiler1 = RuntimeQueryProfiler() - np1 = profiler1.get_or_create(ir1) - np1.row_count = 100 - np1.chunk_count = 5 - np1.decision = "shuffle" - - profiler2 = RuntimeQueryProfiler() - profiler2.get_or_create(ir1).row_count = 150 - profiler2.get_or_create(ir1).chunk_count = 3 - profiler2.get_or_create(ir2).row_count = 200 - profiler2.get_or_create(ir2).chunk_count = 4 - - # Test merge - profiler1.merge(profiler2) - assert profiler1.node_profilers[ir1].row_count == 250 - assert profiler1.node_profilers[ir1].chunk_count == 8 - - # Test _repr_profile_tree output format - partition_info = {ir1: PartitionInfo(count=4)} - output = _repr_profile_tree(ir1, partition_info, profiler1) - assert "rows=250" in output - assert "chunks=8" in output - assert "decision=shuffle" in output - - # Test write_profile_output - output_path = tmp_path / "profile.txt" - write_profile_output(output_path, ir1, partition_info, profiler1) - assert output_path.exists() - content = output_path.read_text() - assert "rows=250" in content - - @pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") def test_profiling_basic_query(tmp_path, df): """Test profiling output with a DataFrameScan query.""" diff --git a/python/cudf_polars/tests/experimental/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/test_runtime_profiler.py new file mode 100644 index 000000000000..488b7f8e134e --- /dev/null +++ b/python/cudf_polars/tests/experimental/test_runtime_profiler.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for RuntimeQueryProfiler and related output functions.""" + +from __future__ import annotations + +import polars as pl + +from cudf_polars.experimental.base import PartitionInfo, RuntimeQueryProfiler +from cudf_polars.experimental.explain import _repr_profile_tree, write_profile_output + + +def test_runtime_profiler_and_output(tmp_path): + """Test RuntimeQueryProfiler, merge, and output formatting.""" + + class MockIR: + children = () + + def __init__(self): + self.schema = {"x": pl.Int64} + + ir1, ir2 = MockIR(), MockIR() + + # Test node profiler accumulation + profiler1 = RuntimeQueryProfiler() + np1 = profiler1.get_or_create(ir1) + np1.row_count = 100 + np1.chunk_count = 5 + np1.decision = "shuffle" + + profiler2 = RuntimeQueryProfiler() + profiler2.get_or_create(ir1).row_count = 150 + profiler2.get_or_create(ir1).chunk_count = 3 + profiler2.get_or_create(ir2).row_count = 200 + profiler2.get_or_create(ir2).chunk_count = 4 + + # Test merge + profiler1.merge(profiler2) + assert profiler1.node_profilers[ir1].row_count == 250 + assert profiler1.node_profilers[ir1].chunk_count == 8 + + # Test _repr_profile_tree output format + partition_info = {ir1: PartitionInfo(count=4)} + output = _repr_profile_tree(ir1, partition_info, profiler1) + assert "rows=250" in output + assert "chunks=8" in output + assert "decision=shuffle" in output + + # Test write_profile_output + output_path = tmp_path / "profile.txt" + write_profile_output(output_path, ir1, partition_info, profiler1) + assert output_path.exists() + content = output_path.read_text() + assert "rows=250" in content From d0bf467d7877dca9871862bdb34ded486b092873 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 Jan 2026 19:11:13 -0800 Subject: [PATCH 12/25] more coverage --- .../cudf_polars/tests/experimental/test_runtime_profiler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/experimental/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/test_runtime_profiler.py index 488b7f8e134e..e66c405057df 100644 --- a/python/cudf_polars/tests/experimental/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/test_runtime_profiler.py @@ -31,8 +31,10 @@ def __init__(self): profiler2 = RuntimeQueryProfiler() profiler2.get_or_create(ir1).row_count = 150 profiler2.get_or_create(ir1).chunk_count = 3 + profiler2.get_or_create(ir1).decision = "shuffle" profiler2.get_or_create(ir2).row_count = 200 - profiler2.get_or_create(ir2).chunk_count = 4 + profiler2.get_or_create(ir2).chunk_count = 3 + profiler2.get_or_create(ir2).add_chunk() # Test merge profiler1.merge(profiler2) From 74737c5a0fa19f9e5c5ccc2e592180ba94388e86 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Sat, 31 Jan 2026 04:56:35 -0800 Subject: [PATCH 13/25] missed a line --- python/cudf_polars/cudf_polars/experimental/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py index 06251fc22385..c51309bb4db1 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/experimental/base.py @@ -438,7 +438,7 @@ def add_chunk(self, *, table: plc.Table | None = None) -> None: 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: + 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 From c44ae02bb53bdb1ad65eb7b4bbe287b212a6a31c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 2 Feb 2026 10:52:47 -0800 Subject: [PATCH 14/25] simplify explain.py logic --- .../cudf_polars/experimental/explain.py | 23 ++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 96683338a105..5df6ae626871 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -210,30 +210,27 @@ def _repr_profile_tree( ) -> str: """Recursively build a tree representation with profiler data.""" header = _repr_ir(ir, offset=offset) - - # Get node profiler if it exists - node_profiler = profiler.node_profilers.get(ir) header = header.rstrip("\n") - # Add actual row count if available - if node_profiler is not None and node_profiler.row_count is not None: - header += f" rows={_fmt_row_count(node_profiler.row_count)}" + # Get node profiler if it exists + if (node_profiler := profiler.node_profilers.get(ir)) is not None: + # Add actual row count if available + if node_profiler.row_count is not None: + header += f" rows={_fmt_row_count(node_profiler.row_count)}" - # Add decision if present - if node_profiler is not None and node_profiler.decision is not None: - header += f" decision={node_profiler.decision}" + # Add decision if present + if node_profiler.decision is not None: + header += f" decision={node_profiler.decision}" - # Add actual chunk count if available - if node_profiler is not None and node_profiler.chunk_count > 0: + # Add actual chunk count header += f" chunks={node_profiler.chunk_count}" - header += "\n" - children_strs = [ _repr_profile_tree(child, partition_info, profiler, 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 From 8800222705b3ff3fba340e606db447672c0d23c3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 2 Feb 2026 11:22:45 -0800 Subject: [PATCH 15/25] reuse code --- .../cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index e6f5c1e38d1d..73abb836f295 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -541,7 +541,9 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + profiler: RuntimeQueryProfiler | None = rec.state.get("profiler") + node_profiler = profiler.get_or_create(ir) if profiler is not None else None if len(ir.children) == 1: # Single-channel default node @@ -561,7 +563,7 @@ def _( channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), preserve_partitioning=preserve_partitioning, - node_profiler=profiler.get_or_create(ir) if profiler else None, + node_profiler=node_profiler, ) ] else: @@ -573,7 +575,7 @@ def _( rec.state["ir_context"], channels[ir].reserve_input_slot(), tuple(channels[c].reserve_output_slot() for c in ir.children), - node_profiler=profiler.get_or_create(ir) if profiler else None, + node_profiler=node_profiler, ) ] From 09a0aac0c82a5269f7da78733cf40cc115a6c712 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 2 Feb 2026 20:49:08 -0800 Subject: [PATCH 16/25] add structlog tracing --- .../experimental/rapidsmpf/core.py | 3 +- .../cudf_polars/experimental/rapidsmpf/io.py | 17 ++-- .../experimental/rapidsmpf/nodes.py | 27 ++++-- .../experimental/rapidsmpf/repartition.py | 12 +-- .../experimental/rapidsmpf/utils.py | 84 +++++++++++++++++-- 5 files changed, 117 insertions(+), 26 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 72c86d6cad44..aaeb2bd2f033 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.base import RuntimeQueryProfiler from cudf_polars.experimental.rapidsmpf.collectives import ReserveOpIDs @@ -246,7 +247,7 @@ def evaluate_pipeline( ) profiler: RuntimeQueryProfiler | None = ( RuntimeQueryProfiler() - if config_options.executor.profiling is not None + if config_options.executor.profiling is not None or LOG_TRACES else None ) nodes, output = generate_network( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py index 28c545888009..3c15a4048160 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -174,7 +174,9 @@ async def dataframescan_node( node_profiler The node profiler for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_out): + async with shutdown_on_error( + context, ch_out, ir=ir, node_profiler=node_profiler + ) as profiler: # Find local partition count. nrows = ir.df.shape()[0] global_count = math.ceil(nrows / rows_per_partition) if nrows > 0 else 0 @@ -224,7 +226,7 @@ async def dataframescan_node( ch_out, ir_context, estimated_chunk_bytes, - node_profiler=node_profiler, + node_profiler=profiler, ) await ch_out.drain(context) return @@ -250,7 +252,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, - node_profiler=node_profiler, + node_profiler=profiler, ) await ch_out.drain(context) @@ -416,7 +418,9 @@ async def scan_node( node_profiler The node profiler for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_out): + async with shutdown_on_error( + context, ch_out, ir=ir, node_profiler=node_profiler + ) as profiler: # Build a list of local Scan operations scans: list[Scan | SplitScan] = [] if plan.flavor == IOPartitionFlavor.SPLIT_FILES: @@ -507,7 +511,7 @@ async def scan_node( ch_out, ir_context, estimated_chunk_bytes, - node_profiler=node_profiler, + node_profiler=profiler, ) await ch_out.drain(context) return @@ -533,7 +537,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, - node_profiler=node_profiler, + node_profiler=profiler, ) await ch_out.drain(context) @@ -707,6 +711,7 @@ def _( # node does not send metadata. metadata_node = metadata_feeder_node( rec.state["context"], + ir, ch_in, ch_out, ChannelMetadata( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index 73abb836f295..df7d8b7a19a4 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -75,7 +75,9 @@ async def default_node_single( ----- 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, ir=ir, node_profiler=node_profiler + ) as profiler: # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) partitioning = None @@ -127,8 +129,8 @@ async def default_node_single( ), context=ir_context, ) - if node_profiler is not None: - node_profiler.add_chunk(table=df.table) + if profiler is not None: + profiler.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -175,7 +177,9 @@ async def default_node_multi( node_profiler Node profiler for collecting runtime statistics. """ - async with shutdown_on_error(context, *chs_in, ch_out): + async with shutdown_on_error( + context, *chs_in, ch_out, ir=ir, node_profiler=node_profiler + ) as profiler: # Merge and forward basic metadata. local_count = 1 duplicated = True @@ -262,8 +266,8 @@ async def default_node_multi( *dfs, context=ir_context, ) - if node_profiler is not None: - node_profiler.add_chunk(table=df.table) + if profiler is not None: + profiler.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -686,6 +690,7 @@ def generate_ir_sub_network_wrapper( @define_py_node() async def metadata_feeder_node( context: Context, + ir: IR, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk], metadata: ChannelMetadata, @@ -698,6 +703,8 @@ async def metadata_feeder_node( ---------- context The rapidsmpf context. + ir + The IR node. ch_in The input channel to pull data from. ch_out @@ -707,12 +714,14 @@ async def metadata_feeder_node( node_profiler Node profiler for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_in, ch_out): + async with shutdown_on_error( + context, ch_in, ch_out, ir=ir, node_profiler=node_profiler + ) as profiler: await send_metadata(ch_out, context, metadata) while (msg := await ch_in.recv(context)) is not None: await ch_out.send(context, msg) - if node_profiler is not None: - node_profiler.chunk_count += 1 + if profiler is not None: + profiler.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 dbc8aeb5f5f1..a8a1de7c49c4 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -80,7 +80,9 @@ async def concatenate_node( node_profiler Node profiler for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_in, ch_out): + async with shutdown_on_error( + context, ch_in, ch_out, ir=ir, node_profiler=node_profiler + ) as profiler: # Receive metadata. input_metadata = await recv_metadata(ch_in, context) nranks = context.comm().nranks @@ -148,8 +150,8 @@ async def concatenate_node( # Extract concatenated result result_table = await allgather.extract_concatenated(stream) - if node_profiler is not None: - node_profiler.add_chunk(table=result_table) + if profiler is not None: + profiler.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. @@ -207,8 +209,8 @@ async def concatenate_node( ), context=ir_context, ) - if node_profiler is not None: - node_profiler.add_chunk(table=df.table) + if profiler is not None: + profiler.add_chunk(table=df.table) await ch_out.send( context, Message( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 3a8d0522be53..f7d50b292fd5 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -5,11 +5,20 @@ from __future__ import annotations import asyncio +import hashlib import operator from contextlib import asynccontextmanager, contextmanager 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, @@ -33,30 +42,95 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.dsl.ir import IR + from cudf_polars.experimental.base import RuntimeNodeProfiler from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator from cudf_polars.typing import DataType +def _stable_ir_id(ir_node: IR) -> int: + """ + Compute a stable identifier for an IR node. + + This identifier is based on the node's content (type + schema + children's IDs) + and is stable across process boundaries. Uses hashlib instead of Python's + built-in hash() because hash() uses a random per-process seed (PYTHONHASHSEED). + + Parameters + ---------- + ir_node + The IR node. + + Returns + ------- + int + A stable 64-bit identifier for this node. + """ + type_name = type(ir_node).__name__ + schema_keys = tuple(ir_node.schema.keys()) + children_ids = tuple(_stable_ir_id(child) for child in ir_node.children) + content = repr((type_name, schema_keys, children_ids)).encode("utf-8") + # Use first 16 hex chars (64 bits) for a more concise ID + return int(hashlib.md5(content).hexdigest()[:16], 16) + + @asynccontextmanager async def shutdown_on_error( - context: Context, *channels: Channel[Any] -) -> AsyncIterator[None]: + context: Context, + *channels: Channel[Any], + ir: IR | None = None, + node_profiler: RuntimeNodeProfiler | None = None, +) -> AsyncIterator[RuntimeNodeProfiler | None]: """ Shutdown on error for rapidsmpf. + This context manager handles channel cleanup on errors and optionally + manages node profiling with structlog tracing integration. + Parameters ---------- context The rapidsmpf context. channels - The channels to shutdown. + The channels to shutdown on error. + ir + The IR node being executed (for tracing). + node_profiler + Node profiler for collecting runtime statistics. + + Yields + ------ + RuntimeNodeProfiler | None + The node profiler (if provided) for use within the context. """ - # TODO: This probably belongs in rapidsmpf. + # Setup tracing if enabled and ir is provided + ir_id: int | None = None + ir_type: str | None = None + if ir is not None and LOG_TRACES: + ir_id = _stable_ir_id(ir) + ir_type = type(ir).__name__ + structlog.contextvars.bind_contextvars(ir_id=ir_id) + try: - yield + yield node_profiler 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: + log = structlog.get_logger() + record: dict[str, Any] = { + "ir_id": ir_id, + "ir_type": ir_type, + } + if node_profiler is not None: + record["chunks"] = node_profiler.chunk_count + if node_profiler.row_count is not None: + record["rows"] = node_profiler.row_count + if node_profiler.decision is not None: + record["decision"] = node_profiler.decision + log.info("Streaming Node", **record) + structlog.contextvars.unbind_contextvars("ir_id") def remap_partitioning( From 1c80b1bb069e2bfda1e2ecf6984c9b8fa5cd0272 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 06:35:03 -0800 Subject: [PATCH 17/25] add structlog events to RuntimeNodeProfiler --- .../cudf_polars/experimental/base.py | 63 +++++++++++++++++-- .../cudf_polars/experimental/rapidsmpf/io.py | 5 +- .../experimental/rapidsmpf/nodes.py | 9 +-- .../experimental/rapidsmpf/repartition.py | 6 +- .../experimental/rapidsmpf/utils.py | 56 ++++------------- .../experimental/test_runtime_profiler.py | 3 + 6 files changed, 85 insertions(+), 57 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py index c51309bb4db1..681196afca29 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/experimental/base.py @@ -6,6 +6,7 @@ import dataclasses import enum +import hashlib from collections import defaultdict from enum import IntEnum from functools import cached_property @@ -408,12 +409,37 @@ def __init__(self) -> None: self.join_info = JoinInfo() +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 RuntimeNodeProfiler: """ Profiler for a single IR node. 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. @@ -422,14 +448,27 @@ class RuntimeNodeProfiler: 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", "row_count") + __slots__ = ( + "chunk_count", + "decision", + "duplicated", + "ir_id", + "ir_type", + "row_count", + ) - def __init__(self) -> None: + 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: """ @@ -442,10 +481,24 @@ def add_chunk(self, *, table: plc.Table | None = None) -> None: 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: RuntimeNodeProfiler) -> None: """Merge another node profiler's stats into this one.""" if other.row_count is not None: - self.row_count = (self.row_count or 0) + other.row_count + if self.duplicated or other.duplicated: + # For duplicated data, take max (don't sum across ranks) + self.row_count = max(self.row_count or 0, other.row_count) + self.duplicated = True + else: + self.row_count = (self.row_count or 0) + other.row_count self.chunk_count += other.chunk_count if other.decision is not None: self.decision = other.decision @@ -475,7 +528,9 @@ def get_or_create(self, ir: IR) -> RuntimeNodeProfiler: was profiled without creating an entry, use `node_profilers.get(ir)`. """ if ir not in self.node_profilers: - self.node_profilers[ir] = RuntimeNodeProfiler() + ir_id = _stable_ir_id(ir) + ir_type = type(ir).__name__ + self.node_profilers[ir] = RuntimeNodeProfiler(ir_id, ir_type) return self.node_profilers[ir] def merge(self, other: RuntimeQueryProfiler) -> None: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py index 3cfc289579ca..c21ee5ac0f8f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -175,7 +175,7 @@ async def dataframescan_node( The node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, ch_out, ir=ir, node_profiler=node_profiler + context, ch_out, node_profiler=node_profiler ) as profiler: # Find local partition count. nrows = ir.df.shape()[0] @@ -419,7 +419,7 @@ async def scan_node( The node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, ch_out, ir=ir, node_profiler=node_profiler + context, ch_out, node_profiler=node_profiler ) as profiler: # Build a list of local Scan operations scans: list[Scan | SplitScan] = [] @@ -711,7 +711,6 @@ def _( # node does not send metadata. metadata_node = metadata_feeder_node( rec.state["context"], - ir, ch_in, ch_out, ChannelMetadata( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index df7d8b7a19a4..85f0e785b1f0 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -76,7 +76,7 @@ async def default_node_single( Chunks are processed in the order they are received. """ async with shutdown_on_error( - context, ch_in, ch_out, ir=ir, node_profiler=node_profiler + context, ch_in, ch_out, node_profiler=node_profiler ) as profiler: # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) @@ -178,7 +178,7 @@ async def default_node_multi( Node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, *chs_in, ch_out, ir=ir, node_profiler=node_profiler + context, *chs_in, ch_out, node_profiler=node_profiler ) as profiler: # Merge and forward basic metadata. local_count = 1 @@ -690,7 +690,6 @@ def generate_ir_sub_network_wrapper( @define_py_node() async def metadata_feeder_node( context: Context, - ir: IR, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk], metadata: ChannelMetadata, @@ -703,8 +702,6 @@ async def metadata_feeder_node( ---------- context The rapidsmpf context. - ir - The IR node. ch_in The input channel to pull data from. ch_out @@ -715,7 +712,7 @@ async def metadata_feeder_node( Node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, ch_in, ch_out, ir=ir, node_profiler=node_profiler + context, ch_in, ch_out, node_profiler=node_profiler ) as profiler: await send_metadata(ch_out, context, metadata) while (msg := await ch_in.recv(context)) is not None: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py index a8a1de7c49c4..83935ceaa127 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -81,7 +81,7 @@ async def concatenate_node( Node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, ch_in, ch_out, ir=ir, node_profiler=node_profiler + context, ch_in, ch_out, node_profiler=node_profiler ) as profiler: # Receive metadata. input_metadata = await recv_metadata(ch_in, context) @@ -138,6 +138,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) + if profiler is not None and output_duplicated: + profiler.set_duplicated() allgather = AllGatherManager(context, collective_id) stream = context.get_stream_from_pool() @@ -172,6 +174,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) + if profiler is not None and output_duplicated: + profiler.set_duplicated() # Local repartitioning seq_num = 0 diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index f7d50b292fd5..3b05b17d8ac6 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -5,7 +5,6 @@ from __future__ import annotations import asyncio -import hashlib import operator from contextlib import asynccontextmanager, contextmanager from functools import reduce @@ -47,37 +46,10 @@ from cudf_polars.typing import DataType -def _stable_ir_id(ir_node: IR) -> int: - """ - Compute a stable identifier for an IR node. - - This identifier is based on the node's content (type + schema + children's IDs) - and is stable across process boundaries. Uses hashlib instead of Python's - built-in hash() because hash() uses a random per-process seed (PYTHONHASHSEED). - - Parameters - ---------- - ir_node - The IR node. - - Returns - ------- - int - A stable 64-bit identifier for this node. - """ - type_name = type(ir_node).__name__ - schema_keys = tuple(ir_node.schema.keys()) - children_ids = tuple(_stable_ir_id(child) for child in ir_node.children) - content = repr((type_name, schema_keys, children_ids)).encode("utf-8") - # Use first 16 hex chars (64 bits) for a more concise ID - return int(hashlib.md5(content).hexdigest()[:16], 16) - - @asynccontextmanager async def shutdown_on_error( context: Context, *channels: Channel[Any], - ir: IR | None = None, node_profiler: RuntimeNodeProfiler | None = None, ) -> AsyncIterator[RuntimeNodeProfiler | None]: """ @@ -92,22 +64,19 @@ async def shutdown_on_error( The rapidsmpf context. channels The channels to shutdown on error. - ir - The IR node being executed (for tracing). node_profiler - Node profiler for collecting runtime statistics. + Optional node profiler for collecting runtime statistics. + If provided and tracing is enabled, structlog events are emitted. Yields ------ RuntimeNodeProfiler | None The node profiler (if provided) for use within the context. """ - # Setup tracing if enabled and ir is provided + # Setup tracing if enabled and profiler has ir_id ir_id: int | None = None - ir_type: str | None = None - if ir is not None and LOG_TRACES: - ir_id = _stable_ir_id(ir) - ir_type = type(ir).__name__ + if node_profiler is not None and node_profiler.ir_id is not None and LOG_TRACES: + ir_id = node_profiler.ir_id structlog.contextvars.bind_contextvars(ir_id=ir_id) try: @@ -118,17 +87,18 @@ async def shutdown_on_error( finally: # Emit structlog event on exit if tracing is enabled if ir_id is not None and LOG_TRACES: + assert node_profiler is not None # ir_id implies node_profiler exists log = structlog.get_logger() record: dict[str, Any] = { "ir_id": ir_id, - "ir_type": ir_type, + "ir_type": node_profiler.ir_type, + "chunks": node_profiler.chunk_count, + "duplicated": node_profiler.duplicated, } - if node_profiler is not None: - record["chunks"] = node_profiler.chunk_count - if node_profiler.row_count is not None: - record["rows"] = node_profiler.row_count - if node_profiler.decision is not None: - record["decision"] = node_profiler.decision + if node_profiler.row_count is not None: + record["rows"] = node_profiler.row_count + if node_profiler.decision is not None: + record["decision"] = node_profiler.decision log.info("Streaming Node", **record) structlog.contextvars.unbind_contextvars("ir_id") diff --git a/python/cudf_polars/tests/experimental/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/test_runtime_profiler.py index e66c405057df..53c5df6a1d23 100644 --- a/python/cudf_polars/tests/experimental/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/test_runtime_profiler.py @@ -19,6 +19,9 @@ class MockIR: 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 profiler accumulation From 13475629a31981bb36e5c13f8c09a04c48d11ca6 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 06:53:14 -0800 Subject: [PATCH 18/25] use 'tracing' name for consistency --- .../cudf_polars/experimental/base.py | 135 +--------------- .../cudf_polars/experimental/explain.py | 35 ++--- .../experimental/rapidsmpf/core.py | 10 +- .../experimental/rapidsmpf/dask.py | 15 +- .../experimental/rapidsmpf/dispatch.py | 9 +- .../cudf_polars/experimental/rapidsmpf/io.py | 53 +++---- .../experimental/rapidsmpf/nodes.py | 37 +++-- .../experimental/rapidsmpf/repartition.py | 10 +- .../experimental/rapidsmpf/tracing.py | 145 ++++++++++++++++++ .../experimental/rapidsmpf/utils.py | 40 ++--- .../experimental/test_runtime_profiler.py | 47 +++--- 11 files changed, 272 insertions(+), 264 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py index 681196afca29..35490c8ce7b9 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/experimental/base.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Multi-partition base classes.""" @@ -6,7 +6,6 @@ import dataclasses import enum -import hashlib from collections import defaultdict from enum import IntEnum from functools import cached_property @@ -15,8 +14,6 @@ if TYPE_CHECKING: from collections.abc import Generator, Iterator, MutableMapping - import pylibcudf as plc - from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR from cudf_polars.dsl.nodebase import Node @@ -409,136 +406,6 @@ def __init__(self) -> None: self.join_info = JoinInfo() -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 RuntimeNodeProfiler: - """ - Profiler for a single IR node. - - 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: RuntimeNodeProfiler) -> None: - """Merge another node profiler's stats into this one.""" - if other.row_count is not None: - if self.duplicated or other.duplicated: - # For duplicated data, take max (don't sum across ranks) - self.row_count = max(self.row_count or 0, other.row_count) - self.duplicated = True - else: - self.row_count = (self.row_count or 0) + other.row_count - self.chunk_count += other.chunk_count - if other.decision is not None: - self.decision = other.decision - - -class RuntimeQueryProfiler: - """ - Profiler for collecting runtime statistics for an entire query. - - Attributes - ---------- - node_profilers - Mapping from each IR node to its node profiler. - """ - - __slots__ = ("node_profilers",) - node_profilers: dict[IR, RuntimeNodeProfiler] - - def __init__(self) -> None: - self.node_profilers = {} - - def get_or_create(self, ir: IR) -> RuntimeNodeProfiler: - """ - Get or create a node profiler for the given IR. - - Use this when setting up profiling for a node. To check if a node - was profiled without creating an entry, use `node_profilers.get(ir)`. - """ - if ir not in self.node_profilers: - ir_id = _stable_ir_id(ir) - ir_type = type(ir).__name__ - self.node_profilers[ir] = RuntimeNodeProfiler(ir_id, ir_type) - return self.node_profilers[ir] - - def merge(self, other: RuntimeQueryProfiler) -> None: - """Merge another query profiler's statistics into this one.""" - for ir, node_profiler in other.node_profilers.items(): - self.get_or_create(ir).merge(node_profiler) - - class IOPartitionFlavor(IntEnum): """Flavor of IO partitioning.""" diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 5df6ae626871..fb4f6ca75ecb 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -31,11 +31,8 @@ import polars as pl from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import ( - PartitionInfo, - RuntimeQueryProfiler, - StatsCollector, - ) + from cudf_polars.experimental.base import PartitionInfo, StatsCollector + from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer def explain_query( @@ -179,7 +176,7 @@ def write_profile_output( profile_output: str | Path, ir: IR, partition_info: MutableMapping[IR, PartitionInfo], - profiler: RuntimeQueryProfiler, + tracer: StreamingQueryTracer, ) -> None: """ Write a post-execution profile showing actual row counts and decisions. @@ -192,41 +189,41 @@ def write_profile_output( The lowered IR root node. partition_info Partition information for the IR nodes. - profiler - The profiler with actual row counts and decisions from execution. + tracer + The tracer with actual row counts and decisions from execution. """ from pathlib import Path - profile_repr = _repr_profile_tree(ir, partition_info, profiler) + profile_repr = _repr_profile_tree(ir, partition_info, tracer) Path(profile_output).write_text(profile_repr) def _repr_profile_tree( ir: IR, partition_info: MutableMapping[IR, PartitionInfo], - profiler: RuntimeQueryProfiler, + tracer: StreamingQueryTracer, *, offset: str = "", ) -> str: - """Recursively build a tree representation with profiler data.""" + """Recursively build a tree representation with tracer data.""" header = _repr_ir(ir, offset=offset) header = header.rstrip("\n") - # Get node profiler if it exists - if (node_profiler := profiler.node_profilers.get(ir)) is not None: + # 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_profiler.row_count is not None: - header += f" rows={_fmt_row_count(node_profiler.row_count)}" + if node_tracer.row_count is not None: + header += f" rows={_fmt_row_count(node_tracer.row_count)}" # Add decision if present - if node_profiler.decision is not None: - header += f" decision={node_profiler.decision}" + if node_tracer.decision is not None: + header += f" decision={node_tracer.decision}" # Add actual chunk count - header += f" chunks={node_profiler.chunk_count}" + header += f" chunks={node_tracer.chunk_count}" children_strs = [ - _repr_profile_tree(child, partition_info, profiler, offset=offset + " ") + _repr_profile_tree(child, partition_info, tracer, offset=offset + " ") for child in ir.children ] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index aaeb2bd2f033..89716b54c997 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -34,13 +34,13 @@ 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.base import RuntimeQueryProfiler from cudf_polars.experimental.rapidsmpf.collectives import ReserveOpIDs from cudf_polars.experimental.rapidsmpf.dispatch import FanoutInfo, lower_ir_node from cudf_polars.experimental.rapidsmpf.nodes import ( 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 @@ -152,7 +152,7 @@ def evaluate_pipeline( rmpf_context: Context | None = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, RuntimeQueryProfiler | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: """ Build and evaluate a RapidsMPF streaming pipeline. @@ -245,8 +245,8 @@ def evaluate_pipeline( metadata_collector: list[ChannelMetadata] | None = ( [] if collect_metadata else None ) - profiler: RuntimeQueryProfiler | None = ( - RuntimeQueryProfiler() + profiler: StreamingQueryTracer | None = ( + StreamingQueryTracer() if config_options.executor.profiling is not None or LOG_TRACES else None ) @@ -435,7 +435,7 @@ def generate_network( ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], metadata_collector: list[ChannelMetadata] | None, - profiler: RuntimeQueryProfiler | None = None, + profiler: StreamingQueryTracer | None = None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py index cb9fe70ceedd..24595f1c7589 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py @@ -22,12 +22,9 @@ from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import ( - PartitionInfo, - RuntimeQueryProfiler, - StatsCollector, - ) + 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): @@ -43,7 +40,7 @@ def __call__( rmpf_context: Context | None = None, *, collect_metadata: bool = False, - ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, RuntimeQueryProfiler | None]: + ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: """Evaluate a pipeline and return the result DataFrame, metadata, and profiler.""" ... @@ -65,7 +62,7 @@ def evaluate_pipeline_dask( collective_id_map: dict[IR, list[int]], *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, RuntimeQueryProfiler | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: """ Evaluate a RapidsMPF streaming pipeline on a Dask cluster. @@ -103,7 +100,7 @@ def evaluate_pipeline_dask( ) dfs: list[pl.DataFrame] = [] metadata_collector: list[ChannelMetadata] = [] - merged_profiler: RuntimeQueryProfiler | None = None + merged_profiler: StreamingQueryTracer | None = None for df, md, profiler in result.values(): dfs.append(df) if md is not None: @@ -127,7 +124,7 @@ def _evaluate_pipeline_dask( dask_worker: Any = None, *, collect_metadata: bool = False, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, RuntimeQueryProfiler | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: """ Build and evaluate a RapidsMPF streaming pipeline. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py index 9403293e1785..3d96bae38379 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py @@ -15,11 +15,8 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import ( - PartitionInfo, - RuntimeQueryProfiler, - 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 @@ -90,7 +87,7 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] - profiler: RuntimeQueryProfiler | None + profiler: 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 c21ee5ac0f8f..dda9e6609815 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -51,14 +51,13 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import ( - ColumnStat, - RuntimeNodeProfiler, - RuntimeQueryProfiler, - StatsCollector, - ) + 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 @@ -149,7 +148,7 @@ async def dataframescan_node( num_producers: int, rows_per_partition: int, estimated_chunk_bytes: int, - node_profiler: RuntimeNodeProfiler | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ DataFrameScan node for rapidsmpf. @@ -171,12 +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_profiler + node_tracer The node profiler for collecting runtime statistics. """ - async with shutdown_on_error( - context, ch_out, node_profiler=node_profiler - ) as profiler: + async with shutdown_on_error(context, ch_out, node_tracer=node_tracer) as profiler: # Find local partition count. nrows = ir.df.shape()[0] global_count = math.ceil(nrows / rows_per_partition) if nrows > 0 else 0 @@ -226,7 +223,7 @@ async def dataframescan_node( ch_out, ir_context, estimated_chunk_bytes, - node_profiler=profiler, + node_tracer=profiler, ) await ch_out.drain(context) return @@ -252,7 +249,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, - node_profiler=profiler, + node_tracer=profiler, ) await ch_out.drain(context) @@ -278,7 +275,7 @@ def _( context = rec.state["context"] ir_context = rec.state["ir_context"] - profiler: RuntimeQueryProfiler | None = rec.state["profiler"] + profiler: StreamingQueryTracer | None = rec.state["profiler"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} nodes: dict[IR, list[Any]] = { ir: [ @@ -290,7 +287,7 @@ def _( num_producers=num_producers, rows_per_partition=rows_per_partition, estimated_chunk_bytes=estimated_chunk_bytes, - node_profiler=profiler.get_or_create(ir) if profiler else None, + node_tracer=profiler.get_or_create(ir) if profiler else None, ) ] } @@ -336,7 +333,7 @@ async def read_chunk( ch_out: Channel[TableChunk], ir_context: IRExecutionContext, estimated_chunk_bytes: int, - node_profiler: RuntimeNodeProfiler | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Read a chunk from disk and send it to the output channel. @@ -356,7 +353,7 @@ 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_profiler + node_tracer The node profiler for collecting runtime statistics. """ with opaque_reservation(context, estimated_chunk_bytes): @@ -365,8 +362,8 @@ async def read_chunk( *scan._non_child_args, context=ir_context, ) - if node_profiler is not None: - node_profiler.add_chunk(table=df.table) + if node_tracer is not None: + node_tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -391,7 +388,7 @@ async def scan_node( plan: IOPartitionPlan, parquet_options: ParquetOptions, estimated_chunk_bytes: int, - node_profiler: RuntimeNodeProfiler | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Scan node for rapidsmpf. @@ -415,12 +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_profiler + node_tracer The node profiler for collecting runtime statistics. """ - async with shutdown_on_error( - context, ch_out, node_profiler=node_profiler - ) as profiler: + async with shutdown_on_error(context, ch_out, node_tracer=node_tracer) as profiler: # Build a list of local Scan operations scans: list[Scan | SplitScan] = [] if plan.flavor == IOPartitionFlavor.SPLIT_FILES: @@ -511,7 +506,7 @@ async def scan_node( ch_out, ir_context, estimated_chunk_bytes, - node_profiler=profiler, + node_tracer=profiler, ) await ch_out.drain(context) return @@ -537,7 +532,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, - node_profiler=profiler, + node_tracer=profiler, ) await ch_out.drain(context) @@ -668,7 +663,7 @@ def _( parquet_options = config_options.parquet_options partition_info = rec.state["partition_info"][ir] num_producers = rec.state["max_io_threads"] - profiler: RuntimeQueryProfiler | None = rec.state["profiler"] + profiler: StreamingQueryTracer | None = rec.state["profiler"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} assert partition_info.io_plan is not None, "Scan node must have a partition plan" @@ -720,7 +715,7 @@ def _( partition_info.count / rec.state["context"].comm().nranks ), ), - node_profiler=profiler.get_or_create(ir) if profiler else None, + node_tracer=profiler.get_or_create(ir) if profiler else None, ) nodes[ir] = [native_node, metadata_node] else: @@ -737,7 +732,7 @@ def _( plan=plan, parquet_options=parquet_options, estimated_chunk_bytes=executor.target_partition_size, - node_profiler=profiler.get_or_create(ir) if profiler else None, + node_tracer=profiler.get_or_create(ir) if profiler 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 85f0e785b1f0..2728d0f47045 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -36,8 +36,11 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.base import RuntimeNodeProfiler, RuntimeQueryProfiler from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import ( + StreamingNodeTracer, + StreamingQueryTracer, + ) @define_py_node() @@ -49,7 +52,7 @@ async def default_node_single( ch_in: Channel[TableChunk], *, preserve_partitioning: bool = False, - node_profiler: RuntimeNodeProfiler | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Single-channel default node for rapidsmpf. @@ -68,7 +71,7 @@ async def default_node_single( The input Channel[TableChunk]. preserve_partitioning Whether to preserve the partitioning metadata of the input chunks. - node_profiler + node_tracer Node profiler for collecting runtime statistics. Notes @@ -76,7 +79,7 @@ async def default_node_single( Chunks are processed in the order they are received. """ async with shutdown_on_error( - context, ch_in, ch_out, node_profiler=node_profiler + context, ch_in, ch_out, node_tracer=node_tracer ) as profiler: # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) @@ -92,6 +95,8 @@ async def default_node_single( duplicated=metadata_in.duplicated, ) await send_metadata(ch_out, context, metadata_out) + if profiler is not None and metadata_in.duplicated: + profiler.set_duplicated() # Recv/send data. seq_num = 0 @@ -154,7 +159,7 @@ async def default_node_multi( chs_in: tuple[Channel[TableChunk], ...], *, partitioning_index: int | None = None, - node_profiler: RuntimeNodeProfiler | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Pointwise node for rapidsmpf. @@ -174,11 +179,11 @@ 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_profiler + node_tracer Node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, *chs_in, ch_out, node_profiler=node_profiler + context, *chs_in, ch_out, node_tracer=node_tracer ) as profiler: # Merge and forward basic metadata. local_count = 1 @@ -202,6 +207,8 @@ async def default_node_multi( duplicated=duplicated, ) await send_metadata(ch_out, context, metadata) + if profiler is not None and duplicated: + profiler.set_duplicated() seq_num = 0 n_children = len(chs_in) @@ -546,8 +553,8 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) - profiler: RuntimeQueryProfiler | None = rec.state.get("profiler") - node_profiler = profiler.get_or_create(ir) if profiler is not None else None + profiler: StreamingQueryTracer | None = rec.state.get("profiler") + node_tracer = profiler.get_or_create(ir) if profiler is not None else None if len(ir.children) == 1: # Single-channel default node @@ -567,7 +574,7 @@ def _( channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), preserve_partitioning=preserve_partitioning, - node_profiler=node_profiler, + node_tracer=node_tracer, ) ] else: @@ -579,7 +586,7 @@ def _( rec.state["ir_context"], channels[ir].reserve_input_slot(), tuple(channels[c].reserve_output_slot() for c in ir.children), - node_profiler=node_profiler, + node_tracer=node_tracer, ) ] @@ -693,7 +700,7 @@ async def metadata_feeder_node( ch_in: Channel[TableChunk], ch_out: Channel[TableChunk], metadata: ChannelMetadata, - node_profiler: RuntimeNodeProfiler | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Forward data with new metadata. @@ -708,13 +715,15 @@ 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_profiler + node_tracer Node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, ch_in, ch_out, node_profiler=node_profiler + context, ch_in, ch_out, node_tracer=node_tracer ) as profiler: await send_metadata(ch_out, context, metadata) + if profiler is not None and metadata.duplicated: + profiler.set_duplicated() while (msg := await ch_in.recv(context)) is not None: await ch_out.send(context, msg) if profiler is not None: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py index 83935ceaa127..af6b860b8d33 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -32,8 +32,8 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import RuntimeNodeProfiler from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import StreamingNodeTracer @define_py_node() @@ -46,7 +46,7 @@ async def concatenate_node( *, output_count: int, collective_id: int, - node_profiler: RuntimeNodeProfiler | None = None, + node_tracer: StreamingNodeTracer | None = None, ) -> None: """ Concatenate node for rapidsmpf. @@ -77,11 +77,11 @@ async def concatenate_node( The expected global number of output chunks. collective_id Pre-allocated collective ID for this operation. - node_profiler + node_tracer Node profiler for collecting runtime statistics. """ async with shutdown_on_error( - context, ch_in, ch_out, node_profiler=node_profiler + context, ch_in, ch_out, node_tracer=node_tracer ) as profiler: # Receive metadata. input_metadata = await recv_metadata(ch_in, context) @@ -267,7 +267,7 @@ def _( channels[ir.children[0]].reserve_output_slot(), output_count=partition_info[ir].count, collective_id=collective_id, - node_profiler=profiler.get_or_create(ir) if profiler else None, + node_tracer=profiler.get_or_create(ir) if profiler 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..926c17effa9a --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py @@ -0,0 +1,145 @@ +# 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.""" + if other.row_count is not None: + if self.duplicated or other.duplicated: + # For duplicated data, take max (don't sum across ranks) + self.row_count = max(self.row_count or 0, other.row_count) + self.duplicated = True + else: + self.row_count = (self.row_count or 0) + other.row_count + self.chunk_count += other.chunk_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 3b05b17d8ac6..2e13d0fad51c 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -41,8 +41,8 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import RuntimeNodeProfiler from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import StreamingNodeTracer from cudf_polars.typing import DataType @@ -50,13 +50,13 @@ async def shutdown_on_error( context: Context, *channels: Channel[Any], - node_profiler: RuntimeNodeProfiler | None = None, -) -> AsyncIterator[RuntimeNodeProfiler | None]: + 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 profiling with structlog tracing integration. + manages node tracing with structlog integration. Parameters ---------- @@ -64,41 +64,41 @@ async def shutdown_on_error( The rapidsmpf context. channels The channels to shutdown on error. - node_profiler - Optional node profiler for collecting runtime statistics. + node_tracer + Optional node tracer for collecting runtime statistics. If provided and tracing is enabled, structlog events are emitted. Yields ------ - RuntimeNodeProfiler | None - The node profiler (if provided) for use within the context. + StreamingNodeTracer | None + The node tracer (if provided) for use within the context. """ - # Setup tracing if enabled and profiler has ir_id + # Setup tracing if enabled and tracer has ir_id ir_id: int | None = None - if node_profiler is not None and node_profiler.ir_id is not None and LOG_TRACES: - ir_id = node_profiler.ir_id + 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 node_profiler + 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_profiler is not None # ir_id implies node_profiler exists + 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_profiler.ir_type, - "chunks": node_profiler.chunk_count, - "duplicated": node_profiler.duplicated, + "ir_type": node_tracer.ir_type, + "chunks": node_tracer.chunk_count, + "duplicated": node_tracer.duplicated, } - if node_profiler.row_count is not None: - record["rows"] = node_profiler.row_count - if node_profiler.decision is not None: - record["decision"] = node_profiler.decision + 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") diff --git a/python/cudf_polars/tests/experimental/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/test_runtime_profiler.py index 53c5df6a1d23..3aecfa566275 100644 --- a/python/cudf_polars/tests/experimental/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/test_runtime_profiler.py @@ -1,17 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for RuntimeQueryProfiler and related output functions.""" +"""Unit tests for StreamingQueryTracer and related output functions.""" from __future__ import annotations import polars as pl -from cudf_polars.experimental.base import PartitionInfo, RuntimeQueryProfiler +from cudf_polars.experimental.base import PartitionInfo from cudf_polars.experimental.explain import _repr_profile_tree, write_profile_output +from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer -def test_runtime_profiler_and_output(tmp_path): - """Test RuntimeQueryProfiler, merge, and output formatting.""" +def test_streaming_query_tracer_and_output(tmp_path): + """Test StreamingQueryTracer, merge, and output formatting.""" class MockIR: children = () @@ -24,36 +25,36 @@ def get_hashable(self): ir1, ir2 = MockIR(), MockIR() - # Test node profiler accumulation - profiler1 = RuntimeQueryProfiler() - np1 = profiler1.get_or_create(ir1) - np1.row_count = 100 - np1.chunk_count = 5 - np1.decision = "shuffle" - - profiler2 = RuntimeQueryProfiler() - profiler2.get_or_create(ir1).row_count = 150 - profiler2.get_or_create(ir1).chunk_count = 3 - profiler2.get_or_create(ir1).decision = "shuffle" - profiler2.get_or_create(ir2).row_count = 200 - profiler2.get_or_create(ir2).chunk_count = 3 - profiler2.get_or_create(ir2).add_chunk() + # 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 - profiler1.merge(profiler2) - assert profiler1.node_profilers[ir1].row_count == 250 - assert profiler1.node_profilers[ir1].chunk_count == 8 + tracer1.merge(tracer2) + assert tracer1.node_tracers[ir1].row_count == 250 + assert tracer1.node_tracers[ir1].chunk_count == 8 # Test _repr_profile_tree output format partition_info = {ir1: PartitionInfo(count=4)} - output = _repr_profile_tree(ir1, partition_info, profiler1) + output = _repr_profile_tree(ir1, partition_info, tracer1) assert "rows=250" in output assert "chunks=8" in output assert "decision=shuffle" in output # Test write_profile_output output_path = tmp_path / "profile.txt" - write_profile_output(output_path, ir1, partition_info, profiler1) + write_profile_output(output_path, ir1, partition_info, tracer1) assert output_path.exists() content = output_path.read_text() assert "rows=250" in content From 96c5739b7719a2bda62ed42fb0a2b2f6bc884f1f Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 07:11:44 -0800 Subject: [PATCH 19/25] adopt 'tracing' terminology --- .../experimental/benchmarks/utils.py | 16 ++--- .../experimental/rapidsmpf/core.py | 32 ++++----- .../experimental/rapidsmpf/dask.py | 20 +++--- .../experimental/rapidsmpf/dispatch.py | 6 +- .../cudf_polars/experimental/rapidsmpf/io.py | 28 ++++---- .../experimental/rapidsmpf/nodes.py | 40 +++++------ .../experimental/rapidsmpf/repartition.py | 24 +++---- .../cudf_polars/cudf_polars/utils/config.py | 51 +++++++------- ...st_runtime_profiler.py => test_tracing.py} | 66 +++++++++++++++---- ...st_runtime_profiler.py => test_tracing.py} | 0 python/cudf_polars/tests/test_config.py | 22 +++---- 11 files changed, 175 insertions(+), 130 deletions(-) rename python/cudf_polars/tests/experimental/rapidsmpf/{test_runtime_profiler.py => test_tracing.py} (55%) rename python/cudf_polars/tests/experimental/{test_runtime_profiler.py => test_tracing.py} (100%) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 61b1d0881c1e..74ec91ade58e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -257,7 +257,7 @@ class RunConfig: collect_traces: bool = False stats_planning: bool dynamic_planning: bool | None = None - profile_output_path: str | None = None + trace_output_path: str | None = None max_io_threads: int native_parquet: bool spill_to_pinned_memory: bool @@ -378,7 +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, - profile_output_path=args.profile_output_path, + trace_output_path=args.trace_output_path, max_io_threads=args.max_io_threads, native_parquet=args.native_parquet, extra_info=args.extra_info, @@ -479,10 +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.profile_output_path is not None: - executor_options["profiling"] = { - "output_path": run_config.profile_output_path - } + if run_config.trace_output_path is not None: + executor_options["tracing"] = {"output_path": run_config.trace_output_path} if ( benchmark @@ -981,11 +979,11 @@ def parse_args( help="Enable dynamic shuffle planning (not yet implemented). ", ) parser.add_argument( - "--profile-output-path", - dest="profile_output_path", + "--trace-output-path", + dest="trace_output_path", type=str, default=None, - help="Path to write profiling output (row counts per node).", + help="Path to write tracing output (row counts per node).", ) parser.add_argument( "--max-io-threads", diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 89716b54c997..e0ad4d2fe4e1 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -109,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, profiler = evaluate_pipeline_dask( + result, metadata_collector, tracer = evaluate_pipeline_dask( evaluate_pipeline, ir, partition_info, @@ -120,7 +120,7 @@ def evaluate_logical_plan( ) else: # Single-process execution: Run locally - result, metadata_collector, profiler = evaluate_pipeline( + result, metadata_collector, tracer = evaluate_pipeline( ir, partition_info, config_options, @@ -129,16 +129,12 @@ def evaluate_logical_plan( collect_metadata=collect_metadata, ) - # Write profiler output if configured - profiling = config_options.executor.profiling - if ( - profiling is not None - and profiling.output_path is not None - and profiler is not None - ): + # 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_profile_output - write_profile_output(profiling.output_path, ir, partition_info, profiler) + write_profile_output(tracing.output_path, ir, partition_info, tracer) return result, metadata_collector @@ -175,7 +171,7 @@ def evaluate_pipeline( Returns ------- - The output DataFrame, metadata collector, and profiler. + 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" @@ -245,9 +241,9 @@ def evaluate_pipeline( metadata_collector: list[ChannelMetadata] | None = ( [] if collect_metadata else None ) - profiler: StreamingQueryTracer | None = ( + tracer: StreamingQueryTracer | None = ( StreamingQueryTracer() - if config_options.executor.profiling is not None or LOG_TRACES + if config_options.executor.tracing is not None or LOG_TRACES else None ) nodes, output = generate_network( @@ -259,7 +255,7 @@ def evaluate_pipeline( ir_context=ir_context, collective_id_map=collective_id_map, metadata_collector=metadata_collector, - profiler=profiler, + tracer=tracer, ) # Run the network @@ -313,7 +309,7 @@ def evaluate_pipeline( if _initial_mr is not None: rmm.mr.set_current_device_resource(_original_mr) - return result, metadata_collector, profiler + return result, metadata_collector, tracer def lower_ir_graph( @@ -435,7 +431,7 @@ def generate_network( ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], metadata_collector: list[ChannelMetadata] | None, - profiler: StreamingQueryTracer | None = None, + tracer: StreamingQueryTracer | None = None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. @@ -460,7 +456,7 @@ 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. - profiler + tracer Profiler for collecting runtime statistics. Returns @@ -494,7 +490,7 @@ def generate_network( "max_io_threads": max_io_threads_local, "stats": stats, "collective_id_map": collective_id_map, - "profiler": profiler, + "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 24595f1c7589..565180d7752b 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dask.py @@ -41,7 +41,7 @@ def __call__( *, collect_metadata: bool = False, ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None, StreamingQueryTracer | None]: - """Evaluate a pipeline and return the result DataFrame, metadata, and profiler.""" + """Evaluate a pipeline and return the result DataFrame, metadata, and tracer.""" ... @@ -85,7 +85,7 @@ def evaluate_pipeline_dask( Returns ------- - The output DataFrame, metadata collector, and merged profiler. + The output DataFrame, metadata collector, and merged tracer. """ client = get_dask_client() result = client.run( @@ -100,18 +100,18 @@ def evaluate_pipeline_dask( ) dfs: list[pl.DataFrame] = [] metadata_collector: list[ChannelMetadata] = [] - merged_profiler: StreamingQueryTracer | None = None - for df, md, profiler 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 profiler is not None: - if merged_profiler is None: - merged_profiler = profiler + if tracer is not None: + if merged_tracer is None: + merged_tracer = tracer else: - merged_profiler.merge(profiler) + merged_tracer.merge(tracer) - return pl.concat(dfs), metadata_collector or None, merged_profiler + return pl.concat(dfs), metadata_collector or None, merged_tracer def _evaluate_pipeline_dask( @@ -151,7 +151,7 @@ def _evaluate_pipeline_dask( Returns ------- - The output DataFrame, metadata collector, and profiler. + 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 3d96bae38379..061a4cda89a7 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py @@ -75,8 +75,8 @@ class GenState(TypedDict): Statistics collector. collective_id_map The mapping of IR nodes to lists of collective IDs. - profiler - Runtime profiler for collecting execution statistics. + tracer + Runtime tracer for collecting execution statistics. """ context: Context @@ -87,7 +87,7 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] - profiler: StreamingQueryTracer | None + 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 dda9e6609815..d79c3dccec42 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -171,9 +171,9 @@ async def dataframescan_node( Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. node_tracer - The node profiler for collecting runtime statistics. + The node tracer for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_out, node_tracer=node_tracer) as profiler: + 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 @@ -223,7 +223,7 @@ async def dataframescan_node( ch_out, ir_context, estimated_chunk_bytes, - node_tracer=profiler, + node_tracer=tracer, ) await ch_out.drain(context) return @@ -249,7 +249,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, - node_tracer=profiler, + node_tracer=tracer, ) await ch_out.drain(context) @@ -275,7 +275,7 @@ def _( context = rec.state["context"] ir_context = rec.state["ir_context"] - profiler: StreamingQueryTracer | None = rec.state["profiler"] + tracer: StreamingQueryTracer | None = rec.state["tracer"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} nodes: dict[IR, list[Any]] = { ir: [ @@ -287,7 +287,7 @@ def _( num_producers=num_producers, rows_per_partition=rows_per_partition, estimated_chunk_bytes=estimated_chunk_bytes, - node_tracer=profiler.get_or_create(ir) if profiler else None, + node_tracer=tracer.get_or_create(ir) if tracer else None, ) ] } @@ -354,7 +354,7 @@ async def read_chunk( Estimated size of the chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. node_tracer - The node profiler for collecting runtime statistics. + The node tracer for collecting runtime statistics. """ with opaque_reservation(context, estimated_chunk_bytes): df = await asyncio.to_thread( @@ -413,9 +413,9 @@ async def scan_node( Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. node_tracer - The node profiler for collecting runtime statistics. + The node tracer for collecting runtime statistics. """ - async with shutdown_on_error(context, ch_out, node_tracer=node_tracer) as profiler: + 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: @@ -506,7 +506,7 @@ async def scan_node( ch_out, ir_context, estimated_chunk_bytes, - node_tracer=profiler, + node_tracer=tracer, ) await ch_out.drain(context) return @@ -532,7 +532,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, - node_tracer=profiler, + node_tracer=tracer, ) await ch_out.drain(context) @@ -663,7 +663,7 @@ def _( parquet_options = config_options.parquet_options partition_info = rec.state["partition_info"][ir] num_producers = rec.state["max_io_threads"] - profiler: StreamingQueryTracer | None = rec.state["profiler"] + 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" @@ -715,7 +715,7 @@ def _( partition_info.count / rec.state["context"].comm().nranks ), ), - node_tracer=profiler.get_or_create(ir) if profiler else None, + node_tracer=tracer.get_or_create(ir) if tracer else None, ) nodes[ir] = [native_node, metadata_node] else: @@ -732,7 +732,7 @@ def _( plan=plan, parquet_options=parquet_options, estimated_chunk_bytes=executor.target_partition_size, - node_tracer=profiler.get_or_create(ir) if profiler else None, + 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 2728d0f47045..8542633d9618 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -72,7 +72,7 @@ async def default_node_single( preserve_partitioning Whether to preserve the partitioning metadata of the input chunks. node_tracer - Node profiler for collecting runtime statistics. + Node tracer for collecting runtime statistics. Notes ----- @@ -80,7 +80,7 @@ async def default_node_single( """ async with shutdown_on_error( context, ch_in, ch_out, node_tracer=node_tracer - ) as profiler: + ) as tracer: # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) partitioning = None @@ -95,8 +95,8 @@ async def default_node_single( duplicated=metadata_in.duplicated, ) await send_metadata(ch_out, context, metadata_out) - if profiler is not None and metadata_in.duplicated: - profiler.set_duplicated() + if tracer is not None and metadata_in.duplicated: + tracer.set_duplicated() # Recv/send data. seq_num = 0 @@ -134,8 +134,8 @@ async def default_node_single( ), context=ir_context, ) - if profiler is not None: - profiler.add_chunk(table=df.table) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -180,11 +180,11 @@ async def default_node_multi( Index of the input channel to preserve partitioning information for. If None, no partitioning information is preserved. node_tracer - Node profiler for collecting runtime statistics. + Node tracer for collecting runtime statistics. """ async with shutdown_on_error( context, *chs_in, ch_out, node_tracer=node_tracer - ) as profiler: + ) as tracer: # Merge and forward basic metadata. local_count = 1 duplicated = True @@ -207,8 +207,8 @@ async def default_node_multi( duplicated=duplicated, ) await send_metadata(ch_out, context, metadata) - if profiler is not None and duplicated: - profiler.set_duplicated() + if tracer is not None and duplicated: + tracer.set_duplicated() seq_num = 0 n_children = len(chs_in) @@ -273,8 +273,8 @@ async def default_node_multi( *dfs, context=ir_context, ) - if profiler is not None: - profiler.add_chunk(table=df.table) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -553,8 +553,8 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) - profiler: StreamingQueryTracer | None = rec.state.get("profiler") - node_tracer = profiler.get_or_create(ir) if profiler is not None else None + 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 @@ -716,18 +716,18 @@ async def metadata_feeder_node( metadata The metadata to add to the output channel. node_tracer - Node profiler for collecting runtime statistics. + Node tracer for collecting runtime statistics. """ async with shutdown_on_error( context, ch_in, ch_out, node_tracer=node_tracer - ) as profiler: + ) as tracer: await send_metadata(ch_out, context, metadata) - if profiler is not None and metadata.duplicated: - profiler.set_duplicated() + 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 profiler is not None: - profiler.chunk_count += 1 + 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 af6b860b8d33..8854065f5231 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -78,11 +78,11 @@ async def concatenate_node( collective_id Pre-allocated collective ID for this operation. node_tracer - Node profiler for collecting runtime statistics. + Node tracer for collecting runtime statistics. """ async with shutdown_on_error( context, ch_in, ch_out, node_tracer=node_tracer - ) as profiler: + ) as tracer: # Receive metadata. input_metadata = await recv_metadata(ch_in, context) nranks = context.comm().nranks @@ -138,8 +138,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) - if profiler is not None and output_duplicated: - profiler.set_duplicated() + if tracer is not None and output_duplicated: + tracer.set_duplicated() allgather = AllGatherManager(context, collective_id) stream = context.get_stream_from_pool() @@ -152,8 +152,8 @@ async def concatenate_node( # Extract concatenated result result_table = await allgather.extract_concatenated(stream) - if profiler is not None: - profiler.add_chunk(table=result_table) + 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. @@ -174,8 +174,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) - if profiler is not None and output_duplicated: - profiler.set_duplicated() + if tracer is not None and output_duplicated: + tracer.set_duplicated() # Local repartitioning seq_num = 0 @@ -213,8 +213,8 @@ async def concatenate_node( ), context=ir_context, ) - if profiler is not None: - profiler.add_chunk(table=df.table) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -257,7 +257,7 @@ def _( collective_id = rec.state["collective_id_map"][ir][0] # Add python node - profiler = rec.state["profiler"] + tracer = rec.state["tracer"] nodes[ir] = [ concatenate_node( rec.state["context"], @@ -267,7 +267,7 @@ def _( channels[ir.children[0]].reserve_output_slot(), output_count=partition_info[ir].count, collective_id=collective_id, - node_tracer=profiler.get_or_create(ir) if profiler else None, + node_tracer=tracer.get_or_create(ir) if tracer else None, ) ] return nodes, channels diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index f751364bed27..71e4a0136822 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -48,7 +48,6 @@ "DynamicPlanningOptions", "InMemoryExecutor", "ParquetOptions", - "ProfilingOptions", "Runtime", "Scheduler", # Deprecated, kept for backward compatibility "ShuffleMethod", @@ -56,6 +55,7 @@ "StatsPlanningOptions", "StreamingExecutor", "StreamingFallbackMode", + "TracingOptions", ] @@ -495,25 +495,29 @@ def __post_init__(self) -> None: # noqa: D105 @dataclasses.dataclass(frozen=True) -class ProfilingOptions: +class TracingOptions: """ - Configuration for query profiling. + Configuration for query tracing. When enabled, the streaming executor collects per-node metrics - (such as row counts) and writes them to a file after execution. - This feature is only available for the "rapidsmpf" runtime. + (such as row counts and algorithm decisions) and writes them to a + file after execution. This feature is only available for the + "rapidsmpf" runtime. - To enable profiling, pass a ``ProfilingOptions`` instance - to ``StreamingExecutor(profiling=...)``. To disable it, pass - ``None`` (the default). + To enable tracing, pass a ``TracingOptions`` instance to + ``StreamingExecutor(tracing=...)``. To disable it, pass ``None`` + (the default). + + To also emit structlog events for each streaming node, set the + environment variable ``CUDF_POLARS_LOG_TRACES=1``. Parameters ---------- output_path - Path to write the profiling results. The output will be in a + Path to write the tracing results. The output will be in a human-readable text format similar to :func:`~cudf_polars.experimental.explain.explain_query`. - If ``None`` (the default), profiling data is collected but not + If ``None`` (the default), tracing data is collected but not written to a file. """ @@ -737,12 +741,15 @@ 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. - profiling - Options controlling query profiling. When set to a - :class:`~cudf_polars.utils.config.ProfilingOptions` instance, - per-node metrics (such as row counts) are collected during execution - and written to the specified output file. When ``None`` (the default), - profiling is disabled. + 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. @@ -854,7 +861,7 @@ class StreamingExecutor: f"{_env_prefix}__SPILL_TO_PINNED_MEMORY", bool, default=False ) ) - profiling: ProfilingOptions | None = None + tracing: TracingOptions | None = None def __post_init__(self) -> None: # noqa: D105 # Check for rapidsmpf runtime @@ -972,13 +979,13 @@ def __post_init__(self) -> None: # noqa: D105 DynamicPlanningOptions(**self.dynamic_planning), ) - # Handle profiling. - # Can be None, dict, or ProfilingOptions - if isinstance(self.profiling, dict): + # Handle tracing. + # Can be None, dict, or TracingOptions + if isinstance(self.tracing, dict): object.__setattr__( self, - "profiling", - ProfilingOptions(**self.profiling), + "tracing", + TracingOptions(**self.tracing), ) if self.cluster == "distributed": diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py similarity index 55% rename from python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py rename to python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py index 57a1f614e035..2169427d1735 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_runtime_profiler.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py @@ -1,9 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -"""Integration tests for runtime profiling with rapidsmpf.""" +"""Integration tests for runtime tracing with rapidsmpf.""" from __future__ import annotations +import subprocess +import sys +import textwrap + import pytest import polars as pl @@ -19,7 +23,7 @@ def get_engine(output_path: str, parquet_options: dict | None = None) -> pl.GPUE executor_options={ "cluster": DEFAULT_CLUSTER, "runtime": DEFAULT_RUNTIME, - "profiling": {"output_path": str(output_path)}, + "tracing": {"output_path": str(output_path)}, "max_rows_per_partition": 10, "target_partition_size": 500, }, @@ -33,9 +37,9 @@ def df(): @pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") -def test_profiling_basic_query(tmp_path, df): - """Test profiling output with a DataFrameScan query.""" - output_path = tmp_path / "dataframe_scan_profile.txt" +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) @@ -47,9 +51,9 @@ def test_profiling_basic_query(tmp_path, df): @pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") -def test_profiling_scan_parquet_python(tmp_path, df): - """Test profiling output with a ScanParquet query.""" - output_path = tmp_path / "scan_parquet_profile.txt" +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() @@ -61,9 +65,9 @@ def test_profiling_scan_parquet_python(tmp_path, df): @pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") -def test_profiling_scan_parquet_native(tmp_path, df): - """Test profiling output with a ScanParquet query.""" - output_path = tmp_path / "scan_parquet_profile.txt" +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() @@ -74,3 +78,43 @@ def test_profiling_scan_parquet_native(tmp_path, df): # 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_runtime_profiler.py b/python/cudf_polars/tests/experimental/test_tracing.py similarity index 100% rename from python/cudf_polars/tests/experimental/test_runtime_profiler.py rename to python/cudf_polars/tests/experimental/test_tracing.py diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 82098476ebc7..b4dff2d6408d 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -953,39 +953,39 @@ def test_dynamic_planning_from_instance() -> None: assert config.executor.dynamic_planning.sample_chunk_count == 2 # default -def test_profiling_options() -> None: - from cudf_polars.utils.config import ProfilingOptions +def test_tracing_options() -> None: + from cudf_polars.utils.config import TracingOptions - # Profiling is disabled (None) by default + # Tracing is disabled (None) by default config = ConfigOptions.from_polars_engine(pl.GPUEngine()) assert config.executor.name == "streaming" - assert config.executor.profiling is None + assert config.executor.tracing is None # Can enable via dict config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"profiling": {"output_path": "/tmp/profile.txt"}}, + executor_options={"tracing": {"output_path": "/tmp/trace.txt"}}, ) ) assert config.executor.name == "streaming" - assert config.executor.profiling is not None - assert config.executor.profiling.output_path == "/tmp/profile.txt" + 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={"profiling": ProfilingOptions()}, + executor_options={"tracing": TracingOptions()}, ) ) assert config.executor.name == "streaming" - assert config.executor.profiling is not None - assert config.executor.profiling.output_path is None + 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"): - ProfilingOptions(output_path="") + TracingOptions(output_path="") def test_parse_memory_resource_config() -> None: From 75db21e2ac175c59803e5248d93799d5c86ee157 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 07:12:38 -0800 Subject: [PATCH 20/25] move import --- python/cudf_polars/tests/test_config.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index b4dff2d6408d..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 @@ -954,8 +955,6 @@ def test_dynamic_planning_from_instance() -> None: def test_tracing_options() -> None: - from cudf_polars.utils.config import TracingOptions - # Tracing is disabled (None) by default config = ConfigOptions.from_polars_engine(pl.GPUEngine()) assert config.executor.name == "streaming" From 20f26e49de7c6974c155b33cb685bd5e866b72e8 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 07:20:11 -0800 Subject: [PATCH 21/25] fixes --- docs/cudf/source/cudf_polars/api.md | 2 +- .../cudf_polars/experimental/explain.py | 18 +++++++++--------- .../cudf_polars/experimental/rapidsmpf/core.py | 4 ++-- .../tests/experimental/test_tracing.py | 12 ++++++------ 4 files changed, 18 insertions(+), 18 deletions(-) diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index 4a779c941bd6..9c0034b9f1df 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -11,7 +11,7 @@ For the most part, the public API of `cudf-polars` is the polars API. DynamicPlanningOptions, InMemoryExecutor, ParquetOptions, - ProfilingOptions, + TracingOptions, Cluster, ShuffleMethod, ShufflerInsertionMethod, diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index fb4f6ca75ecb..564039507ac6 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -172,19 +172,19 @@ def _(ir: Scan, *, offset: str = "") -> str: return _repr_header(offset, label, ir.schema) -def write_profile_output( - profile_output: str | Path, +def write_query_trace( + trace_output: str | Path, ir: IR, partition_info: MutableMapping[IR, PartitionInfo], tracer: StreamingQueryTracer, ) -> None: """ - Write a post-execution profile showing actual row counts and decisions. + Write a post-execution trace showing actual row counts and decisions. Parameters ---------- - profile_output - Path to write the profile file. + trace_output + Path to write the trace file. ir The lowered IR root node. partition_info @@ -194,11 +194,11 @@ def write_profile_output( """ from pathlib import Path - profile_repr = _repr_profile_tree(ir, partition_info, tracer) - Path(profile_output).write_text(profile_repr) + trace_repr = _repr_trace_tree(ir, partition_info, tracer) + Path(trace_output).write_text(trace_repr) -def _repr_profile_tree( +def _repr_trace_tree( ir: IR, partition_info: MutableMapping[IR, PartitionInfo], tracer: StreamingQueryTracer, @@ -223,7 +223,7 @@ def _repr_profile_tree( header += f" chunks={node_tracer.chunk_count}" children_strs = [ - _repr_profile_tree(child, partition_info, tracer, offset=offset + " ") + _repr_trace_tree(child, partition_info, tracer, offset=offset + " ") for child in ir.children ] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index e0ad4d2fe4e1..5d6093247770 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -132,9 +132,9 @@ def evaluate_logical_plan( # 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_profile_output + from cudf_polars.experimental.explain import write_query_trace - write_profile_output(tracing.output_path, ir, partition_info, tracer) + write_query_trace(tracing.output_path, ir, partition_info, tracer) return result, metadata_collector diff --git a/python/cudf_polars/tests/experimental/test_tracing.py b/python/cudf_polars/tests/experimental/test_tracing.py index 3aecfa566275..3d884c091e7b 100644 --- a/python/cudf_polars/tests/experimental/test_tracing.py +++ b/python/cudf_polars/tests/experimental/test_tracing.py @@ -7,7 +7,7 @@ import polars as pl from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.explain import _repr_profile_tree, write_profile_output +from cudf_polars.experimental.explain import _repr_trace_tree, write_query_trace from cudf_polars.experimental.rapidsmpf.tracing import StreamingQueryTracer @@ -45,16 +45,16 @@ def get_hashable(self): assert tracer1.node_tracers[ir1].row_count == 250 assert tracer1.node_tracers[ir1].chunk_count == 8 - # Test _repr_profile_tree output format + # Test _repr_trace_tree output format partition_info = {ir1: PartitionInfo(count=4)} - output = _repr_profile_tree(ir1, partition_info, tracer1) + 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_profile_output - output_path = tmp_path / "profile.txt" - write_profile_output(output_path, ir1, partition_info, tracer1) + # 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 From 72a0916d95fe04247eecf1e047eb09f63b379248 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 07:34:03 -0800 Subject: [PATCH 22/25] clean up docstring --- .../cudf_polars/cudf_polars/utils/config.py | 36 +++++++++++-------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 71e4a0136822..4ff584cc0499 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -497,28 +497,34 @@ def __post_init__(self) -> None: # noqa: D105 @dataclasses.dataclass(frozen=True) class TracingOptions: """ - Configuration for query tracing. + Configuration for coarse-grained streaming-node tracing (rapidsmpf only). - When enabled, the streaming executor collects per-node metrics - (such as row counts and algorithm decisions) and writes them to a - file after execution. This feature is only available for the - "rapidsmpf" runtime. + 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. - To enable tracing, pass a ``TracingOptions`` instance to - ``StreamingExecutor(tracing=...)``. To disable it, pass ``None`` - (the default). + 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 the - environment variable ``CUDF_POLARS_LOG_TRACES=1``. + To also emit structlog events for each streaming node, set + ``CUDF_POLARS_LOG_TRACES=1``. Parameters ---------- output_path - Path to write the tracing results. The output will be in a - human-readable text format similar to - :func:`~cudf_polars.experimental.explain.explain_query`. - If ``None`` (the default), tracing data is collected but not - written to a file. + 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 From 694cf8947cc26eb007cd4f57e3ff4daf1f9d2de3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 10:22:03 -0800 Subject: [PATCH 23/25] partial code review --- .../experimental/rapidsmpf/tracing.py | 24 ++++++++++++------- .../experimental/rapidsmpf/utils.py | 7 ++---- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py index 926c17effa9a..59b5e698f346 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py @@ -98,16 +98,22 @@ def set_duplicated(self, *, duplicated: bool = True) -> None: def merge(self, other: StreamingNodeTracer) -> None: """Merge another node tracer's stats into this one.""" - if other.row_count is not None: - if self.duplicated or other.duplicated: - # For duplicated data, take max (don't sum across ranks) - self.row_count = max(self.row_count or 0, other.row_count) - self.duplicated = True - else: + 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 - self.chunk_count += other.chunk_count - if other.decision is not None: - self.decision = other.decision + if other.decision is not None: + self.decision = other.decision class StreamingQueryTracer: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 2e13d0fad51c..042e7352496d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -136,8 +136,9 @@ 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 hs is None or hs == "inherit": return hs # None or "inherit" passes through unchanged + assert isinstance(hs, HashScheme), "Expected HashScheme" try: new_indices = tuple( new_name_to_idx[old_names[i]] for i in hs.column_indices @@ -153,10 +154,6 @@ def remap_hash_scheme(hs: HashScheme | None | str) -> HashScheme | None | str: 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) From 459f780d9e119e24366c5a69465bb8568ee5ad24 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 10:24:14 -0800 Subject: [PATCH 24/25] partial code review (2) --- .../experimental/rapidsmpf/utils.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 042e7352496d..85d4adad9490 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -136,16 +136,16 @@ 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 hs == "inherit": + 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 - assert isinstance(hs, HashScheme), "Expected 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) new_inter_rank = remap_hash_scheme(partitioning.inter_rank) new_local = remap_hash_scheme(partitioning.local) From 09ad6cc23815281d8de2ac2d10879d0a0abb7588 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 3 Feb 2026 10:28:09 -0800 Subject: [PATCH 25/25] drop unnecessary check --- .../cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 85d4adad9490..8868b1b64472 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -149,11 +149,6 @@ def remap_hash_scheme(hs: HashScheme | None | str) -> HashScheme | None | str: 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 - return Partitioning(inter_rank=new_inter_rank, local=new_local)