support boot from MPI for debug

srcs/cpp/src/peer.cpp

@@ -10,9 +10,18 @@ Peer::Peer()
         fprintf(stderr, "%s failed\n", "GoKungfuInit");
         exit(1);
     }
+    if (std::getenv("KUNGFU_DEBUG_BOOT_FROM_MPI")) {
+        // TODO: call MPI_Init();
+    }
 }
 
-Peer::~Peer() { GoKungfuFinalize(); }
+Peer::~Peer()
+{
+    GoKungfuFinalize();
+    if (std::getenv("KUNGFU_DEBUG_BOOT_FROM_MPI")) {
+        // TODO: call MPI_Finalize();
+    }
+}
 
 bool Peer::Detached() const { return GoKungfuDetached(); }
 

srcs/cpp/src/python/cuda.cpp

@@ -29,6 +29,11 @@ std::vector<int> parse_cuda_visible_devices(const std::string &val)
 int kungfu_get_cuda_index()
 {
     int dev = 0;
+    if (std::getenv("KUNGFU_DEBUG_BOOT_FROM_MPI")) {
+        const char *ptr = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK");
+        if (ptr != nullptr) { dev = std::stoi(ptr); }
+        return dev;
+    }
     {
         const char *ptr = std::getenv("KUNGFU_CUDA_VISIBLE_DEVICES");
         if (ptr != nullptr) { dev = std::stoi(ptr); }

srcs/go/kungfu/env/config.go

@@ -22,6 +22,9 @@ type Config struct {
 }
 
 func ParseConfigFromEnv() (*Config, error) {
+	if _, ok := os.LookupEnv(BootFromMPI); ok {
+		return ParseConfigFromOpenMPIEnv()
+	}
 	if _, ok := os.LookupEnv(SelfSpecEnvKey); !ok {
 		return singleEnv(), nil
 	}

srcs/go/kungfu/env/envs.go

@@ -15,4 +15,6 @@ const (
 	ProcStartTimestamp = `KUNGFU_PROC_START_TIMESTAMP`
 
 	AllowNvLink = `KUNGFU_ALLOW_NVLINK`
+
+	BootFromMPI = `KUNGFU_DEBUG_BOOT_FROM_MPI`
 )

srcs/go/kungfu/env/mpi.go

@@ -0,0 +1,33 @@
+package env
+
+import (
+	"os"
+	"strconv"
+
+	"github.com/lsds/KungFu/srcs/go/kungfu/base"
+	"github.com/lsds/KungFu/srcs/go/plan"
+	"github.com/lsds/KungFu/srcs/go/utils/assert"
+)
+
+// ParseConfigFromOpenMPIEnv is for debug only
+func ParseConfigFromOpenMPIEnv() (*Config, error) {
+	mpiSize, err := strconv.Atoi(os.Getenv(`OMPI_COMM_WORLD_SIZE`))
+	if err != nil {
+		return nil, err
+	}
+	mpiRank, err := strconv.Atoi(os.Getenv(`OMPI_COMM_WORLD_RANK`))
+	if err != nil {
+		return nil, err
+	}
+	hl := plan.HostList{{
+		IPv4:  plan.MustParseIPv4(`127.0.0.1`),
+		Slots: mpiSize,
+	}}
+	peers, err := hl.GenPeerList(mpiSize, plan.DefaultPortRange)
+	assert.OK(err)
+	return &Config{
+		Self:      peers[mpiRank],
+		InitPeers: peers,
+		Strategy:  base.DefaultStrategy,
+	}, nil
+}