Skip to content

feat: customize R generation to include GPUs #3325

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vsoch wants to merge 3 commits into
base: master
Choose a base branch
from
+66 −8
Open

feat: customize R generation to include GPUs #3325

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b9cee29
feat: customize R generation to include GPUs
vsoch Mar 13, 2026
b499369
review: use generateRange function for gpus
vsoch Mar 13, 2026
764a765
review and add gpu trainer example
vsoch Mar 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions pkg/constants/constants.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,9 @@ const (
// Path for Flux curve path
FluxCurveVolumePath = "/curve"

// Ensure MPI has full memory of the host
FluxMemoryVolumeName = "shared-memory"

// emptyDir volume using for complete spack view software
FluxSpackViewVolumeName = "spack-install"

Expand Down
30 changes: 23 additions & 7 deletions pkg/runtime/framework/plugins/flux/flux.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,6 +295,10 @@ func getViewVolumes(configMapName string) []corev1ac.VolumeApplyConfiguration {
spackInstallAC := corev1ac.Volume().
WithName(constants.FluxSpackViewVolumeName).
WithEmptyDir(corev1ac.EmptyDirVolumeSource())
memoryVolumeAC := corev1ac.Volume().
WithName(constants.FluxMemoryVolumeName).
WithEmptyDir(corev1ac.EmptyDirVolumeSource().
WithMedium(corev1.StorageMediumMemory))
fluxVolumeAC := corev1ac.Volume().
Comment on lines +299 to 303
WithEmptyDir(corev1ac.EmptyDirVolumeSource()).
WithName(constants.FluxInstallVolumeName)
Expand All @@ -305,7 +309,7 @@ func getViewVolumes(configMapName string) []corev1ac.VolumeApplyConfiguration {
WithName(configMapName).
WithDefaultMode(0755),
)
return []corev1ac.VolumeApplyConfiguration{*spackInstallAC, *fluxVolumeAC, *cmAC}
return []corev1ac.VolumeApplyConfiguration{*spackInstallAC, *fluxVolumeAC, *cmAC, *memoryVolumeAC}
}

// buildInitScriptConfigMap creates a ConfigMapApplyConfiguration to support server-side Apply
Expand Down Expand Up @@ -402,25 +406,37 @@ func (f *Flux) generateFluxEntrypoint(trainJob *trainer.TrainJob, info *runtime.
// Derive number of tasks
// This may not technically be the number of processes per node,
// but that is all the TrainJob can currently represent.
var tasks string
var tasks int32
var flags string

nodes := *trainJob.Spec.Trainer.NumNodes
if trainJob.Spec.Trainer.NumProcPerNode != nil {
tasks = fmt.Sprintf("-N %d -n %d", nodes, *trainJob.Spec.Trainer.NumProcPerNode*nodes)
tasks = *trainJob.Spec.Trainer.NumProcPerNode
} else {
tasks = fmt.Sprintf("-N %d -n %d", nodes, *info.RuntimePolicy.MLPolicySource.Flux.NumProcPerNode*nodes)
tasks = *info.RuntimePolicy.MLPolicySource.Flux.NumProcPerNode
}
flags = fmt.Sprintf("-N %d -n %d", nodes, tasks*nodes)

// Derive number of GPUs from resources. In Flux, -g is --gpus-per-task
resourcesPerNode := ptr.Deref(runtime.ExtractResourcePerNodeFromRuntime(info), corev1.ResourceRequirements{})
if jobTrainer := trainJob.Spec.Trainer; jobTrainer != nil && jobTrainer.ResourcesPerNode != nil {
resourcesPerNode = ptr.Deref(jobTrainer.ResourcesPerNode, corev1.ResourceRequirements{})
}
gpus := runtime.GetNumGPUPerNode(&resourcesPerNode)

// Resource file for cluster includes GPUs or not
// flux R encode --hosts=${hosts} --cores=0-1 --gpu=0
Rspec := fmt.Sprintf("--cores=0-%d", tasks-1)
if gpus > 0 {
tasks = fmt.Sprintf("%s -g %d", tasks, gpus)
flags = fmt.Sprintf("%s -g %d", flags, gpus)
gpus = gpus - 1
if gpus == 0 {
Rspec = fmt.Sprintf("%s --gpu=0", Rspec)
} else {
Rspec = fmt.Sprintf("%s --gpu=0-%d", Rspec, gpus)
}
}

return fmt.Sprintf(entrypointTemplate, mainHost, tasks)
return fmt.Sprintf(entrypointTemplate, Rspec, mainHost, flags)
}

// generateInitEntrypoint generates the flux entrypoint to prepare flux
Expand Down
2 changes: 1 addition & 1 deletion pkg/runtime/framework/plugins/flux/templates/entrypoint.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ chown -R ${fluxuid} ${curvepath}

# 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
mv /tmp/R ${configroot}/etc/flux/system/R

# Put the state directory in /var/lib on shared view
Expand Down
Loading