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, "wb")
        myPickle = pickle.Pickler(pklFile)
        myPickle.dump(self.cmsProcess)
        pklFile.close()

        outFile = open(common.work_space.jobDir()+name, "w")
        outFile.write("import FWCore.ParameterSet.Config as cms\n")
        outFile.write("import pickle\n")
        outFile.write("process = pickle.load(open('%s', 'rb'))\n" % (name + ".pkl"))
        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.1
Committed:	Thu Apr 15 08:20:31 2010 UTC (15 years ago) by spiga
Content type:	text/x-python
Branch:	CRAB_2_7_1_branch
CVS Tags:	CRAB_2_7_2_pre2
Changes since 1.34:	+8 -10 lines
Log Message:	Fix for https://savannah.cern.ch/bugs/?58045
#	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	spiga	1.34.2.1	pklFileName = common.work_space.jobDir() + name + ".pkl"
69			pklFile = open(pklFileName, "wb")
70	slacapra	1.32	myPickle = pickle.Pickler(pklFile)
71			myPickle.dump(self.cmsProcess)
72			pklFile.close()
73	spiga	1.34.2.1
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	spiga	1.34.2.1	outFile.write("process = pickle.load(open('%s', 'rb'))\n" % (name + ".pkl"))
78	ewv	1.16	outFile.close()
79	gutsche	1.1
80	spiga	1.34.2.1
81	gutsche	1.1	return
82	gutsche	1.6
83	slacapra	1.18	def getTFileService(self):
84			""" Get Output filename from TFileService and return it. If not existing, return None """
85			if not self.cfg.data.services.has_key('TFileService'):
86			return None
87			tFileService = self.cfg.data.services['TFileService']
88			if "fileName" in tFileService.parameterNames_():
89			fileName = getattr(tFileService,'fileName',None).value()
90			return fileName
91			return None
92	ewv	1.22
93	slacapra	1.19	def getPoolOutputModule(self):
94			""" Get Output filename from PoolOutputModule and return it. If not existing, return None """
95	ewv	1.31	outputFinder = PoolOutputFinder()
96			for p in self.cfg.data.endpaths.itervalues():
97			p.visit(outputFinder)
98			return outputFinder.getList()
99	slacapra	1.18
100	ewv	1.28	def getBadFilesSetting(self):
101			setting = False
102			try:
103			if self.cfg.data.source.skipBadFiles.value():
104			setting = True
105			except AttributeError:
106			pass # Either no source or no setting of skipBadFiles
107			return setting
108	ewv	1.31
109			class PoolOutputFinder(object):
110
111			def __init__(self):
112			self._poolList = []
113			def enter(self,visitee):
114			if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule":
115	slacapra	1.34	filename=visitee.fileName.value().split(":")[-1]
116			self._poolList.append(filename)
117	ewv	1.31	def leave(self,visitee):
118			pass
119
120			def getList(self):
121			return self._poolList