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
  parentFileNames = 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','')
  parentFiles     = os.environ.get('ParentFiles','')
  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:
      print "Importing .cfg file"
      cfo = include(fileName)
      cmsProcess = cfo
    except Exception, ex:
      msg =  "The cfg file is not valid, %s\n" % str(ex)
      raise "Error: ",msg
  print "Getting interface"
  cfg = CfgInterface(cmsProcess)

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

  # handle parent files if needed
  if parentFiles:
    parentFiles = parentFiles.replace('\\','')
    parentFiles = parentFiles.replace('"','')
    parentFileNames = parentFiles.split(',')
    inModule.setSecondaryFileNames(*parentFileNames)
    
  # 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.11
Committed:	Thu Jun 5 07:13:31 2008 UTC (16 years, 10 months ago) by spiga
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_3_2_pre1, CRAB_2_4_0_test, CRAB_2_3_1, CRAB_2_3_1_pre6, CRAB_2_3_1_pre5, CRAB_2_3_1_pre4, CRAB_2_3_1_pre3, CRAB_2_3_1_pre2, CRAB_2_3_1_pre1, CRAB_2_3_0, CRAB_2_3_0_pre6, CRAB_2_3_0_pre1, CRAB_2_2_2_pre5, CRAB_2_2_2_pre4, CRAB_2_2_2_pre3
Branch point for:	AnaDataSet, CRAB_2_3_0_br
Changes since 1.10:	+18 -10 lines
Log Message:	fully added support for two tier inputs
#	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	spiga	1.11	inputFileNames = None
46			parentFileNames = None
47			firstRun = 0
48			sourceSeed = 0
49			vtxSeed = 0
50			g4Seed = 0
51			mixSeed = 0
52			debug = False
53			_MAXINT = 900000000
54			maxSeeds = 4 # Kludge, maximum # of seeds that any engine takes
55	ewv	1.1
56			try:
57	ewv	1.3	opts, args = getopt.getopt(argv, "", ["debug", "help"])
58	ewv	1.1	except getopt.GetoptError:
59			print main.__doc__
60			sys.exit(2)
61
62	ewv	1.3	try:
63			CMSSW = os.environ['CMSSW_VERSION']
64			parts = CMSSW.split('_')
65			CMSSW_major = int(parts[1])
66			CMSSW_minor = int(parts[2])
67			CMSSW_patch = int(parts[3])
68			except KeyError, ValueError:
69			msg = "Your environment doesn't specify the CMSSW version or specifies it incorrectly"
70			raise ConfigException(msg)
71
72			# Parse command line options
73	ewv	1.1	for opt, arg in opts :
74			if opt == "--help" :
75			print main.__doc__
76			sys.exit()
77			elif opt == "--debug" :
78			debug = True
79			# Parse remaining parameters
80
81	ewv	1.3	try:
82			fileName = args[0];
83			outFileName = args[1];
84			except IndexError:
85	ewv	1.10	print main.__doc__
86			sys.exit()
87	ewv	1.3
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	spiga	1.11	parentFiles = os.environ.get('ParentFiles','')
97	ewv	1.7	preserveSeeds = os.environ.get('PreserveSeeds','')
98	ewv	1.3	incrementSeeds = os.environ.get('IncrementSeeds','')
99	ewv	1.1
100	ewv	1.7	# Read Input cfg or python cfg file, FUTURE: Get rid cfg mode
101	ewv	1.1
102	ewv	1.6	if fileName.endswith('py'):
103	ewv	1.1	handle = open(fileName, 'r')
104			try: # Nested form for Python < 2.5
105			try:
106			cfo = imp.load_source("pycfg", fileName, handle)
107	ewv	1.6	cmsProcess = cfo.process
108	ewv	1.1	except Exception, ex:
109			msg = "Your pycfg file is not valid python: %s" % str(ex)
110			raise "Error: ",msg
111			finally:
112			handle.close()
113			else:
114			try:
115	ewv	1.10	print "Importing .cfg file"
116	ewv	1.1	cfo = include(fileName)
117	ewv	1.6	cmsProcess = cfo
118	ewv	1.1	except Exception, ex:
119			msg = "The cfg file is not valid, %s\n" % str(ex)
120			raise "Error: ",msg
121	ewv	1.10	print "Getting interface"
122	ewv	1.6	cfg = CfgInterface(cmsProcess)
123	ewv	1.1
124			# Set parameters for job
125	ewv	1.10	print "Setting parameters"
126	ewv	1.6	inModule = cfg.inputSource
127	ewv	1.3	if maxEvents:
128			cfg.maxEvents.setMaxEventsInput(maxEvents)
129	ewv	1.1
130	ewv	1.3	if skipEvents:
131			inModule.setSkipEvents(skipEvents)
132	ewv	1.1
133	ewv	1.3	if inputFiles:
134			inputFiles = inputFiles.replace('\\','')
135			inputFiles = inputFiles.replace('"','')
136			inputFileNames = inputFiles.split(',')
137	ewv	1.1	inModule.setFileNames(*inputFileNames)
138
139	spiga	1.11	# handle parent files if needed
140			if parentFiles:
141			parentFiles = parentFiles.replace('\\','')
142			parentFiles = parentFiles.replace('"','')
143			parentFileNames = parentFiles.split(',')
144			inModule.setSecondaryFileNames(*parentFileNames)
145
146	ewv	1.2	# Pythia parameters
147			if (firstRun):
148			inModule.setFirstRun(firstRun)
149	ewv	1.3
150			incrementSeedList = []
151			preserveSeedList = []
152
153			if incrementSeeds:
154			incrementSeedList = incrementSeeds.split(',')
155			if preserveSeeds:
156			preserveSeedList = preserveSeeds.split(',')
157
158	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
159			if cfg.data.services.has_key('RandomNumberGeneratorService'): # There are random #'s to deal with
160	ewv	1.9	print "RandomNumberGeneratorService found, will attempt to change seeds"
161	ewv	1.8	ranGenerator = cfg.data.services['RandomNumberGeneratorService']
162			ranModules = getattr(ranGenerator,"moduleSeeds",None)
163			if ranModules != None: # Old format present, no matter the CMSSW version
164	ewv	1.9	print "Old-style random number seeds found, will be changed."
165	ewv	1.3	sourceSeed = int(ranGenerator.sourceSeed.value())
166	ewv	1.9	if ('sourceSeed' in preserveSeedList) or ('theSource' in preserveSeedList):
167	ewv	1.3	pass
168	ewv	1.9	elif ('sourceSeed' in incrementSeedList) or ('theSource' in incrementSeedList):
169	ewv	1.3	ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
170			else:
171			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT)))
172
173			for seed in incrementSeedList:
174			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
175			if curSeed:
176			curValue = int(curSeed.value())
177			setattr(ranGenerator.moduleSeeds,seed,CfgTypes.untracked(CfgTypes.uint32(curValue+nJob)))
178			preserveSeedList.append(seed)
179
180	ewv	1.5	for seed in ranGenerator.moduleSeeds.parameterNames_():
181	ewv	1.3	if seed not in preserveSeedList:
182			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
183			if curSeed:
184			curValue = int(curSeed.value())
185			setattr(ranGenerator.moduleSeeds,seed,CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT))))
186	ewv	1.8	elif CMSSW_major > 2 or (CMSSW_major == 2 and CMSSW_minor >= 1): # Treatment for seeds, CMSSW 2_1_x and later
187	ewv	1.9	print "New-style random number seeds found, will be changed."
188	ewv	1.7	from IOMC.RandomEngine.RandomServiceHelper import RandomNumberServiceHelper
189			randSvc = RandomNumberServiceHelper(ranGenerator)
190	ewv	1.3
191	ewv	1.7	# Increment requested seed sets
192			for seedName in incrementSeedList:
193			curSeeds = randSvc.getNamedSeed(seedName)
194			newSeeds = [x+nJob for x in curSeeds]
195			randSvc.setNamedSeed(seedName,*newSeeds)
196			preserveSeedList.append(seedName)
197	ewv	1.2
198	ewv	1.7	# Randomize remaining seeds
199			randSvc.populate(*preserveSeedList)
200	ewv	1.9	else:
201			print "Neither old nor new seed format found!"
202	ewv	1.2
203	ewv	1.8	# End version specific code
204	ewv	1.2
205	ewv	1.7	# Write out new config file in one format or the other, FUTURE: Get rid of cfg mode
206	ewv	1.1	outFile = open(outFileName,"w")
207	ewv	1.7	if outFileName.endswith('py'):
208	ewv	1.3	outFile.write("import FWCore.ParameterSet.Config as cms\n")
209	ewv	1.6	outFile.write(cmsProcess.dumpPython())
210	ewv	1.3	if (debug):
211			print "writeCfg output:"
212			print "import FWCore.ParameterSet.Config as cms"
213	ewv	1.6	print cmsProcess.dumpPython()
214	ewv	1.3	else:
215			outFile.write(str(cfg))
216			if (debug):
217			print "writeCfg output:"
218			print str(cfg)
219	ewv	1.1	outFile.close()
220
221
222			if __name__ == '__main__' :
223			exit_status = main(sys.argv[1:])
224			sys.exit(exit_status)