filefi/027/wCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: wCfg.py,v 1.1 2012/03/30 17:46:35 paus Exp $"
__version__ = "$Revision: 1.1 $"

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
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 pycfg or pickled pycfg file
    - Modify job specific parameters based on environment variables and arguments.xml
    - Write out pickled 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
    # FUTURE: Remove firstRun
    firstLumi  = 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("FirstRun"):
                firstRun = int(elem.getAttribute("FirstRun"))
            if elem.getAttribute("FirstLumi"):
                firstLumi = int(elem.getAttribute("FirstLumi"))

            generator      = str(elem.getAttribute('Generator'))
            inputFiles     = str(elem.getAttribute('InputFiles'))
            parentFiles    = str(elem.getAttribute('ParentFiles'))
            lumis          = str(elem.getAttribute('Lumis'))

  # Read Input python config file

    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()

    cfg = CfgInterface(cmsProcess)

    # Set parameters for job
    print "Setting parameters"
    inModule = cfg.inputSource
    if maxEvents:
        cfg.maxEvents.setMaxEventsInput(maxEvents)

    if skipEvents:
        inModule.setSkipEvents(skipEvents)

    # Set "skip events" for various generators
    if generator == 'comphep':
        cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(firstEvent)
    elif generator == 'lhe':
        cmsProcess.source.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
        cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent+1))
    elif firstEvent != -1: # (Old? Madgraph)
        cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))

    if inputFiles:
        inputFileNames = inputFiles.split(',')
        inModule.setFileNames(*inputFileNames)

    # handle parent files if needed
    if parentFiles:
        parentFileNames = parentFiles.split(',')
        inModule.setSecondaryFileNames(*parentFileNames)

    if lumis:
        if CMSSW_major < 3: # FUTURE: Can remove this check
            print "Cannot skip lumis for CMSSW 2_x"
        else:
            lumiRanges = lumis.split(',')
            inModule.setLumisToProcess(*lumiRanges)

    # Pythia parameters
    if (firstRun):
        inModule.setFirstRun(firstRun)
    if (firstLumi):
        inModule.setFirstLumi(firstLumi)

    # Check if there are random #'s to deal with
    if cfg.data.services.has_key('RandomNumberGeneratorService'):
        print "RandomNumberGeneratorService found, will attempt to change seeds"
        from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
        ranGenerator = cfg.data.services['RandomNumberGeneratorService']
        randSvc = RandomNumberServiceHelper(ranGenerator)

        incrementSeedList = []
        preserveSeedList  = []

        if incrementSeeds:
            incrementSeedList = incrementSeeds.split(',')
        if preserveSeeds:
            preserveSeedList  = preserveSeeds.split(',')

        # 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)

    # Write out new config file
    outFile = open(outFileName,"w")
    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")
    outFile.close()
    if (debug):
        print "writeCfg output (May not be exact):"
        print "import FWCore.ParameterSet.Config as cms"
        print cmsProcess.dumpPython()


if __name__ == '__main__' :
    exit_status = main(sys.argv[1:])
    sys.exit(exit_status)
Revision:	1.1
Committed:	Sat May 5 16:53:24 2012 UTC (13 years ago) by paus
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	Mit_032, Mit_031, Mit_030, Mit_029c, Mit_029b, Mit_030_pre1, Mit_029a, Mit_029, Mit_029_pre1, Mit_028a, Mit_028, Mit_027a, Mit_027, HEAD
Log Message:	Version 027 - complete version for ICHEP 2012.
#	User	Rev	Content
1	paus	1.1	#!/usr/bin/env python
2
3			"""
4			Re-write config file and optionally convert to python
5			"""
6
7			__revision__ = "$Id: wCfg.py,v 1.1 2012/03/30 17:46:35 paus Exp $"
8			__version__ = "$Revision: 1.1 $"
9
10			import getopt
11			import imp
12			import os
13			import pickle
14			import sys
15			import xml.dom.minidom
16
17			from random import SystemRandom
18
19			from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
20			import FWCore.ParameterSet.Types as CfgTypes
21
22			MyRandom = SystemRandom()
23
24			class ConfigException(Exception):
25			"""
26			Exceptions raised by writeCfg
27			"""
28
29			def __init__(self, msg):
30			Exception.__init__(self, msg)
31			self._msg = msg
32			return
33
34			def __str__(self):
35			return self._msg
36
37			def main(argv) :
38			"""
39			writeCfg
40
41			- Read in existing, user supplied pycfg or pickled pycfg file
42			- Modify job specific parameters based on environment variables and arguments.xml
43			- Write out pickled pycfg file
44
45			required parameters: none
46
47			optional parameters:
48			--help : help
49			--debug : debug statements
50
51			"""
52
53			# defaults
54			inputFileNames = None
55			parentFileNames = None
56			debug = False
57			_MAXINT = 900000000
58
59			try:
60			opts, args = getopt.getopt(argv, "", ["debug", "help"])
61			except getopt.GetoptError:
62			print main.__doc__
63			sys.exit(2)
64
65			try:
66			CMSSW = os.environ['CMSSW_VERSION']
67			parts = CMSSW.split('_')
68			CMSSW_major = int(parts[1])
69			CMSSW_minor = int(parts[2])
70			CMSSW_patch = int(parts[3])
71			except (KeyError, ValueError):
72			msg = "Your environment doesn't specify the CMSSW version or specifies it incorrectly"
73			raise ConfigException(msg)
74
75			# Parse command line options
76			for opt, arg in opts :
77			if opt == "--help" :
78			print main.__doc__
79			sys.exit()
80			elif opt == "--debug" :
81			debug = True
82
83			# Parse remaining parameters
84			try:
85			fileName = args[0]
86			outFileName = args[1]
87			except IndexError:
88			print main.__doc__
89			sys.exit()
90
91			# Read in Environment, XML and get optional Parameters
92
93			nJob = int(os.environ.get('NJob', '0'))
94			preserveSeeds = os.environ.get('PreserveSeeds','')
95			incrementSeeds = os.environ.get('IncrementSeeds','')
96
97			# Defaults
98
99			maxEvents = 0
100			skipEvents = 0
101			firstEvent = -1
102			compHEPFirstEvent = 0
103			firstRun = 0
104			# FUTURE: Remove firstRun
105			firstLumi = 0
106
107			dom = xml.dom.minidom.parse(os.environ['RUNTIME_AREA']+'/arguments.xml')
108
109			for elem in dom.getElementsByTagName("Job"):
110			if nJob == int(elem.getAttribute("JobID")):
111			if elem.getAttribute("MaxEvents"):
112			maxEvents = int(elem.getAttribute("MaxEvents"))
113			if elem.getAttribute("SkipEvents"):
114			skipEvents = int(elem.getAttribute("SkipEvents"))
115			if elem.getAttribute("FirstEvent"):
116			firstEvent = int(elem.getAttribute("FirstEvent"))
117			if elem.getAttribute("FirstRun"):
118			firstRun = int(elem.getAttribute("FirstRun"))
119			if elem.getAttribute("FirstLumi"):
120			firstLumi = int(elem.getAttribute("FirstLumi"))
121
122			generator = str(elem.getAttribute('Generator'))
123			inputFiles = str(elem.getAttribute('InputFiles'))
124			parentFiles = str(elem.getAttribute('ParentFiles'))
125			lumis = str(elem.getAttribute('Lumis'))
126
127			# Read Input python config file
128
129			handle = open(fileName, 'r')
130			try: # Nested form for Python < 2.5
131			try:
132			print "Importing .py file"
133			cfo = imp.load_source("pycfg", fileName, handle)
134			cmsProcess = cfo.process
135			except Exception, ex:
136			msg = "Your pycfg file is not valid python: %s" % str(ex)
137			raise ConfigException(msg)
138			finally:
139			handle.close()
140
141			cfg = CfgInterface(cmsProcess)
142
143			# Set parameters for job
144			print "Setting parameters"
145			inModule = cfg.inputSource
146			if maxEvents:
147			cfg.maxEvents.setMaxEventsInput(maxEvents)
148
149			if skipEvents:
150			inModule.setSkipEvents(skipEvents)
151
152			# Set "skip events" for various generators
153			if generator == 'comphep':
154			cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(firstEvent)
155			elif generator == 'lhe':
156			cmsProcess.source.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
157			cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent+1))
158			elif firstEvent != -1: # (Old? Madgraph)
159			cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
160
161			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			if lumis:
171			if CMSSW_major < 3: # FUTURE: Can remove this check
172			print "Cannot skip lumis for CMSSW 2_x"
173			else:
174			lumiRanges = lumis.split(',')
175			inModule.setLumisToProcess(*lumiRanges)
176
177			# Pythia parameters
178			if (firstRun):
179			inModule.setFirstRun(firstRun)
180			if (firstLumi):
181			inModule.setFirstLumi(firstLumi)
182
183			# Check if there are random #'s to deal with
184			if cfg.data.services.has_key('RandomNumberGeneratorService'):
185			print "RandomNumberGeneratorService found, will attempt to change seeds"
186			from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
187			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
188			randSvc = RandomNumberServiceHelper(ranGenerator)
189
190			incrementSeedList = []
191			preserveSeedList = []
192
193			if incrementSeeds:
194			incrementSeedList = incrementSeeds.split(',')
195			if preserveSeeds:
196			preserveSeedList = preserveSeeds.split(',')
197
198			# Increment requested seed sets
199			for seedName in incrementSeedList:
200			curSeeds = randSvc.getNamedSeed(seedName)
201			newSeeds = [x+nJob for x in curSeeds]
202			randSvc.setNamedSeed(seedName, *newSeeds)
203			preserveSeedList.append(seedName)
204
205			# Randomize remaining seeds
206			randSvc.populate(*preserveSeedList)
207
208			# Write out new config file
209			outFile = open(outFileName,"w")
210			outFile.write("import FWCore.ParameterSet.Config as cms\n")
211			outFile.write("import pickle\n")
212			outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(cmsProcess))
213			outFile.write("process = pickle.loads(pickledCfg)\n")
214			outFile.close()
215			if (debug):
216			print "writeCfg output (May not be exact):"
217			print "import FWCore.ParameterSet.Config as cms"
218			print cmsProcess.dumpPython()
219
220
221			if __name__ == '__main__' :
222			exit_status = main(sys.argv[1:])
223			sys.exit(exit_status)