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. Can be simplified as we drop support for old versions
  if CMSSW_major < 2 or (CMSSW_major == 2 and CMSSW_minor == 0): # Treatment for seeds, CMSSW < 2_1_x
    if cfg.data.services.has_key('RandomNumberGeneratorService'):
      ranGenerator = cfg.data.services['RandomNumberGeneratorService']
      ranModules   = ranGenerator.moduleSeeds

      sourceSeed = int(ranGenerator.sourceSeed.value())
      if 'sourceSeed' in preserveSeedList:
        pass
      elif 'sourceSeed' 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))))
  else: # Treatment for  seeds, CMSSW 2_1_x and later
    if cfg.data.services.has_key('RandomNumberGeneratorService'):
      from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper

      ranGenerator = cfg.data.services['RandomNumberGeneratorService']
      randSvc = RandomNumberServiceHelper(ranGenerator)
      # sourceSeed is different from the rest, John says will become "theSource"
      if 'sourceSeed' in preserveSeedList:
        pass
      elif 'sourceSeed' in incrementSeedList:
        sourceSeed = int(ranGenerator.sourceSeed.value())
        ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
      else:
        ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT)))

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

  # 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.7
Committed:	Tue May 6 01:09:58 2008 UTC (16 years, 11 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_2_0_pre18
Changes since 1.6:	+34 -32 lines
Log Message:	Support 2_1_x seeds
#	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.6	# FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
148	ewv	1.7	if CMSSW_major < 2 or (CMSSW_major == 2 and CMSSW_minor == 0): # Treatment for seeds, CMSSW < 2_1_x
149	ewv	1.3	if cfg.data.services.has_key('RandomNumberGeneratorService'):
150			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
151	ewv	1.2	ranModules = ranGenerator.moduleSeeds
152
153	ewv	1.3	sourceSeed = int(ranGenerator.sourceSeed.value())
154			if 'sourceSeed' in preserveSeedList:
155			pass
156			elif 'sourceSeed' in incrementSeedList:
157			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
158			else:
159			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT)))
160
161			for seed in incrementSeedList:
162			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
163			if curSeed:
164			curValue = int(curSeed.value())
165			setattr(ranGenerator.moduleSeeds,seed,CfgTypes.untracked(CfgTypes.uint32(curValue+nJob)))
166			preserveSeedList.append(seed)
167
168	ewv	1.5	for seed in ranGenerator.moduleSeeds.parameterNames_():
169	ewv	1.3	if seed not in preserveSeedList:
170			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
171			if curSeed:
172			curValue = int(curSeed.value())
173			setattr(ranGenerator.moduleSeeds,seed,CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT))))
174	ewv	1.7	else: # Treatment for seeds, CMSSW 2_1_x and later
175			if cfg.data.services.has_key('RandomNumberGeneratorService'):
176			from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
177	ewv	1.3
178	ewv	1.7	ranGenerator = cfg.data.services['RandomNumberGeneratorService']
179			randSvc = RandomNumberServiceHelper(ranGenerator)
180			# sourceSeed is different from the rest, John says will become "theSource"
181			if 'sourceSeed' in preserveSeedList:
182			pass
183			elif 'sourceSeed' in incrementSeedList:
184			sourceSeed = int(ranGenerator.sourceSeed.value())
185			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
186			else:
187			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT)))
188	ewv	1.3
189	ewv	1.7	# Increment requested seed sets
190			for seedName in incrementSeedList:
191			curSeeds = randSvc.getNamedSeed(seedName)
192			newSeeds = [x+nJob for x in curSeeds]
193			randSvc.setNamedSeed(seedName,*newSeeds)
194			preserveSeedList.append(seedName)
195	ewv	1.2
196	ewv	1.7	# Randomize remaining seeds
197			randSvc.populate(*preserveSeedList)
198	ewv	1.2
199	ewv	1.7	# End version specific code
200	ewv	1.2
201	ewv	1.7	# Write out new config file in one format or the other, FUTURE: Get rid of cfg mode
202	ewv	1.1	outFile = open(outFileName,"w")
203	ewv	1.7	if outFileName.endswith('py'):
204	ewv	1.3	outFile.write("import FWCore.ParameterSet.Config as cms\n")
205	ewv	1.6	outFile.write(cmsProcess.dumpPython())
206	ewv	1.3	if (debug):
207			print "writeCfg output:"
208			print "import FWCore.ParameterSet.Config as cms"
209	ewv	1.6	print cmsProcess.dumpPython()
210	ewv	1.3	else:
211			outFile.write(str(cfg))
212			if (debug):
213			print "writeCfg output:"
214			print str(cfg)
215	ewv	1.1	outFile.close()
216
217
218			if __name__ == '__main__' :
219			exit_status = main(sys.argv[1:])
220			sys.exit(exit_status)