Trace node execution in cudf-polars

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -96,6 +96,7 @@ dependencies:
 - sphinx>=8.1.0
 - sphinxcontrib-websupport
 - streamz
+- structlog
 - sysroot_linux-aarch64==2.28
 - typing_extensions>=4.0.0
 - zlib>=1.2.13

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -97,6 +97,7 @@ dependencies:
 - sphinx>=8.1.0
 - sphinxcontrib-websupport
 - streamz
+- structlog
 - sysroot_linux-64==2.28
 - typing_extensions>=4.0.0
 - zlib>=1.2.13

conda/environments/all_cuda-130_arch-aarch64.yaml

@@ -95,6 +95,7 @@ dependencies:
 - sphinx>=8.1.0
 - sphinxcontrib-websupport
 - streamz
+- structlog
 - sysroot_linux-aarch64==2.28
 - typing_extensions>=4.0.0
 - zlib>=1.2.13

conda/environments/all_cuda-130_arch-x86_64.yaml

@@ -96,6 +96,7 @@ dependencies:
 - sphinx>=8.1.0
 - sphinxcontrib-websupport
 - streamz
+- structlog
 - sysroot_linux-64==2.28
 - typing_extensions>=4.0.0
 - zlib>=1.2.13

dependencies.yaml

@@ -13,6 +13,7 @@ files:
       - clang
       - cuda
       - cuda_version
+      - cudf_polars_trace
       - depends_on_cupy
       - depends_on_dask_cuda
       - depends_on_libkvikio
@@ -337,6 +338,15 @@ files:
       - numpy_run
       - test_python_common
       - test_python_cudf_polars
+      - cudf_polars_trace
+  py_trace_cudf_polars:
+    output: pyproject
+    pyproject_dir: python/cudf_polars
+    extras:
+      table: project.optional-dependencies
+      key: trace
+    includes:
+      - cudf_polars_trace
   py_build_dask_cudf:
     output: pyproject
     pyproject_dir: python/dask_cudf
@@ -1214,3 +1224,8 @@ dependencies:
       - output_types: [conda, requirements, pyproject]
         packages:
           - narwhals==2.0.1
+  cudf_polars_trace:
+    common:
+      - output_types: [conda, requirements, pyproject]
+        packages:
+          - structlog

docs/cudf/source/cudf_polars/usage.md

@@ -83,3 +83,42 @@ shape: (3, 3)
 │ Scan               ┆ 813   ┆ 233993 │
 └────────────────────┴───────┴────────┘
 ```
+
+## Tracing
+
+cudf-polars can optionally trace execution of each node in the query plan.
+To enable tracing, set the environment variable ``CUDF_POLARS_LOG_TRACES`` to a
+true value ("1", "true", "y", "yes") before starting your process. This will
+capture and log information about each node before and after it executes, including information on
+
+- The type of the node being executed (e.g. `Scan`, `Select`, `Join`, `Groupby`, etc.)
+- The shape and size (in memory) of each input and output DataFrame
+- The GPU memory usage, as reported by [nvml], before and after executing the node
+- A start and stop clock, which can be used to measure the duration of the node's execution
+
+The implementation uses [structlog] to build log records. You can configure the
+output using structlog's [configuration][structlog-configure] and enrich the
+records with [context variables][structlog-context].
+
+
+```
+>>> df = pl.DataFrame({"a": ["a", "a", "b"], "b": [1, 2, 3]}).lazy()
+>>> df.group_by("a").agg(pl.col("b").min().alias("min"), pl.col("b").max().alias("max")).collect(engine="gpu")
+2025-09-10 07:44:01 [info     ] Execute IR      count_frames_input=0 count_frames_output=1 ... type=DataFrameScan
+2025-09-10 07:44:01 [info     ] Execute IR      count_frames_input=1 count_frames_output=1 ... type=GroupBy
+shape: (2, 3)
+┌─────┬─────┬─────┐
+│ a   ┆ min ┆ max │
+│ --- ┆ --- ┆ --- │
+│ str ┆ i64 ┆ i64 │
+╞═════╪═════╪═════╡
+│ b   ┆ 3   ┆ 3   │
+│ a   ┆ 1   ┆ 2   │
+└─────┴─────┴─────┘
+```
+
+[nvml]: https://developer.nvidia.com/management-library-nvml
+[rmm-stats]: https://docs.rapids.ai/api/rmm/stable/guide/#memory-statistics-and-profiling
+[structlog]: https://www.structlog.org/
+[structlog-configure](https://www.structlog.org/en/stable/configuration.html)
+[structlog-context](https://www.structlog.org/en/stable/contextvars.html)

python/cudf_polars/cudf_polars/callback.py

@@ -22,6 +22,7 @@
 import rmm
 from rmm._cuda import gpu
 
+import cudf_polars.dsl.tracing
 from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN
 from cudf_polars.dsl.translate import Translator
 from cudf_polars.utils.config import _env_get_int, get_total_device_memory
@@ -133,6 +134,12 @@ def set_memory_resource(
                 != 0
             ),
         )
+
+    if (
+        cudf_polars.dsl.tracing.LOG_TRACES
+    ):  # pragma: no cover; requires CUDF_POLARS_LOG_TRACES=1
+        mr = rmm.mr.StatisticsResourceAdaptor(mr)
+
     rmm.mr.set_current_device_resource(mr)
     try:
         yield mr

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -33,7 +33,7 @@
 from cudf_polars.dsl.expressions.base import ExecutionContext
 from cudf_polars.dsl.nodebase import Node
 from cudf_polars.dsl.to_ast import to_ast, to_parquet_filter
-from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars
+from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars
 from cudf_polars.dsl.utils.reshape import broadcast
 from cudf_polars.dsl.utils.windows import range_window_bounds
 from cudf_polars.utils import dtypes
@@ -479,6 +479,7 @@ def fast_count(self) -> int:  # pragma: no cover
         return max(total_rows, 0)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Scan")
     def do_evaluate(
         cls,
@@ -975,6 +976,7 @@ def _write_parquet(
             plc.io.parquet.write_parquet(writer_options)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Sink")
     def do_evaluate(
         cls,
@@ -1034,6 +1036,7 @@ def is_equal(self, other: Self) -> bool:  # noqa: D102
         return False
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Cache")
     def do_evaluate(
         cls, key: int, refcount: int | None, df: DataFrame
@@ -1114,6 +1117,7 @@ def get_hashable(self) -> Hashable:
         )
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="DataFrameScan")
     def do_evaluate(
         cls,
@@ -1173,6 +1177,7 @@ def _is_len_expr(exprs: tuple[expr.NamedExpr, ...]) -> bool:  # pragma: no cover
         return False
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Select")
     def do_evaluate(
         cls,
@@ -1256,6 +1261,7 @@ def __init__(
         self._non_child_args = (self.exprs,)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Reduce")
     def do_evaluate(
         cls,
@@ -1352,6 +1358,7 @@ def __init__(
         )
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Rolling")
     def do_evaluate(
         cls,
@@ -1476,6 +1483,7 @@ def __init__(
         )
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="GroupBy")
     def do_evaluate(
         cls,
@@ -1603,7 +1611,8 @@ def __reduce__(self) -> tuple[Any, ...]:
         tuple[
             str,
             pl_expr.Operator | Iterable[pl_expr.Operator],
-        ],
+        ]
+        | None,
         bool,
         Zlice | None,
         str,
@@ -1625,6 +1634,12 @@ def __init__(
     ) -> None:
         self.schema = schema
         self.predicate = predicate
+        # options[0] is a tuple[str, Operator, ...]
+        # The Operator class can't be pickled, but we don't use it anyway so
+        # just throw that away
+        if options[0] is not None:
+            options = (None, *options[1:])
+
         self.options = options
         self.children = (left, right)
         predicate_wrapper = self.Predicate(predicate)
@@ -1637,20 +1652,21 @@ def __init__(
             raise NotImplementedError(
                 f"Conditional join with predicate {predicate}"
             )  # pragma: no cover; polars never delivers expressions we can't handle
-        self._non_child_args = (predicate_wrapper, zlice, suffix, maintain_order)
+        self._non_child_args = (predicate_wrapper, options)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="ConditionalJoin")
     def do_evaluate(
         cls,
         predicate_wrapper: Predicate,
-        zlice: Zlice | None,
-        suffix: str,
-        maintain_order: Literal["none", "left", "right", "left_right", "right_left"],
+        options: tuple,
         left: DataFrame,
         right: DataFrame,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
+        _, _, zlice, suffix, _, _ = options
+
         lg, rg = plc.join.conditional_inner_join(
             left.table,
             right.table,
@@ -1875,6 +1891,7 @@ def _build_columns(
         return columns
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Join")
     def do_evaluate(
         cls,
@@ -2027,6 +2044,7 @@ def __init__(
         self.children = (df,)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="HStack")
     def do_evaluate(
         cls,
@@ -2092,6 +2110,7 @@ def __init__(
     }
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Distinct")
     def do_evaluate(
         cls,
@@ -2182,6 +2201,7 @@ def __init__(
         self.children = (df,)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Sort")
     def do_evaluate(
         cls,
@@ -2232,6 +2252,7 @@ def __init__(self, schema: Schema, offset: int, length: int | None, df: IR):
         self.children = (df,)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Slice")
     def do_evaluate(cls, offset: int, length: int, df: DataFrame) -> DataFrame:
         """Evaluate and return a dataframe."""
@@ -2253,6 +2274,7 @@ def __init__(self, schema: Schema, mask: expr.NamedExpr, df: IR):
         self.children = (df,)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Filter")
     def do_evaluate(cls, mask_expr: expr.NamedExpr, df: DataFrame) -> DataFrame:
         """Evaluate and return a dataframe."""
@@ -2272,6 +2294,7 @@ def __init__(self, schema: Schema, df: IR):
         self.children = (df,)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Projection")
     def do_evaluate(cls, schema: Schema, df: DataFrame) -> DataFrame:
         """Evaluate and return a dataframe."""
@@ -2305,6 +2328,7 @@ def __init__(self, schema: Schema, key: str, left: IR, right: IR):
         self._non_child_args = (key,)
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="MergeSorted")
     def do_evaluate(cls, key: str, *dfs: DataFrame) -> DataFrame:
         """Evaluate and return a dataframe."""
@@ -2425,6 +2449,7 @@ def get_hashable(self) -> Hashable:
         )
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="MapFunction")
     def do_evaluate(
         cls, schema: Schema, name: str, options: Any, df: DataFrame
@@ -2525,6 +2550,7 @@ def __init__(self, schema: Schema, zlice: Zlice | None, *children: IR):
         schema = self.children[0].schema
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Union")
     def do_evaluate(cls, zlice: Zlice | None, *dfs: DataFrame) -> DataFrame:
         """Evaluate and return a dataframe."""
@@ -2582,6 +2608,7 @@ def _extend_with_nulls(table: plc.Table, *, nrows: int) -> plc.Table:
         )
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="HConcat")
     def do_evaluate(
         cls,
@@ -2627,6 +2654,7 @@ def __init__(self, schema: Schema):
         self.children = ()
 
     @classmethod
+    @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Empty")
     def do_evaluate(cls, schema: Schema) -> DataFrame:  # pragma: no cover
         """Evaluate and return a dataframe."""