diff --git a/src/mpi/init/mpir_init.c b/src/mpi/init/mpir_init.c index 4a2ce338c4c..55e2dcf32be 100644 --- a/src/mpi/init/mpir_init.c +++ b/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(); diff --git a/src/pmi/errnames.txt b/src/pmi/errnames.txt index aa75c6668c8..0514174c901 100644 --- a/src/pmi/errnames.txt +++ b/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 diff --git a/src/util/mpir_pmi.c b/src/util/mpir_pmi.c index 540ae93e6a4..e33176a20a6 100644 --- a/src/util/mpir_pmi.c +++ b/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) diff --git a/src/util/mpir_pmix.inc b/src/util/mpir_pmix.inc index d2255a7a175..9512a15eb34 100644 --- a/src/util/mpir_pmix.inc +++ b/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); diff --git a/test/mpi/session/testlist b/test/mpi/session/testlist index 6352e3bdcc1..54cc45d9d50 100644 --- a/test/mpi/session/testlist +++ b/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