Change kvikio default thread count in cudf-polars to 256 - #23634
Conversation
github-actions
Bot
added
Python
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 (4)
🚧 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 Mergeability Score: 🔵 Low · up to The PR changes the default KvikIO thread count to 256, but an SPMD configuration path can still pass zero or negative values to KvikIO before validation, potentially causing a runtime failure for invalid settings. The change is otherwise mergeable with explicit owner follow-up on that validation gap. 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
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
wence-
left a comment
wence-
left a comment
There was a problem hiding this comment.
I have one concern about the duplication of defaults and parsing everywhere.
Also, I presume that if kvikio has already set up its thread pool then none of these settings take effect?
Finally: does this new default have any negative consequences on non-cloud runs?
| 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) |
There was a problem hiding this comment.
Something seems to be off here. I thought the construction of the executor_options did all of this parsing/validation. And yet here we are re-inspecting the environment.
This seems very fragile, there are now many places I have to remember to handle a potential renaming/rejigging of the meaning/default value/validation.
There was a problem hiding this comment.
This still does the parsing/applying defaults everywhere rather than just relying on the executor config having done it. Under what circumstances will execution_options not have a kvikio_nthreads slot set correctly?
Not as far as I can tell. See the description of #23633 |
No because the the thread pool is reset when we call void defaults::set_thread_pool_nthreads(unsigned int nthreads)
{
KVIKIO_EXPECT(
nthreads > 0, "number of threads must be a positive integer", std::invalid_argument);
thread_pool().reset(nthreads, make_thread_pool_init_task("kvikio"));
}But I am concerned looking at what /**
* @brief Reset the pool with a new number of threads and a new initialization function. Waits for all currently running tasks to be completed, then destroys all threads in the pool and creates a new thread pool with the new number of threads and initialization function. Any tasks that were waiting in the queue before the pool was reset will then be executed by the new threads. If the pool was paused before resetting it, the new pool will be paused as well.
*
* @param num_threads The number of threads to use.
* @param init_task An initialization function to run in each thread before it starts to execute any submitted tasks. The function must take no arguments and have no return value. It will only be executed exactly once, when the thread is first constructed.
*/
void reset(const concurrency_t num_threads, const std::function<void()>& init_task)
{
#ifdef BS_THREAD_POOL_ENABLE_PAUSE
std::unique_lock tasks_lock(tasks_mutex);
const bool was_paused = paused;
paused = true;
tasks_lock.unlock();
#endif
wait();
destroy_threads();
thread_count = determine_thread_count(num_threads);
threads = std::make_unique<std::thread[]>(thread_count);
create_threads(init_task);
#ifdef BS_THREAD_POOL_ENABLE_PAUSE
tasks_lock.lock();
paused = was_paused;
#endif
}So this means
|
When `set_thread_pool_nthreads`, we reset the thread pool which drains and detroys the previous one. We avoid that if the number of threads we are setting is unchanged.
See
```cpp
void defaults::set_thread_pool_nthreads(unsigned int nthreads)
{
KVIKIO_EXPECT(
nthreads > 0, "number of threads must be a positive integer", std::invalid_argument);
thread_pool().reset(nthreads, make_thread_pool_init_task("kvikio"));
}
```
xref NVIDIA/cudf#23634
Authors:
- Matthew Murray (https://github.com/Matt711)
Approvers:
- Mads R. B. Kristensen (https://github.com/madsbk)
URL: #1032
3 participants
Description
Makes the default 256 kvikio threads and makes kvikio python package a required dependency of cudf-polars.
Benchmarks
SF300, Q1 hot times (S3)
SF300, Q1 hot times (NVMe)
Checklist