CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.31 2011/03/02 10:48:46 spiga Exp $"
__version__ = "$Revision: 1.31 $"

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'))
            inputBlocks    = str(elem.getAttribute('InputBlocks'))
            inputFiles     = str(elem.getAttribute('InputFiles'))
            parentFiles    = str(elem.getAttribute('ParentFiles'))
            lumis          = str(elem.getAttribute('Lumis'))

    report(inputBlocks,inputFiles,parentFiles,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 and inModule.sourceType not in ['EmptySource']:
        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']

        ranModules   = getattr(ranGenerator, "moduleSeeds", None)
        oldSource    = getattr(ranGenerator, "sourceSeed",  None)
        if ranModules != None or oldSource != None:
            msg = "Your random number seeds are set in an old,\n"
            msg += "deprecated style. Please change to new style:\n"
            msg += "https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideEDMRandomNumberGeneratorService"
            raise ConfigException(msg)

        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
    pklFileName = outFileName + '.pkl'
    outFile = open(outFileName,"w")
    outFile.write("import FWCore.ParameterSet.Config as cms\n")
    outFile.write("import pickle\n")
    outFile.write("process = pickle.load(open('%s', 'rb'))\n" % pklFileName)
    outFile.close()

    pklFile = open(pklFileName,"wb")
    myPickle = pickle.Pickler(pklFile)
    myPickle.dump(cmsProcess)
    pklFile.close()

    if (debug):
        print "writeCfg output (May not be exact):"
        print "import FWCore.ParameterSet.Config as cms"
        print cmsProcess.dumpPython()


def report( inputBlocks='', inputFiles='', parentFiles='', lumis='' ):
    """
    Writes the 4 parameters to a file, one parameter per line.
    """
    outFile = open('inputsReport.txt',"a")

  #  InputFileList=inputFiles.split(',') 
  #  parentFilesList= parentFiles.split(',')
  #  lumisList= lumis.split(',')

    ## replacing , with ; otherwise report.py will split it as a new parameter
    txt = ''
    txt += 'inputBlocks='+inputBlocks.replace(',',';')+'\n'
    txt += 'inputFiles='+inputFiles.replace(',',';')+'\n'
    txt += 'parentFiles='+parentFiles.replace(',',';')+'\n'
    txt += 'lumisRange='+lumis.replace(',',';')+'\n'
    outFile.write(str(txt))
    outFile.close()
    return


if __name__ == '__main__' :
    exit_status = main(sys.argv[1:])
    sys.exit(exit_status)
Revision:	1.32
Committed:	Wed Mar 23 10:30:23 2011 UTC (14 years, 1 month ago) by spiga
Content type:	text/x-python
Branch:	MAIN
Changes since 1.31:	+2 -3 lines
Log Message:	more info to dashboard
#	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	spiga	1.32	__revision__ = "$Id: writeCfg.py,v 1.31 2011/03/02 10:48:46 spiga Exp $"
8			__version__ = "$Revision: 1.31 $"
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	import FWCore.ParameterSet.Types as CfgTypes
21	ewv	1.1
22	ewv	1.12	MyRandom = SystemRandom()
23	ewv	1.1
24	ewv	1.3	class ConfigException(Exception):
25	ewv	1.12	"""
26			Exceptions raised by writeCfg
27			"""
28
29	ewv	1.3	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	ewv	1.1
37			def main(argv) :
38	ewv	1.12	"""
39			writeCfg
40
41	ewv	1.25	- Read in existing, user supplied pycfg or pickled pycfg file
42	ewv	1.20	- Modify job specific parameters based on environment variables and arguments.xml
43	ewv	1.25	- Write out pickled pycfg file
44	ewv	1.12
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	ewv	1.1
65			try:
66	ewv	1.12	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	ewv	1.2
83	ewv	1.12	# 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	ewv	1.19	# Read in Environment, XML and get optional Parameters
92	ewv	1.12
93			nJob = int(os.environ.get('NJob', '0'))
94	ewv	1.19	preserveSeeds = os.environ.get('PreserveSeeds','')
95			incrementSeeds = os.environ.get('IncrementSeeds','')
96	ewv	1.12
97	ewv	1.18	# Defaults
98
99			maxEvents = 0
100			skipEvents = 0
101			firstEvent = -1
102			compHEPFirstEvent = 0
103			firstRun = 0
104	ewv	1.26	# FUTURE: Remove firstRun
105			firstLumi = 0
106	ewv	1.18
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	ewv	1.26	if elem.getAttribute("FirstLumi"):
120			firstLumi = int(elem.getAttribute("FirstLumi"))
121	ewv	1.18
122	ewv	1.20	generator = str(elem.getAttribute('Generator'))
123	spiga	1.31	inputBlocks = str(elem.getAttribute('InputBlocks'))
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	spiga	1.31	report(inputBlocks,inputFiles,parentFiles,lumis)
129
130	ewv	1.24	# Read Input python config file
131	ewv	1.12
132	ewv	1.24	handle = open(fileName, 'r')
133			try: # Nested form for Python < 2.5
134	ewv	1.12	try:
135	ewv	1.24	print "Importing .py file"
136			cfo = imp.load_source("pycfg", fileName, handle)
137			cmsProcess = cfo.process
138	ewv	1.12	except Exception, ex:
139	ewv	1.24	msg = "Your pycfg file is not valid python: %s" % str(ex)
140	ewv	1.12	raise ConfigException(msg)
141	ewv	1.24	finally:
142			handle.close()
143	ewv	1.12
144			cfg = CfgInterface(cmsProcess)
145
146			# Set parameters for job
147			print "Setting parameters"
148			inModule = cfg.inputSource
149			if maxEvents:
150			cfg.maxEvents.setMaxEventsInput(maxEvents)
151
152	ewv	1.30	if skipEvents and inModule.sourceType not in ['EmptySource']:
153	ewv	1.12	inModule.setSkipEvents(skipEvents)
154	ewv	1.20
155			# Set "skip events" for various generators
156			if generator == 'comphep':
157			cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(firstEvent)
158			elif generator == 'lhe':
159	ewv	1.21	cmsProcess.source.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
160			cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent+1))
161	ewv	1.20	elif firstEvent != -1: # (Old? Madgraph)
162	ewv	1.14	cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
163	ewv	1.20
164	ewv	1.12	if inputFiles:
165			inputFileNames = inputFiles.split(',')
166			inModule.setFileNames(*inputFileNames)
167
168			# handle parent files if needed
169			if parentFiles:
170			parentFileNames = parentFiles.split(',')
171			inModule.setSecondaryFileNames(*parentFileNames)
172
173	ewv	1.22	if lumis:
174	ewv	1.23	if CMSSW_major < 3: # FUTURE: Can remove this check
175			print "Cannot skip lumis for CMSSW 2_x"
176			else:
177			lumiRanges = lumis.split(',')
178			inModule.setLumisToProcess(*lumiRanges)
179	ewv	1.22
180	ewv	1.12	# Pythia parameters
181			if (firstRun):
182			inModule.setFirstRun(firstRun)
183	ewv	1.26	if (firstLumi):
184			inModule.setFirstLumi(firstLumi)
185	ewv	1.12
186	ewv	1.24	# Check if there are random #'s to deal with
187			if cfg.data.services.has_key('RandomNumberGeneratorService'):
188	ewv	1.12	print "RandomNumberGeneratorService found, will attempt to change seeds"
189	ewv	1.24	from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
190	ewv	1.12	ranGenerator = cfg.data.services['RandomNumberGeneratorService']
191	ewv	1.28
192			ranModules = getattr(ranGenerator, "moduleSeeds", None)
193			oldSource = getattr(ranGenerator, "sourceSeed", None)
194			if ranModules != None or oldSource != None:
195			msg = "Your random number seeds are set in an old,\n"
196			msg += "deprecated style. Please change to new style:\n"
197			msg += "https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideEDMRandomNumberGeneratorService"
198			raise ConfigException(msg)
199
200	ewv	1.24	randSvc = RandomNumberServiceHelper(ranGenerator)
201
202			incrementSeedList = []
203			preserveSeedList = []
204	ewv	1.12
205	ewv	1.24	if incrementSeeds:
206			incrementSeedList = incrementSeeds.split(',')
207			if preserveSeeds:
208			preserveSeedList = preserveSeeds.split(',')
209
210			# Increment requested seed sets
211			for seedName in incrementSeedList:
212			curSeeds = randSvc.getNamedSeed(seedName)
213			newSeeds = [x+nJob for x in curSeeds]
214			randSvc.setNamedSeed(seedName, *newSeeds)
215			preserveSeedList.append(seedName)
216	ewv	1.12
217	ewv	1.24	# Randomize remaining seeds
218			randSvc.populate(*preserveSeedList)
219	ewv	1.12
220	ewv	1.24	# Write out new config file
221	ewv	1.27	pklFileName = outFileName + '.pkl'
222	ewv	1.12	outFile = open(outFileName,"w")
223	ewv	1.24	outFile.write("import FWCore.ParameterSet.Config as cms\n")
224			outFile.write("import pickle\n")
225	ewv	1.27	outFile.write("process = pickle.load(open('%s', 'rb'))\n" % pklFileName)
226	ewv	1.12	outFile.close()
227	ewv	1.27
228			pklFile = open(pklFileName,"wb")
229			myPickle = pickle.Pickler(pklFile)
230			myPickle.dump(cmsProcess)
231			pklFile.close()
232
233	ewv	1.24	if (debug):
234			print "writeCfg output (May not be exact):"
235			print "import FWCore.ParameterSet.Config as cms"
236			print cmsProcess.dumpPython()
237	ewv	1.1
238
239	spiga	1.31	def report( inputBlocks='', inputFiles='', parentFiles='', lumis='' ):
240			"""
241			Writes the 4 parameters to a file, one parameter per line.
242			"""
243			outFile = open('inputsReport.txt',"a")
244
245			# InputFileList=inputFiles.split(',')
246			# parentFilesList= parentFiles.split(',')
247			# lumisList= lumis.split(',')
248
249			## replacing , with ; otherwise report.py will split it as a new parameter
250			txt = ''
251			txt += 'inputBlocks='+inputBlocks.replace(',',';')+'\n'
252			txt += 'inputFiles='+inputFiles.replace(',',';')+'\n'
253			txt += 'parentFiles='+parentFiles.replace(',',';')+'\n'
254			txt += 'lumisRange='+lumis.replace(',',';')+'\n'
255			outFile.write(str(txt))
256			outFile.close()
257			return
258
259
260	ewv	1.1	if __name__ == '__main__' :
261			exit_status = main(sys.argv[1:])
262			sys.exit(exit_status)