Change kvikio default thread count in cudf-polars to 256

conda/recipes/cudf-polars/recipe.yaml

@@ -37,6 +37,7 @@ requirements:
     - python
     - pylibcudf =${{ version }}
     - rapidsmpf =${{ minor_version }}
+    - kvikio =${{ minor_version }}
     - cudf-streaming =${{ version }}
     - polars>=1.35,<1.43
     - packaging

dependencies.yaml

@@ -423,6 +423,7 @@ files:
       - depends_on_cuda_bindings
       - depends_on_rapidsmpf
       - depends_on_cudf_streaming
+      - depends_on_kvikio
       - depends_on_pylibcudf
   # TODO(26.10): Remove this alias in favor of dask
   py_run_cudf_polars_experimental:
@@ -1547,6 +1548,31 @@ dependencies:
               use_cuda_wheels: "true"
             packages:
               - cupy-cuda13x[ctk]>=14.0.1,!=14.1.0
+  depends_on_kvikio:
+    common:
+      - output_types: conda
+        packages:
+          - &kvikio_unsuffixed kvikio==26.10.*,>=0.0.0a0
+      - output_types: requirements
+        packages:
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple/
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix:
+              cuda: "12.*"
+              cuda_suffixed: "true"
+            packages:
+              - kvikio-cu12==26.10.*,>=0.0.0a0
+          - matrix:
+              cuda: "13.*"
+              cuda_suffixed: "true"
+            packages:
+              - kvikio-cu13==26.10.*,>=0.0.0a0
+          - matrix:
+              cuda_suffixed: "false"
+            packages:
+              - *kvikio_unsuffixed
   depends_on_libkvikio:
     common:
       - output_types: conda

python/cudf_polars/cudf_polars/dsl/utils/io.py

@@ -9,6 +9,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+import kvikio
+
 import pylibcudf as plc
 
 from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars
@@ -72,15 +74,8 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI
     # For now, we'll just use kvikio to explicitly get the size.
     sizes: list[int | None] = []
 
-    try:  # pragma: no cover; kvikio is optional
-        import kvikio
-    except ImportError:
-        kvikio = None
-
     for path in paths:
-        if (
-            paths and kvikio is not None and plc.io.SourceInfo._is_remote_uri(path)
-        ):  # pragma: no cover; kvikio is optional
+        if paths and plc.io.SourceInfo._is_remote_uri(path):
             # We're OK to use `kvikio.RemoteFile.open` here. It does make an HTTP HEAD
             # request for S3/HTTP endpoints, but that's the entire reason we're running
             # this code. So long as it makes just *one* HTTP request, there's no advantage

python/cudf_polars/cudf_polars/engine/core.py

@@ -185,6 +185,13 @@ class StreamingEngine(pl.GPUEngine):
     destruction and context manager exit must occur on the thread that created
     the instance.
 
+    Creating an engine configures the process-wide kvikio 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("num_threads", ...)`` call. Use the
+    ``kvikio_nthreads`` executor option or the ``KVIKIO_NTHREADS`` environment
+    variable to control the thread count.
+
     Parameters
     ----------
     nranks

python/cudf_polars/cudf_polars/engine/dask.py

@@ -15,6 +15,7 @@
 
 import distributed
 import distributed.system
+import kvikio.defaults
 import pynvml
 import ucxx._lib.libucxx as ucx_api
 
@@ -51,7 +52,11 @@
 )
 from cudf_polars.quent._context import LocalQuentContext
 from cudf_polars.unstable import unstable
-from cudf_polars.utils.config import DaskContext, MemoryResourceConfig
+from cudf_polars.utils.config import (
+    DaskContext,
+    MemoryResourceConfig,
+    resolve_kvikio_nthreads,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -326,6 +331,7 @@ def _setup_worker(
     worker_ids: list[uuid.UUID],
     engine_id: uuid.UUID,
     num_py_executors: int,
+    kvikio_nthreads: int,
     quent_context: cudf_polars.quent.QuentContext | None,
     dask_worker: distributed.Worker | None = None,
 ) -> None:
@@ -361,11 +367,14 @@ def _setup_worker(
         Injected by ``distributed`` when called via :meth:`distributed.Client.run`.
     num_py_executors
         Number of Python executors to use for this worker.
+    kvikio_nthreads
+        Number of kvikio threads to configure on this worker process.
     quent_context
         Quent context to use for this worker, if quent is enabled.
 
     """
     assert dask_worker is not None
+    kvikio.defaults.set("num_threads", kvikio_nthreads)
     options = Options.deserialize(rapidsmpf_options_as_bytes)
     attr = f"_cudf_polars_mp_context_{uid}"
     mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None)
@@ -482,6 +491,7 @@ def _reset_worker(
     rapidsmpf_options_as_bytes: bytes,
     *,
     uid: str,
+    kvikio_nthreads: int,
     dask_worker: distributed.Worker | None = None,
 ) -> None:
     """
@@ -496,10 +506,13 @@ def _reset_worker(
         Serialized :class:`Options` to install.
     uid
         Cluster instance identifier used to look up the per-worker context.
+    kvikio_nthreads
+        Number of kvikio threads to configure on this worker process.
     dask_worker
         Injected by ``distributed`` when called via :meth:`distributed.Client.run`.
     """
     assert dask_worker is not None
+    kvikio.defaults.set("num_threads", kvikio_nthreads)
     attr = f"_cudf_polars_mp_context_{uid}"
     mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None)
     if mp_ctx is None:
@@ -830,6 +843,9 @@ def __init__(
         engine_options: dict[str, Any] | None = None,
     ) -> None:
         executor_options = executor_options or {}
+        executor_options.setdefault(
+            "kvikio_nthreads", resolve_kvikio_nthreads(executor_options)
+        )
         engine_options = engine_options or {}
 
         quent_context: cudf_polars.quent.QuentContext | None = executor_options.get(
@@ -944,6 +960,7 @@ def __init__(
             rapidsmpf_options_as_bytes,
             quent_context=quent_context,
             num_py_executors=executor_options.get("num_py_executors", 8),
+            kvikio_nthreads=executor_options["kvikio_nthreads"],
         )
 
         dask_ctx = DaskContext(
@@ -985,6 +1002,9 @@ def _reset(
             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)
         engine_options = engine_options or {}
 
         rapidsmpf_options_as_bytes = resolve_rapidsmpf_options(
@@ -997,7 +1017,11 @@ def _reset(
         # inside :func:`_reset_worker` synchronizes the teardown across
         # workers.
         ctx.client.run(
-            functools.partial(_reset_worker, uid=ctx.rapidsmpf_id),
+            functools.partial(
+                _reset_worker,
+                uid=ctx.rapidsmpf_id,
+                kvikio_nthreads=executor_options["kvikio_nthreads"],
+            ),
             rapidsmpf_options_as_bytes,
         )
 

python/cudf_polars/cudf_polars/engine/options.py

@@ -345,6 +345,9 @@ class StreamingOptions:
     num_py_executors: int | Unspecified = _opt(
         "executor", "CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS", int
     )
+    kvikio_nthreads: int | Unspecified = _opt(
+        "executor", "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", int
+    )
     max_concurrent_io_tasks: int | Unspecified = _opt(
         "executor", "CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", int
     )

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.defaults
 import ray
 import ray.exceptions
 import ucxx._lib.libucxx as ucx_api
@@ -48,7 +49,11 @@
 from cudf_polars.quent._context import LocalQuentContext
 from cudf_polars.quent._types import Worker
 from cudf_polars.unstable import unstable
-from cudf_polars.utils.config import MemoryResourceConfig, RayContext
+from cudf_polars.utils.config import (
+    MemoryResourceConfig,
+    RayContext,
+    resolve_kvikio_nthreads,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -239,13 +244,15 @@ def __init__(
         nranks: int,
         rapidsmpf_options_as_bytes: bytes,
         num_py_executors: int,
+        kvikio_nthreads: int,
         hardware_binding: HardwareBindingPolicy,
         memory_resource_config: MemoryResourceConfig | None,
         worker_id: uuid.UUID,
         engine: cudf_polars.quent.Engine,
         quent_enabled: bool,
     ) -> None:
         bind_to_gpu(hardware_binding)
+        kvikio.defaults.set("num_threads", kvikio_nthreads)
         memory_resource_config = (
             memory_resource_config or MemoryResourceConfig.default()
         )
@@ -342,7 +349,7 @@ 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) -> None:
+    def reset(self, *, rapidsmpf_options_as_bytes: bytes, kvikio_nthreads: int) -> None:
         """
         Rebuild the streaming Context with new options.
 
@@ -353,9 +360,12 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes) -> None:
         ----------
         rapidsmpf_options_as_bytes
             Serialized :class:`Options` to install.
+        kvikio_nthreads
+            Number of kvikio threads to configure on this worker process.
         """
         if self._ctx is None:
             raise RuntimeError("reset() requires setup_worker() to have run")
+        kvikio.defaults.set("num_threads", kvikio_nthreads)
         assert self._comm is not None
         # Collective: all ranks idle before any rank tears down its Context.
         if self._comm.nranks > 1:
@@ -716,6 +726,9 @@ def __init__(
         num_ranks: int | None = None,
     ) -> None:
         executor_options = executor_options or {}
+        executor_options.setdefault(
+            "kvikio_nthreads", resolve_kvikio_nthreads(executor_options)
+        )
         engine_options = engine_options or {}
         ray_init_options = ray_init_options or {}
 
@@ -795,6 +808,7 @@ def __init__(
                         "int",
                         executor_options.get("num_py_executors", 8),
                     ),
+                    kvikio_nthreads=executor_options["kvikio_nthreads"],
                     hardware_binding=hw_binding,
                     memory_resource_config=mr_config,
                     worker_id=worker_id,
@@ -851,6 +865,9 @@ def _reset(
             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)
         engine_options = engine_options or {}
         rapidsmpf_options_as_bytes = resolve_rapidsmpf_options(
             rapidsmpf_options
@@ -863,6 +880,7 @@ def _reset(
             [
                 rank.reset.remote(
                     rapidsmpf_options_as_bytes=rapidsmpf_options_as_bytes,
+                    kvikio_nthreads=executor_options["kvikio_nthreads"],
                 )
                 for rank in self._rank_actors
             ]

python/cudf_polars/cudf_polars/engine/spmd.py

@@ -11,6 +11,8 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any, cast
 
+import kvikio
+
 import pylibcudf as plc
 import rmm.mr
 from cudf_streaming.partition_utils import (
@@ -58,6 +60,7 @@
     MemoryResourceConfig,
     SPMDContext,
     StreamingExecutor,
+    resolve_kvikio_nthreads,
 )
 
 if TYPE_CHECKING:
@@ -407,7 +410,11 @@ def __init__(
         engine_options: dict[str, Any] | None = None,
     ) -> None:
         executor_options = executor_options or {}
+        executor_options.setdefault(
+            "kvikio_nthreads", resolve_kvikio_nthreads(executor_options)
+        )
         engine_options = engine_options or {}
+
         quent_context: cudf_polars.quent.QuentContext | None = executor_options.get(
             "quent_context"
         )
@@ -423,6 +430,8 @@ def __init__(
         )
         bind_to_gpu(hw_binding)
 
+        kvikio.defaults.set("num_threads", executor_options["kvikio_nthreads"])
+
         self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)
         mr_config: MemoryResourceConfig = engine_options.get(
             "memory_resource_config", MemoryResourceConfig.default()
@@ -599,6 +608,10 @@ def _reset(
             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)
+        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"