python/utils/cmsswPreprocessor.py

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