From 9cf60a172702186f933f39286c817ee5d924d5ae Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 12 Aug 2026 17:58:16 +0000 Subject: [PATCH 1/8] Change kvikio default thread count in cudf-polars to 256 --- dependencies.yaml | 26 +++++++++++ .../cudf_polars/cudf_polars/dsl/utils/io.py | 11 ++--- .../cudf_polars/cudf_polars/engine/options.py | 3 ++ python/cudf_polars/cudf_polars/engine/spmd.py | 15 +++++++ .../cudf_polars/cudf_polars/utils/config.py | 17 ++++++++ python/cudf_polars/pyproject.toml | 1 + .../tests/streaming/test_options.py | 5 +++ python/cudf_polars/tests/test_config.py | 43 +++++++++++++++++++ 8 files changed, 113 insertions(+), 8 deletions(-) diff --git a/dependencies.yaml b/dependencies.yaml index d14a559a977c..be600b713b1d 100644 --- a/dependencies.yaml +++ b/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 diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 75373afeed96..e89825357803 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/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 diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 611129cf2c01..6ab79178af2b 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/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 ) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index bc1fc56fe8e7..fe941a38f62d 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -7,10 +7,13 @@ import contextlib import dataclasses import json +import os import uuid 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 ( @@ -408,6 +411,18 @@ def __init__( ) -> None: executor_options = executor_options or {} engine_options = engine_options or {} + + kvikio_nthreads = int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ) + kvikio.defaults.set("num_threads", kvikio_nthreads) + quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( "quent_context" ) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 185f43b7142f..3513a464bfcd 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -685,6 +685,13 @@ class StreamingExecutor: num_py_executors Maximum number of workers for the Python ThreadPoolExecutor. Default is 8. + kvikio_nthreads + Number of threads in the kvikio thread pool. Defaults to 256. + This can be set via + + - ``executor_options`` passed to ``polars.GPUEngine`` + - the ``CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS`` environment variable + - the ``KVIKIO_NTHREADS`` environment variable (lower precedence) quent_context Quent tracing context. When ``None`` (default), Quent tracing is disabled. Pass a :class:`~cudf_polars.quent.QuentContext` instance to enable tracing. @@ -756,6 +763,14 @@ class StreamingExecutor: f"{_env_prefix}__NUM_PY_EXECUTORS", int, default=8 ) ) + kvikio_nthreads: int = dataclasses.field( + default_factory=lambda: int( + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ) + ) + ) min_device_size: int | None = None spmd_context: SPMDContext | None = None @@ -840,6 +855,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_concurrent_io_tasks must be an int") if not isinstance(self.num_py_executors, int): raise TypeError("num_py_executors must be an int") + if not isinstance(self.kvikio_nthreads, int): + raise TypeError("kvikio_nthreads must be an int") def __hash__(self) -> int: # noqa: D105 # dynamic_planning factory, a dataclass, isn't natively hashable. We'll dump it diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml index 44825a801a42..e8233488d0c3 100644 --- a/python/cudf_polars/pyproject.toml +++ b/python/cudf_polars/pyproject.toml @@ -23,6 +23,7 @@ requires-python = ">=3.11" dependencies = [ "cuda-bindings>=13.0.1,<14.0", "cudf-streaming==26.10.*,>=0.0.0a0", + "kvikio==26.10.*,>=0.0.0a0", "nvidia-ml-py>=12", "packaging", "polars>=1.35,<1.43", diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py index d422c92d5a65..920478bceea2 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -78,6 +78,11 @@ def test_executor_options_max_concurrent_io_tasks() -> None: assert result["max_concurrent_io_tasks"] == 6 +def test_executor_options_kvikio_nthreads() -> None: + result = StreamingOptions(kvikio_nthreads=128).to_executor_options() + assert result["kvikio_nthreads"] == 128 + + @pytest.mark.parametrize("value", [True, False]) def test_executor_options_sink_to_directory(*, value: bool) -> None: result = StreamingOptions(sink_to_directory=value).to_executor_options() diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index dc726970101a..ae1e9f9c0c72 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -801,6 +801,49 @@ def test_num_py_executors_from_env( assert config.executor.num_py_executors == 8 +def test_kvikio_nthreads_default() -> None: + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.kvikio_nthreads == 256 + + +def test_kvikio_nthreads_from_executor_options() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"kvikio_nthreads": 128}, + ) + ) + assert config.executor.kvikio_nthreads == 128 + + +def test_kvikio_nthreads_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with monkeypatch.context() as m: + m.setenv("CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", "64") + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.kvikio_nthreads == 64 + + +def test_kvikio_nthreads_from_kvikio_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with monkeypatch.context() as m: + m.setenv("KVIKIO_NTHREADS", "32") + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.kvikio_nthreads == 32 + + +def test_kvikio_nthreads_cudf_polars_env_takes_precedence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with monkeypatch.context() as m: + m.setenv("CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", "64") + m.setenv("KVIKIO_NTHREADS", "32") + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.kvikio_nthreads == 64 + + def test_dask_sink_to_directory_false_raises() -> None: with pytest.raises( ValueError, match="The dask cluster requires sink_to_directory=True" From ac1d8dd5abd868ce598f05ceab292716fc72385e Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 12 Aug 2026 18:36:42 +0000 Subject: [PATCH 2/8] address reviews --- conda/recipes/cudf-polars/recipe.yaml | 1 + python/cudf_polars/cudf_polars/engine/spmd.py | 36 +++++++++++++------ .../cudf_polars/cudf_polars/utils/config.py | 2 ++ python/cudf_polars/tests/test_config.py | 10 ++++-- 4 files changed, 35 insertions(+), 14 deletions(-) diff --git a/conda/recipes/cudf-polars/recipe.yaml b/conda/recipes/cudf-polars/recipe.yaml index 7bcefc99c3fc..a59fedabd19a 100644 --- a/conda/recipes/cudf-polars/recipe.yaml +++ b/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 diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index fe941a38f62d..19726d8dc8ab 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -412,17 +412,6 @@ def __init__( executor_options = executor_options or {} engine_options = engine_options or {} - kvikio_nthreads = int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ) - kvikio.defaults.set("num_threads", kvikio_nthreads) - quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( "quent_context" ) @@ -438,6 +427,19 @@ def __init__( ) bind_to_gpu(hw_binding) + kvikio_nthreads = int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ) + if kvikio_nthreads <= 0: + raise ValueError(f"kvikio_nthreads must be positive, got {kvikio_nthreads}") + kvikio.defaults.set("num_threads", kvikio_nthreads) + self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) mr_config: MemoryResourceConfig = engine_options.get( "memory_resource_config", MemoryResourceConfig.default() @@ -609,6 +611,18 @@ def _reset( engine_options=engine_options, ) executor_options = executor_options or {} + kvikio_nthreads = int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ) + if kvikio_nthreads <= 0: + raise ValueError(f"kvikio_nthreads must be positive, got {kvikio_nthreads}") + kvikio.defaults.set("num_threads", kvikio_nthreads) existing_executor_options = self.config.get("executor_options", {}) if isinstance(existing_executor_options, dict): existing_quent_context = existing_executor_options.get("quent_context") diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 3513a464bfcd..0ee980f67283 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -857,6 +857,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("num_py_executors must be an int") if not isinstance(self.kvikio_nthreads, int): raise TypeError("kvikio_nthreads must be an int") + if self.kvikio_nthreads <= 0: + raise ValueError("kvikio_nthreads must be positive") def __hash__(self) -> int: # noqa: D105 # dynamic_planning factory, a dataclass, isn't natively hashable. We'll dump it diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index ae1e9f9c0c72..83d7a5bc4819 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -801,9 +801,12 @@ def test_num_py_executors_from_env( assert config.executor.num_py_executors == 8 -def test_kvikio_nthreads_default() -> None: - config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) - assert config.executor.kvikio_nthreads == 256 +def test_kvikio_nthreads_default(monkeypatch: pytest.MonkeyPatch) -> None: + with monkeypatch.context() as m: + m.delenv("CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", raising=False) + m.delenv("KVIKIO_NTHREADS", raising=False) + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert config.executor.kvikio_nthreads == 256 def test_kvikio_nthreads_from_executor_options() -> None: @@ -829,6 +832,7 @@ def test_kvikio_nthreads_from_kvikio_env( monkeypatch: pytest.MonkeyPatch, ) -> None: with monkeypatch.context() as m: + m.delenv("CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", raising=False) m.setenv("KVIKIO_NTHREADS", "32") config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) assert config.executor.kvikio_nthreads == 32 From 9712cf47322096c886d8fb1af3249537751eebce Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 12 Aug 2026 19:19:41 +0000 Subject: [PATCH 3/8] plumb through for dask and ray engines --- python/cudf_polars/cudf_polars/engine/dask.py | 32 ++++++++++++++++++- python/cudf_polars/cudf_polars/engine/ray.py | 26 ++++++++++++++- 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 20b0139650b0..09eebc2e2a36 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.defaults import pynvml import ucxx._lib.libucxx as ucx_api @@ -326,6 +327,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 +363,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 +487,7 @@ def _reset_worker( rapidsmpf_options_as_bytes: bytes, *, uid: str, + kvikio_nthreads: int, dask_worker: distributed.Worker | None = None, ) -> None: """ @@ -496,10 +502,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: @@ -944,6 +953,15 @@ def __init__( rapidsmpf_options_as_bytes, quent_context=quent_context, num_py_executors=executor_options.get("num_py_executors", 8), + kvikio_nthreads=int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ), ) dask_ctx = DaskContext( @@ -997,7 +1015,19 @@ 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=int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ), + ), rapidsmpf_options_as_bytes, ) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index b8c7067119dd..8005fb2ca78a 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.defaults import ray import ray.exceptions import ucxx._lib.libucxx as ucx_api @@ -239,6 +240,7 @@ 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, @@ -246,6 +248,7 @@ def __init__( 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 +345,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 +356,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: @@ -795,6 +801,15 @@ def __init__( "int", executor_options.get("num_py_executors", 8), ), + kvikio_nthreads=int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ), hardware_binding=hw_binding, memory_resource_config=mr_config, worker_id=worker_id, @@ -863,6 +878,15 @@ def _reset( [ rank.reset.remote( rapidsmpf_options_as_bytes=rapidsmpf_options_as_bytes, + kvikio_nthreads=int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ), ) for rank in self._rank_actors ] From 26bd8999a28fb70d9ebaf99904d6bbe2efce041f Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 12 Aug 2026 19:28:13 +0000 Subject: [PATCH 4/8] handle engine reset --- python/cudf_polars/cudf_polars/engine/dask.py | 3 +++ python/cudf_polars/cudf_polars/engine/ray.py | 3 +++ python/cudf_polars/cudf_polars/engine/spmd.py | 13 ++++++++----- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 09eebc2e2a36..8a26dc01b916 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1003,6 +1003,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( diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 8005fb2ca78a..2ed0c53fcdbc 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -866,6 +866,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 diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 19726d8dc8ab..3b15725c0d38 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -611,6 +611,14 @@ def _reset( engine_options=engine_options, ) 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) kvikio_nthreads = int( executor_options.get( "kvikio_nthreads", @@ -623,11 +631,6 @@ def _reset( if kvikio_nthreads <= 0: raise ValueError(f"kvikio_nthreads must be positive, got {kvikio_nthreads}") kvikio.defaults.set("num_threads", kvikio_nthreads) - 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) engine_options = engine_options or {} quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( "quent_context" From a58c6d068648047babbd2e5f6374cd94dddd4f4e Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 13 Aug 2026 13:14:08 +0000 Subject: [PATCH 5/8] DRY --- python/cudf_polars/cudf_polars/engine/dask.py | 26 +++++------------ python/cudf_polars/cudf_polars/engine/ray.py | 26 +++++------------ python/cudf_polars/cudf_polars/engine/spmd.py | 28 ++----------------- .../cudf_polars/cudf_polars/utils/config.py | 20 +++++++++---- 4 files changed, 31 insertions(+), 69 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 8a26dc01b916..34f166ac59df 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -52,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 @@ -953,15 +957,7 @@ def __init__( rapidsmpf_options_as_bytes, quent_context=quent_context, num_py_executors=executor_options.get("num_py_executors", 8), - kvikio_nthreads=int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ), + kvikio_nthreads=resolve_kvikio_nthreads(executor_options), ) dask_ctx = DaskContext( @@ -1021,15 +1017,7 @@ def _reset( functools.partial( _reset_worker, uid=ctx.rapidsmpf_id, - kvikio_nthreads=int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ), + kvikio_nthreads=resolve_kvikio_nthreads(executor_options), ), rapidsmpf_options_as_bytes, ) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 2ed0c53fcdbc..decc8c3d2089 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -49,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 @@ -801,15 +805,7 @@ def __init__( "int", executor_options.get("num_py_executors", 8), ), - kvikio_nthreads=int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ), + kvikio_nthreads=resolve_kvikio_nthreads(executor_options), hardware_binding=hw_binding, memory_resource_config=mr_config, worker_id=worker_id, @@ -881,15 +877,7 @@ def _reset( [ rank.reset.remote( rapidsmpf_options_as_bytes=rapidsmpf_options_as_bytes, - kvikio_nthreads=int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ), + kvikio_nthreads=resolve_kvikio_nthreads(executor_options), ) for rank in self._rank_actors ] diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 3b15725c0d38..135fb7787f02 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -7,7 +7,6 @@ import contextlib import dataclasses import json -import os import uuid from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast @@ -61,6 +60,7 @@ MemoryResourceConfig, SPMDContext, StreamingExecutor, + resolve_kvikio_nthreads, ) if TYPE_CHECKING: @@ -427,18 +427,7 @@ def __init__( ) bind_to_gpu(hw_binding) - kvikio_nthreads = int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ) - if kvikio_nthreads <= 0: - raise ValueError(f"kvikio_nthreads must be positive, got {kvikio_nthreads}") - kvikio.defaults.set("num_threads", kvikio_nthreads) + kvikio.defaults.set("num_threads", resolve_kvikio_nthreads(executor_options)) self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) mr_config: MemoryResourceConfig = engine_options.get( @@ -619,18 +608,7 @@ 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) - kvikio_nthreads = int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ) - if kvikio_nthreads <= 0: - raise ValueError(f"kvikio_nthreads must be positive, got {kvikio_nthreads}") - kvikio.defaults.set("num_threads", kvikio_nthreads) + kvikio.defaults.set("num_threads", resolve_kvikio_nthreads(executor_options)) engine_options = engine_options or {} quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( "quent_context" diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 0ee980f67283..96aa502d9a0b 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -161,6 +161,19 @@ def default_factory() -> T | DefaultT: return default_factory +def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: + """Resolve kvikio thread count from executor options with env var fallback.""" + return int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ) + + def _bool_converter(v: str) -> bool: lowered = v.lower() if lowered in {"true", "yes", "y", "1"}: @@ -764,12 +777,7 @@ class StreamingExecutor: ) ) kvikio_nthreads: int = dataclasses.field( - default_factory=lambda: int( - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ) - ) + default_factory=lambda: resolve_kvikio_nthreads({}) ) min_device_size: int | None = None From 11ee15e5ca9d672ecf56993d0c780ec8abbae228 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 13 Aug 2026 15:16:25 +0000 Subject: [PATCH 6/8] use executor_options[kvikio_nthreads] --- python/cudf_polars/cudf_polars/engine/dask.py | 7 +++++-- python/cudf_polars/cudf_polars/engine/ray.py | 7 +++++-- python/cudf_polars/cudf_polars/engine/spmd.py | 7 +++++-- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 34f166ac59df..649046997880 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -843,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( @@ -957,7 +960,7 @@ def __init__( rapidsmpf_options_as_bytes, quent_context=quent_context, num_py_executors=executor_options.get("num_py_executors", 8), - kvikio_nthreads=resolve_kvikio_nthreads(executor_options), + kvikio_nthreads=executor_options["kvikio_nthreads"], ) dask_ctx = DaskContext( @@ -1017,7 +1020,7 @@ def _reset( functools.partial( _reset_worker, uid=ctx.rapidsmpf_id, - kvikio_nthreads=resolve_kvikio_nthreads(executor_options), + kvikio_nthreads=executor_options["kvikio_nthreads"], ), rapidsmpf_options_as_bytes, ) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index decc8c3d2089..0b4d4078f821 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -726,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 {} @@ -805,7 +808,7 @@ def __init__( "int", executor_options.get("num_py_executors", 8), ), - kvikio_nthreads=resolve_kvikio_nthreads(executor_options), + kvikio_nthreads=executor_options["kvikio_nthreads"], hardware_binding=hw_binding, memory_resource_config=mr_config, worker_id=worker_id, @@ -877,7 +880,7 @@ def _reset( [ rank.reset.remote( rapidsmpf_options_as_bytes=rapidsmpf_options_as_bytes, - kvikio_nthreads=resolve_kvikio_nthreads(executor_options), + kvikio_nthreads=executor_options["kvikio_nthreads"], ) for rank in self._rank_actors ] diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 135fb7787f02..0cc198d7bf3d 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -410,6 +410,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( @@ -427,7 +430,7 @@ def __init__( ) bind_to_gpu(hw_binding) - kvikio.defaults.set("num_threads", resolve_kvikio_nthreads(executor_options)) + kvikio.defaults.set("num_threads", executor_options["kvikio_nthreads"]) self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) mr_config: MemoryResourceConfig = engine_options.get( @@ -608,7 +611,7 @@ 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) - kvikio.defaults.set("num_threads", resolve_kvikio_nthreads(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" From ae08a6f9783c1ef4aa09bdcf1af04298b053860c Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 13 Aug 2026 15:39:42 +0000 Subject: [PATCH 7/8] respect existing kvikio thread count --- python/cudf_polars/cudf_polars/utils/config.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 96aa502d9a0b..00a5e1b263d5 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -29,6 +29,8 @@ import os from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar +import kvikio.defaults + if TYPE_CHECKING: import uuid from collections.abc import Callable @@ -163,15 +165,13 @@ def default_factory() -> T | DefaultT: def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: """Resolve kvikio thread count from executor options with env var fallback.""" - return int( - executor_options.get( - "kvikio_nthreads", - os.environ.get( - "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", - os.environ.get("KVIKIO_NTHREADS", "256"), - ), - ) - ) + if "kvikio_nthreads" in executor_options: + return int(executor_options["kvikio_nthreads"]) + cudf_env = os.environ.get("CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS") + if cudf_env is not None: + return int(cudf_env) + current = kvikio.defaults.get("num_threads") + return current if current > 1 else 256 def _bool_converter(v: str) -> bool: From 403f74acb7bfdd6ca1d9c504a8d1246ff47e3b2e Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 13 Aug 2026 17:08:44 +0000 Subject: [PATCH 8/8] document limitation --- python/cudf_polars/cudf_polars/engine/core.py | 7 ++++ .../cudf_polars/cudf_polars/utils/config.py | 36 +++++++++++++------ 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index b382f12099a4..8ee02917d952 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/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 diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 00a5e1b263d5..0af764cada82 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -29,8 +29,6 @@ import os from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar -import kvikio.defaults - if TYPE_CHECKING: import uuid from collections.abc import Callable @@ -165,13 +163,15 @@ def default_factory() -> T | DefaultT: def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: """Resolve kvikio thread count from executor options with env var fallback.""" - if "kvikio_nthreads" in executor_options: - return int(executor_options["kvikio_nthreads"]) - cudf_env = os.environ.get("CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS") - if cudf_env is not None: - return int(cudf_env) - current = kvikio.defaults.get("num_threads") - return current if current > 1 else 256 + return int( + executor_options.get( + "kvikio_nthreads", + os.environ.get( + "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", + os.environ.get("KVIKIO_NTHREADS", "256"), + ), + ) + ) def _bool_converter(v: str) -> bool: @@ -699,12 +699,26 @@ class StreamingExecutor: Maximum number of workers for the Python ThreadPoolExecutor. Default is 8. kvikio_nthreads - Number of threads in the kvikio thread pool. Defaults to 256. - This can be set via + Number of threads in the kvikio thread pool. Defaults to 256, which is + tuned for cloud object-store IO. This can be set via - ``executor_options`` passed to ``polars.GPUEngine`` - the ``CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS`` environment variable - the ``KVIKIO_NTHREADS`` environment variable (lower precedence) + + .. warning:: + + kvikio uses a single process-wide thread pool. When a streaming + engine is created, it configures that pool to ``kvikio_nthreads`` + threads. This operation blocks until all in-flight kvikio IO in the + process completes and then rebuilds the pool. As a result: + + - Any code in the same process that is using kvikio concurrently at + engine creation time will be disrupted. + - Any ``kvikio.defaults.set("num_threads", N)`` call made before + engine creation will be overridden. Use the ``kvikio_nthreads`` + executor option or ``KVIKIO_NTHREADS`` environment variable + instead. quent_context Quent tracing context. When ``None`` (default), Quent tracing is disabled. Pass a :class:`~cudf_polars.quent.QuentContext` instance to enable tracing.