CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.18 2009/05/04 21:45:43 ewv Exp $"
__version__ = "$Revision: 1.18 $"

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 pycfg file
    - Modify job specific parameters based on environment variables
    - Write out modified cfg or 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("CompHEPFirstEvent"):
                compHEPFirstEvent = int(elem.getAttribute("CompHEPFirstEvent"))
            if elem.getAttribute("FirstRun"):
                firstRun = int(elem.getAttribute("FirstRun"))

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