Serialize engine config in new pdsh benchmark CLI - #22572
Merged
rapids-bot[bot] merged 3 commits intoMay 19, 2026
Merged
Conversation
This updates the cudf-polars benchmarks CLI to serialize the engine configuration. This will let us see what options were *actually* used. Authors: - Tom Augspurger (https://github.com/TomAugspurger) Approvers: - Matthew Roeschke (https://github.com/mroeschke) - Mads R. B. Kristensen (https://github.com/madsbk) URL: NVIDIA#22365 (cherry picked from commit 38f2d41)
This was missed while merging NVIDIA#22365 Authors: - Lawrence Mitchell (https://github.com/wence-) Approvers: - Tom Augspurger (https://github.com/TomAugspurger) - Matthew Murray (https://github.com/Matt711) URL: NVIDIA#22549 (cherry picked from commit f239649)
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 (5)
✅ Files skipped from review due to trivial changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughResolved RapidsMPF options are now stored on StreamingEngine instances to enable serialization in benchmark outputs. Engine-aware serialization extracts configuration from live engine objects for distributed/GPU frontends while using ChangesRapidsMPF options storage and benchmark serialization
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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
🤖 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/experimental/benchmarks/utils_new_frontends.py`:
- Around line 1266-1268: _finalize_benchmark_run is being called with
engine=engine after the engine context has exited, so RunConfig.serialize() may
read a torn/cleared engine; move the _finalize_benchmark_run(...) call (the ones
in run_polars_ray and run_polars_dask) so they execute inside the corresponding
"with ... as engine:" block (before the context manager returns) to ensure the
engine is still alive when RunConfig.serialize()/finalization runs.
In `@python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py`:
- Around line 619-620: _reset() rebuilds worker contexts from newly resolved
options but doesn't update the persisted field self.rapidsmpf_options (and its
serialized bytes), causing stale options to be emitted. In the _reset()
implementation, after calling resolve_rapidsmpf_options(...) and before
rebuilding contexts, assign the result to self.rapidsmpf_options and reserialize
(e.g., update rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() or
the instance field that holds the bytes) so both the in-memory options and their
serialized representation remain synchronized with the newly resolved options
used to recreate worker contexts.
In `@python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py`:
- Around line 553-554: In RayEngine._reset() the resolved options are only
converted to bytes (rapidsmpf_options_as_bytes) and self.rapidsmpf_options is
left stale; change _reset() to call
resolve_rapidsmpf_options(rapidsmpf_options), assign the result back to
self.rapidsmpf_options, then call self.rapidsmpf_options.serialize() (updating
rapidsmpf_options_as_bytes) so the engine state and persisted bytes reflect the
fresh config; reference the resolve_rapidsmpf_options, self.rapidsmpf_options,
rapidsmpf_options_as_bytes and serialize() symbols when making the change.
In `@python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py`:
- Line 353: SPMDEngine stores init-time options in self.rapidsmpf_options but
_reset() resolves options into a local variable and never updates the instance
attribute, causing stale persisted options; modify the _reset() implementation
(and the other occurrences where resolve_rapidsmpf_options(...) is called) to
assign the resolved value back to self.rapidsmpf_options (e.g.,
self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)) and use
that instance attribute everywhere the local resolved variable is currently used
so the engine and serialization always reflect the latest options.
🪄 Autofix (Beta)
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: 33f7fe23-2129-4e84-be24-97416d0a624d
📒 Files selected for processing (5)
python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.pypython/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.pypython/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.pypython/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.pypython/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py
Comment on lines
+1266
to
+1268
| _finalize_benchmark_run( | ||
| args, run_config, validation_failures, query_failures, engine=engine | ||
| ) |
There was a problem hiding this comment.
Serialize before leaving the Ray/Dask engine context.
These calls now pass engine=engine, but they execute after the with ... as engine: block. At that point the engine is shut down, so RunConfig.serialize() may read cleared/invalid engine config and emit incorrect benchmark config.
Suggested fix
- _finalize_benchmark_run(
- args, run_config, validation_failures, query_failures, engine=engine
- )
+ _finalize_benchmark_run(
+ args, run_config, validation_failures, query_failures, engine=engine
+ )Apply the same move for both run_polars_ray and run_polars_dask so finalization happens while the engine is still alive.
Also applies to: 1328-1330
🤖 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/experimental/benchmarks/utils_new_frontends.py`
around lines 1266 - 1268, _finalize_benchmark_run is being called with
engine=engine after the engine context has exited, so RunConfig.serialize() may
read a torn/cleared engine; move the _finalize_benchmark_run(...) call (the ones
in run_polars_ray and run_polars_dask) so they execute inside the corresponding
"with ... as engine:" block (before the context manager returns) to ensure the
engine is still alive when RunConfig.serialize()/finalization runs.
Comment on lines
+619
to
+620
| self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) | ||
| rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() |
There was a problem hiding this comment.
Keep self.rapidsmpf_options synchronized on reset.
This new persisted field is initialized here, but _reset() currently rebuilds worker contexts from newly resolved options without updating self.rapidsmpf_options. After a reset, benchmark serialization will emit stale RapidsMPF options.
Suggested fix
def _reset(
@@
- rapidsmpf_options_as_bytes = resolve_rapidsmpf_options(
- rapidsmpf_options
- ).serialize()
+ self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)
+ rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) | |
| rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() | |
| def _reset( | |
| self, | |
| rapidsmpf_options: dict[str, Any] | None = None, | |
| ) -> None: | |
| """Reset the Dask cluster, optionally updating rapidsmpf options.""" | |
| self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) | |
| rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() | |
| # ... rest of the method |
🤖 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/experimental/rapidsmpf/frontend/dask.py`
around lines 619 - 620, _reset() rebuilds worker contexts from newly resolved
options but doesn't update the persisted field self.rapidsmpf_options (and its
serialized bytes), causing stale options to be emitted. In the _reset()
implementation, after calling resolve_rapidsmpf_options(...) and before
rebuilding contexts, assign the result to self.rapidsmpf_options and reserialize
(e.g., update rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() or
the instance field that holds the bytes) so both the in-memory options and their
serialized representation remain synchronized with the newly resolved options
used to recreate worker contexts.
Comment on lines
+553
to
+554
| self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) | ||
| rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() |
There was a problem hiding this comment.
Update persisted RapidsMPF options in RayEngine._reset().
self.rapidsmpf_options is now part of engine state, but _reset() still computes fresh options only as local bytes. That leaves self.rapidsmpf_options stale and can serialize incorrect post-reset config.
Suggested fix
def _reset(
@@
- rapidsmpf_options_as_bytes = resolve_rapidsmpf_options(
- rapidsmpf_options
- ).serialize()
+ self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)
+ rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize()🤖 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/experimental/rapidsmpf/frontend/ray.py` around
lines 553 - 554, In RayEngine._reset() the resolved options are only converted
to bytes (rapidsmpf_options_as_bytes) and self.rapidsmpf_options is left stale;
change _reset() to call resolve_rapidsmpf_options(rapidsmpf_options), assign the
result back to self.rapidsmpf_options, then call
self.rapidsmpf_options.serialize() (updating rapidsmpf_options_as_bytes) so the
engine state and persisted bytes reflect the fresh config; reference the
resolve_rapidsmpf_options, self.rapidsmpf_options, rapidsmpf_options_as_bytes
and serialize() symbols when making the change.
| @@ -350,7 +350,7 @@ def __init__( | |||
| ) | |||
| bind_to_gpu(hw_binding) | |||
|
|
|||
| rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) | |||
| self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) | |||
There was a problem hiding this comment.
Persisted options become stale after _reset().
SPMDEngine now uses self.rapidsmpf_options for init-time construction, but _reset() rebuilds the context from a local resolved variable and does not refresh the instance attribute. That can cause benchmark serialization to report old options.
Suggested fix
- rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)
+ self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)
@@
- self._ctx = Context.from_options(self._comm.logger, self._mr, rapidsmpf_options)
+ self._ctx = Context.from_options(
+ self._comm.logger, self._mr, self.rapidsmpf_options
+ )Also applies to: 364-364, 369-369, 383-383
🤖 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/experimental/rapidsmpf/frontend/spmd.py` at
line 353, SPMDEngine stores init-time options in self.rapidsmpf_options but
_reset() resolves options into a local variable and never updates the instance
attribute, causing stale persisted options; modify the _reset() implementation
(and the other occurrences where resolve_rapidsmpf_options(...) is called) to
assign the resolved value back to self.rapidsmpf_options (e.g.,
self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options)) and use
that instance attribute everywhere the local resolved variable is currently used
so the engine and serialization always reflect the latest options.
bdice
approved these changes
May 19, 2026
…ize-config-backport
TomAugspurger
added
bug
Contributor
Author
|
/merge |
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
4 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
This backports a pair of commits for the cudf-polars benchmarking CLI. We're currently running benchmarks against both release/26.06 and main.