CRAB/python/PsetManipulator.py

#!/usr/bin/env python

import os
import common
import imp
import pickle

from crab_util import *
from crab_exceptions import *
from crab_logger import Logger

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

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

class PsetManipulator:
    def __init__(self, pset):
        """
        Read in Pset object and initialize
        """

        self.pset = pset
        #convert Pset
        from FWCore.ParameterSet.Config import include
        common.logger.debug(3,"PsetManipulator::__init__: PSet file = "+self.pset)
        # FUTURE: Can drop cfg mode for CMSSW < 2_1_x
        if self.pset.endswith('py'):
            handle = open(self.pset, 'r')
            try:   # Nested form for Python < 2.5
                try:
                    self.cfo = imp.load_source("pycfg", self.pset, handle)
                    self.cmsProcess = self.cfo.process
                except Exception, ex:
                    msg = "Your config file is not valid python: %s" % str(ex)
                    raise CrabException(msg)
            finally:
                handle.close()
        else:
            try:
                self.cfo = include(self.pset)
                self.cmsProcess = self.cfo
            except Exception, ex:
                msg =  "Your cfg file is not valid, %s\n" % str(ex)
                msg += "  https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideCrabFaq#Problem_with_ParameterSet_parsin\n"
                msg += "  may help you understand the problem."
                raise CrabException(msg)
        self.cfg = CfgInterface(self.cmsProcess)

    def maxEvent(self, maxEv):
        """
        Set max event in the standalone untracked module
        """
        self.cfg.maxEvents.setMaxEventsInput(maxEv)
        return

    def skipEvent(self, skipEv):
        """
        Set max event in the standalone untracked module
        """
        self.cfg.inputSource.setSkipEvents(skipEv)
        return

    def psetWriter(self, name):
        """
        Write out modified CMSSW.cfg
        """

        # FUTURE: Can drop cfg mode for CMSSW < 2_1_x
        outFile = open(common.work_space.jobDir()+name,"w")
        if name.endswith('py'):
            outFile.write("import FWCore.ParameterSet.Config as cms\n")
            outFile.write("import pickle\n")
            outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(self.cmsProcess))
            outFile.write("process = pickle.loads(pickledCfg)\n")
        else:
            outFile.write(self.cfg.data.dumpConfig())
        outFile.close()

        return

    def getTFileService(self):
        """ Get Output filename from TFileService and return it. If not existing, return None """
        if not self.cfg.data.services.has_key('TFileService'):
            return None
        tFileService = self.cfg.data.services['TFileService']
        if "fileName" in tFileService.parameterNames_():
            fileName = getattr(tFileService,'fileName',None).value()
            return fileName
        return None

    def getPoolOutputModule(self):
        """ Get Output filename from PoolOutputModule and return it. If not existing, return None """
        if not self.cfg.data.outputModules:
            return None
        poolOutputModule = self.cfg.data.outputModules
        for out in poolOutputModule:
            return poolOutputModule[out].fileName.value()

Revision:	1.27
Committed:	Mon May 4 19:10:16 2009 UTC (15 years, 11 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_6_0_pre2, CRAB_2_6_0_pre1
Changes since 1.26:	+4 -8 lines
Log Message:	Use pickled config file and problems with triple-quote commenting.
#	User	Rev	Content
1	gutsche	1.1	#!/usr/bin/env python
2	ewv	1.12
3	slacapra	1.8	import os
4	gutsche	1.1	import common
5	ewv	1.16	import imp
6	ewv	1.27	import pickle
7	ewv	1.16
8	gutsche	1.1	from crab_util import *
9			from crab_exceptions import *
10	ewv	1.12	from crab_logger import Logger
11
12			from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
13			from FWCore.ParameterSet.DictTypes import SortedKeysDict
14			from FWCore.ParameterSet.Modules import Service
15			from FWCore.ParameterSet.Types import *
16	gutsche	1.1
17	ewv	1.12	import FWCore.ParameterSet.Types as CfgTypes
18			import FWCore.ParameterSet.Modules as CfgModules
19	gutsche	1.1
20			class PsetManipulator:
21			def __init__(self, pset):
22	ewv	1.12	"""
23			Read in Pset object and initialize
24	gutsche	1.1	"""
25
26			self.pset = pset
27			#convert Pset
28	ewv	1.12	from FWCore.ParameterSet.Config import include
29			common.logger.debug(3,"PsetManipulator::__init__: PSet file = "+self.pset)
30	ewv	1.26	# FUTURE: Can drop cfg mode for CMSSW < 2_1_x
31	ewv	1.16	if self.pset.endswith('py'):
32	slacapra	1.15	handle = open(self.pset, 'r')
33			try: # Nested form for Python < 2.5
34			try:
35			self.cfo = imp.load_source("pycfg", self.pset, handle)
36	ewv	1.17	self.cmsProcess = self.cfo.process
37	slacapra	1.15	except Exception, ex:
38	ewv	1.16	msg = "Your config file is not valid python: %s" % str(ex)
39	slacapra	1.15	raise CrabException(msg)
40			finally:
41			handle.close()
42			else:
43	ewv	1.14	try:
44	slacapra	1.15	self.cfo = include(self.pset)
45	ewv	1.17	self.cmsProcess = self.cfo
46	ewv	1.14	except Exception, ex:
47	slacapra	1.15	msg = "Your cfg file is not valid, %s\n" % str(ex)
48	slacapra	1.21	msg += " https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideCrabFaq#Problem_with_ParameterSet_parsin\n"
49	slacapra	1.15	msg += " may help you understand the problem."
50			raise CrabException(msg)
51	ewv	1.17	self.cfg = CfgInterface(self.cmsProcess)
52	spiga	1.7
53	gutsche	1.1	def maxEvent(self, maxEv):
54	ewv	1.12	"""
55			Set max event in the standalone untracked module
56			"""
57			self.cfg.maxEvents.setMaxEventsInput(maxEv)
58	gutsche	1.1	return
59
60	ewv	1.24	def skipEvent(self, skipEv):
61			"""
62			Set max event in the standalone untracked module
63			"""
64			self.cfg.inputSource.setSkipEvents(skipEv)
65			return
66
67	gutsche	1.1	def psetWriter(self, name):
68	ewv	1.12	"""
69			Write out modified CMSSW.cfg
70			"""
71	gutsche	1.1
72	ewv	1.16	# FUTURE: Can drop cfg mode for CMSSW < 2_1_x
73			outFile = open(common.work_space.jobDir()+name,"w")
74			if name.endswith('py'):
75	slacapra	1.18	outFile.write("import FWCore.ParameterSet.Config as cms\n")
76	ewv	1.27	outFile.write("import pickle\n")
77			outFile.write("pickledCfg=\"\"\"%s\"\"\"\n" % pickle.dumps(self.cmsProcess))
78			outFile.write("process = pickle.loads(pickledCfg)\n")
79	ewv	1.16	else:
80	ewv	1.22	outFile.write(self.cfg.data.dumpConfig())
81	ewv	1.16	outFile.close()
82	gutsche	1.1
83			return
84	gutsche	1.6
85	slacapra	1.18	def getTFileService(self):
86			""" Get Output filename from TFileService and return it. If not existing, return None """
87			if not self.cfg.data.services.has_key('TFileService'):
88			return None
89			tFileService = self.cfg.data.services['TFileService']
90			if "fileName" in tFileService.parameterNames_():
91			fileName = getattr(tFileService,'fileName',None).value()
92			return fileName
93			return None
94	ewv	1.22
95	slacapra	1.19	def getPoolOutputModule(self):
96			""" Get Output filename from PoolOutputModule and return it. If not existing, return None """
97	slacapra	1.20	if not self.cfg.data.outputModules:
98			return None
99			poolOutputModule = self.cfg.data.outputModules
100			for out in poolOutputModule:
101			return poolOutputModule[out].fileName.value()
102	slacapra	1.18