Full structlog-based tracing

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)

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:

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.

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,13 +342,17 @@ 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(
             scan.do_evaluate,
             *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(

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,17 +690,23 @@ 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
         The output channel to forward data to and add metadata to.
     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)
 
 

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(