Merge pull request #2022 from rhc54/topic/nnodes

Provide the number of nodes in the job

opal/mca/pmix/pmix_types.h

@@ -111,6 +111,7 @@ BEGIN_C_DECLS
 #define OPAL_PMIX_LOCAL_SIZE                    "pmix.local.size"       // (uint32_t) #procs in this job on this node
 #define OPAL_PMIX_NODE_SIZE                     "pmix.node.size"        // (uint32_t) #procs across all jobs on this node
 #define OPAL_PMIX_MAX_PROCS                     "pmix.max.size"         // (uint32_t) max #procs for this job
+#define OPAL_PMIX_NUM_NODES                     "pmix.num.nodes"        // (uint32_t) #nodes in this nspace
 
 /* topology info */
 #define OPAL_PMIX_NET_TOPO                      "pmix.ntopo"            // (char*) xml-representation of network topology

orte/mca/ess/pmi/ess_pmi_module.c

@@ -214,6 +214,13 @@ static int rte_init(void)
         orte_process_info.num_local_peers = 0;
     }
 
+    /* get number of nodes in the job */
+    OPAL_MODEX_RECV_VALUE_OPTIONAL(ret, OPAL_PMIX_NUM_NODES,
+                                   &wildcard_rank, &u32ptr, OPAL_UINT32);
+    if (OPAL_SUCCESS == ret) {
+        orte_process_info.num_nodes = u32;
+    }
+
     /* setup transport keys in case the MPI layer needs them -
      * we can use the jobfam and stepid as unique keys
      * because they are unique values assigned by the RM

orte/orted/pmix/pmix_server_register_fns.c

@@ -96,7 +96,7 @@ int orte_pmix_server_register_nspace(orte_job_t *jdata)
     map = jdata->map;
     for (i=0; i < map->nodes->size; i++) {
         micro = NULL;
-        if (NULL != (node = (orte_node_t*)opal_pointer_array_get_item(jdata->map->nodes, i))) {
+        if (NULL != (node = (orte_node_t*)opal_pointer_array_get_item(map->nodes, i))) {
             opal_argv_append_nosize(&list, node->name);
             /* assemble all the ranks for this job that are on this node */
             for (k=0; k < node->procs->size; k++) {
@@ -185,6 +185,13 @@ int orte_pmix_server_register_nspace(orte_job_t *jdata)
     kv->data.uint32 = mynode->num_procs;
     opal_list_append(info, &kv->super);
 
+    /* pass the number of nodes in the job */
+    kv = OBJ_NEW(opal_value_t);
+    kv->key = strdup(OPAL_PMIX_NUM_NODES);
+    kv->type = OPAL_UINT32;
+    kv->data.uint32 = map->num_nodes;
+    opal_list_append(info, &kv->super);
+
     /* univ size */
     kv = OBJ_NEW(opal_value_t);
     kv->key = strdup(OPAL_PMIX_UNIV_SIZE);