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, template, output_troncate_flag): ## added by Matty
        # pattern -> action
        ### FEDE added modify_report FOR DBS OUTPUT PUBLICATION
        self.actions = {
            'title'                       : self.title_,
            '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_,
            'clean_env'                   : self.cleanEnv_,
            'check_output_limit'          : self.checkOut_
            }

        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

        self.output_troncate_flag = output_troncate_flag
        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")

        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.taskDB.dict('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

    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 checkOut_(self):
        """
        With glite check if the output is too big
        """
        txt = "\n"
        if self.output_troncate_flag == 1:
            limit = 55000000 ##52 MB
            jbt = common.job_list.type()
            txt = jbt.checkOut(limit)
        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 > 1: # Use py files for CMSSW 2_0_x and up
          txt += "echo  $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.pycfg\n"
          txt += "python $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.pycfg\n"
        else:
          txt += "echo  $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.cfg\n"
          txt += "python $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.cfg\n"

        return txt
Revision:	1.21
Committed:	Tue Feb 26 22:24:20 2008 UTC (17 years, 2 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_1_1, CRAB_2_1_1_pre3, CRAB_2_1_1_pre1
Branch point for:	CRAB_2_1_2_br, CRAB_2_1_1_pre2
Changes since 1.20:	+31 -6 lines
Log Message:	Python config file, remove reliance on sed
#	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	mcinquil	1.14	def __init__(self, template, output_troncate_flag): ## added by Matty
12	nsmirnov	1.1	# pattern -> action
13	fanzago	1.11	### FEDE added modify_report FOR DBS OUTPUT PUBLICATION
14	nsmirnov	1.1	self.actions = {
15	nsmirnov	1.3	'title' : self.title_,
16			'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	mcinquil	1.14	'clean_env' : self.cleanEnv_,
26			'check_output_limit' : self.checkOut_
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			self.output_troncate_flag = output_troncate_flag
38	ewv	1.21	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	slacapra	1.19
48	nsmirnov	1.1	return
49
50			def setAction(self, pattern, action):
51			self.actions[pattern] = action
52			return
53	ewv	1.21
54	slacapra	1.10	def modifyTemplateScript(self):
55	nsmirnov	1.1	"""
56			Create a script from scratch.
57			"""
58	ewv	1.21
59	nsmirnov	1.1	tpl = open(self.template, 'r')
60	slacapra	1.10	script = open(common.taskDB.dict('ScriptName'),'w')
61	nsmirnov	1.1
62			for line in tpl:
63			if len(line) > 6 and line[:6] == '#CRAB ':
64			act_str = string.strip(line[6:])
65			try:
66			action = self.actions[act_str]
67			except KeyError:
68			continue
69
70			if action:
71			txt = action()
72			script.write(txt)
73			pass
74			else:
75			script.write(line)
76			pass
77			else:
78			script.write(line)
79			pass
80			pass
81
82			script.close()
83			tpl.close()
84			return
85
86	nsmirnov	1.2	def title_(self):
87	nsmirnov	1.1	txt = '# This script was generated by '+common.prog_name
88			txt += ' (version '+common.prog_version_str+').\n'
89			return txt
90	ewv	1.21
91	nsmirnov	1.3	def setupSchedulerEnvironment_(self):
92			"""
93			Returns part of a job script which does scheduler-specific work.
94			"""
95			txt = common.scheduler.wsSetupEnvironment()
96			return txt
97	nsmirnov	1.1
98	nsmirnov	1.3	def setupJobTypeEnvironment_(self):
99	nsmirnov	1.1	"""
100	fanzago	1.7	Returns part of a job script which does jobtype-specific work.
101	nsmirnov	1.1	"""
102	nsmirnov	1.3	jbt = common.job_list.type()
103			txt = jbt.wsSetupEnvironment(self.nj)
104			return txt
105	ewv	1.21
106	nsmirnov	1.2	def buildExe_(self):
107	nsmirnov	1.1	"""
108	nsmirnov	1.2	Returns part of a job script which builds the binary executable.
109	nsmirnov	1.1	"""
110	nsmirnov	1.2	jbt = common.job_list.type()
111	nsmirnov	1.1
112	nsmirnov	1.3	txt = jbt.wsBuildExe(self.nj)
113	nsmirnov	1.1
114	nsmirnov	1.2	job = common.job_list[self.nj]
115	nsmirnov	1.1	exe = job.type().executableName()
116
117	nsmirnov	1.2	txt += 'executable='+exe+'\n'
118			return txt
119
120			def runExe_(self):
121			"""
122			Returns part of a job script which executes the application.
123			"""
124	gutsche	1.9	job = common.job_list[self.nj]
125			args = job.type().executableArgs()
126			return '$executable '+args+'\n'
127	nsmirnov	1.1
128	nsmirnov	1.3	def renameOutput_(self):
129			"""
130			Returns part of a job script which renames output files.
131			"""
132			jbt = common.job_list.type()
133			txt = '\n'
134			txt += jbt.wsRenameOutput(self.nj)
135			return txt
136	fanzago	1.5
137	fanzago	1.8	def copyInput_(self):
138			"""
139			Returns part of a job script which copies input files from SE.
140			"""
141			txt = common.scheduler.wsCopyInput()
142			return txt
143
144	fanzago	1.5	def copyOutput_(self):
145			"""
146			Returns part of a job script which copies output files to SE.
147			"""
148			txt = common.scheduler.wsCopyOutput()
149			return txt
150
151	fanzago	1.11	def modifyReport_(self):
152			"""
153			Returns part of a job script which modifies the FrameworkJobReport.
154			"""
155			jbt = common.job_list.type()
156			txt = jbt.modifyReport(self.nj)
157			return txt
158	fanzago	1.12
159			def cleanEnv_(self):
160			"""
161			In OSG environment this function removes the WORKING_DIR
162			"""
163			jbt = common.job_list.type()
164			txt = jbt.cleanEnv()
165			return txt
166	mcinquil	1.14
167			def checkOut_(self):
168			"""
169			With glite check if the output is too big
170			"""
171			txt = "\n"
172			if self.output_troncate_flag == 1:
173	mcinquil	1.18	limit = 55000000 ##52 MB
174	mcinquil	1.14	jbt = common.job_list.type()
175			txt = jbt.checkOut(limit)
176			return txt
177	ewv	1.21	def rewriteCMSSWcfg_(self):
178			"""
179			Returns part of the script that runs writeCfg.py on the WN
180			"""
181			# FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
182			txt = "# Rewrite cfg for this job\n"
183
184			if self.CMSSW_major > 1: # Use py files for CMSSW 2_0_x and up
185			txt += "echo $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.pycfg\n"
186			txt += "python $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.pycfg\n"
187			else:
188			txt += "echo $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.cfg\n"
189			txt += "python $RUNTIME_AREA/writeCfg.py --debug pset.cfg pset.cfg\n"
190
191			return txt