1515 * Copyright (c) 2009 Sun Microsystems, Inc. All rights reserved.
1616 * Copyright (c) 2011-2015 Los Alamos National Security, LLC. All rights
1717 * reserved.
18- * Copyright (c) 2013-2017 Intel, Inc. All rights reserved.
18+ * Copyright (c) 2013-2019 Intel, Inc. All rights reserved.
1919 * Copyright (c) 2014-2017 Research Organization for Information Science
2020 * and Technology (RIST). All rights reserved.
2121 * Copyright (c) 2018 Amazon.com, Inc. or its affiliates. All Rights reserved.
@@ -406,9 +406,43 @@ int ompi_dpm_connect_accept(ompi_communicator_t *comm, int root,
406406 goto exit ;
407407 }
408408 if (0 < opal_list_get_size (& ilist )) {
409+ uint32_t * peer_ranks = NULL ;
410+ int prn , nprn ;
411+ char * val , * mycpuset ;
412+ uint16_t u16 ;
413+ opal_process_name_t wildcard_rank ;
409414 /* convert the list of new procs to a proc_t array */
410415 new_proc_list = (ompi_proc_t * * )calloc (opal_list_get_size (& ilist ),
411416 sizeof (ompi_proc_t * ));
417+ /* get the list of local peers for the new procs */
418+ cd = (ompi_dpm_proct_caddy_t * )opal_list_get_first (& ilist );
419+ proc = cd -> p ;
420+ wildcard_rank .jobid = proc -> super .proc_name .jobid ;
421+ wildcard_rank .vpid = OMPI_NAME_WILDCARD -> vpid ;
422+ /* retrieve the local peers */
423+ OPAL_MODEX_RECV_VALUE_OPTIONAL (rc , OPAL_PMIX_LOCAL_PEERS ,
424+ & wildcard_rank , & val , OPAL_STRING );
425+ if (OPAL_SUCCESS == rc && NULL != val ) {
426+ char * * peers = opal_argv_split (val , ',' );
427+ free (val );
428+ nprn = opal_argv_count (peers );
429+ peer_ranks = (uint32_t * )calloc (nprn , sizeof (uint32_t ));
430+ for (prn = 0 ; NULL != peers [prn ]; prn ++ ) {
431+ peer_ranks [prn ] = strtoul (peers [prn ], NULL , 10 );
432+ }
433+ opal_argv_free (peers );
434+ }
435+
436+ /* get my locality string */
437+ val = NULL ;
438+ OPAL_MODEX_RECV_VALUE_OPTIONAL (rc , OPAL_PMIX_LOCALITY_STRING ,
439+ OMPI_PROC_MY_NAME , & val , OPAL_STRING );
440+ if (OPAL_SUCCESS == rc && NULL != val ) {
441+ mycpuset = val ;
442+ } else {
443+ mycpuset = NULL ;
444+ }
445+
412446 i = 0 ;
413447 OPAL_LIST_FOREACH (cd , & ilist , ompi_dpm_proct_caddy_t ) {
414448 opal_value_t * kv ;
@@ -418,15 +452,38 @@ int ompi_dpm_connect_accept(ompi_communicator_t *comm, int root,
418452 * OPAL_PMIX_LOCALITY and OPAL_PMIX_HOSTNAME. since we can live without
419453 * them, we are just fine */
420454 ompi_proc_complete_init_single (proc );
421- /* save the locality for later */
422- kv = OBJ_NEW (opal_value_t );
423- kv -> key = strdup (OPAL_PMIX_LOCALITY );
424- kv -> type = OPAL_UINT16 ;
425- kv -> data .uint16 = proc -> super .proc_flags ;
426- opal_pmix .store_local (& proc -> super .proc_name , kv );
427- OBJ_RELEASE (kv ); // maintain accounting
455+ /* if this proc is local, then get its locality */
456+ if (NULL != peer_ranks ) {
457+ for (prn = 0 ; prn < nprn ; prn ++ ) {
458+ if (peer_ranks [prn ] == proc -> super .proc_name .vpid ) {
459+ /* get their locality string */
460+ val = NULL ;
461+ OPAL_MODEX_RECV_VALUE_IMMEDIATE (rc , OPAL_PMIX_LOCALITY_STRING ,
462+ & proc -> super .proc_name , & val , OPAL_STRING );
463+ if (OPAL_SUCCESS == rc && NULL != val ) {
464+ u16 = opal_hwloc_compute_relative_locality (mycpuset , val );
465+ free (val );
466+ } else {
467+ /* all we can say is that it shares our node */
468+ u16 = OPAL_PROC_ON_CLUSTER | OPAL_PROC_ON_CU | OPAL_PROC_ON_NODE ;
469+ }
470+ proc -> super .proc_flags = u16 ;
471+ /* save the locality for later */
472+ kv = OBJ_NEW (opal_value_t );
473+ kv -> key = strdup (OPAL_PMIX_LOCALITY );
474+ kv -> type = OPAL_UINT16 ;
475+ kv -> data .uint16 = proc -> super .proc_flags ;
476+ opal_pmix .store_local (& proc -> super .proc_name , kv );
477+ OBJ_RELEASE (kv ); // maintain accounting
478+ break ;
479+ }
480+ }
481+ }
428482 ++ i ;
429483 }
484+ if (NULL != mycpuset ) {
485+ free (mycpuset );
486+ }
430487 /* call add_procs on the new ones */
431488 rc = MCA_PML_CALL (add_procs (new_proc_list , opal_list_get_size (& ilist )));
432489 free (new_proc_list );
0 commit comments