CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.21 2009/06/19 18:46:54 ewv Exp $"
__version__ = "$Revision: 1.21 $"

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'))
            lumis          = str(elem.getAttribute('Lumis'))

  # 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)

    if lumis:
        lumiRanges = lumis.split(',')
        inModule.setLumisToProcess(*lumiRanges)

    # 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.22
Committed:	Wed Jul 29 21:20:03 2009 UTC (15 years, 9 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
Changes since 1.21:	+7 -2 lines
Log Message:	Set lumis to process for ADS
#	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.22	__revision__ = "$Id: writeCfg.py,v 1.21 2009/06/19 18:46:54 ewv Exp $"
8			__version__ = "$Revision: 1.21 $"
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.22	lumis = str(elem.getAttribute('Lumis'))
123	ewv	1.12
124			# Read Input cfg or python cfg file, FUTURE: Get rid cfg mode
125
126			if fileName.endswith('py'):
127			handle = open(fileName, 'r')
128			try: # Nested form for Python < 2.5
129			try:
130			print "Importing .py file"
131			cfo = imp.load_source("pycfg", fileName, handle)
132			cmsProcess = cfo.process
133			except Exception, ex:
134			msg = "Your pycfg file is not valid python: %s" % str(ex)
135			raise ConfigException(msg)
136			finally:
137			handle.close()
138			else:
139			try:
140			print "Importing .cfg file"
141			cfo = include(fileName)
142			cmsProcess = cfo
143			except Exception, ex:
144			msg = "The cfg file is not valid, %s\n" % str(ex)
145			raise ConfigException(msg)
146
147			cfg = CfgInterface(cmsProcess)
148
149			# Set parameters for job
150			print "Setting parameters"
151			inModule = cfg.inputSource
152			if maxEvents:
153			cfg.maxEvents.setMaxEventsInput(maxEvents)
154
155			if skipEvents:
156			inModule.setSkipEvents(skipEvents)
157	ewv	1.20
158			# Set "skip events" for various generators
159			if generator == 'comphep':
160			cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(firstEvent)
161			elif generator == 'lhe':
162	ewv	1.21	cmsProcess.source.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
163			cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent+1))
164	ewv	1.20	elif firstEvent != -1: # (Old? Madgraph)
165	ewv	1.14	cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
166	ewv	1.20
167	ewv	1.12	if inputFiles:
168			inputFileNames = inputFiles.split(',')
169			inModule.setFileNames(*inputFileNames)
170
171			# handle parent files if needed
172			if parentFiles:
173			parentFileNames = parentFiles.split(',')
174			inModule.setSecondaryFileNames(*parentFileNames)
175
176	ewv	1.22	if lumis:
177			lumiRanges = lumis.split(',')
178			inModule.setLumisToProcess(*lumiRanges)
179
180	ewv	1.12	# Pythia parameters
181			if (firstRun):
182			inModule.setFirstRun(firstRun)
183
184			incrementSeedList = []
185			preserveSeedList = []
186
187			if incrementSeeds:
188			incrementSeedList = incrementSeeds.split(',')
189			if preserveSeeds:
190			preserveSeedList = preserveSeeds.split(',')
191
192			# FUTURE: This function tests the CMSSW version and presence of old-style seed specification.
193			# Reduce when we drop support for old versions
194			if cfg.data.services.has_key('RandomNumberGeneratorService'): # There are random #'s to deal with
195			print "RandomNumberGeneratorService found, will attempt to change seeds"
196			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
197			ranModules = getattr(ranGenerator, "moduleSeeds", None)
198			oldSource = getattr(ranGenerator, "sourceSeed", None)
199			if ranModules != None or oldSource != None: # Old format present, no matter the CMSSW version
200			print "Old-style random number seeds found, will be changed."
201			if oldSource != None:
202			sourceSeed = int(ranGenerator.sourceSeed.value())
203			if ('sourceSeed' in preserveSeedList) or ('theSource' in preserveSeedList):
204			pass
205			elif ('sourceSeed' in incrementSeedList) or ('theSource' in incrementSeedList):
206			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
207			else:
208			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1, _MAXINT)))
209
210			for seed in incrementSeedList:
211			curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
212			if curSeed:
213			curValue = int(curSeed.value())
214			setattr(ranGenerator.moduleSeeds, seed, CfgTypes.untracked(CfgTypes.uint32(curValue+nJob)))
215			preserveSeedList.append(seed)
216
217			if ranModules != None:
218			for seed in ranGenerator.moduleSeeds.parameterNames_():
219			if seed not in preserveSeedList:
220			curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
221			if curSeed:
222			curValue = int(curSeed.value())
223			setattr(ranGenerator.moduleSeeds, seed,
224			CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1,_MAXINT))))
225			elif CMSSW_major > 2 or (CMSSW_major == 2 and CMSSW_minor >= 1): # Treatment for seeds, CMSSW 2_1_x and later
226			print "New-style random number seeds found, will be changed."
227			from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
228			randSvc = RandomNumberServiceHelper(ranGenerator)
229
230			# Increment requested seed sets
231			for seedName in incrementSeedList:
232			curSeeds = randSvc.getNamedSeed(seedName)
233			newSeeds = [x+nJob for x in curSeeds]
234			randSvc.setNamedSeed(seedName, *newSeeds)
235			preserveSeedList.append(seedName)
236
237			# Randomize remaining seeds
238			randSvc.populate(*preserveSeedList)
239			else:
240			print "Neither old nor new seed format found!"
241
242			# End version specific code
243
244			# Write out new config file in one format or the other, FUTURE: Get rid of cfg mode
245			outFile = open(outFileName,"w")
246			if outFileName.endswith('py'):
247			outFile.write("import FWCore.ParameterSet.Config as cms\n")
248	ewv	1.17	outFile.write("import pickle\n")
249			outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(cmsProcess))
250			outFile.write("process = pickle.loads(pickledCfg)\n")
251	ewv	1.12	if (debug):
252	ewv	1.17	print "writeCfg output (May not be exact):"
253	ewv	1.12	print "import FWCore.ParameterSet.Config as cms"
254			print cmsProcess.dumpPython()
255	ewv	1.9	else:
256	ewv	1.13	outFile.write(cfg.data.dumpConfig())
257	ewv	1.12	if (debug):
258			print "writeCfg output:"
259	ewv	1.13	print str(cfg.data.dumpConfig())
260	ewv	1.12	outFile.close()
261	ewv	1.1
262
263			if __name__ == '__main__' :
264			exit_status = main(sys.argv[1:])
265			sys.exit(exit_status)