diff --git a/lib/spack/spack/build_systems/cached_cmake.py b/lib/spack/spack/build_systems/cached_cmake.py index f3a365e264e..ca754edf094 100644 --- a/lib/spack/spack/build_systems/cached_cmake.py +++ b/lib/spack/spack/build_systems/cached_cmake.py @@ -2,9 +2,10 @@ # # SPDX-License-Identifier: (Apache-2.0 OR MIT) import collections.abc +import enum import os import re -from typing import Tuple +from typing import Optional, Tuple import llnl.util.filesystem as fs import llnl.util.tty as tty @@ -13,6 +14,7 @@ import spack.spec import spack.util.prefix from spack.directives import depends_on +from spack.util.executable import which_string from .cmake import CMakeBuilder, CMakePackage @@ -178,6 +180,64 @@ def initconfig_compiler_entries(self): return entries + class Scheduler(enum.Enum): + LSF = enum.auto() + SLURM = enum.auto() + FLUX = enum.auto() + + def get_scheduler(self) -> Optional[Scheduler]: + spec = self.pkg.spec + + # Check for Spectrum-mpi, which always uses LSF or LSF MPI variant + if spec.satisfies("^spectrum-mpi") or spec["mpi"].satisfies("schedulers=lsf"): + return self.Scheduler.LSF + + # Check for Slurm MPI variants + slurm_checks = ["+slurm", "schedulers=slurm", "process_managers=slurm"] + if any(spec["mpi"].satisfies(variant) for variant in slurm_checks): + return self.Scheduler.SLURM + + # TODO improve this when MPI implementations support flux + # Do this check last to avoid using a flux wrapper present next to Slurm/ LSF schedulers + if which_string("flux") is not None: + return self.Scheduler.FLUX + + return None + + def get_mpi_exec(self) -> Optional[str]: + spec = self.pkg.spec + scheduler = self.get_scheduler() + + if scheduler == self.Scheduler.LSF: + return which_string("lrun") + + elif scheduler == self.Scheduler.SLURM: + if spec["mpi"].external: + return which_string("srun") + else: + return os.path.join(spec["slurm"].prefix.bin, "srun") + + elif scheduler == self.Scheduler.FLUX: + flux = which_string("flux") + return f"{flux};run" if flux else None + + elif hasattr(spec["mpi"].package, "mpiexec"): + return spec["mpi"].package.mpiexec + + else: + mpiexec = os.path.join(spec["mpi"].prefix.bin, "mpirun") + if not os.path.exists(mpiexec): + mpiexec = os.path.join(spec["mpi"].prefix.bin, "mpiexec") + return mpiexec + + def get_mpi_exec_num_proc(self) -> str: + scheduler = self.get_scheduler() + + if scheduler in [self.Scheduler.FLUX, self.Scheduler.LSF, self.Scheduler.SLURM]: + return "-n" + else: + return "-np" + def initconfig_mpi_entries(self): spec = self.pkg.spec @@ -197,27 +257,10 @@ def initconfig_mpi_entries(self): if hasattr(spec["mpi"], "mpifc"): entries.append(cmake_cache_path("MPI_Fortran_COMPILER", spec["mpi"].mpifc)) - # Check for slurm - using_slurm = False - slurm_checks = ["+slurm", "schedulers=slurm", "process_managers=slurm"] - if any(spec["mpi"].satisfies(variant) for variant in slurm_checks): - using_slurm = True - # Determine MPIEXEC - if using_slurm: - if spec["mpi"].external: - # Heuristic until we have dependents on externals - mpiexec = "/usr/bin/srun" - else: - mpiexec = os.path.join(spec["slurm"].prefix.bin, "srun") - elif hasattr(spec["mpi"].package, "mpiexec"): - mpiexec = spec["mpi"].package.mpiexec - else: - mpiexec = os.path.join(spec["mpi"].prefix.bin, "mpirun") - if not os.path.exists(mpiexec): - mpiexec = os.path.join(spec["mpi"].prefix.bin, "mpiexec") + mpiexec = self.get_mpi_exec() - if not os.path.exists(mpiexec): + if mpiexec is None or not os.path.exists(mpiexec.split(";")[0]): msg = "Unable to determine MPIEXEC, %s tests may fail" % self.pkg.name entries.append("# {0}\n".format(msg)) tty.warn(msg) @@ -230,10 +273,7 @@ def initconfig_mpi_entries(self): entries.append(cmake_cache_path("MPIEXEC", mpiexec)) # Determine MPIEXEC_NUMPROC_FLAG - if using_slurm: - entries.append(cmake_cache_string("MPIEXEC_NUMPROC_FLAG", "-n")) - else: - entries.append(cmake_cache_string("MPIEXEC_NUMPROC_FLAG", "-np")) + entries.append(cmake_cache_string("MPIEXEC_NUMPROC_FLAG", self.get_mpi_exec_num_proc())) return entries diff --git a/var/spack/repos/builtin/packages/axom/package.py b/var/spack/repos/builtin/packages/axom/package.py index 7c67a3176b6..c9ee774813a 100644 --- a/var/spack/repos/builtin/packages/axom/package.py +++ b/var/spack/repos/builtin/packages/axom/package.py @@ -8,7 +8,6 @@ from os.path import join as pjoin from spack.package import * -from spack.util.executable import which_string def get_spec_path(spec, package_name, path_replacements={}, use_bin=False): @@ -452,19 +451,6 @@ def initconfig_mpi_entries(self): entries.append(cmake_cache_option("ENABLE_MPI", True)) if spec["mpi"].name == "spectrum-mpi": entries.append(cmake_cache_string("BLT_MPI_COMMAND_APPEND", "mpibind")) - - # Replace /usr/bin/srun path with srun flux wrapper path on TOSS 4 - # TODO: Remove this logic by adding `using_flux` case in - # spack/lib/spack/spack/build_systems/cached_cmake.py:196 and remove hard-coded - # path to srun in same file. - if "toss_4" in self._get_sys_type(spec): - srun_wrapper = which_string("srun") - mpi_exec_index = [ - index for index, entry in enumerate(entries) if "MPIEXEC_EXECUTABLE" in entry - ] - if mpi_exec_index: - del entries[mpi_exec_index[0]] - entries.append(cmake_cache_path("MPIEXEC_EXECUTABLE", srun_wrapper)) else: entries.append(cmake_cache_option("ENABLE_MPI", False))