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):
        # pattern -> action
        ### FEDE aggiunte ultime due azioni
        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_
            }
        
        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
        return

    def setAction(self, pattern, action):
        self.actions[pattern] = action
        return
    
    def modifyTemplateScript(self, nj):
        """
        Create a script from scratch.
        """
        self.nj = nj
        
        tpl = open(self.template, 'r')
        script = open(common.job_list[nj].scriptFilename(), '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.
        """
        return '$executable\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 = common.scheduler.wsRegisterOutput()
        return txt

Revision:	1.8
Committed:	Mon Feb 20 18:02:06 2006 UTC (19 years, 2 months ago) by fanzago
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	pre_cmssw_integration_20060527, CRAB_1_1_0, CRAB_1_1_0_pre4, CRAB_1_1_0_pre3, CRAB_1_1_0_pre1, CRAB_1_0_7, CRAB_1_0_7_pre1, CRAB_1_0_6, CRAB_1_0_5
Changes since 1.7:	+8 -1 lines
Log Message:	added new function CopyInput
#	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			def __init__(self, template):
12			# pattern -> action
13	fanzago	1.5	### FEDE aggiunte ultime due azioni
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			'register_output' : self.registerOutput_
24	nsmirnov	1.1	}
25
26	slacapra	1.6	if os.path.isfile("./"+template):
27			self.template = "./"+template
28			elif os.getenv('CRABDIR') and os.path.isfile(os.getenv('CRABDIR')+'/python/'+template):
29			self.template = os.getenv('CRABDIR')+'/python/'+template
30			else:
31			raise CrabException("No crab_template.sh found!")
32	nsmirnov	1.2	self.nj = -1 # current job number
33	nsmirnov	1.1	return
34
35			def setAction(self, pattern, action):
36			self.actions[pattern] = action
37			return
38
39			def modifyTemplateScript(self, nj):
40			"""
41			Create a script from scratch.
42			"""
43	nsmirnov	1.2	self.nj = nj
44
45	nsmirnov	1.1	tpl = open(self.template, 'r')
46			script = open(common.job_list[nj].scriptFilename(), 'w')
47
48			for line in tpl:
49			if len(line) > 6 and line[:6] == '#CRAB ':
50			act_str = string.strip(line[6:])
51			try:
52			action = self.actions[act_str]
53			except KeyError:
54			continue
55
56			if action:
57			txt = action()
58			script.write(txt)
59			pass
60			else:
61			script.write(line)
62			pass
63			else:
64			script.write(line)
65			pass
66			pass
67
68			script.close()
69			tpl.close()
70			return
71
72	nsmirnov	1.2	def title_(self):
73	nsmirnov	1.1	txt = '# This script was generated by '+common.prog_name
74			txt += ' (version '+common.prog_version_str+').\n'
75			return txt
76
77	nsmirnov	1.3	def setupSchedulerEnvironment_(self):
78			"""
79			Returns part of a job script which does scheduler-specific work.
80			"""
81			txt = common.scheduler.wsSetupEnvironment()
82			return txt
83	nsmirnov	1.1
84	nsmirnov	1.3	def setupJobTypeEnvironment_(self):
85	nsmirnov	1.1	"""
86	fanzago	1.7	Returns part of a job script which does jobtype-specific work.
87	nsmirnov	1.1	"""
88	nsmirnov	1.3	jbt = common.job_list.type()
89			txt = jbt.wsSetupEnvironment(self.nj)
90			return txt
91
92	nsmirnov	1.2	def buildExe_(self):
93	nsmirnov	1.1	"""
94	nsmirnov	1.2	Returns part of a job script which builds the binary executable.
95	nsmirnov	1.1	"""
96	nsmirnov	1.2	jbt = common.job_list.type()
97	nsmirnov	1.1
98	nsmirnov	1.3	txt = jbt.wsBuildExe(self.nj)
99	nsmirnov	1.1
100	nsmirnov	1.2	job = common.job_list[self.nj]
101	nsmirnov	1.1	exe = job.type().executableName()
102
103	nsmirnov	1.2	txt += 'executable='+exe+'\n'
104			return txt
105
106			def runExe_(self):
107			"""
108			Returns part of a job script which executes the application.
109			"""
110			return '$executable\n'
111	nsmirnov	1.1
112	nsmirnov	1.3	def renameOutput_(self):
113			"""
114			Returns part of a job script which renames output files.
115			"""
116			jbt = common.job_list.type()
117			txt = '\n'
118			txt += jbt.wsRenameOutput(self.nj)
119			return txt
120	fanzago	1.5
121	fanzago	1.8	def copyInput_(self):
122			"""
123			Returns part of a job script which copies input files from SE.
124			"""
125			txt = common.scheduler.wsCopyInput()
126			return txt
127
128	fanzago	1.5	def copyOutput_(self):
129			"""
130			Returns part of a job script which copies output files to SE.
131			"""
132			txt = common.scheduler.wsCopyOutput()
133			return txt
134
135			def registerOutput_(self):
136			"""
137			Returns part of a job script which registers output files to RLS catalog.
138			"""
139			txt = common.scheduler.wsRegisterOutput()
140			return txt
141