From c41eb9e8b594a3a6ccd973f38edb80554c1e31e1 Mon Sep 17 00:00:00 2001
From: Ken Raffenetti <raffenet@mcs.anl.gov>
Date: Fri, 11 Aug 2023 10:23:37 -0500
Subject: [PATCH] Add MPICH HPC environment detection

Default MPICH builds use the Hydra process manager (mpiexec) which sets
PMI_RANK in the application environment. Update GetUniquePathFromEnv()
test accordingly.

Signed-off-by: Ken Raffenetti <raffenet@mcs.anl.gov>
---
 src/base/sysinfo.cc               |  6 ++++++
 src/tests/unique_path_unittest.cc | 27 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/src/base/sysinfo.cc b/src/base/sysinfo.cc
index 022f607..a8c1cb6 100644
--- a/src/base/sysinfo.cc
+++ b/src/base/sysinfo.cc
@@ -249,6 +249,12 @@ static std::tuple<bool, const char*, const char*> QueryHPCEnvironment() {
     return {true, "", ""};
   }
 
+  // Check for Hydra process manager (MPICH)
+  envval = getenv("PMI_RANK");
+  if (envval != nullptr && *envval != 0) {
+    return {false, ".rank-", envval};
+  }
+
   // No HPC environment was detected
   return {false, "", ""};
 }
diff --git a/src/tests/unique_path_unittest.cc b/src/tests/unique_path_unittest.cc
index 4b3788c..eb65257 100644
--- a/src/tests/unique_path_unittest.cc
+++ b/src/tests/unique_path_unittest.cc
@@ -181,11 +181,38 @@ void testOMPI() {
   EXPECT_EQ(expectedPath, GetTestPath());
 }
 
+// MPICH is another type of MPI environment that we detect. We
+// expect .rank-${PMI_RANK} to be appended when we detect this
+// environment.
+void testMPICH() {
+  WithEnv rank("PMI_RANK", "5");
+  WithEnv withTestVar(TEST_VAR, TEST_VAL);
+
+  const auto expectedParent = TEST_VAL ".rank-5";
+  const auto expectedChild = AppendPID(expectedParent);
+
+  // Test parent case (will set the child flag)
+  EXPECT_EQ(expectedParent, GetTestPath());
+
+  // Test child case
+  EXPECT_EQ(expectedChild, GetTestPath());
+
+  withTestVar.Reset();
+  WithEnv withForced(TEST_VAR "_USE_PID", "1");
+
+  // Test parent case (will set the child flag)
+  EXPECT_EQ(expectedChild, GetTestPath());
+
+  // Test child case
+  EXPECT_EQ(expectedChild, GetTestPath());
+}
+
 int main(int argc, char** argv) {
   testDefault();
   testPMIx();
   testSlurm();
   testOMPI();
+  testMPICH();
 
   printf("PASS\n");
   return 0;