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.

Sign up for GitHub

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

Move SLURM to separate module #528

Merged
merged 3 commits into from
Dec 20, 2024
Merged

Move SLURM to separate module #528

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0b9ee35
Move SLURM to separate module
jan-janssen Dec 20, 2024
79b588c
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Dec 20, 2024
b4a8a37
clean up
jan-janssen Dec 20, 2024
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: 2 additions & 4 deletions executorlib/interactive/executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
InteractiveStepExecutor,
execute_tasks_with_dependencies,
)
from executorlib.interactive.slurm import SrunSpawner
from executorlib.standalone.inputcheck import (
check_command_line_argument_lst,
check_executor,
Expand All @@ -17,10 +18,7 @@
check_pmi,
validate_number_of_cores,
)
from executorlib.standalone.interactive.spawner import (
MpiExecSpawner,
SrunSpawner,
)
from executorlib.standalone.interactive.spawner import MpiExecSpawner
from executorlib.standalone.plot import (
draw,
generate_nodes_and_edges,
Expand Down
94 changes: 94 additions & 0 deletions executorlib/interactive/slurm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
from typing import Optional

from executorlib.standalone.interactive.spawner import SubprocessSpawner

SLURM_COMMAND = "srun"


class SrunSpawner(SubprocessSpawner):
def __init__(
self,
cwd: Optional[str] = None,
cores: int = 1,
threads_per_core: int = 1,
gpus_per_core: int = 0,
openmpi_oversubscribe: bool = False,
slurm_cmd_args: list[str] = [],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Avoid using mutable default arguments.
Using mutable lists as default argument values can lead to unexpected behavior when multiple instances share the same list reference.

You can fix this by setting the default to None and initializing it in the method body:

-        slurm_cmd_args: list[str] = [],
+        slurm_cmd_args: Optional[list[str]] = None,

And inside the constructor:

 if slurm_cmd_args is None:
     slurm_cmd_args = []

Committable suggestion skipped: line range outside the PR's diff.

🧰 Tools
🪛 Ruff (0.8.2)

16-16: Do not use mutable data structures for argument defaults

Replace with None; initialize within function

(B006)

):
"""
Srun interface implementation.

Args:
cwd (str, optional): The current working directory. Defaults to None.
cores (int, optional): The number of cores to use. Defaults to 1.
threads_per_core (int, optional): The number of threads per core. Defaults to 1.
gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
"""
super().__init__(
cwd=cwd,
cores=cores,
openmpi_oversubscribe=openmpi_oversubscribe,
threads_per_core=threads_per_core,
)
self._gpus_per_core = gpus_per_core
self._slurm_cmd_args = slurm_cmd_args

def generate_command(self, command_lst: list[str]) -> list[str]:
"""
Generate the command list for the Srun interface.

Args:
command_lst (list[str]): The command list.

Returns:
list[str]: The generated command list.
"""
command_prepend_lst = generate_slurm_command(
cores=self._cores,
cwd=self._cwd,
threads_per_core=self._threads_per_core,
gpus_per_core=self._gpus_per_core,
openmpi_oversubscribe=self._openmpi_oversubscribe,
slurm_cmd_args=self._slurm_cmd_args,
)
return super().generate_command(
command_lst=command_prepend_lst + command_lst,
)


def generate_slurm_command(
cores: int,
cwd: str,
threads_per_core: int = 1,
gpus_per_core: int = 0,
openmpi_oversubscribe: bool = False,
slurm_cmd_args: list[str] = [],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Same mutable default argument concern.
This function parameter also defines a mutable list as the default value, raising the same risk of shared state across multiple calls.

Use the same approach as above to ensure each invocation gets its own list instance:

-    slurm_cmd_args: list[str] = [],
+    slurm_cmd_args: Optional[list[str]] = None,

and initialize within the function body:

 if slurm_cmd_args is None:
     slurm_cmd_args = []

Committable suggestion skipped: line range outside the PR's diff.

🧰 Tools
🪛 Ruff (0.8.2)

67-67: Do not use mutable data structures for argument defaults

Replace with None; initialize within function

(B006)

) -> list[str]:
"""
Generate the command list for the SLURM interface.

Args:
cores (int): The number of cores.
cwd (str): The current working directory.
threads_per_core (int, optional): The number of threads per core. Defaults to 1.
gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].

Returns:
list[str]: The generated command list.
"""
command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
if cwd is not None:
command_prepend_lst += ["-D", cwd]
if threads_per_core > 1:
command_prepend_lst += ["--cpus-per-task" + str(threads_per_core)]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Potential typo: missing equals sign in '--cpus-per-task'.
SLURM typically expects an equals sign for this argument to parse the integer value properly.

Proposed fix:

-        command_prepend_lst += ["--cpus-per-task" + str(threads_per_core)]
+        command_prepend_lst += ["--cpus-per-task=" + str(threads_per_core)]
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
command_prepend_lst += ["--cpus-per-task" + str(threads_per_core)]
command_prepend_lst += ["--cpus-per-task=" + str(threads_per_core)]

if gpus_per_core > 0:
command_prepend_lst += ["--gpus-per-task=" + str(gpus_per_core)]
if openmpi_oversubscribe:
command_prepend_lst += ["--oversubscribe"]
if len(slurm_cmd_args) > 0:
command_prepend_lst += slurm_cmd_args
return command_prepend_lst
3 changes: 1 addition & 2 deletions executorlib/standalone/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
interface_send,
interface_shutdown,
)
from executorlib.standalone.interactive.spawner import MpiExecSpawner, SrunSpawner
from executorlib.standalone.interactive.spawner import MpiExecSpawner
from executorlib.standalone.thread import RaisingThread

__all__ = [
Expand All @@ -18,5 +18,4 @@
interface_receive,
RaisingThread,
MpiExecSpawner,
SrunSpawner,
]
90 changes: 0 additions & 90 deletions executorlib/standalone/interactive/spawner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from typing import Optional

MPI_COMMAND = "mpiexec"
SLURM_COMMAND = "srun"


class BaseSpawner(ABC):
Expand Down Expand Up @@ -147,59 +146,6 @@ def generate_command(self, command_lst: list[str]) -> list[str]:
)


class SrunSpawner(SubprocessSpawner):
def __init__(
self,
cwd: Optional[str] = None,
cores: int = 1,
threads_per_core: int = 1,
gpus_per_core: int = 0,
openmpi_oversubscribe: bool = False,
slurm_cmd_args: list[str] = [],
):
"""
Srun interface implementation.

Args:
cwd (str, optional): The current working directory. Defaults to None.
cores (int, optional): The number of cores to use. Defaults to 1.
threads_per_core (int, optional): The number of threads per core. Defaults to 1.
gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
"""
super().__init__(
cwd=cwd,
cores=cores,
openmpi_oversubscribe=openmpi_oversubscribe,
threads_per_core=threads_per_core,
)
self._gpus_per_core = gpus_per_core
self._slurm_cmd_args = slurm_cmd_args

def generate_command(self, command_lst: list[str]) -> list[str]:
"""
Generate the command list for the Srun interface.

Args:
command_lst (list[str]): The command list.

Returns:
list[str]: The generated command list.
"""
command_prepend_lst = generate_slurm_command(
cores=self._cores,
cwd=self._cwd,
threads_per_core=self._threads_per_core,
gpus_per_core=self._gpus_per_core,
openmpi_oversubscribe=self._openmpi_oversubscribe,
slurm_cmd_args=self._slurm_cmd_args,
)
return super().generate_command(
command_lst=command_prepend_lst + command_lst,
)


def generate_mpiexec_command(
cores: int, openmpi_oversubscribe: bool = False
) -> list[str]:
Expand All @@ -220,39 +166,3 @@ def generate_mpiexec_command(
if openmpi_oversubscribe:
command_prepend_lst += ["--oversubscribe"]
return command_prepend_lst


def generate_slurm_command(
cores: int,
cwd: str,
threads_per_core: int = 1,
gpus_per_core: int = 0,
openmpi_oversubscribe: bool = False,
slurm_cmd_args: list[str] = [],
) -> list[str]:
"""
Generate the command list for the SLURM interface.

Args:
cores (int): The number of cores.
cwd (str): The current working directory.
threads_per_core (int, optional): The number of threads per core. Defaults to 1.
gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].

Returns:
list[str]: The generated command list.
"""
command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
if cwd is not None:
command_prepend_lst += ["-D", cwd]
if threads_per_core > 1:
command_prepend_lst += ["--cpus-per-task" + str(threads_per_core)]
if gpus_per_core > 0:
command_prepend_lst += ["--gpus-per-task=" + str(gpus_per_core)]
if openmpi_oversubscribe:
command_prepend_lst += ["--oversubscribe"]
if len(slurm_cmd_args) > 0:
command_prepend_lst += slurm_cmd_args
return command_prepend_lst
3 changes: 2 additions & 1 deletion tests/test_shared_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@
import unittest

from executorlib.standalone.interactive.backend import parse_arguments
from executorlib.standalone.interactive.spawner import SrunSpawner, MpiExecSpawner
from executorlib.standalone.interactive.spawner import MpiExecSpawner
from executorlib.interactive.slurm import SrunSpawner


class TestParser(unittest.TestCase):
Expand Down
Loading