CRAB/python/PsetManipulator.py

#!/usr/bin/env python

import os
import common
import imp
from crab_util import *
from crab_exceptions import *
from crab_logger import Logger

from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
from FWCore.ParameterSet.DictTypes import SortedKeysDict
from FWCore.ParameterSet.Modules   import Service
from FWCore.ParameterSet.Types     import *

import FWCore.ParameterSet.Types   as CfgTypes
import FWCore.ParameterSet.Modules as CfgModules

class PsetManipulator:
    def __init__(self, pset):
        """
        Read in Pset object and initialize
        """

        self.pset = pset
        #convert Pset
        from FWCore.ParameterSet.Config import include
        common.logger.debug(3,"PsetManipulator::__init__: PSet file = "+self.pset)
        if (self.pset.endswith('py') or self.pset.endswith('pycfg') ):
          handle = open(self.pset, 'r')
          try:   # Nested form for Python < 2.5
            try:
              self.cfo = imp.load_source("pycfg", self.pset, handle)
            except Exception, ex:
              msg = "Your pycfg file is not valid python: %s" % str(ex)
              raise CrabException(msg)
          finally:
              handle.close()
          self.cfg = CfgInterface(self.cfo.process)
        else:
          try:
            self.cfo = include(self.pset)
            self.cfg = CfgInterface(self.cfo)
          except Exception, ex:
            msg =  "Your cfg file is not valid, %s\n" % str(ex)
            msg += "  https://twiki.cern.ch/twiki/bin/view/CMS/CrabFaq#Problem_with_ParameterSet_parsin\n"
            msg += "  may help you understand the problem."
            raise CrabException(msg)
    def inputModule(self, source):
        """
        Set  vString Filenames key
        """
        # set input module
        inModule = self.cfg.inputSource
        inModule.setFileNames(source)
        return

    def pythiaSeed(self,seed):
        """
        Set pythia seed key
        """
        ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
        ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(seed))
        return

    def vtxSeed(self,vtxSeed):
        """
        Set vtx seed key
        """
        ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
        ranModules   = ranGenerator.moduleSeeds
        # set seed
        ranModules.VtxSmeared = CfgTypes.untracked(CfgTypes.uint32(vtxSeed))
        return

    def g4Seed(self,g4Seed):
        """
        Set g4 seed key
        """
        ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
        ranModules   = ranGenerator.moduleSeeds
        # set seed
        ranModules.g4SimHits = CfgTypes.untracked(CfgTypes.uint32(g4Seed))
        return

    def mixSeed(self,mixSeed):
        """
        Set mix seed key
        """
        ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
        ranModules   = ranGenerator.moduleSeeds
        ranModules.mix = CfgTypes.untracked(CfgTypes.uint32(mixSeed))
        return

    def pythiaFirstRun(self, firstrun):
        """
        Set firstRun
        """
        inModule = self.cfg.inputSource
        inModule.setFirstRun(firstrun)   ## Add Daniele
        return

    def maxEvent(self, maxEv):
        """
        Set max event in the standalone untracked module
        """
        self.cfg.maxEvents.setMaxEventsInput(maxEv)
        return

    def skipEvent(self, skipEv):
        """
        Set skipEvents
        """
        inModule = self.cfg.inputSource
        inModule.setSkipEvents(skipEv)   ## Add Daniele
        return

    def outputModule(self, output):

        #set output module
        outModule = self.cfg.outputModules['out']
        outModule.setFileName('file:'+str(output))

        return

    def psetWriter(self, name):
        """
        Write out modified CMSSW.cfg
        """

        file1 = open(common.work_space.jobDir()+name,"w")
        file1.write(str(self.cfg))
        file1.close()

        return

    def addCrabFJR(self,name):
        """
        _addCrabFJR_
        add CRAB specific FrameworkJobReport (FJR)
        if a FJR already exists in input CMSSW parameter-set, add a second one.
        This code is not needed for CMSSW >= 1.5.x and is non-functional in CMSSW >= 1.7.x.
        It should be removed at some point in the future.
        """

        # Check if MessageLogger service already exists in configuration. If not, add it
        svcs = self.cfg.data.services
        if not svcs.has_key('MessageLogger'):
            self.cfg.data.add_(CfgModules.Service("MessageLogger"))

        messageLogger = self.cfg.data.services['MessageLogger']

        # Add fwkJobReports to Message logger if it doesn't exist
        if "fwkJobReports" not in messageLogger.parameterNames_():
            messageLogger.fwkJobReports = CfgTypes.untracked(CfgTypes.vstring())

        # should figure out how to remove "name" if it is there.

        if name not in messageLogger.fwkJobReports:
            messageLogger.fwkJobReports.append(name)

        return
Revision:	1.14
Committed:	Fri Dec 7 21:41:20 2007 UTC (17 years, 4 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_0_4, CRAB_2_0_4_pre2, CRAB_2_0_4_pre1
Changes since 1.13:	+21 -3 lines
Log Message:	Allow user to supply py/pycfg as config file
#	User	Rev	Content
1	gutsche	1.1	#!/usr/bin/env python
2	ewv	1.12
3	slacapra	1.8	import os
4	gutsche	1.1	import common
5	ewv	1.14	import imp
6	gutsche	1.1	from crab_util import *
7			from crab_exceptions import *
8	ewv	1.12	from crab_logger import Logger
9
10			from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
11			from FWCore.ParameterSet.DictTypes import SortedKeysDict
12			from FWCore.ParameterSet.Modules import Service
13			from FWCore.ParameterSet.Types import *
14	gutsche	1.1
15	ewv	1.12	import FWCore.ParameterSet.Types as CfgTypes
16			import FWCore.ParameterSet.Modules as CfgModules
17	gutsche	1.1
18			class PsetManipulator:
19			def __init__(self, pset):
20	ewv	1.12	"""
21			Read in Pset object and initialize
22	gutsche	1.1	"""
23
24			self.pset = pset
25			#convert Pset
26	ewv	1.12	from FWCore.ParameterSet.Config import include
27			common.logger.debug(3,"PsetManipulator::__init__: PSet file = "+self.pset)
28	ewv	1.14	if (self.pset.endswith('py') or self.pset.endswith('pycfg') ):
29			handle = open(self.pset, 'r')
30			try: # Nested form for Python < 2.5
31			try:
32			self.cfo = imp.load_source("pycfg", self.pset, handle)
33			except Exception, ex:
34			msg = "Your pycfg file is not valid python: %s" % str(ex)
35			raise CrabException(msg)
36			finally:
37			handle.close()
38			self.cfg = CfgInterface(self.cfo.process)
39			else:
40			try:
41			self.cfo = include(self.pset)
42			self.cfg = CfgInterface(self.cfo)
43			except Exception, ex:
44			msg = "Your cfg file is not valid, %s\n" % str(ex)
45			msg += " https://twiki.cern.ch/twiki/bin/view/CMS/CrabFaq#Problem_with_ParameterSet_parsin\n"
46			msg += " may help you understand the problem."
47			raise CrabException(msg)
48	gutsche	1.1	def inputModule(self, source):
49	ewv	1.12	"""
50			Set vString Filenames key
51	gutsche	1.1	"""
52			# set input module
53			inModule = self.cfg.inputSource
54			inModule.setFileNames(source)
55			return
56	ewv	1.12
57	slacapra	1.4	def pythiaSeed(self,seed):
58	ewv	1.12	"""
59	slacapra	1.4	Set pythia seed key
60	slacapra	1.2	"""
61	ewv	1.12	ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
62			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(seed))
63			return
64	slacapra	1.4
65	slacapra	1.10	def vtxSeed(self,vtxSeed):
66	ewv	1.12	"""
67	slacapra	1.4	Set vtx seed key
68			"""
69	ewv	1.12	ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
70			ranModules = ranGenerator.moduleSeeds
71	slacapra	1.4	# set seed
72	ewv	1.12	ranModules.VtxSmeared = CfgTypes.untracked(CfgTypes.uint32(vtxSeed))
73			return
74	slacapra	1.10
75			def g4Seed(self,g4Seed):
76	ewv	1.12	"""
77	slacapra	1.10	Set g4 seed key
78			"""
79	ewv	1.12	ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
80			ranModules = ranGenerator.moduleSeeds
81	slacapra	1.10	# set seed
82	ewv	1.12	ranModules.g4SimHits = CfgTypes.untracked(CfgTypes.uint32(g4Seed))
83			return
84	slacapra	1.10
85			def mixSeed(self,mixSeed):
86	ewv	1.12	"""
87	slacapra	1.10	Set mix seed key
88			"""
89	ewv	1.12	ranGenerator = self.cfg.data.services['RandomNumberGeneratorService']
90			ranModules = ranGenerator.moduleSeeds
91			ranModules.mix = CfgTypes.untracked(CfgTypes.uint32(mixSeed))
92			return
93	gutsche	1.1
94	spiga	1.7	def pythiaFirstRun(self, firstrun):
95	ewv	1.12	"""
96			Set firstRun
97			"""
98	spiga	1.7	inModule = self.cfg.inputSource
99	ewv	1.12	inModule.setFirstRun(firstrun) ## Add Daniele
100	spiga	1.7	return
101
102	gutsche	1.1	def maxEvent(self, maxEv):
103	ewv	1.12	"""
104			Set max event in the standalone untracked module
105			"""
106			self.cfg.maxEvents.setMaxEventsInput(maxEv)
107	gutsche	1.1	return
108
109	gutsche	1.5	def skipEvent(self, skipEv):
110	ewv	1.12	"""
111			Set skipEvents
112			"""
113	gutsche	1.5	inModule = self.cfg.inputSource
114	ewv	1.12	inModule.setSkipEvents(skipEv) ## Add Daniele
115	gutsche	1.5	return
116
117	gutsche	1.1	def outputModule(self, output):
118
119			#set output module
120			outModule = self.cfg.outputModules['out']
121			outModule.setFileName('file:'+str(output))
122
123			return
124
125			def psetWriter(self, name):
126	ewv	1.12	"""
127			Write out modified CMSSW.cfg
128			"""
129	gutsche	1.1
130			file1 = open(common.work_space.jobDir()+name,"w")
131	ewv	1.12	file1.write(str(self.cfg))
132	gutsche	1.1	file1.close()
133
134			return
135	gutsche	1.6
136			def addCrabFJR(self,name):
137			"""
138	ewv	1.12	_addCrabFJR_
139			add CRAB specific FrameworkJobReport (FJR)
140	ewv	1.13	if a FJR already exists in input CMSSW parameter-set, add a second one.
141			This code is not needed for CMSSW >= 1.5.x and is non-functional in CMSSW >= 1.7.x.
142			It should be removed at some point in the future.
143	ewv	1.12	"""
144
145			# Check if MessageLogger service already exists in configuration. If not, add it
146			svcs = self.cfg.data.services
147			if not svcs.has_key('MessageLogger'):
148			self.cfg.data.add_(CfgModules.Service("MessageLogger"))
149	gutsche	1.6
150	ewv	1.12	messageLogger = self.cfg.data.services['MessageLogger']
151	gutsche	1.6
152	ewv	1.12	# Add fwkJobReports to Message logger if it doesn't exist
153			if "fwkJobReports" not in messageLogger.parameterNames_():
154			messageLogger.fwkJobReports = CfgTypes.untracked(CfgTypes.vstring())
155	gutsche	1.6
156	ewv	1.12	# should figure out how to remove "name" if it is there.
157	gutsche	1.6
158	ewv	1.12	if name not in messageLogger.fwkJobReports:
159			messageLogger.fwkJobReports.append(name)
160	gutsche	1.6
161	ewv	1.12	return