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

  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'))
  inputFiles = os.environ.get('InputFiles','')
  firstRun   = int(os.environ.get('FirstRun','0'))
  nJob       = int(os.environ.get('NJob','0'))
  preserveSeeds = os.environ.get('PreserveSeeds','')
  incrementSeeds = os.environ.get('IncrementSeeds','')

# Read Input cfg or python cfg file

  if (fileName.endswith('py') or fileName.endswith('pycfg') ):
    handle = open(fileName, 'r')
    try:   # Nested form for Python < 2.5
      try:
        cfo = imp.load_source("pycfg", fileName, handle)
      except Exception, ex:
        msg = "Your pycfg file is not valid python: %s" % str(ex)
        raise "Error: ",msg
    finally:
        handle.close()
    cfg = CfgInterface(cfo.process)
  else:
    try:
      cfo = include(fileName)
      cfg = CfgInterface(cfo)
    except Exception, ex:
      msg =  "The cfg file is not valid, %s\n" % str(ex)
      raise "Error: ",msg
  inModule = cfg.inputSource

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

  # FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
  # Pythia parameters
  if (firstRun):
    inModule.setFirstRun(firstRun)

  incrementSeedList = []
  preserveSeedList  = []

  if incrementSeeds:
    incrementSeedList = incrementSeeds.split(',')
  if preserveSeeds:
    preserveSeedList  = preserveSeeds.split(',')

  if CMSSW_major < 3: # True for now, should be < 2 when really ready
  # Treatment for seeds, CMSSW < 2_0_x
    if cfg.data.services.has_key('RandomNumberGeneratorService'):
      ranGenerator = cfg.data.services['RandomNumberGeneratorService']
      ranModules   = ranGenerator.moduleSeeds

      _MAXINT = 900000000

      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.parameters().keys():
        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_0_x
#from RandomService import RandomSeedService


    # This code not currently working because randSvc is not part of the actual configuration file

    # Translate old format to new format first
    randSvc = RandomSeedService()
    try:
      ranGenerator = cfg.data.services['RandomNumberGeneratorService']
      ranModules   = ranGenerator.moduleSeeds
      for seed in ranGenerator.moduleSeeds.parameters().keys():
        curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
        curValue = int(curSeed.value())
        setattr(randSvc,seed,CfgTypes.PSet())
        curPSet = getattr(randSvc,seed,None)
        curPSet.initialSeed = CfgTypes.untracked(CfgTypes.uint32(curValue))
      del ranGenerator.moduleSeeds # Get rid of seeds in old format
# Doesn't work, filter is false      randSvc.populate()

    except:
      print "Problems converting old seeds to new format"

  # Write out new config file in one format or the other

  outFile = open(outFileName,"w")
  if (outFileName.endswith('py') or outFileName.endswith('pycfg') ):
    outFile.write("import FWCore.ParameterSet.Config as cms\n")
    outFile.write(cfo.dumpPython())
    if (debug):
      print "writeCfg output:"
      print "import FWCore.ParameterSet.Config as cms"
      print cfo.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.3
Committed:	Tue Feb 26 22:21:19 2008 UTC (17 years, 2 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_2_0_pre4, CRAB_2_2_0_pre2, CRAB_2_1_1, CRAB_2_1_1_pre3, CRAB_2_2_0_pre1, CRAB_2_1_1_pre1
Branch point for:	CRAB_2_1_2_br, CRAB_2_1_1_pre2
Changes since 1.2:	+118 -56 lines
Log Message:	Python config file, remove reliance on sed
#	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.1
53			try:
54	ewv	1.3	opts, args = getopt.getopt(argv, "", ["debug", "help"])
55	ewv	1.1	except getopt.GetoptError:
56			print main.__doc__
57			sys.exit(2)
58
59	ewv	1.3	try:
60			CMSSW = os.environ['CMSSW_VERSION']
61			parts = CMSSW.split('_')
62			CMSSW_major = int(parts[1])
63			CMSSW_minor = int(parts[2])
64			CMSSW_patch = int(parts[3])
65			except KeyError, ValueError:
66			msg = "Your environment doesn't specify the CMSSW version or specifies it incorrectly"
67			raise ConfigException(msg)
68
69			# Parse command line options
70	ewv	1.1	for opt, arg in opts :
71			if opt == "--help" :
72			print main.__doc__
73			sys.exit()
74			elif opt == "--debug" :
75			debug = True
76
77			# Parse remaining parameters
78
79	ewv	1.3	try:
80			fileName = args[0];
81			outFileName = args[1];
82			except IndexError:
83			print main.__doc__
84			sys.exit()
85
86			# Optional Parameters
87
88			maxEvents = int(os.environ.get('MaxEvents','0'))
89			skipEvents = int(os.environ.get('SkipEvents','0'))
90			inputFiles = os.environ.get('InputFiles','')
91			firstRun = int(os.environ.get('FirstRun','0'))
92			nJob = int(os.environ.get('NJob','0'))
93			preserveSeeds = os.environ.get('PreserveSeeds','')
94			incrementSeeds = os.environ.get('IncrementSeeds','')
95	ewv	1.1
96	ewv	1.3	# Read Input cfg or python cfg file
97	ewv	1.1
98			if (fileName.endswith('py') or fileName.endswith('pycfg') ):
99			handle = open(fileName, 'r')
100			try: # Nested form for Python < 2.5
101			try:
102			cfo = imp.load_source("pycfg", fileName, handle)
103			except Exception, ex:
104			msg = "Your pycfg file is not valid python: %s" % str(ex)
105			raise "Error: ",msg
106			finally:
107			handle.close()
108			cfg = CfgInterface(cfo.process)
109			else:
110			try:
111			cfo = include(fileName)
112			cfg = CfgInterface(cfo)
113			except Exception, ex:
114			msg = "The cfg file is not valid, %s\n" % str(ex)
115			raise "Error: ",msg
116	ewv	1.3	inModule = cfg.inputSource
117	ewv	1.1
118			# Set parameters for job
119	ewv	1.3	if maxEvents:
120			cfg.maxEvents.setMaxEventsInput(maxEvents)
121	ewv	1.1
122	ewv	1.3	if skipEvents:
123			inModule.setSkipEvents(skipEvents)
124	ewv	1.1
125	ewv	1.3	if inputFiles:
126			inputFiles = inputFiles.replace('\\','')
127			inputFiles = inputFiles.replace('"','')
128			inputFileNames = inputFiles.split(',')
129	ewv	1.1	inModule.setFileNames(*inputFileNames)
130
131	ewv	1.3	# FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
132	ewv	1.2	# Pythia parameters
133			if (firstRun):
134			inModule.setFirstRun(firstRun)
135	ewv	1.3
136			incrementSeedList = []
137			preserveSeedList = []
138
139			if incrementSeeds:
140			incrementSeedList = incrementSeeds.split(',')
141			if preserveSeeds:
142			preserveSeedList = preserveSeeds.split(',')
143
144			if CMSSW_major < 3: # True for now, should be < 2 when really ready
145			# Treatment for seeds, CMSSW < 2_0_x
146			if cfg.data.services.has_key('RandomNumberGeneratorService'):
147			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
148	ewv	1.2	ranModules = ranGenerator.moduleSeeds
149
150	ewv	1.3	_MAXINT = 900000000
151
152			sourceSeed = int(ranGenerator.sourceSeed.value())
153			if 'sourceSeed' in preserveSeedList:
154			pass
155			elif 'sourceSeed' in incrementSeedList:
156			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(sourceSeed+nJob))
157			else:
158			ranGenerator.sourceSeed = CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT)))
159
160			for seed in incrementSeedList:
161			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
162			if curSeed:
163			curValue = int(curSeed.value())
164			setattr(ranGenerator.moduleSeeds,seed,CfgTypes.untracked(CfgTypes.uint32(curValue+nJob)))
165			preserveSeedList.append(seed)
166
167			for seed in ranGenerator.moduleSeeds.parameters().keys():
168			if seed not in preserveSeedList:
169			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
170			if curSeed:
171			curValue = int(curSeed.value())
172			setattr(ranGenerator.moduleSeeds,seed,CfgTypes.untracked(CfgTypes.uint32(_inst.randint(1,_MAXINT))))
173			else:
174			# Treatment for seeds, CMSSW => 2_0_x
175			#from RandomService import RandomSeedService
176
177
178			# This code not currently working because randSvc is not part of the actual configuration file
179	ewv	1.2
180	ewv	1.3	# Translate old format to new format first
181			randSvc = RandomSeedService()
182			try:
183			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
184			ranModules = ranGenerator.moduleSeeds
185			for seed in ranGenerator.moduleSeeds.parameters().keys():
186			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
187			curValue = int(curSeed.value())
188			setattr(randSvc,seed,CfgTypes.PSet())
189			curPSet = getattr(randSvc,seed,None)
190			curPSet.initialSeed = CfgTypes.untracked(CfgTypes.uint32(curValue))
191			del ranGenerator.moduleSeeds # Get rid of seeds in old format
192			# Doesn't work, filter is false randSvc.populate()
193	ewv	1.2
194	ewv	1.3	except:
195			print "Problems converting old seeds to new format"
196	ewv	1.2
197	ewv	1.3	# Write out new config file in one format or the other
198	ewv	1.1
199			outFile = open(outFileName,"w")
200	ewv	1.3	if (outFileName.endswith('py') or outFileName.endswith('pycfg') ):
201			outFile.write("import FWCore.ParameterSet.Config as cms\n")
202			outFile.write(cfo.dumpPython())
203			if (debug):
204			print "writeCfg output:"
205			print "import FWCore.ParameterSet.Config as cms"
206			print cfo.dumpPython()
207			else:
208			outFile.write(str(cfg))
209			if (debug):
210			print "writeCfg output:"
211			print str(cfg)
212	ewv	1.1	outFile.close()
213
214
215			if __name__ == '__main__' :
216			exit_status = main(sys.argv[1:])
217			sys.exit(exit_status)
218