Add new API to make it easier to handled installed MPI implementations

go.mod

@@ -5,5 +5,4 @@ go 1.13
 require (
 	github.com/gvallee/go_exec v0.0.1
 	github.com/gvallee/go_util v1.0.1
-	github.com/gvallee/kv v1.0.0
 )

go.sum

@@ -1,4 +1,4 @@
 github.com/gvallee/go_exec v0.0.1 h1:yNP5/fTWnnym4wT17JIEBiRAxOEhclOePdwxCxWsEZ8=
 github.com/gvallee/go_exec v0.0.1/go.mod h1:4AwegK9oPhkgwkd0rjlTwxRw//8cW4pPcCSFLZ6+LZg=
+github.com/gvallee/go_util v1.0.1 h1:Ch/PpAlHrHNmL2Upaxif/Nt4CqtaazDyTXh5fIhutJo=
 github.com/gvallee/go_util v1.0.1/go.mod h1:fTexpwdH/n05Ziu0TXJIQsr7E+46QpBxNdeOOsyC0/s=
-github.com/gvallee/kv v1.0.0/go.mod h1:sfSclfFfLV+Y+9e9FayIbBUOtvbt1779S6q52bSSU5E=

internal/pkg/mpich/mpich.go

@@ -5,23 +5,32 @@
 
 package mpich
 
+import "fmt"
+
 const (
 	// VersionTag is the tag used to refer to the MPI version in MPICH template(s)
 	VersionTag = "MPICHVERSION"
 	// URLTag is the tag used to refer to the MPI URL in MPICH template(s)
 	URLTag = "MPICHURL"
 	// TarballTag is the tag used to refer to the MPI tarball in MPICH template(s)
 	TarballTag = "MPICHTARBALL"
+	// ID is the internal ID for MPICH
+	ID = "mpich"
 )
 
-// MPICHGetExtraMpirunArgs returns the extra mpirun arguments required by MPICH for a specific configuration
-func MPICHGetExtraMpirunArgs() []string {
+// GetExtraMpirunArgs returns the extra mpirun arguments required by MPICH for a specific configuration
+func GetExtraMpirunArgs() []string {
 	var extraArgs []string
 	return extraArgs
 }
 
-// MPICHGetConfigureExtraArgs returns the extra arguments required to configure MPICH
-func MPICHGetConfigureExtraArgs() []string {
+// GetConfigureExtraArgs returns the extra arguments required to configure MPICH
+func GetConfigureExtraArgs() []string {
 	var extraArgs []string
 	return extraArgs
 }
+
+// DetectFromDir tries to figure out which version of MPICH is installed in a given directory
+func DetectFromDir(dir string) (string, string, error) {
+	return "", "", fmt.Errorf("not implemented")
+}

internal/pkg/openmpi/open_test.go

@@ -0,0 +1,21 @@
+// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
+// This software is licensed under a 3-clause BSD license. Please consult the
+// LICENSE.md file distributed with the sources of this project regarding your
+// rights to use or distribute this software.
+
+package openmpi
+
+import "testing"
+
+func TestParseOMPIInfoOutputForVersion(t *testing.T) {
+	output := "Open MPI v3.0.4\n\nhttp://www.open-mpi.org/community/help/\n"
+	expectedResult := "3.0.4"
+
+	version, err := parseOmpiInfoOutputForVersion(output)
+	if err != nil {
+		t.Fatalf("parseOmpiInfoOutputForVersion() failed: %s", err)
+	}
+	if version != expectedResult {
+		t.Fatalf("parseOmpiInfoOutputForVersion() returned %s instead of %s", version, expectedResult)
+	}
+}

internal/pkg/openmpi/openmpi.go

@@ -6,7 +6,14 @@
 package openmpi
 
 import (
+	"fmt"
+	"log"
+	"path/filepath"
+	"strings"
+
+	"github.com/gvallee/go_exec/pkg/advexec"
 	"github.com/gvallee/go_hpc_jobmgr/internal/pkg/sys"
+	"github.com/gvallee/go_util/pkg/util"
 )
 
 const (
@@ -18,18 +25,53 @@ const (
 
 	// TarballTag is the tag used to refer to the MPI tarball in Open MPI template(s)
 	TarballTag = "OMPITARBALL"
+
+	// ID is the internal ID for Open MPI
+	ID = "openmpi"
 )
 
 // GetExtraMpirunArgs returns the set of arguments required for the mpirun command for the target platform
 func GetExtraMpirunArgs(sys *sys.Config) []string {
 	var extraArgs []string
-	/*
-		if sys.IBEnabled {
-			extraArgs = append(extraArgs, "--mca")
-			extraArgs = append(extraArgs, "btl")
-			extraArgs = append(extraArgs, "openib,self,vader")
-		}
-	*/
-
+	// By default we always prefer UCX rather than openib
+	extraArgs = append(extraArgs, "--mca")
+	extraArgs = append(extraArgs, "btl")
+	extraArgs = append(extraArgs, "^openib")
+	extraArgs = append(extraArgs, "--mca")
+	extraArgs = append(extraArgs, "pml")
+	extraArgs = append(extraArgs, "ucx")
 	return extraArgs
 }
+
+func parseOmpiInfoOutputForVersion(output string) (string, error) {
+	lines := strings.Split(output, "\n")
+	if !strings.HasPrefix(lines[0], "Open MPI") {
+		return "", fmt.Errorf("invalid output format")
+	}
+	version := strings.TrimLeft(lines[0], "Open MPI v")
+	version = strings.TrimRight(version, "\n")
+	return version, nil
+}
+
+// DetectFromDir tries to figure out which version of OpenMPI is installed in a given directory
+func DetectFromDir(dir 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)
+	}
+
+	var versionCmd advexec.Advcmd
+	versionCmd.BinPath = targetBin
+	versionCmd.CmdArgs = append(versionCmd.CmdArgs, "--version")
+	res := versionCmd.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 ID, version, nil
+}

pkg/implem/implem.go

@@ -5,12 +5,19 @@
 
 package implem
 
+import (
+	"fmt"
+
+	"github.com/gvallee/go_hpc_jobmgr/internal/pkg/mpich"
+	"github.com/gvallee/go_hpc_jobmgr/internal/pkg/openmpi"
+)
+
 const (
 	// OMPI is the identifier for Open MPI
-	OMPI = "openmpi"
+	OMPI = openmpi.ID
 
 	// MPICH is the identifier for MPICH
-	MPICH = "mpich"
+	MPICH = mpich.ID
 )
 
 // Info gathers all data about a specific MPI implementation
@@ -33,3 +40,25 @@ func IsMPI(i *Info) bool {
 
 	return false
 }
+
+// Load figures out the details about a specific implementation of MPI so it can be used later on.
+// Users have the option to specify:
+// - the install directory, then the function figures out the implementation and version
+// - 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 {
+	if i.InstallDir != "" && (i.ID == "" || i.Version == "") {
+		var err error
+		i.ID, i.Version, err = openmpi.DetectFromDir(i.InstallDir)
+		if err == nil {
+			return nil
+		}
+		i.ID, i.Version, err = mpich.DetectFromDir(i.InstallDir)
+		if err == nil {
+			return nil
+		}
+		return fmt.Errorf("unable to detect MPI implementation from %s", i.InstallDir)
+	}
+	return nil
+}

pkg/mpi/mpi.go

@@ -86,3 +86,4 @@ func Detect() (*implem.Info, error) {
 
 	return mpiInfo, nil
 }
+