PML/UCX: simple mpi test hangs if using UCX PML in singleton mode

[hppritcha]

The UCX PML doesn't work in singleton mode.

To reproduce do these steps

• Build Open MPI main with UCX support.
• Build hello_c in the examples folder.
• run hello_c in singleton mode:

the process will hang in MPI_Finalize.

traceback shows:

#0  0x00007ffff6d8365a in opal_common_ucx_mca_pmix_fence (worker=0x87e310) at common_ucx.c:451
#1  0x00007ffff6d83a3c in opal_common_ucx_del_procs (procs=0x7fffec000e70, count=1, my_rank=0, max_disconnect=1, worker=0x87e310) at common_ucx.c:531
#2  0x00007ffff79bd10a in mca_pml_ucx_del_procs (procs=0x7fffec000ff0, nprocs=1) at pml_ucx.c:567
#3  0x00007ffff76e8d66 in ompi_mpi_instance_cleanup_pml () at instance/instance.c:189
#4  0x00007ffff6d0f566 in opal_finalize_cleanup_domain (domain=0x7ffff70480e0 <opal_init_domain>) at runtime/opal_finalize_core.c:129
#5  0x00007ffff6cfe757 in opal_finalize () at runtime/opal_finalize.c:56
#6  0x00007ffff76e5886 in ompi_rte_finalize () at runtime/ompi_rte.c:1045
#7  0x00007ffff76eb6e1 in ompi_mpi_instance_finalize_common () at instance/instance.c:951
#8  0x00007ffff76eb958 in ompi_mpi_instance_finalize (instance=0x7ffff7d93ac0 <ompi_mpi_instance_default>) at instance/instance.c:996
#9  0x00007ffff76dfd36 in ompi_mpi_finalize () at runtime/ompi_mpi_finalize.c:294
#10 0x00007ffff7731f1e in PMPI_Finalize () at finalize_generated.c:53
#11 0x00000000004008b8 in main (argc=1, argv=0x7ffffffefc08) at hello_c.c:23

This problem can be fixed by patching the UCX PML for the singleton case. Alternatively, the PMIx_Fence_nb implementation can be patched to invoke the cbfunc (if non null) in the quick exit for singleton case:

     /* if we are a singleton, there is nothing to do */
     if (pmix_client_globals.singleton) {
-        return PMIX_SUCCESS;
+        rc = PMIX_SUCCESS;
+        if(NULL != cbfunc) {
+            cbfunc(rc, cbdata);
+        }
+        return rc;
     }

[rhc54]

What version of PMIx are you looking at? Technically, it should not return success unless it is going to execute the cbfunc (but must do so after returning), but should return PMIX_OPERATION_SUCCEEDED to indicate that it was locally completed. I can fix that upstream, though you may need a corresponding change to OMPI to handle the "operation succeeded" return.

Alternatively, since OMPI has an "I am a singleton" flag - you could just have OMPI not call fence when a singleton. IIRC, that may be done elsewhere.

[hppritcha]

open pmix at d30c15e65f3e6f154807424dd428698fe9fbad3a

[rhc54]

Okay, back a little ways but the code in that section hasn't changed. I'll fix it so that it conforms to the Standard. I checked the doc and confirmed that my prior description of the behavior is correct.

[rhc54]

I've got the fix ready, but it occurs to me that it might break OMPI where you use Fence_nb. Per the Standard, Fence_nb is to return PMIX_OPERATION_SUCCEEDED when the proc is a singleton, or

indicating that the request was immediately processed and returned
success -the cbfunc will not be called. This can occur if the collective
involved only processes on the local node

You want to check and see where it is used to make sure OMPI's ready?

[hppritcha]

I patched my version of openpmix in my work area and made the relevant change in ucx pml and it works now.

[rhc54]

Understood - but my change is going to wind up in the 5.0.10 release which hits OMPI v5, and in v6.1 which hits OMPI v6. So we need to see how much impact it has.

[rhc54]

FWIW: there are several places in OMPI that do not follow the PMIx Standard. I'm working on a patch now. Does raise some questions about what to do with already-released OMPI v5 versions as they will fail with the PMIx revisions. We previously ran into this with the Slurm folks (who had a Standard violation on the server side) - required some "news" explanations of version compatibility.

I'll attend the Monday RM meeting to discuss. Meantime, I'll post the patch in a bit.

[rhc54]

@janjust The opal_common_ucx_mca_pmix_fence_nb function in common_ucx.c doesn't actually do anything. It calls PMIx_Fence_nb, but then doesn't wait for the fence to complete. In fact, nowhere does anyone check if the fence has completed - which means there really is no fence being done. No sync is accomplished.

Can you confirm that you actually want this fence to act as a fence? It is only used by oshmem/mca/spml/ucx/spml_ucx_component.c. You could eliminate the nb version and change the spml component to call the blocking version (opal_common_ucx_mca_pmix_fence). Or you could make it a true non-blocking version and have someplace wait for the callback to be executed. Or you could just eliminate the call altogether - as I said, right now it actually isn't doing anything.

[janjust]

@gleon99 I would suggest we just call the blocking version, do we need it non blocking?

[janjust]

hmm wait, so I think we need a non-blocking fence indeed, before we can enter end point tear down on any active worker, we need to make sure they progress and are flushed, so indeed, we need a non-blocking fence,

[rhc54]

Makes sense - though you might want to look at your blocking version. It calls PMIx_Fence_nb and then cycles progress until it completes. Might be what you need.

[janjust]

diff --git a/opal/mca/common/ucx/common_ucx.c b/opal/mca/common/ucx/common_ucx.c
index 62a4703870..365dd91cdc 100644
--- a/opal/mca/common/ucx/common_ucx.c
+++ b/opal/mca/common/ucx/common_ucx.c
@@ -429,11 +429,6 @@ void opal_common_ucx_mca_proc_added(void)
 #endif
 }
 
-OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence_nb(int *fenced)
-{
-    return PMIx_Fence_nb(NULL, 0, NULL, 0, opal_common_ucx_mca_fence_complete_cb, (void *) fenced);
-}
-
 OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence(ucp_worker_h worker)
 {
     volatile int fenced = 0;
diff --git a/opal/mca/common/ucx/common_ucx.h b/opal/mca/common/ucx/common_ucx.h
index 13e6989186..78dc3ff62e 100644
--- a/opal/mca/common/ucx/common_ucx.h
+++ b/opal/mca/common/ucx/common_ucx.h
@@ -120,7 +120,6 @@ OPAL_DECLSPEC opal_common_ucx_support_level_t opal_common_ucx_support_level(ucp_
 OPAL_DECLSPEC void opal_common_ucx_mca_proc_added(void);
 OPAL_DECLSPEC void opal_common_ucx_empty_complete_cb(void *request, ucs_status_t status);
 OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence(ucp_worker_h worker);
-OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence_nb(int *fenced);
 OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, size_t count,
                                             size_t my_rank, size_t max_disconnect,
                                             ucp_worker_h worker);
diff --git a/oshmem/mca/spml/ucx/spml_ucx_component.c b/oshmem/mca/spml/ucx/spml_ucx_component.c
index 66c8863008..3eef24268b 100644
--- a/oshmem/mca/spml/ucx/spml_ucx_component.c
+++ b/oshmem/mca/spml/ucx/spml_ucx_component.c
@@ -25,6 +25,7 @@
 
 #include "opal/util/opal_environ.h"
 #include "opal/runtime/opal_progress_threads.h"
+#include "opal/mca/pmix/pmix-internal.h"
 
 static int mca_spml_ucx_component_register(void);
 static int mca_spml_ucx_component_open(void);
@@ -464,9 +465,13 @@ static void mca_spml_ucx_ctx_fini(mca_spml_ucx_ctx_t *ctx)
     }
 }
 
+static void mca_spml_ucx_fence_complete_cb(int status, void *fenced)
+{
+    *(volatile int *) fenced = 1;
+}
+
 static int mca_spml_ucx_component_fini(void)
 {
-    volatile int fenced = 0;
     int i;
     int ret = OSHMEM_SUCCESS;
 
@@ -497,29 +502,41 @@ static int mca_spml_ucx_component_fini(void)
         _ctx_cleanup(mca_spml_ucx.idle_array.ctxs[i]);
     }
 
+    /* Start non-blocking fence - we need to progress all workers before teardown */
+    {
+        volatile int fenced = 0;
+        ret = PMIx_Fence_nb(NULL, 0, NULL, 0, mca_spml_ucx_fence_complete_cb, (void *) &fenced);
+        if (ret != PMIX_SUCCESS) {
+            SPML_UCX_WARN("pmix fence failed: %s", PMIx_Error_string(ret));
+            /* In case of pmix fence failure just continue cleanup */
+            fenced = 1;
+        }
 
-    ret = opal_common_ucx_mca_pmix_fence_nb(&fenced);
-    if (ret != PMIX_SUCCESS) {
-        SPML_UCX_WARN("pmix fence failed: %s", PMIx_Error_string(ret));
-        /* In case of pmix fence failure just continue cleanup */
-        fenced = 1;
-    }
+        /* Progress all workers to flush pending operations and call opal_progress 
+         * to allow PMIx callback to fire */
+        while (!fenced) {
+            /* Progress all active context workers */
+            for (i = 0; i < mca_spml_ucx.active_array.ctxs_count; i++) {
+                ucp_worker_progress(mca_spml_ucx.active_array.ctxs[i]->ucp_worker[0]);
+            }
 
-    while (!fenced) {
-        for (i = 0; i < mca_spml_ucx.active_array.ctxs_count; i++) {
-            ucp_worker_progress(mca_spml_ucx.active_array.ctxs[i]->ucp_worker[0]);
-        }
+            /* Progress all idle context workers */
+            for (i = 0; i < mca_spml_ucx.idle_array.ctxs_count; i++) {
+                ucp_worker_progress(mca_spml_ucx.idle_array.ctxs[i]->ucp_worker[0]);
+            }
+            
+            /* Progress all default context workers */
+            for (i = 0; i < (signed int)mca_spml_ucx.ucp_workers; i++) {
+                ucp_worker_progress(mca_spml_ucx_ctx_default.ucp_worker[i]);
+            }
 
-        for (i = 0; i < mca_spml_ucx.idle_array.ctxs_count; i++) {
-            ucp_worker_progress(mca_spml_ucx.idle_array.ctxs[i]->ucp_worker[0]);
-        }
-        
-        for (i = 0; i < (signed int)mca_spml_ucx.ucp_workers; i++) {
-            ucp_worker_progress(mca_spml_ucx_ctx_default.ucp_worker[i]);
-        }
+            /* Progress auxiliary context worker if it exists */
+            if (mca_spml_ucx.aux_ctx != NULL) {
+                ucp_worker_progress(mca_spml_ucx.aux_ctx->ucp_worker[0]);
+            }
 
-        if (mca_spml_ucx.aux_ctx != NULL) {
-            ucp_worker_progress(mca_spml_ucx.aux_ctx->ucp_worker[0]);
+            /* Call opal_progress() to allow PMIx event library to cycle and fire callback */
+            opal_progress();
         }
     }