Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Raise error when srun is not available #206

Merged
merged 3 commits into from
Nov 7, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion pympipool/__init__.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
from ._version import get_versions
from pympipool.mpi.executor import PyMPIExecutor
from pympipool.slurm.executor import PySlurmExecutor

try: # The PyFluxExecutor requires flux-core to be installed.
from pympipool.flux.executor import PyFluxExecutor
except ImportError:
pass

try: # The PySlurmExecutor requires the srun command to be available.
from pympipool.slurm.executor import PySlurmExecutor
except ImportError:
pass

__version__ = get_versions()["version"]
del get_versions
8 changes: 6 additions & 2 deletions pympipool/shared/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,10 @@
import subprocess


MPI_COMMAND = "mpiexec"
SLURM_COMMAND = "srun"


class BaseInterface(ABC):
def __init__(self, cwd, cores=1, oversubscribe=False):
self._cwd = cwd
Expand Down Expand Up @@ -94,7 +98,7 @@ def generate_command(self, command_lst):


def generate_mpiexec_command(cores, oversubscribe=False):
command_prepend_lst = ["mpiexec", "-n", str(cores)]
command_prepend_lst = [MPI_COMMAND, "-n", str(cores)]
if oversubscribe:
command_prepend_lst += ["--oversubscribe"]
return command_prepend_lst
Expand All @@ -103,7 +107,7 @@ def generate_mpiexec_command(cores, oversubscribe=False):
def generate_slurm_command(
cores, cwd, threads_per_core=1, gpus_per_core=0, oversubscribe=False
):
command_prepend_lst = ["srun", "-n", str(cores)]
command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
if cwd is not None:
command_prepend_lst += ["-D", cwd]
if threads_per_core > 1:
Expand Down
10 changes: 9 additions & 1 deletion pympipool/slurm/executor.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,21 @@
import shutil
import subprocess


from pympipool.shared.executorbase import (
cloudpickle_register,
execute_parallel_tasks,
ExecutorBase,
executor_broker,
)
from pympipool.shared.interface import SrunInterface
from pympipool.shared.interface import SrunInterface, SLURM_COMMAND
from pympipool.shared.thread import RaisingThread


if shutil.which(SLURM_COMMAND) is None:
raise ImportError("SLURM command " + SLURM_COMMAND + " not found.")


class PySlurmExecutor(ExecutorBase):
"""
Args:
Expand Down
Loading