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 Service
from FWCore.ParameterSet.Types     import *
from FWCore.ParameterSet.Modules   import OutputModule

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("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.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":
            self._poolList.append(visitee.fileName.value())
    def leave(self,visitee):
        pass

    def getList(self):
        return self._poolList
Revision:	1.32
Committed:	Tue Nov 3 17:53:04 2009 UTC (15 years, 5 months ago) by slacapra
Content type:	text/x-python
Branch:	MAIN
Changes since 1.31:	+14 -9 lines
Log Message:	fix problem in pickling pset.py with multicrab
#	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			from FWCore.ParameterSet.Modules import Service
14			from FWCore.ParameterSet.Types import *
15	ewv	1.31	from FWCore.ParameterSet.Modules import OutputModule
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	spiga	1.29	common.logger.debug("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	slacapra	1.32	Write out modified CMSSW.py
70	ewv	1.12	"""
71	gutsche	1.1
72	slacapra	1.32	pklFileName=common.work_space.jobDir()+name+".pkl"
73			pklFile = open(pklFileName,"w")
74			myPickle = pickle.Pickler(pklFile)
75			myPickle.dump(self.cmsProcess)
76			pklFile.close()
77			pklFile = open(pklFileName,"rb")
78	ewv	1.16	outFile = open(common.work_space.jobDir()+name,"w")
79	slacapra	1.32	outFile.write("import FWCore.ParameterSet.Config as cms\n")
80			outFile.write("import pickle\n")
81			outFile.write("pickledCfg=\"\"\"")
82			outFile.write(pklFile.read())
83			outFile.write("\"\"\"\n")
84			outFile.write("process = pickle.loads(pickledCfg)\n")
85			pklFile.close()
86	ewv	1.16	outFile.close()
87	gutsche	1.1
88			return
89	gutsche	1.6
90	slacapra	1.18	def getTFileService(self):
91			""" Get Output filename from TFileService and return it. If not existing, return None """
92			if not self.cfg.data.services.has_key('TFileService'):
93			return None
94			tFileService = self.cfg.data.services['TFileService']
95			if "fileName" in tFileService.parameterNames_():
96			fileName = getattr(tFileService,'fileName',None).value()
97			return fileName
98			return None
99	ewv	1.22
100	slacapra	1.19	def getPoolOutputModule(self):
101			""" Get Output filename from PoolOutputModule and return it. If not existing, return None """
102	ewv	1.31	outputFinder = PoolOutputFinder()
103			for p in self.cfg.data.endpaths.itervalues():
104			p.visit(outputFinder)
105			return outputFinder.getList()
106	slacapra	1.18
107	ewv	1.28	def getBadFilesSetting(self):
108			setting = False
109			try:
110			if self.cfg.data.source.skipBadFiles.value():
111			setting = True
112			except AttributeError:
113			pass # Either no source or no setting of skipBadFiles
114			return setting
115	ewv	1.31
116			class PoolOutputFinder(object):
117
118			def __init__(self):
119			self._poolList = []
120			def enter(self,visitee):
121			if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule":
122			self._poolList.append(visitee.fileName.value())
123			def leave(self,visitee):
124			pass
125
126			def getList(self):
127			return self._poolList