CRAB/python/writeCfg.py

#!/usr/bin/env python

"""
Re-write config file and optionally convert to python
"""

__revision__ = "$Id: writeCfg.py,v 1.20 2009/06/17 20:58:08 ewv Exp $"
__version__ = "$Revision: 1.20 $"

import getopt
import imp
import os
import pickle
import sys
import xml.dom.minidom

from random import SystemRandom

from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
from FWCore.ParameterSet.Config                       import include
import FWCore.ParameterSet.Types as CfgTypes

MyRandom  = SystemRandom()

class ConfigException(Exception):
    """
    Exceptions raised by writeCfg
    """

    def __init__(self, msg):
        Exception.__init__(self, msg)
        self._msg = msg
        return

    def __str__(self):
        return self._msg

def main(argv) :
    """
    writeCfg

    - Read in existing, user supplied cfg or pickled pycfg file
    - Modify job specific parameters based on environment variables and arguments.xml
    - Write out modified cfg or pickled pycfg file

    required parameters: none

    optional parameters:
    --help             :       help
    --debug            :       debug statements

    """

    # defaults
    inputFileNames  = None
    parentFileNames = None
    debug           = False
    _MAXINT         = 900000000

    try:
        opts, args = getopt.getopt(argv, "", ["debug", "help"])
    except getopt.GetoptError:
        print main.__doc__
        sys.exit(2)

    try:
        CMSSW  = os.environ['CMSSW_VERSION']
        parts = CMSSW.split('_')
        CMSSW_major = int(parts[1])
        CMSSW_minor = int(parts[2])
        CMSSW_patch = int(parts[3])
    except (KeyError, ValueError):
        msg = "Your environment doesn't specify the CMSSW version or specifies it incorrectly"
        raise ConfigException(msg)

    # Parse command line options
    for opt, arg in opts :
        if opt  == "--help" :
            print main.__doc__
            sys.exit()
        elif opt == "--debug" :
            debug = True

    # Parse remaining parameters
    try:
        fileName    = args[0]
        outFileName = args[1]
    except IndexError:
        print main.__doc__
        sys.exit()

  # Read in Environment, XML and get optional Parameters

    nJob       = int(os.environ.get('NJob',      '0'))
    preserveSeeds  = os.environ.get('PreserveSeeds','')
    incrementSeeds = os.environ.get('IncrementSeeds','')

  # Defaults

    maxEvents  = 0
    skipEvents = 0
    firstEvent = -1
    compHEPFirstEvent = 0
    firstRun   = 0

    dom = xml.dom.minidom.parse(os.environ['RUNTIME_AREA']+'/arguments.xml')

    for elem in dom.getElementsByTagName("Job"):
        if nJob == int(elem.getAttribute("JobID")):
            if elem.getAttribute("MaxEvents"):
                maxEvents = int(elem.getAttribute("MaxEvents"))
            if elem.getAttribute("SkipEvents"):
                skipEvents = int(elem.getAttribute("SkipEvents"))
            if elem.getAttribute("FirstEvent"):
                firstEvent = int(elem.getAttribute("FirstEvent"))
            if elem.getAttribute("FirstRun"):
                firstRun = int(elem.getAttribute("FirstRun"))

            generator      = str(elem.getAttribute('Generator'))
            inputFiles     = str(elem.getAttribute('InputFiles'))
            parentFiles    = str(elem.getAttribute('ParentFiles'))

  # Read Input cfg or python cfg file, FUTURE: Get rid cfg mode

    if fileName.endswith('py'):
        handle = open(fileName, 'r')
        try:   # Nested form for Python < 2.5
            try:
                print "Importing .py file"
                cfo = imp.load_source("pycfg", fileName, handle)
                cmsProcess = cfo.process
            except Exception, ex:
                msg = "Your pycfg file is not valid python: %s" % str(ex)
                raise ConfigException(msg)
        finally:
            handle.close()
    else:
        try:
            print "Importing .cfg file"
            cfo = include(fileName)
            cmsProcess = cfo
        except Exception, ex:
            msg =  "The cfg file is not valid, %s\n" % str(ex)
            raise ConfigException(msg)

    cfg = CfgInterface(cmsProcess)

    # Set parameters for job
    print "Setting parameters"
    inModule = cfg.inputSource
    if maxEvents:
        cfg.maxEvents.setMaxEventsInput(maxEvents)

    if skipEvents:
        inModule.setSkipEvents(skipEvents)

    # Set "skip events" for various generators
    if generator == 'comphep':
        cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(firstEvent)
    elif generator == 'lhe':
        cmsProcess.source.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
        cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent+1))
    elif firstEvent != -1: # (Old? Madgraph)
        cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))

    if inputFiles:
        inputFileNames = inputFiles.split(',')
        inModule.setFileNames(*inputFileNames)

    # handle parent files if needed
    if parentFiles:
        parentFileNames = parentFiles.split(',')
        inModule.setSecondaryFileNames(*parentFileNames)

    # Pythia parameters
    if (firstRun):
        inModule.setFirstRun(firstRun)

    incrementSeedList = []
    preserveSeedList  = []

    if incrementSeeds:
        incrementSeedList = incrementSeeds.split(',')
    if preserveSeeds:
        preserveSeedList  = preserveSeeds.split(',')

    # FUTURE: This function tests the CMSSW version and presence of old-style seed specification.
    # Reduce when we drop support for old versions
    if cfg.data.services.has_key('RandomNumberGeneratorService'): # There are random #'s to deal with
        print "RandomNumberGeneratorService found, will attempt to change seeds"
        ranGenerator = cfg.data.services['RandomNumberGeneratorService']
        ranModules   = getattr(ranGenerator, "moduleSeeds", None)
        oldSource    = getattr(ranGenerator, "sourceSeed",  None)
        if ranModules != None or oldSource != None:     # Old format present, no matter the CMSSW version
            print "Old-style random number seeds found, will be changed."
            if oldSource != None:
                sourceSeed = int(ranGenerator.sourceSeed.value())
                if ('sourceSeed' in preserveSeedList) or ('theSource' in preserveSeedList):
                    pass
                elif ('sourceSeed' in incrementSeedList) or ('theSource' in incrementSeedList):
                    ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
                else:
                    ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1, _MAXINT)))

            for seed in incrementSeedList:
                curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
                if curSeed:
                    curValue = int(curSeed.value())
                    setattr(ranGenerator.moduleSeeds, seed, CfgTypes.untracked(CfgTypes.uint32(curValue+nJob)))
                    preserveSeedList.append(seed)

            if ranModules != None:
                for seed in ranGenerator.moduleSeeds.parameterNames_():
                    if seed not in preserveSeedList:
                        curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
                        if curSeed:
                            curValue = int(curSeed.value())
                            setattr(ranGenerator.moduleSeeds, seed,
                                    CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1,_MAXINT))))
        elif CMSSW_major > 2 or (CMSSW_major == 2 and CMSSW_minor >= 1): # Treatment for  seeds, CMSSW 2_1_x and later
            print "New-style random number seeds found, will be changed."
            from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
            randSvc = RandomNumberServiceHelper(ranGenerator)

            # Increment requested seed sets
            for seedName in incrementSeedList:
                curSeeds = randSvc.getNamedSeed(seedName)
                newSeeds = [x+nJob for x in curSeeds]
                randSvc.setNamedSeed(seedName, *newSeeds)
                preserveSeedList.append(seedName)

            # Randomize remaining seeds
            randSvc.populate(*preserveSeedList)
        else:
            print "Neither old nor new seed format found!"

      # End version specific code

    # Write out new config file in one format or the other, FUTURE: Get rid of cfg mode
    outFile = open(outFileName,"w")
    if outFileName.endswith('py'):
        outFile.write("import FWCore.ParameterSet.Config as cms\n")
        outFile.write("import pickle\n")
        outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(cmsProcess))
        outFile.write("process = pickle.loads(pickledCfg)\n")
        if (debug):
            print "writeCfg output (May not be exact):"
            print "import FWCore.ParameterSet.Config as cms"
            print cmsProcess.dumpPython()
    else:
        outFile.write(cfg.data.dumpConfig())
        if (debug):
            print "writeCfg output:"
            print str(cfg.data.dumpConfig())
    outFile.close()


if __name__ == '__main__' :
    exit_status = main(sys.argv[1:])
    sys.exit(exit_status)
Revision:	1.21
Committed:	Fri Jun 19 18:46:54 2009 UTC (15 years, 10 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_6_1_pre4, CRAB_2_6_1_pre3, CRAB_2_6_1_pre2, CRAB_2_6_1_pre1, CRAB_2_6_1, CRAB_2_6_0, CRAB_2_6_0_pre14, CRAB_2_6_0_pre13, CRAB_2_6_0_pre12
Changes since 1.20:	+4 -3 lines
Log Message:	Do LHE right, add setting of firstEvent
#	User	Rev	Content
1	ewv	1.1	#!/usr/bin/env python
2
3	ewv	1.12	"""
4			Re-write config file and optionally convert to python
5			"""
6
7	ewv	1.21	__revision__ = "$Id: writeCfg.py,v 1.20 2009/06/17 20:58:08 ewv Exp $"
8			__version__ = "$Revision: 1.20 $"
9	ewv	1.13
10	ewv	1.18	import getopt
11	ewv	1.1	import imp
12			import os
13	ewv	1.17	import pickle
14	ewv	1.18	import sys
15			import xml.dom.minidom
16
17	ewv	1.2	from random import SystemRandom
18	ewv	1.1
19			from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
20	ewv	1.12	from FWCore.ParameterSet.Config import include
21			import FWCore.ParameterSet.Types as CfgTypes
22	ewv	1.1
23	ewv	1.12	MyRandom = SystemRandom()
24	ewv	1.1
25	ewv	1.3	class ConfigException(Exception):
26	ewv	1.12	"""
27			Exceptions raised by writeCfg
28			"""
29
30	ewv	1.3	def __init__(self, msg):
31			Exception.__init__(self, msg)
32			self._msg = msg
33			return
34
35			def __str__(self):
36			return self._msg
37	ewv	1.1
38			def main(argv) :
39	ewv	1.12	"""
40			writeCfg
41
42	ewv	1.20	- Read in existing, user supplied cfg or pickled pycfg file
43			- Modify job specific parameters based on environment variables and arguments.xml
44			- Write out modified cfg or pickled pycfg file
45	ewv	1.12
46			required parameters: none
47
48			optional parameters:
49			--help : help
50			--debug : debug statements
51
52			"""
53
54			# defaults
55			inputFileNames = None
56			parentFileNames = None
57			debug = False
58			_MAXINT = 900000000
59
60			try:
61			opts, args = getopt.getopt(argv, "", ["debug", "help"])
62			except getopt.GetoptError:
63			print main.__doc__
64			sys.exit(2)
65	ewv	1.1
66			try:
67	ewv	1.12	CMSSW = os.environ['CMSSW_VERSION']
68			parts = CMSSW.split('_')
69			CMSSW_major = int(parts[1])
70			CMSSW_minor = int(parts[2])
71			CMSSW_patch = int(parts[3])
72			except (KeyError, ValueError):
73			msg = "Your environment doesn't specify the CMSSW version or specifies it incorrectly"
74			raise ConfigException(msg)
75
76			# Parse command line options
77			for opt, arg in opts :
78			if opt == "--help" :
79			print main.__doc__
80			sys.exit()
81			elif opt == "--debug" :
82			debug = True
83	ewv	1.2
84	ewv	1.12	# Parse remaining parameters
85			try:
86			fileName = args[0]
87			outFileName = args[1]
88			except IndexError:
89			print main.__doc__
90			sys.exit()
91
92	ewv	1.19	# Read in Environment, XML and get optional Parameters
93	ewv	1.12
94			nJob = int(os.environ.get('NJob', '0'))
95	ewv	1.19	preserveSeeds = os.environ.get('PreserveSeeds','')
96			incrementSeeds = os.environ.get('IncrementSeeds','')
97	ewv	1.12
98	ewv	1.18	# Defaults
99
100			maxEvents = 0
101			skipEvents = 0
102			firstEvent = -1
103			compHEPFirstEvent = 0
104			firstRun = 0
105
106			dom = xml.dom.minidom.parse(os.environ['RUNTIME_AREA']+'/arguments.xml')
107
108			for elem in dom.getElementsByTagName("Job"):
109			if nJob == int(elem.getAttribute("JobID")):
110			if elem.getAttribute("MaxEvents"):
111			maxEvents = int(elem.getAttribute("MaxEvents"))
112			if elem.getAttribute("SkipEvents"):
113			skipEvents = int(elem.getAttribute("SkipEvents"))
114			if elem.getAttribute("FirstEvent"):
115			firstEvent = int(elem.getAttribute("FirstEvent"))
116			if elem.getAttribute("FirstRun"):
117			firstRun = int(elem.getAttribute("FirstRun"))
118
119	ewv	1.20	generator = str(elem.getAttribute('Generator'))
120	ewv	1.18	inputFiles = str(elem.getAttribute('InputFiles'))
121			parentFiles = str(elem.getAttribute('ParentFiles'))
122	ewv	1.12
123			# Read Input cfg or python cfg file, FUTURE: Get rid cfg mode
124
125			if fileName.endswith('py'):
126			handle = open(fileName, 'r')
127			try: # Nested form for Python < 2.5
128			try:
129			print "Importing .py file"
130			cfo = imp.load_source("pycfg", fileName, handle)
131			cmsProcess = cfo.process
132			except Exception, ex:
133			msg = "Your pycfg file is not valid python: %s" % str(ex)
134			raise ConfigException(msg)
135			finally:
136			handle.close()
137			else:
138			try:
139			print "Importing .cfg file"
140			cfo = include(fileName)
141			cmsProcess = cfo
142			except Exception, ex:
143			msg = "The cfg file is not valid, %s\n" % str(ex)
144			raise ConfigException(msg)
145
146			cfg = CfgInterface(cmsProcess)
147
148			# Set parameters for job
149			print "Setting parameters"
150			inModule = cfg.inputSource
151			if maxEvents:
152			cfg.maxEvents.setMaxEventsInput(maxEvents)
153
154			if skipEvents:
155			inModule.setSkipEvents(skipEvents)
156	ewv	1.20
157			# Set "skip events" for various generators
158			if generator == 'comphep':
159			cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(firstEvent)
160			elif generator == 'lhe':
161	ewv	1.21	cmsProcess.source.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
162			cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent+1))
163	ewv	1.20	elif firstEvent != -1: # (Old? Madgraph)
164	ewv	1.14	cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
165	ewv	1.20
166	ewv	1.12	if inputFiles:
167			inputFileNames = inputFiles.split(',')
168			inModule.setFileNames(*inputFileNames)
169
170			# handle parent files if needed
171			if parentFiles:
172			parentFileNames = parentFiles.split(',')
173			inModule.setSecondaryFileNames(*parentFileNames)
174
175			# Pythia parameters
176			if (firstRun):
177			inModule.setFirstRun(firstRun)
178
179			incrementSeedList = []
180			preserveSeedList = []
181
182			if incrementSeeds:
183			incrementSeedList = incrementSeeds.split(',')
184			if preserveSeeds:
185			preserveSeedList = preserveSeeds.split(',')
186
187			# FUTURE: This function tests the CMSSW version and presence of old-style seed specification.
188			# Reduce when we drop support for old versions
189			if cfg.data.services.has_key('RandomNumberGeneratorService'): # There are random #'s to deal with
190			print "RandomNumberGeneratorService found, will attempt to change seeds"
191			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
192			ranModules = getattr(ranGenerator, "moduleSeeds", None)
193			oldSource = getattr(ranGenerator, "sourceSeed", None)
194			if ranModules != None or oldSource != None: # Old format present, no matter the CMSSW version
195			print "Old-style random number seeds found, will be changed."
196			if oldSource != None:
197			sourceSeed = int(ranGenerator.sourceSeed.value())
198			if ('sourceSeed' in preserveSeedList) or ('theSource' in preserveSeedList):
199			pass
200			elif ('sourceSeed' in incrementSeedList) or ('theSource' in incrementSeedList):
201			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
202			else:
203			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1, _MAXINT)))
204
205			for seed in incrementSeedList:
206			curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
207			if curSeed:
208			curValue = int(curSeed.value())
209			setattr(ranGenerator.moduleSeeds, seed, CfgTypes.untracked(CfgTypes.uint32(curValue+nJob)))
210			preserveSeedList.append(seed)
211
212			if ranModules != None:
213			for seed in ranGenerator.moduleSeeds.parameterNames_():
214			if seed not in preserveSeedList:
215			curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
216			if curSeed:
217			curValue = int(curSeed.value())
218			setattr(ranGenerator.moduleSeeds, seed,
219			CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1,_MAXINT))))
220			elif CMSSW_major > 2 or (CMSSW_major == 2 and CMSSW_minor >= 1): # Treatment for seeds, CMSSW 2_1_x and later
221			print "New-style random number seeds found, will be changed."
222			from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
223			randSvc = RandomNumberServiceHelper(ranGenerator)
224
225			# Increment requested seed sets
226			for seedName in incrementSeedList:
227			curSeeds = randSvc.getNamedSeed(seedName)
228			newSeeds = [x+nJob for x in curSeeds]
229			randSvc.setNamedSeed(seedName, *newSeeds)
230			preserveSeedList.append(seedName)
231
232			# Randomize remaining seeds
233			randSvc.populate(*preserveSeedList)
234			else:
235			print "Neither old nor new seed format found!"
236
237			# End version specific code
238
239			# Write out new config file in one format or the other, FUTURE: Get rid of cfg mode
240			outFile = open(outFileName,"w")
241			if outFileName.endswith('py'):
242			outFile.write("import FWCore.ParameterSet.Config as cms\n")
243	ewv	1.17	outFile.write("import pickle\n")
244			outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(cmsProcess))
245			outFile.write("process = pickle.loads(pickledCfg)\n")
246	ewv	1.12	if (debug):
247	ewv	1.17	print "writeCfg output (May not be exact):"
248	ewv	1.12	print "import FWCore.ParameterSet.Config as cms"
249			print cmsProcess.dumpPython()
250	ewv	1.9	else:
251	ewv	1.13	outFile.write(cfg.data.dumpConfig())
252	ewv	1.12	if (debug):
253			print "writeCfg output:"
254	ewv	1.13	print str(cfg.data.dumpConfig())
255	ewv	1.12	outFile.close()
256	ewv	1.1
257
258			if __name__ == '__main__' :
259			exit_status = main(sys.argv[1:])
260			sys.exit(exit_status)