CRAB/python/Splitter.py


__revision__ = "$Id: writeCfg.py,v 1.22 2009/07/29 21:20:03 ewv Exp $"
__version__ = "$Revision: 1.22 $"

import common
from sets import Set
from crab_exceptions import *
from crab_util import *

from WMCore.DataStructs.File import File
from WMCore.DataStructs.Fileset import Fileset
from WMCore.DataStructs.Run import Run
from WMCore.DataStructs.Subscription import Subscription
from WMCore.DataStructs.Workflow import Workflow
from WMCore.JobSplitting.SplitterFactory import SplitterFactory
from WMCore.SiteScreening.BlackWhiteListParser import SEBlackWhiteListParser

class JobSplitter:
    def __init__( self, cfg_params,  args ):
        self.cfg_params = cfg_params
        self.args=args

        self.lumisPerJob = -1
        self.totalNLumis = 0
        self.theNumberOfJobs = 0
        self.limitNJobs = False
        self.limitTotalLumis = False
        self.limitJobLumis = False

        #self.maxEvents
        # init BlackWhiteListParser
        self.seWhiteList = cfg_params.get('GRID.se_white_list',[])
        seBlackList = cfg_params.get('GRID.se_black_list',[])
        self.blackWhiteListParser = SEBlackWhiteListParser(self.seWhiteList, seBlackList, common.logger())


    def checkUserSettings(self):
        ## Events per job
        if self.cfg_params.has_key('CMSSW.events_per_job'):
            self.eventsPerJob =int( self.cfg_params['CMSSW.events_per_job'])
            self.selectEventsPerJob = 1
        else:
            self.eventsPerJob = -1
            self.selectEventsPerJob = 0

        ## number of jobs
        if self.cfg_params.has_key('CMSSW.number_of_jobs'):
            self.theNumberOfJobs =int( self.cfg_params['CMSSW.number_of_jobs'])
            self.selectNumberOfJobs = 1
        else:
            self.theNumberOfJobs = 0
            self.selectNumberOfJobs = 0

        if self.cfg_params.has_key('CMSSW.total_number_of_events'):
            self.total_number_of_events = int(self.cfg_params['CMSSW.total_number_of_events'])
            self.selectTotalNumberEvents = 1
            if self.selectNumberOfJobs  == 1:
                if (self.total_number_of_events != -1) and int(self.total_number_of_events) < int(self.theNumberOfJobs):
                    msg = 'Must specify at least one event per job. total_number_of_events > number_of_jobs '
                    raise CrabException(msg)
        else:
            self.total_number_of_events = 0
            self.selectTotalNumberEvents = 0


    def checkLumiSettings(self):
        """
        Check to make sure the user has specified enough information to
        perform splitting by Lumis to run the job
        """
        settings = 0
        if self.cfg_params.has_key('CMSSW.lumis_per_job'):
            self.lumisPerJob =int( self.cfg_params['CMSSW.lumis_per_job'])
            self.limitJobLumis = True
            settings += 1

        if self.cfg_params.has_key('CMSSW.number_of_jobs'):
            self.theNumberOfJobs =int( self.cfg_params['CMSSW.number_of_jobs'])
            self.limitNJobs = True
            settings += 1

        if self.cfg_params.has_key('CMSSW.total_number_of_lumis'):
            self.totalNLumis = int(self.cfg_params['CMSSW.total_number_of_lumis'])
            self.limitTotalLumis = (self.totalNLumis != -1)
            settings += 1

        if settings != 2:
            msg = 'When running on analysis datasets you must specify two and only two of:\n'
            msg = '  number_of_jobs, lumis_per_job, total_number_of_lumis'
            raise CrabException(msg)
        if self.limitNJobs and self.limitJobLumis:
            self.limitTotalLumis = True
            self.totalNLumis = self.lumisPerJob * self.theNumberOfJobs


    def ComputeSubBlockSites( self, blockSites ):
        """
        """
        sub_blockSites = {}
        for k,v in blockSites.iteritems():
            sites=self.blackWhiteListParser.checkWhiteList(v)
            if sites : sub_blockSites[k]=v
        if len(sub_blockSites) < 1:
            msg = 'WARNING: the sites %s is not hosting any part of data.'%self.seWhiteList
            raise CrabException(msg)
        return sub_blockSites

########################################################################
    def jobSplittingByEvent( self ):
        """
        Perform job splitting. Jobs run over an integer number of files
        and no more than one block.
        ARGUMENT: blockSites: dictionary with blocks as keys and list of host sites as values
        REQUIRES: self.selectTotalNumberEvents, self.selectEventsPerJob, self.selectNumberofJobs,
                  self.total_number_of_events, self.eventsPerJob, self.theNumberOfJobs,
                  self.maxEvents, self.filesbyblock
        SETS: jobDestination - Site destination(s) for each job (a list of lists)
              self.total_number_of_jobs - Total # of jobs
              self.list_of_args - File(s) job will run on (a list of lists)
        """

        jobDestination=[]
        self.checkUserSettings()
        if ( (self.selectTotalNumberEvents + self.selectEventsPerJob + self.selectNumberOfJobs) != 2 ):
            msg = 'Must define exactly two of total_number_of_events, events_per_job, or number_of_jobs.'
            raise CrabException(msg)

        blockSites = self.args['blockSites']
        pubdata = self.args['pubdata']
        filesbyblock=pubdata.getFiles()

        self.eventsbyblock=pubdata.getEventsPerBlock()
        self.eventsbyfile=pubdata.getEventsPerFile()
        self.parentFiles=pubdata.getParent()

        ## get max number of events
        self.maxEvents=pubdata.getMaxEvents()

        self.useParent = int(self.cfg_params.get('CMSSW.use_parent',0))
        noBboundary = int(self.cfg_params.get('CMSSW.no_block_boundary',0))

        if noBboundary == 1:
            if self.total_number_of_events== -1:
                msg = 'You are selecting no_block_boundary=1 which does not allow to set total_number_of_events=-1\n'
                msg +='\tYou shoud get the number of event from DBS web interface and use it for your configuration.'
                raise CrabException(msg)
            if len(self.seWhiteList) == 0 or  len(self.seWhiteList.split(',')) != 1:
                msg = 'You are selecting no_block_boundary=1 which requires to choose one and only one site.\n'
                msg += "\tPlease set se_white_list with the site's storage element name."
                raise  CrabException(msg)
            blockSites = self.ComputeSubBlockSites(blockSites)

        # ---- Handle the possible job splitting configurations ---- #
        if (self.selectTotalNumberEvents):
            totalEventsRequested = self.total_number_of_events
        if (self.selectEventsPerJob):
            eventsPerJobRequested = self.eventsPerJob
            if (self.selectNumberOfJobs):
                totalEventsRequested = self.theNumberOfJobs * self.eventsPerJob

        # If user requested all the events in the dataset
        if (totalEventsRequested == -1):
            eventsRemaining=self.maxEvents
        # If user requested more events than are in the dataset
        elif (totalEventsRequested > self.maxEvents):
            eventsRemaining = self.maxEvents
            common.logger.info("Requested "+str(self.total_number_of_events)+ " events, but only "+str(self.maxEvents)+" events are available.")
        # If user requested less events than are in the dataset
        else:
            eventsRemaining = totalEventsRequested

        # If user requested more events per job than are in the dataset
        if (self.selectEventsPerJob and eventsPerJobRequested > self.maxEvents):
            eventsPerJobRequested = self.maxEvents

        # For user info at end
        totalEventCount = 0

        if (self.selectTotalNumberEvents and self.selectNumberOfJobs):
            eventsPerJobRequested = int(eventsRemaining/self.theNumberOfJobs)

        if (self.selectNumberOfJobs):
            common.logger.info("May not create the exact number_of_jobs requested.")

        # old... to remove Daniele
        totalNumberOfJobs = 999999999

        blocks = blockSites.keys()
        blockCount = 0
        # Backup variable in case self.maxEvents counted events in a non-included block
        numBlocksInDataset = len(blocks)

        jobCount = 0
        list_of_lists = []

        # list tracking which jobs are in which jobs belong to which block
        jobsOfBlock = {}

        parString = ""
        pString = ""
        filesEventCount = 0
        msg=''

        # ---- Iterate over the blocks in the dataset until ---- #
        # ---- we've met the requested total # of events    ---- #
        while ( (eventsRemaining > 0) and (blockCount < numBlocksInDataset) and (jobCount < totalNumberOfJobs)):
            block = blocks[blockCount]
            blockCount += 1
            if block not in jobsOfBlock.keys() :
                jobsOfBlock[block] = []

            if self.eventsbyblock.has_key(block) :
                numEventsInBlock = self.eventsbyblock[block]
                common.logger.debug('Events in Block File '+str(numEventsInBlock))

                files = filesbyblock[block]
                numFilesInBlock = len(files)
                if (numFilesInBlock <= 0):
                    continue
                fileCount = 0
                if noBboundary == 0: # DD
                    # ---- New block => New job ---- #
                    parString = ""
                    pString=""
                    # counter for number of events in files currently worked on
                    filesEventCount = 0
                # flag if next while loop should touch new file
                newFile = 1
                # job event counter
                jobSkipEventCount = 0

                # ---- Iterate over the files in the block until we've met the requested ---- #
                # ---- total # of events or we've gone over all the files in this block  ---- #
                msg='\n'
                while ( (eventsRemaining > 0) and (fileCount < numFilesInBlock) and (jobCount < totalNumberOfJobs) ):
                    file = files[fileCount]
                    if self.useParent==1:
                        parent = self.parentFiles[file]
                        common.logger.log(10-1, "File "+str(file)+" has the following parents: "+str(parent))
                    if newFile :
                        try:
                            numEventsInFile = self.eventsbyfile[file]
                            common.logger.log(10-1, "File "+str(file)+" has "+str(numEventsInFile)+" events")
                            # increase filesEventCount
                            filesEventCount += numEventsInFile
                            # Add file to current job
                            parString +=  file + ','
                            if self.useParent==1:
                                for f in parent :
                                    pString += f  + ','
                            newFile = 0
                        except KeyError:
                            common.logger.info("File "+str(file)+" has unknown number of events: skipping")

                    eventsPerJobRequested = min(eventsPerJobRequested, eventsRemaining)
                    # if less events in file remain than eventsPerJobRequested
                    if ( filesEventCount - jobSkipEventCount < eventsPerJobRequested):
                        if noBboundary == 1: ## DD
                            newFile = 1
                            fileCount += 1
                        else:
                            # if last file in block
                            if ( fileCount == numFilesInBlock-1 ) :
                                # end job using last file, use remaining events in block
                                # close job and touch new file
                                fullString = parString[:-1]
                                if self.useParent==1:
                                    fullParentString = pString[:-1]
                                    list_of_lists.append([fullString,fullParentString,str(-1),str(jobSkipEventCount)])
                                else:
                                    list_of_lists.append([fullString,str(-1),str(jobSkipEventCount)])
                                msg += "Job %s can run over %s  events (last file in block).\n"%(str(jobCount+1), str(filesEventCount - jobSkipEventCount))
                                jobDestination.append(blockSites[block])
                                msg += "Job %s Destination: %s\n"%(str(jobCount+1),str(SE2CMS(jobDestination[jobCount])))
                                # fill jobs of block dictionary
                                jobsOfBlock[block].append(jobCount+1)
                                # reset counter
                                jobCount = jobCount + 1
                                totalEventCount = totalEventCount + filesEventCount - jobSkipEventCount
                                eventsRemaining = eventsRemaining - filesEventCount + jobSkipEventCount
                                jobSkipEventCount = 0
                                # reset file
                                pString = ""
                                parString = ""
                                filesEventCount = 0
                                newFile = 1
                                fileCount += 1
                            else :
                                # go to next file
                                newFile = 1
                                fileCount += 1
                    # if events in file equal to eventsPerJobRequested
                    elif ( filesEventCount - jobSkipEventCount == eventsPerJobRequested ) :
                        # close job and touch new file
                        fullString = parString[:-1]
                        if self.useParent==1:
                            fullParentString = pString[:-1]
                            list_of_lists.append([fullString,fullParentString,str(eventsPerJobRequested),str(jobSkipEventCount)])
                        else:
                            list_of_lists.append([fullString,str(eventsPerJobRequested),str(jobSkipEventCount)])
                        msg += "Job %s can run over %s events.\n"%(str(jobCount+1),str(eventsPerJobRequested))
                        jobDestination.append(blockSites[block])
                        msg+= "Job %s Destination: %s\n"%(str(jobCount+1),str(SE2CMS(jobDestination[jobCount])))
                        jobsOfBlock[block].append(jobCount+1)
                        # reset counter
                        jobCount = jobCount + 1
                        totalEventCount = totalEventCount + eventsPerJobRequested
                        eventsRemaining = eventsRemaining - eventsPerJobRequested
                        jobSkipEventCount = 0
                        # reset file
                        pString = ""
                        parString = ""
                        filesEventCount = 0
                        newFile = 1
                        fileCount += 1

                    # if more events in file remain than eventsPerJobRequested
                    else :
                        # close job but don't touch new file
                        fullString = parString[:-1]
                        if self.useParent==1:
                            fullParentString = pString[:-1]
                            list_of_lists.append([fullString,fullParentString,str(eventsPerJobRequested),str(jobSkipEventCount)])
                        else:
                            list_of_lists.append([fullString,str(eventsPerJobRequested),str(jobSkipEventCount)])
                        msg += "Job %s can run over %s events.\n"%(str(jobCount+1),str(eventsPerJobRequested))
                        jobDestination.append(blockSites[block])
                        msg+= "Job %s Destination: %s\n"%(str(jobCount+1),str(SE2CMS(jobDestination[jobCount])))
                        jobsOfBlock[block].append(jobCount+1)
                        # increase counter
                        jobCount = jobCount + 1
                        totalEventCount = totalEventCount + eventsPerJobRequested
                        eventsRemaining = eventsRemaining - eventsPerJobRequested
                        # calculate skip events for last file
                        # use filesEventCount (contains several files), jobSkipEventCount and eventsPerJobRequest
                        jobSkipEventCount = eventsPerJobRequested - (filesEventCount - jobSkipEventCount - self.eventsbyfile[file])
                        # remove all but the last file
                        filesEventCount = self.eventsbyfile[file]
                        pString_tmp=''
                        if self.useParent==1:
                            for f in parent : pString_tmp +=  f + ','
                        pString =  pString_tmp
                        parString =  file + ','
                    pass # END if
                pass # END while (iterate over files in the block)
        pass # END while (iterate over blocks in the dataset)
        common.logger.debug(msg)
        self.ncjobs = self.total_number_of_jobs = jobCount
        if (eventsRemaining > 0 and jobCount < totalNumberOfJobs ):
            common.logger.info("Could not run on all requested events because some blocks not hosted at allowed sites.")
        common.logger.info(str(jobCount)+" job(s) can run on "+str(totalEventCount)+" events.\n")

        # skip check on  block with no sites  DD
        if noBboundary == 0 : self.checkBlockNoSite(blocks,jobsOfBlock,blockSites)

       # prepare dict output
        dictOut = {}
        dictOut['params']= ['InputFiles','MaxEvents','SkipEvents']
        if self.useParent: dictOut['params']= ['InputFiles','ParentFiles','MaxEvents','SkipEvents']
        dictOut['args'] = list_of_lists
        dictOut['jobDestination'] = jobDestination
        dictOut['njobs']=self.total_number_of_jobs

        return dictOut

        # keep trace of block with no sites to print a warning at the end

    def checkBlockNoSite(self,blocks,jobsOfBlock,blockSites):
        # screen output
        screenOutput = "List of jobs and available destination sites:\n\n"
        noSiteBlock = []
        bloskNoSite = []
        allBlock = []

        blockCounter = 0
        for block in blocks:
            if block in jobsOfBlock.keys() :
                blockCounter += 1
                allBlock.append( blockCounter )
                sites=self.blackWhiteListParser.checkWhiteList(self.blackWhiteListParser.checkBlackList(blockSites[block],[block]),[block])
                screenOutput += "Block %5i: jobs %20s: sites: %s\n" % (blockCounter,spanRanges(jobsOfBlock[block]),
                    ', '.join(SE2CMS(sites)))
                if len(sites) == 0:
                    noSiteBlock.append( spanRanges(jobsOfBlock[block]) )
                    bloskNoSite.append( blockCounter )

        common.logger.info(screenOutput)
        if len(noSiteBlock) > 0 and len(bloskNoSite) > 0:
            msg = 'WARNING: No sites are hosting any part of data for block:\n                '
            virgola = ""
            if len(bloskNoSite) > 1:
                virgola = ","
            for block in bloskNoSite:
                msg += ' ' + str(block) + virgola
            msg += '\n\t\tRelated jobs:\n                 '
            virgola = ""
            if len(noSiteBlock) > 1:
                virgola = ","
            for range_jobs in noSiteBlock:
                msg += str(range_jobs) + virgola
            msg += '\n\t\twill not be submitted and this block of data can not be analyzed!\n'
            if self.cfg_params.has_key('GRID.se_white_list'):
                msg += '\tWARNING: SE White List: '+self.cfg_params['GRID.se_white_list']+'\n'
                msg += '\t(Hint: By whitelisting you force the job to run at this particular site(s).\n'
                msg += '\tPlease check if the dataset is available at this site!)'
            if self.cfg_params.has_key('GRID.ce_white_list'):
                msg += '\tWARNING: CE White List: '+self.cfg_params['GRID.ce_white_list']+'\n'
                msg += '\t(Hint: By whitelisting you force the job to run at this particular site(s).\n'
                msg += '\tPlease check if the dataset is available at this site!)\n'

            common.logger.info(msg)

        if bloskNoSite == allBlock:
            raise CrabException('No jobs created')

        return


########################################################################
    def jobSplittingByRun(self):
        """
        """

        self.checkUserSettings()
        blockSites = self.args['blockSites']
        pubdata = self.args['pubdata']

        if self.selectNumberOfJobs == 0 :
            self.theNumberOfJobs = 9999999
        blocks = {}
        runList = []
        thefiles = Fileset(name='FilesToSplit')
        fileList = pubdata.getListFiles()
        for f in fileList:
            block = f['Block']['Name']
            try:
                f['Block']['StorageElementList'].extend(blockSites[block])
            except:
                continue
            wmbsFile = File(f['LogicalFileName'])
            [ wmbsFile['locations'].add(x) for x in blockSites[block] ]
            wmbsFile['block'] = block
            runNum = f['RunsList'][0]['RunNumber']
            runList.append(runNum)
            myRun = Run(runNumber=runNum)
            wmbsFile.addRun( myRun )
            thefiles.addFile(
                wmbsFile
                )

        work = Workflow()
        subs = Subscription(
        fileset = thefiles,
        workflow = work,
        split_algo = 'RunBased',
        type = "Processing")
        splitter = SplitterFactory()
        jobfactory = splitter(subs)

        #loop over all runs
        set = Set(runList)
        list_of_lists = []
        jobDestination = []
        count = 0
        for jobGroup in  jobfactory():
            if count <  self.theNumberOfJobs:
                res = self.getJobInfo(jobGroup)
                parString = ''
                for file in res['lfns']:
                    parString += file + ','
                fullString = parString[:-1]
                list_of_lists.append([fullString,str(-1),str(0)])
                #need to check single file location
                jobDestination.append(res['locations'])
                count +=1
       # prepare dict output
        dictOut = {}
        dictOut['params']= ['InputFiles','MaxEvents','SkipEvents']
        dictOut['args'] = list_of_lists
        dictOut['jobDestination'] = jobDestination
        dictOut['njobs']=count

        return dictOut

    def getJobInfo( self,jobGroup ):
        res = {}
        lfns = []
        locations = []
        tmp_check=0
        for job in jobGroup.jobs:
            for file in job.getFiles():
                lfns.append(file['lfn'])
                for loc in file['locations']:
                    if tmp_check < 1 :
                        locations.append(loc)
                tmp_check = tmp_check + 1
                ### qui va messo il check per la locations
        res['lfns'] = lfns
        res['locations'] = locations
        return res

########################################################################
    def prepareSplittingNoInput(self):
        """
        """
        if (self.selectEventsPerJob):
            common.logger.info('Required '+str(self.eventsPerJob)+' events per job ')
        if (self.selectNumberOfJobs):
            common.logger.info('Required '+str(self.theNumberOfJobs)+' jobs in total ')
        if (self.selectTotalNumberEvents):
            common.logger.info('Required '+str(self.total_number_of_events)+' events in total ')

        if (self.total_number_of_events < 0):
            msg='Cannot split jobs per Events with "-1" as total number of events'
            raise CrabException(msg)

        if (self.selectEventsPerJob):
            if (self.selectTotalNumberEvents):
                self.total_number_of_jobs = int(self.total_number_of_events/self.eventsPerJob)
            elif(self.selectNumberOfJobs) :
                self.total_number_of_jobs =self.theNumberOfJobs
                self.total_number_of_events =int(self.theNumberOfJobs*self.eventsPerJob)

        elif (self.selectNumberOfJobs) :
            self.total_number_of_jobs = self.theNumberOfJobs
            self.eventsPerJob = int(self.total_number_of_events/self.total_number_of_jobs)


    def jobSplittingNoInput(self):
        """
        Perform job splitting based on number of event per job
        """
        common.logger.debug('Splitting per events')
        self.checkUserSettings()
        jobDestination=[]
        if ( (self.selectTotalNumberEvents + self.selectEventsPerJob + self.selectNumberOfJobs) != 2 ):
            msg = 'Must define exactly two of total_number_of_events, events_per_job, or number_of_jobs.'
            raise CrabException(msg)

        managedGenerators =self.args['managedGenerators']
        generator = self.args['generator']
        firstRun = self.cfg_params.get('CMSSW.first_run',None)

        self.prepareSplittingNoInput()

        common.logger.debug('N jobs  '+str(self.total_number_of_jobs))

        # is there any remainder?
        check = int(self.total_number_of_events) - (int(self.total_number_of_jobs)*self.eventsPerJob)

        common.logger.debug('Check  '+str(check))

        common.logger.info(str(self.total_number_of_jobs)+' jobs can be created, each for '+str(self.eventsPerJob)+' for a total of '+str(self.total_number_of_jobs*self.eventsPerJob)+' events')
        if check > 0:
            common.logger.info('Warning: asked '+str(self.total_number_of_events)+' but can do only '+str(int(self.total_number_of_jobs)*self.eventsPerJob))

        # argument is seed number.$i
        self.list_of_args = []
        for i in range(self.total_number_of_jobs):
            ## Since there is no input, any site is good
            jobDestination.append([""]) #must be empty to write correctly the xml
            args=[]
            if (firstRun):
                ## pythia first run
                args.append(str(firstRun)+str(i))
            if (generator in managedGenerators):
               args.append(generator)
               if (generator == 'comphep' and i == 0):
                    # COMPHEP is brain-dead and wants event #'s like 1,100,200,300
                    args.append('1')
               else:
                    args.append(str(i*self.eventsPerJob))
            args.append(str(self.eventsPerJob))
            self.list_of_args.append(args)
       # prepare dict output

        dictOut = {}
        dictOut['params'] = ['MaxEvents']
        if (firstRun):
            dictOut['params'] = ['FirstRun','MaxEvents']
            if ( generator in managedGenerators ) :
                dictOut['params'] = ['FirstRun', 'Generator', 'FirstEvent', 'MaxEvents']
        else:
            if (generator in managedGenerators) :
                dictOut['params'] = ['Generator', 'FirstEvent', 'MaxEvents']
        dictOut['args'] = self.list_of_args
        dictOut['jobDestination'] = jobDestination
        dictOut['njobs']=self.total_number_of_jobs

        return dictOut


    def jobSplittingForScript(self):
        """
        Perform job splitting based on number of job
        """
        self.checkUserSettings()
        if (self.selectNumberOfJobs == 0):
            msg = 'must specify  number_of_jobs.'
            raise crabexception(msg)
        jobDestination = []
        common.logger.debug('Splitting per job')
        common.logger.info('Required '+str(self.theNumberOfJobs)+' jobs in total ')

#        self.total_number_of_jobs = self.theNumberOfJobs

        self.prepareSplittingNoInput()

        common.logger.debug('N jobs  '+str(self.total_number_of_jobs))

        common.logger.info(str(self.total_number_of_jobs)+' jobs can be created')

        # argument is seed number.$i
        self.list_of_args = []
        for i in range(self.total_number_of_jobs):
            args=[]
            jobDestination.append([""])
            if self.eventsPerJob != 0 :
                args.append(str(self.eventsPerJob))
                self.list_of_args.append(args)

       # prepare dict output
        dictOut = {}
        dictOut['params'] = ['MaxEvents']
        dictOut['args'] =  self.list_of_args
        dictOut['jobDestination'] = jobDestination
        dictOut['njobs']=self.total_number_of_jobs
        return dictOut


    def jobSplittingByLumi(self):
        """
        Split task into jobs by Lumi section paying attention to which
        lumis should be run (according to the analysis dataset).
        This uses WMBS job splitting which does not split files over jobs
        so the job will have AT LEAST as many lumis as requested, perhaps
        more
        """

        common.logger.debug('Splitting by Lumi')
        self.checkLumiSettings()

        blockSites = self.args['blockSites']
        pubdata = self.args['pubdata']

        lumisPerFile  = pubdata.getLumis()

        # Make the list of WMBS files for job splitter
        fileList = pubdata.getListFiles()
        thefiles = Fileset(name='FilesToSplit')
        for jobFile in fileList:
            block = jobFile['Block']['Name']
            try:
                jobFile['Block']['StorageElementList'].extend(blockSites[block])
            except:
                continue
            wmbsFile = File(jobFile['LogicalFileName'])
            [ wmbsFile['locations'].add(x) for x in blockSites[block] ]
            wmbsFile['block'] = block
            for lumi in lumisPerFile[jobFile['LogicalFileName']]:
                wmbsFile.addRun(Run(lumi[0], lumi[1]))
            thefiles.addFile(wmbsFile)

        # Create the factory and workflow
        work = Workflow()
        subs = Subscription(fileset    = thefiles,    workflow = work,
                            split_algo = 'LumiBased', type     = "Processing")
        splitter = SplitterFactory()
        jobFactory = splitter(subs)

        list_of_lists = []
        jobDestination = []
        jobCount = 0
        lumisCreated = 0

        if not self.limitJobLumis:
            self.lumisPerJob = pubdata.getMaxLumis() // self.theNumberOfJobs + 1
            common.logger.info('Each job will process about %s lumis.' %
                                self.lumisPerJob)

        for jobGroup in  jobFactory(lumis_per_job = self.lumisPerJob):
            for job in jobGroup.jobs:
                if (self.limitNJobs and jobCount >= self.theNumberOfJobs):
                    common.logger.info('Limit on number of jobs reached.')
                    break
                if (self.limitTotalLumis and lumisCreated >= self.totalNLumis):
                    common.logger.info('Limit on number of lumis reached.')
                    break
                lumis = []
                lfns  = []
                locations = []
                firstFile = True
                # Collect information from all the files
                for jobFile in job.getFiles():
                    if firstFile:  # Get locations from first file in the job
                        for loc in jobFile['locations']:
                            locations.append(loc)
                        firstFile = False
                    # Accumulate Lumis from all files
                    for lumiList in jobFile['runs']:
                        theRun = lumiList.run
                        for theLumi in list(lumiList):
                            lumis.append( (theRun, theLumi) )

                    lfns.append(jobFile['lfn'])
                fileString = ','.join(lfns)
                lumiString = compressLumiString(lumis)
                list_of_lists.append([fileString, str(-1), str(0), lumiString])

                jobDestination.append(locations)
                jobCount += 1
                lumisCreated += len(lumis)
                common.logger.debug('Job %s will run on %s files and %s lumis '
                    % (jobCount, len(lfns), len(lumis) ))

        common.logger.info('%s jobs created to run on %s lumis' %
                              (jobCount, lumisCreated))

        # Prepare dict output matching back to non-WMBS job creation
        dictOut = {}
        dictOut['params'] = ['InputFiles', 'MaxEvents', 'SkipEvents', 'Lumis']
        dictOut['args'] = list_of_lists
        dictOut['jobDestination'] = jobDestination
        dictOut['njobs'] = jobCount

        return dictOut


    def Algos(self):
        """
        Define key splittingType matrix
        """
        SplitAlogs = {
                     'EventBased'           : self.jobSplittingByEvent,
                     'RunBased'             : self.jobSplittingByRun,
                     'LumiBased'            : self.jobSplittingByLumi,
                     'NoInput'              : self.jobSplittingNoInput,
                     'ForScript'            : self.jobSplittingForScript
                     }
        return SplitAlogs


def compressLumiString(lumis):
    """
    Turn a list of 2-tuples of run/lumi numbers into a list of the format
    R1:L1,R2:L2-R3:L3 which is acceptable to CMSSW LumiBlockRange variable
    """

    lumis.sort()
    parts = []
    startRange = None
    endRange = None

    for lumiBlock in lumis:
        if not startRange: # This is the first one
            startRange = lumiBlock
            endRange = lumiBlock
        elif lumiBlock == endRange: # Same Lumi (different files?)
            pass
        elif lumiBlock[0] == endRange[0] and lumiBlock[1] == endRange[1] + 1: # This is a continuation
            endRange = lumiBlock
        else: # This is the start of a new range
            part = ':'.join(map(str, startRange))
            if startRange != endRange:
                part += '-' + ':'.join(map(str, endRange))
            parts.append(part)
            startRange = lumiBlock
            endRange = lumiBlock

    # Put out what's left
    if startRange:
        part = ':'.join(map(str, startRange))
        if startRange != endRange:
            part += '-' + ':'.join(map(str, endRange))
        parts.append(part)

    output = ','.join(parts)
    return output


Revision:	1.27
Committed:	Thu Jul 30 18:45:44 2009 UTC (15 years, 9 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
Changes since 1.26:	+87 -30 lines
Log Message:	Add config parameters for lumi-based splitting