python/utils/cmsswPreprocessor.py

0001 import os
0002 import sys
0003 import re
0004 import imp
0005 import timeit
0006 import subprocess
0007 from math import ceil
0008 from PhysicsTools.HeppyCore.framework.config import CFG
0009 from PhysicsTools.Heppy.utils.edmUtils import edmFileLs
0010 class CmsswPreprocessor :
0011     def __init__(self,configFile,command="cmsRun", addOrigAsSecondary=True, prefetch=False, options={}) :
0012         self.configFile=configFile
0013         self.command=command
0014         self.addOrigAsSecondary=addOrigAsSecondary
0015         self.prefetch=prefetch
0016         self.garbageFiles=[]
0017         self.options=options
0018
0019     def prefetchOneXrootdFile(self,fname):
0020         tmpdir = os.environ['TMPDIR'] if 'TMPDIR' in os.environ else "/tmp"
0021         rndchars  = "".join([hex(ord(i))[2:] for i in os.urandom(8)])
0022         localfile = "%s/%s-%s.root" % (tmpdir, os.path.basename(fname).replace(".root",""), rndchars)
0023         try:
0024             print("Fetching %s to local path %s " % (fname,localfile))
0025             start = timeit.default_timer()
0026             subprocess.check_output(["xrdcp","-f","-N",fname,localfile])
0027             print("Time used for transferring the file locally: %s s" % (timeit.default_timer() - start))
0028             return (localfile,True)
0029         except:
0030             print("Could not save file locally, will run from remote")
0031             if os.path.exists(localfile): os.remove(localfile) # delete in case of incomplete transfer
0032             return (fname,False)
0033     def maybePrefetchFiles(self,component):
0034         newfiles = []
0035         component._preprocessor_tempFiles = []
0036         for fn in component.files:
0037             if self.prefetch and fn.startswith("root://"):
0038                 (newfile,istemp) = self.prefetchOneXrootdFile(fn)
0039                 newfiles.append(newfile)
0040                 if istemp:
0041                     component._preprocessor_tempFiles.append(newfile)
0042             else:
0043                 newfiles.append(fn)
0044         component.files = newfiles
0045     def endLoop(self,component):
0046         for fname in component._preprocessor_tempFiles:
0047             print("Removing local cache file ",fname)
0048             os.remove(fname)
0049         component._preprocessor_tempFiles = []
0050     def run(self,component,wd,firstEvent,nEvents):
0051         if firstEvent != 0: raise RuntimeError("The preprocessor can't skip events at the moment")
0052         fineSplitIndex, fineSplitFactor = getattr(component, 'fineSplit', (1,1))
0053         if fineSplitFactor > 1:
0054             if len(component.files) != 1:
0055                 raise RuntimeError("Any component with fineSplit > 1 is supposed to have just a single file, while %s has %s" % (component.name, component.files))
0056             evtsInFile = edmFileLs(component.files[0])['events']
0057             if nEvents in (None, -1) or nEvents > evtsInFile: nEvents =  evtsInFile
0058             nEvents = int(ceil(nEvents/float(fineSplitFactor)))
0059             firstEvent = fineSplitIndex * nEvents
0060             # Now we will run on these events, and the output will contain only those
0061             # Thus, we switch off fine-split in the component
0062             component.fineSplit = (1,1)
0063         if nEvents is None:
0064             nEvents = -1
0065         self.maybePrefetchFiles(component)
0066         cmsswConfig = imp.load_source("cmsRunProcess",os.path.expandvars(self.configFile))
0067         inputfiles= []
0068         for fn in component.files :
0069             if not re.match("file:.*",fn) and not re.match("root:.*",fn) :
0070                 fn="file:"+fn
0071             inputfiles.append(fn)
0072
0073         # Four cases:
0074         # - no options, cmsswConfig with initialize function
0075         #     run initialize with default parameters
0076         # - filled options, cmsswConfig with initialize function
0077         #     pass on options to initialize
0078         # - no options, classic cmsswConfig
0079         #     legacy mode
0080         # - filled options, classic cmsswConfig
0081         #     legacy mode but warn that options are not passed on
0082
0083         if hasattr(cmsswConfig, "initialize"):
0084             if len(self.options) == 0:
0085                 cmsswConfig.process = cmsswConfig.initialize()
0086             else:
0087                 cmsswConfig.process = cmsswConfig.initialize(**self.options)
0088         else:
0089             if len(self.options) == 0:
0090                 pass
0091             else:
0092                 print("WARNING: cmsswPreprocessor received options but can't pass on to cmsswConfig")
0093
0094         cmsswConfig.process.source.fileNames = inputfiles
0095         # cmsRun will not create the output file if maxEvents==0, leading to crash of the analysis downstream.
0096         # Thus, we set nEvents = 1 if the input file is empty (the output file will be empty as well).
0097         cmsswConfig.process.maxEvents.input = 1 if (fineSplitFactor>1 and nEvents==0) else nEvents
0098         cmsswConfig.process.source.skipEvents = cmsswConfig.cms.untracked.uint32(0 if (fineSplitFactor>1 and nEvents==0) else firstEvent)
0099         #fixme: implement skipEvent / firstevent
0100
0101         outfilename=wd+"/cmsswPreProcessing.root"
0102         # for outName in cmsswConfig.process.endpath.moduleNames():
0103         for module in cmsswConfig.process.endpaths.viewvalues():
0104             for outName in module.moduleNames():
0105                 out = getattr(cmsswConfig.process,outName)
0106             if not hasattr(out,"fileName"): continue
0107             out.fileName = outfilename
0108
0109         if not hasattr(component,"options"):
0110             component.options = CFG(name="postCmsrunOptions")
0111         #use original as primary and new as secondary
0112         #component.options.inputFiles= component.files
0113         #component.options.secondaryInputFiles=[outfilename]
0114
0115         #use new as primary and original as secondary
0116         if self.addOrigAsSecondary:
0117             component.options.secondaryInputFiles= component.files
0118         component.options.inputFiles=[outfilename]
0119         component.files=[outfilename]
0120
0121         configfile=wd+"/cmsRun_config.py"
0122         f = open(configfile, 'w')
0123         f.write(cmsswConfig.process.dumpPython())
0124         f.close()
0125         runstring="%s %s >& %s/cmsRun.log" % (self.command,configfile,wd)
0126         print("Running pre-processor: %s " %runstring)
0127         ret=os.system(runstring)
0128         if ret != 0:
0129             print("CMSRUN failed")
0130             exit(ret)
0131         return component