[colleeneb]

Pull Request Description

On Aurora, we noticed that after an upgrade from PMIx 4.2.9 to PMIx 5.0.8, the time for MPI_Init at large scales increased. Most of the extra time was spent in PMIx_Resolve_peers. All credit to @anguyen-intel for tracking this down.

From comparing PMIx 4.2.9 and 5.0.8 we found that how the ranks-per-node were discovered changed. PMIx 4.2.9 did a lookup with PMIx_Get based on what PMIx_Init found. PMIx 5.0.8 queries a server if the client is connected. mpich calls PMIx_Resolve_peers once per node, so the cost increases with number of nodes.

To work around this temporarily, this adds an env MPIR_PMIX_USE_GET_LOCAL_PEERS which when set uses pmix_get_local_peers to look up the ranks on a node based on the PMIx 4.2.9 way (see the 4.2.9 way: https://github.com/openpmix/openpmix/blob/eebd966ae14a47b2298d678734015761306c3bd3/src/client/pmix_client.c#L1654-L1689 and https://github.com/openpmix/openpmix/blob/eebd966ae14a47b2298d678734015761306c3bd3/src/client/pmix_client.c#L1563-L1567, this code mirrors it but adds a PMIX_OPTIONAL to ensure it's only a local lookup). If the env is not set or the local lookup fails, PMIx_Resolve_peers is called as it currently is and there's no change to behavior.

As an example of timings, at 1024 nodes with 96 ppn, MPI_Init takes ~35s with the PMIx 5.0.8, and ~8s with this PR.

I went with an env to toggle on since this is a workaround, but definitely open to other suggestions. If this should just be for the aurora_test branch that works too. Thanks!

Author Checklist

• Provide Description
  Particularly focus on why, not what. Reference background, issues, test failures, xfail entries, etc.
• Commits Follow Good Practice
  Commits are self-contained and do not do two things at once.
  Commit message is of the form: module: short description
  Commit message explains what's in the commit.
• Passes All Tests
  Whitespace checker. Warnings test. Additional tests via comments.
• Contribution Agreement
  For non-Argonne authors, check contribution agreement.
  If necessary, request an explicit comment from your companies PR approval manager.

Welcome, New Contributors!

Thank you for contributing to MPICH. If this is your first contribution, please consider:

• Subscribing to the developer mailing list, devel@lists.mpich.org, where design discussions and announcements take place.
• Joining our weekly MPICH developer call, every Thursday at 9:00 AM US Central Time. The meeting agenda and meeting link are announced on the devel mailing list on Wednesday.
• Discussing your PR on the devel mailing list or at the developer call, to help maintainers review it.

[hzhou]

Use a CVAR, e.g. MPIR_CVAR_PMIX_RESOLVE_PEERS_WORKAROUND. Reference the CVAR block at the top of src/util/mpir_pmi.c for examples.