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.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
Committed:	Thu Dec 17 14:45:02 2009 UTC (15 years, 4 months ago) by slacapra
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_7_2_pre1, CRAB_2_7_1, fede_170310, CRAB_2_7_1_pre12, CRAB_2_7_1_pre11, CRAB_2_7_1_pre10, CRAB_2_7_1_pre9, CRAB_LumiMask, CRAB_2_7_lumi, from_LimiMask, CRAB_2_7_1_pre8, CRAB_2_7_1_pre6, CRAB_2_7_1_pre5, CRAB_2_7_1_wmbs_pre4, CRAB_2_7_1_pre4, CRAB_2_7_1_pre3
Branch point for:	CRAB_multiout, CRAB_2_7_1_branch
Changes since 1.33:	+2 -1 lines
Log Message:	add protection against protocol: in PoolOutputModule.filename
#	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			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