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_,
            '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 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.29
Committed:	Tue May 27 13:32:44 2008 UTC (16 years, 11 months ago) by spiga
Content type:	text/x-python
Branch:	MAIN
Changes since 1.28:	+1 -2 lines
Log Message:	fixed bugs for pset_debug functionality
#	User	Rev	Content
1	nsmirnov	1.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	ewv	1.21	import Scram
8	slacapra	1.6	import string,os
9	nsmirnov	1.1
10			class ScriptWriter:
11	spiga	1.29	def __init__(self, cfg_params, template):
12	nsmirnov	1.1	# pattern -> action
13			self.actions = {
14	nsmirnov	1.3	'title' : self.title_,
15	fanzago	1.23	'untar_software' : self.untarSoftware_,
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	fanzago	1.12	'modify_report' : self.modifyReport_,
25	spiga	1.27	'func_exit' : self.func_exit_
26	nsmirnov	1.1	}
27	ewv	1.21
28	slacapra	1.6	if os.path.isfile("./"+template):
29			self.template = "./"+template
30			elif os.getenv('CRABDIR') and os.path.isfile(os.getenv('CRABDIR')+'/python/'+template):
31			self.template = os.getenv('CRABDIR')+'/python/'+template
32			else:
33			raise CrabException("No crab_template.sh found!")
34	nsmirnov	1.2	self.nj = -1 # current job number
35	mcinquil	1.14
36	ewv	1.21	try:
37			self.scram = Scram.Scram(None)
38			self.CMSSWversion = self.scram.getSWVersion()
39			parts = self.CMSSWversion.split('_')
40			self.CMSSW_major = int(parts[1])
41			self.CMSSW_minor = int(parts[2])
42			self.CMSSW_patch = int(parts[3])
43			except:
44			raise CrabException("Could not determine CMSSW version")
45	spiga	1.28	self.debug_pset=''
46			debug = cfg_params.get('USER.debug_pset',False)
47			if debug: self.debug_pset='--debug'
48
49	nsmirnov	1.1	return
50
51			def setAction(self, pattern, action):
52			self.actions[pattern] = action
53			return
54	ewv	1.21
55	slacapra	1.10	def modifyTemplateScript(self):
56	nsmirnov	1.1	"""
57			Create a script from scratch.
58			"""
59	ewv	1.21
60	nsmirnov	1.1	tpl = open(self.template, 'r')
61	spiga	1.22	script = open(common._db.queryTask('scriptName'),'w')
62	nsmirnov	1.1
63			for line in tpl:
64			if len(line) > 6 and line[:6] == '#CRAB ':
65			act_str = string.strip(line[6:])
66			try:
67			action = self.actions[act_str]
68			except KeyError:
69			continue
70
71			if action:
72			txt = action()
73			script.write(txt)
74			pass
75			else:
76			script.write(line)
77			pass
78			else:
79			script.write(line)
80			pass
81			pass
82
83			script.close()
84			tpl.close()
85			return
86
87	nsmirnov	1.2	def title_(self):
88	nsmirnov	1.1	txt = '# This script was generated by '+common.prog_name
89			txt += ' (version '+common.prog_version_str+').\n'
90			return txt
91	ewv	1.26
92	fanzago	1.23
93			### FEDE ###
94	ewv	1.26
95	fanzago	1.23	def untarSoftware_(self):
96			"""
97			Returns part of a job script which untar CMSSW software.
98			"""
99			jbt = common.job_list.type()
100
101			txt = jbt.wsUntarSoftware(self.nj)
102
103			#txt += 'executable='+exe+'\n'
104			return txt
105	ewv	1.26
106	fanzago	1.23	###########################################
107	ewv	1.21
108	nsmirnov	1.3	def setupSchedulerEnvironment_(self):
109			"""
110			Returns part of a job script which does scheduler-specific work.
111			"""
112			txt = common.scheduler.wsSetupEnvironment()
113			return txt
114	nsmirnov	1.1
115	nsmirnov	1.3	def setupJobTypeEnvironment_(self):
116	nsmirnov	1.1	"""
117	fanzago	1.7	Returns part of a job script which does jobtype-specific work.
118	nsmirnov	1.1	"""
119	nsmirnov	1.3	jbt = common.job_list.type()
120			txt = jbt.wsSetupEnvironment(self.nj)
121			return txt
122	ewv	1.21
123	nsmirnov	1.2	def buildExe_(self):
124	nsmirnov	1.1	"""
125	nsmirnov	1.2	Returns part of a job script which builds the binary executable.
126	nsmirnov	1.1	"""
127	nsmirnov	1.2	jbt = common.job_list.type()
128	nsmirnov	1.1
129	nsmirnov	1.3	txt = jbt.wsBuildExe(self.nj)
130	nsmirnov	1.1
131	nsmirnov	1.2	job = common.job_list[self.nj]
132	nsmirnov	1.1	exe = job.type().executableName()
133
134	nsmirnov	1.2	txt += 'executable='+exe+'\n'
135			return txt
136
137			def runExe_(self):
138			"""
139			Returns part of a job script which executes the application.
140			"""
141	gutsche	1.9	job = common.job_list[self.nj]
142			args = job.type().executableArgs()
143			return '$executable '+args+'\n'
144	nsmirnov	1.1
145	nsmirnov	1.3	def renameOutput_(self):
146			"""
147			Returns part of a job script which renames output files.
148			"""
149			jbt = common.job_list.type()
150			txt = '\n'
151			txt += jbt.wsRenameOutput(self.nj)
152			return txt
153	fanzago	1.5
154	fanzago	1.8	def copyInput_(self):
155			"""
156			Returns part of a job script which copies input files from SE.
157			"""
158			txt = common.scheduler.wsCopyInput()
159			return txt
160
161	fanzago	1.5	def copyOutput_(self):
162			"""
163			Returns part of a job script which copies output files to SE.
164			"""
165			txt = common.scheduler.wsCopyOutput()
166			return txt
167
168	fanzago	1.11	def modifyReport_(self):
169			"""
170			Returns part of a job script which modifies the FrameworkJobReport.
171			"""
172			jbt = common.job_list.type()
173			txt = jbt.modifyReport(self.nj)
174			return txt
175	fanzago	1.12
176			def cleanEnv_(self):
177			"""
178			In OSG environment this function removes the WORKING_DIR
179			"""
180			jbt = common.job_list.type()
181			txt = jbt.cleanEnv()
182			return txt
183	mcinquil	1.14
184	spiga	1.27	def func_exit_(self):
185	mcinquil	1.14	"""
186	spiga	1.27	Returns part of a job script which does scheduler-specific
187			output checks and management.
188	mcinquil	1.14	"""
189	spiga	1.27	txt = common.scheduler.wsExitFunc()
190	mcinquil	1.14	return txt
191	spiga	1.22
192	ewv	1.21	def rewriteCMSSWcfg_(self):
193			"""
194			Returns part of the script that runs writeCfg.py on the WN
195			"""
196			# FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
197			txt = "# Rewrite cfg for this job\n"
198
199	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
200	spiga	1.28	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.py pset.py\n"
201			txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.py pset.py\n"
202	ewv	1.26	elif self.CMSSW_major >= 2: # cfg in, py out for 2_0_x
203	spiga	1.28	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.py\n"
204			txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.py\n"
205	ewv	1.26	else: # cfg in, cfg out for 1_x_y
206	spiga	1.28	txt += "echo $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.cfg\n"
207			txt += "python $RUNTIME_AREA/writeCfg.py "+str(self.debug_pset)+" pset.cfg pset.cfg\n"
208	ewv	1.21
209			return txt