CRAB/python/ScriptWriter.py

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