Coverage for python / lsst / ctrl / bps / parsl / sites / slurm.py: 20%
67 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-03 01:39 -0800
1# This file is part of ctrl_bps_parsl.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org) and the LSST DESC (https://www.lsstdesc.org/).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <https://www.gnu.org/licenses/>.
28from typing import TYPE_CHECKING, Any
30from parsl.executors import HighThroughputExecutor
31from parsl.executors.base import ParslExecutor
32from parsl.providers import SlurmProvider
34from ..configuration import get_bps_config_value, get_workflow_name
35from ..site import SiteConfig
37if TYPE_CHECKING:
38 from ..job import ParslJob
40__all__ = ("Slurm", "TripleSlurm")
43Kwargs = dict[str, Any]
46class Slurm(SiteConfig):
47 """Configuration for generic Slurm cluster.
49 This can be used directly as the site configuration for a Slurm cluster by
50 setting the BPS config, e.g.:
52 .. code-block:: yaml
54 computeSite: slurm
55 site:
56 slurm:
57 class: lsst.ctrl.bps.parsl.sites.Slurm
58 nodes: 3
59 cores_per_node: 20
60 walltime: "00:59:00" # Note: always quote walltime in YAML
62 Alternatively, it can be used as a base class for Slurm cluster
63 configurations.
65 The following BPS configuration parameters are recognised (and required
66 unless there is a default mentioned here, or provided by a subclass):
68 - ``nodes`` (`int`): number of nodes for each Slurm job.
69 - ``cores_per_node`` (`int`): number of cores per node for each Slurm job;
70 by default we use all cores on the node.
71 - ``walltime`` (`str`): time limit for each Slurm job.
72 - ``mem_per_node`` (`int`): memory per node (GB) for each Slurm job; by
73 default we use whatever Slurm gives us.
74 - ``qos`` (`str`): quality of service to request for each Slurm job; by
75 default we use whatever Slurm gives us.
76 - ``singleton`` (`bool`): allow only one job to run at a time; by default
77 ``False``.
78 - ``account`` (`str`): account to use for Slurm jobs.
79 - ``scheduler_options`` (`str`): text to prepend to the Slurm submission
80 script (each line usually starting with ``#SBATCH``).
81 """
83 def __init__(self, *args, **kwargs):
84 # Have BPS-defined resource requests for each job passed to executor.
85 kwargs["resource_list"] = ["priority"]
86 super().__init__(*args, **kwargs)
88 def make_executor(
89 self,
90 label: str,
91 *,
92 nodes: int | None = None,
93 cores_per_node: int | None = None,
94 walltime: str | None = None,
95 mem_per_node: int | None = None,
96 mem_per_worker: float | None = None,
97 qos: str | None = None,
98 constraint: str | None = None,
99 singleton: bool = False,
100 scheduler_options: str | None = None,
101 provider_options: Kwargs | None = None,
102 executor_options: Kwargs | None = None,
103 ) -> ParslExecutor:
104 """Return an executor for running on a Slurm cluster.
106 Parameters
107 ----------
108 label : `str`
109 Label for executor.
110 nodes : `int`, optional
111 Default number of nodes for each Slurm job.
112 cores_per_node : `int`, optional
113 Default number of cores per node for each Slurm job.
114 walltime : `str`, optional
115 Default time limit for each Slurm job.
116 mem_per_node : `float`, optional
117 Memory per node (GB) to request for each Slurm job.
118 mem_per_worker : `float`, optional
119 Minimum memory per worker (GB), limited by the executor.
120 qos : `str`, optional
121 Quality of service for each Slurm job.
122 constraint : `str`, optional
123 Node feature(s) to require for each Slurm job.
124 singleton : `bool`, optional
125 Wether to allow only a single Slurm job to run at a time.
126 scheduler_options : `str`, optional
127 ``#SBATCH`` directives to prepend to the Slurm submission script.
128 provider_options : `dict`, optional
129 Additional arguments for `SlurmProvider` constructor.
130 executor_options : `dict`, optional
131 Additional arguments for `HighThroughputExecutor` constructor.
133 Returns
134 -------
135 executor : `HighThroughputExecutor`
136 Executor for Slurm jobs.
137 """
138 nodes = get_bps_config_value(self.site, "nodes", int, nodes, required=True)
139 cores_per_node = get_bps_config_value(self.site, "cores_per_node", int, cores_per_node)
140 walltime = get_bps_config_value(self.site, "walltime", str, walltime, required=True)
141 mem_per_node = get_bps_config_value(self.site, "mem_per_node", int, mem_per_node)
142 qos = get_bps_config_value(self.site, "qos", str, qos)
143 singleton = get_bps_config_value(self.site, "singleton", bool, singleton)
144 account = get_bps_config_value(self.site, "account", str)
145 scheduler_options = get_bps_config_value(self.site, "scheduler_options", str, scheduler_options)
147 job_name = get_workflow_name(self.config)
148 if scheduler_options is None:
149 scheduler_options = ""
150 else:
151 scheduler_options += "\n"
152 scheduler_options += f"#SBATCH --job-name={job_name}\n"
153 if qos:
154 scheduler_options += f"#SBATCH --qos={qos}\n"
155 if constraint:
156 scheduler_options += f"#SBATCH --constraint={constraint}\n"
157 if singleton:
158 # The following SBATCH directives allow only a single slurm job
159 # (parsl block) with our job_name to run at once. This means we can
160 # have one job running, and one already in the queue when the first
161 # exceeds the walltime limit. More backups could be achieved with a
162 # larger value of max_blocks. This only allows one job to be
163 # actively running at once, so that needs to be sized appropriately
164 # by the user.
165 scheduler_options += "#SBATCH --dependency=singleton\n"
166 return HighThroughputExecutor(
167 label,
168 provider=SlurmProvider(
169 nodes_per_block=nodes,
170 cores_per_node=cores_per_node,
171 mem_per_node=mem_per_node,
172 walltime=walltime,
173 account=account,
174 scheduler_options=scheduler_options,
175 **(provider_options or {}),
176 ),
177 mem_per_worker=mem_per_worker,
178 address=self.get_address(),
179 **(executor_options or {}),
180 )
182 def get_executors(self) -> list[ParslExecutor]:
183 """Get a list of executors to be used in processing.
185 Each executor should have a unique ``label``.
186 """
187 return [self.make_executor("slurm")]
189 def select_executor(self, job: "ParslJob") -> str:
190 """Get the ``label`` of the executor to use to execute a job.
192 Parameters
193 ----------
194 job : `ParslJob`
195 Job to be executed.
197 Returns
198 -------
199 label : `str`
200 Label of executor to use to execute ``job``.
201 """
202 return "slurm"
205class TripleSlurm(Slurm):
206 """Configuration for running jobs on a Slurm cluster with three levels.
208 Parameters
209 ----------
210 *args : `~typing.Any`
211 Parameters forwarded to base class constructor.
212 **kwargs : `~typing.Any`
213 Keyword arguments passed to base class constructor.
215 Notes
216 -----
217 The three levels are useful for having workers with different amount of
218 available memory (and this is how executors are selected, by default),
219 though other uses are possible.
221 The following BPS configuration parameters are recognised, overriding the
222 defaults:
224 - ``nodes`` (`int`): number of nodes for each Slurm job.
225 - ``cores_per_node`` (`int`): number of cores per node for each Slurm job;
226 by default we use all cores on the node.
227 - ``walltime`` (`str`): time limit for each Slurm job; setting this would
228 override each of the ``small_walltime``, ``medium_walltime`` and
229 ``large_walltime`` values.
230 - ``mem_per_node`` (`float`): memory per node for each Slurm job; by
231 default we use whatever Slurm gives us.
232 - ``qos`` (`str`): quality of service to request for each Slurm job; by
233 default we use whatever Slurm gives us.
234 - ``small_memory`` (`float`): memory per worker (GB) for each 'small' Slurm
235 job.
236 - ``medium_memory`` (`float`): memory per worker (GB) for each 'medium'
237 Slurm job.
238 - ``large_memory`` (`float`): memory per worker (GB) for each 'large' Slurm
239 job.
240 - ``small_walltime`` (`str`): time limit for each 'small' Slurm job.
241 - ``medium_walltime`` (`str`): time limit for each 'medium' Slurm job.
242 - ``large_walltime`` (`str`): time limit for each 'large' Slurm job.
243 """
245 def __init__(self, *args, **kwargs):
246 super().__init__(*args, **kwargs)
247 self.small_memory = get_bps_config_value(self.site, "small_memory", float, 2.0)
248 self.medium_memory = get_bps_config_value(self.site, "medium_memory", float, 4.0)
249 self.large_memory = get_bps_config_value(self.site, "large_memory", float, 8.0)
250 self.small_walltime = get_bps_config_value(self.site, "small_walltime", str, "10:00:00")
251 self.medium_walltime = get_bps_config_value(self.site, "medium_walltime", str, "10:00:00")
252 self.large_walltime = get_bps_config_value(self.site, "large_walltime", str, "40:00:00")
254 def get_executors(
255 self,
256 small_options: Kwargs | None = None,
257 medium_options: Kwargs | None = None,
258 large_options: Kwargs | None = None,
259 **common_options,
260 ) -> list[ParslExecutor]:
261 """Get a list of executors to be used in processing.
263 We create three executors, with different walltime and memory per
264 worker.
266 Parameters
267 ----------
268 small_options : kwargs
269 Options for ``make_executor`` for small executor.
270 medium_options : kwargs
271 Options for ``make_executor`` for medium executor.
272 large_options : kwargs
273 Options for ``make_executor`` for large executor.
274 **common_options
275 Common options for ``make_executor`` for each of the executors.
276 """
277 if small_options is None:
278 small_options = {}
279 if medium_options is None:
280 medium_options = {}
281 if large_options is None:
282 large_options = {}
284 small_options["walltime"] = small_options.get("walltime", self.small_walltime)
285 medium_options["walltime"] = medium_options.get("walltime", self.medium_walltime)
286 large_options["walltime"] = large_options.get("walltime", self.large_walltime)
288 small_options["mem_per_worker"] = small_options.get("mem_per_worker", self.small_memory)
289 medium_options["mem_per_worker"] = medium_options.get("mem_per_worker", self.medium_memory)
290 large_options["mem_per_worker"] = large_options.get("mem_per_worker", self.large_memory)
292 return [
293 self.make_executor("small", **small_options, **(common_options or {})),
294 self.make_executor("medium", **medium_options, **(common_options or {})),
295 self.make_executor("large", **large_options, **(common_options or {})),
296 ]
298 def select_executor(self, job: "ParslJob") -> str:
299 """Get the ``label`` of the executor to use to execute a job.
301 This implementation only looks at the requested memory.
303 Parameters
304 ----------
305 job : `ParslJob`
306 Job to be executed.
308 Returns
309 -------
310 label : `str`
311 Label of executor to use to execute ``job``.
312 """
313 memory = job.generic.request_memory / 1024 # GB
314 if memory <= self.small_memory:
315 return "small"
316 if memory <= self.medium_memory:
317 return "medium"
318 return "large"