slurm: fix session-stack lifetime and other plm/ras correctness bugs

.gitignore

@@ -190,6 +190,7 @@ examples/dynamic-dep
 examples/abort
 examples/elastic
 examples/sessionctrl
+examples/slurm/elastic-coordination
 
 src/sys/powerpc/atomic-32.s
 src/sys/powerpc/atomic-64.s

examples/slurm/README

@@ -10,8 +10,8 @@ elastic-coordination.c:
 A Slurm-specific PMIx client for exercising dynamic allocation changes. It
 accepts interactive requests to extend by node count and to shrink by node
 count, allocation ID, or node list. After every request it queries PRRTE's
-allocation state, displays the corresponding Slurm jobs, and uses prun to
-verify reachability across the active DVM.
+allocation state and displays the corresponding Slurm jobs; a separate `dvm`
+command uses prun to verify reachability across the active DVM.
 
 The client treats the PMIx allocation callback as request acceptance and uses
 allocation queries to observe the resulting state. It does not require the
@@ -31,8 +31,7 @@ Then run inside a Slurm allocation with one initial node:
 
     export PRTE_MCA_prte_elastic_mode=1
     salloc -N1
-    srun -N1 -n1 \
-        prterun -np 1 $HOME/slurm-elastic-coordination
+    prterun -np 1 $HOME/slurm-elastic-coordination
 
 At the coord prompt, available commands are:
 
@@ -41,9 +40,15 @@ At the coord prompt, available commands are:
     shrink alloc SLURM_JOB_ID
     shrink list NODE1,NODE2
     prun [ARGS...]
+    dvm
     refresh
     quit
 
+`refresh` re-queries and prints PRRTE's allocation state. `dvm` does the
+same and then also runs a `prun`-based reachability probe (one process
+per node, echoing back its hostname) to confirm every daemon in the DVM
+is actually reachable.
+
 The program invokes squeue and prun for diagnostics, so both commands must be
 available in PATH.
 

examples/slurm/elastic-coordination.c

@@ -672,15 +672,27 @@ static void show_state(void)
     if (0 == query_snapshot(&snap)) {
         trace_point("show_state: printing PRRTE snapshot");
         print_snapshot(&snap);
-        trace_point("show_state: running reachability probe");
-        run_prun_hostname(&snap);
         snapshot_free(&snap);
     } else {
         printf("\nPRRTE allocation query failed\n");
     }
     trace_point("show_state: end");
 }
 
+static void handle_dvm(void)
+{
+    alloc_snapshot_t snap;
+
+    trace_point("dvm: querying PRRTE snapshot");
+    if (0 == query_snapshot(&snap)) {
+        trace_point("dvm: running reachability probe");
+        run_prun_hostname(&snap);
+        snapshot_free(&snap);
+    } else {
+        printf("\nPRRTE allocation query failed\n");
+    }
+}
+
 static pmix_status_t submit_alloc_request(pmix_alloc_directive_t directive,
                                           pmix_info_t *info, size_t ninfo,
                                           const char *description)
@@ -957,6 +969,7 @@ static void usage(const char *argv0)
     printf("  shrink alloc SLURM_JOB_ID\n");
     printf("  shrink list NODE1,NODE2\n");
     printf("  prun [ARGS...]\n");
+    printf("  dvm\n");
     printf("  refresh\n");
     printf("  quit\n");
 }
@@ -1042,6 +1055,10 @@ int main(int argc, char **argv)
             show_state();
             continue;
         }
+        if (0 == strcmp(cmd, "dvm")) {
+            handle_dvm();
+            continue;
+        }
         if (0 == strcmp(cmd, "extend")) {
             uint64_t count;
 

src/mca/plm/slurm/AGENTS.md

@@ -158,14 +158,6 @@ the daemons, not srun).
 - **Environment purge in the child is mandatory** — SLURM forwards the
   full environment; leaving `PMIX_`/`PRTE_` vars in breaks tool
   connections and duplicates command-line settings.
-- **Known wart: the "are we using the whole allocation?" test compares a
-  count against a capacity.** `map->num_new_daemons < session->nodes->size`
-  reads `size` from a `pmix_pointer_array_t`, which is the *allocated*
-  slot count, not the number of nodes in the session. The effect today is
-  benign (the explicit `--nodes`/`--nodelist` is emitted more often than
-  intended, which srun accepts), but anyone fixing it must re-test the
-  elastic grow path on a real SLURM system — dropping the explicit node
-  list changes which nodes srun picks.
 - Multi-node behavior that does not need SLURM (tree-spawn, throttling,
   the prted command line) is covered by
   [`contrib/dockerswarm`](../../../../contrib/dockerswarm/); the

src/mca/plm/slurm/plm_slurm_module.c

@@ -274,6 +274,7 @@ static void launch_daemons(int fd, short args, void *cbdata)
     prte_state_caddy_t *state = (prte_state_caddy_t *) cbdata;
     uint32_t job_id = UINT32_MAX;
     prte_session_t *session = NULL;
+    int32_t num_session_nodes;
     PRTE_HIDE_UNUSED_PARAMS(fd, args);
 
     PMIX_ACQUIRE_OBJECT(state);
@@ -451,10 +452,17 @@ static void launch_daemons(int fd, short args, void *cbdata)
     pmix_argv_append(&argc, &argv, tmp);
     free(tmp);
 
+    num_session_nodes = 0;
+    for (n = 0; n < session->nodes->size; n++) {
+        if (NULL != pmix_pointer_array_get_item(session->nodes, n)) {
+            ++num_session_nodes;
+        }
+    }
+
     /* if we are using all nodes in the job, then srun doesn't
      * require any further arguments
      */
-    if (map->num_new_daemons < session->nodes->size) {
+    if (map->num_new_daemons < num_session_nodes) {
         pmix_asprintf(&tmp, "--nodes=%lu", (unsigned long) map->num_new_daemons);
         pmix_argv_append(&argc, &argv, tmp);
         free(tmp);

src/mca/ras/slurm/AGENTS.md

@@ -178,13 +178,19 @@ rules are easy to get wrong and the failures are silent leaks.
   `pmix_pointer_array_t` frees the array, never its items — destructing
   it without releasing first leaks one reference per detached node on
   every partial shrink.
-- **`prte_slurm_session_stack` items hold a *borrowed* `prte_session_t
-  *`** (no retain). The stack is the authority for "is this allocation
-  still ours to manage": `release_allocation` consults it from the
-  session destructor to decide whether to `scancel`, which is why the
-  full-release path removes the stack item *before* releasing the
-  session. Keep those two operations in that order, and never leave a
-  stack item pointing at a session someone else may release.
+- **`prte_slurm_session_stack` retains a reference on each session it
+  holds** (released automatically by `prte_session_stack_item_t`'s
+  destructor), independent of the session's own creation reference —
+  removal sites just do `pmix_list_remove_item` + `PMIX_RELEASE(item)`.
+- The creation reference is separate: `rollback_session` and
+  `shrink_complete`'s full-release branch also release the session
+  itself, ending the allocation's lifecycle. `prte_ras_slurm_drain_session_stack`
+  (`finalize()`) releases only what the stack retained and leaves the
+  creation reference alone — releasing it too could double-free if
+  whatever else holds it releases its own copy later.
+- No more reactive `release_allocation` hook: the stack's own reference
+  means `session_des` can't run until the item is already removed, so
+  `scancel` is decided explicitly at each removal site instead.
 - `prte_ras_slurm_rollback_session` relies on `session_des` clearing
   `prte_sessions[index]`, so the released session leaves no dangling
   global entry. It deliberately does not `scancel` — the `complete:`

src/mca/ras/slurm/ras_slurm.h

@@ -72,7 +72,7 @@ int prte_ras_slurm_modify_release_init(void);
 int prte_ras_slurm_modify_release_finalize(void);
 int prte_ras_slurm_serve_release_req(prte_pmix_server_req_t *req);
 void prte_ras_slurm_shrink_complete(prte_shrink_campaign_t *campaign);
-int prte_ras_slurm_release_allocation(prte_session_t *session);
+void prte_ras_slurm_drain_session_stack(void);
 
 /* Common modify extend/release features */
 int prte_ras_slurm_kill_job(const char *slurm_jobid, char *err_msg, size_t err_msg_size);

src/mca/ras/slurm/ras_slurm_modify_cancel.c

@@ -135,13 +135,10 @@ int prte_ras_slurm_add_pending_req(const char *request_id, const char *slurm_job
         return PRTE_ERR_OUT_OF_RESOURCE;
     }
 
-    int pmix_err = pmix_pointer_array_add(&pending_reqs, pending_req);
-
-    if (0 > pmix_err) {
+    if (0 > pmix_pointer_array_add(&pending_reqs, pending_req)) {
         prte_ras_slurm_pending_req_free(pending_req);
-        err = prte_pmix_convert_status(pmix_err);
-        PRTE_ERROR_LOG(err);
-        return err;
+        PRTE_ERROR_LOG(PRTE_ERR_OUT_OF_RESOURCE);
+        return PRTE_ERR_OUT_OF_RESOURCE;
     }
 
     return PRTE_SUCCESS;

src/mca/ras/slurm/ras_slurm_modify_common.c

@@ -19,8 +19,8 @@
 /*
  * Cancel a Slurm job using scancel.
  *
- * If scancel returns an error, the first line of stderr/stdout output is copied
- * into err_msg. On success, err_msg is cleared.
+ * If scancel returns an error, the stderr/stdout output is copied
+ * into err_msg (truncated to buffer size). On success, err_msg is cleared.
  *
  * @param[in]  slurm_jobid   Null-terminated Slurm job ID string to cancel.
  * @param[out] err_msg       Buffer for error message, or NULL.

src/mca/ras/slurm/ras_slurm_modify_extend.c

@@ -705,8 +705,6 @@ static int prte_ras_slurm_add_reused_nodes_to_session(const char *slurm_jobid,
         return PRTE_ERR_BAD_PARAM;
     }
 
-    int err = PRTE_SUCCESS;
-    int pmix_err;
     int added = 0;
     prte_session_t *session;
     prte_session_stack_item_t *item;
@@ -724,11 +722,9 @@ static int prte_ras_slurm_add_reused_nodes_to_session(const char *slurm_jobid,
         }
 
         PMIX_RETAIN(node);
-        pmix_err = pmix_pointer_array_add(session->nodes, node);
-        if (0 > pmix_err) {
+        if (0 > pmix_pointer_array_add(session->nodes, node)) {
             PMIX_RELEASE(node);
-            err = prte_pmix_convert_status(pmix_err);
-            return err;
+            return PRTE_ERR_OUT_OF_RESOURCE;
         }
         added++;
     }

src/mca/ras/slurm/ras_slurm_modify_release.c

@@ -452,43 +452,38 @@ static int prte_ras_slurm_complete_release_request(prte_pmix_server_req_t *req)
 }
 
 /**
- * @brief Release a Slurm-backed session allocation.
+ * @brief Release every session still on the stack at component finalize,
+ * scancel'ing any that PRRTE added dynamically.
  *
- * @param[in] session Session being destroyed.
+ * Only releases the stack's own reference (via PMIX_RELEASE(item)) -
+ * never the session's creation reference, to avoid a double-free if
+ * whatever else holds it releases its own copy later.
  */
-int prte_ras_slurm_release_allocation(prte_session_t *session)
+void prte_ras_slurm_drain_session_stack(void)
 {
+    prte_session_stack_item_t *item, *next;
     char err_msg[PRTE_SLURM_ERR_STR_MAX_LEN + 1] = {0};
     int err;
 
-    if (NULL == session || NULL == session->alloc_refid) {
-        return PRTE_ERR_TAKE_NEXT_OPTION;
-    }
-
-    /* The session stack is the authoritative record of allocations whose
-     * lifetime is still managed by this component.  A normal full release
-     * removes the entry before destroying the session and cancels the copied
-     * Slurm job ID explicitly afterward.  Do not cancel it a second time from
-     * the session destructor. */
-    if (NULL == prte_ras_slurm_find_session_item_by_alloc_id(session->alloc_refid)) {
-        return PRTE_ERR_TAKE_NEXT_OPTION;
-    }
+    PMIX_LIST_FOREACH_SAFE(item, next, prte_slurm_session_stack, prte_session_stack_item_t) {
+        if (NULL == item->session || NULL == item->session->alloc_refid) {
+            continue;
+        }
 
-    if (!prte_ras_slurm_session_is_dynamic(session)) {
-        /* This allocation was not added by PRRTE, so we do not manage its lifetime. */
-        return PRTE_SUCCESS;
-    }
+        if (prte_ras_slurm_session_is_dynamic(item->session)) {
+            err = prte_ras_slurm_kill_job(item->session->alloc_refid, err_msg,
+                                          sizeof(err_msg));
+            if (PRTE_ERR_SLURM_CANCEL_FAILURE == err) {
+                pmix_output(0, "ras:slurm:drain_session_stack: failed to kill job %s: %s.",
+                            item->session->alloc_refid, err_msg);
+            } else if (PRTE_SUCCESS != err) {
+                PRTE_ERROR_LOG(err);
+            }
+        }
 
-    err = prte_ras_slurm_kill_job(session->alloc_refid, err_msg,
-                                  sizeof(err_msg));
-    if (PRTE_ERR_SLURM_CANCEL_FAILURE == err) {
-        pmix_output(0, "ras:slurm:release_allocation: failed to kill job %s: %s.",
-                    session->alloc_refid, err_msg);
-    } else if (PRTE_SUCCESS != err) {
-        PRTE_ERROR_LOG(err);
+        pmix_list_remove_item(prte_slurm_session_stack, &item->super);
+        PMIX_RELEASE(item);
     }
-
-    return err;
 }
 
 /**
@@ -540,13 +535,11 @@ void prte_ras_slurm_shrink_complete(prte_shrink_campaign_t *campaign)
                 continue;
             }
 
-            /* Stop managing the allocation before releasing the session.  Its
-             * destructor calls prte_ras_slurm_release_allocation(), which uses
-             * this stack as the authority for whether it should issue scancel.
-             * Removing the item first therefore lets the session be completely
-             * torn down before the external Slurm job is killed. */
+            /* Stop managing the allocation before killing the external Slurm
+             * job. We own ending this allocation's lifecycle, so release
+             * both the stack's reference (via session_item) and the
+             * session's own creation reference. */
             pmix_list_remove_item(prte_slurm_session_stack, &session_item->super);
-            session_item->session = NULL;
             prte_num_allocated_nodes -= node_count;
             PMIX_RELEASE(session);
             PMIX_RELEASE(session_item);
@@ -1271,6 +1264,14 @@ static int prte_ras_slurm_remove_nodes_by_count(prte_pmix_server_req_t *req, uin
                 nodes_to_rem_from_session = removable_count;
             }
 
+            if (0 >= nodes_to_rem_from_session) {
+                /* Shouldn't happen: removable_count > 0 for any session
+                 * find_releasable_session returns. Guard anyway. */
+                err = PRTE_ERR_NOT_FOUND;
+                PRTE_ERROR_LOG(err);
+                goto cleanup;
+            }
+
             /* Case 2: keep the Slurm job alive and shrink it to its survivors. */
             char *protected_node = NULL;
 

src/mca/ras/slurm/ras_slurm_module.c

@@ -80,7 +80,6 @@ prte_ras_base_module_t prte_ras_slurm_module = {
     .allocate = prte_ras_slurm_allocate,
     .modify = modify,
     .shrink_complete = prte_ras_slurm_shrink_complete,
-    .release_allocation = prte_ras_slurm_release_allocation,
     .finalize = prte_ras_slurm_finalize
 };
 
@@ -91,10 +90,16 @@ static int prte_ras_slurm_parse_range(char *base, char *range, char ***nodelist)
 
 pmix_list_t *prte_slurm_session_stack = NULL;
 
+static void session_stack_item_des(prte_session_stack_item_t *p)
+{
+    if (NULL != p->session) {
+        PMIX_RELEASE(p->session);
+    }
+}
 PMIX_CLASS_INSTANCE(prte_session_stack_item_t,
                     pmix_list_item_t,
                     NULL,
-                    NULL);
+                    session_stack_item_des);
 
 static bool check_taint(char *name, char *evar)
 {
@@ -327,6 +332,7 @@ static int prte_ras_slurm_finalize(void)
     prte_ras_slurm_modify_cancel_finalize();
     prte_ras_slurm_modify_release_finalize();
     if (NULL != prte_slurm_session_stack) {
+        prte_ras_slurm_drain_session_stack();
         PMIX_RELEASE(prte_slurm_session_stack);
         prte_slurm_session_stack = NULL;
     }
@@ -927,6 +933,7 @@ int prte_ras_slurm_assign_new_session(const char *slurm_jobid, const char *user_
         goto cleanup;
     }
 
+    PMIX_RETAIN(session);
     item->session = session;
     item->nodes_in_session = pmix_list_get_size(node_list);
 

src/runtime/prte_globals.c

@@ -1122,8 +1122,11 @@ static void session_des(prte_session_t *s)
         PMIX_RELEASE(s->owner_job);
         s->owner_job = NULL;
     }
-    // remove this from the global array
-    if (0 <= s->index) {
+    // remove this from the global array, if it still exists - a session
+    // whose last reference is released after prte_sessions itself has
+    // already been torn down (e.g. a ras component's finalize, which
+    // runs after prte_finalize's own session sweep) has nothing to clear
+    if (0 <= s->index && NULL != prte_sessions) {
         pmix_pointer_array_set_item(prte_sessions, s->index, NULL);
     }
 }

test/unit/ras/test_ras.c

@@ -790,13 +790,15 @@ static int test_slurm_allocation(void)
 
     mod = (prte_ras_base_module_t *) module;
     /* the vtable contract, on the module the component actually hands out.
-     * slurm is the only component implementing the elastic completion hooks
-     * today; if another grows them, the base cycles every module for both. */
+     * slurm is the only component implementing shrink_complete today; if
+     * another grows it, the base cycles every module for it. release_allocation
+     * is unset - the session stack now retains its own reference and drains
+     * itself at finalize, so the reactive hook is no longer needed. */
     CHECK("slurm contract: allocate", NULL != mod->allocate);
     CHECK("slurm contract: init", NULL != mod->init);
     CHECK("slurm contract: modify", NULL != mod->modify);
     CHECK("slurm contract: shrink_complete", NULL != mod->shrink_complete);
-    CHECK("slurm contract: release_allocation", NULL != mod->release_allocation);
+    CHECK("slurm contract: release_allocation unset", NULL == mod->release_allocation);
     if (NULL == mod->allocate) {
         return failures;
     }