Default cudf-polars to the kvikio MULTI_POLL backend

python/cudf_polars/cudf_polars/engine/core.py

@@ -286,12 +286,22 @@ class StreamingEngine(pl.GPUEngine):
     destruction and context manager exit must occur on the thread that created
     the instance.
 
-    Creating an engine sets the kvikio remote I/O backend to ``EASY_THREADPOOL``
-    and configures its thread pool (default 256 threads). Because kvikio's pool
-    is a global singleton, this blocks any concurrent kvikio IO in the process
-    until in-flight IO completes and overrides any prior ``kvikio.defaults.set(...)``
-    calls. Use the ``kvikio_nthreads`` executor option or the ``KVIKIO_NTHREADS``
-    environment variable to control the thread count.
+    Creating an engine sets the kvikio remote I/O backend to
+    ``kvikio.RemoteIOBackend.MULTI_POLL`` by default (see the
+    ``kvikio_remote_io_backend`` executor option), along with the
+    ``kvikio_task_size`` executor option (16 MiB under ``MULTI_POLL``, 64 MiB
+    under ``EASY_THREADPOOL``). Because kvikio's configuration is a global
+    singleton, this overrides any prior ``kvikio.defaults.set(...)`` calls made
+    in the process. When the backend is ``EASY_THREADPOOL``, engine creation
+    also configures kvikio's thread pool (default 256 threads), which blocks
+    any concurrent kvikio IO in the process until in-flight IO completes. Use
+    the ``kvikio_nthreads`` executor option or the ``KVIKIO_NTHREADS``
+    environment variable to control the thread count. Under ``MULTI_POLL``,
+    cudf-polars does not resolve a thread-pool size at all (kvikio itself may
+    still honor ``KVIKIO_NTHREADS`` via its own deferred default); remote I/O
+    concurrency is instead controlled by the ``kvikio_reactor_count``,
+    ``kvikio_reactor_dispatch``, and ``kvikio_request_ceiling`` executor
+    options.
 
     Parameters
     ----------

python/cudf_polars/cudf_polars/engine/dask.py

@@ -15,6 +15,7 @@
 
 import distributed
 import distributed.system
+import kvikio
 import pynvml
 import ucxx._lib.libucxx as ucx_api
 
@@ -58,8 +59,7 @@
     DaskContext,
     MemoryResourceConfig,
     configure_kvikio,
-    resolve_kvikio_nthreads,
-    resolve_kvikio_statistics,
+    resolve_kvikio_executor_options,
 )
 
 if TYPE_CHECKING:
@@ -338,8 +338,14 @@ def _setup_worker(
     worker_ids: list[uuid.UUID],
     engine_id: uuid.UUID,
     num_py_executors: int,
-    kvikio_nthreads: int,
+    kvikio_nthreads: int | None,
     kvikio_statistics: bool,
+    kvikio_remote_io_backend: kvikio.RemoteIOBackend,
+    kvikio_task_size: int,
+    kvikio_bounce_buffer_bytes: int,
+    kvikio_reactor_count: int,
+    kvikio_reactor_dispatch: kvikio.RemoteReactorDispatch,
+    kvikio_request_ceiling: int,
     quent_context: cudf_polars.quent.QuentContext | None,
     dask_worker: distributed.Worker | None = None,
 ) -> None:
@@ -376,9 +382,22 @@ def _setup_worker(
     num_py_executors
         Number of Python executors to use for this worker.
     kvikio_nthreads
-        Number of kvikio threads to configure on this worker process.
+        Number of kvikio threads to configure on this worker process. ``None``
+        defers to kvikio's own built-in default.
     kvikio_statistics
         Whether to collect KvikIO I/O statistics on this worker.
+    kvikio_remote_io_backend
+        The kvikio remote I/O backend to configure on this worker process.
+    kvikio_task_size
+        Size, in bytes, of the kvikio task size to configure on this worker process.
+    kvikio_bounce_buffer_bytes
+        Size, in bytes, of the kvikio bounce buffer to configure on this worker process.
+    kvikio_reactor_count
+        Number of ``MULTI_POLL`` reactor threads to configure on this worker process.
+    kvikio_reactor_dispatch
+        ``MULTI_POLL`` reactor dispatch policy to configure on this worker process.
+    kvikio_request_ceiling
+        ``MULTI_POLL`` concurrent-request ceiling to configure on this worker process.
     quent_context
         Quent context to use for this worker, if quent is enabled.
 
@@ -391,7 +410,15 @@ def _setup_worker(
     if mp_ctx is None:
         # Non-root worker: create communicator now.
         bind_to_gpu(hardware_binding)
-        configure_kvikio(kvikio_nthreads)
+        configure_kvikio(
+            kvikio_nthreads,
+            remote_io_backend=kvikio_remote_io_backend,
+            task_size=kvikio_task_size,
+            bounce_buffer_bytes=kvikio_bounce_buffer_bytes,
+            reactor_count=kvikio_reactor_count,
+            reactor_dispatch=kvikio_reactor_dispatch,
+            request_ceiling=kvikio_request_ceiling,
+        )
         memory_resource_config = (
             memory_resource_config or MemoryResourceConfig.default()
         )
@@ -412,7 +439,15 @@ def _setup_worker(
         base_mr = mp_ctx.base_mr
         comm = mp_ctx.comm
         statistics = mp_ctx.statistics
-        configure_kvikio(kvikio_nthreads)
+        configure_kvikio(
+            kvikio_nthreads,
+            remote_io_backend=kvikio_remote_io_backend,
+            task_size=kvikio_task_size,
+            bounce_buffer_bytes=kvikio_bounce_buffer_bytes,
+            reactor_count=kvikio_reactor_count,
+            reactor_dispatch=kvikio_reactor_dispatch,
+            request_ceiling=kvikio_request_ceiling,
+        )
 
     barrier(comm)
     worker_id = worker_ids[comm.rank]
@@ -508,8 +543,14 @@ def _reset_worker(
     rapidsmpf_options_as_bytes: bytes,
     *,
     uid: str,
-    kvikio_nthreads: int,
+    kvikio_nthreads: int | None,
     kvikio_statistics: bool,
+    kvikio_remote_io_backend: kvikio.RemoteIOBackend,
+    kvikio_task_size: int,
+    kvikio_bounce_buffer_bytes: int,
+    kvikio_reactor_count: int,
+    kvikio_reactor_dispatch: kvikio.RemoteReactorDispatch,
+    kvikio_request_ceiling: int,
     dask_worker: distributed.Worker | None = None,
 ) -> None:
     """
@@ -525,14 +566,35 @@ def _reset_worker(
     uid
         Cluster instance identifier used to look up the per-worker context.
     kvikio_nthreads
-        Number of kvikio threads to configure on this worker process.
+        Number of kvikio threads to configure on this worker process. ``None``
+        defers to kvikio's own built-in default.
     kvikio_statistics
         Whether to collect KvikIO I/O statistics on this worker.
+    kvikio_remote_io_backend
+        The kvikio remote I/O backend to configure on this worker process.
+    kvikio_task_size
+        Size, in bytes, of the kvikio task size to configure on this worker process.
+    kvikio_bounce_buffer_bytes
+        Size, in bytes, of the kvikio bounce buffer to configure on this worker process.
+    kvikio_reactor_count
+        Number of ``MULTI_POLL`` reactor threads to configure on this worker process.
+    kvikio_reactor_dispatch
+        ``MULTI_POLL`` reactor dispatch policy to configure on this worker process.
+    kvikio_request_ceiling
+        ``MULTI_POLL`` concurrent-request ceiling to configure on this worker process.
     dask_worker
         Injected by ``distributed`` when called via :meth:`distributed.Client.run`.
     """
     assert dask_worker is not None
-    configure_kvikio(kvikio_nthreads)
+    configure_kvikio(
+        kvikio_nthreads,
+        remote_io_backend=kvikio_remote_io_backend,
+        task_size=kvikio_task_size,
+        bounce_buffer_bytes=kvikio_bounce_buffer_bytes,
+        reactor_count=kvikio_reactor_count,
+        reactor_dispatch=kvikio_reactor_dispatch,
+        request_ceiling=kvikio_request_ceiling,
+    )
     attr = f"_cudf_polars_mp_context_{uid}"
     mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None)
     if mp_ctx is None:
@@ -951,13 +1013,7 @@ def __init__(
         executor_options: dict[str, Any] | None = None,
         engine_options: dict[str, Any] | None = None,
     ) -> None:
-        executor_options = executor_options or {}
-        executor_options.setdefault(
-            "kvikio_nthreads", resolve_kvikio_nthreads(executor_options)
-        )
-        executor_options.setdefault(
-            "kvikio_statistics", resolve_kvikio_statistics(executor_options)
-        )
+        executor_options = resolve_kvikio_executor_options(executor_options or {})
         engine_options = engine_options or {}
 
         quent_context: cudf_polars.quent.QuentContext | None = executor_options.get(
@@ -1074,6 +1130,12 @@ def __init__(
             num_py_executors=executor_options.get("num_py_executors", 8),
             kvikio_nthreads=executor_options["kvikio_nthreads"],
             kvikio_statistics=executor_options["kvikio_statistics"],
+            kvikio_remote_io_backend=executor_options["kvikio_remote_io_backend"],
+            kvikio_task_size=executor_options["kvikio_task_size"],
+            kvikio_bounce_buffer_bytes=executor_options["kvikio_bounce_buffer_bytes"],
+            kvikio_reactor_count=executor_options["kvikio_reactor_count"],
+            kvikio_reactor_dispatch=executor_options["kvikio_reactor_dispatch"],
+            kvikio_request_ceiling=executor_options["kvikio_request_ceiling"],
         )
 
         dask_ctx = DaskContext(
@@ -1111,16 +1173,16 @@ def _reset(
         )
         executor_options = executor_options or {}
         existing_executor_options = self.config.get("executor_options", {})
-        if isinstance(existing_executor_options, dict):
-            existing_quent_context = existing_executor_options.get("quent_context")
-            if existing_quent_context is not None:
-                executor_options.setdefault("quent_context", existing_quent_context)
-            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)
-        )
+        if not isinstance(existing_executor_options, dict):
+            existing_executor_options = {}
+        existing_quent_context = existing_executor_options.get("quent_context")
+        if existing_quent_context is not None:
+            executor_options.setdefault("quent_context", existing_quent_context)
+        if "kvikio_nthreads" in existing_executor_options:
+            executor_options.setdefault(
+                "kvikio_nthreads", existing_executor_options["kvikio_nthreads"]
+            )
+        executor_options = resolve_kvikio_executor_options(executor_options)
         engine_options = engine_options or {}
 
         self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)
@@ -1137,6 +1199,14 @@ def _reset(
                 uid=ctx.rapidsmpf_id,
                 kvikio_nthreads=executor_options["kvikio_nthreads"],
                 kvikio_statistics=executor_options["kvikio_statistics"],
+                kvikio_remote_io_backend=executor_options["kvikio_remote_io_backend"],
+                kvikio_task_size=executor_options["kvikio_task_size"],
+                kvikio_bounce_buffer_bytes=executor_options[
+                    "kvikio_bounce_buffer_bytes"
+                ],
+                kvikio_reactor_count=executor_options["kvikio_reactor_count"],
+                kvikio_reactor_dispatch=executor_options["kvikio_reactor_dispatch"],
+                kvikio_request_ceiling=executor_options["kvikio_request_ceiling"],
             ),
             rapidsmpf_options_as_bytes,
         )

python/cudf_polars/cudf_polars/engine/options.py

@@ -12,6 +12,8 @@
 import textwrap
 from typing import TYPE_CHECKING, Any, Literal
 
+import kvikio
+
 from rapidsmpf.config import Options
 from rapidsmpf.utils.string import parse_boolean
 
@@ -42,6 +44,14 @@
 ]
 
 
+def _parse_remote_io_backend(value: str) -> kvikio.RemoteIOBackend:
+    return kvikio.RemoteIOBackend[value.upper()]
+
+
+def _parse_reactor_dispatch(value: str) -> kvikio.RemoteReactorDispatch:
+    return kvikio.RemoteReactorDispatch[value.upper()]
+
+
 def _opt(
     category: str,
     env_var: str | None = None,
@@ -208,6 +218,42 @@ class StreamingOptions:
         Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS``.
         Default: ``False``.
         Category: executor.
+    kvikio_remote_io_backend
+        The kvikio remote I/O backend.
+        Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_REMOTE_IO_BACKEND``.
+        Default: ``kvikio.RemoteIOBackend.MULTI_POLL``.
+        Category: executor.
+    kvikio_task_size
+        Size, in bytes, of the chunks kvikio splits reads into for parallel
+        dispatch. Applies to local and remote I/O under both backends.
+        Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_TASK_SIZE``.
+        Default: ``16777216`` (16 MiB) for the ``MULTI_POLL`` backend,
+        ``67108864`` (64 MiB) for ``EASY_THREADPOOL``.
+        Category: executor.
+    kvikio_bounce_buffer_bytes
+        Size, in bytes, of the kvikio bounce buffer used to stage host memory
+        for device-memory transfers. Applies to local and remote I/O under
+        both backends (not specific to ``MULTI_POLL``).
+        Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_BOUNCE_BUFFER_BYTES``.
+        Default: ``16777216`` (16 MiB).
+        Category: executor.
+    kvikio_reactor_count
+        Number of reactor threads used by the ``MULTI_POLL`` remote I/O backend.
+        Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_REACTOR_COUNT``.
+        Default: ``24``.
+        Category: executor.
+    kvikio_reactor_dispatch
+        How sub-ranges of one read are distributed across reactor threads under
+        the ``MULTI_POLL`` remote I/O backend.
+        Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_REACTOR_DISPATCH``.
+        Default: ``kvikio.RemoteReactorDispatch.PER_CHUNK``.
+        Category: executor.
+    kvikio_request_ceiling
+        Maximum number of concurrent in-flight requests across all reactor
+        threads under the ``MULTI_POLL`` remote I/O backend. 0 means unlimited.
+        Env: ``CUDF_POLARS__EXECUTOR__KVIKIO_REQUEST_CEILING``.
+        Default: ``256``.
+        Category: executor.
     max_concurrent_io_tasks
         Maximum concurrent IO tasks for each scan node.
         Env: ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS``.
@@ -342,6 +388,28 @@ class StreamingOptions:
     kvikio_statistics: bool | Unspecified = _opt(
         "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_STATISTICS", parse_boolean
     )
+    kvikio_remote_io_backend: kvikio.RemoteIOBackend | Unspecified = _opt(
+        "executor",
+        "CUDF_POLARS__EXECUTOR__KVIKIO_REMOTE_IO_BACKEND",
+        _parse_remote_io_backend,
+    )
+    kvikio_task_size: int | Unspecified = _opt(
+        "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_TASK_SIZE", int
+    )
+    kvikio_bounce_buffer_bytes: int | Unspecified = _opt(
+        "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_BOUNCE_BUFFER_BYTES", int
+    )
+    kvikio_reactor_count: int | Unspecified = _opt(
+        "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_REACTOR_COUNT", int
+    )
+    kvikio_reactor_dispatch: kvikio.RemoteReactorDispatch | Unspecified = _opt(
+        "executor",
+        "CUDF_POLARS__EXECUTOR__KVIKIO_REACTOR_DISPATCH",
+        _parse_reactor_dispatch,
+    )
+    kvikio_request_ceiling: int | Unspecified = _opt(
+        "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_REQUEST_CEILING", int
+    )
     max_concurrent_io_tasks: int | dict[str, int] | Unspecified | None = _opt(
         "executor",
         "CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS",