Coverage for python/lsst/ctrl/execute/allocator.py: 19%

197 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-10-10 01:46 -0700

1#!/usr/bin/env python 

2 

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# 

24 

25import logging 

26import os 

27import pwd 

28import sys 

29from datetime import datetime 

30from string import Template 

31 

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 

36 

37_LOG = logging.getLogger(__name__) 

38 

39 

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. 

44 

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 

53 

54 Raises 

55 ------ 

56 TypeError 

57 If the condorInfoFileName is the wrong type. 

58 """ 

59 

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 

74 

75 condorInfoConfig = CondorInfoConfig() 

76 condorInfoConfig.loadFromStream(ResourcePath(condorInfoFileName).read()) 

77 

78 self.platform = platform 

79 

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() 

123 

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. 

127 

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 

146 

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 

156 

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()) 

165 

166 self.defaults["QUEUE"] = allocationConfig.platform.queue 

167 self.defaults["EMAIL_NOTIFICATION"] = allocationConfig.platform.email 

168 self.defaults["HOST_NAME"] = allocationConfig.platform.loginHostName 

169 

170 self.defaults["UTILITY_PATH"] = allocationConfig.platform.utilityPath 

171 

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) 

176 

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" 

181 

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" 

186 

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 

192 

193 self.uniqueIdentifier = self.createUniqueIdentifier() 

194 

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 ) 

202 

203 self.submitFileName = os.path.join(self.configDir, f"alloc_{self.uniqueIdentifier}.{suffix}") 

204 

205 self.condorConfigFileName = os.path.join(self.configDir, f"condor_{self.uniqueIdentifier}.config") 

206 

207 self.defaults["GENERATED_CONFIG"] = os.path.basename(self.condorConfigFileName) 

208 self.defaults["CONFIGURATION_ID"] = self.uniqueIdentifier 

209 return allocationConfig 

210 

211 def createSubmitFile(self, inputFile): 

212 """Creates a batch submit file using the file "input" as a Template 

213 

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 

224 

225 def createCondorConfigFile(self, input): 

226 """Creates a Condor config file using the file "input" as a Template 

227 

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 

236 

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. 

240 

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 

258 

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 

264 

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 

270 

271 def getUserName(self): 

272 """Accessor for USER_NAME 

273 @return the value of USER_NAME 

274 """ 

275 return self.getParameter("USER_NAME") 

276 

277 def getUserHome(self): 

278 """Accessor for USER_HOME 

279 @return the value of USER_HOME 

280 """ 

281 return self.getParameter("USER_HOME") 

282 

283 def getUserScratch(self): 

284 """Accessor for USER_SCRATCH 

285 @return the value of USER_SCRATCH 

286 """ 

287 return self.getParameter("USER_SCRATCH") 

288 

289 def getHostName(self): 

290 """Accessor for HOST_NAME 

291 @return the value of HOST_NAME 

292 """ 

293 return self.getParameter("HOST_NAME") 

294 

295 def getUtilityPath(self): 

296 """Accessor for UTILITY_PATH 

297 @return the value of UTILITY_PATH 

298 """ 

299 return self.getParameter("UTILITY_PATH") 

300 

301 def getScratchDirectory(self): 

302 """Accessor for SCRATCH_DIR 

303 @return the value of SCRATCH_DIR 

304 """ 

305 return self.getParameter("SCRATCH_DIR") 

306 

307 def getLocalScratchDirectory(self): 

308 """Accessor for LOCAL_SCRATCH 

309 @return the value of LOCAL_SCRATCH 

310 """ 

311 return self.getParameter("LOCAL_SCRATCH") 

312 

313 def getNodeSetName(self): 

314 """Accessor for NODE_SET 

315 @return the value of NODE_SET 

316 """ 

317 return self.getParameter("NODE_SET") 

318 

319 def getNodes(self): 

320 """Accessor for NODE_COUNT 

321 @return the value of NODE_COUNT 

322 """ 

323 return self.getParameter("NODE_COUNT") 

324 

325 def getMemoryPerCore(self): 

326 """Accessor for MemoryPerCore 

327 @return the value of MemoryPerCore 

328 """ 

329 return self.getParameter("MEMPERCORE") 

330 

331 def getAllowedAutoGlideins(self): 

332 """Accessor for AllowedAutoGlideins 

333 @return the value of AllowedAuto 

334 """ 

335 return self.getParameter("ALLOWEDAUTO") 

336 

337 def getQOS(self): 

338 """Accessor for QOS 

339 @return the value of QOS 

340 """ 

341 return self.getParameter("QOS") 

342 

343 def getCPUs(self): 

344 """Accessor for CPUS 

345 @return the value of CPUS 

346 """ 

347 return self.getParameter("CPUS") 

348 

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") 

358 

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") 

364 

365 def getWallClock(self): 

366 """Accessor for WALL_CLOCK 

367 @return the value of WALL_CLOCK 

368 """ 

369 return self.getParameter("WALL_CLOCK") 

370 

371 def getScheduler(self): 

372 """Accessor for SCHEDULER 

373 @return the value of SCHEDULER 

374 """ 

375 return self.getParameter("SCHEDULER") 

376 

377 def getReservation(self): 

378 """Accessor for RESERVATION 

379 @return the value of RESERVATION 

380 """ 

381 return self.getParameter("RESERVATION") 

382 

383 def getExclusive(self): 

384 """Accessor for EXCLUSIVE 

385 @return the value of EXCLUSIVE 

386 """ 

387 return self.getParameter("EXCLUSIVE") 

388 

389 def getExcluser(self): 

390 """Accessor for EXCLUSER 

391 @return the value of EXCLUSER 

392 """ 

393 return self.getParameter("EXCLUSER") 

394 

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 

405 

406 def printNodeSetInfo(self): 

407 nodes = self.getNodes() 

408 cpus = self.getCPUs() 

409 wallClock = self.getWallClock() 

410 nodeString = "" 

411 

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()) 

434 

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 

468 

469 def submit(self): 

470 """Submit the glidein jobs to the Batch system.""" 

471 raise NotImplementedError