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..35e0baa7ffe6 100644 --- a/docs/cudf/source/cudf_polars/profiling.md +++ b/docs/cudf/source/cudf_polars/profiling.md @@ -58,6 +58,72 @@ print(total) ``` +## I/O Statistics + +`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(kvikio_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 +`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: + +```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, and [busy time and bandwidth][kvikio-busy] +for how the busy figures are measured. + +### 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. + + ## GPU Profiling For streaming queries, we recommend profiling with [NVIDIA NSight Systems][nsight]. `cudf-polars` @@ -184,6 +250,8 @@ 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/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 3a58ef3ac1c3..07fe1750d648 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,106 @@ def reset_statistics_from_options( return statistics +def make_kvikio_monitor(*, enabled: bool) -> kvikio.SummaryMonitor | None: + """ + Create a kvikio I/O monitor if ``enabled``. + + Parameters + ---------- + enabled + Whether to count, from the ``kvikio_statistics`` executor option. + + 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 enabled: + return None + return kvikio.SummaryMonitor() + + +def reset_kvikio_monitor( + monitor: kvikio.SummaryMonitor | None, *, enabled: bool +) -> kvikio.SummaryMonitor | None: + """ + Bring a kvikio I/O monitor into line with a new ``enabled`` setting. + + Parameters + ---------- + monitor + The rank's existing monitor, if it has one. + enabled + Whether to count, from the ``kvikio_statistics`` executor option. + + 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 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 +419,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..7e3d1fa0c20b 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, reset_statistics_from_options, resolve_rapidsmpf_options, + take_io_summary, ) from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, @@ -56,6 +60,7 @@ DaskContext, MemoryResourceConfig, resolve_kvikio_nthreads, + resolve_kvikio_statistics, ) if TYPE_CHECKING: @@ -136,6 +141,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( @@ -332,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: @@ -369,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. @@ -436,6 +445,7 @@ def _setup_worker( quent_worker=quent_worker, quent_logger=quent_logger, statistics=statistics, + kvikio_monitor=make_kvikio_monitor(enabled=kvikio_statistics), ) setattr(dask_worker, attr, mp_ctx) if mp_ctx.quent_logger is not None: @@ -463,6 +473,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() @@ -492,6 +507,7 @@ def _reset_worker( *, uid: str, kvikio_nthreads: int, + kvikio_statistics: bool, dask_worker: distributed.Worker | None = None, ) -> None: """ @@ -508,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`. """ @@ -534,6 +552,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( + 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 ) @@ -631,6 +652,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], @@ -902,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( @@ -1017,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( @@ -1061,11 +1114,13 @@ 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 {} - 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 @@ -1077,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, ) @@ -1170,6 +1226,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/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 3c2c513f0b1b..b519d2aba9d1 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 ) @@ -696,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", diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 38b088756926..424483999609 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, reset_statistics_from_options, resolve_rapidsmpf_options, + take_io_summary, ) from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, @@ -53,6 +57,7 @@ MemoryResourceConfig, RayContext, resolve_kvikio_nthreads, + resolve_kvikio_statistics, ) if TYPE_CHECKING: @@ -248,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, @@ -269,6 +275,7 @@ def __init__( rapidsmpf_options_as_bytes ) self._rapidsmpf_statistics = Statistics.from_options(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, @@ -352,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. @@ -365,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") @@ -382,6 +397,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( + self._kvikio_monitor, enabled=kvikio_statistics + ) assert self._base_mr is not None self._ctx = Context.from_options( self._comm.logger, @@ -410,6 +428,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 @@ -465,6 +487,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, @@ -749,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 {} @@ -829,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, @@ -888,10 +932,12 @@ 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 {} - 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 @@ -901,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 ] @@ -1030,6 +1077,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..38d712f9bbbf 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, reset_statistics_from_options, resolve_rapidsmpf_options, + take_io_summary, ) from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, @@ -61,6 +65,7 @@ SPMDContext, StreamingExecutor, resolve_kvikio_nthreads, + resolve_kvikio_statistics, ) if TYPE_CHECKING: @@ -413,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( @@ -463,6 +471,10 @@ def __init__( self._py_executor: ThreadPoolExecutor | None = None self._store_uid = uuid.uuid4().hex exit_stack = contextlib.ExitStack() + 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 # a typed config object here. @@ -534,6 +546,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. @@ -611,12 +629,16 @@ 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( "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: @@ -632,6 +654,9 @@ def _reset( self._comm.progress_thread.statistics, rapidsmpf_options ) statistics.clear() + self._kvikio_monitor = reset_kvikio_monitor( + self._kvikio_monitor, enabled=executor_options["kvikio_statistics"] + ) self._ctx = Context.from_options( self._comm.logger, self._base_mr, rapidsmpf_options, statistics @@ -789,6 +814,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..571b0cb491c7 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/print_results_file.py @@ -0,0 +1,270 @@ +# 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 + +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." +) + +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, without interpreting it. + + Parameters + ---------- + path + Path to a file written by a benchmark runner's ``--output``. + + Returns + ------- + One decoded line per run, in the order they were appended. + + Raises + ------ + ValueError + If the file holds 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[SuccessRecord | FailedRecord]: + """ + 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["records"].items(), key=lambda kv: int(kv[0])): + yield from map(record_from_dict, 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["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 + 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: SuccessRecord | FailedRecord, +) -> 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. + """ + if not isinstance(record, SuccessRecord) or record.io_summaries is None: + return {} + return {int(rank): s for rank, s in record.io_summaries.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: + # The run was made without `--rapidsmpf-statistics`, so no rank counted. + 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__": + 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 d64949b2e4e7..d9437dfce820 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -252,6 +252,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 +264,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 +273,7 @@ def new( iteration=iteration, duration=duration, statistics=statistics, + io_summaries=io_summaries, traces=traces, ) @@ -463,6 +466,47 @@ 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. + """ + status = data["status"] + if status == "success": + validation = data.get("validation_result") + return SuccessRecord( + query=data["query"], + iteration=data["iteration"], + duration=data["duration"], + statistics=data.get("statistics"), + io_summaries=data.get("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.""" @@ -935,6 +979,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 + # 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( q_id: int, iteration: int, @@ -960,6 +1021,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 +1054,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/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/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..c967d1aefc45 --- /dev/null +++ b/python/cudf_polars/tests/streaming/benchmarks/test_print_results_file.py @@ -0,0 +1,220 @@ +# 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, +) +from cudf_polars.streaming.benchmarks.utils import SuccessRecord + +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 _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]]] = {} + 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": {}, "1": {}}, [0, 1]), + # A run made without statistics has nothing to report. + (None, []), + ], +) +def test_io_summaries_keys(raw: Any, expected: list[int]) -> None: + """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_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: + """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(_failed())) + + +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 65cf2c398a1c..c7dfec71931d 100644 --- a/python/cudf_polars/tests/streaming/test_statistics.py +++ b/python/cudf_polars/tests/streaming/test_statistics.py @@ -1,19 +1,23 @@ # 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 from typing import TYPE_CHECKING +import kvikio 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 @@ -28,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 + ), ) @@ -58,3 +64,75 @@ 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 + + # 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( + 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_independent_of_rapidsmpf_statistics( + streaming_engine_factory: Callable[..., StreamingEngine], + scan_query: pl.LazyFrame, +) -> None: + """``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( + 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() == {} + + # 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)) + assert sum(s.num_ops for s in summaries.values()) > 0