review: use generateRange function for gpus

Also update volume to be added to container. This still
is pending testing on AWS!

Signed-off-by: vsoch <vsoch@users.noreply.github.com>

Author: vsoch

pkg/constants/constants.go

@@ -116,6 +116,7 @@ const (
 
 	// Ensure MPI has full memory of the host
 	FluxMemoryVolumeName = "shared-memory"
+	FluxMemoryVolumePath = "/dev/shm"
 
 	// emptyDir volume using for complete spack view software
 	FluxSpackViewVolumeName = "spack-install"

pkg/runtime/framework/plugins/flux/flux.go

@@ -201,6 +201,7 @@ func (f *Flux) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) e
 					*corev1ac.VolumeMount().WithName(constants.FluxSpackViewVolumeName).WithMountPath(constants.FluxSpackViewVolumePath),
 					*corev1ac.VolumeMount().WithName(configMapName).WithMountPath(constants.FluxConfigVolumeName).WithReadOnly(true),
 					*corev1ac.VolumeMount().WithName(constants.FluxCurveVolumeName).WithMountPath(constants.FluxCurveVolumePath).WithReadOnly(true),
+					*corev1ac.VolumeMount().WithName(constants.FluxMemoryVolumeName).WithMountPath(constants.FluxMemoryVolumePath).WithReadOnly(true),
 				)
 			}
 		}
@@ -429,12 +430,8 @@ func (f *Flux) generateFluxEntrypoint(trainJob *trainer.TrainJob, info *runtime.
 	Rspec := fmt.Sprintf("--cores=0-%d", tasks-1)
 	if gpus > 0 {
 		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)
-		}
+		gpuSpec := generateRange(int32(gpus), 0)
+		Rspec = fmt.Sprintf("%s --gpu=%s", Rspec, gpuSpec)
 	}
 	return fmt.Sprintf(entrypointTemplate, Rspec, mainHost, flags)
 }