1 |
|
#!/usr/bin/env python |
2 |
< |
|
3 |
< |
import os, sys, commands |
2 |
> |
|
3 |
> |
import os |
4 |
|
import common |
5 |
+ |
import imp |
6 |
+ |
import pickle |
7 |
+ |
|
8 |
|
from crab_util import * |
9 |
|
from crab_exceptions import * |
10 |
|
|
11 |
< |
# Crabpydir=commands.getoutput('which crab') |
12 |
< |
# Topdir=string.replace(Crabpydir,'/python/crab','') |
13 |
< |
# sys.path.append(Topdir+'/PsetCode') |
14 |
< |
|
15 |
< |
from cmsconfig import cmsconfig |
16 |
< |
from CfgInterface import CfgInterface |
11 |
> |
from ProdCommon.CMSConfigTools.ConfigAPI.CfgInterface import CfgInterface |
12 |
> |
# FIXME: Cleanup includes from FWCore. Most of this is not needed. |
13 |
> |
#from FWCore.ParameterSet.Config import include |
14 |
> |
from FWCore.ParameterSet.DictTypes import SortedKeysDict |
15 |
> |
from FWCore.ParameterSet.Modules import OutputModule |
16 |
> |
from FWCore.ParameterSet.Modules import Service |
17 |
> |
from FWCore.ParameterSet.Types import * |
18 |
> |
|
19 |
> |
import FWCore.ParameterSet.Types as CfgTypes |
20 |
> |
import FWCore.ParameterSet.Modules as CfgModules |
21 |
> |
import FWCore.ParameterSet.Config as cms |
22 |
|
|
23 |
|
class PsetManipulator: |
24 |
|
def __init__(self, pset): |
25 |
< |
""" |
26 |
< |
Convert Pset in Python format |
19 |
< |
and initialize |
25 |
> |
""" |
26 |
> |
Read in Pset object and initialize |
27 |
|
""" |
28 |
|
|
29 |
|
self.pset = pset |
30 |
< |
#convert Pset |
31 |
< |
self.pyPset = os.path.basename(pset) |
32 |
< |
cmd = 'EdmConfigToPython > '+common.work_space.shareDir()+self.pyPset+'py < '+ self.pset |
33 |
< |
#cmd = 'EdmConfigToPython > '+common.work_space.shareDir()+self.pset+'py < '+ self.pset |
34 |
< |
cmd_out = runCommand(cmd) |
35 |
< |
if cmd_out != '': |
36 |
< |
msg = 'Could not convert Pset.cfg into python Dictionary \n' |
37 |
< |
msg1= ' Did you do eval `scramv1 runtime ...` from your CMSSW working area ?' |
38 |
< |
raise CrabException(msg+msg1) |
30 |
> |
|
31 |
> |
common.logger.debug("PsetManipulator::__init__: PSet file = "+self.pset) |
32 |
> |
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 |
> |
self.cmsProcess = self.cfo.process |
37 |
> |
except Exception, ex: |
38 |
> |
msg = "Your config file is not valid python: %s" % str(ex) |
39 |
> |
raise CrabException(msg) |
40 |
> |
finally: |
41 |
> |
handle.close() |
42 |
> |
|
43 |
> |
self.cfg = CfgInterface(self.cmsProcess) |
44 |
> |
try: # Quiet the output |
45 |
> |
if self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery.value() < 100: |
46 |
> |
self.cfg.data.MessageLogger.cerr.FwkReport.reportEvery = cms.untracked.int32(100) |
47 |
> |
except AttributeError: |
48 |
|
pass |
33 |
– |
|
34 |
– |
par = file(common.work_space.shareDir()+self.pyPset+'py').read() |
35 |
– |
# par = file(common.work_space.shareDir()+self.pset+'py').read() |
36 |
– |
|
37 |
– |
# get PSet |
38 |
– |
self.cfg = CfgInterface(par,True) |
39 |
– |
|
40 |
– |
|
41 |
– |
def inputModule(self, source): |
42 |
– |
""" Clean String FileName if there |
43 |
– |
and add vString Filenames key |
44 |
– |
""" |
45 |
– |
# set input module |
46 |
– |
inModule = self.cfg.inputSource |
47 |
– |
inModule.cleanStringFileNames() ## Add Daniele |
48 |
– |
inModule.setFileNames(source) |
49 |
– |
return |
49 |
|
|
50 |
|
def maxEvent(self, maxEv): |
51 |
< |
""" """ |
52 |
< |
# set input module |
53 |
< |
inModule = self.cfg.inputSource |
54 |
< |
inModule.cleanMaxEvent() |
56 |
< |
inModule.setMaxEvents(maxEv) ## Add Daniele |
51 |
> |
""" |
52 |
> |
Set max event in the standalone untracked module |
53 |
> |
""" |
54 |
> |
self.cfg.maxEvents.setMaxEventsInput(maxEv) |
55 |
|
return |
56 |
|
|
57 |
< |
|
58 |
< |
|
59 |
< |
def outputModule(self, output): |
60 |
< |
|
61 |
< |
#set output module |
62 |
< |
outModule = self.cfg.outputModules['out'] |
65 |
< |
outModule.setFileName('file:'+str(output)) |
66 |
< |
|
57 |
> |
def skipEvent(self, skipEv): |
58 |
> |
""" |
59 |
> |
Set max event in the standalone untracked module |
60 |
> |
""" |
61 |
> |
if self.cfg.inputSource.sourceType not in ['EmptySource']: |
62 |
> |
self.cfg.inputSource.setSkipEvents(skipEv) |
63 |
|
return |
64 |
|
|
65 |
|
def psetWriter(self, name): |
66 |
+ |
""" |
67 |
+ |
Write out modified CMSSW.py |
68 |
+ |
""" |
69 |
|
|
70 |
< |
configObject = cmsconfig(str(self.cfg)) |
70 |
> |
pklFileName = common.work_space.jobDir() + name + ".pkl" |
71 |
> |
pklFile = open(pklFileName, "wb") |
72 |
> |
myPickle = pickle.Pickler(pklFile) |
73 |
> |
myPickle.dump(self.cmsProcess) |
74 |
> |
pklFile.close() |
75 |
> |
|
76 |
> |
outFile = open(common.work_space.jobDir()+name, "w") |
77 |
> |
outFile.write("import FWCore.ParameterSet.Config as cms\n") |
78 |
> |
outFile.write("import pickle\n") |
79 |
> |
outFile.write("process = pickle.load(open('%s', 'rb'))\n" % (name + ".pkl")) |
80 |
> |
outFile.close() |
81 |
|
|
73 |
– |
file1 = open(common.work_space.jobDir()+name,"w") |
74 |
– |
file1.write(str(configObject.asConfigurationString())) |
75 |
– |
file1.close() |
82 |
|
|
83 |
|
return |
84 |
+ |
|
85 |
+ |
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 |
+ |
|
95 |
+ |
def getPoolOutputModule(self): |
96 |
+ |
""" Get Output filename from PoolOutputModule and return it. If not existing, return None """ |
97 |
+ |
outputFinder = PoolOutputFinder() |
98 |
+ |
for p in self.cfg.data.endpaths.itervalues(): |
99 |
+ |
p.visit(outputFinder) |
100 |
+ |
return outputFinder.getDict() |
101 |
+ |
#return outputFinder.getList() |
102 |
+ |
|
103 |
+ |
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 |
+ |
|
112 |
+ |
class PoolOutputFinder(object): |
113 |
+ |
|
114 |
+ |
def __init__(self): |
115 |
+ |
self._poolList = [] |
116 |
+ |
self._poolDict = {} |
117 |
+ |
|
118 |
+ |
def enter(self,visitee): |
119 |
+ |
if isinstance(visitee,OutputModule) and visitee.type_() == "PoolOutputModule": |
120 |
+ |
filename=visitee.fileName.value().split(":")[-1] |
121 |
+ |
self._poolList.append(filename) |
122 |
+ |
|
123 |
+ |
try: |
124 |
+ |
selectEvents = visitee.SelectEvents.SelectEvents.value() |
125 |
+ |
except AttributeError: |
126 |
+ |
selectEvents = None |
127 |
+ |
try: |
128 |
+ |
dataset = visitee.dataset.filterName.value() |
129 |
+ |
except AttributeError: |
130 |
+ |
dataset = None |
131 |
+ |
self._poolDict.update({filename:{'dataset':dataset, 'selectEvents':selectEvents}}) |
132 |
+ |
|
133 |
+ |
def leave(self,visitee): |
134 |
+ |
pass |
135 |
+ |
|
136 |
+ |
def getList(self): |
137 |
+ |
return self._poolList |
138 |
+ |
|
139 |
+ |
def getDict(self): |
140 |
+ |
#### FEDE FOR MULTI #### |
141 |
+ |
return self._poolDict |