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.Config    import include
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
            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.getDict()
        #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 = []
        self._poolDict = {}
    def enter(self,visitee):
        if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule":
            filename = visitee.fileName.value().split(":")[-1]
            self._poolList.append(filename)

            try:
                selectEvents = visitee.SelectEvents.SelectEvents.value()
            except AttributeError:
                selectEvents = None
            try:
                dataset = visitee.dataset.filterName.value()
            except AttributeError:
                dataset = None
            self._poolDict.update({filename:{'dataset':dataset, 'selectEvents':selectEvents}})

    def leave(self,visitee):
        pass

    def getList(self):
        return self._poolList
    def getDict(self):
        #### FOR MULTI PUBBLICATION
        return self._poolDict
Revision:	1.34.4.1
Committed:	Fri Mar 19 13:35:18 2010 UTC (15 years, 1 month ago) by fanzago
Content type:	text/x-python
Branch:	CRAB_multiout
Changes since 1.34:	+18 -2 lines
Log Message:	added changes done by eric
#	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	ewv	1.33	from FWCore.ParameterSet.Config import include
13	ewv	1.12	from FWCore.ParameterSet.DictTypes import SortedKeysDict
14	ewv	1.33	from FWCore.ParameterSet.Modules import OutputModule
15	ewv	1.12	from FWCore.ParameterSet.Modules import Service
16			from FWCore.ParameterSet.Types import *
17	gutsche	1.1
18	ewv	1.12	import FWCore.ParameterSet.Types as CfgTypes
19			import FWCore.ParameterSet.Modules as CfgModules
20	ewv	1.33	import FWCore.ParameterSet.Config as cms
21	gutsche	1.1
22			class PsetManipulator:
23			def __init__(self, pset):
24	ewv	1.12	"""
25			Read in Pset object and initialize
26	gutsche	1.1	"""
27
28			self.pset = pset
29	ewv	1.33
30	spiga	1.29	common.logger.debug("PsetManipulator::__init__: PSet file = "+self.pset)
31	ewv	1.33	handle = open(self.pset, 'r')
32			try: # Nested form for Python < 2.5
33	ewv	1.14	try:
34	ewv	1.33	self.cfo = imp.load_source("pycfg", self.pset, handle)
35			self.cmsProcess = self.cfo.process
36	ewv	1.14	except Exception, ex:
37	ewv	1.33	msg = "Your config file is not valid python: %s" % str(ex)
38	slacapra	1.15	raise CrabException(msg)
39	ewv	1.33	finally:
40			handle.close()
41
42	ewv	1.17	self.cfg = CfgInterface(self.cmsProcess)
43	ewv	1.33	try: # Quiet the output
44			self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery = cms.untracked.int32(100)
45			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	slacapra	1.32	pklFileName=common.work_space.jobDir()+name+".pkl"
68			pklFile = open(pklFileName,"w")
69			myPickle = pickle.Pickler(pklFile)
70			myPickle.dump(self.cmsProcess)
71			pklFile.close()
72			pklFile = open(pklFileName,"rb")
73	ewv	1.16	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			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	fanzago	1.34.4.1	return outputFinder.getDict()
101			#return outputFinder.getList()
102	slacapra	1.18
103	ewv	1.28	def getBadFilesSetting(self):
104			setting = False
105			try:
106			if self.cfg.data.source.skipBadFiles.value():
107			setting = True
108			except AttributeError:
109			pass # Either no source or no setting of skipBadFiles
110			return setting
111	ewv	1.31
112			class PoolOutputFinder(object):
113
114			def __init__(self):
115			self._poolList = []
116	fanzago	1.34.4.1	self._poolDict = {}
117	ewv	1.31	def enter(self,visitee):
118			if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule":
119	fanzago	1.34.4.1	filename = visitee.fileName.value().split(":")[-1]
120	slacapra	1.34	self._poolList.append(filename)
121	fanzago	1.34.4.1
122			try:
123			selectEvents = visitee.SelectEvents.SelectEvents.value()
124			except AttributeError:
125			selectEvents = None
126			try:
127			dataset = visitee.dataset.filterName.value()
128			except AttributeError:
129			dataset = None
130			self._poolDict.update({filename:{'dataset':dataset, 'selectEvents':selectEvents}})
131
132	ewv	1.31	def leave(self,visitee):
133			pass
134
135			def getList(self):
136			return self._poolList
137	fanzago	1.34.4.1	def getDict(self):
138			#### FOR MULTI PUBBLICATION
139			return self._poolDict