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
#	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
8	import string,os
9
10	class ScriptWriter:
11	def __init__(self, template):
12	# pattern -> action
13	### FEDE aggiunte ultime due azioni
14	self.actions = {
15	'title' : self.title_,
16	'setup_scheduler_environment' : self.setupSchedulerEnvironment_,
17	'setup_jobtype_environment' : self.setupJobTypeEnvironment_,
18	'copy_input' : self.copyInput_,
19	'build_executable' : self.buildExe_,
20	'run_executable' : self.runExe_,
21	'rename_output' : self.renameOutput_,
22	'copy_output' : self.copyOutput_,
23	'register_output' : self.registerOutput_
24	}
25
26	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	self.nj = -1 # current job number
33	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	self.nj = nj
44
45	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	def title_(self):
73	txt = '# This script was generated by '+common.prog_name
74	txt += ' (version '+common.prog_version_str+').\n'
75	return txt
76
77	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
84	def setupJobTypeEnvironment_(self):
85	"""
86	Returns part of a job script which does jobtype-specific work.
87	"""
88	jbt = common.job_list.type()
89	txt = jbt.wsSetupEnvironment(self.nj)
90	return txt
91
92	def buildExe_(self):
93	"""
94	Returns part of a job script which builds the binary executable.
95	"""
96	jbt = common.job_list.type()
97
98	txt = jbt.wsBuildExe(self.nj)
99
100	job = common.job_list[self.nj]
101	exe = job.type().executableName()
102
103	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
112	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
121	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	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