Allow the specification of an environment when detecting an MPI install

Signed-off-by: Geoffroy Vallee <[email protected]>
gvallee · Aug 11, 2021 · b91e791 · b91e791
1 parent 552f260
commit b91e791
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 12 deletions.
diff --git a/internal/pkg/mpich/mpich.go b/internal/pkg/mpich/mpich.go
@@ -31,6 +31,6 @@ func GetConfigureExtraArgs() []string {
 }
 
 // DetectFromDir tries to figure out which version of MPICH is installed in a given directory
-func DetectFromDir(dir string) (string, string, error) {
+func DetectFromDir(dir string, env []string) (string, string, error) {
 	return "", "", fmt.Errorf("not implemented")
 }
diff --git a/internal/pkg/openmpi/openmpi.go b/internal/pkg/openmpi/openmpi.go
@@ -58,7 +58,7 @@ func parseOmpiInfoOutputForVersion(output string) (string, error) {
 }
 
 // DetectFromDir tries to figure out which version of OpenMPI is installed in a given directory
-func DetectFromDir(dir string) (string, string, error) {
+func DetectFromDir(dir string, env []string) (string, string, error) {
 	targetBin := filepath.Join(dir, "bin", "ompi_info")
 	if !util.FileExists(targetBin) {
 		return "", "", fmt.Errorf("%s does not exist, not an OpenMPI implementation", targetBin)
@@ -67,22 +67,30 @@ func DetectFromDir(dir string) (string, string, error) {
 	var versionCmd advexec.Advcmd
 	versionCmd.BinPath = targetBin
 	versionCmd.CmdArgs = append(versionCmd.CmdArgs, "--version")
-	newLDPath := filepath.Join(dir, "lib") + ":$LD_LIBRARY_PATH"
-	newPath := filepath.Join(dir, "bin") + ":$PATH"
-	versionCmd.Env = []string{"LD_LIBRARY_PATH=" + newLDPath, "PATH=" + newPath}
+	versionCmd.Env = env
+	if env == nil {
+		newLDPath := filepath.Join(dir, "lib") + ":$LD_LIBRARY_PATH"
+		newPath := filepath.Join(dir, "bin") + ":$PATH"
+		versionCmd.Env = append(versionCmd.Env, "LD_LIBRARY_PATH="+newLDPath)
+		versionCmd.Env = append(versionCmd.Env, "PATH="+newPath)
+	}
 	res := versionCmd.Run()
 	if res.Err != nil {
-		// If it fails we try with OPAL_PREFIX set
-		versionCmd.Env = append(versionCmd.Env, "OPAL_PREFIX="+dir)
-		res = versionCmd.Run()
+		// If it fails we try with OPAL_PREFIX set. We create a new command to avoid "exec: already started" issues.
+		var versionCmdWithOpalPrefix advexec.Advcmd
+		versionCmdWithOpalPrefix.BinPath = versionCmd.BinPath
+		versionCmdWithOpalPrefix.CmdArgs = versionCmd.CmdArgs
+		versionCmdWithOpalPrefix.Env = versionCmd.Env
+		versionCmdWithOpalPrefix.Env = append(versionCmd.Env, "OPAL_PREFIX="+dir)
+		res = versionCmdWithOpalPrefix.Run()
 		if res.Err != nil {
 			log.Printf("unable to run ompi_info: %s; stdout: %s; stderr: %s", res.Err, res.Stdout, res.Stderr)
 			return "", "", res.Err
 		}
 	}
 	version, err := parseOmpiInfoOutputForVersion(res.Stdout)
 	if err != nil {
-		return "", "", err
+		return "", "", fmt.Errorf("parseOmpiInfoOutputForVersion() failed - %w", err)
 	}
 
 	return ID, version, nil

diff --git a/internal/pkg/openmpi/open_test.go → internal/pkg/openmpi/openmpi_test.go b/internal/pkg/openmpi/open_test.go → internal/pkg/openmpi/openmpi_test.go
diff --git a/pkg/implem/implem.go b/pkg/implem/implem.go
@@ -47,14 +47,14 @@ func IsMPI(i *Info) bool {
 // - the implementation (e.g., openmpi) and the function figures out where it is installed
 // - a few other combinations of these to provide a flexible way to handle various implementation of MPI
 // If no suitable implementation can be found, the function returns an error
-func (i *Info) Load() error {
+func (i *Info) Load(env []string) error {
 	if i.InstallDir != "" && (i.ID == "" || i.Version == "") {
 		var err error
-		i.ID, i.Version, err = openmpi.DetectFromDir(i.InstallDir)
+		i.ID, i.Version, err = openmpi.DetectFromDir(i.InstallDir, env)
 		if err == nil {
 			return nil
 		}
-		i.ID, i.Version, err = mpich.DetectFromDir(i.InstallDir)
+		i.ID, i.Version, err = mpich.DetectFromDir(i.InstallDir, env)
 		if err == nil {
 			return nil
 		}