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 ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
from FWCore.ParameterSet.DictTypes import SortedKeysDict
from FWCore.ParameterSet.Modules   import OutputModule
from FWCore.ParameterSet.Modules   import Service
from FWCore.ParameterSet.Types     import *

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

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

        self.pset = pset

        common.logger.debug("PsetManipulator::__init__: PSet file = "+self.pset)
        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()

        self.cfg = CfgInterface(self.cmsProcess)
        try: # Quiet the output
            if self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery.value() < 100:
                self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery = cms.untracked.int32(100)
        except AttributeError:
            pass

    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.py
        """

        pklFileName=common.work_space.jobDir()+name+".pkl"
        pklFile = open(pklFileName,"w")
        myPickle = pickle.Pickler(pklFile)
        myPickle.dump(self.cmsProcess)
        pklFile.close()
        pklFile = open(pklFileName,"rb")
        outFile = open(common.work_space.jobDir()+name,"w")
        outFile.write("import FWCore.ParameterSet.Config as cms\n")
        outFile.write("import pickle\n")
        outFile.write("pickledCfg=\"\"\"")
        outFile.write(pklFile.read())
        outFile.write("\"\"\"\n")
        outFile.write("process = pickle.loads(pickledCfg)\n")
        pklFile.close()
        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 """
        outputFinder = PoolOutputFinder()
        for p  in self.cfg.data.endpaths.itervalues():
            p.visit(outputFinder)
        return outputFinder.getList()

    def getBadFilesSetting(self):
        setting = False
        try:
            if self.cfg.data.source.skipBadFiles.value():
                setting = True
        except AttributeError:
            pass # Either no source or no setting of skipBadFiles
        return setting

class PoolOutputFinder(object):

    def __init__(self):
        self._poolList = []
    def enter(self,visitee):
        if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule":
            filename=visitee.fileName.value().split(":")[-1]
            self._poolList.append(filename)
    def leave(self,visitee):
        pass

    def getList(self):
        return self._poolList
Revision:	1.34.2.3
Committed:	Mon Apr 26 16:13:42 2010 UTC (15 years ago) by ewv
Content type:	text/x-python
Branch:	CRAB_2_7_1_branch
CVS Tags:	CRAB_2_7_2_p1, CRAB_2_7_1_branch_firstMERGE, CRAB_2_7_2
Changes since 1.34.2.2:	+0 -1 lines
Log Message:	Fix for 3.7 on the branch
#	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
11			from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface
12			from FWCore.ParameterSet.DictTypes import SortedKeysDict
13	ewv	1.33	from FWCore.ParameterSet.Modules import OutputModule
14	ewv	1.12	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	ewv	1.33	import FWCore.ParameterSet.Config as cms
20	gutsche	1.1
21			class PsetManipulator:
22			def __init__(self, pset):
23	ewv	1.12	"""
24			Read in Pset object and initialize
25	gutsche	1.1	"""
26
27			self.pset = pset
28	ewv	1.33
29	spiga	1.29	common.logger.debug("PsetManipulator::__init__: PSet file = "+self.pset)
30	ewv	1.33	handle = open(self.pset, 'r')
31			try: # Nested form for Python < 2.5
32	ewv	1.14	try:
33	ewv	1.33	self.cfo = imp.load_source("pycfg", self.pset, handle)
34			self.cmsProcess = self.cfo.process
35	ewv	1.14	except Exception, ex:
36	ewv	1.33	msg = "Your config file is not valid python: %s" % str(ex)
37	slacapra	1.15	raise CrabException(msg)
38	ewv	1.33	finally:
39			handle.close()
40
41	ewv	1.17	self.cfg = CfgInterface(self.cmsProcess)
42	ewv	1.33	try: # Quiet the output
43	spiga	1.34.2.1	if self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery.value() < 100:
44			self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery = cms.untracked.int32(100)
45	ewv	1.33	except AttributeError:
46			pass
47	spiga	1.7
48	gutsche	1.1	def maxEvent(self, maxEv):
49	ewv	1.12	"""
50			Set max event in the standalone untracked module
51			"""
52			self.cfg.maxEvents.setMaxEventsInput(maxEv)
53	gutsche	1.1	return
54
55	ewv	1.24	def skipEvent(self, skipEv):
56			"""
57			Set max event in the standalone untracked module
58			"""
59			self.cfg.inputSource.setSkipEvents(skipEv)
60			return
61
62	gutsche	1.1	def psetWriter(self, name):
63	ewv	1.12	"""
64	slacapra	1.32	Write out modified CMSSW.py
65	ewv	1.12	"""
66	gutsche	1.1
67	ewv	1.34.2.2	pklFileName=common.work_space.jobDir()+name+".pkl"
68			pklFile = open(pklFileName,"w")
69	slacapra	1.32	myPickle = pickle.Pickler(pklFile)
70			myPickle.dump(self.cmsProcess)
71			pklFile.close()
72	ewv	1.34.2.2	pklFile = open(pklFileName,"rb")
73			outFile = open(common.work_space.jobDir()+name,"w")
74	slacapra	1.32	outFile.write("import FWCore.ParameterSet.Config as cms\n")
75			outFile.write("import pickle\n")
76	ewv	1.34.2.2	outFile.write("pickledCfg=\"\"\"")
77			outFile.write(pklFile.read())
78			outFile.write("\"\"\"\n")
79			outFile.write("process = pickle.loads(pickledCfg)\n")
80			pklFile.close()
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	ewv	1.31	outputFinder = PoolOutputFinder()
98			for p in self.cfg.data.endpaths.itervalues():
99			p.visit(outputFinder)
100			return outputFinder.getList()
101	slacapra	1.18
102	ewv	1.28	def getBadFilesSetting(self):
103			setting = False
104			try:
105			if self.cfg.data.source.skipBadFiles.value():
106			setting = True
107			except AttributeError:
108			pass # Either no source or no setting of skipBadFiles
109			return setting
110	ewv	1.31
111			class PoolOutputFinder(object):
112
113			def __init__(self):
114			self._poolList = []
115			def enter(self,visitee):
116			if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule":
117	slacapra	1.34	filename=visitee.fileName.value().split(":")[-1]
118			self._poolList.append(filename)
119	ewv	1.31	def leave(self,visitee):
120			pass
121
122			def getList(self):
123			return self._poolList