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_,
            '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_pset=''
        debug = cfg_params.get('USER.debug_pset',False)
        if debug: self.debug_pset='--debug'
 
        return

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

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

        tpl = open(self.template, 'r')
        script = open(common._db.queryTask('scriptName'),'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()
        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 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()
        return '$executable '+args+'\n'

    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.modifyReport(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_pset)+" pset.py pset.py\n"
          txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" 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_pset)+" pset.cfg pset.py\n"
          txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.py\n"
        else:                                                                         # cfg in, cfg out for 1_x_y
          txt += "echo  $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.cfg\n"
          txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.cfg\n"

        return txt
Revision:	1.31
Committed:	Tue May 27 14:36:11 2008 UTC (16 years, 11 months ago) by spiga
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	PRODCOMMON_0_10_7_testCS2
Changes since 1.30:	+1 -1 lines
Log Message:	wrong argument
#	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	'setup_scheduler_environment' : self.setupSchedulerEnvironment_,
17	'setup_jobtype_environment' : self.setupJobTypeEnvironment_,
18	'copy_input' : self.copyInput_,
19	'rewrite_cmssw_cfg' : self.rewriteCMSSWcfg_,
20	'build_executable' : self.buildExe_,
21	'run_executable' : self.runExe_,
22	'rename_output' : self.renameOutput_,
23	'copy_output' : self.copyOutput_,
24	'parse_report' : self.parseReport_,
25	'modify_report' : self.modifyReport_,
26	'func_exit' : self.func_exit_
27	}
28
29	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	self.nj = -1 # current job number
36
37	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	self.debug_pset=''
47	debug = cfg_params.get('USER.debug_pset',False)
48	if debug: self.debug_pset='--debug'
49
50	return
51
52	def setAction(self, pattern, action):
53	self.actions[pattern] = action
54	return
55
56	def modifyTemplateScript(self):
57	"""
58	Create a script from scratch.
59	"""
60
61	tpl = open(self.template, 'r')
62	script = open(common._db.queryTask('scriptName'),'w')
63
64	for line in tpl:
65	if len(line) > 6 and line[:6] == '#CRAB ':
66	act_str = string.strip(line[6:])
67	try:
68	action = self.actions[act_str]
69	except KeyError:
70	continue
71
72	if action:
73	txt = action()
74	script.write(txt)
75	pass
76	else:
77	script.write(line)
78	pass
79	else:
80	script.write(line)
81	pass
82	pass
83
84	script.close()
85	tpl.close()
86	return
87
88	def title_(self):
89	txt = '# This script was generated by '+common.prog_name
90	txt += ' (version '+common.prog_version_str+').\n'
91	return txt
92
93
94	### FEDE ###
95
96	def untarSoftware_(self):
97	"""
98	Returns part of a job script which untar CMSSW software.
99	"""
100	jbt = common.job_list.type()
101
102	txt = jbt.wsUntarSoftware(self.nj)
103
104	#txt += 'executable='+exe+'\n'
105	return txt
106
107	###########################################
108
109	def setupSchedulerEnvironment_(self):
110	"""
111	Returns part of a job script which does scheduler-specific work.
112	"""
113	txt = common.scheduler.wsSetupEnvironment()
114	return txt
115
116	def setupJobTypeEnvironment_(self):
117	"""
118	Returns part of a job script which does jobtype-specific work.
119	"""
120	jbt = common.job_list.type()
121	txt = jbt.wsSetupEnvironment(self.nj)
122	return txt
123
124	def buildExe_(self):
125	"""
126	Returns part of a job script which builds the binary executable.
127	"""
128	jbt = common.job_list.type()
129
130	txt = jbt.wsBuildExe(self.nj)
131
132	job = common.job_list[self.nj]
133	exe = job.type().executableName()
134
135	txt += 'executable='+exe+'\n'
136	return txt
137
138	def runExe_(self):
139	"""
140	Returns part of a job script which executes the application.
141	"""
142	job = common.job_list[self.nj]
143	args = job.type().executableArgs()
144	return '$executable '+args+'\n'
145
146	def renameOutput_(self):
147	"""
148	Returns part of a job script which renames output files.
149	"""
150	jbt = common.job_list.type()
151	txt = '\n'
152	txt += jbt.wsRenameOutput(self.nj)
153	return txt
154
155	def copyInput_(self):
156	"""
157	Returns part of a job script which copies input files from SE.
158	"""
159	txt = common.scheduler.wsCopyInput()
160	return txt
161
162	def copyOutput_(self):
163	"""
164	Returns part of a job script which copies output files to SE.
165	"""
166	txt = common.scheduler.wsCopyOutput()
167	return txt
168
169	def parseReport_(self):
170	"""
171	Returns part of a job script which parse the FrameworkJobReport.
172	"""
173	jbt = common.job_list.type()
174	txt = jbt.wsParseFJR()
175	return txt
176
177	def modifyReport_(self):
178	"""
179	Returns part of a job script which modifies the FrameworkJobReport.
180	"""
181	jbt = common.job_list.type()
182	txt = jbt.modifyReport(self.nj)
183	return txt
184
185	def cleanEnv_(self):
186	"""
187	In OSG environment this function removes the WORKING_DIR
188	"""
189	jbt = common.job_list.type()
190	txt = jbt.cleanEnv()
191	return txt
192
193	def func_exit_(self):
194	"""
195	Returns part of a job script which does scheduler-specific
196	output checks and management.
197	"""
198	txt = common.scheduler.wsExitFunc()
199	return txt
200
201	def rewriteCMSSWcfg_(self):
202	"""
203	Returns part of the script that runs writeCfg.py on the WN
204	"""
205	# FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
206	txt = "# Rewrite cfg for this job\n"
207
208	if (self.CMSSW_major >= 2 and self.CMSSW_minor >= 1) or self.CMSSW_major > 2: # py in, py out for 2_1_x
209	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.py pset.py\n"
210	txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.py pset.py\n"
211	elif self.CMSSW_major >= 2: # cfg in, py out for 2_0_x
212	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.py\n"
213	txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.py\n"
214	else: # cfg in, cfg out for 1_x_y
215	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.cfg\n"
216	txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.cfg\n"
217
218	return txt