CRAB/python/ScriptWriter.py

from WorkSpace import WorkSpace
from JobList import JobList
from Scheduler import Scheduler
from crab_exceptions import *
import common
import Scram
import string,os

class ScriptWriter:
    def __init__(self, cfg_params, template):
        # pattern -> action
        self.actions = {
            'title'                       : self.title_,
            'untar_software'              : self.untarSoftware_,
            'initial_environment'         : self.initialEnvironment_,
            'setup_scheduler_environment' : self.setupSchedulerEnvironment_,
            'setup_jobtype_environment'   : self.setupJobTypeEnvironment_,
            'copy_input'                  : self.copyInput_,
            'rewrite_cmssw_cfg'           : self.rewriteCMSSWcfg_,
            'build_executable'            : self.buildExe_,
            'run_executable'              : self.runExe_,
            'rename_output'               : self.renameOutput_,
            'copy_output'                 : self.copyOutput_,
            'parse_report'                : self.parseReport_,
            'modify_report'               : self.modifyReport_,
            'func_exit'                   : self.func_exit_
            }

        if os.path.isfile("./"+template):
            self.template = "./"+template
        elif os.getenv('CRABDIR') and os.path.isfile(os.getenv('CRABDIR')+'/python/'+template):
            self.template = os.getenv('CRABDIR')+'/python/'+template
        else:
            raise CrabException("No crab_template.sh found!")
        self.nj = -1     # current job number

        try:
            self.scram = Scram.Scram(None)
            self.CMSSWversion = self.scram.getSWVersion()
            parts = self.CMSSWversion.split('_')
            self.CMSSW_major = int(parts[1])
            self.CMSSW_minor = int(parts[2])
            self.CMSSW_patch = int(parts[3])
        except:
            raise CrabException("Could not determine CMSSW version")
        self.debug_wrapper=''
        debug = int(cfg_params.get('USER.debug_wrapper',0))
        if debug==1 : self.debug_wrapper='--debug'

        self.scriptName = cfg_params.get('CRAB.jobtype').upper()+'.sh'  

        return

    def setAction(self, pattern, action):
        self.actions[pattern] = action
        return

    def modifyTemplateScript(self):
        """
        Create a script from scratch.
        """

        tpl = open(self.template, 'r')
 
        wrapper_fullPath = common.work_space.jobDir()+self.scriptName
        script = open(wrapper_fullPath,'w')

        for line in tpl:
            if len(line) > 6 and line[:6] == '#CRAB ':
                act_str = string.strip(line[6:])
                try:
                    action = self.actions[act_str]
                except KeyError:
                    continue

                if action:
                    txt = action()
                    script.write(txt)
                    pass
                else:
                    script.write(line)
                pass
            else:
                script.write(line)
                pass
            pass

        script.close()
        tpl.close()
        os.chmod(wrapper_fullPath, 0744)
        return

    def title_(self):
        txt = '# This script was generated by '+common.prog_name
        txt += ' (version '+common.prog_version_str+').\n'
        return txt


    ### FEDE ###

    def untarSoftware_(self):
        """
        Returns part of a job script which untar CMSSW software.
        """
        jbt = common.job_list.type()

        txt = jbt.wsUntarSoftware(self.nj)

        #txt += 'executable='+exe+'\n'
        return txt

    ###########################################

    def setupSchedulerEnvironment_(self):
        """
        Returns part of a job script which does scheduler-specific work.
        """
        txt = common.scheduler.wsSetupEnvironment()
        return txt

    def initialEnvironment_(self):
        """
        Returns part of a job script which does scheduler-specific work.
        """
        txt = common.scheduler.wsInitialEnvironment()
        return txt

    def setupJobTypeEnvironment_(self):
        """
        Returns part of a job script which does jobtype-specific work.
        """
        jbt = common.job_list.type()
        txt = jbt.wsSetupEnvironment(self.nj)
        return txt

    def buildExe_(self):
        """
        Returns part of a job script which builds the binary executable.
        """
        jbt = common.job_list.type()

        txt = jbt.wsBuildExe(self.nj)

        job = common.job_list[self.nj]
        exe = job.type().executableName()

        txt += 'executable='+exe+'\n'
        return txt

    def runExe_(self):
        """
        Returns part of a job script which executes the application.
        """
        job = common.job_list[self.nj]
        args = job.type().executableArgs()

        txt  = ''
        txt += 'CPU_INFOS=-1 \n'
        # NO carriage return for this line #Fabio 
        txt += '/usr/bin/time -f \"%U %S %P\" -o cpu_timing.txt '
        txt += '$executable '+args+'\n'
        return txt

    def renameOutput_(self):
        """
        Returns part of a job script which renames output files.
        """
        jbt = common.job_list.type()
        txt = '\n'
        txt += jbt.wsRenameOutput(self.nj)
        return txt

    def copyInput_(self):
        """
        Returns part of a job script which copies input files from SE.
        """
        txt = common.scheduler.wsCopyInput()
        return txt

    def copyOutput_(self):
        """
        Returns part of a job script which copies output files to SE.
        """
        txt = common.scheduler.wsCopyOutput()
        return txt

    def parseReport_(self):
        """
        Returns part of a job script which parse the FrameworkJobReport.
        """
        jbt = common.job_list.type()
        txt = jbt.wsParseFJR()
        return txt

    def modifyReport_(self):
        """
        Returns part of a job script which modifies the FrameworkJobReport.
        """
        jbt = common.job_list.type()
        txt = jbt.wsModifyReport(self.nj)
        return txt

    def cleanEnv_(self):
        """
        In OSG environment this function removes the WORKING_DIR
        """
        jbt = common.job_list.type()
        txt = jbt.cleanEnv()
        return txt

    def func_exit_(self):
        """
        Returns part of a job script which does scheduler-specific
        output checks and management.
        """
        txt = common.scheduler.wsExitFunc()
        return txt

    def rewriteCMSSWcfg_(self):
        """
        Returns part of the script that runs writeCfg.py on the WN
        """
        # FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
        txt = "# Rewrite cfg for this job\n"

        if (self.CMSSW_major >= 2 and self.CMSSW_minor >= 1) or self.CMSSW_major > 2: #  py in,  py out for 2_1_x
          txt += "echo  $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.py pset.py\n"
          txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.py pset.py\n"
        elif self.CMSSW_major >= 2:                                                   # cfg in,  py out for 2_0_x
          txt += "echo  $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.py\n"
          txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.py\n"
        else:                                                                         # cfg in, cfg out for 1_x_y
          txt += "echo  $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.cfg\n"
          txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.cfg\n"

        return txt
Revision:	1.43
Committed:	Tue May 26 10:23:01 2009 UTC (15 years, 11 months ago) by spiga
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_6_0_pre10, CRAB_2_6_0_pre9, CRAB_2_6_0_pre8, CRAB_2_6_0_pre7, CRAB_2_6_0_pre6, CRAB_2_6_0_pre5, CRAB_2_6_0_pre4, CRAB_2_6_0_pre3
Changes since 1.42:	+0 -1 lines
Log Message:	adapting code to logging usage. (crab_logger removed)
#	User	Rev	Content
1	nsmirnov	1.1	from WorkSpace import WorkSpace
2			from JobList import JobList
3			from Scheduler import Scheduler
4			from crab_exceptions import *
5			import common
6	ewv	1.21	import Scram
7	slacapra	1.6	import string,os
8	nsmirnov	1.1
9			class ScriptWriter:
10	ewv	1.32	def __init__(self, cfg_params, template):
11	nsmirnov	1.1	# pattern -> action
12			self.actions = {
13	nsmirnov	1.3	'title' : self.title_,
14	fanzago	1.23	'untar_software' : self.untarSoftware_,
15	ewv	1.32	'initial_environment' : self.initialEnvironment_,
16	nsmirnov	1.3	'setup_scheduler_environment' : self.setupSchedulerEnvironment_,
17			'setup_jobtype_environment' : self.setupJobTypeEnvironment_,
18	fanzago	1.8	'copy_input' : self.copyInput_,
19	ewv	1.21	'rewrite_cmssw_cfg' : self.rewriteCMSSWcfg_,
20	nsmirnov	1.3	'build_executable' : self.buildExe_,
21			'run_executable' : self.runExe_,
22			'rename_output' : self.renameOutput_,
23	fanzago	1.5	'copy_output' : self.copyOutput_,
24	spiga	1.30	'parse_report' : self.parseReport_,
25	fanzago	1.12	'modify_report' : self.modifyReport_,
26	spiga	1.27	'func_exit' : self.func_exit_
27	nsmirnov	1.1	}
28	ewv	1.21
29	slacapra	1.6	if os.path.isfile("./"+template):
30			self.template = "./"+template
31			elif os.getenv('CRABDIR') and os.path.isfile(os.getenv('CRABDIR')+'/python/'+template):
32			self.template = os.getenv('CRABDIR')+'/python/'+template
33			else:
34			raise CrabException("No crab_template.sh found!")
35	nsmirnov	1.2	self.nj = -1 # current job number
36	mcinquil	1.14
37	ewv	1.21	try:
38			self.scram = Scram.Scram(None)
39			self.CMSSWversion = self.scram.getSWVersion()
40			parts = self.CMSSWversion.split('_')
41			self.CMSSW_major = int(parts[1])
42			self.CMSSW_minor = int(parts[2])
43			self.CMSSW_patch = int(parts[3])
44			except:
45			raise CrabException("Could not determine CMSSW version")
46	spiga	1.42	self.debug_wrapper=''
47	fanzago	1.41	debug = int(cfg_params.get('USER.debug_wrapper',0))
48			if debug==1 : self.debug_wrapper='--debug'
49	spiga	1.37
50			self.scriptName = cfg_params.get('CRAB.jobtype').upper()+'.sh'
51
52	nsmirnov	1.1	return
53
54			def setAction(self, pattern, action):
55			self.actions[pattern] = action
56			return
57	ewv	1.21
58	slacapra	1.10	def modifyTemplateScript(self):
59	nsmirnov	1.1	"""
60			Create a script from scratch.
61			"""
62	ewv	1.21
63	nsmirnov	1.1	tpl = open(self.template, 'r')
64	spiga	1.38
65			wrapper_fullPath = common.work_space.jobDir()+self.scriptName
66			script = open(wrapper_fullPath,'w')
67	nsmirnov	1.1
68			for line in tpl:
69			if len(line) > 6 and line[:6] == '#CRAB ':
70			act_str = string.strip(line[6:])
71			try:
72			action = self.actions[act_str]
73			except KeyError:
74			continue
75
76			if action:
77			txt = action()
78			script.write(txt)
79			pass
80			else:
81			script.write(line)
82			pass
83			else:
84			script.write(line)
85			pass
86			pass
87
88			script.close()
89			tpl.close()
90	spiga	1.39	os.chmod(wrapper_fullPath, 0744)
91	nsmirnov	1.1	return
92
93	nsmirnov	1.2	def title_(self):
94	nsmirnov	1.1	txt = '# This script was generated by '+common.prog_name
95			txt += ' (version '+common.prog_version_str+').\n'
96			return txt
97	ewv	1.26
98	fanzago	1.23
99			### FEDE ###
100	ewv	1.26
101	fanzago	1.23	def untarSoftware_(self):
102			"""
103			Returns part of a job script which untar CMSSW software.
104			"""
105			jbt = common.job_list.type()
106
107			txt = jbt.wsUntarSoftware(self.nj)
108
109			#txt += 'executable='+exe+'\n'
110			return txt
111	ewv	1.26
112	fanzago	1.23	###########################################
113	ewv	1.21
114	nsmirnov	1.3	def setupSchedulerEnvironment_(self):
115			"""
116			Returns part of a job script which does scheduler-specific work.
117			"""
118			txt = common.scheduler.wsSetupEnvironment()
119			return txt
120	nsmirnov	1.1
121	ewv	1.32	def initialEnvironment_(self):
122			"""
123			Returns part of a job script which does scheduler-specific work.
124			"""
125			txt = common.scheduler.wsInitialEnvironment()
126			return txt
127
128	nsmirnov	1.3	def setupJobTypeEnvironment_(self):
129	nsmirnov	1.1	"""
130	fanzago	1.7	Returns part of a job script which does jobtype-specific work.
131	nsmirnov	1.1	"""
132	nsmirnov	1.3	jbt = common.job_list.type()
133			txt = jbt.wsSetupEnvironment(self.nj)
134			return txt
135	ewv	1.21
136	nsmirnov	1.2	def buildExe_(self):
137	nsmirnov	1.1	"""
138	nsmirnov	1.2	Returns part of a job script which builds the binary executable.
139	nsmirnov	1.1	"""
140	nsmirnov	1.2	jbt = common.job_list.type()
141	nsmirnov	1.1
142	nsmirnov	1.3	txt = jbt.wsBuildExe(self.nj)
143	nsmirnov	1.1
144	nsmirnov	1.2	job = common.job_list[self.nj]
145	nsmirnov	1.1	exe = job.type().executableName()
146
147	nsmirnov	1.2	txt += 'executable='+exe+'\n'
148			return txt
149
150			def runExe_(self):
151			"""
152			Returns part of a job script which executes the application.
153			"""
154	gutsche	1.9	job = common.job_list[self.nj]
155			args = job.type().executableArgs()
156	farinafa	1.34
157			txt = ''
158	spiga	1.40	txt += 'CPU_INFOS=-1 \n'
159	farinafa	1.34	# NO carriage return for this line #Fabio
160	farinafa	1.35	txt += '/usr/bin/time -f \"%U %S %P\" -o cpu_timing.txt '
161	farinafa	1.34	txt += '$executable '+args+'\n'
162			return txt
163	nsmirnov	1.1
164	nsmirnov	1.3	def renameOutput_(self):
165			"""
166			Returns part of a job script which renames output files.
167			"""
168			jbt = common.job_list.type()
169			txt = '\n'
170			txt += jbt.wsRenameOutput(self.nj)
171			return txt
172	fanzago	1.5
173	fanzago	1.8	def copyInput_(self):
174			"""
175			Returns part of a job script which copies input files from SE.
176			"""
177			txt = common.scheduler.wsCopyInput()
178			return txt
179
180	fanzago	1.5	def copyOutput_(self):
181			"""
182			Returns part of a job script which copies output files to SE.
183			"""
184			txt = common.scheduler.wsCopyOutput()
185			return txt
186
187	spiga	1.30	def parseReport_(self):
188			"""
189			Returns part of a job script which parse the FrameworkJobReport.
190			"""
191			jbt = common.job_list.type()
192	spiga	1.31	txt = jbt.wsParseFJR()
193	spiga	1.30	return txt
194
195	fanzago	1.11	def modifyReport_(self):
196			"""
197			Returns part of a job script which modifies the FrameworkJobReport.
198			"""
199			jbt = common.job_list.type()
200	spiga	1.36	txt = jbt.wsModifyReport(self.nj)
201	fanzago	1.11	return txt
202	fanzago	1.12
203			def cleanEnv_(self):
204			"""
205			In OSG environment this function removes the WORKING_DIR
206			"""
207			jbt = common.job_list.type()
208			txt = jbt.cleanEnv()
209			return txt
210	mcinquil	1.14
211	spiga	1.27	def func_exit_(self):
212	mcinquil	1.14	"""
213	ewv	1.32	Returns part of a job script which does scheduler-specific
214	spiga	1.27	output checks and management.
215	mcinquil	1.14	"""
216	spiga	1.27	txt = common.scheduler.wsExitFunc()
217	mcinquil	1.14	return txt
218	spiga	1.22
219	ewv	1.21	def rewriteCMSSWcfg_(self):
220			"""
221			Returns part of the script that runs writeCfg.py on the WN
222			"""
223			# FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
224			txt = "# Rewrite cfg for this job\n"
225
226	ewv	1.26	if (self.CMSSW_major >= 2 and self.CMSSW_minor >= 1) or self.CMSSW_major > 2: # py in, py out for 2_1_x
227	spiga	1.33	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.py pset.py\n"
228			txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.py pset.py\n"
229	ewv	1.26	elif self.CMSSW_major >= 2: # cfg in, py out for 2_0_x
230	spiga	1.33	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.py\n"
231			txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.py\n"
232	ewv	1.26	else: # cfg in, cfg out for 1_x_y
233	spiga	1.33	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.cfg\n"
234			txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_wrapper)+" pset.cfg pset.cfg\n"
235	ewv	1.21
236			return txt