Cleanup attachment to scheduler

If a scheduler request is received and we are not already attached to the scheduler, then use PMIx_tool_attach_to_server to setup the connection. Note that the function will simply return if we are already attached, but PRRTE doesn't know about it yet. Signed-off-by: Ralph Castain <[email protected]>
openpmix · Dec 19, 2023 · 25c2ca4 · 25c2ca4
1 parent 6c5318a
commit 25c2ca4
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 4 deletions.
diff --git a/src/prted/pmix/pmix_server_gen.c b/src/prted/pmix/pmix_server_gen.c
@@ -525,6 +525,9 @@ static void _toolconn(int sd, short args, void *cbdata)
     pmix_data_buffer_t *buf;
     prte_plm_cmd_flag_t command = PRTE_PLM_ALLOC_JOBID_CMD;
     pmix_status_t xrc;
+    bool primary = false;
+    bool nspace_given = false;
+    bool rank_given = false;
     PRTE_HIDE_UNUSED_PARAMS(sd, args);
 
     PMIX_ACQUIRE_OBJECT(cd);
@@ -560,8 +563,10 @@ static void _toolconn(int sd, short args, void *cbdata)
                 }
             } else if (PMIX_CHECK_KEY(&cd->info[n], PMIX_NSPACE)) {
                 PMIX_LOAD_NSPACE(cd->target.nspace, cd->info[n].value.data.string);
+                nspace_given = true;
             } else if (PMIX_CHECK_KEY(&cd->info[n], PMIX_RANK)) {
                 cd->target.rank = cd->info[n].value.data.rank;
+                rank_given = true;
             } else if (PMIX_CHECK_KEY(&cd->info[n], PMIX_HOSTNAME)) {
                 cd->operation = strdup(cd->info[n].value.data.string);
             } else if (PMIX_CHECK_KEY(&cd->info[n], PMIX_CMD_LINE)) {
@@ -570,6 +575,8 @@ static void _toolconn(int sd, short args, void *cbdata)
                 cd->launcher = PMIX_INFO_TRUE(&cd->info[n]);
             } else if (PMIX_CHECK_KEY(&cd->info[n], PMIX_SERVER_SCHEDULER)) {
                 cd->scheduler = PMIX_INFO_TRUE(&cd->info[n]);
+            } else if (PMIX_CHECK_KEY(&cd->info[n], PMIX_PRIMARY_SERVER)) {
+                primary = PMIX_INFO_TRUE(&cd->info[n]);
             } else if (PMIX_CHECK_KEY(&cd->info[n], PMIX_PROC_PID)) {
                 PMIX_VALUE_GET_NUMBER(xrc, &cd->info[n].value, cd->pid, pid_t);
                 if (PMIX_SUCCESS != xrc) {
@@ -599,11 +606,25 @@ static void _toolconn(int sd, short args, void *cbdata)
         } else {
             /* mark that the scheduler has attached to us */
             prte_pmix_server_globals.scheduler_connected = true;
+            // the scheduler always self-assigns its ID
+            if (!nspace_given || !rank_given) {
+                cd->toolcbfunc(PMIX_ERR_NOT_SUPPORTED, NULL, cd->cbdata);
+                PMIX_RELEASE(cd);
+                return;
+            }
             PMIX_LOAD_PROCID(&prte_pmix_server_globals.scheduler,
                              cd->target.nspace, cd->target.rank);
-            /* we cannot immediately set the scheduler to be our
-             * PMIx server as the PMIx library hasn't finished
-             * recording it */
+            rc = PMIX_SUCCESS;
+
+            if (!primary) {
+                /* we cannot immediately set the scheduler to be our
+                 * PMIx server as the PMIx library hasn't finished
+                 * recording it */
+                goto complete;
+            }
+            // it has been recorded in the library, so record it here
+            prte_pmix_server_globals.scheduler_set_as_server = true;
+            goto complete;
         }
     }
 
@@ -652,6 +673,8 @@ static void _toolconn(int sd, short args, void *cbdata)
     if (PMIX_SUCCESS != rc) {
         rc = prte_pmix_convert_rc(rc);
     }
+
+complete:
     if (NULL != cd->toolcbfunc) {
         cd->toolcbfunc(rc, &cd->target, cd->cbdata);
     }

diff --git a/src/tools/psched/psched.h b/src/tools/psched/psched.h
@@ -129,6 +129,8 @@ typedef struct {
     // allocation request info
     pmix_proc_t requestor;
     pmix_alloc_directive_t directive;
+    // whether the data is a local copy
+    bool copy;
     // original info keys
     pmix_info_t *data;
     size_t ndata;

diff --git a/src/tools/psched/state.c b/src/tools/psched/state.c
@@ -315,6 +315,7 @@ PMIX_CLASS_INSTANCE(psched_state_t,
 static void req_con(psched_req_t *p)
 {
     PMIx_Load_procid(&p->requestor, NULL, PMIX_RANK_INVALID);
+    p->copy = false;  // data is not a local copy
     p->data = NULL;
     p->ndata = 0;
     p->user_refid = NULL;
@@ -341,7 +342,7 @@ static void req_con(psched_req_t *p)
 }
 static void req_des(psched_req_t *p)
 {
-    if (NULL != p->data) {
+    if (NULL != p->data && p->copy) {
         PMIx_Info_free(p->data, p->ndata);
     }
     if (NULL != p->user_refid) {