Coverage for python / lsst / ctrl / execute / slurmPlugin.py: 8%
276 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-13 01:44 -0800
1# This file is part of ctrl_execute.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.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/>.
28import hashlib
29import logging
30import math
31import os
32import subprocess
33import sys
34import time
35from pathlib import Path
36from string import Template
38import htcondor
40from lsst.ctrl.bps.htcondor import condor_q
41from lsst.ctrl.execute.allocator import Allocator
42from lsst.ctrl.execute.findPackageFile import find_package_file
43from lsst.resources import ResourcePath, ResourcePathExpression
45_LOG = logging.getLogger(__name__)
48class SlurmPlugin(Allocator):
49 @staticmethod
50 def countSlurmJobs(jobname, jobstates):
51 """Check Slurm queue for Glideins of given states
53 Parameters
54 ----------
55 jobname : `string`
56 Slurm jobname to be searched for via squeue.
57 jobstates : `string`
58 Slurm jobstates to be searched for via squeue.
60 Returns
61 -------
62 numberOfJobs : `int`
63 The number of Slurm jobs detected via squeue.
64 """
65 batcmd = f"squeue --noheader --states={jobstates} --name={jobname} | wc -l"
66 _LOG.debug("The squeue command is %s", batcmd)
67 time.sleep(3)
68 try:
69 resultPD = subprocess.check_output(batcmd, shell=True)
70 except subprocess.CalledProcessError as e:
71 _LOG.error(e.output)
72 numberOfJobs = int(resultPD.decode("UTF-8"))
73 return numberOfJobs
75 @staticmethod
76 def countIdleSlurmJobs(jobname):
77 """Check Slurm queue for Idle Glideins
79 Parameters
80 ----------
81 jobname : `string`
82 Slurm jobname to be searched for via squeue.
84 Returns
85 -------
86 numberOfJobs : `int`
87 The number of Slurm jobs detected via squeue.
88 """
89 _LOG.info("Checking if idle Slurm job %s exists:", jobname)
90 numberOfJobs = SlurmPlugin.countSlurmJobs(jobname, jobstates="PD")
91 return numberOfJobs
93 @staticmethod
94 def countRunningSlurmJobs(jobname):
95 """Check Slurm queue for Running Glideins
97 Parameters
98 ----------
99 jobname : `string`
100 Slurm jobname to be searched for via squeue.
102 Returns
103 -------
104 numberOfJobs : `int`
105 The number of Slurm jobs detected via squeue.
106 """
107 _LOG.info("Checking if running Slurm job %s exists:", jobname)
108 numberOfJobs = SlurmPlugin.countSlurmJobs(jobname, jobstates="R")
109 return numberOfJobs
111 def createFilesFromTemplates(self):
112 """Create the Slurm submit, script, and htcondor config files
114 Returns
115 -------
116 generatedSlurmFile : `str`
117 name of the Slurm job description file
118 """
120 scratchDirParam = self.getScratchDirectory()
121 template = Template(scratchDirParam)
122 template.substitute(USER_HOME=self.getUserHome())
124 # create the slurm submit file
125 slurmName = find_package_file("generic.slurm.template", kind="templates", platform=self.platform)
126 generatedSlurmFile = self.createSubmitFile(slurmName)
128 # create the condor configuration file
129 condorFile = find_package_file(
130 "glidein_condor_config.template", kind="templates", platform=self.platform
131 )
132 self.createCondorConfigFile(condorFile)
134 # create the script that the slurm submit file calls
135 allocationName = find_package_file("allocation.sh.template", kind="templates", platform=self.platform)
136 self.createAllocationFile(allocationName)
138 _LOG.debug("The generated Slurm submit file is %s", generatedSlurmFile)
140 return generatedSlurmFile
142 def submit(self):
143 """Submit the glidein jobs to the Batch system."""
144 configName = find_package_file("slurmConfig.py", platform=self.platform)
146 self.loadSlurm(configName)
147 verbose = self.isVerbose()
148 auto = self.isAuto()
150 cpus = self.getCPUs()
151 memoryPerCore = self.getMemoryPerCore()
152 totalMemory = cpus * memoryPerCore
154 # run the sbatch command
155 template = Template(self.getLocalScratchDirectory())
156 localScratchDir = Path(template.substitute(USER_SCRATCH=self.getUserScratch()))
157 slurmSubmitDir = localScratchDir / self.defaults["DATE_STRING"]
158 localScratchDir.mkdir(exist_ok=True)
159 slurmSubmitDir.mkdir(exist_ok=True)
160 os.chdir(slurmSubmitDir)
161 _LOG.debug(
162 "The working local scratch directory localScratchDir is %s ",
163 localScratchDir,
164 )
166 auser = self.getUserName()
167 jobname = f"glide_{auser}"
168 _LOG.debug("The unix user name is %s", auser)
169 _LOG.debug("The Slurm job name for the glidein jobs is %s", jobname)
170 _LOG.debug("The user home directory is %s", self.getUserHome())
172 if auto:
173 self.glideinsFromJobPressure()
174 else:
175 generatedSlurmFile = self.createFilesFromTemplates()
176 cmd = f"sbatch --mem {totalMemory} {generatedSlurmFile}"
177 nodes = self.getNodes()
178 # In this case 'nodes' is the Target.
180 # Limit number of cores to be <= 8000 which 500 16-core glideins
181 # allowed auto glideins is 500
182 allowedAutoGlideins = self.getAllowedAutoGlideins()
183 # auto glidein size is 16
184 autoSize = self.getAutoCPUs()
185 targetedCores = nodes * cpus
186 coreLimit = allowedAutoGlideins * autoSize
187 if targetedCores > coreLimit:
188 # Reduce number of nodes because of threshold
189 nodes = int(coreLimit / cpus)
190 _LOG.info("Reducing number of glideins because of core limit threshold")
191 _LOG.debug("coreLimit %d", coreLimit)
192 _LOG.debug("glidein size %d", cpus)
193 _LOG.info("New number of glideins %d", nodes)
195 _LOG.info("Targeting %d glidein(s) for the computing pool/set.", nodes)
196 batcmd = "".join(["squeue --noheader --name=", jobname, " | wc -l"])
197 _LOG.debug("The squeue command is: %s", batcmd)
198 try:
199 result = subprocess.check_output(batcmd, shell=True)
200 except subprocess.CalledProcessError as e:
201 _LOG.error(e.output)
202 strResult = result.decode("UTF-8")
204 _LOG.info("Detected this number of preexisting glidein jobs: %d", int(strResult))
206 numberToAdd = nodes - int(strResult)
207 _LOG.info("The number of glidein jobs to submit now is %d", numberToAdd)
209 for glide in range(0, numberToAdd):
210 _LOG.info("Submitting glidein %d", glide)
211 exitCode = self.runCommand(cmd, verbose)
212 if exitCode != 0:
213 _LOG.error("error running %s", cmd)
214 sys.exit(exitCode)
216 def loadSlurm(self, name):
217 if self.opts.reservation is not None:
218 self.defaults["RESERVATION"] = f"#SBATCH --reservation {self.opts.reservation}"
219 else:
220 self.defaults["RESERVATION"] = ""
222 if self.opts.exclude is not None:
223 self.defaults["EXCLUDE"] = f"#SBATCH --exclude {self.opts.exclude}"
224 else:
225 self.defaults["EXCLUDE"] = ""
227 if self.opts.nodelist is not None:
228 self.defaults["NODELIST"] = f"#SBATCH --nodelist {self.opts.nodelist}"
229 else:
230 self.defaults["NODELIST"] = ""
232 if self.opts.exclusive is not None:
233 self.defaults["EXCLUSIVE"] = "#SBATCH --exclusive"
234 else:
235 self.defaults["EXCLUSIVE"] = ""
237 if self.opts.exclusiveUser is not None:
238 self.defaults["EXCLUSER"] = "#SBATCH --exclusive=user"
239 else:
240 self.defaults["EXCLUSER"] = ""
242 if self.opts.qos:
243 self.defaults["QOS"] = f"#SBATCH --qos {self.opts.qos}"
244 else:
245 self.defaults["QOS"] = ""
247 allocationConfig = self.loadAllocationConfig(name, "slurm")
249 template = Template(allocationConfig.platform.scratchDirectory)
250 scratchDir = template.substitute(USER_SCRATCH=self.getUserScratch())
251 self.defaults["SCRATCH_DIR"] = scratchDir
253 self.allocationFileName = Path(self.configDir) / f"allocation_{self.uniqueIdentifier}.sh"
254 self.defaults["GENERATED_ALLOCATE_SCRIPT"] = self.allocationFileName.name
256 if self.opts.openfiles is None:
257 self.defaults["OPEN_FILES"] = 20480
258 else:
259 self.defaults["OPEN_FILES"] = self.opts.openfiles
261 # For partitionable slots the classad 'Cpus' shows how many cpus
262 # remain to be allocated. Thus for a slot running jobs the value
263 # of Rank of TotalCpus - Cpus will increase with the load.
264 # Because higher Rank is preferred, loaded slots are favored.
265 if self.opts.packnodes is None:
266 self.defaults["PACK_BLOCK"] = "#"
267 else:
268 self.defaults["PACK_BLOCK"] = "Rank = TotalCpus - Cpus"
270 # handle dynamic slot block template:
271 # 1) if it isn't specified, just put a comment in its place
272 # 2) if it's specified, but without a filename, use the default
273 # 3) if it's specified with a filename, use that.
274 if self.opts.dynamic is None:
275 self.defaults["DYNAMIC_SLOTS_BLOCK"] = "#"
276 return
278 if self.opts.dynamic == "__default__":
279 dynamicSlotsName = find_package_file(
280 "dynamic_slots.template", kind="templates", platform=self.platform
281 )
282 else:
283 dynamicSlotsName = ResourcePath(self.opts.dynamic)
285 with dynamicSlotsName.open() as f:
286 lines = f.readlines()
287 block = ""
288 for line in lines:
289 block += line
290 self.defaults["DYNAMIC_SLOTS_BLOCK"] = block
292 def createAllocationFile(self, input: ResourcePathExpression):
293 """Creates Allocation script file using the file "input" as a Template
295 Returns
296 -------
297 outfile : `str`
298 The newly created file name
299 """
300 outfile = self.createFile(input, self.allocationFileName)
301 _LOG.debug("Wrote new Slurm job allocation bash script to %s", outfile)
302 os.chmod(outfile, 0o755)
303 return outfile
305 def glideinsFromJobPressure(self):
306 """Determine and submit the glideins needed from job pressure."""
308 verbose = self.isVerbose()
309 cpus = self.getCPUs()
310 autoCPUs = self.getAutoCPUs()
311 minAutoCPUs = self.getMinAutoCPUs()
312 if cpus >= minAutoCPUs:
313 autoCPUs = cpus
314 memoryPerCore = self.getMemoryPerCore()
315 memoryLimit = autoCPUs * memoryPerCore
316 auser = self.getUserName()
318 # projection contains the job classads to be returned.
319 # These include the cpu and memory profile of each job,
320 # in the form of RequestCpus and RequestMemory
321 projection = [
322 "ClusterId",
323 "ProcId",
324 "JobStatus",
325 "Owner",
326 "RequestCpus",
327 "JobUniverse",
328 "RequestMemory",
329 ]
330 owner = f'(Owner=="{auser}")'
331 # query for idle jobs
332 jstat = f"(JobStatus=={htcondor.JobStatus.IDLE})"
333 # query for vanilla universe
334 # JobUniverse constants are in htcondor C++
335 # UNIVERSE = { 1: "Standard", ..., 5: "Vanilla", ... }
336 juniv = "(JobUniverse==5)"
338 # The constraint determines that the jobs to be returned belong to
339 # the current user (Owner) and are Idle vanilla universe jobs.
340 full_constraint = f"{owner} && {jstat} && {juniv}"
341 _LOG.info("Auto: Query for htcondor jobs.")
342 _LOG.debug("full_constraint %s", full_constraint)
343 try:
344 condorq_data = condor_q(
345 constraint=full_constraint,
346 projection=projection,
347 )
349 except Exception as exc:
350 raise type(exc)("Problem querying condor schedd for jobs") from None
352 if not condorq_data:
353 _LOG.info("Auto: No HTCondor Jobs detected.")
354 return
356 generatedSlurmFile = self.createFilesFromTemplates()
357 condorq_large = []
358 condorq_small = []
359 schedd_name, condorq_full = condorq_data.popitem()
361 _LOG.info("Auto: Search for Large htcondor jobs.")
362 for jid, ajob in condorq_full.items():
363 thisCpus = ajob["RequestCpus"]
364 if isinstance(ajob["RequestMemory"], int):
365 thisEvalMemory = ajob["RequestMemory"]
366 else:
367 thisEvalMemory = ajob["RequestMemory"].eval()
368 _LOG.debug("Making an evaluation %s", thisEvalMemory)
369 # Search for jobs that are Large jobs
370 # thisCpus > 16 or thisEvalMemory > 16*4096
371 ajob["RequestMemoryEval"] = thisEvalMemory
372 if thisEvalMemory > memoryLimit or thisCpus > autoCPUs:
373 _LOG.info("Appending a Large Job %s", jid)
374 condorq_large.append(ajob)
375 else:
376 condorq_small.append(ajob)
378 if not condorq_large:
379 _LOG.info("Auto: no Large jobs detected.")
380 else:
381 _LOG.info("Auto: detected Large jobs")
382 for ajob in condorq_large:
383 _LOG.debug("\n%d.%d", ajob["ClusterId"], ajob["ProcId"])
384 _LOG.debug("%s", ajob)
385 thisMemory = ajob["RequestMemoryEval"]
386 useCores = ajob["RequestCpus"]
387 clusterid = ajob["ClusterId"]
388 procid = ajob["ProcId"]
389 job_label = f"{clusterid}_{procid}_{thisMemory}"
390 if useCores < autoCPUs:
391 useCores = autoCPUs
392 hash = hashlib.sha1(job_label.encode("UTF-8")).hexdigest()
393 shash = hash[:6]
394 jobname = f"{auser}_{shash}"
395 _LOG.debug("jobname %s", jobname)
396 # Check if Job exists Idle in the queue
397 numberJobname = SlurmPlugin.countIdleSlurmJobs(jobname)
398 if numberJobname > 0:
399 _LOG.info("Job %s already exists, do not submit", jobname)
400 continue
401 cpuopt = f"--cpus-per-task {useCores}"
402 memopt = f"--mem {thisMemory}"
403 jobopt = f"-J {jobname}"
404 cmd = f"sbatch {cpuopt} {memopt} {jobopt} {generatedSlurmFile}"
405 _LOG.debug(cmd)
406 _LOG.info(
407 "Submitting Large glidein for %d.%d",
408 ajob["ClusterId"],
409 ajob["ProcId"],
410 )
411 time.sleep(3)
412 exitCode = self.runCommand(cmd, verbose)
413 if exitCode != 0:
414 _LOG.error("error running %s", cmd)
415 sys.exit(exitCode)
417 if not condorq_small:
418 _LOG.info("Auto: no small Jobs detected.")
419 else:
420 _LOG.info("Auto: summarize small jobs.")
421 maxNumberOfGlideins = self.getNodes()
422 maxAllowedNumberOfGlideins = self.getAllowedAutoGlideins()
423 _LOG.debug("maxNumberOfGlideins %d", maxNumberOfGlideins)
424 _LOG.debug("maxAllowedNumberOfGlideins %d", maxAllowedNumberOfGlideins)
425 # The number of cores for the small glideins is capped at 8000
426 # Corresponds to maxAllowedNumberOfGlideins = 500 16-core glideins
427 if maxNumberOfGlideins > maxAllowedNumberOfGlideins:
428 maxNumberOfGlideins = maxAllowedNumberOfGlideins
429 _LOG.info("Reducing Small Glidein limit due to threshold.")
430 #
431 # In the following loop we calculate the number of cores
432 # required by the set of small jobs. This calculation utilizes
433 # the requested cpus for a job, but also checks the requested
434 # memory and counts an effective core for each 'memoryPerCore'
435 # of memory (by default the 4GB per core of S3DF Slurm scheduler).
436 totalCores = 0
437 for ajob in condorq_small:
438 requestedCpus = ajob["RequestCpus"]
439 # if isinstance(ajob["RequestMemory"], int):
440 # requestedMemory = ajob["RequestMemory"]
441 # else:
442 # requestedMemory = ajob["RequestMemoryEval"]
443 # logging.debug("Using RequestMemoryEval")
444 requestedMemory = ajob["RequestMemoryEval"]
445 totalCores = totalCores + requestedCpus
446 _LOG.debug("small: jobid %d.%d", ajob["ClusterId"], ajob["ProcId"])
447 _LOG.debug("\tRequestCpus %d", requestedCpus)
448 _LOG.debug("\tCurrent value of totalCores %d", totalCores)
449 neededCpus = requestedMemory / memoryPerCore
450 if neededCpus > requestedCpus:
451 _LOG.debug("\t\tNeed to Add More:")
452 _LOG.debug("\t\tRequestMemory is %d", requestedMemory)
453 _LOG.debug("\t\tRatio to %d MB is %d", memoryPerCore, neededCpus)
454 totalCores = totalCores + (neededCpus - requestedCpus)
455 _LOG.debug("\t\tCurrent value of totalCores %d", totalCores)
457 _LOG.info("small: The final TotalCores is %d", totalCores)
459 # The number of Glideins needed to service the detected Idle jobs
460 # is "numberOfGlideins"
461 numberOfGlideins = math.ceil(totalCores / autoCPUs)
462 _LOG.info("small: Number for detected jobs is %d", numberOfGlideins)
464 jobname = f"glide_{auser}"
466 # Check Slurm queue Running glideins
467 existingGlideinsRunning = SlurmPlugin.countRunningSlurmJobs(jobname)
469 # Check Slurm queue Idle Glideins
470 existingGlideinsIdle = SlurmPlugin.countIdleSlurmJobs(jobname)
472 _LOG.debug("small: existingGlideinsRunning %d", existingGlideinsRunning)
473 _LOG.debug("small: existingGlideinsIdle %d", existingGlideinsIdle)
475 # The number of Glideins needed to service the detected
476 # Idle jobs is "numberOfGlideins" less the existing Idle glideins
477 numberOfGlideinsReduced = numberOfGlideins - existingGlideinsIdle
478 _LOG.debug("small: Target Number to submit %d", numberOfGlideinsReduced)
480 # The maximum number of Glideins that we can submit with
481 # the imposed threshold (maxNumberOfGlideins)
482 # is maxSubmitGlideins
483 existingGlideins = existingGlideinsRunning + existingGlideinsIdle
484 maxSubmitGlideins = maxNumberOfGlideins - existingGlideins
485 _LOG.debug("small: maxNumberOfGlideins %d", maxNumberOfGlideins)
486 _LOG.debug("small: maxSubmitGlideins %d", maxSubmitGlideins)
488 # Reduce the number of Glideins to submit if threshold exceeded
489 if numberOfGlideinsReduced > maxSubmitGlideins:
490 numberOfGlideinsReduced = maxSubmitGlideins
491 _LOG.info("small: Reducing due to threshold.")
492 _LOG.debug("small: Number of Glideins to submit is %d", numberOfGlideinsReduced)
494 cpuopt = f"--cpus-per-task {autoCPUs}"
495 memopt = f"--mem {memoryLimit}"
496 jobopt = f"-J {jobname}"
497 cmd = f"sbatch {cpuopt} {memopt} {jobopt} {generatedSlurmFile}"
498 _LOG.debug(cmd)
499 for glide in range(0, numberOfGlideinsReduced):
500 _LOG.info("Submitting glidein %s", glide)
501 exitCode = self.runCommand(cmd, verbose)
502 if exitCode != 0:
503 _LOG.error("error running %s", cmd)
504 sys.exit(exitCode)
506 return