Coverage for python / lsst / ctrl / execute / allocator.py: 19%
197 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-08 00:38 -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.loadFromStream(ResourcePath(condorInfoFileName).read())
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.load()
127 def createUniqueIdentifier(self):
128 """Creates a unique file identifier, based on the user's name
129 and the time at which this method is invoked.
131 Returns
132 -------
133 ident : `str`
134 the new identifier
135 """
136 # This naming scheme follows the conventions used for creating
137 # RUNID names. We've found this allows these files to be more
138 # easily located and shared with other users when debugging
139 # The tempfile.mkstemp method restricts the file to only the user,
140 # and does not guarantee a file name can that easily be identified.
141 now = datetime.now()
142 self.defaults["DATE_STRING"] = f"{now.year:02d}_{now.month:02d}{now.day:02d}"
143 username = pwd.getpwuid(os.geteuid()).pw_name
144 ident = (
145 f"{username}_{now.year:02d}_{now.month:02d}{now.day:02d}_"
146 f"{now.hour:02d}{now.minute:02d}{now.second:02d}"
147 )
148 return ident
150 def load(self):
151 """Loads all values from configuration and command line overrides into
152 data structures suitable for use by the TemplateWriter object.
153 """
154 tempLocalScratch = Template(self.configuration.platform.localScratch)
155 self.defaults["LOCAL_SCRATCH"] = tempLocalScratch.substitute(
156 USER_SCRATCH=self.defaults["USER_SCRATCH"]
157 )
158 self.defaults["SCHEDULER"] = self.configuration.platform.scheduler
160 def loadAllocationConfig(self, name: ResourcePathExpression, suffix):
161 """Loads all values from allocationConfig and command line overrides
162 into data structures suitable for use by the TemplateWriter object.
163 """
164 if not (name_ := ResourcePath(name)).exists():
165 raise RuntimeError(f"{name_} was not found.")
166 allocationConfig = AllocationConfig()
167 allocationConfig.loadFromStream(name_.read())
169 self.defaults["QUEUE"] = allocationConfig.platform.queue
170 self.defaults["EMAIL_NOTIFICATION"] = allocationConfig.platform.email
171 self.defaults["HOST_NAME"] = allocationConfig.platform.loginHostName
173 self.defaults["UTILITY_PATH"] = allocationConfig.platform.utilityPath
175 if self.opts.glideinShutdown is None:
176 self.defaults["GLIDEIN_SHUTDOWN"] = str(allocationConfig.platform.glideinShutdown)
177 else:
178 self.defaults["GLIDEIN_SHUTDOWN"] = str(self.opts.glideinShutdown)
180 if self.opts.outputLog is not None:
181 self.defaults["OUTPUT_LOG"] = self.opts.outputLog
182 else:
183 self.defaults["OUTPUT_LOG"] = "glide.out"
185 if self.opts.errorLog is not None:
186 self.defaults["ERROR_LOG"] = self.opts.errorLog
187 else:
188 self.defaults["ERROR_LOG"] = "glide.err"
190 # This is the TOTAL number of cores in the job, not just the total
191 # of the cores you intend to use. In other words, the total available
192 # on a machine, times the number of machines.
193 totalCoresPerNode = allocationConfig.platform.totalCoresPerNode
194 self.commandLineDefaults["TOTAL_CORE_COUNT"] = self.opts.nodeCount * totalCoresPerNode
196 self.uniqueIdentifier = self.createUniqueIdentifier()
198 # write these pbs and config files to {LOCAL_DIR}/configs
199 self.configDir = os.path.join(
200 self.defaults["LOCAL_SCRATCH"],
201 self.defaults["DATE_STRING"],
202 self.uniqueIdentifier,
203 "configs",
204 )
206 self.submitFileName = os.path.join(self.configDir, f"alloc_{self.uniqueIdentifier}.{suffix}")
208 self.condorConfigFileName = os.path.join(self.configDir, f"condor_{self.uniqueIdentifier}.config")
210 self.defaults["GENERATED_CONFIG"] = os.path.basename(self.condorConfigFileName)
211 self.defaults["CONFIGURATION_ID"] = self.uniqueIdentifier
212 return allocationConfig
214 def createSubmitFile(self, inputFile):
215 """Creates a batch submit file using the file "input" as a Template
217 Returns
218 -------
219 outfile : `str`
220 The newly created file name
221 """
222 if not os.path.exists(self.configDir):
223 os.makedirs(self.configDir)
224 outfile = self.createFile(inputFile, self.submitFileName)
225 _LOG.debug("Wrote new Slurm submit file to %s", outfile)
226 return outfile
228 def createCondorConfigFile(self, input):
229 """Creates a Condor config file using the file "input" as a Template
231 Returns
232 -------
233 outfile : `str`
234 The newly created file name
235 """
236 outfile = self.createFile(input, self.condorConfigFileName)
237 _LOG.debug("Wrote new condor configuration file to %s", outfile)
238 return outfile
240 def createFile(self, input: ResourcePathExpression, output: ResourcePathExpression):
241 """Creates a new file, using "input" as a Template, and writes the
242 new file to output.
244 Returns
245 -------
246 outfile : `str`
247 The newly created file name
248 """
249 _LOG.debug("Creating file from template using %s", input)
250 template = TemplateWriter()
251 # Uses the associative arrays of "defaults" and "commandLineDefaults"
252 # to write out the new file from the template.
253 # The commandLineDefaults override values in "defaults"
254 substitutes = self.defaults.copy()
255 for key in self.commandLineDefaults:
256 val = self.commandLineDefaults[key]
257 if val is not None:
258 substitutes[key] = self.commandLineDefaults[key]
259 template.rewrite(input, output, substitutes)
260 return output
262 def isVerbose(self):
263 """Status of the verbose flag
264 @return True if the flag was set, False otherwise
265 """
266 return self.opts.verbose
268 def isAuto(self):
269 """Status of the auto flag
270 @return True if the flag was set, False otherwise
271 """
272 return self.opts.auto
274 def getUserName(self):
275 """Accessor for USER_NAME
276 @return the value of USER_NAME
277 """
278 return self.getParameter("USER_NAME")
280 def getUserHome(self):
281 """Accessor for USER_HOME
282 @return the value of USER_HOME
283 """
284 return self.getParameter("USER_HOME")
286 def getUserScratch(self):
287 """Accessor for USER_SCRATCH
288 @return the value of USER_SCRATCH
289 """
290 return self.getParameter("USER_SCRATCH")
292 def getHostName(self):
293 """Accessor for HOST_NAME
294 @return the value of HOST_NAME
295 """
296 return self.getParameter("HOST_NAME")
298 def getUtilityPath(self):
299 """Accessor for UTILITY_PATH
300 @return the value of UTILITY_PATH
301 """
302 return self.getParameter("UTILITY_PATH")
304 def getScratchDirectory(self):
305 """Accessor for SCRATCH_DIR
306 @return the value of SCRATCH_DIR
307 """
308 return self.getParameter("SCRATCH_DIR")
310 def getLocalScratchDirectory(self):
311 """Accessor for LOCAL_SCRATCH
312 @return the value of LOCAL_SCRATCH
313 """
314 return self.getParameter("LOCAL_SCRATCH")
316 def getNodeSetName(self):
317 """Accessor for NODE_SET
318 @return the value of NODE_SET
319 """
320 return self.getParameter("NODE_SET")
322 def getNodes(self):
323 """Accessor for NODE_COUNT
324 @return the value of NODE_COUNT
325 """
326 return self.getParameter("NODE_COUNT")
328 def getMemoryPerCore(self):
329 """Accessor for MemoryPerCore
330 @return the value of MemoryPerCore
331 """
332 return self.getParameter("MEMPERCORE")
334 def getAllowedAutoGlideins(self):
335 """Accessor for AllowedAutoGlideins
336 @return the value of AllowedAuto
337 """
338 return self.getParameter("ALLOWEDAUTO")
340 def getQOS(self):
341 """Accessor for QOS
342 @return the value of QOS
343 """
344 return self.getParameter("QOS")
346 def getCPUs(self):
347 """Accessor for CPUS
348 @return the value of CPUS
349 """
350 return self.getParameter("CPUS")
352 def getAutoCPUs(self):
353 """Size of standard glideins for allocateNodes auto
354 @return the value of autoCPUs
355 """
356 if self.getParameter("EXCLUSIVE"):
357 peakcpus = self.configuration.platform.peakcpus
358 return peakcpus
359 else:
360 return self.getParameter("AUTOCPUS")
362 def getMinAutoCPUs(self):
363 """Minimum Size of standard glideins for allocateNodes auto
364 @return the value of minAutoCPUs
365 """
366 return self.getParameter("MINAUTOCPUS")
368 def getWallClock(self):
369 """Accessor for WALL_CLOCK
370 @return the value of WALL_CLOCK
371 """
372 return self.getParameter("WALL_CLOCK")
374 def getScheduler(self):
375 """Accessor for SCHEDULER
376 @return the value of SCHEDULER
377 """
378 return self.getParameter("SCHEDULER")
380 def getReservation(self):
381 """Accessor for RESERVATION
382 @return the value of RESERVATION
383 """
384 return self.getParameter("RESERVATION")
386 def getExclusive(self):
387 """Accessor for EXCLUSIVE
388 @return the value of EXCLUSIVE
389 """
390 return self.getParameter("EXCLUSIVE")
392 def getExcluser(self):
393 """Accessor for EXCLUSER
394 @return the value of EXCLUSER
395 """
396 return self.getParameter("EXCLUSER")
398 def getParameter(self, value):
399 """Accessor for generic value
400 @return None if value is not set. Otherwise, use the command line
401 override (if set), or the default Config value
402 """
403 if value in self.commandLineDefaults:
404 return self.commandLineDefaults[value]
405 if value in self.defaults:
406 return self.defaults[value]
407 return None
409 def printNodeSetInfo(self):
410 nodes = self.getNodes()
411 cpus = self.getCPUs()
412 wallClock = self.getWallClock()
413 nodeString = ""
415 if int(nodes) > 1:
416 nodeString = "s"
417 if self.opts.dynamic is None:
418 print(
419 f"{nodes} glidein{nodeString} will be allocated on "
420 f"{self.platform} using default dynamic slots configuration."
421 )
422 print(f"There will be {cpus} cores per glidein and a maximum time limit of {wallClock}")
423 elif self.opts.dynamic == "__default__":
424 print(
425 f"{nodes} glidein{nodeString} will be allocated on {self.platform} "
426 "using default dynamic slots configuration."
427 )
428 print(f"There will be {cpus} cores per glidein and a maximum time limit of {wallClock}")
429 else:
430 print(
431 f"{nodes} node{nodeString} will be allocated on {self.platform} "
432 f"using dynamic slot block specified in '{self.opts.dynamic}'"
433 )
434 print(f"There will be {cpus} cores per node and maximum time limit of {wallClock}")
435 print("Node set name:")
436 print(self.getNodeSetName())
438 def runCommand(self, cmd, verbose):
439 cmd_split = cmd.split()
440 pid = os.fork()
441 if not pid:
442 # Methods of file transfer and login may
443 # produce different output, depending on how
444 # the "gsi" utilities are used. The user can
445 # either use grid proxies or ssh, and gsiscp/gsissh
446 # does the right thing. Since the output will be
447 # different in either case anything potentially parsing this
448 # output (like drpRun), would have to go through extra
449 # steps to deal with this output, and which ultimately
450 # end up not being useful. So we optinally close the i/o output
451 # of the executing command down.
452 #
453 # stdin/stdio/stderr is treated specially
454 # by python, so we have to close down
455 # both the python objects and the
456 # underlying c implementations
457 if not verbose:
458 # close python i/o
459 sys.stdin.close()
460 sys.stdout.close()
461 sys.stderr.close()
462 # close C's i/o
463 os.close(0)
464 os.close(1)
465 os.close(2)
466 os.execvp(cmd_split[0], cmd_split)
467 pid, status = os.wait()
468 # high order bits are status, low order bits are signal.
469 exitCode = (status & 0xFF00) >> 8
470 return exitCode
472 def submit(self):
473 """Submit the glidein jobs to the Batch system."""
474 raise NotImplementedError