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
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
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 os
26import sys
27from string import Template
29from lsst.ctrl.execute.allocator import Allocator
32class SlurmPlugin(Allocator):
34 def submit(self, platform, platformPkgDir):
36 configName = os.path.join(platformPkgDir, "etc", "config", "slurmConfig.py")
38 self.loadSlurm(configName, platformPkgDir)
39 verbose = self.isVerbose()
41 # create the fully-resolved scratch directory string
42 scratchDirParam = self.getScratchDirectory()
43 template = Template(scratchDirParam)
44 template.substitute(USER_HOME=self.getUserHome())
46 # create the slurm submit file
47 slurmName = os.path.join(platformPkgDir, "etc", "templates", "generic.slurm.template")
48 generatedSlurmFile = self.createSubmitFile(slurmName)
50 # create the condor configuration file
51 condorFile = os.path.join(platformPkgDir, "etc", "templates", "glidein_condor_config.template")
52 self.createCondorConfigFile(condorFile)
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)
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)
70 # print node set information
71 self.printNodeSetInfo()
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"] = ""
79 allocationConfig = self.loadAllocationConfig(name, "slurm")
81 template = Template(allocationConfig.platform.scratchDirectory)
82 scratchDir = template.substitute(USER_NAME=self.getUserName())
83 self.defaults["SCRATCH_DIR"] = scratchDir
85 self.allocationFileName = os.path.join(self.configDir, "allocation_%s.sh" % self.uniqueIdentifier)
86 self.defaults["GENERATED_ALLOCATE_SCRIPT"] = os.path.basename(self.allocationFileName)
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
97 if self.opts.dynamic == "__default__":
98 dynamicSlotsName = os.path.join(platformPkgDir, "etc", "templates", "dynamic_slots.template")
99 else:
100 dynamicSlotsName = self.opts.dynamic
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
109 def createAllocationFile(self, input):
110 """Creates Allocation script file using the file "input" as a Template
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