Coverage for python / lsst / ctrl / execute / allocator.py: 19%
200 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-01 01:59 -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 logging
29import os
30import pwd
31import sys
32from datetime import datetime
33from string import Template
35from lsst.ctrl.execute.allocationConfig import AllocationConfig
36from lsst.ctrl.execute.condorInfoConfig import CondorInfoConfig
37from lsst.ctrl.execute.templateWriter import TemplateWriter
38from lsst.resources import ResourcePath, ResourcePathExpression
40_LOG = logging.getLogger(__name__)
43class Allocator:
44 """A class which consolidates allocation pex_config information with
45 override information (obtained from the command line) and produces a
46 PBS file using these values.
48 Parameters
49 ----------
50 platform : `str`
51 the name of the platform to execute on
52 opts : `Config`
53 Config object containing options
54 condorInfoFileName : `lsst.resources.ResourcePathExpression`
55 Name of the file containing Config information
57 Raises
58 ------
59 TypeError
60 If the condorInfoFileName is the wrong type.
61 """
63 def __init__(
64 self,
65 platform: str,
66 opts,
67 configuration,
68 condorInfoFileName: ResourcePathExpression,
69 ):
70 """Constructor
71 @param platform: target platform for PBS submission
72 @param opts: options to override
73 """
74 self.opts = opts
75 self.defaults = {}
76 self.configuration = configuration
78 condorInfoConfig = CondorInfoConfig()
79 condorInfoConfig.load(condorInfoFileName)
81 self.platform = platform
83 # Look up the user's name and home and scratch directory in the
84 # $HOME/.lsst/condor-info.py file
85 user_name = None
86 user_home = None
87 user_scratch = None
88 for name in condorInfoConfig.platform:
89 if name == self.platform:
90 user_name = condorInfoConfig.platform[name].user.name
91 user_home = condorInfoConfig.platform[name].user.home
92 user_scratch = condorInfoConfig.platform[name].user.scratch
93 if user_scratch is None and "SCRATCH" in os.environ:
94 user_scratch = os.environ["SCRATCH"]
95 if user_name is None:
96 raise RuntimeError(
97 f"error: {condorInfoFileName} does not specify user name for platform == {self.platform}"
98 )
99 if user_home is None:
100 raise RuntimeError(
101 f"error: {condorInfoFileName} does not specify user home for platform == {self.platform}"
102 )
103 if user_scratch is None:
104 raise RuntimeError(
105 f"error: {condorInfoFileName} does not specify user scratch for platform == {self.platform}"
106 )
107 self.defaults["USER_NAME"] = user_name
108 self.defaults["USER_HOME"] = user_home
109 self.defaults["USER_SCRATCH"] = user_scratch
110 self.commandLineDefaults = {}
111 self.commandLineDefaults["NODE_COUNT"] = self.opts.nodeCount
112 self.commandLineDefaults["COLLECTOR"] = self.opts.collector
113 self.commandLineDefaults["CPORT"] = self.opts.collectorport
114 if self.opts.exclusive:
115 self.commandLineDefaults["CPUS"] = self.configuration.platform.peakcpus
116 else:
117 self.commandLineDefaults["CPUS"] = self.opts.cpus
118 self.commandLineDefaults["WALL_CLOCK"] = self.opts.maximumWallClock
119 self.commandLineDefaults["ACCOUNT"] = self.opts.account
120 self.commandLineDefaults["MEMPERCORE"] = 4096
121 self.commandLineDefaults["ALLOWEDAUTO"] = 500
122 self.commandLineDefaults["AUTOCPUS"] = 16
123 self.commandLineDefaults["MINAUTOCPUS"] = 15
124 self.commandLineDefaults["QUEUE"] = self.opts.queue
125 self.commandLineDefaults["NODESET"] = self.opts.nodeset
126 self.load()
128 def createUniqueIdentifier(self):
129 """Creates a unique file identifier, based on the user's name
130 and the time at which this method is invoked.
132 Returns
133 -------
134 ident : `str`
135 the new identifier
136 """
137 # This naming scheme follows the conventions used for creating
138 # RUNID names. We've found this allows these files to be more
139 # easily located and shared with other users when debugging
140 # The tempfile.mkstemp method restricts the file to only the user,
141 # and does not guarantee a file name can that easily be identified.
142 now = datetime.now()
143 self.defaults["DATE_STRING"] = f"{now.year:02d}_{now.month:02d}{now.day:02d}"
144 username = pwd.getpwuid(os.geteuid()).pw_name
145 ident = (
146 f"{username}_{now.year:02d}_{now.month:02d}{now.day:02d}_"
147 f"{now.hour:02d}{now.minute:02d}{now.second:02d}"
148 )
149 return ident
151 def load(self):
152 """Loads all values from configuration and command line overrides into
153 data structures suitable for use by the TemplateWriter object.
154 """
155 tempLocalScratch = Template(self.configuration.platform.localScratch)
156 self.defaults["LOCAL_SCRATCH"] = tempLocalScratch.substitute(
157 USER_SCRATCH=self.defaults["USER_SCRATCH"]
158 )
159 self.defaults["SCHEDULER"] = self.configuration.platform.scheduler
161 def loadAllocationConfig(self, name: ResourcePathExpression, suffix):
162 """Loads all values from allocationConfig and command line overrides
163 into data structures suitable for use by the TemplateWriter object.
164 """
165 if not (name_ := ResourcePath(name)).exists():
166 raise RuntimeError(f"{name_} was not found.")
167 allocationConfig = AllocationConfig()
168 allocationConfig.load(name_)
170 self.defaults["QUEUE"] = allocationConfig.platform.queue
171 self.defaults["EMAIL_NOTIFICATION"] = allocationConfig.platform.email
172 self.defaults["HOST_NAME"] = allocationConfig.platform.loginHostName
174 self.defaults["UTILITY_PATH"] = allocationConfig.platform.utilityPath
176 if self.opts.glideinShutdown is None:
177 self.defaults["GLIDEIN_SHUTDOWN"] = str(allocationConfig.platform.glideinShutdown)
178 else:
179 self.defaults["GLIDEIN_SHUTDOWN"] = str(self.opts.glideinShutdown)
181 if self.opts.outputLog is not None:
182 self.defaults["OUTPUT_LOG"] = self.opts.outputLog
183 else:
184 self.defaults["OUTPUT_LOG"] = "glide.out"
186 if self.opts.errorLog is not None:
187 self.defaults["ERROR_LOG"] = self.opts.errorLog
188 else:
189 self.defaults["ERROR_LOG"] = "glide.err"
191 # This is the TOTAL number of cores in the job, not just the total
192 # of the cores you intend to use. In other words, the total available
193 # on a machine, times the number of machines.
194 totalCoresPerNode = allocationConfig.platform.totalCoresPerNode
195 self.commandLineDefaults["TOTAL_CORE_COUNT"] = self.opts.nodeCount * totalCoresPerNode
197 self.uniqueIdentifier = self.createUniqueIdentifier()
199 # write these pbs and config files to {LOCAL_DIR}/configs
200 self.configDir = os.path.join(
201 self.defaults["LOCAL_SCRATCH"],
202 self.defaults["DATE_STRING"],
203 self.uniqueIdentifier,
204 "configs",
205 )
207 self.submitFileName = os.path.join(self.configDir, f"alloc_{self.uniqueIdentifier}.{suffix}")
209 self.condorConfigFileName = os.path.join(self.configDir, f"condor_{self.uniqueIdentifier}.config")
211 self.defaults["GENERATED_CONFIG"] = os.path.basename(self.condorConfigFileName)
212 self.defaults["CONFIGURATION_ID"] = self.uniqueIdentifier
213 return allocationConfig
215 def createSubmitFile(self, inputFile):
216 """Creates a batch submit file using the file "input" as a Template
218 Returns
219 -------
220 outfile : `str`
221 The newly created file name
222 """
223 if not os.path.exists(self.configDir):
224 os.makedirs(self.configDir)
225 outfile = self.createFile(inputFile, self.submitFileName)
226 _LOG.debug("Wrote new Slurm submit file to %s", outfile)
227 return outfile
229 def createCondorConfigFile(self, input):
230 """Creates a Condor config file using the file "input" as a Template
232 Returns
233 -------
234 outfile : `str`
235 The newly created file name
236 """
237 outfile = self.createFile(input, self.condorConfigFileName)
238 _LOG.debug("Wrote new condor configuration file to %s", outfile)
239 return outfile
241 def createFile(self, input: ResourcePathExpression, output: ResourcePathExpression):
242 """Creates a new file, using "input" as a Template, and writes the
243 new file to output.
245 Returns
246 -------
247 outfile : `str`
248 The newly created file name
249 """
250 _LOG.debug("Creating file from template using %s", input)
251 template = TemplateWriter()
252 # Uses the associative arrays of "defaults" and "commandLineDefaults"
253 # to write out the new file from the template.
254 # The commandLineDefaults override values in "defaults"
255 substitutes = self.defaults.copy()
256 for key in self.commandLineDefaults:
257 val = self.commandLineDefaults[key]
258 if val is not None:
259 substitutes[key] = self.commandLineDefaults[key]
260 template.rewrite(input, output, substitutes)
261 return output
263 def isVerbose(self):
264 """Status of the verbose flag
265 @return True if the flag was set, False otherwise
266 """
267 return self.opts.verbose
269 def isAuto(self):
270 """Status of the auto flag
271 @return True if the flag was set, False otherwise
272 """
273 return self.opts.auto
275 def getUserName(self):
276 """Accessor for USER_NAME
277 @return the value of USER_NAME
278 """
279 return self.getParameter("USER_NAME")
281 def getUserHome(self):
282 """Accessor for USER_HOME
283 @return the value of USER_HOME
284 """
285 return self.getParameter("USER_HOME")
287 def getUserScratch(self):
288 """Accessor for USER_SCRATCH
289 @return the value of USER_SCRATCH
290 """
291 return self.getParameter("USER_SCRATCH")
293 def getHostName(self):
294 """Accessor for HOST_NAME
295 @return the value of HOST_NAME
296 """
297 return self.getParameter("HOST_NAME")
299 def getUtilityPath(self):
300 """Accessor for UTILITY_PATH
301 @return the value of UTILITY_PATH
302 """
303 return self.getParameter("UTILITY_PATH")
305 def getScratchDirectory(self):
306 """Accessor for SCRATCH_DIR
307 @return the value of SCRATCH_DIR
308 """
309 return self.getParameter("SCRATCH_DIR")
311 def getLocalScratchDirectory(self):
312 """Accessor for LOCAL_SCRATCH
313 @return the value of LOCAL_SCRATCH
314 """
315 return self.getParameter("LOCAL_SCRATCH")
317 def getNodeSetName(self):
318 """Accessor for NODE_SET
319 @return the value of NODE_SET
320 """
321 return self.getParameter("NODE_SET")
323 def getNodes(self):
324 """Accessor for NODE_COUNT
325 @return the value of NODE_COUNT
326 """
327 return self.getParameter("NODE_COUNT")
329 def getMemoryPerCore(self):
330 """Accessor for MemoryPerCore
331 @return the value of MemoryPerCore
332 """
333 return self.getParameter("MEMPERCORE")
335 def getAllowedAutoGlideins(self):
336 """Accessor for AllowedAutoGlideins
337 @return the value of AllowedAuto
338 """
339 return self.getParameter("ALLOWEDAUTO")
341 def getQOS(self):
342 """Accessor for QOS
343 @return the value of QOS
344 """
345 return self.getParameter("QOS")
347 def getCPUs(self):
348 """Accessor for CPUS
349 @return the value of CPUS
350 """
351 return self.getParameter("CPUS")
353 def getAutoCPUs(self):
354 """Size of standard glideins for allocateNodes auto
355 @return the value of autoCPUs
356 """
357 if self.getParameter("EXCLUSIVE"):
358 peakcpus = self.configuration.platform.peakcpus
359 return peakcpus
360 else:
361 return self.getParameter("AUTOCPUS")
363 def getMinAutoCPUs(self):
364 """Minimum Size of standard glideins for allocateNodes auto
365 @return the value of minAutoCPUs
366 """
367 return self.getParameter("MINAUTOCPUS")
369 def getWallClock(self):
370 """Accessor for WALL_CLOCK
371 @return the value of WALL_CLOCK
372 """
373 return self.getParameter("WALL_CLOCK")
375 def getScheduler(self):
376 """Accessor for SCHEDULER
377 @return the value of SCHEDULER
378 """
379 return self.getParameter("SCHEDULER")
381 def getReservation(self):
382 """Accessor for RESERVATION
383 @return the value of RESERVATION
384 """
385 return self.getParameter("RESERVATION")
387 def getExclusive(self):
388 """Accessor for EXCLUSIVE
389 @return the value of EXCLUSIVE
390 """
391 return self.getParameter("EXCLUSIVE")
393 def getExcluser(self):
394 """Accessor for EXCLUSER
395 @return the value of EXCLUSER
396 """
397 return self.getParameter("EXCLUSER")
399 def getNodeset(self):
400 """Accessor for NODESET
401 @return the value of NODESET
402 """
403 return self.getParameter("NODESET")
405 def getParameter(self, value):
406 """Accessor for generic value
407 @return None if value is not set. Otherwise, use the command line
408 override (if set), or the default Config value
409 """
410 if value in self.commandLineDefaults:
411 return self.commandLineDefaults[value]
412 if value in self.defaults:
413 return self.defaults[value]
414 return None
416 def printNodeSetInfo(self):
417 nodes = self.getNodes()
418 cpus = self.getCPUs()
419 wallClock = self.getWallClock()
420 nodeString = ""
422 if int(nodes) > 1:
423 nodeString = "s"
424 if self.opts.dynamic is None:
425 print(
426 f"{nodes} glidein{nodeString} will be allocated on "
427 f"{self.platform} using default dynamic slots configuration."
428 )
429 print(f"There will be {cpus} cores per glidein and a maximum time limit of {wallClock}")
430 elif self.opts.dynamic == "__default__":
431 print(
432 f"{nodes} glidein{nodeString} will be allocated on {self.platform} "
433 "using default dynamic slots configuration."
434 )
435 print(f"There will be {cpus} cores per glidein and a maximum time limit of {wallClock}")
436 else:
437 print(
438 f"{nodes} node{nodeString} will be allocated on {self.platform} "
439 f"using dynamic slot block specified in '{self.opts.dynamic}'"
440 )
441 print(f"There will be {cpus} cores per node and maximum time limit of {wallClock}")
442 print("Node set name:")
443 print(self.getNodeSetName())
445 def runCommand(self, cmd, verbose):
446 cmd_split = cmd.split()
447 pid = os.fork()
448 if not pid:
449 # Methods of file transfer and login may
450 # produce different output, depending on how
451 # the "gsi" utilities are used. The user can
452 # either use grid proxies or ssh, and gsiscp/gsissh
453 # does the right thing. Since the output will be
454 # different in either case anything potentially parsing this
455 # output (like drpRun), would have to go through extra
456 # steps to deal with this output, and which ultimately
457 # end up not being useful. So we optinally close the i/o output
458 # of the executing command down.
459 #
460 # stdin/stdio/stderr is treated specially
461 # by python, so we have to close down
462 # both the python objects and the
463 # underlying c implementations
464 if not verbose:
465 # close python i/o
466 sys.stdin.close()
467 sys.stdout.close()
468 sys.stderr.close()
469 # close C's i/o
470 os.close(0)
471 os.close(1)
472 os.close(2)
473 os.execvp(cmd_split[0], cmd_split)
474 pid, status = os.wait()
475 # high order bits are status, low order bits are signal.
476 exitCode = (status & 0xFF00) >> 8
477 return exitCode
479 def submit(self):
480 """Submit the glidein jobs to the Batch system."""
481 raise NotImplementedError