CRAB/python/writeCfg.py

#!/usr/bin/env python

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

__revision__ = "$Id: writeCfg.py,v 1.24 2009/08/27 12:42:23 ewv Exp $"
__version__ = "$Revision: 1.24 $"

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

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

            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)

    # 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
Committed:	Fri Aug 28 18:56:43 2009 UTC (15 years, 8 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_7_0, CRAB_2_7_0_pre8, CRAB_2_7_0_pre7, CRAB_2_7_0_pre6, CRAB_2_7_0_pre5, CRAB_2_7_0_pre4, CRAB_2_7_0_pre3, CRAB_2_7_0_pre2, CRAB_2_7_0_pre1
Branch point for:	Lumi2_8
Changes since 1.24:	+4 -4 lines
Log Message:	Update some docs and comments now that CMSSW.cfg support is gone
#	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	__revision__ = "$Id: writeCfg.py,v 1.24 2009/08/27 12:42:23 ewv Exp $"
8			__version__ = "$Revision: 1.24 $"
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
106			dom = xml.dom.minidom.parse(os.environ['RUNTIME_AREA']+'/arguments.xml')
107
108			for elem in dom.getElementsByTagName("Job"):
109			if nJob == int(elem.getAttribute("JobID")):
110			if elem.getAttribute("MaxEvents"):
111			maxEvents = int(elem.getAttribute("MaxEvents"))
112			if elem.getAttribute("SkipEvents"):
113			skipEvents = int(elem.getAttribute("SkipEvents"))
114			if elem.getAttribute("FirstEvent"):
115			firstEvent = int(elem.getAttribute("FirstEvent"))
116			if elem.getAttribute("FirstRun"):
117			firstRun = int(elem.getAttribute("FirstRun"))
118
119	ewv	1.20	generator = str(elem.getAttribute('Generator'))
120	ewv	1.18	inputFiles = str(elem.getAttribute('InputFiles'))
121			parentFiles = str(elem.getAttribute('ParentFiles'))
122	ewv	1.22	lumis = str(elem.getAttribute('Lumis'))
123	ewv	1.12
124	ewv	1.24	# Read Input python config file
125	ewv	1.12
126	ewv	1.24	handle = open(fileName, 'r')
127			try: # Nested form for Python < 2.5
128	ewv	1.12	try:
129	ewv	1.24	print "Importing .py file"
130			cfo = imp.load_source("pycfg", fileName, handle)
131			cmsProcess = cfo.process
132	ewv	1.12	except Exception, ex:
133	ewv	1.24	msg = "Your pycfg file is not valid python: %s" % str(ex)
134	ewv	1.12	raise ConfigException(msg)
135	ewv	1.24	finally:
136			handle.close()
137	ewv	1.12
138			cfg = CfgInterface(cmsProcess)
139
140			# Set parameters for job
141			print "Setting parameters"
142			inModule = cfg.inputSource
143			if maxEvents:
144			cfg.maxEvents.setMaxEventsInput(maxEvents)
145
146			if skipEvents:
147			inModule.setSkipEvents(skipEvents)
148	ewv	1.20
149			# Set "skip events" for various generators
150			if generator == 'comphep':
151			cmsProcess.source.CompHEPFirstEvent = CfgTypes.int32(firstEvent)
152			elif generator == 'lhe':
153	ewv	1.21	cmsProcess.source.skipEvents = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
154			cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent+1))
155	ewv	1.20	elif firstEvent != -1: # (Old? Madgraph)
156	ewv	1.14	cmsProcess.source.firstEvent = CfgTypes.untracked(CfgTypes.uint32(firstEvent))
157	ewv	1.20
158	ewv	1.12	if inputFiles:
159			inputFileNames = inputFiles.split(',')
160			inModule.setFileNames(*inputFileNames)
161
162			# handle parent files if needed
163			if parentFiles:
164			parentFileNames = parentFiles.split(',')
165			inModule.setSecondaryFileNames(*parentFileNames)
166
167	ewv	1.22	if lumis:
168	ewv	1.23	if CMSSW_major < 3: # FUTURE: Can remove this check
169			print "Cannot skip lumis for CMSSW 2_x"
170			else:
171			lumiRanges = lumis.split(',')
172			inModule.setLumisToProcess(*lumiRanges)
173	ewv	1.22
174	ewv	1.12	# Pythia parameters
175			if (firstRun):
176			inModule.setFirstRun(firstRun)
177
178	ewv	1.24	# Check if there are random #'s to deal with
179			if cfg.data.services.has_key('RandomNumberGeneratorService'):
180	ewv	1.12	print "RandomNumberGeneratorService found, will attempt to change seeds"
181	ewv	1.24	from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
182	ewv	1.12	ranGenerator = cfg.data.services['RandomNumberGeneratorService']
183	ewv	1.24	randSvc = RandomNumberServiceHelper(ranGenerator)
184
185			incrementSeedList = []
186			preserveSeedList = []
187	ewv	1.12
188	ewv	1.24	if incrementSeeds:
189			incrementSeedList = incrementSeeds.split(',')
190			if preserveSeeds:
191			preserveSeedList = preserveSeeds.split(',')
192
193			# Increment requested seed sets
194			for seedName in incrementSeedList:
195			curSeeds = randSvc.getNamedSeed(seedName)
196			newSeeds = [x+nJob for x in curSeeds]
197			randSvc.setNamedSeed(seedName, *newSeeds)
198			preserveSeedList.append(seedName)
199	ewv	1.12
200	ewv	1.24	# Randomize remaining seeds
201			randSvc.populate(*preserveSeedList)
202	ewv	1.12
203	ewv	1.24	# Write out new config file
204	ewv	1.12	outFile = open(outFileName,"w")
205	ewv	1.24	outFile.write("import FWCore.ParameterSet.Config as cms\n")
206			outFile.write("import pickle\n")
207			outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(cmsProcess))
208			outFile.write("process = pickle.loads(pickledCfg)\n")
209	ewv	1.12	outFile.close()
210	ewv	1.24	if (debug):
211			print "writeCfg output (May not be exact):"
212			print "import FWCore.ParameterSet.Config as cms"
213			print cmsProcess.dumpPython()
214	ewv	1.1
215
216			if __name__ == '__main__' :
217			exit_status = main(sys.argv[1:])
218			sys.exit(exit_status)