CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.16 2009/03/05 16:49:56 spiga Exp $"
__version__ = "$Revision: 1.16 $"

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