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
            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.2
Committed:	Thu Apr 15 21:54:29 2010 UTC (15 years ago) by ewv
Content type:	text/x-python
Branch:	CRAB_2_7_1_branch
CVS Tags:	CRAB_2_7_2_pre4, CRAB_2_7_2_pre3
Changes since 1.34.2.1:	+9 -6 lines
Log Message:	Heartbeat changes but not binary pickle to 271_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	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	spiga	1.34.2.1	if self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery.value() < 100:
45			self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery = cms.untracked.int32(100)
46	ewv	1.33	except AttributeError:
47			pass
48	spiga	1.7
49	gutsche	1.1	def maxEvent(self, maxEv):
50	ewv	1.12	"""
51			Set max event in the standalone untracked module
52			"""
53			self.cfg.maxEvents.setMaxEventsInput(maxEv)
54	gutsche	1.1	return
55
56	ewv	1.24	def skipEvent(self, skipEv):
57			"""
58			Set max event in the standalone untracked module
59			"""
60			self.cfg.inputSource.setSkipEvents(skipEv)
61			return
62
63	gutsche	1.1	def psetWriter(self, name):
64	ewv	1.12	"""
65	slacapra	1.32	Write out modified CMSSW.py
66	ewv	1.12	"""
67	gutsche	1.1
68	ewv	1.34.2.2	pklFileName=common.work_space.jobDir()+name+".pkl"
69			pklFile = open(pklFileName,"w")
70	slacapra	1.32	myPickle = pickle.Pickler(pklFile)
71			myPickle.dump(self.cmsProcess)
72			pklFile.close()
73	ewv	1.34.2.2	pklFile = open(pklFileName,"rb")
74			outFile = open(common.work_space.jobDir()+name,"w")
75	slacapra	1.32	outFile.write("import FWCore.ParameterSet.Config as cms\n")
76			outFile.write("import pickle\n")
77	ewv	1.34.2.2	outFile.write("pickledCfg=\"\"\"")
78			outFile.write(pklFile.read())
79			outFile.write("\"\"\"\n")
80			outFile.write("process = pickle.loads(pickledCfg)\n")
81			pklFile.close()
82	ewv	1.16	outFile.close()
83	gutsche	1.1
84			return
85	gutsche	1.6
86	slacapra	1.18	def getTFileService(self):
87			""" Get Output filename from TFileService and return it. If not existing, return None """
88			if not self.cfg.data.services.has_key('TFileService'):
89			return None
90			tFileService = self.cfg.data.services['TFileService']
91			if "fileName" in tFileService.parameterNames_():
92			fileName = getattr(tFileService,'fileName',None).value()
93			return fileName
94			return None
95	ewv	1.22
96	slacapra	1.19	def getPoolOutputModule(self):
97			""" Get Output filename from PoolOutputModule and return it. If not existing, return None """
98	ewv	1.31	outputFinder = PoolOutputFinder()
99			for p in self.cfg.data.endpaths.itervalues():
100			p.visit(outputFinder)
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			def enter(self,visitee):
117			if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule":
118	slacapra	1.34	filename=visitee.fileName.value().split(":")[-1]
119			self._poolList.append(filename)
120	ewv	1.31	def leave(self,visitee):
121			pass
122
123			def getList(self):
124			return self._poolList