@@ -201,9 +201,16 @@ static pmix_status_t hash_assign_module(pmix_info_t *info, size_t ninfo,
201201 return PMIX_SUCCESS ;
202202}
203203
204+ /* Define a bitmask to track what information may not have
205+ * been provided but is computable from other info */
206+ #define PMIX_HASH_PROC_DATA 0x00000001
207+ #define PMIX_HASH_JOB_SIZE 0x00000002
208+ #define PMIX_HASH_MAX_PROCS 0x00000004
209+ #define PMIX_HASH_NUM_NODES 0x00000008
210+
204211static pmix_status_t store_map (pmix_hash_table_t * ht ,
205212 char * * nodes , char * * ppn ,
206- bool procdataprovided )
213+ uint32_t flags )
207214{
208215 pmix_status_t rc ;
209216 pmix_value_t * val ;
@@ -213,6 +220,8 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
213220 bool updated ;
214221 pmix_kval_t * kp2 ;
215222 char * * procs ;
223+ uint32_t totalprocs = 0 ;
224+ bool localldr ;
216225
217226 pmix_output_verbose (2 , pmix_gds_base_framework .framework_output ,
218227 "[%s:%d] gds:hash:store_map" ,
@@ -224,6 +233,22 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
224233 return PMIX_ERR_BAD_PARAM ;
225234 }
226235
236+ /* if they didn't provide the number of nodes, then
237+ * compute it from the list of nodes */
238+ if (!(PMIX_HASH_NUM_NODES & flags )) {
239+ kp2 = PMIX_NEW (pmix_kval_t );
240+ kp2 -> key = strdup (PMIX_NUM_NODES );
241+ kp2 -> value = (pmix_value_t * )malloc (sizeof (pmix_value_t ));
242+ kp2 -> value -> type = PMIX_UINT32 ;
243+ kp2 -> value -> data .uint32 = pmix_argv_count (nodes );
244+ if (PMIX_SUCCESS != (rc = pmix_hash_store (ht , PMIX_RANK_WILDCARD , kp2 ))) {
245+ PMIX_ERROR_LOG (rc );
246+ PMIX_RELEASE (kp2 );
247+ return rc ;
248+ }
249+ PMIX_RELEASE (kp2 ); // maintain acctg
250+ }
251+
227252 for (n = 0 ; NULL != nodes [n ]; n ++ ) {
228253 /* check and see if we already have data for this node */
229254 val = NULL ;
@@ -241,18 +266,22 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
241266 }
242267 iptr = (pmix_info_t * )val -> data .darray -> array ;
243268 updated = false;
269+ localldr = false;
244270 for (m = 0 ; m < val -> data .darray -> size ; m ++ ) {
245- if (0 == strncmp ( iptr [m ]. key , PMIX_LOCAL_PEERS , PMIX_MAX_KEYLEN )) {
271+ if (PMIX_CHECK_KEY ( & iptr [m ], PMIX_LOCAL_PEERS )) {
246272 /* we will update this entry */
247273 if (NULL != iptr [m ].value .data .string ) {
248274 free (iptr [m ].value .data .string );
249275 }
250276 iptr [m ].value .data .string = strdup (ppn [n ]);
251- updated = true;
252- break ;
277+ updated = true; // no need to add the local_peers to the array
278+ } else if (PMIX_CHECK_KEY (& iptr [m ], PMIX_LOCALLDR )) {
279+ rank = strtoul (ppn [n ], NULL , 10 );
280+ iptr [m ].value .data .rank = rank ;
281+ localldr = true; // no need to add localldr to the array
253282 }
254283 }
255- if (!updated ) {
284+ if (!updated || ! localldr ) {
256285 /* append this entry to the current data */
257286 kp2 = PMIX_NEW (pmix_kval_t );
258287 if (NULL == kp2 ) {
@@ -271,7 +300,18 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
271300 return PMIX_ERR_NOMEM ;
272301 }
273302 kp2 -> value -> data .darray -> type = PMIX_INFO ;
274- kp2 -> value -> data .darray -> size = val -> data .darray -> size + 1 ;
303+ /* if we didn't update the local leader, then we will
304+ * add it here */
305+ m = 0 ;
306+ if (!localldr ) {
307+ kp2 -> value -> data .darray -> size = val -> data .darray -> size + 1 ;
308+ ++ m ;
309+ }
310+ /* if they didn't update the local peers, then we add it here */
311+ if (!updated ) {
312+ kp2 -> value -> data .darray -> size = val -> data .darray -> size + 1 ;
313+ ++ m ;
314+ }
275315 PMIX_INFO_CREATE (info , kp2 -> value -> data .darray -> size );
276316 if (NULL == info ) {
277317 PMIX_RELEASE (kp2 );
@@ -281,7 +321,15 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
281321 for (m = 0 ; m < val -> data .darray -> size ; m ++ ) {
282322 PMIX_INFO_XFER (& info [m ], & iptr [m ]);
283323 }
284- PMIX_INFO_LOAD (& info [kp2 -> value -> data .darray -> size - 1 ], PMIX_LOCAL_PEERS , ppn [n ], PMIX_STRING );
324+ if (!updated ) {
325+ PMIX_INFO_LOAD (& info [kp2 -> value -> data .darray -> size - m ], PMIX_LOCAL_PEERS , ppn [n ], PMIX_STRING );
326+ -- m ;
327+ }
328+ if (!localldr ) {
329+ rank = strtoul (ppn [n ], NULL , 10 );
330+ PMIX_INFO_LOAD (& info [kp2 -> value -> data .darray -> size - m ], PMIX_LOCALLDR , & rank , PMIX_PROC_RANK );
331+ -- m ;
332+ }
285333 kp2 -> value -> data .darray -> array = info ;
286334 if (PMIX_SUCCESS != (rc = pmix_hash_store (ht , PMIX_RANK_WILDCARD , kp2 ))) {
287335 PMIX_ERROR_LOG (rc );
@@ -309,14 +357,16 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
309357 return PMIX_ERR_NOMEM ;
310358 }
311359 kp2 -> value -> data .darray -> type = PMIX_INFO ;
312- PMIX_INFO_CREATE (info , 1 );
360+ PMIX_INFO_CREATE (info , 2 );
313361 if (NULL == info ) {
314362 PMIX_RELEASE (kp2 );
315363 return PMIX_ERR_NOMEM ;
316364 }
317365 PMIX_INFO_LOAD (& info [0 ], PMIX_LOCAL_PEERS , ppn [n ], PMIX_STRING );
366+ rank = strtoul (ppn [n ], NULL , 10 );
367+ PMIX_INFO_LOAD (& info [1 ], PMIX_LOCALLDR , & rank , PMIX_PROC_RANK );
318368 kp2 -> value -> data .darray -> array = info ;
319- kp2 -> value -> data .darray -> size = 1 ;
369+ kp2 -> value -> data .darray -> size = 2 ;
320370 if (PMIX_SUCCESS != (rc = pmix_hash_store (ht , PMIX_RANK_WILDCARD , kp2 ))) {
321371 PMIX_ERROR_LOG (rc );
322372 PMIX_RELEASE (kp2 );
@@ -327,6 +377,7 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
327377 /* split the list of procs so we can store their
328378 * individual location data */
329379 procs = pmix_argv_split (ppn [n ], ',' );
380+ totalprocs += pmix_argv_count (procs );
330381 for (m = 0 ; NULL != procs [m ]; m ++ ) {
331382 /* store the hostname for each proc */
332383 kp2 = PMIX_NEW (pmix_kval_t );
@@ -342,7 +393,7 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
342393 return rc ;
343394 }
344395 PMIX_RELEASE (kp2 ); // maintain acctg
345- if (!procdataprovided ) {
396+ if (!( PMIX_HASH_PROC_DATA & flags ) ) {
346397 /* add an entry for the nodeid */
347398 kp2 = PMIX_NEW (pmix_kval_t );
348399 kp2 -> key = strdup (PMIX_NODEID );
@@ -403,6 +454,41 @@ static pmix_status_t store_map(pmix_hash_table_t *ht,
403454 }
404455 PMIX_RELEASE (kp2 ); // maintain acctg
405456
457+ /* if they didn't provide the job size, compute it as
458+ * being the number of provided procs (i.e., size of
459+ * ppn list) */
460+ if (!(PMIX_HASH_JOB_SIZE & flags )) {
461+ kp2 = PMIX_NEW (pmix_kval_t );
462+ kp2 -> key = strdup (PMIX_JOB_SIZE );
463+ kp2 -> value = (pmix_value_t * )malloc (sizeof (pmix_value_t ));
464+ kp2 -> value -> type = PMIX_UINT32 ;
465+ kp2 -> value -> data .uint32 = totalprocs ;
466+ if (PMIX_SUCCESS != (rc = pmix_hash_store (ht , PMIX_RANK_WILDCARD , kp2 ))) {
467+ PMIX_ERROR_LOG (rc );
468+ PMIX_RELEASE (kp2 );
469+ return rc ;
470+ }
471+ PMIX_RELEASE (kp2 ); // maintain acctg
472+ }
473+
474+ /* if they didn't provide a value for max procs, just
475+ * assume it is the same as the number of procs in the
476+ * job and store it */
477+ if (!(PMIX_HASH_MAX_PROCS & flags )) {
478+ kp2 = PMIX_NEW (pmix_kval_t );
479+ kp2 -> key = strdup (PMIX_MAX_PROCS );
480+ kp2 -> value = (pmix_value_t * )malloc (sizeof (pmix_value_t ));
481+ kp2 -> value -> type = PMIX_UINT32 ;
482+ kp2 -> value -> data .uint32 = totalprocs ;
483+ if (PMIX_SUCCESS != (rc = pmix_hash_store (ht , PMIX_RANK_WILDCARD , kp2 ))) {
484+ PMIX_ERROR_LOG (rc );
485+ PMIX_RELEASE (kp2 );
486+ return rc ;
487+ }
488+ PMIX_RELEASE (kp2 ); // maintain acctg
489+ }
490+
491+
406492 return PMIX_SUCCESS ;
407493}
408494
@@ -419,7 +505,7 @@ pmix_status_t hash_cache_job_info(struct pmix_namespace_t *ns,
419505 pmix_rank_t rank ;
420506 pmix_status_t rc = PMIX_SUCCESS ;
421507 size_t n , j , size , len ;
422- bool procdataprovided = false ;
508+ uint32_t flags = 0 ;
423509
424510 pmix_output_verbose (2 , pmix_gds_base_framework .framework_output ,
425511 "[%s:%d] gds:hash:cache_job_info for nspace %s" ,
@@ -482,7 +568,7 @@ pmix_status_t hash_cache_job_info(struct pmix_namespace_t *ns,
482568 goto release ;
483569 }
484570 } else if (0 == strcmp (info [n ].key , PMIX_PROC_DATA )) {
485- procdataprovided = true ;
571+ flags |= PMIX_HASH_PROC_DATA ;
486572 /* an array of data pertaining to a specific proc */
487573 if (PMIX_DATA_ARRAY != info [n ].value .type ) {
488574 PMIX_ERROR_LOG (PMIX_ERR_BAD_PARAM );
@@ -572,9 +658,15 @@ pmix_status_t hash_cache_job_info(struct pmix_namespace_t *ns,
572658 goto release ;
573659 }
574660 PMIX_RELEASE (kp2 ); // maintain acctg
575- /* if this is the job size, then store it */
576- if (0 == strncmp (info [n ].key , PMIX_JOB_SIZE , PMIX_MAX_KEYLEN )) {
661+ /* if this is the job size, then store it in
662+ * the nptr tracker and flag that we were given it */
663+ if (PMIX_CHECK_KEY (& info [n ], PMIX_JOB_SIZE )) {
577664 nptr -> nprocs = info [n ].value .data .uint32 ;
665+ flags |= PMIX_HASH_JOB_SIZE ;
666+ } else if (PMIX_CHECK_KEY (& info [n ], PMIX_NUM_NODES )) {
667+ flags |= PMIX_HASH_NUM_NODES ;
668+ } else if (PMIX_CHECK_KEY (& info [n ], PMIX_MAX_PROCS )) {
669+ flags |= PMIX_HASH_MAX_PROCS ;
578670 }
579671 }
580672 }
@@ -612,7 +704,7 @@ pmix_status_t hash_cache_job_info(struct pmix_namespace_t *ns,
612704 goto release ;
613705 }
614706
615- if (PMIX_SUCCESS != (rc = store_map (ht , nodes , procs , procdataprovided ))) {
707+ if (PMIX_SUCCESS != (rc = store_map (ht , nodes , procs , flags ))) {
616708 PMIX_ERROR_LOG (rc );
617709 goto release ;
618710 }
0 commit comments