From 8987e952c8169c30f2df398b6b9fc7c7cf05c5ef Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 20 Aug 2026 10:41:25 +0200 Subject: [PATCH 1/9] Report kvikio I/O statistics per rank in the streaming engines --- docs/cudf/source/conf.py | 1 + docs/cudf/source/cudf_polars/api.md | 8 +- docs/cudf/source/cudf_polars/profiling.md | 82 ++++++ python/cudf_polars/cudf_polars/engine/core.py | 126 +++++++++ python/cudf_polars/cudf_polars/engine/dask.py | 59 ++++ python/cudf_polars/cudf_polars/engine/ray.py | 50 ++++ python/cudf_polars/cudf_polars/engine/spmd.py | 43 +++ .../benchmarks/print_results_file.py | 261 ++++++++++++++++++ .../cudf_polars/streaming/benchmarks/utils.py | 27 ++ .../tests/streaming/test_statistics.py | 78 +++++- 10 files changed, 730 insertions(+), 5 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py diff --git a/docs/cudf/source/conf.py b/docs/cudf/source/conf.py index ec0b2c798fba..a252db9b2886 100644 --- a/docs/cudf/source/conf.py +++ b/docs/cudf/source/conf.py @@ -695,6 +695,7 @@ def on_missing_reference(app, env, node, contnode): ("ref.*", ".*pandas.*"), # External libs without configured intersphinx inventories. ("py:.*", r"rapidsmpf(\..*)?"), + ("py:.*", r"kvikio(\..*)?"), ("py:.*", r"ray(\..*)?"), ("py:.*", r"distributed(\..*)?"), ("py:.*", r"dask_cuda(\..*)?"), diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index 5c099b8b7177..7a2a17d4a629 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -9,15 +9,15 @@ multi-GPU engines. ```{eval-rst} .. autoclass:: cudf_polars.engine.ray.RayEngine - :members: from_options, gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks + :members: from_options, gather_cluster_info, gather_statistics, global_statistics, gather_io_summary, shutdown, nranks :show-inheritance: .. autoclass:: cudf_polars.engine.dask.DaskEngine - :members: from_options, gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks + :members: from_options, gather_cluster_info, gather_statistics, global_statistics, gather_io_summary, shutdown, nranks :show-inheritance: .. autoclass:: cudf_polars.engine.spmd.SPMDEngine - :members: from_options, gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks, rank, comm, context + :members: from_options, gather_cluster_info, gather_statistics, global_statistics, gather_io_summary, shutdown, nranks, rank, comm, context :show-inheritance: .. autoclass:: cudf_polars.engine.default_singleton_engine.DefaultSingletonEngine @@ -29,7 +29,7 @@ The engine classes share a common base class: ```{eval-rst} .. autoclass:: cudf_polars.engine.core.StreamingEngine - :members: gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks + :members: gather_cluster_info, gather_statistics, global_statistics, gather_io_summary, shutdown, nranks :show-inheritance: .. autoclass:: cudf_polars.engine.core.ClusterInfo diff --git a/docs/cudf/source/cudf_polars/profiling.md b/docs/cudf/source/cudf_polars/profiling.md index e7f22529a2e9..22d461fe2912 100644 --- a/docs/cudf/source/cudf_polars/profiling.md +++ b/docs/cudf/source/cudf_polars/profiling.md @@ -58,6 +58,87 @@ print(total) ``` +## I/O Statistics + +The same `statistics=True` also turns on [KvikIO I/O statistics][kvikio-stats] on every rank, +which report what storage did. `gather_io_summary()` returns one `kvikio.Summary` per rank, +keyed by rank index: + +```python +import polars as pl +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.ray import RayEngine + +opts = StreamingOptions(statistics=True) + +with RayEngine.from_options(opts) as engine: + pl.scan_parquet("/data/*.parquet").collect(engine=engine) + + for rank, summary in engine.gather_io_summary().items(): + print(f"--- rank {rank} ---") + print(summary) +``` + +`clear=True` restarts each rank's measured span after reading, scoping the next gather to +whatever follows. A rank that is not counting is absent, so the result is empty unless +`statistics=True` is set. That is distinct from a zeroed summary, which means the rank was +counting and did no I/O. + +Printing a summary gives KvikIO's own report: + +``` +KvikIO I/O summary + wall time 122.55 ms + busy time 18.40 ms (15.02 % of the wall time) + busy bandwidth 66.44 MB/s + operations 12 (12 read, 0 write) + mean duration 3.90 ms + bytes 1.17 MiB of 1.17 MiB requested (1.17 MiB read, 0 B written) + errors 0 + backend POSIX 1.17 MiB in 12 ops, 46.83 ms, 26.11 MB/s + backend GDS unused + backend MMAP unused + backend REMOTE_HTTP unused + backend REMOTE_HDFS unused +``` + +Every row is also an attribute, `s.bytes_read`, `s.busy_ns` and so on. See the +[KvikIO reference][kvikio-stats] for the full set. + +**Busy time** is the one worth explaining. It counts only the stretches with at least one read in +flight, so overlapping reads count once and the gaps between them count as idle. Rates divided by +it measure the storage rather than the query, which is why a query that reads for 10 ms and then +computes for 90 ms is not reported as ten times slower than its disks really are. + +### What is and is not counted + +Counting happens per process, so what a summary covers depends on what else shares that +process. With {class}`~cudf_polars.engine.ray.RayEngine` and +{class}`~cudf_polars.engine.dask.DaskEngine` each rank has a process to itself, so a summary +covers only cudf-polars[^shared-worker]. With {class}`~cudf_polars.engine.spmd.SPMDEngine` +cudf-polars shares your script's process, so KvikIO operations your own code performs are +counted too. + +Some I/O never reaches the monitor. On a system with working GDS the cuFile asynchronous API +reports nothing, the batch API reports nothing, and anything cudf-polars reads outside KvikIO is +invisible. + +[^shared-worker]: A Dask worker can host more than one rank if you run several engines, or other + Dask work, against one cluster. Neither is a recommended setup, and the summaries would be + mixed together. + + +### In the benchmarks + +The PDS benchmark runners record per-rank I/O summaries alongside the streaming statistics when +`--rapidsmpf-statistics` is passed. Each record carries an `io_summaries` dict keyed by rank, so I/O skew stays queryable across a whole sweep: + +```bash +python -m cudf_polars.streaming.benchmarks.pdsh \ + --path /data/tpch/ --output results.json --frontend ray \ + --rapidsmpf-statistics +``` + ## GPU Profiling For streaming queries, we recommend profiling with [NVIDIA NSight Systems][nsight]. `cudf-polars` @@ -184,6 +265,7 @@ shape: (2, 3) [nsight]: https://developer.nvidia.com/nsight-systems [nvtx]: https://nvidia.github.io/NVTX/ +[kvikio-stats]: https://docs.rapids.ai/api/kvikio/nightly/statistics/ [rapidsmpf-stats]: https://docs.rapids.ai/api/rapidsmpf/nightly/statistics/ [structlog]: https://www.structlog.org/en/stable/ [structlog-configure]: https://www.structlog.org/en/stable/configuration.html diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 3a58ef3ac1c3..b55ddbdf136a 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeVar import cuda.core +import kvikio import polars as pl @@ -96,6 +97,107 @@ def reset_statistics_from_options( return statistics +def make_kvikio_monitor(options: Options) -> kvikio.SummaryMonitor | None: + """ + Create a kvikio I/O monitor if statistics are enabled in ``options``. + + Parameters + ---------- + options + Options providing the enabled setting. The same + ``RAPIDSMPF_STATISTICS`` knob that gates rapidsmpf statistics. + + Returns + ------- + kvikio.SummaryMonitor + A monitor, already counting, if statistics are enabled. + None + If they are not. + + Notes + ----- + kvikio has no enable/disable: the existence of a monitor is what turns + counting on for the process, so "disabled" means "no monitor". + + With :class:`~cudf_polars.engine.spmd.SPMDEngine` the monitor counts every + thread's kvikio I/O in the script's process, so user code performing kvikio + reads is counted too. It cannot attribute I/O to a particular query. + """ + if not Statistics.from_options(options).enabled: + return None + return kvikio.SummaryMonitor() + + +def reset_kvikio_monitor_from_options( + monitor: kvikio.SummaryMonitor | None, options: Options +) -> kvikio.SummaryMonitor | None: + """ + Bring a kvikio I/O monitor into line with ``options``. + + Parameters + ---------- + monitor + The rank's existing monitor, if it has one. + options + Options providing the new enabled setting. + + Returns + ------- + kvikio.SummaryMonitor + The reset or newly created monitor, if statistics are enabled. + None + If they are not, in which case any existing monitor has been stopped. + """ + if not Statistics.from_options(options).enabled: + if monitor is not None: + monitor.stop() + return None + if monitor is None: + return kvikio.SummaryMonitor() + monitor.reset() + return monitor + + +def take_io_summary( + monitor: kvikio.SummaryMonitor | None, *, clear: bool +) -> kvikio.Summary | None: + """ + Read a rank's I/O totals, optionally restarting the measured span. + + Parameters + ---------- + monitor + The rank's monitor, or ``None`` if it is not counting. + clear + If ``True``, reset the monitor after reading, so the returned summary + is the last word on the span that just ended. + + Returns + ------- + kvikio.Summary + The totals so far. + None + If ``monitor`` is ``None``. + + Notes + ----- + ``None`` means "this rank was not counting", which is not the same as a + zeroed summary meaning "this rank did no I/O". + + ``get()`` and ``reset()`` are two separate calls, so an operation + completing between them is counted in the returned summary but dropped + from the next one. Use ``kvikio.Summary.since(previous)`` if you need + gapless differencing. + """ + if monitor is None: + return None + # Read before the reset, so the returned summary still carries the span. + summary = monitor.get() + if clear: + monitor.reset() + return summary + + def resolve_rapidsmpf_options(rapidsmpf_options: Options | None) -> Options: """ Resolve ``rapidsmpf_options`` and apply cross-frontend defaults. @@ -318,6 +420,30 @@ def gather_statistics(self, *, clear: bool = False) -> list[Statistics]: """ raise NotImplementedError + def gather_io_summary(self, *, clear: bool = False) -> dict[int, kvikio.Summary]: + """ + Collect kvikio I/O statistics from every rank. + + Parameters + ---------- + clear + If ``True``, restart each rank's measured span after reading, so + the next call describes only what followed this one. + + Returns + ------- + A :class:`kvikio.Summary` per rank, keyed by rank index and in rank + order. A rank that is not counting is absent, so the result is empty + unless the ``statistics`` option is enabled. + + Examples + -------- + >>> for rank, summary in engine.gather_io_summary().items(): # doctest: +SKIP + ... print(f"--- rank {rank} ---") + ... print(summary) + """ + raise NotImplementedError + def global_statistics(self, *, clear: bool = False) -> Statistics: """ Collect statistics from every rank and merge them into a single global statistics. diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 86a409a7413e..f21b96ba0a43 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -15,6 +15,7 @@ import distributed import distributed.system +import kvikio import kvikio.defaults import pynvml import ucxx._lib.libucxx as ucx_api @@ -39,8 +40,11 @@ check_reserved_keys, drop_if_replicated, evaluate_on_rank, + make_kvikio_monitor, + reset_kvikio_monitor_from_options, reset_statistics_from_options, resolve_rapidsmpf_options, + take_io_summary, ) from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, @@ -136,6 +140,7 @@ class _WorkerContext: quent_worker: cudf_polars.quent._types.Worker statistics: Statistics mr: RmmResourceAdaptor | None = None # set after `Context` is built (below). + kvikio_monitor: kvikio.SummaryMonitor | None = None def _worker_evaluate_persisted( @@ -436,6 +441,7 @@ def _setup_worker( quent_worker=quent_worker, quent_logger=quent_logger, statistics=statistics, + kvikio_monitor=make_kvikio_monitor(options), ) setattr(dask_worker, attr, mp_ctx) if mp_ctx.quent_logger is not None: @@ -478,6 +484,9 @@ def _teardown_worker( if mp_ctx.ctx is not None: mp_ctx.ctx.shutdown() finally: + if mp_ctx.kvikio_monitor is not None: + mp_ctx.kvikio_monitor.stop() + mp_ctx.kvikio_monitor = None mp_ctx.ctx = None mp_ctx.comm = None mp_ctx.base_mr = None @@ -534,6 +543,9 @@ def _reset_worker( options = Options.deserialize(rapidsmpf_options_as_bytes) mp_ctx.statistics = reset_statistics_from_options(mp_ctx.statistics, options) mp_ctx.statistics.clear() + mp_ctx.kvikio_monitor = reset_kvikio_monitor_from_options( + mp_ctx.kvikio_monitor, options + ) mp_ctx.ctx = Context.from_options( mp_ctx.comm.logger, mp_ctx.base_mr, options, mp_ctx.statistics ) @@ -631,6 +643,34 @@ def _get_statistics( return mp_ctx.comm.rank, stats +def _get_io_summary( + *, clear: bool, uid: str, dask_worker: distributed.Worker | None = None +) -> tuple[int, kvikio.Summary | None]: + """ + Return this worker's ``(rank, Summary)`` pair of kvikio I/O totals. + + The rank is used on the client to produce a rank-ordered list. + + Parameters + ---------- + clear + If ``True``, restart this worker's measured span after reading. + uid + Cluster instance identifier used to look up the per-worker context. + dask_worker + Injected by ``distributed`` when called via :meth:`distributed.Client.run`. + + Returns + ------- + Pair of ``(rank, Summary)``, the summary being ``None`` if this worker is + not counting. + """ + assert dask_worker is not None + mp_ctx: _WorkerContext = getattr(dask_worker, f"_cudf_polars_mp_context_{uid}") + assert mp_ctx.comm is not None + return mp_ctx.comm.rank, take_io_summary(mp_ctx.kvikio_monitor, clear=clear) + + def _worker_evaluate( ir: IR, config_options: ConfigOptions[StreamingExecutor], @@ -1170,6 +1210,25 @@ def gather_statistics(self, *, clear: bool = False) -> list[Statistics]: """ return list(self._run_by_rank(_get_statistics, clear=clear).values()) + def gather_io_summary(self, *, clear: bool = False) -> dict[int, kvikio.Summary]: + """ + Collect kvikio I/O statistics from every rank via ``client.run``. + + Parameters + ---------- + clear + If ``True``, restart each rank's measured span after reading. + + Returns + ------- + A :class:`kvikio.Summary` per rank, keyed by rank index, omitting + ranks that are not counting. + """ + summaries = self._run_by_rank(_get_io_summary, clear=clear) + return { + rank: summary for rank, summary in summaries.items() if summary is not None + } + def shutdown(self) -> None: """ Shut down all Dask workers' GPU resources. diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 38b088756926..fe3c8ae3b0a9 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -10,6 +10,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast +import kvikio import kvikio.defaults import ray import ray.exceptions @@ -35,8 +36,11 @@ check_reserved_keys, drop_if_replicated, evaluate_on_rank, + make_kvikio_monitor, + reset_kvikio_monitor_from_options, reset_statistics_from_options, resolve_rapidsmpf_options, + take_io_summary, ) from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, @@ -269,6 +273,7 @@ def __init__( rapidsmpf_options_as_bytes ) self._rapidsmpf_statistics = Statistics.from_options(self._rapidsmpf_options) + self._kvikio_monitor = make_kvikio_monitor(self._rapidsmpf_options) self._nranks: int = nranks self._py_executor = ThreadPoolExecutor( max_workers=num_py_executors, @@ -382,6 +387,9 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes, kvikio_nthreads: int) -> N self._rapidsmpf_statistics, self._rapidsmpf_options ) self._rapidsmpf_statistics.clear() + self._kvikio_monitor = reset_kvikio_monitor_from_options( + self._kvikio_monitor, self._rapidsmpf_options + ) assert self._base_mr is not None self._ctx = Context.from_options( self._comm.logger, @@ -420,6 +428,9 @@ def shutdown(self) -> None: if self._ctx is not None: self._ctx.shutdown() finally: + if self._kvikio_monitor is not None: + self._kvikio_monitor.stop() + self._kvikio_monitor = None self._ctx = None self._comm = None self._mr = None @@ -465,6 +476,24 @@ def get_statistics(self, *, clear: bool = False) -> tuple[int, Statistics]: return self._comm.rank, detached return self._comm.rank, stats + def get_io_summary( + self, *, clear: bool = False + ) -> tuple[int, kvikio.Summary | None]: + """ + Return this rank's index and its kvikio I/O totals. + + Parameters + ---------- + clear + If ``True``, restart this rank's measured span after reading. + + Returns + ------- + This rank's index, and its totals or ``None`` if it is not counting. + """ + assert self._comm is not None + return self._comm.rank, take_io_summary(self._kvikio_monitor, clear=clear) + def evaluate_polars_ir( self, ir: IR, @@ -1030,6 +1059,27 @@ def gather_statistics(self, *, clear: bool = False) -> list[Statistics]: ).values() ) + def gather_io_summary(self, *, clear: bool = False) -> dict[int, kvikio.Summary]: + """ + Collect kvikio I/O statistics from every rank via Ray. + + Parameters + ---------- + clear + If ``True``, restart each rank's measured span after reading. + + Returns + ------- + A :class:`kvikio.Summary` per rank, keyed by rank index, omitting + ranks that are not counting. + """ + summaries = self._gather_by_rank( + [rank.get_io_summary.remote(clear=clear) for rank in self.rank_actors] + ) + return { + rank: summary for rank, summary in summaries.items() if summary is not None + } + def shutdown(self) -> None: """ Shut down all rank actors and release resources. diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 10db61bd388c..fbb375dc3ace 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -11,6 +11,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast +import kvikio import kvikio.defaults import pylibcudf as plc @@ -40,8 +41,11 @@ all_gather_host_data, check_reserved_keys, evaluate_on_rank, + make_kvikio_monitor, + reset_kvikio_monitor_from_options, reset_statistics_from_options, resolve_rapidsmpf_options, + take_io_summary, ) from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, @@ -463,6 +467,8 @@ def __init__( self._py_executor: ThreadPoolExecutor | None = None self._store_uid = uuid.uuid4().hex exit_stack = contextlib.ExitStack() + self._kvikio_monitor = make_kvikio_monitor(self.rapidsmpf_options) + exit_stack.callback(self._stop_kvikio_monitor) # TODO: there's no reason our API needs a plain dict[str, Any] rather than # a typed config object here. @@ -534,6 +540,12 @@ def __init__( exit_stack.close() raise + def _stop_kvikio_monitor(self) -> None: + """Stop this rank's kvikio monitor if any; called from exit-stack.""" + if self._kvikio_monitor is not None: + self._kvikio_monitor.stop() + self._kvikio_monitor = None + def _cleanup_ctx(self) -> None: """ Shut down the current ``self._ctx`` if any; called from exit-stack. @@ -632,6 +644,9 @@ def _reset( self._comm.progress_thread.statistics, rapidsmpf_options ) statistics.clear() + self._kvikio_monitor = reset_kvikio_monitor_from_options( + self._kvikio_monitor, rapidsmpf_options + ) self._ctx = Context.from_options( self._comm.logger, self._base_mr, rapidsmpf_options, statistics @@ -789,6 +804,34 @@ def gather_statistics(self, *, clear: bool = False) -> list[Statistics]: self.context.statistics().clear() return [Statistics.deserialize(r) for r in results] + def gather_io_summary(self, *, clear: bool = False) -> dict[int, kvikio.Summary]: + """ + Collect kvikio I/O statistics from every rank via an all-gather. + + This is a collective operation, every rank must call it. + + Parameters + ---------- + clear + If ``True``, restart each rank's measured span after reading. + + Returns + ------- + A :class:`kvikio.Summary` per rank, keyed by rank index, omitting + ranks that are not counting. + """ + summary = take_io_summary(self._kvikio_monitor, clear=clear) + # A rank that is not counting sends nothing, which is distinguishable + # from a zeroed summary because the latter is a fixed, non-empty size. + data = b"" if summary is None else summary.serialize() + with reserve_op_id() as op_id: + results = all_gather_host_data(self.comm, self.context.br(), op_id, data) + return { + rank: kvikio.Summary.deserialize(r) + for rank, r in enumerate(results) + if r != b"" + } + def shutdown(self) -> None: """ Shut down the engine and release all owned resources. diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py new file mode 100644 index 000000000000..19999f0b7e54 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Read a saved benchmark results file and print it in human-readable form. + +The benchmark runners write one JSON object per line to ``--output``, appending +across runs. This reads such a file back and prints the query timings, and the +per-rank I/O summaries when the run was made with ``--rapidsmpf-statistics``. + + python -m cudf_polars.streaming.benchmarks.print_results_file results.json +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from statistics import mean +from typing import TYPE_CHECKING, Any + +from rapidsmpf.utils.string import format_bytes + +_DESCRIPTION = ( + "Read a saved benchmark results file and print the query timings, and the " + "per-rank I/O summaries when the run was made with --rapidsmpf-statistics." +) + +if TYPE_CHECKING: + from collections.abc import Iterator + + +def load_runs(path: Path) -> list[dict[str, Any]]: + """ + Read every run from a results file. + + Parameters + ---------- + path + Path to a file written by a benchmark runner's ``--output``. + + Returns + ------- + One dictionary per run, in the order they were appended. + + Raises + ------ + ValueError + If the file contains no runs. + """ + with path.open() as f: + runs = [json.loads(line) for line in f if line.strip()] + if not runs: + raise ValueError(f"{path} contains no runs") + return runs + + +def iter_records(run: dict[str, Any]) -> Iterator[dict[str, Any]]: + """ + Yield every per-iteration record of a run, ordered by query then iteration. + + Parameters + ---------- + run + A single run, as returned by :func:`load_runs`. + + Yields + ------ + The run's per-iteration records. + """ + for _, records in sorted(run.get("records", {}).items(), key=lambda kv: int(kv[0])): + yield from records + + +def print_header(run: dict[str, Any]) -> None: + """ + Print the run's identifying configuration. + + Parameters + ---------- + run + A single run, as returned by :func:`load_runs`. + """ + print(f"run : {run.get('run_id')} ({run.get('timestamp')})") + print(f"engine : {run.get('engine_name')} frontend={run.get('frontend')}") + print(f"dataset : {run.get('dataset_path')} scale={run.get('scale_factor')}") + print(f"workers : {run.get('n_workers')} iterations={run.get('iterations')}") + + +def print_timings(run: dict[str, Any]) -> None: + """ + Print min, max and mean duration for each query. + + Parameters + ---------- + run + A single run, as returned by :func:`load_runs`. + """ + print("\nTimings") + print(f" {'query':>6} {'iters':>5} {'min':>9} {'max':>9} {'mean':>9}") + total = 0.0 + for query, records in sorted( + run.get("records", {}).items(), key=lambda kv: int(kv[0]) + ): + durations = [r["duration"] for r in records if r.get("status") == "success"] + if not durations: + print(f" {query:>6} {'-':>5} {'no successful iterations':>31}") + continue + total += mean(durations) + print( + f" {query:>6} {len(durations):>5} {min(durations):>8.4f}s " + f"{max(durations):>8.4f}s {mean(durations):>8.4f}s" + ) + if total > 0: + print(f" {'total':>6} {'':>5} {'':>9} {'':>9} {total:>8.4f}s") + + +def _io_summaries(record: dict[str, Any]) -> dict[int, dict[str, Any]]: + """ + Return a record's I/O summaries keyed by rank. + + Parameters + ---------- + record + One per-iteration record. + + Returns + ------- + The summaries, empty if the iteration recorded none. + """ + raw = record.get("io_summaries") + if not raw: + return {} + return {int(rank): s for rank, s in raw.items()} + + +def print_io_summaries(run: dict[str, Any]) -> None: + """ + Print the per-rank I/O summaries, one row per rank per iteration. + + Kept per rank rather than totalled, since the point of per-rank I/O + statistics is to expose skew that a total would hide. + + Parameters + ---------- + run + A single run, as returned by :func:`load_runs`. + """ + rows = [ + (r["query"], r["iteration"], rank, s) + for r in iter_records(run) + for rank, s in sorted(_io_summaries(r).items()) + ] + if not rows: + # Either the run predates I/O statistics, or it was made without + # `--rapidsmpf-statistics`, in which case no rank counts anything. + print("\nI/O: not collected (run with --rapidsmpf-statistics)") + return + + print("\nI/O per rank") + print( + f" {'query':>6} {'iter':>4} {'rank':>4} {'ops':>7} {'read':>11} " + f"{'busy':>9} {'busy%':>6} {'bandwidth':>11} backends" + ) + for query, iteration, rank, s in rows: + backends = ( + ",".join( + name + for name, totals in s.get("by_backend", {}).items() + if totals.get("num_ops") + ) + or "-" + ) + print( + f" {query:>6} {iteration:>4} {rank:>4} {s['num_ops']:>7} " + f"{format_bytes(s['bytes_read']):>11} {s['busy_ns'] / 1e6:>7.1f}ms " + f"{s['busy_fraction'] * 100:>5.1f}% " + f"{s['busy_bytes_per_sec'] / 1e9:>7.2f}GB/s {backends}" + ) + + # Skew is what the per-rank view is for, so say it outright rather than + # leaving it to be eyeballed across rows. + by_iteration: dict[tuple[Any, Any], list[int]] = {} + for query, iteration, _, s in rows: + by_iteration.setdefault((query, iteration), []).append(s["bytes_read"]) + shared = {k: v for k, v in by_iteration.items() if len(v) > 1} + + # A rank that read nothing while its peers did is the most extreme skew there + # is, and it has no finite ratio, so it is reported rather than skipped. + starved = sorted(k for k, v in shared.items() if min(v) == 0 < max(v)) + if starved: + where = ", ".join(f"q{q} iter {i}" for q, i in starved) + print(f"\n ranks that read nothing while peers did: {where}") + + worst = max( + ((max(v) / min(v), k) for k, v in shared.items() if min(v) > 0), + default=None, + ) + if worst is not None: + ratio, (query, iteration) = worst + label = "widest finite read skew" if starved else "widest read skew" + print(f"\n {label}: {ratio:.2f}x (query {query}, iteration {iteration})") + + +def main(args: argparse.Namespace) -> None: + """ + Print the requested runs. + + Parameters + ---------- + args + Parsed command-line arguments. + """ + runs = load_runs(args.path) + selected = runs if args.all else runs[-1:] + if not args.all and len(runs) > 1: + print(f"({len(runs)} runs in {args.path}, showing the last, --all for every)\n") + + for i, run in enumerate(selected): + if i: + print() + print("=" * 78) + print_header(run) + print("=" * 78) + print_timings(run) + if not args.no_io: + print_io_summaries(run) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """ + Parse command-line arguments. + + Parameters + ---------- + argv + Argument list, or ``None`` to read ``sys.argv``. + + Returns + ------- + The parsed arguments. + """ + parser = argparse.ArgumentParser(description=_DESCRIPTION) + parser.add_argument( + "path", type=Path, help="Results file written by a runner's --output." + ) + parser.add_argument( + "--all", + action="store_true", + help="Print every run in the file, not just the last one.", + ) + parser.add_argument( + "--no-io", + action="store_true", + help="Skip the per-rank I/O summaries.", + ) + return parser.parse_args(argv) + + +if __name__ == "__main__": + main(parse_args()) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index d64949b2e4e7..91692db91462 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -67,6 +67,8 @@ pynvml = None try: + from rapidsmpf.statistics import Statistics + import cudf_polars.dsl.tracing import cudf_polars.quent from cudf_polars.dsl.ir import IRExecutionContext @@ -252,6 +254,7 @@ class SuccessRecord: iteration: int duration: float statistics: dict[str, Any] | None = None + io_summaries: dict[str, dict[str, Any]] | None = None traces: list[dict[str, Any]] | None = None validation_result: ValidationResult | None = None status: Literal["success"] = "success" @@ -263,6 +266,7 @@ def new( iteration: int, duration: float, statistics: dict[str, Any] | None = None, + io_summaries: dict[str, dict[str, Any]] | None = None, traces: list[dict[str, Any]] | None = None, ) -> SuccessRecord: """Create a Record from plain data.""" @@ -271,6 +275,7 @@ def new( iteration=iteration, duration=duration, statistics=statistics, + io_summaries=io_summaries, traces=traces, ) @@ -935,6 +940,23 @@ def _collect_statistics(engine: pl.GPUEngine | None) -> dict[str, Any] | None: return engine.global_statistics(clear=True).to_dict() +def _collect_io_summaries( + engine: pl.GPUEngine | None, +) -> dict[str, dict[str, Any]] | None: + """Gather + clear kvikio I/O statistics, keyed by rank.""" + if engine is None: + return None + if not isinstance(engine, StreamingEngine): + return None + if not Statistics.from_options(engine.rapidsmpf_options).enabled: + return None + # String keys, since the record is written as JSON. + return { + str(rank): dataclasses.asdict(summary) + for rank, summary in engine.gather_io_summary(clear=True).items() + } + + def run_polars_query_iteration( q_id: int, iteration: int, @@ -960,6 +982,10 @@ def run_polars_query_iteration( # Once we support polars 1.40, we should remove this result = result.with_columns(*result_casts) + # I/O first: gathering it is itself a RapidsMPF collective, so doing it after + # `_collect_statistics` would leave those events in the freshly cleared + # counters and report them against the next iteration. + io_summaries = _collect_io_summaries(engine) statistics = _collect_statistics(engine) if expected is not None: @@ -989,6 +1015,7 @@ def run_polars_query_iteration( iteration=iteration, duration=duration, statistics=statistics, + io_summaries=io_summaries, validation_result=validation_result, ) diff --git a/python/cudf_polars/tests/streaming/test_statistics.py b/python/cudf_polars/tests/streaming/test_statistics.py index 65cf2c398a1c..adf1702fe36c 100644 --- a/python/cudf_polars/tests/streaming/test_statistics.py +++ b/python/cudf_polars/tests/streaming/test_statistics.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for ``gather_statistics`` / ``global_statistics`` on streaming engines.""" +"""Tests for statistics gathering on streaming engines: rapidsmpf and kvikio.""" from __future__ import annotations @@ -8,12 +8,15 @@ import pytest +import polars as pl + from rapidsmpf.statistics import Statistics from cudf_polars.engine.options import StreamingOptions if TYPE_CHECKING: from collections.abc import Callable + from pathlib import Path from cudf_polars.engine.core import StreamingEngine @@ -58,3 +61,76 @@ def test_statistics(engine: StreamingEngine) -> None: # rank and clear the stats before the allgather event loop is # removed. So we might see an event loop stat, but no other. assert s.list_stat_names() == [] or s.list_stat_names() == ["event-loop-total"] + + +@pytest.fixture +def scan_query(tmp_path: Path) -> pl.LazyFrame: + """A parquet scan, so that the engine actually performs I/O.""" + path = tmp_path / "data.parquet" + pl.DataFrame({"a": range(1000), "b": range(1000)}).write_parquet(path) + return pl.scan_parquet(path) + + +def test_io_summary(engine: StreamingEngine, scan_query: pl.LazyFrame) -> None: + """gather_io_summary reports what each rank read.""" + scan_query.collect(engine=engine) + + # Statistics are enabled on this fixture, so every rank is counting. + gathered = engine.gather_io_summary() + assert sorted(gathered) == list(range(engine.nranks)) + summaries = list(gathered.values()) + # Only some ranks may have been given a file to read, but not none of them. + assert sum(s.num_ops for s in summaries) > 0 + assert sum(s.bytes_read for s in summaries) > 0 + + for s in summaries: + # A summary is self-consistent, whether or not this rank did any I/O. + assert s.busy_ns <= s.wall_ns + assert s.bytes_read + s.bytes_written == s.bytes_transferred + assert sum(b["num_ops"] for b in s.by_backend.values()) == s.num_ops + assert s.num_reads + s.num_writes == s.num_ops + assert s.num_errors == 0 + # kvikio renders the report, so this is a check that it is reachable + # per rank rather than a check of its content. + assert "KvikIO I/O summary" in str(s) + + +def test_io_summary_clear_starts_a_new_span( + engine: StreamingEngine, scan_query: pl.LazyFrame +) -> None: + """clear=True returns the totals so far and restarts the span.""" + scan_query.collect(engine=engine) + + before = engine.gather_io_summary(clear=True).values() + assert sum(s.num_ops for s in before) > 0 + + # Nothing between the clear and this gather, so the new span is empty. + after = engine.gather_io_summary() + assert len(after) == engine.nranks + assert sum(s.num_ops for s in after.values()) == 0 + assert sum(s.bytes_transferred for s in after.values()) == 0 + + +def test_io_summary_is_off_without_statistics( + streaming_engine_factory: Callable[..., StreamingEngine], + scan_query: pl.LazyFrame, +) -> None: + """No monitor is created when statistics are off. One is created again when they are turned back on.""" + engine = streaming_engine_factory( + StreamingOptions(statistics=False, max_rows_per_partition=10) + ) + scan_query.collect(engine=engine) + + # An absent rank is "this rank was not counting", which is not the same as + # a zeroed summary meaning "this rank did no I/O". + assert engine.gather_io_summary() == {} + + # Turning statistics back on has to build a fresh monitor, the previous one + # having been stopped rather than paused. + engine = streaming_engine_factory( + StreamingOptions(statistics=True, max_rows_per_partition=10) + ) + scan_query.collect(engine=engine) + summaries = engine.gather_io_summary() + assert sorted(summaries) == list(range(engine.nranks)) + assert sum(s.num_ops for s in summaries.values()) > 0 From ba4362863f235486a305b7517decfa5ff8086ece Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 24 Aug 2026 09:08:54 +0200 Subject: [PATCH 2/9] tests --- docs/cudf/source/cudf_polars/profiling.md | 2 +- python/cudf_polars/cudf_polars/engine/dask.py | 5 +- python/cudf_polars/cudf_polars/engine/ray.py | 5 +- python/cudf_polars/cudf_polars/engine/spmd.py | 1 + .../benchmarks/print_results_file.py | 4 +- .../tests/streaming/benchmarks/__init__.py | 2 + .../benchmarks/test_print_results_file.py | 201 ++++++++++++++++++ .../tests/streaming/test_statistics.py | 5 + 8 files changed, 217 insertions(+), 8 deletions(-) create mode 100644 python/cudf_polars/tests/streaming/benchmarks/__init__.py create mode 100644 python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py diff --git a/docs/cudf/source/cudf_polars/profiling.md b/docs/cudf/source/cudf_polars/profiling.md index 22d461fe2912..5f6bbf4c866d 100644 --- a/docs/cudf/source/cudf_polars/profiling.md +++ b/docs/cudf/source/cudf_polars/profiling.md @@ -86,7 +86,7 @@ counting and did no I/O. Printing a summary gives KvikIO's own report: -``` +```text KvikIO I/O summary wall time 122.55 ms busy time 18.40 ms (15.02 % of the wall time) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index f21b96ba0a43..42b22131ac20 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1103,9 +1103,8 @@ def _reset( executor_options.setdefault("kvikio_nthreads", existing_kvikio_nthreads) engine_options = engine_options or {} - rapidsmpf_options_as_bytes = resolve_rapidsmpf_options( - rapidsmpf_options - ).serialize() + self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) + rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() ctx = self._dask_context # Reset all worker Contexts collectively. ``client.run`` blocks diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index fe3c8ae3b0a9..fc5a4f96365c 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -918,9 +918,8 @@ def _reset( if existing_kvikio_nthreads is not None: executor_options.setdefault("kvikio_nthreads", existing_kvikio_nthreads) engine_options = engine_options or {} - rapidsmpf_options_as_bytes = resolve_rapidsmpf_options( - rapidsmpf_options - ).serialize() + self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) + rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() # Reset all actor Contexts collectively. ``ray.get`` blocks until # every actor's reset returns; the per-actor barrier inside diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index fbb375dc3ace..9917799476d1 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -629,6 +629,7 @@ def _reset( "quent_context" ) rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) + self.rapidsmpf_options = rapidsmpf_options # Collective: synchronize all ranks before tearing down the Context. if self._comm.nranks > 1: diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py index 19999f0b7e54..93b10a0b845d 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py @@ -129,7 +129,9 @@ def _io_summaries(record: dict[str, Any]) -> dict[int, dict[str, Any]]: The summaries, empty if the iteration recorded none. """ raw = record.get("io_summaries") - if not raw: + if not isinstance(raw, dict): + # Absent, or written by a version that shaped it differently. The file + # comes from disk, so an unexpected shape is skipped rather than raised. return {} return {int(rank): s for rank, s in raw.items()} diff --git a/python/cudf_polars/tests/streaming/benchmarks/__init__.py b/python/cudf_polars/tests/streaming/benchmarks/__init__.py new file mode 100644 index 000000000000..d51c4fe1e089 --- /dev/null +++ b/python/cudf_polars/tests/streaming/benchmarks/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py b/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py new file mode 100644 index 000000000000..bad5beb35aed --- /dev/null +++ b/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the benchmark results-file printer.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any + +import pytest + +from cudf_polars.streaming.benchmarks.print_results_file import ( + _io_summaries, + iter_records, + load_runs, + main, + parse_args, +) + +if TYPE_CHECKING: + from collections.abc import Callable + from pathlib import Path + + +def _summary(**overrides: Any) -> dict[str, Any]: + """One rank's I/O summary.""" + return { + "num_ops": 10, + "bytes_read": 1024, + "busy_ns": 1_000_000, + "busy_fraction": 0.5, + "busy_bytes_per_sec": 1e9, + "by_backend": {"POSIX": {"num_ops": 10}, "GDS": {"num_ops": 0}}, + **overrides, + } + + +def _record(**overrides: Any) -> dict[str, Any]: + """One successful iteration.""" + return { + "query": 1, + "iteration": 0, + "duration": 0.25, + "status": "success", + **overrides, + } + + +def _run(*records: dict[str, Any], **overrides: Any) -> dict[str, Any]: + """A run holding ``records``, grouped by query as the runners write them.""" + by_query: dict[str, list[dict[str, Any]]] = {} + for record in records: + by_query.setdefault(str(record["query"]), []).append(record) + return { + "run_id": "abc", + "timestamp": "2026-01-01T00:00:00+00:00", + "engine_name": "cudf-polars", + "frontend": "ray", + "dataset_path": "/data", + "scale_factor": 10, + "n_workers": 2, + "iterations": 1, + "records": by_query, + **overrides, + } + + +@pytest.fixture +def report(tmp_path: Path, capsys: pytest.CaptureFixture) -> Callable[..., str]: + """Write runs as the NDJSON the runners append, print them, return the output.""" + + def run(*runs: dict[str, Any], args: tuple[str, ...] = ()) -> str: + path = tmp_path / "out.json" + path.write_text("".join(f"{json.dumps(r)}\n" for r in runs)) + main(parse_args([str(path), *args])) + return capsys.readouterr().out + + return run + + +def test_load_runs_reads_every_line(tmp_path: Path) -> None: + """Each line is a run, and blank lines are skipped.""" + path = tmp_path / "out.json" + path.write_text( + f"{json.dumps(_run(run_id='first'))}\n\n{json.dumps(_run(run_id='second'))}\n" + ) + assert [r["run_id"] for r in load_runs(path)] == ["first", "second"] + + +def test_load_runs_rejects_an_empty_file(tmp_path: Path) -> None: + """An empty file is an error rather than an empty report.""" + path = tmp_path / "empty.json" + path.write_text("") + with pytest.raises(ValueError, match="no runs"): + load_runs(path) + + +def test_iter_records_orders_by_query() -> None: + """Records come out grouped by query, whatever order the keys are in.""" + run = _run(_record(query=3), _record(query=1), _record(query=1, iteration=1)) + assert [(r["query"], r["iteration"]) for r in iter_records(run)] == [ + (1, 0), + (1, 1), + (3, 0), + ] + + +@pytest.mark.parametrize( + "raw, expected", + [ + # JSON object keys are strings, so they are cast back to ranks. + ({"0": None, "1": None}, [0, 1]), + # A run without statistics has nothing to report. + (None, []), + # The file comes from disk, so a surprise is skipped rather than raised. + ([{}], []), + ], +) +def test_io_summaries_keys(raw: Any, expected: list[int]) -> None: + """Rank keys are integers, and anything unexpected yields nothing.""" + assert sorted(_io_summaries({"io_summaries": raw})) == expected + + +def test_prints_timings_and_io(report: Callable[..., str]) -> None: + """The last run is printed, with a row per rank per iteration.""" + out = report(_run(_record(io_summaries={"0": _summary(), "1": _summary()}))) + assert "Timings" in out + assert "0.2500s" in out + # One row per rank, naming the backend that carried the work. + assert out.count("POSIX") == 2 + + +def test_reports_when_io_was_not_collected(report: Callable[..., str]) -> None: + """A run made without --rapidsmpf-statistics says so rather than printing nothing.""" + out = report(_run(_record())) + assert "not collected" in out + assert "I/O per rank" not in out + + +def test_failed_iterations_are_left_out_of_the_timings( + report: Callable[..., str], +) -> None: + """A query whose every iteration failed has no min, max or mean to report.""" + assert "no successful iterations" in report(_run(_record(status="error"))) + + +def test_a_rank_that_read_nothing_is_called_out(report: Callable[..., str]) -> None: + """Zero bytes on one rank has no finite ratio, so it gets its own line.""" + out = report( + _run( + _record( + io_summaries={ + "0": _summary(bytes_read=4096), + "1": _summary( + num_ops=0, + bytes_read=0, + busy_ns=0, + busy_fraction=0.0, + busy_bytes_per_sec=0.0, + by_backend={"POSIX": {"num_ops": 0}}, + ), + } + ) + ) + ) + assert "ranks that read nothing while peers did: q1 iter 0" in out + + +def test_read_skew_is_reported_as_a_ratio(report: Callable[..., str]) -> None: + """With every rank reading, the spread is the max over the min.""" + out = report( + _run( + _record( + io_summaries={ + "0": _summary(bytes_read=4000), + "1": _summary(bytes_read=1000), + } + ) + ) + ) + assert "widest read skew: 4.00x" in out + + +def test_all_prints_every_run(report: Callable[..., str]) -> None: + """Without --all only the last run is printed, since --output appends.""" + runs = (_run(run_id="older"), _run(run_id="newer")) + + out = report(*runs) + assert "newer" in out + assert "older" not in out + + out = report(*runs, args=("--all",)) + assert "older" in out + assert "newer" in out + + +def test_no_io_skips_the_io_section(report: Callable[..., str]) -> None: + """--no-io leaves the timings alone.""" + out = report(_run(_record(io_summaries={"0": _summary()})), args=("--no-io",)) + assert "Timings" in out + assert "I/O per rank" not in out diff --git a/python/cudf_polars/tests/streaming/test_statistics.py b/python/cudf_polars/tests/streaming/test_statistics.py index adf1702fe36c..9d3a4c3547ef 100644 --- a/python/cudf_polars/tests/streaming/test_statistics.py +++ b/python/cudf_polars/tests/streaming/test_statistics.py @@ -125,11 +125,16 @@ def test_io_summary_is_off_without_statistics( # a zeroed summary meaning "this rank did no I/O". assert engine.gather_io_summary() == {} + # A reset leaves `rapidsmpf_options` describing the reset, which is what + # callers read to decide whether anything is being counted. + assert not Statistics.from_options(engine.rapidsmpf_options).enabled + # Turning statistics back on has to build a fresh monitor, the previous one # having been stopped rather than paused. engine = streaming_engine_factory( StreamingOptions(statistics=True, max_rows_per_partition=10) ) + assert Statistics.from_options(engine.rapidsmpf_options).enabled scan_query.collect(engine=engine) summaries = engine.gather_io_summary() assert sorted(summaries) == list(range(engine.nranks)) From c2dde72c701f5f57a84ffa327061c09b7ae04ec6 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 24 Aug 2026 09:57:41 +0200 Subject: [PATCH 3/9] teardown fix --- python/cudf_polars/cudf_polars/engine/dask.py | 8 +++++--- python/cudf_polars/cudf_polars/engine/ray.py | 7 ++++--- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 42b22131ac20..5b8bacafa3d4 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -469,6 +469,11 @@ def _teardown_worker( mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None) traces = [] if mp_ctx is not None: + # First, so that a failure below cannot leave it counting. The monitor is + # process-global and the worker outlives this teardown. + if mp_ctx.kvikio_monitor is not None: + mp_ctx.kvikio_monitor.stop() + mp_ctx.kvikio_monitor = None if mp_ctx.quent_worker is not None and mp_ctx.quent_logger is not None: mp_ctx.quent_logger.emit(mp_ctx.quent_worker._exit()) traces = mp_ctx.quent_logger.drain() @@ -484,9 +489,6 @@ def _teardown_worker( if mp_ctx.ctx is not None: mp_ctx.ctx.shutdown() finally: - if mp_ctx.kvikio_monitor is not None: - mp_ctx.kvikio_monitor.stop() - mp_ctx.kvikio_monitor = None mp_ctx.ctx = None mp_ctx.comm = None mp_ctx.base_mr = None diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index fc5a4f96365c..2862d7354280 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -418,6 +418,10 @@ def shutdown(self) -> None: Raises `ray.exceptions.RayActorError`. """ + # First, so that a failure below cannot leave it counting. + if self._kvikio_monitor is not None: + self._kvikio_monitor.stop() + self._kvikio_monitor = None self._py_executor.shutdown(wait=True, cancel_futures=True) # Release resources in dependency order before exit_actor() terminates # the process. Shut down the Context explicitly on the same thread @@ -428,9 +432,6 @@ def shutdown(self) -> None: if self._ctx is not None: self._ctx.shutdown() finally: - if self._kvikio_monitor is not None: - self._kvikio_monitor.stop() - self._kvikio_monitor = None self._ctx = None self._comm = None self._mr = None From 1bfeb256cf308d1ea2d7ed8776e99b7997e68b80 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 25 Aug 2026 09:24:00 +0200 Subject: [PATCH 4/9] Cleanup based on review by @TomAugspurger --- docs/cudf/source/cudf_polars/profiling.md | 20 ++------ .../benchmarks/print_results_file.py | 47 +++++++++++-------- .../cudf_polars/streaming/benchmarks/utils.py | 47 +++++++++++++++++++ .../benchmarks/test_print_results_file.py | 41 ++++++++++++---- .../tests/streaming/test_statistics.py | 14 ++---- 5 files changed, 114 insertions(+), 55 deletions(-) diff --git a/docs/cudf/source/cudf_polars/profiling.md b/docs/cudf/source/cudf_polars/profiling.md index 5f6bbf4c866d..317259b64737 100644 --- a/docs/cudf/source/cudf_polars/profiling.md +++ b/docs/cudf/source/cudf_polars/profiling.md @@ -103,12 +103,8 @@ KvikIO I/O summary ``` Every row is also an attribute, `s.bytes_read`, `s.busy_ns` and so on. See the -[KvikIO reference][kvikio-stats] for the full set. - -**Busy time** is the one worth explaining. It counts only the stretches with at least one read in -flight, so overlapping reads count once and the gaps between them count as idle. Rates divided by -it measure the storage rather than the query, which is why a query that reads for 10 ms and then -computes for 90 ms is not reported as ten times slower than its disks really are. +[KvikIO reference][kvikio-stats] for the full set, and [busy time and bandwidth][kvikio-busy] +for how the busy figures are measured. ### What is and is not counted @@ -128,17 +124,6 @@ invisible. mixed together. -### In the benchmarks - -The PDS benchmark runners record per-rank I/O summaries alongside the streaming statistics when -`--rapidsmpf-statistics` is passed. Each record carries an `io_summaries` dict keyed by rank, so I/O skew stays queryable across a whole sweep: - -```bash -python -m cudf_polars.streaming.benchmarks.pdsh \ - --path /data/tpch/ --output results.json --frontend ray \ - --rapidsmpf-statistics -``` - ## GPU Profiling For streaming queries, we recommend profiling with [NVIDIA NSight Systems][nsight]. `cudf-polars` @@ -266,6 +251,7 @@ shape: (2, 3) [nsight]: https://developer.nvidia.com/nsight-systems [nvtx]: https://nvidia.github.io/NVTX/ [kvikio-stats]: https://docs.rapids.ai/api/kvikio/nightly/statistics/ +[kvikio-busy]: https://docs.rapids.ai/api/kvikio/nightly/statistics/#busy-time-and-bandwidth [rapidsmpf-stats]: https://docs.rapids.ai/api/rapidsmpf/nightly/statistics/ [structlog]: https://www.structlog.org/en/stable/ [structlog-configure]: https://www.structlog.org/en/stable/configuration.html diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py index 93b10a0b845d..571b0cb491c7 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py @@ -21,6 +21,8 @@ from rapidsmpf.utils.string import format_bytes +from cudf_polars.streaming.benchmarks.utils import SuccessRecord, record_from_dict + _DESCRIPTION = ( "Read a saved benchmark results file and print the query timings, and the " "per-rank I/O summaries when the run was made with --rapidsmpf-statistics." @@ -29,10 +31,12 @@ if TYPE_CHECKING: from collections.abc import Iterator + from cudf_polars.streaming.benchmarks.utils import FailedRecord + def load_runs(path: Path) -> list[dict[str, Any]]: """ - Read every run from a results file. + Read every run from a results file, without interpreting it. Parameters ---------- @@ -41,12 +45,12 @@ def load_runs(path: Path) -> list[dict[str, Any]]: Returns ------- - One dictionary per run, in the order they were appended. + One decoded line per run, in the order they were appended. Raises ------ ValueError - If the file contains no runs. + If the file holds no runs. """ with path.open() as f: runs = [json.loads(line) for line in f if line.strip()] @@ -55,7 +59,7 @@ def load_runs(path: Path) -> list[dict[str, Any]]: return runs -def iter_records(run: dict[str, Any]) -> Iterator[dict[str, Any]]: +def iter_records(run: dict[str, Any]) -> Iterator[SuccessRecord | FailedRecord]: """ Yield every per-iteration record of a run, ordered by query then iteration. @@ -68,8 +72,8 @@ def iter_records(run: dict[str, Any]) -> Iterator[dict[str, Any]]: ------ The run's per-iteration records. """ - for _, records in sorted(run.get("records", {}).items(), key=lambda kv: int(kv[0])): - yield from records + for _, records in sorted(run["records"].items(), key=lambda kv: int(kv[0])): + yield from map(record_from_dict, records) def print_header(run: dict[str, Any]) -> None: @@ -99,10 +103,12 @@ def print_timings(run: dict[str, Any]) -> None: print("\nTimings") print(f" {'query':>6} {'iters':>5} {'min':>9} {'max':>9} {'mean':>9}") total = 0.0 - for query, records in sorted( - run.get("records", {}).items(), key=lambda kv: int(kv[0]) - ): - durations = [r["duration"] for r in records if r.get("status") == "success"] + for query, records in sorted(run["records"].items(), key=lambda kv: int(kv[0])): + durations = [ + r.duration + for r in map(record_from_dict, records) + if isinstance(r, SuccessRecord) + ] if not durations: print(f" {query:>6} {'-':>5} {'no successful iterations':>31}") continue @@ -115,7 +121,9 @@ def print_timings(run: dict[str, Any]) -> None: print(f" {'total':>6} {'':>5} {'':>9} {'':>9} {total:>8.4f}s") -def _io_summaries(record: dict[str, Any]) -> dict[int, dict[str, Any]]: +def _io_summaries( + record: SuccessRecord | FailedRecord, +) -> dict[int, dict[str, Any]]: """ Return a record's I/O summaries keyed by rank. @@ -128,12 +136,9 @@ def _io_summaries(record: dict[str, Any]) -> dict[int, dict[str, Any]]: ------- The summaries, empty if the iteration recorded none. """ - raw = record.get("io_summaries") - if not isinstance(raw, dict): - # Absent, or written by a version that shaped it differently. The file - # comes from disk, so an unexpected shape is skipped rather than raised. + if not isinstance(record, SuccessRecord) or record.io_summaries is None: return {} - return {int(rank): s for rank, s in raw.items()} + return {int(rank): s for rank, s in record.io_summaries.items()} def print_io_summaries(run: dict[str, Any]) -> None: @@ -149,13 +154,12 @@ def print_io_summaries(run: dict[str, Any]) -> None: A single run, as returned by :func:`load_runs`. """ rows = [ - (r["query"], r["iteration"], rank, s) + (r.query, r.iteration, rank, s) for r in iter_records(run) for rank, s in sorted(_io_summaries(r).items()) ] if not rows: - # Either the run predates I/O statistics, or it was made without - # `--rapidsmpf-statistics`, in which case no rank counts anything. + # The run was made without `--rapidsmpf-statistics`, so no rank counted. print("\nI/O: not collected (run with --rapidsmpf-statistics)") return @@ -260,4 +264,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: if __name__ == "__main__": - main(parse_args()) + try: + main(parse_args()) + except ValueError as e: + raise SystemExit(f"error: {e}") from None diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 91692db91462..84adff5d2bd7 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -468,6 +468,53 @@ def _infer_scale_factor(name: str, path: str | Path, suffix: str) -> int | float raise ValueError(f"Invalid benchmark script name: '{name}'.") +def record_from_dict(data: dict[str, Any]) -> SuccessRecord | FailedRecord: + """ + Read one iteration record back from its serialized form. + + Parameters + ---------- + data + One entry of a run's ``records``. + + Returns + ------- + The record, typed by its ``status``. + + Raises + ------ + ValueError + If the status is unrecognized, or the summaries predate rank keying. + """ + status = data["status"] + if status == "success": + io_summaries = data.get("io_summaries") + if io_summaries is not None and not isinstance(io_summaries, dict): + raise ValueError( + "An iteration's io_summaries is not keyed by rank. This results " + "file was written by an incompatible version of the benchmarks." + ) + validation = data.get("validation_result") + return SuccessRecord( + query=data["query"], + iteration=data["iteration"], + duration=data["duration"], + statistics=data.get("statistics"), + io_summaries=io_summaries, + traces=data.get("traces"), + validation_result=( + ValidationResult(**validation) if validation is not None else None + ), + ) + if status == "error": + return FailedRecord( + query=data["query"], + iteration=data["iteration"], + traceback=data["traceback"], + ) + raise ValueError(f"Unrecognized iteration status: {status!r}") + + @dataclasses.dataclass(kw_only=True) class RunConfig: """Benchmark run configuration for SPMD / Ray / DuckDB frontends.""" diff --git a/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py b/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py index bad5beb35aed..6151c3b2a5d9 100644 --- a/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py +++ b/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py @@ -16,6 +16,7 @@ main, parse_args, ) +from cudf_polars.streaming.benchmarks.utils import SuccessRecord if TYPE_CHECKING: from collections.abc import Callable @@ -46,6 +47,17 @@ def _record(**overrides: Any) -> dict[str, Any]: } +def _failed(**overrides: Any) -> dict[str, Any]: + """One failed iteration.""" + return { + "query": 1, + "iteration": 0, + "status": "error", + "traceback": "boom", + **overrides, + } + + def _run(*records: dict[str, Any], **overrides: Any) -> dict[str, Any]: """A run holding ``records``, grouped by query as the runners write them.""" by_query: dict[str, list[dict[str, Any]]] = {} @@ -98,7 +110,7 @@ def test_load_runs_rejects_an_empty_file(tmp_path: Path) -> None: def test_iter_records_orders_by_query() -> None: """Records come out grouped by query, whatever order the keys are in.""" run = _run(_record(query=3), _record(query=1), _record(query=1, iteration=1)) - assert [(r["query"], r["iteration"]) for r in iter_records(run)] == [ + assert [(r.query, r.iteration) for r in iter_records(run)] == [ (1, 0), (1, 1), (3, 0), @@ -109,16 +121,29 @@ def test_iter_records_orders_by_query() -> None: "raw, expected", [ # JSON object keys are strings, so they are cast back to ranks. - ({"0": None, "1": None}, [0, 1]), - # A run without statistics has nothing to report. + ({"0": {}, "1": {}}, [0, 1]), + # A run made without statistics has nothing to report. (None, []), - # The file comes from disk, so a surprise is skipped rather than raised. - ([{}], []), ], ) def test_io_summaries_keys(raw: Any, expected: list[int]) -> None: - """Rank keys are integers, and anything unexpected yields nothing.""" - assert sorted(_io_summaries({"io_summaries": raw})) == expected + """Rank keys are integers, and a run without statistics yields nothing.""" + record = SuccessRecord(query=1, iteration=0, duration=0.5, io_summaries=raw) + assert sorted(_io_summaries(record)) == expected + + +def test_an_incompatible_run_is_rejected(report: Callable[..., str]) -> None: + """Older files keyed I/O summaries differently, and are not read.""" + with pytest.raises(ValueError, match="incompatible version"): + report(_run(_record(io_summaries=[{}]))) + + +def test_an_unreadable_older_run_does_not_stop_the_report( + report: Callable[..., str], +) -> None: + """The default path shows the last run, so an older bad one is not touched.""" + out = report(_run(_record(io_summaries=[{}]), run_id="old"), _run(run_id="new")) + assert "new" in out def test_prints_timings_and_io(report: Callable[..., str]) -> None: @@ -141,7 +166,7 @@ def test_failed_iterations_are_left_out_of_the_timings( report: Callable[..., str], ) -> None: """A query whose every iteration failed has no min, max or mean to report.""" - assert "no successful iterations" in report(_run(_record(status="error"))) + assert "no successful iterations" in report(_run(_failed())) def test_a_rank_that_read_nothing_is_called_out(report: Callable[..., str]) -> None: diff --git a/python/cudf_polars/tests/streaming/test_statistics.py b/python/cudf_polars/tests/streaming/test_statistics.py index 9d3a4c3547ef..668f7fc0acb5 100644 --- a/python/cudf_polars/tests/streaming/test_statistics.py +++ b/python/cudf_polars/tests/streaming/test_statistics.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING +import kvikio import pytest import polars as pl @@ -83,16 +84,9 @@ def test_io_summary(engine: StreamingEngine, scan_query: pl.LazyFrame) -> None: assert sum(s.num_ops for s in summaries) > 0 assert sum(s.bytes_read for s in summaries) > 0 - for s in summaries: - # A summary is self-consistent, whether or not this rank did any I/O. - assert s.busy_ns <= s.wall_ns - assert s.bytes_read + s.bytes_written == s.bytes_transferred - assert sum(b["num_ops"] for b in s.by_backend.values()) == s.num_ops - assert s.num_reads + s.num_writes == s.num_ops - assert s.num_errors == 0 - # kvikio renders the report, so this is a check that it is reachable - # per rank rather than a check of its content. - assert "KvikIO I/O summary" in str(s) + # kvikio owns Summary and its invariants, so the only thing to check here is + # that each rank handed back one intact. + assert all(isinstance(s, kvikio.Summary) for s in summaries) def test_io_summary_clear_starts_a_new_span( From 98a396a36a1ff6cab5009294584e3b453dd15506 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 25 Aug 2026 09:46:23 +0200 Subject: [PATCH 5/9] kvikio_statistics --- docs/cudf/source/cudf_polars/profiling.md | 12 ++++---- python/cudf_polars/cudf_polars/engine/core.py | 23 +++++++------- python/cudf_polars/cudf_polars/engine/dask.py | 23 +++++++++++--- .../cudf_polars/cudf_polars/engine/options.py | 9 ++++++ python/cudf_polars/cudf_polars/engine/ray.py | 28 +++++++++++++---- python/cudf_polars/cudf_polars/engine/spmd.py | 17 ++++++++--- .../cudf_polars/streaming/benchmarks/utils.py | 10 +++---- .../cudf_polars/cudf_polars/utils/config.py | 14 +++++++++ .../tests/streaming/test_statistics.py | 30 ++++++++++--------- python/pylibcudf/pyproject.toml | 2 +- 10 files changed, 116 insertions(+), 52 deletions(-) diff --git a/docs/cudf/source/cudf_polars/profiling.md b/docs/cudf/source/cudf_polars/profiling.md index 317259b64737..35e0baa7ffe6 100644 --- a/docs/cudf/source/cudf_polars/profiling.md +++ b/docs/cudf/source/cudf_polars/profiling.md @@ -60,16 +60,16 @@ print(total) ## I/O Statistics -The same `statistics=True` also turns on [KvikIO I/O statistics][kvikio-stats] on every rank, -which report what storage did. `gather_io_summary()` returns one `kvikio.Summary` per rank, -keyed by rank index: +`kvikio_statistics=True` turns on [KvikIO I/O statistics][kvikio-stats] on every rank, which +report what storage did. It is separate from `statistics`, so you can collect either on its own. +`gather_io_summary()` returns one `kvikio.Summary` per rank, keyed by rank index: ```python import polars as pl from cudf_polars.engine.options import StreamingOptions from cudf_polars.engine.ray import RayEngine -opts = StreamingOptions(statistics=True) +opts = StreamingOptions(kvikio_statistics=True) with RayEngine.from_options(opts) as engine: pl.scan_parquet("/data/*.parquet").collect(engine=engine) @@ -81,8 +81,8 @@ with RayEngine.from_options(opts) as engine: `clear=True` restarts each rank's measured span after reading, scoping the next gather to whatever follows. A rank that is not counting is absent, so the result is empty unless -`statistics=True` is set. That is distinct from a zeroed summary, which means the rank was -counting and did no I/O. +`kvikio_statistics=True` is set. That is distinct from a zeroed summary, which means the rank +was counting and did no I/O. Printing a summary gives KvikIO's own report: diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index b55ddbdf136a..07fe1750d648 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -97,15 +97,14 @@ def reset_statistics_from_options( return statistics -def make_kvikio_monitor(options: Options) -> kvikio.SummaryMonitor | None: +def make_kvikio_monitor(*, enabled: bool) -> kvikio.SummaryMonitor | None: """ - Create a kvikio I/O monitor if statistics are enabled in ``options``. + Create a kvikio I/O monitor if ``enabled``. Parameters ---------- - options - Options providing the enabled setting. The same - ``RAPIDSMPF_STATISTICS`` knob that gates rapidsmpf statistics. + enabled + Whether to count, from the ``kvikio_statistics`` executor option. Returns ------- @@ -123,23 +122,23 @@ def make_kvikio_monitor(options: Options) -> kvikio.SummaryMonitor | None: thread's kvikio I/O in the script's process, so user code performing kvikio reads is counted too. It cannot attribute I/O to a particular query. """ - if not Statistics.from_options(options).enabled: + if not enabled: return None return kvikio.SummaryMonitor() -def reset_kvikio_monitor_from_options( - monitor: kvikio.SummaryMonitor | None, options: Options +def reset_kvikio_monitor( + monitor: kvikio.SummaryMonitor | None, *, enabled: bool ) -> kvikio.SummaryMonitor | None: """ - Bring a kvikio I/O monitor into line with ``options``. + Bring a kvikio I/O monitor into line with a new ``enabled`` setting. Parameters ---------- monitor The rank's existing monitor, if it has one. - options - Options providing the new enabled setting. + enabled + Whether to count, from the ``kvikio_statistics`` executor option. Returns ------- @@ -148,7 +147,7 @@ def reset_kvikio_monitor_from_options( None If they are not, in which case any existing monitor has been stopped. """ - if not Statistics.from_options(options).enabled: + if not enabled: if monitor is not None: monitor.stop() return None diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 5b8bacafa3d4..7e3d1fa0c20b 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -41,7 +41,7 @@ drop_if_replicated, evaluate_on_rank, make_kvikio_monitor, - reset_kvikio_monitor_from_options, + reset_kvikio_monitor, reset_statistics_from_options, resolve_rapidsmpf_options, take_io_summary, @@ -60,6 +60,7 @@ DaskContext, MemoryResourceConfig, resolve_kvikio_nthreads, + resolve_kvikio_statistics, ) if TYPE_CHECKING: @@ -337,6 +338,7 @@ def _setup_worker( engine_id: uuid.UUID, num_py_executors: int, kvikio_nthreads: int, + kvikio_statistics: bool, quent_context: cudf_polars.quent.QuentContext | None, dask_worker: distributed.Worker | None = None, ) -> None: @@ -374,6 +376,8 @@ def _setup_worker( Number of Python executors to use for this worker. kvikio_nthreads Number of kvikio threads to configure on this worker process. + kvikio_statistics + Whether to collect KvikIO I/O statistics on this worker. quent_context Quent context to use for this worker, if quent is enabled. @@ -441,7 +445,7 @@ def _setup_worker( quent_worker=quent_worker, quent_logger=quent_logger, statistics=statistics, - kvikio_monitor=make_kvikio_monitor(options), + kvikio_monitor=make_kvikio_monitor(enabled=kvikio_statistics), ) setattr(dask_worker, attr, mp_ctx) if mp_ctx.quent_logger is not None: @@ -503,6 +507,7 @@ def _reset_worker( *, uid: str, kvikio_nthreads: int, + kvikio_statistics: bool, dask_worker: distributed.Worker | None = None, ) -> None: """ @@ -519,6 +524,8 @@ def _reset_worker( Cluster instance identifier used to look up the per-worker context. kvikio_nthreads Number of kvikio threads to configure on this worker process. + kvikio_statistics + Whether to collect KvikIO I/O statistics on this worker. dask_worker Injected by ``distributed`` when called via :meth:`distributed.Client.run`. """ @@ -545,8 +552,8 @@ def _reset_worker( options = Options.deserialize(rapidsmpf_options_as_bytes) mp_ctx.statistics = reset_statistics_from_options(mp_ctx.statistics, options) mp_ctx.statistics.clear() - mp_ctx.kvikio_monitor = reset_kvikio_monitor_from_options( - mp_ctx.kvikio_monitor, options + mp_ctx.kvikio_monitor = reset_kvikio_monitor( + mp_ctx.kvikio_monitor, enabled=kvikio_statistics ) mp_ctx.ctx = Context.from_options( mp_ctx.comm.logger, mp_ctx.base_mr, options, mp_ctx.statistics @@ -944,6 +951,9 @@ def __init__( executor_options.setdefault( "kvikio_nthreads", resolve_kvikio_nthreads(executor_options) ) + executor_options.setdefault( + "kvikio_statistics", resolve_kvikio_statistics(executor_options) + ) engine_options = engine_options or {} quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( @@ -1059,6 +1069,7 @@ def __init__( quent_context=quent_context, num_py_executors=executor_options.get("num_py_executors", 8), kvikio_nthreads=executor_options["kvikio_nthreads"], + kvikio_statistics=executor_options["kvikio_statistics"], ) dask_ctx = DaskContext( @@ -1103,6 +1114,9 @@ def _reset( existing_kvikio_nthreads = existing_executor_options.get("kvikio_nthreads") if existing_kvikio_nthreads is not None: executor_options.setdefault("kvikio_nthreads", existing_kvikio_nthreads) + executor_options.setdefault( + "kvikio_statistics", resolve_kvikio_statistics(executor_options) + ) engine_options = engine_options or {} self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) @@ -1118,6 +1132,7 @@ def _reset( _reset_worker, uid=ctx.rapidsmpf_id, kvikio_nthreads=executor_options["kvikio_nthreads"], + kvikio_statistics=executor_options["kvikio_statistics"], ), rapidsmpf_options_as_bytes, ) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 3c2c513f0b1b..6b79b88f3db8 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -201,6 +201,12 @@ class StreamingOptions: Env: ``CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS``. Default: ``8``. Category: executor. + kvikio_statistics + Collect KvikIO I/O statistics, reachable through + :meth:`~cudf_polars.engine.core.StreamingEngine.gather_io_summary`. + Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS``. + Default: ``False``. + Category: executor. max_concurrent_io_tasks Maximum concurrent IO tasks for each scan node. Env: ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS``. @@ -329,6 +335,9 @@ class StreamingOptions: kvikio_nthreads: int | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", int ) + kvikio_statistics: bool | Unspecified = _opt( + "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS", parse_boolean + ) max_concurrent_io_tasks: int | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", int ) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 2862d7354280..424483999609 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -37,7 +37,7 @@ drop_if_replicated, evaluate_on_rank, make_kvikio_monitor, - reset_kvikio_monitor_from_options, + reset_kvikio_monitor, reset_statistics_from_options, resolve_rapidsmpf_options, take_io_summary, @@ -57,6 +57,7 @@ MemoryResourceConfig, RayContext, resolve_kvikio_nthreads, + resolve_kvikio_statistics, ) if TYPE_CHECKING: @@ -252,6 +253,7 @@ def __init__( rapidsmpf_options_as_bytes: bytes, num_py_executors: int, kvikio_nthreads: int, + kvikio_statistics: bool, hardware_binding: HardwareBindingPolicy, memory_resource_config: MemoryResourceConfig | None, worker_id: uuid.UUID, @@ -273,7 +275,7 @@ def __init__( rapidsmpf_options_as_bytes ) self._rapidsmpf_statistics = Statistics.from_options(self._rapidsmpf_options) - self._kvikio_monitor = make_kvikio_monitor(self._rapidsmpf_options) + self._kvikio_monitor = make_kvikio_monitor(enabled=kvikio_statistics) self._nranks: int = nranks self._py_executor = ThreadPoolExecutor( max_workers=num_py_executors, @@ -357,7 +359,13 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: self._mr = self._ctx.br().device_mr_adaptor() rmm.mr.set_current_device_resource(self._mr) - def reset(self, *, rapidsmpf_options_as_bytes: bytes, kvikio_nthreads: int) -> None: + def reset( + self, + *, + rapidsmpf_options_as_bytes: bytes, + kvikio_nthreads: int, + kvikio_statistics: bool, + ) -> None: """ Rebuild the streaming Context with new options. @@ -370,6 +378,8 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes, kvikio_nthreads: int) -> N Serialized :class:`Options` to install. kvikio_nthreads Number of kvikio threads to configure on this worker process. + kvikio_statistics + Whether to collect KvikIO I/O statistics on this rank. """ if self._ctx is None: raise RuntimeError("reset() requires setup_worker() to have run") @@ -387,8 +397,8 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes, kvikio_nthreads: int) -> N self._rapidsmpf_statistics, self._rapidsmpf_options ) self._rapidsmpf_statistics.clear() - self._kvikio_monitor = reset_kvikio_monitor_from_options( - self._kvikio_monitor, self._rapidsmpf_options + self._kvikio_monitor = reset_kvikio_monitor( + self._kvikio_monitor, enabled=kvikio_statistics ) assert self._base_mr is not None self._ctx = Context.from_options( @@ -779,6 +789,9 @@ def __init__( executor_options.setdefault( "kvikio_nthreads", resolve_kvikio_nthreads(executor_options) ) + executor_options.setdefault( + "kvikio_statistics", resolve_kvikio_statistics(executor_options) + ) engine_options = engine_options or {} ray_init_options = ray_init_options or {} @@ -859,6 +872,7 @@ def __init__( executor_options.get("num_py_executors", 8), ), kvikio_nthreads=executor_options["kvikio_nthreads"], + kvikio_statistics=executor_options["kvikio_statistics"], hardware_binding=hw_binding, memory_resource_config=mr_config, worker_id=worker_id, @@ -918,6 +932,9 @@ def _reset( existing_kvikio_nthreads = existing_executor_options.get("kvikio_nthreads") if existing_kvikio_nthreads is not None: executor_options.setdefault("kvikio_nthreads", existing_kvikio_nthreads) + executor_options.setdefault( + "kvikio_statistics", resolve_kvikio_statistics(executor_options) + ) engine_options = engine_options or {} self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() @@ -930,6 +947,7 @@ def _reset( rank.reset.remote( rapidsmpf_options_as_bytes=rapidsmpf_options_as_bytes, kvikio_nthreads=executor_options["kvikio_nthreads"], + kvikio_statistics=executor_options["kvikio_statistics"], ) for rank in self._rank_actors ] diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 9917799476d1..38d712f9bbbf 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -42,7 +42,7 @@ check_reserved_keys, evaluate_on_rank, make_kvikio_monitor, - reset_kvikio_monitor_from_options, + reset_kvikio_monitor, reset_statistics_from_options, resolve_rapidsmpf_options, take_io_summary, @@ -65,6 +65,7 @@ SPMDContext, StreamingExecutor, resolve_kvikio_nthreads, + resolve_kvikio_statistics, ) if TYPE_CHECKING: @@ -417,6 +418,9 @@ def __init__( executor_options.setdefault( "kvikio_nthreads", resolve_kvikio_nthreads(executor_options) ) + executor_options.setdefault( + "kvikio_statistics", resolve_kvikio_statistics(executor_options) + ) engine_options = engine_options or {} quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( @@ -467,7 +471,9 @@ def __init__( self._py_executor: ThreadPoolExecutor | None = None self._store_uid = uuid.uuid4().hex exit_stack = contextlib.ExitStack() - self._kvikio_monitor = make_kvikio_monitor(self.rapidsmpf_options) + self._kvikio_monitor = make_kvikio_monitor( + enabled=executor_options["kvikio_statistics"] + ) exit_stack.callback(self._stop_kvikio_monitor) # TODO: there's no reason our API needs a plain dict[str, Any] rather than @@ -623,6 +629,9 @@ def _reset( existing_kvikio_nthreads = existing_executor_options.get("kvikio_nthreads") if existing_kvikio_nthreads is not None: executor_options.setdefault("kvikio_nthreads", existing_kvikio_nthreads) + executor_options.setdefault( + "kvikio_statistics", resolve_kvikio_statistics(executor_options) + ) kvikio.defaults.set("num_threads", executor_options["kvikio_nthreads"]) engine_options = engine_options or {} quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( @@ -645,8 +654,8 @@ def _reset( self._comm.progress_thread.statistics, rapidsmpf_options ) statistics.clear() - self._kvikio_monitor = reset_kvikio_monitor_from_options( - self._kvikio_monitor, rapidsmpf_options + self._kvikio_monitor = reset_kvikio_monitor( + self._kvikio_monitor, enabled=executor_options["kvikio_statistics"] ) self._ctx = Context.from_options( diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 84adff5d2bd7..c49a145ffe4e 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -67,8 +67,6 @@ pynvml = None try: - from rapidsmpf.statistics import Statistics - import cudf_polars.dsl.tracing import cudf_polars.quent from cudf_polars.dsl.ir import IRExecutionContext @@ -995,13 +993,13 @@ def _collect_io_summaries( return None if not isinstance(engine, StreamingEngine): return None - if not Statistics.from_options(engine.rapidsmpf_options).enabled: - return None - # String keys, since the record is written as JSON. - return { + # String keys, since the record is written as JSON. Empty when no rank is + # counting, which the report reads as "not collected". + summaries = { str(rank): dataclasses.asdict(summary) for rank, summary in engine.gather_io_summary(clear=True).items() } + return summaries or None def run_polars_query_iteration( diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 315aa0c8f8cd..55b5863f99bc 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -198,6 +198,17 @@ def default_factory() -> T | DefaultT: return default_factory +def resolve_kvikio_statistics(executor_options: dict[str, Any]) -> bool: + """Resolve whether kvikio I/O statistics are collected, with env var fallback.""" + value = executor_options.get( + "kvikio_statistics", + os.environ.get("CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS", "False"), + ) + if isinstance(value, bool): + return value + return value.strip().lower() in {"1", "true", "yes", "y", "on"} + + def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: """Resolve kvikio thread count from executor options with env var fallback.""" return int( @@ -833,6 +844,9 @@ class StreamingExecutor: kvikio_nthreads: int = dataclasses.field( default_factory=lambda: resolve_kvikio_nthreads({}) ) + kvikio_statistics: bool = dataclasses.field( + default_factory=lambda: resolve_kvikio_statistics({}) + ) min_device_size: int | None = None spmd_context: SPMDContext | None = None diff --git a/python/cudf_polars/tests/streaming/test_statistics.py b/python/cudf_polars/tests/streaming/test_statistics.py index 668f7fc0acb5..02e90c8f9939 100644 --- a/python/cudf_polars/tests/streaming/test_statistics.py +++ b/python/cudf_polars/tests/streaming/test_statistics.py @@ -32,9 +32,11 @@ def engine( streaming_engine_factory: Callable[..., StreamingEngine], ) -> StreamingEngine: - """Yield each supported streaming engine with statistics enabled.""" + """Yield each supported streaming engine with both kinds of statistics on.""" return streaming_engine_factory( - StreamingOptions(statistics=True, max_rows_per_partition=10), + StreamingOptions( + statistics=True, kvikio_statistics=True, max_rows_per_partition=10 + ), ) @@ -105,30 +107,30 @@ def test_io_summary_clear_starts_a_new_span( assert sum(s.bytes_transferred for s in after.values()) == 0 -def test_io_summary_is_off_without_statistics( +def test_io_summary_is_independent_of_rapidsmpf_statistics( streaming_engine_factory: Callable[..., StreamingEngine], scan_query: pl.LazyFrame, ) -> None: - """No monitor is created when statistics are off. One is created again when they are turned back on.""" + """``kvikio_statistics`` gates I/O counting on its own, and can be reset.""" + # RapidsMPF statistics on, kvikio off. One does not imply the other. engine = streaming_engine_factory( - StreamingOptions(statistics=False, max_rows_per_partition=10) + StreamingOptions( + statistics=True, kvikio_statistics=False, max_rows_per_partition=10 + ) ) scan_query.collect(engine=engine) - + assert engine.gather_statistics()[0].enabled # An absent rank is "this rank was not counting", which is not the same as # a zeroed summary meaning "this rank did no I/O". assert engine.gather_io_summary() == {} - # A reset leaves `rapidsmpf_options` describing the reset, which is what - # callers read to decide whether anything is being counted. - assert not Statistics.from_options(engine.rapidsmpf_options).enabled - - # Turning statistics back on has to build a fresh monitor, the previous one - # having been stopped rather than paused. + # Turning it on has to build a fresh monitor, the previous one having been + # stopped rather than paused. engine = streaming_engine_factory( - StreamingOptions(statistics=True, max_rows_per_partition=10) + StreamingOptions( + statistics=False, kvikio_statistics=True, max_rows_per_partition=10 + ) ) - assert Statistics.from_options(engine.rapidsmpf_options).enabled scan_query.collect(engine=engine) summaries = engine.gather_io_summary() assert sorted(summaries) == list(range(engine.nranks)) diff --git a/python/pylibcudf/pyproject.toml b/python/pylibcudf/pyproject.toml index 3a6f0dfe2745..ce91c3211eaf 100644 --- a/python/pylibcudf/pyproject.toml +++ b/python/pylibcudf/pyproject.toml @@ -37,7 +37,7 @@ classifiers = [ [project.optional-dependencies] test = [ - "cupy-cuda13x[ctk]>=14.0.1,!=14.1.0", + "cupy-cuda13x>=14.0.1,!=14.1.0", "fastavro>=0.22.9", "mmh3", "nanoarrow", From edf4465feb49b125b85655a741e171c6681a971a Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 25 Aug 2026 10:08:17 +0200 Subject: [PATCH 6/9] Based on review by @pentschev --- .../cudf_polars/streaming/benchmarks/utils.py | 10 ++-------- .../streaming/benchmarks/test_print_results_file.py | 6 ------ python/cudf_polars/tests/streaming/test_statistics.py | 9 +++++---- 3 files changed, 7 insertions(+), 18 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index c49a145ffe4e..d9437dfce820 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -482,23 +482,17 @@ def record_from_dict(data: dict[str, Any]) -> SuccessRecord | FailedRecord: Raises ------ ValueError - If the status is unrecognized, or the summaries predate rank keying. + If the status is unrecognized. """ status = data["status"] if status == "success": - io_summaries = data.get("io_summaries") - if io_summaries is not None and not isinstance(io_summaries, dict): - raise ValueError( - "An iteration's io_summaries is not keyed by rank. This results " - "file was written by an incompatible version of the benchmarks." - ) validation = data.get("validation_result") return SuccessRecord( query=data["query"], iteration=data["iteration"], duration=data["duration"], statistics=data.get("statistics"), - io_summaries=io_summaries, + io_summaries=data.get("io_summaries"), traces=data.get("traces"), validation_result=( ValidationResult(**validation) if validation is not None else None diff --git a/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py b/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py index 6151c3b2a5d9..c967d1aefc45 100644 --- a/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py +++ b/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py @@ -132,12 +132,6 @@ def test_io_summaries_keys(raw: Any, expected: list[int]) -> None: assert sorted(_io_summaries(record)) == expected -def test_an_incompatible_run_is_rejected(report: Callable[..., str]) -> None: - """Older files keyed I/O summaries differently, and are not read.""" - with pytest.raises(ValueError, match="incompatible version"): - report(_run(_record(io_summaries=[{}]))) - - def test_an_unreadable_older_run_does_not_stop_the_report( report: Callable[..., str], ) -> None: diff --git a/python/cudf_polars/tests/streaming/test_statistics.py b/python/cudf_polars/tests/streaming/test_statistics.py index 02e90c8f9939..c7dfec71931d 100644 --- a/python/cudf_polars/tests/streaming/test_statistics.py +++ b/python/cudf_polars/tests/streaming/test_statistics.py @@ -111,7 +111,7 @@ def test_io_summary_is_independent_of_rapidsmpf_statistics( streaming_engine_factory: Callable[..., StreamingEngine], scan_query: pl.LazyFrame, ) -> None: - """``kvikio_statistics`` gates I/O counting on its own, and can be reset.""" + """``kvikio_statistics`` gates I/O counting on its own, and survives a reset.""" # RapidsMPF statistics on, kvikio off. One does not imply the other. engine = streaming_engine_factory( StreamingOptions( @@ -124,13 +124,14 @@ def test_io_summary_is_independent_of_rapidsmpf_statistics( # a zeroed summary meaning "this rank did no I/O". assert engine.gather_io_summary() == {} - # Turning it on has to build a fresh monitor, the previous one having been - # stopped rather than paused. - engine = streaming_engine_factory( + # The factory resets the one shared engine, so this exercises the monitor + # being created on an engine that was running without one. + reset = streaming_engine_factory( StreamingOptions( statistics=False, kvikio_statistics=True, max_rows_per_partition=10 ) ) + assert reset is engine scan_query.collect(engine=engine) summaries = engine.gather_io_summary() assert sorted(summaries) == list(range(engine.nranks)) From 4da1037a6afdaf06f95a45ca385fa78139d09149 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 25 Aug 2026 10:12:04 +0200 Subject: [PATCH 7/9] --kvikio-statistics --- python/cudf_polars/cudf_polars/engine/options.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 6b79b88f3db8..b519d2aba9d1 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -705,6 +705,16 @@ def _add_cli_args(parser: argparse.ArgumentParser) -> None: Env: CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS. Built-in default: 8."""), ) + g.add_argument( + "--kvikio-statistics", + dest="kvikio_statistics", + default=None, + action=argparse.BooleanOptionalAction, + help=textwrap.dedent("""\ + Collect KvikIO I/O statistics, reported per rank. + Env: CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS. + Built-in default: false."""), + ) g.add_argument( "--max-concurrent-io-tasks", dest="max_concurrent_io_tasks", From 41dff8234e834cb358498cb82a3b38fb2319d29b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 25 Aug 2026 10:16:05 +0200 Subject: [PATCH 8/9] style --- python/pylibcudf/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pylibcudf/pyproject.toml b/python/pylibcudf/pyproject.toml index ce91c3211eaf..3a6f0dfe2745 100644 --- a/python/pylibcudf/pyproject.toml +++ b/python/pylibcudf/pyproject.toml @@ -37,7 +37,7 @@ classifiers = [ [project.optional-dependencies] test = [ - "cupy-cuda13x>=14.0.1,!=14.1.0", + "cupy-cuda13x[ctk]>=14.0.1,!=14.1.0", "fastavro>=0.22.9", "mmh3", "nanoarrow", From 38ab334cb3a49ceebd5c9fbb51be6f06858c6990 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 25 Aug 2026 14:58:56 +0200 Subject: [PATCH 9/9] based on reviews --- python/cudf_polars/cudf_polars/utils/config.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 55b5863f99bc..38d9bef261d1 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -200,13 +200,12 @@ def default_factory() -> T | DefaultT: def resolve_kvikio_statistics(executor_options: dict[str, Any]) -> bool: """Resolve whether kvikio I/O statistics are collected, with env var fallback.""" - value = executor_options.get( - "kvikio_statistics", - os.environ.get("CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS", "False"), - ) - if isinstance(value, bool): - return value - return value.strip().lower() in {"1", "true", "yes", "y", "on"} + value = executor_options.get("kvikio_statistics") + if value is None: + value = os.environ.get("CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS") + if value is None: + return False + return value if isinstance(value, bool) else _bool_converter(value) def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: @@ -845,7 +844,9 @@ class StreamingExecutor: default_factory=lambda: resolve_kvikio_nthreads({}) ) kvikio_statistics: bool = dataclasses.field( - default_factory=lambda: resolve_kvikio_statistics({}) + default_factory=_make_default_factory( + f"{_env_prefix}__KVIKIO_STATISTICS", _bool_converter, default=False + ) ) min_device_size: int | None = None