From 761a5ebdc75413643e2aae9d50791812290818e3 Mon Sep 17 00:00:00 2001
From: Ralph Castain <rhc@pmix.org>
Date: Thu, 22 Feb 2024 15:59:54 -0700
Subject: [PATCH] Fix finding "prte" for singleton comm_spawn

Copy the search code from mpirun to use when starting
"prte" in support of comm_spawn. This respects things
like OPAL_PREFIX.

Signed-off-by: Ralph Castain <rhc@pmix.org>
---
 ompi/dpm/dpm.c | 57 ++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 55 insertions(+), 2 deletions(-)

diff --git a/ompi/dpm/dpm.c b/ompi/dpm/dpm.c
index 25f60586893..f10e6194f4a 100644
--- a/ompi/dpm/dpm.c
+++ b/ompi/dpm/dpm.c
@@ -20,7 +20,7 @@
  * Copyright (c) 2014-2020 Research Organization for Information Science
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2018      Amazon.com, Inc. or its affiliates.  All Rights reserved.
- * Copyright (c) 2021-2022 Nanook Consulting.  All rights reserved.
+ * Copyright (c) 2021-2024 Nanook Consulting  All rights reserved.
  * Copyright (c) 2018-2022 Triad National Security, LLC. All rights
  *                         reserved.
  * Copyright (c) 2022      IBM Corporation.  All rights reserved.
@@ -53,6 +53,7 @@
 #include "opal/util/show_help.h"
 #include "opal/util/printf.h"
 #include "opal/mca/hwloc/base/base.h"
+#include "opal/mca/installdirs/base/base.h"
 #include "opal/mca/pmix/base/base.h"
 
 #include "ompi/communicator/communicator.h"
@@ -1974,6 +1975,46 @@ static void set_handler_default(int sig)
     sigaction(sig, &act, (struct sigaction *)0);
 }
 
+static char *find_prte(void)
+{
+    char *filename = NULL;
+#if !OMPI_USING_INTERNAL_PRRTE
+    char *prrte_prefix = NULL;
+#endif
+
+    /* 1) Did the user tell us exactly where to find prte? */
+    filename = getenv("OMPI_PRTERUN");
+    if (NULL != filename) {
+        return filename;
+    }
+
+#if OMPI_USING_INTERNAL_PRRTE
+    /* 2) If using internal PRRTE, use our bindir.  Note that this
+     * will obey OPAL_PREFIX and OPAL_DESTDIR */
+    opal_asprintf(&filename, "%s%sprte", opal_install_dirs.bindir, OPAL_PATH_SEP);
+    return filename;
+#else
+
+    /* 3) Look in ${PRTE_PREFIX}/bin */
+    prrte_prefix = getenv("PRTE_PREFIX");
+    if (NULL != prrte_prefix) {
+        opal_asprintf(&filename, "%s%sbin%sprte", prrte_prefix, OPAL_PATH_SEP, OPAL_PATH_SEP);
+        return filename;
+    }
+
+    /* 4) See if configure told us where to look, if set */
+#if defined(OMPI_PRTERUN_PATH)
+    return strdup(OMPI_PRTERUN_PATH);
+#else
+
+    /* 5) Use path search */
+    filename = opal_find_absolute_path("prte");
+
+    return filename;
+#endif
+#endif
+}
+
 static int start_dvm(char **hostfiles, char **dash_host)
 {
     pmix_status_t pret;
@@ -1987,11 +2028,23 @@ static int start_dvm(char **hostfiles, char **dash_host)
     pmix_info_t info;
     int buffer_length, num_chars_read, chunk;
     char *uri;
+    char *opal_prefix = getenv("OPAL_PREFIX");
+
+    /* as a special case, if OPAL_PREFIX was set and either PRRTE or
+     * PMIx are internal builds, set their prefix variables as well */
+    if (NULL != opal_prefix) {
+#if OMPI_USING_INTERNAL_PRRTE
+        setenv("PRTE_PREFIX", opal_prefix, 1);
+#endif
+#if OPAL_USING_INTERNAL_PMIX
+        setenv("PMIX_PREFIX", opal_prefix, 1);
+#endif
+    }
 
     /* find the prte binary using the install_dirs support - this also
      * checks to ensure that we can see this executable and it *is* executable by us
      */
-    cmd = opal_find_absolute_path("prte");
+    cmd = find_prte();
     if (NULL == cmd) {
         /* guess we couldn't do it - best to abort */
         OMPI_ERROR_LOG(OMPI_ERROR);