diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 6b63b31c38ba..b6e1a6961d35 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -1,10 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Utilities for tracing and monitoring IR execution.""" from __future__ import annotations +import enum import functools import os import time @@ -50,6 +51,14 @@ from cudf_polars.dsl import ir +class Scope(str, enum.Enum): + """Scope values for structured logging.""" + + PLAN = "plan" + ACTOR = "actor" + EVALUATE_IR_NODE = "evaluate_ir_node" + + @functools.cache def _getpid() -> int: # pragma: no cover # Gets called for each IR.do_evaluate node, so we'll cache it. @@ -199,8 +208,9 @@ def wrapper( before | after | { + "scope": Scope.EVALUATE_IR_NODE.value, "overhead_duration": (before_end - before_start) - + (after_end - after_start) + + (after_end - after_start), } ) log.info("Execute IR", **record) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 9d59650b9688..906413450668 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -45,6 +45,7 @@ try: from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.dsl.tracing import Scope from cudf_polars.dsl.translate import Translator from cudf_polars.experimental.explain import explain_query from cudf_polars.experimental.parallel import evaluate_streaming @@ -1134,13 +1135,17 @@ def gather_logs() -> str: return logger.handlers[0].stream.getvalue() # type: ignore[attr-defined] if client is not None: - all_logs = "\n".join(client.run(gather_logs).values()) + # Gather logs from both client (for Query Plan) and workers + worker_logs = "\n".join(client.run(gather_logs).values()) + client_logs = gather_logs() + all_logs = client_logs + "\n" + worker_logs else: all_logs = gather_logs() parsed_logs = [json.loads(log) for log in all_logs.splitlines() if log] # Some other log records can end up in here. Filter those out. - parsed_logs = [log for log in parsed_logs if log["event"] == "Execute IR"] + scope_values = {s.value for s in Scope} + parsed_logs = [log for log in parsed_logs if log.get("scope") in scope_values] # Now we want to augment the existing Records with the trace data. def group_key(x: dict) -> int: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 3d4ee653f77a..376ce43c2a29 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -39,6 +39,7 @@ generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.experimental.rapidsmpf.tracing import log_query_plan 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 @@ -96,6 +97,9 @@ def evaluate_logical_plan( # Lower the IR graph on the client process (for now). ir, partition_info, stats = lower_ir_graph(ir, config_options) + # Log the query plan structure for tracing (no-op if tracing disabled) + log_query_plan(ir) + # Reserve shuffle IDs for the entire pipeline execution with ReserveOpIDs(ir) as collective_id_map: # Build and execute the 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 02e3aab748de..fcb2ea4906ad 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -51,6 +51,7 @@ from cudf_polars.experimental.base import ColumnStat, StatsCollector from cudf_polars.experimental.dispatch import LowerIRTransformer from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer from cudf_polars.utils.config import ParquetOptions @@ -163,7 +164,7 @@ async def dataframescan_node( Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. """ - async with shutdown_on_error(context, ch_out): + async with shutdown_on_error(context, ch_out, trace_ir=ir) as tracer: # Find local partition count. nrows = ir.df.shape()[0] global_count = math.ceil(nrows / rows_per_partition) if nrows > 0 else 0 @@ -213,6 +214,7 @@ async def dataframescan_node( ch_out, ir_context, estimated_chunk_bytes, + tracer=tracer, ) await ch_out.drain(context) return @@ -238,6 +240,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, + tracer=tracer, ) await ch_out.drain(context) @@ -319,6 +322,7 @@ async def read_chunk( ch_out: Channel[TableChunk], ir_context: IRExecutionContext, estimated_chunk_bytes: int, + tracer: ActorTracer | None = None, ) -> None: """ Read a chunk from disk and send it to the output channel. @@ -338,6 +342,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. + tracer + The actor tracer for collecting runtime statistics. """ with opaque_reservation(context, estimated_chunk_bytes): df = await asyncio.to_thread( @@ -345,6 +351,8 @@ async def read_chunk( *scan._non_child_args, context=ir_context, ) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -393,7 +401,7 @@ async def scan_node( Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. """ - async with shutdown_on_error(context, ch_out): + async with shutdown_on_error(context, ch_out, trace_ir=ir) as tracer: # Build a list of local Scan operations scans: list[Scan | SplitScan] = [] if plan.flavor == IOPartitionFlavor.SPLIT_FILES: @@ -484,6 +492,7 @@ async def scan_node( ch_out, ir_context, estimated_chunk_bytes, + tracer=tracer, ) await ch_out.drain(context) return @@ -509,6 +518,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ch_out, ir_context, estimated_chunk_bytes, + tracer=tracer, ) await ch_out.drain(context) @@ -681,6 +691,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 4330cdb176fe..6f8d3c8f179f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -71,7 +71,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): + async with shutdown_on_error(context, ch_in, ch_out, trace_ir=ir) as tracer: # Recv/send metadata. metadata_in = await recv_metadata(ch_in, context) partitioning = None @@ -86,6 +86,8 @@ async def default_node_single( duplicated=metadata_in.duplicated, ) await send_metadata(ch_out, context, metadata_out) + if tracer is not None and metadata_in.duplicated: + tracer.set_duplicated() # Recv/send data. seq_num = 0 @@ -123,6 +125,8 @@ async def default_node_single( ), context=ir_context, ) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -166,7 +170,7 @@ async def default_node_multi( Index of the input channel to preserve partitioning information for. If None, no partitioning information is preserved. """ - async with shutdown_on_error(context, *chs_in, ch_out): + async with shutdown_on_error(context, *chs_in, ch_out, trace_ir=ir) as tracer: # Merge and forward basic metadata. local_count = 1 duplicated = True @@ -189,6 +193,8 @@ async def default_node_multi( duplicated=duplicated, ) await send_metadata(ch_out, context, metadata) + if tracer is not None and duplicated: + tracer.set_duplicated() seq_num = 0 n_children = len(chs_in) @@ -253,6 +259,8 @@ async def default_node_multi( *dfs, context=ir_context, ) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( @@ -670,6 +678,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, @@ -681,6 +690,8 @@ async def metadata_feeder_node( ---------- context The rapidsmpf context. + ir + The IR node (for tracing). ch_in The input channel to pull data from. ch_out @@ -688,10 +699,14 @@ async def metadata_feeder_node( metadata The metadata to add to the output channel. """ - async with shutdown_on_error(context, ch_in, ch_out): + async with shutdown_on_error(context, ch_in, ch_out, trace_ir=ir) as tracer: await send_metadata(ch_out, context, metadata) + if tracer is not None and metadata.duplicated: + tracer.set_duplicated() while (msg := await ch_in.recv(context)) is not None: await ch_out.send(context, msg) + if tracer is not None: + tracer.chunk_count += 1 await ch_out.drain(context) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py index 380b8a68e0ca..2b0c7503b24d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py @@ -76,7 +76,7 @@ async def concatenate_node( collective_id Pre-allocated collective ID for this operation. """ - async with shutdown_on_error(context, ch_in, ch_out): + async with shutdown_on_error(context, ch_in, ch_out, trace_ir=ir) as tracer: # Receive metadata. input_metadata = await recv_metadata(ch_in, context) nranks = context.comm().nranks @@ -132,6 +132,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) + if tracer is not None and output_duplicated: + tracer.set_duplicated() allgather = AllGatherManager(context, collective_id) stream = context.get_stream_from_pool() @@ -144,6 +146,8 @@ async def concatenate_node( # Extract concatenated result result_table = await allgather.extract_concatenated(stream) + if tracer is not None: + tracer.add_chunk(table=result_table) # If no chunks were gathered, result_table has 0 columns. # We need to create an empty table with the correct schema. @@ -164,6 +168,8 @@ async def concatenate_node( duplicated=output_duplicated, ) await send_metadata(ch_out, context, metadata) + if tracer is not None and output_duplicated: + tracer.set_duplicated() # Local repartitioning seq_num = 0 @@ -201,6 +207,8 @@ async def concatenate_node( ), context=ir_context, ) + if tracer is not None: + tracer.add_chunk(table=df.table) await ch_out.send( context, Message( 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..31ee9c113bb2 --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py @@ -0,0 +1,135 @@ +# 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 + +from cudf_polars.dsl.tracing import LOG_TRACES, Scope +from cudf_polars.dsl.traversal import traversal + +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 ActorTracer: + """ + Tracer for a single streaming actor (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. + + Parameters + ---------- + table + The table to record. + """ + 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.""" + self.duplicated = duplicated + + +def log_query_plan(ir: IR) -> None: + """ + Log the IR tree structure as a structlog event. + + This should be called once on the client process after lowering, + before distributed execution begins. The structure can be used + by post-processing tools to reconstruct annotated plans. + + Parameters + ---------- + ir + The root IR node of the lowered query plan. + + Notes + ----- + This function is a no-op if ``CUDF_POLARS_LOG_TRACES`` is not set. + """ + if not LOG_TRACES: + return + + import structlog + + nodes = [ + { + "ir_id": _stable_ir_id(node), + "ir_type": type(node).__name__, + "children_ir_ids": [_stable_ir_id(c) for c in node.children], + } + for node in traversal([ir]) + ] + + log = structlog.get_logger() + log.info("Query Plan", scope=Scope.PLAN.value, nodes=nodes) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 3a8d0522be53..6627628288f0 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -10,6 +10,14 @@ from functools import reduce from typing import TYPE_CHECKING, Any +from cudf_polars.dsl.tracing import LOG_TRACES, Scope + +try: + import structlog + import structlog.contextvars +except ImportError: + pass + from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.cudf.channel_metadata import ( ChannelMetadata, @@ -34,29 +42,73 @@ from cudf_polars.dsl.ir import IR from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer from cudf_polars.typing import DataType @asynccontextmanager async def shutdown_on_error( - context: Context, *channels: Channel[Any] -) -> AsyncIterator[None]: + context: Context, + *channels: Channel[Any], + trace_ir: IR | None = None, +) -> AsyncIterator[ActorTracer | None]: """ Shutdown on error for rapidsmpf. + This context manager handles channel cleanup on errors and optionally + emits structlog tracing events when LOG_TRACES is enabled. + Parameters ---------- context The rapidsmpf context. channels - The channels to shutdown. + The channels to shutdown on error. + trace_ir + Optional IR node to enable tracing for this streaming actor. + When provided and LOG_TRACES is enabled, an ActorTracer + is yielded for collecting stats, and a structlog event is + emitted on exit. + + Yields + ------ + ActorTracer | None + An actor tracer for collecting stats (if tracing enabled), else None. """ - # TODO: This probably belongs in rapidsmpf. + # Create tracer only if LOG_TRACES is enabled and IR is provided + tracer: ActorTracer | None = None + if LOG_TRACES and trace_ir is not None: + from cudf_polars.experimental.rapidsmpf.tracing import ( + ActorTracer, + _stable_ir_id, + ) + + ir_id = _stable_ir_id(trace_ir) + ir_type = type(trace_ir).__name__ + tracer = ActorTracer(ir_id, ir_type) + structlog.contextvars.bind_contextvars(actor_ir_id=ir_id, actor_ir_type=ir_type) + try: - yield + yield tracer except BaseException: await asyncio.gather(*(ch.shutdown(context) for ch in channels)) raise + finally: + if tracer is not None: + log = structlog.get_logger() + record: dict[str, Any] = { + "scope": Scope.ACTOR.value, + "actor_ir_id": tracer.ir_id, + "actor_ir_type": tracer.ir_type, + "chunk_count": tracer.chunk_count, + "duplicated": tracer.duplicated, + } + if tracer.row_count is not None: + record["rows"] = tracer.row_count + if tracer.decision is not None: + record["decision"] = tracer.decision + log.info("Streaming Actor", **record) + structlog.contextvars.unbind_contextvars("actor_ir_id", "actor_ir_type") def remap_partitioning( @@ -92,27 +144,19 @@ def remap_partitioning( new_name_to_idx = {name: i for i, name in enumerate(new_schema.keys())} def remap_hash_scheme(hs: HashScheme | None | str) -> HashScheme | None | str: - if hs is None or isinstance(hs, str): + if isinstance(hs, HashScheme): + try: + new_indices = tuple( + new_name_to_idx[old_names[i]] for i in hs.column_indices + ) + except (IndexError, KeyError): + return None # Column missing in old or new schema + return HashScheme(new_indices, hs.modulus) + else: return hs # None or "inherit" passes through unchanged - try: - new_indices = tuple( - new_name_to_idx[old_names[i]] for i in hs.column_indices - ) - except (IndexError, KeyError): - return None # Column missing in old or new schema - return HashScheme(new_indices, hs.modulus) new_inter_rank = remap_hash_scheme(partitioning.inter_rank) new_local = remap_hash_scheme(partitioning.local) - - # If inter_rank partitioning was invalidated, the whole partitioning is invalid - if isinstance(partitioning.inter_rank, HashScheme) and new_inter_rank is None: - return None - - # If only local partitioning was invalidated, we can still use inter_rank - if isinstance(partitioning.local, HashScheme) and new_local is None: - new_local = None - return Partitioning(inter_rank=new_inter_rank, local=new_local) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py new file mode 100644 index 000000000000..74b2b56f98f9 --- /dev/null +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_tracing.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Integration tests for structlog tracing with rapidsmpf.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap + +import pytest + +from cudf_polars.testing.asserts import DEFAULT_CLUSTER, DEFAULT_RUNTIME + + +@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 Actor' 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) + """) + + # Build environment with tracing enabled + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + + result = subprocess.check_output( + [sys.executable, "-c", code], env=env, stderr=subprocess.STDOUT + ) + + # Check for Streaming Actor events emitted by shutdown_on_error + assert b"Streaming Actor" in result + assert b"scope=actor" in result or b"'scope': 'actor'" in result + assert b"actor_ir_id=" in result + assert b"actor_ir_type=" in result + assert b"chunk_count=" in result + + +@pytest.mark.skipif( + DEFAULT_RUNTIME != "rapidsmpf", reason="Requires 'rapidsmpf' runtime." +) +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_structlog_contains_expected_ir_types(): + """Test that structlog output contains expected IR types for a query.""" + 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 = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + + result = subprocess.check_output( + [sys.executable, "-c", code], env=env, stderr=subprocess.STDOUT + ) + + # Check for expected IR types in the query + assert b"ir_type=DataFrameScan" in result + assert b"ir_type=Filter" in result + assert b"ir_type=GroupBy" in result + assert b"ir_type=Repartition" in result + + +@pytest.mark.skipif( + DEFAULT_RUNTIME != "rapidsmpf", reason="Requires 'rapidsmpf' runtime." +) +@pytest.mark.skipif(DEFAULT_CLUSTER != "single", reason="Requires 'single' cluster.") +def test_structlog_disabled_by_default(): + """Test that structlog does NOT emit events when CUDF_POLARS_LOG_TRACES is not set.""" + code = textwrap.dedent("""\ + import polars as pl + import rmm + + df = pl.DataFrame({"x": range(10), "y": ["a", "b"] * 5}) + q = df.lazy().filter(pl.col("x") > 5) + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + executor_options={ + "cluster": "single", + "runtime": "rapidsmpf", + "max_rows_per_partition": 5, + }, + memory_resource=rmm.mr.ManagedMemoryResource(), + ) + q.collect(engine=engine) + """) + + # Environment WITHOUT CUDF_POLARS_LOG_TRACES + env = os.environ.copy() + env.pop("CUDF_POLARS_LOG_TRACES", None) + + result = subprocess.check_output( + [sys.executable, "-c", code], env=env, stderr=subprocess.STDOUT + ) + + # Should NOT see Streaming Actor events + assert b"Streaming Actor" not in result diff --git a/python/cudf_polars/tests/test_tracing.py b/python/cudf_polars/tests/test_tracing.py index 1928617810c4..f67fcda89a45 100644 --- a/python/cudf_polars/tests/test_tracing.py +++ b/python/cudf_polars/tests/test_tracing.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 from __future__ import annotations @@ -75,3 +75,45 @@ def test_import_without_structlog(monkeypatch: pytest.MonkeyPatch) -> None: # And we can run a query without error q = pl.DataFrame({"a": [1, 2, 3]}).lazy().select(pl.col("a").sum()) q.collect(engine="gpu") + + +@pytest.mark.skipif( + cudf_polars.testing.asserts.DEFAULT_RUNTIME != "rapidsmpf", + reason="Requires 'rapidsmpf' runtime.", +) +def test_log_query_plan() -> None: + """Test that log_query_plan emits a Query Plan event.""" + import os + + code = textwrap.dedent("""\ + import polars as pl + import rmm + + df = pl.DataFrame({"x": range(10), "y": ["a", "b"] * 5}) + q = df.lazy().filter(pl.col("x") > 5).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": 5, + }, + memory_resource=rmm.mr.ManagedMemoryResource(), + ) + q.collect(engine=engine) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + + result = subprocess.check_output( + [sys.executable, "-c", code], env=env, stderr=subprocess.STDOUT + ) + + # Check for Query Plan event + assert b"Query Plan" in result + assert b"scope=plan" in result or b"'scope': 'plan'" in result + assert b"ir_id" in result + assert b"ir_type" in result + assert b"children_ir_ids" in result