Add flux executor shutdown #479
Conversation
WalkthroughThe changes in this pull request involve modifications to the Jupyter notebook Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Executor
participant Task
User->>Executor: Initialize with backend="flux"
Executor->>Task: Submit task
Task-->>Executor: Task completed
Executor-->>User: Return results
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
for more information, see https://pre-commit.ci
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
executorlib/interactive/flux.py (1)
103-105: Good addition of executor cleanup, but could be more robust.The addition of flux executor shutdown is a good improvement for resource cleanup. However, consider these enhancements:
Consider applying this diff to make the shutdown more robust:
# shutdown flux executor if self._flux_executor is not None: - self._flux_executor.shutdown(wait=wait) + try: + self._flux_executor.shutdown(wait=wait) + except Exception as e: + # Log but don't re-raise as we want to continue cleanup + print(f"Warning: Error during flux executor shutdown: {e}") + finally: + self._flux_executor = None # Clear referenceThis change:
- Adds error handling to prevent shutdown failures from blocking cleanup
- Cleans up the executor reference after shutdown
- Maintains the existing behavior while being more resilient
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (2)
executorlib/interactive/flux.py(1 hunks)notebooks/examples.ipynb(8 hunks)
🔇 Additional comments (6)
executorlib/interactive/flux.py (1)
103-105: Verify shutdown behavior in nested executor scenarios.
The shutdown implementation looks correct, but given that the class supports nested executors (via flux_executor_nesting), we should verify the shutdown behavior in nested scenarios.
Let's check for any special handling of nested executors in the codebase:
✅ Verification successful
Let me gather more information about the nested executor handling.
Based on the search results, I can now generate the final response:
Shutdown implementation is correct for nested executors
The current shutdown implementation is appropriate because:
- Nested executors are only used during job submission via
JobspecV1.from_nest_command()and don't create a hierarchy of executor instances - The
_flux_executoris the only executor instance that needs cleanup, regardless of nesting - The nesting flag (
flux_executor_nesting) only affects the job specification during bootup and doesn't impact the shutdown sequence
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for nested executor handling patterns
rg -l "flux_executor_nesting|nest_command" --type py
# Look for any existing shutdown handling of nested executors
ast-grep --pattern 'if self._flux_executor_nesting {
$$$
}'
Length of output: 240
Script:
#!/bin/bash
# Check the implementation details of flux executor nesting
rg -A 10 "flux_executor_nesting" --type py
# Look for any initialization or setup of nested executors
ast-grep --pattern 'def __init__($$$) {
$$$
self._flux_executor_nesting = $_
$$$
}'
# Check for any other shutdown related patterns with nesting
rg -A 5 "shutdown.*nest|nest.*shutdown" --type py
Length of output: 16361
notebooks/examples.ipynb (5)
73-75: LGTM! Simplified initialization pattern.
The simplified initialization improves readability while maintaining the same functionality.
108-120: LGTM! Consistent parallel execution setup.
The parallel execution setup is correctly maintained with the simplified initialization pattern.
526-533: LGTM! GPU configuration is properly set.
The GPU allocation is correctly configured with the simplified initialization pattern.
638-651: LGTM! Clear demonstration of function dependencies.
The example effectively demonstrates chained computations with proper resource allocation.
468-475: LGTM! Verify MPI configuration.
The MPI configuration is correct with the simplified initialization. The pmix mode is properly set for OpenMPI compatibility.
✅ Verification successful
Let me gather more information about the MPI configuration and pmix usage.
MPI configuration is correctly set for OpenMPI v5 compatibility
The code's pmix configuration is correct and well-documented. The documentation in docs/installation.md explicitly states that OpenMPI version 5 requires pmix mode, and the example notebook correctly sets flux_executor_pmi_mode="pmix". The environment also includes the required flux-pmix plugin (version 0.5.0) in the dependencies.
Key validations:
- Documentation confirms OpenMPI v5 requires pmix mode
- Required
flux-pmixplugin is included in dependencies - Input validation in
inputcheck.pyconfirms pmix is a valid PMI mode for Flux - Executor properly documents the pmix requirement for OpenMPI v5 in its parameters
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Verify OpenMPI version to confirm pmix compatibility
rg -l "OpenMPI" | xargs -I {} rg "OpenMPI.*5\." {} || echo "OpenMPI version not found or not version 5.x"
Length of output: 153
Script:
#!/bin/bash
# Search for MPI and pmix related configurations
rg -A 5 -B 5 "pmix|PMI"
# Also check for any MPI configuration documentation
fd -e md -e rst -e txt | xargs rg -l "MPI|mpi"
Length of output: 12225
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (2)
notebooks/examples.ipynb (2)
332-339: LGTM: Data handling example is well structuredThe example clearly demonstrates data handling with preloaded values through
init_function. The explicitblock_allocation=Trueparameter helps emphasize its importance for this functionality.Consider adding a comment explaining that
block_allocation=Trueis required for theinit_functionto work effectively.
630-643: LGTM: Coupled functions example effectively demonstrates dependenciesThe example clearly shows how to chain function executions using futures while maintaining proper resource control through individual resource dictionaries.
Consider adding a comment explaining that the resource specifications can be different for each task in the chain, as this is a powerful feature for optimizing resource usage.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (1)
notebooks/examples.ipynb(8 hunks)
🔇 Additional comments (3)
notebooks/examples.ipynb (3)
72-74: LGTM: API simplification improves usability
The change simplifies the API by removing the nested FluxExecutor context manager while maintaining the same functionality through the backend="flux" parameter.
106-118: LGTM: Resource assignment example maintains clarity
The example effectively demonstrates resource assignment while using the simplified API. The code structure is clean and the resource specification is clear.
462-469: LGTM: MPI example properly handles resource requirements
The example correctly demonstrates MPI usage with proper resource specification and PMI mode configuration for OpenMPI 5 compatibility.
Summary by CodeRabbit
New Features
Executorclass.Bug Fixes