Coverage for python/lsst/ctrl/execute/slurmPlugin.py: 31%

Shortcuts on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

58 statements  

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 os 

26import sys 

27from string import Template 

28 

29from lsst.ctrl.execute.allocator import Allocator 

30 

31 

32class SlurmPlugin(Allocator): 

33 

34 def submit(self, platform, platformPkgDir): 

35 

36 configName = os.path.join(platformPkgDir, "etc", "config", "slurmConfig.py") 

37 

38 self.loadSlurm(configName, platformPkgDir) 

39 verbose = self.isVerbose() 

40 

41 # create the fully-resolved scratch directory string 

42 scratchDirParam = self.getScratchDirectory() 

43 template = Template(scratchDirParam) 

44 template.substitute(USER_HOME=self.getUserHome()) 

45 

46 # create the slurm submit file 

47 slurmName = os.path.join(platformPkgDir, "etc", "templates", "generic.slurm.template") 

48 generatedSlurmFile = self.createSubmitFile(slurmName) 

49 

50 # create the condor configuration file 

51 condorFile = os.path.join(platformPkgDir, "etc", "templates", "glidein_condor_config.template") 

52 self.createCondorConfigFile(condorFile) 

53 

54 # create the script that the slurm submit file calls 

55 allocationName = os.path.join(platformPkgDir, "etc", "templates", "allocation.sh.template") 

56 self.createAllocationFile(allocationName) 

57 

58 # run the sbatch command 

59 template = Template(self.getLocalScratchDirectory()) 

60 localScratchDir = template.substitute(USER_NAME=self.getUserName()) 

61 if not os.path.exists(localScratchDir): 

62 os.mkdir(localScratchDir) 

63 os.chdir(localScratchDir) 

64 cmd = "sbatch %s" % generatedSlurmFile 

65 exitCode = self.runCommand(cmd, verbose) 

66 if exitCode != 0: 

67 print("error running %s" % cmd) 

68 sys.exit(exitCode) 

69 

70 # print node set information 

71 self.printNodeSetInfo() 

72 

73 def loadSlurm(self, name, platformPkgDir): 

74 if self.opts.reservation is not None: 74 ↛ 75line 74 didn't jump to line 75, because the condition on line 74 was never true

75 self.defaults["RESERVATION"] = "#SBATCH --reservation=%s" % self.opts.reservation 

76 else: 

77 self.defaults["RESERVATION"] = "" 

78 

79 allocationConfig = self.loadAllocationConfig(name, "slurm") 

80 

81 template = Template(allocationConfig.platform.scratchDirectory) 

82 scratchDir = template.substitute(USER_NAME=self.getUserName()) 

83 self.defaults["SCRATCH_DIR"] = scratchDir 

84 

85 self.allocationFileName = os.path.join(self.configDir, "allocation_%s.sh" % self.uniqueIdentifier) 

86 self.defaults["GENERATED_ALLOCATE_SCRIPT"] = os.path.basename(self.allocationFileName) 

87 

88 # handle dynamic slot block template: 

89 # 1) if it isn't specified, just put a comment in it's place 

90 # 2) if it's specified, but without a filename, use the default 

91 # 3) if it's specified with a filename, use that. 

92 dynamicSlotsName = None 

93 if self.opts.dynamic is None: 93 ↛ 97line 93 didn't jump to line 97, because the condition on line 93 was never false

94 self.defaults["DYNAMIC_SLOTS_BLOCK"] = "#" 

95 return 

96 

97 if self.opts.dynamic == "__default__": 

98 dynamicSlotsName = os.path.join(platformPkgDir, "etc", "templates", "dynamic_slots.template") 

99 else: 

100 dynamicSlotsName = self.opts.dynamic 

101 

102 with open(dynamicSlotsName) as f: 

103 lines = f.readlines() 

104 block = "" 

105 for line in lines: 

106 block += line 

107 self.defaults["DYNAMIC_SLOTS_BLOCK"] = block 

108 

109 def createAllocationFile(self, input): 

110 """Creates Allocation script file using the file "input" as a Template 

111 

112 Returns 

113 ------- 

114 outfile : `str` 

115 The newly created file name 

116 """ 

117 outfile = self.createFile(input, self.allocationFileName) 

118 if self.opts.verbose: 

119 print("wrote new allocation script file to %s" % outfile) 

120 os.chmod(outfile, 0o755) 

121 return outfile