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'):
    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 < 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.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_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(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.6
Committed:	Wed Apr 30 19:41:29 2008 UTC (17 years ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_2_0_pre17, CRAB_2_2_0_pre16
Changes since 1.5:	+8 -8 lines
Log Message:	Finalize support for 2_1_x and supplying python config file in crab.cfg
#	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	ewv	1.6	if fileName.endswith('py'):
99	ewv	1.1	handle = open(fileName, 'r')
100			try: # Nested form for Python < 2.5
101			try:
102			cfo = imp.load_source("pycfg", fileName, handle)
103	ewv	1.6	cmsProcess = cfo.process
104	ewv	1.1	except Exception, ex:
105			msg = "Your pycfg file is not valid python: %s" % str(ex)
106			raise "Error: ",msg
107			finally:
108			handle.close()
109			else:
110			try:
111			cfo = include(fileName)
112	ewv	1.6	cmsProcess = cfo
113	ewv	1.1	except Exception, ex:
114			msg = "The cfg file is not valid, %s\n" % str(ex)
115			raise "Error: ",msg
116	ewv	1.6	cfg = CfgInterface(cmsProcess)
117	ewv	1.1
118			# Set parameters for job
119	ewv	1.6	inModule = cfg.inputSource
120	ewv	1.3	if maxEvents:
121			cfg.maxEvents.setMaxEventsInput(maxEvents)
122	ewv	1.1
123	ewv	1.3	if skipEvents:
124			inModule.setSkipEvents(skipEvents)
125	ewv	1.1
126	ewv	1.3	if inputFiles:
127			inputFiles = inputFiles.replace('\\','')
128			inputFiles = inputFiles.replace('"','')
129			inputFileNames = inputFiles.split(',')
130	ewv	1.1	inModule.setFileNames(*inputFileNames)
131
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	ewv	1.6	# FUTURE: This function tests the CMSSW version. Can be simplified as we drop support for old versions
145	ewv	1.3	if CMSSW_major < 3: # True for now, should be < 2 when really ready
146			# Treatment for seeds, CMSSW < 2_0_x
147			if cfg.data.services.has_key('RandomNumberGeneratorService'):
148			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
149	ewv	1.2	ranModules = ranGenerator.moduleSeeds
150
151	ewv	1.3	_MAXINT = 900000000
152
153			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			else:
175			# Treatment for seeds, CMSSW => 2_0_x
176	ewv	1.5	#from RandomService import RandomSeedService
177	ewv	1.3
178
179			# This code not currently working because randSvc is not part of the actual configuration file
180	ewv	1.2
181	ewv	1.3	# Translate old format to new format first
182			randSvc = RandomSeedService()
183			try:
184			ranGenerator = cfg.data.services['RandomNumberGeneratorService']
185			ranModules = ranGenerator.moduleSeeds
186			for seed in ranGenerator.moduleSeeds.parameters().keys():
187			curSeed = getattr(ranGenerator.moduleSeeds,seed,None)
188			curValue = int(curSeed.value())
189			setattr(randSvc,seed,CfgTypes.PSet())
190			curPSet = getattr(randSvc,seed,None)
191			curPSet.initialSeed = CfgTypes.untracked(CfgTypes.uint32(curValue))
192			del ranGenerator.moduleSeeds # Get rid of seeds in old format
193			# Doesn't work, filter is false randSvc.populate()
194	ewv	1.2
195	ewv	1.3	except:
196			print "Problems converting old seeds to new format"
197	ewv	1.2
198	ewv	1.3	# Write out new config file in one format or the other
199	ewv	1.1	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	ewv	1.6	outFile.write(cmsProcess.dumpPython())
203	ewv	1.3	if (debug):
204			print "writeCfg output:"
205			print "import FWCore.ParameterSet.Config as cms"
206	ewv	1.6	print cmsProcess.dumpPython()
207	ewv	1.3	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