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-05 01:27 -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/>. 

27 

28from typing import TYPE_CHECKING, Any 

29 

30from parsl.executors import HighThroughputExecutor 

31from parsl.executors.base import ParslExecutor 

32from parsl.providers import SlurmProvider 

33 

34from ..configuration import get_bps_config_value, get_workflow_name 

35from ..site import SiteConfig 

36 

37if TYPE_CHECKING: 

38 from ..job import ParslJob 

39 

40__all__ = ("Slurm", "TripleSlurm") 

41 

42 

43Kwargs = dict[str, Any] 

44 

45 

46class Slurm(SiteConfig): 

47 """Configuration for generic Slurm cluster. 

48 

49 Parameters 

50 ---------- 

51 *args : `~typing.Any` 

52 Parameters forwarded to base class constructor. 

53 **kwargs : `~typing.Any` 

54 Keyword arguments passed to base class constructor, augmented by 

55 the ``resource_list`` argument. 

56 

57 Notes 

58 ----- 

59 This can be used directly as the site configuration for a Slurm cluster by 

60 setting the BPS config, e.g.: 

61 

62 .. code-block:: yaml 

63 

64 computeSite: slurm 

65 site: 

66 slurm: 

67 class: lsst.ctrl.bps.parsl.sites.Slurm 

68 nodes: 3 

69 cores_per_node: 20 

70 walltime: "00:59:00" # Note: always quote walltime in YAML 

71 

72 Alternatively, it can be used as a base class for Slurm cluster 

73 configurations. 

74 

75 The following BPS configuration parameters are recognised (and required 

76 unless there is a default mentioned here, or provided by a subclass): 

77 

78 - ``nodes`` (`int`): number of nodes for each Slurm job. 

79 - ``cores_per_node`` (`int`): number of cores per node for each Slurm job; 

80 by default we use all cores on the node. 

81 - ``walltime`` (`str`): time limit for each Slurm job. 

82 - ``mem_per_node`` (`int`): memory per node (GB) for each Slurm job; by 

83 default we use whatever Slurm gives us. 

84 - ``qos`` (`str`): quality of service to request for each Slurm job; by 

85 default we use whatever Slurm gives us. 

86 - ``singleton`` (`bool`): allow only one job to run at a time; by default 

87 ``False``. 

88 - ``account`` (`str`): account to use for Slurm jobs. 

89 - ``scheduler_options`` (`str`): text to prepend to the Slurm submission 

90 script (each line usually starting with ``#SBATCH``). 

91 """ 

92 

93 def __init__(self, *args, **kwargs): 

94 # Have BPS-defined resource requests for each job passed to executor. 

95 kwargs["resource_list"] = ["priority"] 

96 super().__init__(*args, **kwargs) 

97 

98 def make_executor( 

99 self, 

100 label: str, 

101 *, 

102 nodes: int | None = None, 

103 cores_per_node: int | None = None, 

104 walltime: str | None = None, 

105 mem_per_node: int | None = None, 

106 mem_per_worker: float | None = None, 

107 qos: str | None = None, 

108 constraint: str | None = None, 

109 singleton: bool = False, 

110 scheduler_options: str | None = None, 

111 provider_options: Kwargs | None = None, 

112 executor_options: Kwargs | None = None, 

113 ) -> ParslExecutor: 

114 """Return an executor for running on a Slurm cluster. 

115 

116 Parameters 

117 ---------- 

118 label : `str` 

119 Label for executor. 

120 nodes : `int`, optional 

121 Default number of nodes for each Slurm job. 

122 cores_per_node : `int`, optional 

123 Default number of cores per node for each Slurm job. 

124 walltime : `str`, optional 

125 Default time limit for each Slurm job. 

126 mem_per_node : `float`, optional 

127 Memory per node (GB) to request for each Slurm job. 

128 mem_per_worker : `float`, optional 

129 Minimum memory per worker (GB), limited by the executor. 

130 qos : `str`, optional 

131 Quality of service for each Slurm job. 

132 constraint : `str`, optional 

133 Node feature(s) to require for each Slurm job. 

134 singleton : `bool`, optional 

135 Wether to allow only a single Slurm job to run at a time. 

136 scheduler_options : `str`, optional 

137 ``#SBATCH`` directives to prepend to the Slurm submission script. 

138 provider_options : `dict`, optional 

139 Additional arguments for `parsl.providers.SlurmProvider` 

140 constructor. 

141 executor_options : `dict`, optional 

142 Additional arguments for `parsl.executors.HighThroughputExecutor` 

143 constructor. 

144 

145 Returns 

146 ------- 

147 executor : `parsl.executors.HighThroughputExecutor` 

148 Executor for Slurm jobs. 

149 """ 

150 nodes = get_bps_config_value(self.site, "nodes", int, nodes, required=True) 

151 cores_per_node = get_bps_config_value(self.site, "cores_per_node", int, cores_per_node) 

152 walltime = get_bps_config_value(self.site, "walltime", str, walltime, required=True) 

153 mem_per_node = get_bps_config_value(self.site, "mem_per_node", int, mem_per_node) 

154 qos = get_bps_config_value(self.site, "qos", str, qos) 

155 singleton = get_bps_config_value(self.site, "singleton", bool, singleton) 

156 account = get_bps_config_value(self.site, "account", str) 

157 scheduler_options = get_bps_config_value(self.site, "scheduler_options", str, scheduler_options) 

158 

159 job_name = get_workflow_name(self.config) 

160 if scheduler_options is None: 

161 scheduler_options = "" 

162 else: 

163 scheduler_options += "\n" 

164 scheduler_options += f"#SBATCH --job-name={job_name}\n" 

165 if qos: 

166 scheduler_options += f"#SBATCH --qos={qos}\n" 

167 if constraint: 

168 scheduler_options += f"#SBATCH --constraint={constraint}\n" 

169 if singleton: 

170 # The following SBATCH directives allow only a single slurm job 

171 # (parsl block) with our job_name to run at once. This means we can 

172 # have one job running, and one already in the queue when the first 

173 # exceeds the walltime limit. More backups could be achieved with a 

174 # larger value of max_blocks. This only allows one job to be 

175 # actively running at once, so that needs to be sized appropriately 

176 # by the user. 

177 scheduler_options += "#SBATCH --dependency=singleton\n" 

178 return HighThroughputExecutor( 

179 label, 

180 provider=SlurmProvider( 

181 nodes_per_block=nodes, 

182 cores_per_node=cores_per_node, 

183 mem_per_node=mem_per_node, 

184 walltime=walltime, 

185 account=account, 

186 scheduler_options=scheduler_options, 

187 **(provider_options or {}), 

188 ), 

189 mem_per_worker=mem_per_worker, 

190 address=self.get_address(), 

191 **(executor_options or {}), 

192 ) 

193 

194 def get_executors(self) -> list[ParslExecutor]: 

195 """Get a list of executors to be used in processing. 

196 

197 Each executor should have a unique ``label``. 

198 """ 

199 return [self.make_executor("slurm")] 

200 

201 def select_executor(self, job: "ParslJob") -> str: 

202 """Get the ``label`` of the executor to use to execute a job. 

203 

204 Parameters 

205 ---------- 

206 job : `lsst.ctrl.bps.parsl.ParslJob` 

207 Job to be executed. 

208 

209 Returns 

210 ------- 

211 label : `str` 

212 Label of executor to use to execute ``job``. 

213 """ 

214 return "slurm" 

215 

216 

217class TripleSlurm(Slurm): 

218 """Configuration for running jobs on a Slurm cluster with three levels. 

219 

220 Parameters 

221 ---------- 

222 *args : `~typing.Any` 

223 Parameters forwarded to base class constructor. 

224 **kwargs : `~typing.Any` 

225 Keyword arguments passed to base class constructor. 

226 

227 Notes 

228 ----- 

229 The three levels are useful for having workers with different amount of 

230 available memory (and this is how executors are selected, by default), 

231 though other uses are possible. 

232 

233 The following BPS configuration parameters are recognised, overriding the 

234 defaults: 

235 

236 - ``nodes`` (`int`): number of nodes for each Slurm job. 

237 - ``cores_per_node`` (`int`): number of cores per node for each Slurm job; 

238 by default we use all cores on the node. 

239 - ``walltime`` (`str`): time limit for each Slurm job; setting this would 

240 override each of the ``small_walltime``, ``medium_walltime`` and 

241 ``large_walltime`` values. 

242 - ``mem_per_node`` (`float`): memory per node for each Slurm job; by 

243 default we use whatever Slurm gives us. 

244 - ``qos`` (`str`): quality of service to request for each Slurm job; by 

245 default we use whatever Slurm gives us. 

246 - ``small_memory`` (`float`): memory per worker (GB) for each 'small' Slurm 

247 job. 

248 - ``medium_memory`` (`float`): memory per worker (GB) for each 'medium' 

249 Slurm job. 

250 - ``large_memory`` (`float`): memory per worker (GB) for each 'large' Slurm 

251 job. 

252 - ``small_walltime`` (`str`): time limit for each 'small' Slurm job. 

253 - ``medium_walltime`` (`str`): time limit for each 'medium' Slurm job. 

254 - ``large_walltime`` (`str`): time limit for each 'large' Slurm job. 

255 """ 

256 

257 def __init__(self, *args, **kwargs): 

258 super().__init__(*args, **kwargs) 

259 self.small_memory = get_bps_config_value(self.site, "small_memory", float, 2.0) 

260 self.medium_memory = get_bps_config_value(self.site, "medium_memory", float, 4.0) 

261 self.large_memory = get_bps_config_value(self.site, "large_memory", float, 8.0) 

262 self.small_walltime = get_bps_config_value(self.site, "small_walltime", str, "10:00:00") 

263 self.medium_walltime = get_bps_config_value(self.site, "medium_walltime", str, "10:00:00") 

264 self.large_walltime = get_bps_config_value(self.site, "large_walltime", str, "40:00:00") 

265 

266 def get_executors( 

267 self, 

268 small_options: Kwargs | None = None, 

269 medium_options: Kwargs | None = None, 

270 large_options: Kwargs | None = None, 

271 **common_options, 

272 ) -> list[ParslExecutor]: 

273 """Get a list of executors to be used in processing. 

274 

275 We create three executors, with different walltime and memory per 

276 worker. 

277 

278 Parameters 

279 ---------- 

280 small_options : `dict` [`str`, `typing.Any`] 

281 Options for ``make_executor`` for small executor. 

282 medium_options : `dict` [`str`, `typing.Any`] 

283 Options for ``make_executor`` for medium executor. 

284 large_options : `dict` [`str`, `typing.Any`] 

285 Options for ``make_executor`` for large executor. 

286 **common_options 

287 Common options for ``make_executor`` for each of the executors. 

288 """ 

289 if small_options is None: 

290 small_options = {} 

291 if medium_options is None: 

292 medium_options = {} 

293 if large_options is None: 

294 large_options = {} 

295 

296 small_options["walltime"] = small_options.get("walltime", self.small_walltime) 

297 medium_options["walltime"] = medium_options.get("walltime", self.medium_walltime) 

298 large_options["walltime"] = large_options.get("walltime", self.large_walltime) 

299 

300 small_options["mem_per_worker"] = small_options.get("mem_per_worker", self.small_memory) 

301 medium_options["mem_per_worker"] = medium_options.get("mem_per_worker", self.medium_memory) 

302 large_options["mem_per_worker"] = large_options.get("mem_per_worker", self.large_memory) 

303 

304 return [ 

305 self.make_executor("small", **small_options, **(common_options or {})), 

306 self.make_executor("medium", **medium_options, **(common_options or {})), 

307 self.make_executor("large", **large_options, **(common_options or {})), 

308 ] 

309 

310 def select_executor(self, job: "ParslJob") -> str: 

311 """Get the ``label`` of the executor to use to execute a job. 

312 

313 This implementation only looks at the requested memory. 

314 

315 Parameters 

316 ---------- 

317 job : `lsst.ctrl.bps.parsl.ParslJob` 

318 Job to be executed. 

319 

320 Returns 

321 ------- 

322 label : `str` 

323 Label of executor to use to execute ``job``. 

324 """ 

325 memory = job.generic.request_memory / 1024 # GB 

326 if memory <= self.small_memory: 

327 return "small" 

328 if memory <= self.medium_memory: 

329 return "medium" 

330 return "large"