init: use MPIR_CVAR_FINALIZE_ATEXIT to bypass re-init

src/mpi/init/mpir_init.c

@@ -102,6 +102,18 @@
       description : >-
         Skip MPIR_pmi_barrier() in MPI_Init
 
+    - name        : MPIR_CVAR_FINALIZE_ATEXIT
+      category    : COLLECTIVE
+      type        : boolean
+      default     : false
+      class       : none
+      verbosity   : MPI_T_VERBOSITY_USER_BASIC
+      scope       : MPI_T_SCOPE_LOCAL
+      description : >-
+        If true, delay MPI_Finalize until exit via atexit hook. This may be
+        necessary to support applications that uses MPI sessions and require
+        multiple re-init via MPI_Session_init/finalize.
+
 === END_MPI_T_CVAR_INFO_BLOCK ===
 */
 
@@ -148,6 +160,12 @@ int MPIR_Init_impl(int *argc, char ***argv)
     return mpi_errno;
 }
 
+static void finalize_atexit(void)
+{
+#define DUMMY_SESSION_PTR ((void *) 1)
+    MPII_Finalize(DUMMY_SESSION_PTR);
+}
+
 int MPII_Init_thread(int *argc, char ***argv, int user_required, int *provided,
                      MPIR_Session ** p_session_ptr)
 {
@@ -167,6 +185,7 @@ int MPII_Init_thread(int *argc, char ***argv, int user_required, int *provided,
     if (init_counter > 1) {
         goto fn_exit;
     }
+
     /**********************************************************************/
     /* Section 1: base components that other components rely on.
      * These need to be initialized first.  They have strong
@@ -184,6 +203,12 @@ int MPII_Init_thread(int *argc, char ***argv, int user_required, int *provided,
     mpi_errno = MPIR_T_env_init();
     MPIR_ERR_CHECK(mpi_errno);
 
+    if (MPIR_CVAR_FINALIZE_ATEXIT) {
+        init_counter++;
+        err = atexit(finalize_atexit);
+        MPIR_ERR_CHKANDJUMP1(err != 0, mpi_errno, MPI_ERR_OTHER, "**atexit", "**atexit %d", err);
+    }
+
     MPIR_Err_init();
     MPII_pre_init_dbg_logging(argc, argv);
     MPII_pre_init_memory_tracing();

src/pmi/errnames.txt

@@ -133,5 +133,5 @@
 #
 # PMI finalize exit handler registration
 #
-**atexit_pmi_finalize: Registration of PMI finalize function in exit handler failed
-**atexit_pmi_finalize %d: Registration of PMI finalize function in exit handler failed with %d
+**atexit: Registration of atexit handler failed
+**atexit %d: Registration of atexit handler failed with %d

src/util/mpir_pmi.c

@@ -217,10 +217,6 @@ int MPIR_pmi_init(void)
 {
     int mpi_errno = MPI_SUCCESS;
     static bool pmi_connected = false;
-    static int init_count = 0;
-
-    /* track init_count to differentiate re-init world_id */
-    init_count++;
 
     if (finalize_pending > 0) {
         finalize_pending--;
@@ -245,19 +241,20 @@ int MPIR_pmi_init(void)
 
     unsigned world_id = 0;
     if (pmi_kvs_name) {
-        char buf[1024];
-        snprintf(buf, 1024, "%s-%d", pmi_kvs_name, init_count);
-        HASH_FNV(buf, strlen(buf), world_id);
+        if (!strcmp(pmi_kvs_name, "singinit") || !strcmp(pmi_kvs_name, "0")) {
+            world_id = getpid();
+        } else {
+            HASH_FNV(pmi_kvs_name, strlen(pmi_kvs_name), world_id);
+        }
     }
 
-    if (!pmi_connected) {
+    if (!MPIR_CVAR_FINALIZE_ATEXIT && !pmi_connected) {
         /* Register finalization of PM connection in exit handler */
         mpi_errno = atexit(MPIR_pmi_finalize_on_exit);
         MPIR_ERR_CHKANDJUMP1(mpi_errno != 0, mpi_errno, MPI_ERR_OTHER,
-                             "**atexit_pmi_finalize", "**atexit_pmi_finalize %d", mpi_errno);
-
-        pmi_connected = true;
+                             "**atexit", "**atexit %d", mpi_errno);
     }
+    pmi_connected = true;
 
     int world_idx = MPIR_add_world(pmi_kvs_name, size);
     MPIR_Assertp(world_idx == 0);
@@ -331,8 +328,12 @@ void MPIR_pmi_finalize(void)
     free_hwloc_topology();
 #endif
 
-    /* delay PMI_Finalize to the exit hook */
-    finalize_pending++;
+    if (MPIR_CVAR_FINALIZE_ATEXIT) {
+        SWITCH_PMI(pmi1_exit(), pmi2_exit(), pmix_exit());
+    } else {
+        /* delay PMI_Finalize to the exit hook */
+        finalize_pending++;
+    }
 }
 
 void MPIR_pmi_abort(int exit_code, const char *error_msg)

src/util/mpir_pmix.inc

@@ -23,21 +23,6 @@ static pmix_proc_t pmix_proc;
 static pmix_proc_t pmix_wcproc;
 static pmix_proc_t pmix_parent;
 
-/* PMIx does not handle duplicate keys reliably. In a session re-init case, the 2nd round
- * business card exchange may see previous values due to key collisions.
- */
-/* Use pmix_init_count to track and differentiate each re-init epoch.
- * 1. skip PMIx_Init if pmix_init_count > 1
- * 2. prefix kvs keys with pmix_init_count to avoid key collisions.
- */
-static int pmix_init_count = 0;
-static const char *prefix_key(const char *key)
-{
-    static char buf[PMIX_MAX_KEYLEN];
-    snprintf(buf, PMIX_MAX_KEYLEN, "%s-%d", key, pmix_init_count);
-    return buf;
-}
-
 static void pmix_not_supported(const char *elem, char *error_str, int len);
 static int pmix_add_to_info(MPIR_Info * info_ptr, const char *key, const char *pmix_key,
                             MPIR_Info * target_ptr, int *key_found, size_t * counter, char **value);
@@ -69,6 +54,7 @@ static int pmix_init(int *has_parent, int *rank, int *size, int *appnum)
 
     /* Since we only call PMIx_Finalize once at `atexit` handler, we need prevent
      * calling PMIx_Init multiple times. */
+    static int pmix_init_count = 0;
     pmix_init_count++;
     if (pmix_init_count == 1) {
         pmi_errno = PMIx_Init(&pmix_proc, NULL, 0);
@@ -194,7 +180,7 @@ static int pmix_put(const char *key, const char *val)
     pmix_value_t value;
     value.type = PMIX_STRING;
     value.data.string = (char *) val;
-    pmi_errno = PMIx_Put(PMIX_GLOBAL, prefix_key(key), &value);
+    pmi_errno = PMIx_Put(PMIX_GLOBAL, key, &value);
     MPIR_ERR_CHKANDJUMP1(pmi_errno != PMIX_SUCCESS, mpi_errno, MPI_ERR_OTHER,
                          "**pmix_put", "**pmix_put %d", pmi_errno);
     pmi_errno = PMIx_Commit();
@@ -213,13 +199,13 @@ static int pmix_get(int src, const char *key, char *val, int val_size)
 
     pmix_value_t *pvalue;
     if (src < 0) {
-        pmi_errno = PMIx_Get(&pmix_wcproc, prefix_key(key), NULL, 0, &pvalue);
+        pmi_errno = PMIx_Get(&pmix_wcproc, key, NULL, 0, &pvalue);
     } else {
         pmix_proc_t proc;
         PMIX_PROC_CONSTRUCT(&proc);
         proc.rank = src;
 
-        pmi_errno = PMIx_Get(&proc, prefix_key(key), NULL, 0, &pvalue);
+        pmi_errno = PMIx_Get(&proc, key, NULL, 0, &pvalue);
     }
     MPIR_ERR_CHKANDJUMP1(pmi_errno != PMIX_SUCCESS, mpi_errno, MPI_ERR_OTHER,
                          "**pmix_get", "**pmix_get %d", pmi_errno);
@@ -250,8 +236,6 @@ static bool pmix_get_jobattr(const char *key, char *valbuf)
     /* translate MPICH key to PMIx standard format */
     if (strcmp(key, "PMI_process_mapping") == 0) {
         key = PMIX_ANL_MAP;
-    } else {
-        goto fn_exit;
     }
 
     /* if this is a non-reserved key, we want to make sure not to block
@@ -269,7 +253,6 @@ static bool pmix_get_jobattr(const char *key, char *valbuf)
     }
     PMIX_INFO_FREE(info, 1);
 
-  fn_exit:
     return found;
 }
 
@@ -400,7 +383,7 @@ static int pmix_optimized_put(const char *key, const char *val, int is_local)
     pmix_value_t value;
     value.type = PMIX_STRING;
     value.data.string = (char *) val;
-    pmi_errno = PMIx_Put(is_local ? PMIX_LOCAL : PMIX_GLOBAL, prefix_key(key), &value);
+    pmi_errno = PMIx_Put(is_local ? PMIX_LOCAL : PMIX_GLOBAL, key, &value);
     MPIR_ERR_CHKANDJUMP1(pmi_errno != PMIX_SUCCESS, mpi_errno, MPI_ERR_OTHER,
                          "**pmix_put", "**pmix_put %d", pmi_errno);
     pmi_errno = PMIx_Commit();
@@ -426,7 +409,7 @@ static int pmix_put_binary(const char *key, const char *buf, int bufsize, int is
     value.type = PMIX_BYTE_OBJECT;
     value.data.bo.bytes = (char *) buf;
     value.data.bo.size = bufsize;
-    pmi_errno = PMIx_Put(is_local ? PMIX_LOCAL : PMIX_GLOBAL, prefix_key(key), &value);
+    pmi_errno = PMIx_Put(is_local ? PMIX_LOCAL : PMIX_GLOBAL, key, &value);
     MPIR_ERR_CHKANDJUMP1(pmi_errno != PMIX_SUCCESS, mpi_errno, MPI_ERR_OTHER,
                          "**pmix_put", "**pmix_put %d", pmi_errno);
     pmi_errno = PMIx_Commit();
@@ -447,13 +430,13 @@ static int pmix_get_binary(int src, const char *key, char *buf, int *p_size, int
     int bufsize ATTRIBUTE((unused)) = *p_size;
     pmix_value_t *pvalue;
     if (src < 0) {
-        pmi_errno = PMIx_Get(&pmix_wcproc, prefix_key(key), NULL, 0, &pvalue);
+        pmi_errno = PMIx_Get(&pmix_wcproc, key, NULL, 0, &pvalue);
     } else {
         pmix_proc_t proc;
         PMIX_PROC_CONSTRUCT(&proc);
         proc.rank = src;
 
-        pmi_errno = PMIx_Get(&proc, prefix_key(key), NULL, 0, &pvalue);
+        pmi_errno = PMIx_Get(&proc, key, NULL, 0, &pvalue);
     }
     MPIR_ERR_CHKANDJUMP1(pmi_errno != PMIX_SUCCESS, mpi_errno, MPI_ERR_OTHER,
                          "**pmix_get", "**pmix_get %d", pmi_errno);

test/mpi/session/testlist

@@ -1,13 +1,13 @@
 session 4
-session_mult_init 4
-session_mult_init 4 arg=5
+session_mult_init 4 -env=MPIR_CVAR_FINALIZE_AT_EXIT=1
+session_mult_init 4 arg=5 -env=MPIR_CVAR_FINALIZE_AT_EXIT=1
 session_mult_init_with_world 4
 session_mult_init_with_world 4 arg=5
-session_re_init 4
+session_re_init 4 -env=MPIR_CVAR_FINALIZE_AT_EXIT=1
 session_mod_group 2
 session_mod_group 4
-session_mod_group_re_init 2
-session_mod_group_re_init 4
+session_mod_group_re_init 2 -env=MPIR_CVAR_FINALIZE_AT_EXIT=1
+session_mod_group_re_init 4 -env=MPIR_CVAR_FINALIZE_AT_EXIT=1
 session_psets 1
 session_psets 4
 session_self 1