CRAB/python/writeCfg.py

#!/usr/bin/env python

import sys, getopt, string
import imp
import os
import random
from random import SystemRandom
_inst  = SystemRandom()

from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
from FWCore.ParameterSet.DictTypes import SortedKeysDict
from FWCore.ParameterSet.Modules   import Service
from FWCore.ParameterSet.Types     import *
from FWCore.ParameterSet.Config    import include

import FWCore.ParameterSet.Types   as CfgTypes
import FWCore.ParameterSet.Modules as CfgModules

class ConfigException(Exception):
    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
  firstRun       = 0
  sourceSeed     = 0
  vtxSeed        = 0
  g4Seed         = 0
  mixSeed        = 0
  debug          = False
  _MAXINT        = 900000000
  maxSeeds       = 4         # Kludge, maximum # of seeds that any engine takes

  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','')
  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:
        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 "Error: ",msg
    finally:
        handle.close()
  else:
    try:
      cfo = include(fileName)
      cmsProcess = cfo
    except Exception, ex:
      msg =  "The cfg file is not valid, %s\n" % str(ex)
      raise "Error: ",msg
  cfg = CfgInterface(cmsProcess)

  # Set parameters for job
  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)

  # 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)
    if ranModules != None:              # Old format present, no matter the CMSSW version
      print "Old-style random number seeds found, will be changed."
      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(_inst.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)

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