CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.15 2008/12/08 21:57:20 ewv Exp $"
__version__ = "$Revision: 1.15 $"

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','-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(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.16
Committed:	Thu Mar 5 16:49:56 2009 UTC (16 years, 1 month ago) by spiga
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_5_1, CRAB_2_5_1_pre4, CRAB_2_5_1_pre3, CRAB_2_5_1_pre2, CRAB_2_5_1_pre1, CRAB_2_5_0, CRAB_2_5_0_pre7, CRAB_2_5_0_pre6
Changes since 1.15:	+4 -4 lines
Log Message:	allow to set firstEvent = 0
#	Content
1	#!/usr/bin/env python
2
3	"""
4	Re-write config file and optionally convert to python
5	"""
6
7	__revision__ = "$Id: writeCfg.py,v 1.15 2008/12/08 21:57:20 ewv Exp $"
8	__version__ = "$Revision: 1.15 $"
9
10	import sys, getopt
11	import imp
12	import os
13	from random import SystemRandom
14
15	from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
16	from FWCore.ParameterSet.Config import include
17	import FWCore.ParameterSet.Types as CfgTypes
18
19	MyRandom = SystemRandom()
20
21	class ConfigException(Exception):
22	"""
23	Exceptions raised by writeCfg
24	"""
25
26	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
34	def main(argv) :
35	"""
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
62	try:
63	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
80	# 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	firstEvent = int(os.environ.get('FirstEvent','-1'))
93	compHEPFirstEvent = int(os.environ.get('CompHEPFirstEvent','0'))
94	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	if firstEvent != -1:
136	cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
137	if compHEPFirstEvent:
138	cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(compHEPFirstEvent)
139	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	else:
226	outFile.write(cfg.data.dumpConfig())
227	if (debug):
228	print "writeCfg output:"
229	print str(cfg.data.dumpConfig())
230	outFile.close()
231
232
233	if __name__ == '__main__' :
234	exit_status = main(sys.argv[1:])
235	sys.exit(exit_status)