Coverage for python / lsst / ctrl / execute / slurmPlugin.py: 7%
293 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-07 01:25 -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 anodeset = self.getNodeset()
168 if anodeset is None:
169 jobname = f"glide_{auser}"
170 else:
171 jobname = f"{anodeset}glide_{auser}"
172 _LOG.debug("The unix user name is %s", auser)
173 _LOG.debug("The Slurm job name for the glidein jobs is %s", jobname)
174 _LOG.debug("The user home directory is %s", self.getUserHome())
176 if auto:
177 self.glideinsFromJobPressure()
178 else:
179 generatedSlurmFile = self.createFilesFromTemplates()
180 cmd = f"sbatch --mem {totalMemory} {generatedSlurmFile}"
181 nodes = self.getNodes()
182 # In this case 'nodes' is the Target.
184 # Limit number of cores to be <= 8000 which 500 16-core glideins
185 # allowed auto glideins is 500
186 allowedAutoGlideins = self.getAllowedAutoGlideins()
187 # auto glidein size is 16
188 autoSize = self.getAutoCPUs()
189 targetedCores = nodes * cpus
190 coreLimit = allowedAutoGlideins * autoSize
191 if targetedCores > coreLimit:
192 # Reduce number of nodes because of threshold
193 nodes = int(coreLimit / cpus)
194 _LOG.info("Reducing number of glideins because of core limit threshold")
195 _LOG.debug("coreLimit %d", coreLimit)
196 _LOG.debug("glidein size %d", cpus)
197 _LOG.info("New number of glideins %d", nodes)
199 _LOG.info("Targeting %d glidein(s) for the computing pool/set.", nodes)
200 batcmd = "".join(["squeue --noheader --name=", jobname, " | wc -l"])
201 _LOG.debug("The squeue command is: %s", batcmd)
202 try:
203 result = subprocess.check_output(batcmd, shell=True)
204 except subprocess.CalledProcessError as e:
205 _LOG.error(e.output)
206 strResult = result.decode("UTF-8")
208 _LOG.info("Detected this number of preexisting glidein jobs: %d", int(strResult))
210 numberToAdd = nodes - int(strResult)
211 _LOG.info("The number of glidein jobs to submit now is %d", numberToAdd)
213 for glide in range(0, numberToAdd):
214 _LOG.info("Submitting glidein %d", glide)
215 exitCode = self.runCommand(cmd, verbose)
216 if exitCode != 0:
217 _LOG.error("error running %s", cmd)
218 sys.exit(exitCode)
220 def loadSlurm(self, name):
221 if self.opts.reservation is not None:
222 self.defaults["RESERVATION"] = f"#SBATCH --reservation {self.opts.reservation}"
223 else:
224 self.defaults["RESERVATION"] = ""
226 if self.opts.exclude is not None:
227 self.defaults["EXCLUDE"] = f"#SBATCH --exclude {self.opts.exclude}"
228 else:
229 self.defaults["EXCLUDE"] = ""
231 if self.opts.nodelist is not None:
232 self.defaults["NODELIST"] = f"#SBATCH --nodelist {self.opts.nodelist}"
233 else:
234 self.defaults["NODELIST"] = ""
236 if self.opts.exclusive is not None:
237 self.defaults["EXCLUSIVE"] = "#SBATCH --exclusive"
238 else:
239 self.defaults["EXCLUSIVE"] = ""
241 if self.opts.exclusiveUser is not None:
242 self.defaults["EXCLUSER"] = "#SBATCH --exclusive=user"
243 else:
244 self.defaults["EXCLUSER"] = ""
246 if self.opts.qos:
247 self.defaults["QOS"] = f"#SBATCH --qos {self.opts.qos}"
248 else:
249 self.defaults["QOS"] = ""
251 allocationConfig = self.loadAllocationConfig(name, "slurm")
253 template = Template(allocationConfig.platform.scratchDirectory)
254 scratchDir = template.substitute(USER_SCRATCH=self.getUserScratch())
255 self.defaults["SCRATCH_DIR"] = scratchDir
257 self.allocationFileName = Path(self.configDir) / f"allocation_{self.uniqueIdentifier}.sh"
258 self.defaults["GENERATED_ALLOCATE_SCRIPT"] = self.allocationFileName.name
260 if self.opts.openfiles is None:
261 self.defaults["OPEN_FILES"] = 20480
262 else:
263 self.defaults["OPEN_FILES"] = self.opts.openfiles
265 if self.opts.nodeset is None:
266 self.defaults["NODESET_BLOCK"] = "#"
267 self.defaults["NODESET"] = ""
268 else:
269 self.defaults["NODESET_BLOCK"] = f'Nodeset = "{self.opts.nodeset}"'
270 self.defaults["NODESET"] = f"{self.opts.nodeset}"
272 # For partitionable slots the classad 'Cpus' shows how many cpus
273 # remain to be allocated. Thus for a slot running jobs the value
274 # of Rank of TotalCpus - Cpus will increase with the load.
275 # Because higher Rank is preferred, loaded slots are favored.
276 if self.opts.packnodes is None:
277 self.defaults["PACK_BLOCK"] = "#"
278 else:
279 self.defaults["PACK_BLOCK"] = "Rank = TotalCpus - Cpus"
281 # handle dynamic slot block template:
282 # 1) if it isn't specified, just put a comment in its place
283 # 2) if it's specified, but without a filename, use the default
284 # 3) if it's specified with a filename, use that.
285 if self.opts.dynamic is None:
286 self.defaults["DYNAMIC_SLOTS_BLOCK"] = "#"
287 return
289 if self.opts.dynamic == "__default__":
290 dynamicSlotsName = find_package_file(
291 "dynamic_slots.template", kind="templates", platform=self.platform
292 )
293 else:
294 dynamicSlotsName = ResourcePath(self.opts.dynamic)
296 with dynamicSlotsName.open() as f:
297 lines = f.readlines()
298 block = ""
299 for line in lines:
300 block += line
301 self.defaults["DYNAMIC_SLOTS_BLOCK"] = block
303 def createAllocationFile(self, input: ResourcePathExpression):
304 """Creates Allocation script file using the file "input" as a Template
306 Returns
307 -------
308 outfile : `str`
309 The newly created file name
310 """
311 outfile = self.createFile(input, self.allocationFileName)
312 _LOG.debug("Wrote new Slurm job allocation bash script to %s", outfile)
313 os.chmod(outfile, 0o755)
314 return outfile
316 def glideinsFromJobPressure(self):
317 """Determine and submit the glideins needed from job pressure."""
319 verbose = self.isVerbose()
320 cpus = self.getCPUs()
321 autoCPUs = self.getAutoCPUs()
322 minAutoCPUs = self.getMinAutoCPUs()
323 if cpus >= minAutoCPUs:
324 autoCPUs = cpus
325 memoryPerCore = self.getMemoryPerCore()
326 memoryLimit = autoCPUs * memoryPerCore
327 auser = self.getUserName()
328 anodeset = self.getNodeset()
330 # projection contains the job classads to be returned.
331 # These include the cpu and memory profile of each job,
332 # in the form of RequestCpus and RequestMemory
333 projection = [
334 "ClusterId",
335 "ProcId",
336 "JobStatus",
337 "Owner",
338 "RequestCpus",
339 "JobUniverse",
340 "RequestMemory",
341 ]
342 owner = f'(Owner=="{auser}")'
343 # query for idle jobs
344 jstat = f"(JobStatus=={htcondor.JobStatus.IDLE})"
345 # query for vanilla universe
346 # JobUniverse constants are in htcondor C++
347 # UNIVERSE = { 1: "Standard", ..., 5: "Vanilla", ... }
348 juniv = "(JobUniverse==5)"
350 # The constraint determines that the jobs to be returned belong to
351 # the current user (Owner) and are Idle vanilla universe jobs.
352 full_constraint = f"{owner} && {jstat} && {juniv}"
353 if anodeset is None:
354 full_constraint += " && (JobNodeset is None)"
355 else:
356 jnodeset = f'(JobNodeset=="{anodeset}")'
357 full_constraint += f" && {jnodeset}"
358 _LOG.info("Auto: Query for htcondor jobs.")
359 _LOG.debug("full_constraint %s", full_constraint)
360 try:
361 condorq_data = condor_q(
362 constraint=full_constraint,
363 projection=projection,
364 )
366 except Exception as exc:
367 raise type(exc)("Problem querying condor schedd for jobs") from None
369 if not condorq_data:
370 _LOG.info("Auto: No HTCondor Jobs detected.")
371 return
373 generatedSlurmFile = self.createFilesFromTemplates()
374 condorq_large = []
375 condorq_small = []
376 schedd_name, condorq_full = condorq_data.popitem()
378 _LOG.info("Auto: Search for Large htcondor jobs.")
379 for jid, ajob in condorq_full.items():
380 thisCpus = ajob["RequestCpus"]
381 if isinstance(ajob["RequestMemory"], int):
382 thisEvalMemory = ajob["RequestMemory"]
383 else:
384 thisEvalMemory = ajob["RequestMemory"].eval()
385 _LOG.debug("Making an evaluation %s", thisEvalMemory)
386 # Search for jobs that are Large jobs
387 # thisCpus > 16 or thisEvalMemory > 16*4096
388 ajob["RequestMemoryEval"] = thisEvalMemory
389 if thisEvalMemory > memoryLimit or thisCpus > autoCPUs:
390 _LOG.info("Appending a Large Job %s", jid)
391 condorq_large.append(ajob)
392 else:
393 condorq_small.append(ajob)
395 if not condorq_large:
396 _LOG.info("Auto: no Large jobs detected.")
397 else:
398 _LOG.info("Auto: detected Large jobs")
399 for ajob in condorq_large:
400 _LOG.debug("\n%d.%d", ajob["ClusterId"], ajob["ProcId"])
401 _LOG.debug("%s", ajob)
402 thisMemory = ajob["RequestMemoryEval"]
403 useCores = ajob["RequestCpus"]
404 clusterid = ajob["ClusterId"]
405 procid = ajob["ProcId"]
406 job_label = f"{clusterid}_{procid}_{thisMemory}"
407 if useCores < autoCPUs:
408 useCores = autoCPUs
409 hash = hashlib.sha1(job_label.encode("UTF-8")).hexdigest()
410 shash = hash[:6]
411 if anodeset is None:
412 jobname = f"{auser}_{shash}"
413 else:
414 jobname = f"{anodeset}{auser}_{shash}"
415 _LOG.debug("jobname %s", jobname)
416 # Check if Job exists Idle in the queue
417 numberJobname = SlurmPlugin.countIdleSlurmJobs(jobname)
418 if numberJobname > 0:
419 _LOG.info("Job %s already exists, do not submit", jobname)
420 continue
421 cpuopt = f"--cpus-per-task {useCores}"
422 memopt = f"--mem {thisMemory}"
423 jobopt = f"-J {jobname}"
424 cmd = f"sbatch {cpuopt} {memopt} {jobopt} {generatedSlurmFile}"
425 _LOG.debug(cmd)
426 _LOG.info(
427 "Submitting Large glidein for %d.%d",
428 ajob["ClusterId"],
429 ajob["ProcId"],
430 )
431 time.sleep(3)
432 exitCode = self.runCommand(cmd, verbose)
433 if exitCode != 0:
434 _LOG.error("error running %s", cmd)
435 sys.exit(exitCode)
437 if not condorq_small:
438 _LOG.info("Auto: no small Jobs detected.")
439 else:
440 _LOG.info("Auto: summarize small jobs.")
441 maxNumberOfGlideins = self.getNodes()
442 maxAllowedNumberOfGlideins = self.getAllowedAutoGlideins()
443 _LOG.debug("maxNumberOfGlideins %d", maxNumberOfGlideins)
444 _LOG.debug("maxAllowedNumberOfGlideins %d", maxAllowedNumberOfGlideins)
445 # The number of cores for the small glideins is capped at 8000
446 # Corresponds to maxAllowedNumberOfGlideins = 500 16-core glideins
447 if maxNumberOfGlideins > maxAllowedNumberOfGlideins:
448 maxNumberOfGlideins = maxAllowedNumberOfGlideins
449 _LOG.info("Reducing Small Glidein limit due to threshold.")
450 #
451 # In the following loop we calculate the number of cores
452 # required by the set of small jobs. This calculation utilizes
453 # the requested cpus for a job, but also checks the requested
454 # memory and counts an effective core for each 'memoryPerCore'
455 # of memory (by default the 4GB per core of S3DF Slurm scheduler).
456 totalCores = 0
457 for ajob in condorq_small:
458 requestedCpus = ajob["RequestCpus"]
459 # if isinstance(ajob["RequestMemory"], int):
460 # requestedMemory = ajob["RequestMemory"]
461 # else:
462 # requestedMemory = ajob["RequestMemoryEval"]
463 # logging.debug("Using RequestMemoryEval")
464 requestedMemory = ajob["RequestMemoryEval"]
465 totalCores = totalCores + requestedCpus
466 _LOG.debug("small: jobid %d.%d", ajob["ClusterId"], ajob["ProcId"])
467 _LOG.debug("\tRequestCpus %d", requestedCpus)
468 _LOG.debug("\tCurrent value of totalCores %d", totalCores)
469 neededCpus = requestedMemory / memoryPerCore
470 if neededCpus > requestedCpus:
471 _LOG.debug("\t\tNeed to Add More:")
472 _LOG.debug("\t\tRequestMemory is %d", requestedMemory)
473 _LOG.debug("\t\tRatio to %d MB is %d", memoryPerCore, neededCpus)
474 totalCores = totalCores + (neededCpus - requestedCpus)
475 _LOG.debug("\t\tCurrent value of totalCores %d", totalCores)
477 _LOG.info("small: The final TotalCores is %d", totalCores)
479 # The number of Glideins needed to service the detected Idle jobs
480 # is "numberOfGlideins"
481 numberOfGlideins = math.ceil(totalCores / autoCPUs)
482 _LOG.info("small: Number for detected jobs is %d", numberOfGlideins)
484 if anodeset is None:
485 jobname = f"glide_{auser}"
486 else:
487 jobname = f"{anodeset}glide_{auser}"
489 # Check Slurm queue Running glideins
490 existingGlideinsRunning = SlurmPlugin.countRunningSlurmJobs(jobname)
492 # Check Slurm queue Idle Glideins
493 existingGlideinsIdle = SlurmPlugin.countIdleSlurmJobs(jobname)
495 _LOG.debug("small: existingGlideinsRunning %d", existingGlideinsRunning)
496 _LOG.debug("small: existingGlideinsIdle %d", existingGlideinsIdle)
498 # The number of Glideins needed to service the detected
499 # Idle jobs is "numberOfGlideins" less the existing Idle glideins
500 numberOfGlideinsReduced = numberOfGlideins - existingGlideinsIdle
501 _LOG.debug("small: Target Number to submit %d", numberOfGlideinsReduced)
503 # The maximum number of Glideins that we can submit with
504 # the imposed threshold (maxNumberOfGlideins)
505 # is maxSubmitGlideins
506 existingGlideins = existingGlideinsRunning + existingGlideinsIdle
507 maxSubmitGlideins = maxNumberOfGlideins - existingGlideins
508 _LOG.debug("small: maxNumberOfGlideins %d", maxNumberOfGlideins)
509 _LOG.debug("small: maxSubmitGlideins %d", maxSubmitGlideins)
511 # Reduce the number of Glideins to submit if threshold exceeded
512 if numberOfGlideinsReduced > maxSubmitGlideins:
513 numberOfGlideinsReduced = maxSubmitGlideins
514 _LOG.info("small: Reducing due to threshold.")
515 _LOG.debug("small: Number of Glideins to submit is %d", numberOfGlideinsReduced)
517 cpuopt = f"--cpus-per-task {autoCPUs}"
518 memopt = f"--mem {memoryLimit}"
519 jobopt = f"-J {jobname}"
520 cmd = f"sbatch {cpuopt} {memopt} {jobopt} {generatedSlurmFile}"
521 _LOG.debug(cmd)
522 for glide in range(0, numberOfGlideinsReduced):
523 _LOG.info("Submitting glidein %s", glide)
524 exitCode = self.runCommand(cmd, verbose)
525 if exitCode != 0:
526 _LOG.error("error running %s", cmd)
527 sys.exit(exitCode)
529 return