Add support for PMIX_MEM_ALLOC_KIND

Add a new cmd line option that corresponds to this
attribute. Add the attribute to the prun payload.
When received, it will default to including in the
job info for the spawned job. Add query support for it.

Signed-off-by: Ralph Castain <rhc@pmix.org>
(cherry picked from commit 3957789)

Author: rhc54

src/mca/schizo/prte/schizo_prte.c

@@ -196,6 +196,7 @@ static struct option prterunoptions[] = {
     PMIX_OPTION_SHORT_DEFINE(PRTE_CLI_PRELOAD_BIN, PMIX_ARG_NONE, 's'),
     PMIX_OPTION_DEFINE(PRTE_CLI_DO_NOT_AGG_HELP, PMIX_ARG_NONE),
     PMIX_OPTION_DEFINE(PRTE_CLI_FWD_ENVIRON, PMIX_ARG_OPTIONAL),
+    PMIX_OPTION_DEFINE(PRTE_CLI_MEM_ALLOC_KIND, PMIX_ARG_REQD),
 
     // output options
     PMIX_OPTION_DEFINE(PRTE_CLI_OUTPUT, PMIX_ARG_REQD),
@@ -310,6 +311,7 @@ static struct option prunoptions[] = {
     PMIX_OPTION_SHORT_DEFINE(PRTE_CLI_PRELOAD_BIN, PMIX_ARG_NONE, 's'),
     PMIX_OPTION_DEFINE(PRTE_CLI_DO_NOT_AGG_HELP, PMIX_ARG_NONE),
     PMIX_OPTION_DEFINE(PRTE_CLI_FWD_ENVIRON, PMIX_ARG_OPTIONAL),
+    PMIX_OPTION_DEFINE(PRTE_CLI_MEM_ALLOC_KIND, PMIX_ARG_REQD),
 
     // output options
     PMIX_OPTION_DEFINE(PRTE_CLI_OUTPUT, PMIX_ARG_REQD),

src/prted/pmix/pmix_server_queries.c

@@ -90,7 +90,8 @@ static void _query(int sd, short args, void *cbdata)
     pmix_proc_info_t *procinfo;
     pmix_data_array_t dry;
     prte_proc_t *proct;
-    pmix_proc_t *proc;
+    pmix_proc_t *proc, pproc;
+    pmix_info_t info;
     size_t sz;
     PRTE_HIDE_UNUSED_PARAMS(sd, args);
 
@@ -823,6 +824,27 @@ static void _query(int sd, short args, void *cbdata)
                 }
 #endif
 
+#ifdef PMIX_MEM_ALLOC_KIND
+            } else if (0 == strcmp(q->keys[n], PMIX_MEM_ALLOC_KIND)) {
+                pmix_value_t *value;
+                jdata = prte_get_job_data_object(jobid);
+                if (NULL == jdata) {
+                    ret = PMIX_ERR_NOT_FOUND;
+                    goto done;
+                }
+                PMIX_LOAD_PROCID(&pproc, jobid, PMIX_RANK_WILDCARD);
+                PMIX_INFO_LOAD(&info, PMIX_IMMEDIATE, NULL, PMIX_BOOL);
+                ret = PMIx_Get(&pproc, PMIX_MEM_ALLOC_KIND, &info, 1, (void**)&value);
+                if (PMIX_SUCCESS != ret) {
+                    goto done;
+                }
+                PMIX_INFO_LIST_ADD(rc, results, PMIX_MEM_ALLOC_KIND, value->data.string, PMIX_STRING);
+                PMIX_VALUE_RELEASE(value);
+                if (PMIX_SUCCESS != rc) {
+                    PMIX_ERROR_LOG(rc);
+                    goto done;
+                }
+#endif
             } else {
                 fprintf(stderr, "Query for unrecognized attribute: %s\n", q->keys[n]);
             }

src/prted/prun_common.c

@@ -661,6 +661,13 @@ int prun_common(pmix_cli_result_t *results,
         PMIX_INFO_LIST_ADD(ret, jinfo, PMIX_LOG_AGG, &flag, PMIX_BOOL);
     }
 
+#ifdef PMIX_MEM_ALLOC_KIND
+    opt = pmix_cmd_line_get_param(results, PRTE_CLI_MEM_ALLOC_KIND);
+    if (NULL != opt) {
+        PMIX_INFO_LIST_ADD(ret, jinfo, PMIX_MEM_ALLOC_KIND, opt->values[0], PMIX_STRING);
+    }
+#endif
+
     /* give the schizo components a chance to add to the job info */
     schizo->job_info(results, jinfo);
 

src/util/prte_cmd_line.h

@@ -15,7 +15,7 @@
  * Copyright (c) 2016-2017 Los Alamos National Security, LLC. All rights
  *                         reserved.
  * Copyright (c) 2017-2022 IBM Corporation.  All rights reserved.
- * Copyright (c) 2021-2023 Nanook Consulting.  All rights reserved.
+ * Copyright (c) 2021-2024 Nanook Consulting  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -108,6 +108,7 @@ BEGIN_C_DECLS
 #define PRTE_CLI_SET_CWD_SESSION        "set-cwd-to-session-dir"    // none
 #define PRTE_CLI_ENABLE_RECOVERY        "enable-recovery"           // none
 #define PRTE_CLI_DISABLE_RECOVERY       "disable-recovery"          // none
+#define PRTE_CLI_MEM_ALLOC_KIND			"memory-alloc-kinds"        // required
 
 // Placement options
 #define PRTE_CLI_MAPBY                  "map-by"                    // required