CRAB/python/writeCfg.py

#!/usr/bin/env python

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

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'))
    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 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(str(cfg))
        if (debug):
            print "writeCfg output:"
            print str(cfg)
    outFile.close()


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