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/dependencies.yaml b/dependencies.yaml index 35044134c583..9216a90ac0c9 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/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/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 20b0139650b0..649046997880 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 @@ -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, ) 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/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index b8c7067119dd..0b4d4078f821 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 @@ -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,6 +244,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 +252,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 +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 ] diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index bc1fc56fe8e7..0cc198d7bf3d 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/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" diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 185f43b7142f..0af764cada82 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"}: @@ -685,6 +698,27 @@ 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, 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. @@ -756,6 +790,9 @@ class StreamingExecutor: f"{_env_prefix}__NUM_PY_EXECUTORS", int, default=8 ) ) + kvikio_nthreads: int = dataclasses.field( + default_factory=lambda: resolve_kvikio_nthreads({}) + ) min_device_size: int | None = None spmd_context: SPMDContext | None = None @@ -840,6 +877,10 @@ 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") + 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/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..83d7a5bc4819 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -801,6 +801,53 @@ def test_num_py_executors_from_env( assert config.executor.num_py_executors == 8 +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: + 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.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 + + +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"