13249826981

Committed 10 Feb 2025 08:28PM UTC coverage: 95.735%. Remained the same

Build # 13249826981

Build Type

Pull #569

github

Committed by

web-flow

Commit Message

Merge cbe50ab45 into f7de7f7ca

Pull Request Pull Request #569: [bug] SLURM: Fix CPUs per task

Run Details

0 of 1 new or added line in 1 file covered. (0.0%)

1100 of 1149 relevant lines covered (95.74%)

0.96 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.67

/executorlib/interactive/slurm.py

import os
from typing import Optional

from executorlib.standalone.interactive.spawner import SubprocessSpawner

SLURM_COMMAND = "srun"


def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):
    cores_total = int(os.environ["SLURM_NTASKS"]) * int(
        os.environ["SLURM_CPUS_PER_TASK"]
    )
    cores_requested = max_workers * cores * threads_per_core
    if cores_total < cores_requested:
        raise ValueError(
            "The number of requested cores is larger than the available cores "
            + str(cores_total)
            + " < "
            + str(cores_requested)
        )


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: Optional[list[str]] = None,
    ):
        """
        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: Optional[str],
    threads_per_core: int = 1,
    gpus_per_core: int = 0,
    openmpi_oversubscribe: bool = False,
    slurm_cmd_args: Optional[list[str]] = None,
) -> 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 slurm_cmd_args is not None and len(slurm_cmd_args) > 0:
        command_prepend_lst += slurm_cmd_args
    return command_prepend_lst

1	import os	1✔
2	from typing import Optional	1✔
3
4	from executorlib.standalone.interactive.spawner import SubprocessSpawner	1✔
5
6	SLURM_COMMAND = "srun"	1✔
7
8
9	def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):	1✔
10	cores_total = int(os.environ["SLURM_NTASKS"]) * int(	1✔
11	os.environ["SLURM_CPUS_PER_TASK"]
12	)
13	cores_requested = max_workers * cores * threads_per_core	1✔
14	if cores_total < cores_requested:	1✔
15	raise ValueError(	1✔
16	"The number of requested cores is larger than the available cores "
17	+ str(cores_total)
18	+ " < "
19	+ str(cores_requested)
20	)
21
22
23	class SrunSpawner(SubprocessSpawner):	1✔
24	def __init__(	1✔
25	self,
26	cwd: Optional[str] = None,
27	cores: int = 1,
28	threads_per_core: int = 1,
29	gpus_per_core: int = 0,
30	openmpi_oversubscribe: bool = False,
31	slurm_cmd_args: Optional[list[str]] = None,
32	):
33	"""
34	Srun interface implementation.
35
36	Args:
37	cwd (str, optional): The current working directory. Defaults to None.
38	cores (int, optional): The number of cores to use. Defaults to 1.
39	threads_per_core (int, optional): The number of threads per core. Defaults to 1.
40	gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
41	openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
42	slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
43	"""
44	super().__init__(	1✔
45	cwd=cwd,
46	cores=cores,
47	openmpi_oversubscribe=openmpi_oversubscribe,
48	threads_per_core=threads_per_core,
49	)
50	self._gpus_per_core = gpus_per_core	1✔
51	self._slurm_cmd_args = slurm_cmd_args	1✔
52
53	def generate_command(self, command_lst: list[str]) -> list[str]:	1✔
54	"""
55	Generate the command list for the Srun interface.
56
57	Args:
58	command_lst (list[str]): The command list.
59
60	Returns:
61	list[str]: The generated command list.
62	"""
63	command_prepend_lst = generate_slurm_command(	1✔
64	cores=self._cores,
65	cwd=self._cwd,
66	threads_per_core=self._threads_per_core,
67	gpus_per_core=self._gpus_per_core,
68	openmpi_oversubscribe=self._openmpi_oversubscribe,
69	slurm_cmd_args=self._slurm_cmd_args,
70	)
71	return super().generate_command(	1✔
72	command_lst=command_prepend_lst + command_lst,
73	)
74
75
76	def generate_slurm_command(	1✔
77	cores: int,
78	cwd: Optional[str],
79	threads_per_core: int = 1,
80	gpus_per_core: int = 0,
81	openmpi_oversubscribe: bool = False,
82	slurm_cmd_args: Optional[list[str]] = None,
83	) -> list[str]:
84	"""
85	Generate the command list for the SLURM interface.
86
87	Args:
88	cores (int): The number of cores.
89	cwd (str): The current working directory.
90	threads_per_core (int, optional): The number of threads per core. Defaults to 1.
91	gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
92	openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
93	slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
94
95	Returns:
96	list[str]: The generated command list.
97	"""
98	command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]	1✔
99	if cwd is not None:	1✔
100	command_prepend_lst += ["-D", cwd]	1✔
101	if threads_per_core > 1:	1✔
NEW 102	command_prepend_lst += ["--cpus-per-task=" + str(threads_per_core)]	×
103	if gpus_per_core > 0:	1✔
104	command_prepend_lst += ["--gpus-per-task=" + str(gpus_per_core)]	1✔
105	if openmpi_oversubscribe:	1✔
106	command_prepend_lst += ["--oversubscribe"]	1✔
107	if slurm_cmd_args is not None and len(slurm_cmd_args) > 0:	1✔
108	command_prepend_lst += slurm_cmd_args	1✔
109	return command_prepend_lst	1✔

pyiron / executorlib / 13249826981

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous