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 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_,
            'build_executable'            : self.buildExe_,
            'run_executable'              : self.runExe_,
            'rename_output'               : self.renameOutput_,
            'copy_output'                 : self.copyOutput_,
            #'register_output'             : self.registerOutput_,
            '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

        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 registerOutput_(self):
    #    """
    #    Returns part of a job script which registers output files to RLS catalog.
    #    """
    #    txt = ''
    #    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 ##105 MB
            jbt = common.job_list.type()
            txt = jbt.checkOut(limit)
        return txt
Revision:	1.17
Committed:	Fri Nov 30 17:38:21 2007 UTC (17 years, 5 months ago) by mcinquil
Content type:	text/x-python
Branch:	MAIN
Changes since 1.16:	+1 -1 lines
Log Message:	Output sandbox limit reduced to 50MB (was 100MB)
#	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
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	nsmirnov	1.3	'build_executable' : self.buildExe_,
20			'run_executable' : self.runExe_,
21			'rename_output' : self.renameOutput_,
22	fanzago	1.5	'copy_output' : self.copyOutput_,
23	fanzago	1.15	#'register_output' : self.registerOutput_,
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
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
39	nsmirnov	1.1	return
40
41			def setAction(self, pattern, action):
42			self.actions[pattern] = action
43			return
44
45	slacapra	1.10	def modifyTemplateScript(self):
46	nsmirnov	1.1	"""
47			Create a script from scratch.
48			"""
49	nsmirnov	1.2
50	nsmirnov	1.1	tpl = open(self.template, 'r')
51	slacapra	1.10	script = open(common.taskDB.dict('ScriptName'),'w')
52	nsmirnov	1.1
53			for line in tpl:
54			if len(line) > 6 and line[:6] == '#CRAB ':
55			act_str = string.strip(line[6:])
56			try:
57			action = self.actions[act_str]
58			except KeyError:
59			continue
60
61			if action:
62			txt = action()
63			script.write(txt)
64			pass
65			else:
66			script.write(line)
67			pass
68			else:
69			script.write(line)
70			pass
71			pass
72
73			script.close()
74			tpl.close()
75			return
76
77	nsmirnov	1.2	def title_(self):
78	nsmirnov	1.1	txt = '# This script was generated by '+common.prog_name
79			txt += ' (version '+common.prog_version_str+').\n'
80			return txt
81
82	nsmirnov	1.3	def setupSchedulerEnvironment_(self):
83			"""
84			Returns part of a job script which does scheduler-specific work.
85			"""
86			txt = common.scheduler.wsSetupEnvironment()
87			return txt
88	nsmirnov	1.1
89	nsmirnov	1.3	def setupJobTypeEnvironment_(self):
90	nsmirnov	1.1	"""
91	fanzago	1.7	Returns part of a job script which does jobtype-specific work.
92	nsmirnov	1.1	"""
93	nsmirnov	1.3	jbt = common.job_list.type()
94			txt = jbt.wsSetupEnvironment(self.nj)
95			return txt
96
97	nsmirnov	1.2	def buildExe_(self):
98	nsmirnov	1.1	"""
99	nsmirnov	1.2	Returns part of a job script which builds the binary executable.
100	nsmirnov	1.1	"""
101	nsmirnov	1.2	jbt = common.job_list.type()
102	nsmirnov	1.1
103	nsmirnov	1.3	txt = jbt.wsBuildExe(self.nj)
104	nsmirnov	1.1
105	nsmirnov	1.2	job = common.job_list[self.nj]
106	nsmirnov	1.1	exe = job.type().executableName()
107
108	nsmirnov	1.2	txt += 'executable='+exe+'\n'
109			return txt
110
111			def runExe_(self):
112			"""
113			Returns part of a job script which executes the application.
114			"""
115	gutsche	1.9	job = common.job_list[self.nj]
116			args = job.type().executableArgs()
117			return '$executable '+args+'\n'
118	nsmirnov	1.1
119	nsmirnov	1.3	def renameOutput_(self):
120			"""
121			Returns part of a job script which renames output files.
122			"""
123			jbt = common.job_list.type()
124			txt = '\n'
125			txt += jbt.wsRenameOutput(self.nj)
126			return txt
127	fanzago	1.5
128	fanzago	1.8	def copyInput_(self):
129			"""
130			Returns part of a job script which copies input files from SE.
131			"""
132			txt = common.scheduler.wsCopyInput()
133			return txt
134
135	fanzago	1.5	def copyOutput_(self):
136			"""
137			Returns part of a job script which copies output files to SE.
138			"""
139			txt = common.scheduler.wsCopyOutput()
140			return txt
141
142	fanzago	1.15	#def registerOutput_(self):
143			# """
144			# Returns part of a job script which registers output files to RLS catalog.
145			# """
146			# txt = ''
147			# return txt
148	fanzago	1.5
149	fanzago	1.11	def modifyReport_(self):
150			"""
151			Returns part of a job script which modifies the FrameworkJobReport.
152			"""
153			jbt = common.job_list.type()
154			txt = jbt.modifyReport(self.nj)
155			return txt
156	fanzago	1.12
157			def cleanEnv_(self):
158			"""
159			In OSG environment this function removes the WORKING_DIR
160			"""
161			jbt = common.job_list.type()
162			txt = jbt.cleanEnv()
163			return txt
164	mcinquil	1.14
165			def checkOut_(self):
166			"""
167			With glite check if the output is too big
168			"""
169			txt = "\n"
170			if self.output_troncate_flag == 1:
171	mcinquil	1.17	limit = 55000000 ##105 MB
172	mcinquil	1.14	jbt = common.job_list.type()
173			txt = jbt.checkOut(limit)
174			return txt