CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.19 2009/05/05 15:29:38 ewv Exp $"
__version__ = "$Revision: 1.19 $"

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.LHESource.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
    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.20
Committed:	Wed Jun 17 20:58:08 2009 UTC (15 years, 10 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_6_0_pre11
Changes since 1.19:	+14 -10 lines
Log Message:	Restore generators other than pythia, add lhe
#	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.20	__revision__ = "$Id: writeCfg.py,v 1.19 2009/05/05 15:29:38 ewv Exp $"
8			__version__ = "$Revision: 1.19 $"
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			cmsProcess.LHESource.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
162			elif firstEvent != -1: # (Old? Madgraph)
163	ewv	1.14	cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
164	ewv	1.20
165	ewv	1.12	if inputFiles:
166			inputFileNames = inputFiles.split(',')
167			inModule.setFileNames(*inputFileNames)
168
169			# handle parent files if needed
170			if parentFiles:
171			parentFileNames = parentFiles.split(',')
172			inModule.setSecondaryFileNames(*parentFileNames)
173
174			# Pythia parameters
175			if (firstRun):
176			inModule.setFirstRun(firstRun)
177
178			incrementSeedList = []
179			preserveSeedList = []
180
181			if incrementSeeds:
182			incrementSeedList = incrementSeeds.split(',')
183			if preserveSeeds:
184			preserveSeedList = preserveSeeds.split(',')
185
186			# FUTURE: This function tests the CMSSW version and presence of old-style seed specification.
187			# Reduce when we drop support for old versions
188			if cfg.data.services.has_key('RandomNumberGeneratorService'): # There are random #'s to deal with
189			print "RandomNumberGeneratorService found, will attempt to change seeds"
190			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
191			ranModules = getattr(ranGenerator, "moduleSeeds", None)
192			oldSource = getattr(ranGenerator, "sourceSeed", None)
193			if ranModules != None or oldSource != None: # Old format present, no matter the CMSSW version
194			print "Old-style random number seeds found, will be changed."
195			if oldSource != None:
196			sourceSeed = int(ranGenerator.sourceSeed.value())
197			if ('sourceSeed' in preserveSeedList) or ('theSource' in preserveSeedList):
198			pass
199			elif ('sourceSeed' in incrementSeedList) or ('theSource' in incrementSeedList):
200			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
201			else:
202			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1, _MAXINT)))
203
204			for seed in incrementSeedList:
205			curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
206			if curSeed:
207			curValue = int(curSeed.value())
208			setattr(ranGenerator.moduleSeeds, seed, CfgTypes.untracked(CfgTypes.uint32(curValue+nJob)))
209			preserveSeedList.append(seed)
210
211			if ranModules != None:
212			for seed in ranGenerator.moduleSeeds.parameterNames_():
213			if seed not in preserveSeedList:
214			curSeed = getattr(ranGenerator.moduleSeeds, seed, None)
215			if curSeed:
216			curValue = int(curSeed.value())
217			setattr(ranGenerator.moduleSeeds, seed,
218			CfgTypes.untracked(CfgTypes.uint32(MyRandom.randint(1,_MAXINT))))
219			elif CMSSW_major > 2 or (CMSSW_major == 2 and CMSSW_minor >= 1): # Treatment for seeds, CMSSW 2_1_x and later
220			print "New-style random number seeds found, will be changed."
221			from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
222			randSvc = RandomNumberServiceHelper(ranGenerator)
223
224			# Increment requested seed sets
225			for seedName in incrementSeedList:
226			curSeeds = randSvc.getNamedSeed(seedName)
227			newSeeds = [x+nJob for x in curSeeds]
228			randSvc.setNamedSeed(seedName, *newSeeds)
229			preserveSeedList.append(seedName)
230
231			# Randomize remaining seeds
232			randSvc.populate(*preserveSeedList)
233			else:
234			print "Neither old nor new seed format found!"
235
236			# End version specific code
237
238			# Write out new config file in one format or the other, FUTURE: Get rid of cfg mode
239			outFile = open(outFileName,"w")
240			if outFileName.endswith('py'):
241			outFile.write("import FWCore.ParameterSet.Config as cms\n")
242	ewv	1.17	outFile.write("import pickle\n")
243			outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(cmsProcess))
244			outFile.write("process = pickle.loads(pickledCfg)\n")
245	ewv	1.12	if (debug):
246	ewv	1.17	print "writeCfg output (May not be exact):"
247	ewv	1.12	print "import FWCore.ParameterSet.Config as cms"
248			print cmsProcess.dumpPython()
249	ewv	1.9	else:
250	ewv	1.13	outFile.write(cfg.data.dumpConfig())
251	ewv	1.12	if (debug):
252			print "writeCfg output:"
253	ewv	1.13	print str(cfg.data.dumpConfig())
254	ewv	1.12	outFile.close()
255	ewv	1.1
256
257			if __name__ == '__main__' :
258			exit_status = main(sys.argv[1:])
259			sys.exit(exit_status)