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