Change kvikio default thread count in cudf-polars to 256 - #23634
Open
Matt711 wants to merge 6 commits into
Open
Conversation
github-actions
Bot
added
Python
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR adds KvikIO as a cuDF Polars dependency, uses it for remote Parquet handling, and configures its thread count across SPMD, Dask, and Ray streaming engines. ChangesKvikIO configuration
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
python/cudf_polars/cudf_polars/engine/options.py (1)
348-350: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument
kvikio_nthreadsinStreamingOptions.The field is public, but the
Parameterssection does not describe it. Add the default, precedence, and both environment variable names so the public documentation matches the new configuration surface.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/engine/options.py` around lines 348 - 350, Document the public kvikio_nthreads field in StreamingOptions’ Parameters section, including its default value, configuration precedence, and both supported environment variable names. Keep the documentation aligned with the existing _opt("executor", "CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", int) configuration behavior.python/cudf_polars/tests/streaming/test_options.py (1)
81-83: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftTest the engine-side effect in addition to option forwarding.
This test can pass even when
SPMDEngine.__init__does not apply the value or_resetleaves the old KvikIO pool active. Add a focused test for default, explicit, and reset behavior aroundkvikio.defaults.set.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/tests/streaming/test_options.py` around lines 81 - 83, Extend test_executor_options_kvikio_nthreads to verify the engine-side behavior, not only the executor-options mapping: mock or spy on kvikio.defaults.set, cover default and explicit kvikio_nthreads values during SPMDEngine initialization, and verify _reset restores the default setting. Preserve the existing forwarding assertion while asserting each expected set call and value.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@dependencies.yaml`:
- Line 426: Add kvikio =${{ minor_version }} to the run dependency list in the
cudf-polars conda recipe, alongside the existing runtime dependencies, so the
generated package includes KvikIO required by cudf_polars.dsl.utils.io.
In `@python/cudf_polars/cudf_polars/engine/spmd.py`:
- Around line 415-424: Centralize KvikIO thread-pool configuration in the
SPMDEngine lifecycle: validate the requested kvikio_nthreads during
construction, reject a second engine’s conflicting value, and invoke
kvikio.defaults.set only after validation succeeds. Update _reset to apply the
accepted kvikio_nthreads value to KvikIO as well as self.config, ensuring
existing engines cannot be silently reconfigured.
In `@python/cudf_polars/cudf_polars/utils/config.py`:
- Around line 766-773: Validate kvikio_nthreads as strictly positive in
StreamingExecutor.__post_init__ before applying the setting, and add the same
guard in SPMDEngine.__init__ before its direct kvikio.defaults.set call. Reject
zero and negative values while preserving existing handling for valid positive
integers.
In `@python/cudf_polars/tests/test_config.py`:
- Around line 804-806: Update test_kvikio_nthreads_default and
test_kvikio_nthreads_from_kvikio_env to clear the relevant KvikIO thread
environment variables with monkeypatch.delenv(..., raising=False) before
evaluating defaults or setting the variable under test, ensuring ambient CI
values cannot affect either assertion.
---
Nitpick comments:
In `@python/cudf_polars/cudf_polars/engine/options.py`:
- Around line 348-350: Document the public kvikio_nthreads field in
StreamingOptions’ Parameters section, including its default value, configuration
precedence, and both supported environment variable names. Keep the
documentation aligned with the existing _opt("executor",
"CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS", int) configuration behavior.
In `@python/cudf_polars/tests/streaming/test_options.py`:
- Around line 81-83: Extend test_executor_options_kvikio_nthreads to verify the
engine-side behavior, not only the executor-options mapping: mock or spy on
kvikio.defaults.set, cover default and explicit kvikio_nthreads values during
SPMDEngine initialization, and verify _reset restores the default setting.
Preserve the existing forwarding assertion while asserting each expected set
call and value.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0af71e53-3485-424f-919e-dde06318b9f3
📒 Files selected for processing (8)
dependencies.yamlpython/cudf_polars/cudf_polars/dsl/utils/io.pypython/cudf_polars/cudf_polars/engine/options.pypython/cudf_polars/cudf_polars/engine/spmd.pypython/cudf_polars/cudf_polars/utils/config.pypython/cudf_polars/pyproject.tomlpython/cudf_polars/tests/streaming/test_options.pypython/cudf_polars/tests/test_config.py
mroeschke
reviewed
Aug 12, 2026
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudf_polars/cudf_polars/engine/dask.py`:
- Around line 1018-1030: Update the _reset worker configuration in
python/cudf_polars/cudf_polars/engine/dask.py lines 1018-1030 and
python/cudf_polars/cudf_polars/engine/ray.py lines 881-889 so omitted
executor_options retain the engine’s existing kvikio_nthreads value before
environment/default fallbacks; explicit reset values must override the retained
value. The python/cudf_polars/cudf_polars/engine/spmd.py lines 614-620 site
requires no direct change unless its corresponding reset handling also omits the
retained value.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7a1ce993-d800-4ac1-b480-08b20d5aae6e
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/engine/dask.pypython/cudf_polars/cudf_polars/engine/ray.py
mroeschke
approved these changes
Aug 12, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Makes the default 256 kvikio threads and makes kvikio python package a required dependency of cudf-polars.
Checklist