CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.25 2009/08/28 18:56:43 ewv Exp $"
__version__ = "$Revision: 1.25 $"

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