CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.14 2008/12/03 17:51:04 ewv Exp $"
__version__ = "$Revision: 1.14 $"

import sys, getopt
import imp
import os
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()

  # Optional Parameters

    maxEvents  = int(os.environ.get('MaxEvents', '0'))
    skipEvents = int(os.environ.get('SkipEvents','0'))
    firstEvent = int(os.environ.get('FirstEvent','0'))
    compHEPFirstEvent = int(os.environ.get('CompHEPFirstEvent','0'))
    firstRun   = int(os.environ.get('FirstRun',  '0'))
    nJob       = int(os.environ.get('NJob',      '0'))

    inputFiles     = os.environ.get('InputFiles','')
    parentFiles    = os.environ.get('ParentFiles','')
    preserveSeeds  = os.environ.get('PreserveSeeds','')
    incrementSeeds = os.environ.get('IncrementSeeds','')

  # 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:
        cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
    if compHEPFirstEvent:
        cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(compHEPFirstEvent)
    if inputFiles:
        inputFiles = inputFiles.replace('\\','')
        inputFiles = inputFiles.replace('"','')
        inputFileNames = inputFiles.split(',')
        inModule.setFileNames(*inputFileNames)

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