Report kvikio I/O statistics per rank in the streaming engines

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(\..*)?"),

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

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:
+
+```text
+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

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.