feat: customize R generation to include GPUs#3325
feat: customize R generation to include GPUs#3325vsoch wants to merge 3 commits intokubeflow:masterfrom
Conversation
Flux detection of GPUs depends on hwloc plugins, and a newer version. To get around any edge cases where the dependencies are missing, we can easily generate the R from the expected resource spec (cores and gpus). We can also add a shared memory mount to ensure the job MPI gets all available shared memory of the host. The current default is 64M (automatic from container runtime) and it can have implications for MPI performance. Signed-off-by: vsoch <vsoch@users.noreply.github.com>
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
There was a problem hiding this comment.
Pull request overview
Updates the Flux runtime plugin to generate a Flux resource file (flux R encode) based on derived per-node CPU/GPU topology and introduces a new in-memory EmptyDir volume intended for shared memory.
Changes:
- Parameterize
templates/entrypoint.shsoflux R encodecan receive a computed resource spec (cores/GPU ranges). - Refactor Flux entrypoint generation to build explicit Flux flags (
-N/-nand optional-g) and derive anRspecstring for resource encoding. - Add a new
shared-memoryvolume constant and include a memory-backed EmptyDir volume in Flux view volumes.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| pkg/runtime/framework/plugins/flux/templates/entrypoint.sh | Accepts a formatted resource spec argument for flux R encode. |
| pkg/runtime/framework/plugins/flux/flux.go | Builds Rspec/flags for Flux execution and adds a memory-backed EmptyDir volume to the pod spec. |
| pkg/constants/constants.go | Introduces FluxMemoryVolumeName constant for the new shared-memory volume. |
You can also share your feedback on Copilot code review. Take the survey.
| memoryVolumeAC := corev1ac.Volume(). | ||
| WithName(constants.FluxMemoryVolumeName). | ||
| WithEmptyDir(corev1ac.EmptyDirVolumeSource(). | ||
| WithMedium(corev1.StorageMediumMemory)) | ||
| fluxVolumeAC := corev1ac.Volume(). |
| } else { | ||
| tasks = fmt.Sprintf("-N %d -n %d", nodes, *info.RuntimePolicy.MLPolicySource.Flux.NumProcPerNode*nodes) | ||
| tasks = *info.RuntimePolicy.MLPolicySource.Flux.NumProcPerNode | ||
| } |
| // Path for Flux curve path | ||
| FluxCurveVolumePath = "/curve" | ||
|
|
||
| // Ensure MPI has full memory of the host |
| # Generate host resources | ||
| hosts=$(cat ${configroot}/etc/flux/system/hostlist) | ||
| flux R encode --hosts=${hosts} --local > /tmp/R | ||
| flux R encode --hosts=${hosts} %s > /tmp/R |
Also update volume to be added to container. This still is pending testing on AWS! Signed-off-by: vsoch <vsoch@users.noreply.github.com>
|
@andreyvelich I'm going for a quick run and will test with GPU when I am back! Would you like an example added to the examples/flux directory that uses GPUs? My plan is to test on AWS, lammps with GPU. |
Sure, let's add another example in Flux subdirectory. |
|
GPU working is a go! 
I have some changes I will push shortly, and looks like I need to look into those failing tests. |
Signed-off-by: vsoch <vsoch@users.noreply.github.com>
What this PR does / why we need it:
Flux detection of GPUs depends on hwloc plugins, and a newer version. To get around any edge cases where the dependencies are missing, we can easily generate the R from the expected resource spec (cores and gpus). We can also add a shared memory mount to ensure the job MPI gets all available shared memory of the host. The current default is 64M (automatic from container runtime) and it can have implications for MPI performance.
This will close #3321
I would like to test this with GPUs on AWS today before we do any kind of merge. Thank you!
Which issue(s) this PR fixes (optional, in
Fixes #<issue number>, #<issue number>, ...format, will close the issue(s) when PR gets merged):Fixes #3321
Checklist: