Coverage for python/lsst/ctrl/execute/allocator.py: 19%
197 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-11 01:45 -0700
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-11 01:45 -0700
1#!/usr/bin/env python
3#
4# LSST Data Management System
5# Copyright 2008-2016 LSST Corporation.
6#
7# This product includes software developed by the
8# LSST Project (http://www.lsst.org/).
9#
10# This program is free software: you can redistribute it and/or modify
11# it under the terms of the GNU General Public License as published by
12# the Free Software Foundation, either version 3 of the License, or
13# (at your option) any later version.
14#
15# This program is distributed in the hope that it will be useful,
16# but WITHOUT ANY WARRANTY; without even the implied warranty of
17# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
18# GNU General Public License for more details.
19#
20# You should have received a copy of the LSST License Statement and
21# the GNU General Public License along with this program. If not,
22# see <http://www.lsstcorp.org/LegalNotices/>.
23#
25import logging
26import os
27import pwd
28import sys
29from datetime import datetime
30from string import Template
32from lsst.ctrl.execute.allocationConfig import AllocationConfig
33from lsst.ctrl.execute.condorInfoConfig import CondorInfoConfig
34from lsst.ctrl.execute.templateWriter import TemplateWriter
35from lsst.resources import ResourcePath, ResourcePathExpression
37_LOG = logging.getLogger(__name__)
40class Allocator:
41 """A class which consolidates allocation pex_config information with
42 override information (obtained from the command line) and produces a
43 PBS file using these values.
45 Parameters
46 ----------
47 platform : `str`
48 the name of the platform to execute on
49 opts : `Config`
50 Config object containing options
51 condorInfoFileName : `lsst.resources.ResourcePathExpression`
52 Name of the file containing Config information
54 Raises
55 ------
56 TypeError
57 If the condorInfoFileName is the wrong type.
58 """
60 def __init__(
61 self,
62 platform: str,
63 opts,
64 configuration,
65 condorInfoFileName: ResourcePathExpression,
66 ):
67 """Constructor
68 @param platform: target platform for PBS submission
69 @param opts: options to override
70 """
71 self.opts = opts
72 self.defaults = {}
73 self.configuration = configuration
75 condorInfoConfig = CondorInfoConfig()
76 condorInfoConfig.loadFromStream(ResourcePath(condorInfoFileName).read())
78 self.platform = platform
80 # Look up the user's name and home and scratch directory in the
81 # $HOME/.lsst/condor-info.py file
82 user_name = None
83 user_home = None
84 user_scratch = None
85 for name in condorInfoConfig.platform:
86 if name == self.platform:
87 user_name = condorInfoConfig.platform[name].user.name
88 user_home = condorInfoConfig.platform[name].user.home
89 user_scratch = condorInfoConfig.platform[name].user.scratch
90 if user_scratch is None and "SCRATCH" in os.environ:
91 user_scratch = os.environ["SCRATCH"]
92 if user_name is None:
93 raise RuntimeError(
94 f"error: {condorInfoFileName} does not specify user name for platform == {self.platform}"
95 )
96 if user_home is None:
97 raise RuntimeError(
98 f"error: {condorInfoFileName} does not specify user home for platform == {self.platform}"
99 )
100 if user_scratch is None:
101 raise RuntimeError(
102 f"error: {condorInfoFileName} does not specify user scratch for platform == {self.platform}"
103 )
104 self.defaults["USER_NAME"] = user_name
105 self.defaults["USER_HOME"] = user_home
106 self.defaults["USER_SCRATCH"] = user_scratch
107 self.commandLineDefaults = {}
108 self.commandLineDefaults["NODE_COUNT"] = self.opts.nodeCount
109 self.commandLineDefaults["COLLECTOR"] = self.opts.collector
110 self.commandLineDefaults["CPORT"] = self.opts.collectorport
111 if self.opts.exclusive:
112 self.commandLineDefaults["CPUS"] = self.configuration.platform.peakcpus
113 else:
114 self.commandLineDefaults["CPUS"] = self.opts.cpus
115 self.commandLineDefaults["WALL_CLOCK"] = self.opts.maximumWallClock
116 self.commandLineDefaults["ACCOUNT"] = self.opts.account
117 self.commandLineDefaults["MEMPERCORE"] = 4096
118 self.commandLineDefaults["ALLOWEDAUTO"] = 500
119 self.commandLineDefaults["AUTOCPUS"] = 16
120 self.commandLineDefaults["MINAUTOCPUS"] = 15
121 self.commandLineDefaults["QUEUE"] = self.opts.queue
122 self.load()
124 def createUniqueIdentifier(self):
125 """Creates a unique file identifier, based on the user's name
126 and the time at which this method is invoked.
128 Returns
129 -------
130 ident : `str`
131 the new identifier
132 """
133 # This naming scheme follows the conventions used for creating
134 # RUNID names. We've found this allows these files to be more
135 # easily located and shared with other users when debugging
136 # The tempfile.mkstemp method restricts the file to only the user,
137 # and does not guarantee a file name can that easily be identified.
138 now = datetime.now()
139 self.defaults["DATE_STRING"] = f"{now.year:02d}_{now.month:02d}{now.day:02d}"
140 username = pwd.getpwuid(os.geteuid()).pw_name
141 ident = (
142 f"{username}_{now.year:02d}_{now.month:02d}{now.day:02d}_"
143 f"{now.hour:02d}{now.minute:02d}{now.second:02d}"
144 )
145 return ident
147 def load(self):
148 """Loads all values from configuration and command line overrides into
149 data structures suitable for use by the TemplateWriter object.
150 """
151 tempLocalScratch = Template(self.configuration.platform.localScratch)
152 self.defaults["LOCAL_SCRATCH"] = tempLocalScratch.substitute(
153 USER_SCRATCH=self.defaults["USER_SCRATCH"]
154 )
155 self.defaults["SCHEDULER"] = self.configuration.platform.scheduler
157 def loadAllocationConfig(self, name: ResourcePathExpression, suffix):
158 """Loads all values from allocationConfig and command line overrides
159 into data structures suitable for use by the TemplateWriter object.
160 """
161 if not (name_ := ResourcePath(name)).exists():
162 raise RuntimeError(f"{name_} was not found.")
163 allocationConfig = AllocationConfig()
164 allocationConfig.loadFromStream(name_.read())
166 self.defaults["QUEUE"] = allocationConfig.platform.queue
167 self.defaults["EMAIL_NOTIFICATION"] = allocationConfig.platform.email
168 self.defaults["HOST_NAME"] = allocationConfig.platform.loginHostName
170 self.defaults["UTILITY_PATH"] = allocationConfig.platform.utilityPath
172 if self.opts.glideinShutdown is None:
173 self.defaults["GLIDEIN_SHUTDOWN"] = str(allocationConfig.platform.glideinShutdown)
174 else:
175 self.defaults["GLIDEIN_SHUTDOWN"] = str(self.opts.glideinShutdown)
177 if self.opts.outputLog is not None:
178 self.defaults["OUTPUT_LOG"] = self.opts.outputLog
179 else:
180 self.defaults["OUTPUT_LOG"] = "glide.out"
182 if self.opts.errorLog is not None:
183 self.defaults["ERROR_LOG"] = self.opts.errorLog
184 else:
185 self.defaults["ERROR_LOG"] = "glide.err"
187 # This is the TOTAL number of cores in the job, not just the total
188 # of the cores you intend to use. In other words, the total available
189 # on a machine, times the number of machines.
190 totalCoresPerNode = allocationConfig.platform.totalCoresPerNode
191 self.commandLineDefaults["TOTAL_CORE_COUNT"] = self.opts.nodeCount * totalCoresPerNode
193 self.uniqueIdentifier = self.createUniqueIdentifier()
195 # write these pbs and config files to {LOCAL_DIR}/configs
196 self.configDir = os.path.join(
197 self.defaults["LOCAL_SCRATCH"],
198 self.defaults["DATE_STRING"],
199 self.uniqueIdentifier,
200 "configs",
201 )
203 self.submitFileName = os.path.join(self.configDir, f"alloc_{self.uniqueIdentifier}.{suffix}")
205 self.condorConfigFileName = os.path.join(self.configDir, f"condor_{self.uniqueIdentifier}.config")
207 self.defaults["GENERATED_CONFIG"] = os.path.basename(self.condorConfigFileName)
208 self.defaults["CONFIGURATION_ID"] = self.uniqueIdentifier
209 return allocationConfig
211 def createSubmitFile(self, inputFile):
212 """Creates a batch submit file using the file "input" as a Template
214 Returns
215 -------
216 outfile : `str`
217 The newly created file name
218 """
219 if not os.path.exists(self.configDir):
220 os.makedirs(self.configDir)
221 outfile = self.createFile(inputFile, self.submitFileName)
222 _LOG.debug("Wrote new Slurm submit file to %s", outfile)
223 return outfile
225 def createCondorConfigFile(self, input):
226 """Creates a Condor config file using the file "input" as a Template
228 Returns
229 -------
230 outfile : `str`
231 The newly created file name
232 """
233 outfile = self.createFile(input, self.condorConfigFileName)
234 _LOG.debug("Wrote new condor configuration file to %s", outfile)
235 return outfile
237 def createFile(self, input: ResourcePathExpression, output: ResourcePathExpression):
238 """Creates a new file, using "input" as a Template, and writes the
239 new file to output.
241 Returns
242 -------
243 outfile : `str`
244 The newly created file name
245 """
246 _LOG.debug("Creating file from template using %s", input)
247 template = TemplateWriter()
248 # Uses the associative arrays of "defaults" and "commandLineDefaults"
249 # to write out the new file from the template.
250 # The commandLineDefaults override values in "defaults"
251 substitutes = self.defaults.copy()
252 for key in self.commandLineDefaults:
253 val = self.commandLineDefaults[key]
254 if val is not None:
255 substitutes[key] = self.commandLineDefaults[key]
256 template.rewrite(input, output, substitutes)
257 return output
259 def isVerbose(self):
260 """Status of the verbose flag
261 @return True if the flag was set, False otherwise
262 """
263 return self.opts.verbose
265 def isAuto(self):
266 """Status of the auto flag
267 @return True if the flag was set, False otherwise
268 """
269 return self.opts.auto
271 def getUserName(self):
272 """Accessor for USER_NAME
273 @return the value of USER_NAME
274 """
275 return self.getParameter("USER_NAME")
277 def getUserHome(self):
278 """Accessor for USER_HOME
279 @return the value of USER_HOME
280 """
281 return self.getParameter("USER_HOME")
283 def getUserScratch(self):
284 """Accessor for USER_SCRATCH
285 @return the value of USER_SCRATCH
286 """
287 return self.getParameter("USER_SCRATCH")
289 def getHostName(self):
290 """Accessor for HOST_NAME
291 @return the value of HOST_NAME
292 """
293 return self.getParameter("HOST_NAME")
295 def getUtilityPath(self):
296 """Accessor for UTILITY_PATH
297 @return the value of UTILITY_PATH
298 """
299 return self.getParameter("UTILITY_PATH")
301 def getScratchDirectory(self):
302 """Accessor for SCRATCH_DIR
303 @return the value of SCRATCH_DIR
304 """
305 return self.getParameter("SCRATCH_DIR")
307 def getLocalScratchDirectory(self):
308 """Accessor for LOCAL_SCRATCH
309 @return the value of LOCAL_SCRATCH
310 """
311 return self.getParameter("LOCAL_SCRATCH")
313 def getNodeSetName(self):
314 """Accessor for NODE_SET
315 @return the value of NODE_SET
316 """
317 return self.getParameter("NODE_SET")
319 def getNodes(self):
320 """Accessor for NODE_COUNT
321 @return the value of NODE_COUNT
322 """
323 return self.getParameter("NODE_COUNT")
325 def getMemoryPerCore(self):
326 """Accessor for MemoryPerCore
327 @return the value of MemoryPerCore
328 """
329 return self.getParameter("MEMPERCORE")
331 def getAllowedAutoGlideins(self):
332 """Accessor for AllowedAutoGlideins
333 @return the value of AllowedAuto
334 """
335 return self.getParameter("ALLOWEDAUTO")
337 def getQOS(self):
338 """Accessor for QOS
339 @return the value of QOS
340 """
341 return self.getParameter("QOS")
343 def getCPUs(self):
344 """Accessor for CPUS
345 @return the value of CPUS
346 """
347 return self.getParameter("CPUS")
349 def getAutoCPUs(self):
350 """Size of standard glideins for allocateNodes auto
351 @return the value of autoCPUs
352 """
353 if self.getParameter("EXCLUSIVE"):
354 peakcpus = self.configuration.platform.peakcpus
355 return peakcpus
356 else:
357 return self.getParameter("AUTOCPUS")
359 def getMinAutoCPUs(self):
360 """Minimum Size of standard glideins for allocateNodes auto
361 @return the value of minAutoCPUs
362 """
363 return self.getParameter("MINAUTOCPUS")
365 def getWallClock(self):
366 """Accessor for WALL_CLOCK
367 @return the value of WALL_CLOCK
368 """
369 return self.getParameter("WALL_CLOCK")
371 def getScheduler(self):
372 """Accessor for SCHEDULER
373 @return the value of SCHEDULER
374 """
375 return self.getParameter("SCHEDULER")
377 def getReservation(self):
378 """Accessor for RESERVATION
379 @return the value of RESERVATION
380 """
381 return self.getParameter("RESERVATION")
383 def getExclusive(self):
384 """Accessor for EXCLUSIVE
385 @return the value of EXCLUSIVE
386 """
387 return self.getParameter("EXCLUSIVE")
389 def getExcluser(self):
390 """Accessor for EXCLUSER
391 @return the value of EXCLUSER
392 """
393 return self.getParameter("EXCLUSER")
395 def getParameter(self, value):
396 """Accessor for generic value
397 @return None if value is not set. Otherwise, use the command line
398 override (if set), or the default Config value
399 """
400 if value in self.commandLineDefaults:
401 return self.commandLineDefaults[value]
402 if value in self.defaults:
403 return self.defaults[value]
404 return None
406 def printNodeSetInfo(self):
407 nodes = self.getNodes()
408 cpus = self.getCPUs()
409 wallClock = self.getWallClock()
410 nodeString = ""
412 if int(nodes) > 1:
413 nodeString = "s"
414 if self.opts.dynamic is None:
415 print(
416 f"{nodes} glidein{nodeString} will be allocated on "
417 f"{self.platform} using default dynamic slots configuration."
418 )
419 print(f"There will be {cpus} cores per glidein and a maximum time limit of {wallClock}")
420 elif self.opts.dynamic == "__default__":
421 print(
422 f"{nodes} glidein{nodeString} will be allocated on {self.platform} "
423 "using default dynamic slots configuration."
424 )
425 print(f"There will be {cpus} cores per glidein and a maximum time limit of {wallClock}")
426 else:
427 print(
428 f"{nodes} node{nodeString} will be allocated on {self.platform} "
429 f"using dynamic slot block specified in '{self.opts.dynamic}'"
430 )
431 print(f"There will be {cpus} cores per node and maximum time limit of {wallClock}")
432 print("Node set name:")
433 print(self.getNodeSetName())
435 def runCommand(self, cmd, verbose):
436 cmd_split = cmd.split()
437 pid = os.fork()
438 if not pid:
439 # Methods of file transfer and login may
440 # produce different output, depending on how
441 # the "gsi" utilities are used. The user can
442 # either use grid proxies or ssh, and gsiscp/gsissh
443 # does the right thing. Since the output will be
444 # different in either case anything potentially parsing this
445 # output (like drpRun), would have to go through extra
446 # steps to deal with this output, and which ultimately
447 # end up not being useful. So we optinally close the i/o output
448 # of the executing command down.
449 #
450 # stdin/stdio/stderr is treated specially
451 # by python, so we have to close down
452 # both the python objects and the
453 # underlying c implementations
454 if not verbose:
455 # close python i/o
456 sys.stdin.close()
457 sys.stdout.close()
458 sys.stderr.close()
459 # close C's i/o
460 os.close(0)
461 os.close(1)
462 os.close(2)
463 os.execvp(cmd_split[0], cmd_split)
464 pid, status = os.wait()
465 # high order bits are status, low order bits are signal.
466 exitCode = (status & 0xFF00) >> 8
467 return exitCode
469 def submit(self):
470 """Submit the glidein jobs to the Batch system."""
471 raise NotImplementedError