wmglobalqueue 2.4.5.1__py3-none-any.whl

WMCore/WorkQueue/Policy/Start/Dataset.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python
+"""
+WorkQueue splitting by dataset, it creates one single
+workqueue element per dataset.
+
+This policy is specifically used by DQMHarvest workflows, which requires
+some special handling based on run information. Nonetheless, trying to
+make it generic enough that could be used by other spec types.
+"""
+
+import logging
+from math import ceil
+from WMCore import Lexicon
+from WMCore.WorkQueue.Policy.Start.StartPolicyInterface import StartPolicyInterface
+from WMCore.WorkQueue.WorkQueueExceptions import WorkQueueWMSpecError
+from WMCore.WorkQueue.WorkQueueUtils import makeLocationsList
+
+
+class Dataset(StartPolicyInterface):
+    """Split elements into datasets"""
+
+    def __init__(self, **args):
+        StartPolicyInterface.__init__(self, **args)
+        self.args.setdefault('SliceType', 'NumberOfRuns')
+        self.args.setdefault('SliceSize', 1)
+        self.lumiType = "NumberOfLumis"
+        self.sites = []
+
+    def split(self):
+        """Apply policy to spec"""
+        work = set() if self.args['SliceType'] == 'NumberOfRuns' else 0
+        numFiles = 0
+        numEvents = 0
+        numLumis = 0
+        datasetPath = self.initialTask.getInputDatasetPath()
+
+        # dataset splitting can't have its data selection overridden
+        if self.data and list(self.data) != [datasetPath]:
+            raise RuntimeError("Can't provide different data to split with")
+
+        blocks = self.validBlocks(self.initialTask, self.dbs())
+        if not blocks:
+            return
+
+        for block in blocks:
+            if self.args['SliceType'] == 'NumberOfRuns':
+                work = work.union(block[self.args['SliceType']])
+            else:
+                work += float(block[self.args['SliceType']])
+            numLumis += int(block[self.lumiType])
+            numFiles += int(block['NumberOfFiles'])
+            numEvents += int(block['NumberOfEvents'])
+
+        if self.args['SliceType'] == 'NumberOfRuns':
+            numJobs = ceil(len(work) / float(self.args['SliceSize']))
+        else:
+            numJobs = ceil(float(work) / float(self.args['SliceSize']))
+
+        # parentage
+        parentFlag = True if self.initialTask.parentProcessingFlag() else False
+
+        self.newQueueElement(Inputs={datasetPath: self.data.get(datasetPath, [])},
+                             ParentFlag=parentFlag,
+                             NumberOfLumis=numLumis,
+                             NumberOfFiles=numFiles,
+                             NumberOfEvents=numEvents,
+                             Jobs=numJobs,
+                             NoInputUpdate=self.initialTask.getTrustSitelists().get('trustlists'),
+                             NoPileupUpdate=self.initialTask.getTrustSitelists().get('trustPUlists')
+                            )
+
+    def validate(self):
+        """Check args and spec work with block splitting"""
+        StartPolicyInterface.validateCommon(self)
+        if not self.initialTask.inputDataset():
+            raise WorkQueueWMSpecError(self.wmspec, 'No input dataset')
+
+    def validBlocks(self, task, dbs):
+        """Return blocks that pass the input data restriction"""
+        datasetPath = task.getInputDatasetPath()
+        Lexicon.dataset(datasetPath)  # check dataset name
+        validBlocks = []
+
+        blockWhiteList = task.inputBlockWhitelist()
+        blockBlackList = task.inputBlockBlacklist()
+        runWhiteList = task.inputRunWhitelist()
+        runBlackList = task.inputRunBlacklist()
+        lumiMask = task.getLumiMask()
+        if lumiMask:
+            maskedBlocks = self.getMaskedBlocks(task, dbs, datasetPath)
+
+        for blockName in dbs.listFileBlocks(datasetPath):
+            # check block restrictions
+            if blockWhiteList and blockName not in blockWhiteList:
+                continue
+            if blockName in blockBlackList:
+                continue
+
+            blockSummary = dbs.getDBSSummaryInfo(block=blockName)
+            if int(blockSummary.get('NumberOfFiles', 0)) == 0:
+                logging.warning("Block %s being rejected for lack of valid files to process", blockName)
+                self.badWork.append(blockName)
+                continue
+
+            if self.args['SliceType'] == 'NumberOfRuns':
+                blockSummary['NumberOfRuns'] = dbs.listRuns(block=blockName)
+
+            # check lumi restrictions
+            if lumiMask:
+                if blockName not in maskedBlocks:
+                    logging.warning("Block %s doesn't pass the lumi mask constraints", blockName)
+                    self.rejectedWork.append(blockName)
+                    continue
+
+                acceptedLumiCount = sum([len(maskedBlocks[blockName][lfn].getLumis()) for lfn in maskedBlocks[blockName]])
+                ratioAccepted = 1. * acceptedLumiCount / float(blockSummary['NumberOfLumis'])
+                maskedRuns = [maskedBlocks[blockName][lfn].getRuns() for lfn in maskedBlocks[blockName]]
+                acceptedRuns = set(lumiMask.getRuns()).intersection(set().union(*maskedRuns))
+
+                blockSummary['NumberOfFiles'] = len(maskedBlocks[blockName])
+                blockSummary['NumberOfEvents'] = float(blockSummary['NumberOfEvents']) * ratioAccepted
+                blockSummary[self.lumiType] = acceptedLumiCount
+                blockSummary['NumberOfRuns'] = acceptedRuns
+            # check run restrictions
+            elif runWhiteList or runBlackList:
+                runs = set(dbs.listRuns(block=blockName))
+                # multi run blocks need special account, requires more DBS calls
+                recalculateLumiCounts = True if len(runs) > 1 else False
+
+                # apply blacklist and whitelist
+                runs = runs.difference(runBlackList)
+                if runWhiteList:
+                    runs = runs.intersection(runWhiteList)
+                # any runs left are ones we will run on, if none ignore block
+                if not runs:
+                    logging.warning("Block %s doesn't pass the runs constraints", blockName)
+                    self.rejectedWork.append(blockName)
+                    continue
+
+                if recalculateLumiCounts:
+                    # Recalculate the number of files, lumis and ~events accepted
+                    acceptedLumiCount = 0
+                    acceptedEventCount = 0
+                    acceptedFileCount = 0
+                    fileInfo = dbs.listFilesInBlock(fileBlockName=blockName)
+
+                    for fileEntry in fileInfo:
+                        acceptedFile = False
+                        for lumiInfo in fileEntry['LumiList']:
+                            if lumiInfo['RunNumber'] in runs:
+                                acceptedFile = True
+                                acceptedLumiCount += len(lumiInfo['LumiSectionNumber'])
+                        if acceptedFile:
+                            acceptedFileCount += 1
+                            acceptedEventCount += fileEntry['NumberOfEvents']
+
+                else:
+                    acceptedLumiCount = blockSummary["NumberOfLumis"]
+                    acceptedFileCount = blockSummary['NumberOfFiles']
+                    acceptedEventCount = blockSummary['NumberOfEvents']
+
+                blockSummary[self.lumiType] = acceptedLumiCount
+                blockSummary['NumberOfFiles'] = acceptedFileCount
+                blockSummary['NumberOfEvents'] = acceptedEventCount
+                blockSummary['NumberOfRuns'] = runs
+
+            validBlocks.append(blockSummary)
+
+        # all needed blocks present at these sites
+        if task.getTrustSitelists().get('trustlists'):
+            siteWhitelist = task.siteWhitelist()
+            siteBlacklist = task.siteBlacklist()
+            self.sites = makeLocationsList(siteWhitelist, siteBlacklist)
+            self.data[datasetPath] = self.sites
+        else:
+            # as the container is dealt as a whole, resolve the container location
+            locations = set(self.blockLocationRucioPhedex(datasetPath))
+            self.data[datasetPath] = list(set(self.cric.PNNstoPSNs(locations)))
+
+        return validBlocks

WMCore/WorkQueue/Policy/Start/MonteCarlo.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python
+"""
+WorkQueue splitting by block
+
+"""
+from __future__ import division
+
+import os
+from WMCore.WorkQueue.Policy.Start.StartPolicyInterface import StartPolicyInterface
+from WMCore.WorkQueue.WorkQueueExceptions import WorkQueueWMSpecError, WorkQueueNoWorkError
+from WMCore.DataStructs.Mask import Mask
+from copy import copy
+from math import ceil, floor
+
+__all__ = []
+
+
+class MonteCarlo(StartPolicyInterface):
+    """Split elements into blocks"""
+    def __init__(self, **args):
+        StartPolicyInterface.__init__(self, **args)
+
+
+    def split(self):
+        """Apply policy to spec"""
+        # if not specified take standard defaults
+        self.args.setdefault('SliceType', 'NumberOfEvents')
+        self.args.setdefault('SliceSize', 1000)  # events per job
+        self.args.setdefault('SubSliceType', 'NumberOfEventsPerLumi')
+        self.args.setdefault('SubSliceSize', self.args['SliceSize'])  # events per lumi
+        self.args.setdefault('MaxJobsPerElement', 1000)  # jobs per WQE
+        self.args.setdefault('MaxLumisPerElement', os.environ.get('MAX_LUMIS_PER_WQE'))
+        self.args.setdefault('blowupFactor', 1.0)  # Estimate of additional jobs following tasks.
+                                                  # Total WQE tasks will be Jobs*(1+blowupFactor)
+        noInputUpdate = self.initialTask.getTrustSitelists().get('trustlists')
+        noPileupUpdate = self.initialTask.getTrustSitelists().get('trustPUlists')
+
+        if not self.mask:
+            self.mask = Mask(FirstRun=1,
+                             FirstLumi=self.initialTask.getFirstLumi(),
+                             FirstEvent=self.initialTask.getFirstEvent(),
+                             LastRun=1,
+                             LastEvent=self.initialTask.getFirstEvent() +
+                                             self.initialTask.totalEvents() - 1)
+        mask = Mask(**self.mask)
+
+        # First let's initialize some parameters
+        lumis_per_job = ceil(self.args['SliceSize'] / self.args['SubSliceSize'])
+        totalLumisPerElement = int(self.args['MaxJobsPerElement']) * lumis_per_job
+
+        if self.args['MaxLumisPerElement'] and totalLumisPerElement > int(self.args['MaxLumisPerElement']):
+            # If there are too many lumis in the WQ element. reduce the number of jobs per element
+            self.args['MaxJobsPerElement'] = floor(int(self.args['MaxLumisPerElement']) / lumis_per_job)
+            if self.args['MaxJobsPerElement'] == 0:
+                raise WorkQueueWMSpecError(self.wmspec, """Too many lumis (%s) in a job:
+                                           Change 'SliceSize' / 'SubSliceSize'""" % lumis_per_job)
+
+        stepSize = int(self.args['SliceSize']) * int(self.args['MaxJobsPerElement'])
+        total = mask['LastEvent'] - mask['FirstEvent'] + 1
+        lastAllowedEvent = mask['LastEvent']
+        eventsAccounted = 0
+
+        while eventsAccounted < total:
+            current = mask['FirstEvent'] + stepSize - 1  # inclusive range
+            if current > lastAllowedEvent:
+                current = lastAllowedEvent
+            mask['LastEvent'] = current
+
+            # Calculate the job splitting without actually doing it
+            # number of lumis is calculated by events number and SubSliceSize which is events per lumi
+            # So if there no exact division between events per job and events per lumi
+            # it takes the ceiling of the value.
+            # Therefore total lumis can't be calculated from total events / SubSliceSize
+            # It has to be caluated by adding the lumis_per_job * number of jobs
+            nEvents = mask['LastEvent'] - mask['FirstEvent'] + 1
+            nLumis = floor(nEvents / self.args['SliceSize']) * lumis_per_job
+            remainingLumis = ceil(nEvents % self.args['SliceSize'] / self.args['SubSliceSize'])
+            nLumis += remainingLumis
+            jobs = ceil(nEvents / self.args['SliceSize'])
+
+            mask['LastLumi'] = mask['FirstLumi'] + int(nLumis) - 1  # inclusive range
+            self.newQueueElement(WMSpec=self.wmspec,
+                                 NumberOfLumis=nLumis,
+                                 NumberOfEvents=nEvents,
+                                 Jobs=jobs,
+                                 Mask=copy(mask),
+                                 NoInputUpdate=noInputUpdate,
+                                 NoPileupUpdate=noPileupUpdate,
+                                 blowupFactor=self.args['blowupFactor'])
+
+
+            if mask['LastEvent'] > (2 ** 32 - 1):
+                # This is getting tricky, to ensure consecutive
+                # events numbers we must calculate where the jobSplitter
+                # will restart the firstEvent to 1 for the last time
+                # in the newly created unit
+                internalEvents = mask['FirstEvent']
+                accumulatedEvents = internalEvents
+                breakPoint = internalEvents
+
+                while accumulatedEvents < mask['LastEvent']:
+                    if (internalEvents + self.args['SliceSize'] - 1) > (2 ** 32 - 1):
+                        internalEvents = 1
+                        breakPoint = accumulatedEvents
+                    else:
+                        internalEvents += self.args['SliceSize']
+                        accumulatedEvents += self.args['SliceSize']
+
+                leftoverEvents = mask['LastEvent'] - breakPoint + 1
+                mask['FirstEvent'] = leftoverEvents + 1
+
+            else:
+                mask['FirstEvent'] = mask['LastEvent'] + 1
+
+            mask['FirstLumi'] = mask['LastLumi'] + 1
+            eventsAccounted += stepSize
+            lastAllowedEvent = (total - eventsAccounted) + mask['FirstEvent'] - 1
+
+
+    def validate(self):
+        """Check args and spec work with block splitting"""
+        StartPolicyInterface.validateCommon(self)
+
+        if self.initialTask.totalEvents() < 1:
+            raise WorkQueueNoWorkError(self.wmspec, 'Invalid total events selection: %s' % str(self.initialTask.totalEvents()))
+
+        if self.mask and self.mask['LastEvent'] < self.mask['FirstEvent']:
+            raise WorkQueueWMSpecError(self.wmspec, "Invalid start & end events")
+
+        if self.mask and self.mask['LastLumi'] < self.mask['FirstLumi']:
+            raise WorkQueueWMSpecError(self.wmspec, "Invalid start & end lumis")

WMCore/WorkQueue/Policy/Start/ResubmitBlock.py

@@ -0,0 +1,171 @@
+#!/usr/bin/env python
+"""
+WorkQueue splitting for Resubmission workflows
+
+In this case we can't be agnostic of the specific splitting algorithm
+for the top level task. Each algorithm may require a different way of generating
+ACDC blocks.
+
+Current implementations:
+
+- SingleChunk: Harvest, ParentlessMergeBySize, MinFileBased, EventAwareLumiBased, LumiBased
+- FixedSizeChunks: FileBased, FixedDelay, MergeBySize, RunBased, SizeBased, SplitFileBased, TwoFileAndEventBased, TwoFileBased
+
+ACDC unsupported:
+
+- WMBSMergeBySize
+- SiblingProcessingBased
+
+"""
+import json
+from math import ceil
+from WMCore.WorkQueue.Policy.Start.StartPolicyInterface import StartPolicyInterface
+from WMCore.WorkQueue.WorkQueueExceptions import WorkQueueWMSpecError
+from WMCore.WorkQueue.WorkQueueUtils import makeLocationsList
+from WMCore.WorkQueue.DataStructs.ACDCBlock import ACDCBlock
+from WMCore.ACDC.DataCollectionService import DataCollectionService
+
+
+class ResubmitBlock(StartPolicyInterface):
+    """Split elements into blocks"""
+
+    def __init__(self, **args):
+        StartPolicyInterface.__init__(self, **args)
+        self.args.setdefault('SliceType', 'NumberOfFiles')
+        self.args.setdefault('SliceSize', 1)
+        self.args.setdefault('SplittingAlgo', 'LumiBased')
+        self.lumiType = "NumberOfLumis"
+
+        # Define how to handle the different splitting algorithms
+        self.algoMapping = {'Harvest': self.singleChunk,
+                            'ParentlessMergeBySize': self.singleChunk,
+                            'MinFileBased': self.singleChunk,
+                            'LumiBased': self.singleChunk,
+                            'EventAwareLumiBased': self.singleChunk,
+                            'EventBased': self.singleChunk}
+        self.unsupportedAlgos = ['WMBSMergeBySize', 'SiblingProcessingBased']
+        self.defaultAlgo = self.fixedSizeChunk
+        self.sites = []
+
+    def split(self):
+        """Apply policy to spec"""
+        # Prepare a site list in case we need it
+        siteWhitelist = self.initialTask.siteWhitelist()
+        siteBlacklist = self.initialTask.siteBlacklist()
+        self.sites = makeLocationsList(siteWhitelist, siteBlacklist)
+
+        for block in self.validBlocks(self.initialTask):
+            parentList = {}
+            parentFlag = False
+            if self.initialTask.parentProcessingFlag():
+                parentFlag = True
+                parentList[block["Name"]] = block['Sites']
+
+            self.newQueueElement(Inputs={block['Name']: block['Sites']},
+                                 ParentFlag=parentFlag,
+                                 ParentData=parentList,
+                                 NumberOfLumis=block[self.lumiType],
+                                 NumberOfFiles=block['NumberOfFiles'],
+                                 NumberOfEvents=block['NumberOfEvents'],
+                                 Jobs=ceil(float(block[self.args['SliceType']]) /
+                                           float(self.args['SliceSize'])),
+                                 ACDC=block['ACDC'],
+                                 NoInputUpdate=self.initialTask.getTrustSitelists().get('trustlists'),
+                                 NoPileupUpdate=self.initialTask.getTrustSitelists().get('trustPUlists')
+                                )
+
+    def validate(self):
+        """Check args and spec work with block splitting"""
+        StartPolicyInterface.validateCommon(self)
+
+    def validBlocks(self, task):
+        """Return blocks that pass the input data restriction according
+           to the splitting algorithm"""
+        validBlocks = []
+
+        acdcInfo = task.getInputACDC()
+        if not acdcInfo:
+            raise WorkQueueWMSpecError(self.wmspec, 'No acdc section for %s' % task.getPathName())
+        acdc = DataCollectionService(acdcInfo["server"], acdcInfo["database"])
+
+        self.logger.info("Policy self.data variable content: %s", self.data)
+        if self.data:
+            acdcBlockSplit = ACDCBlock.splitBlockName(next(iter(self.data)))
+        else:
+            # if self.data is not passed, assume the data is input dataset from the spec
+            acdcBlockSplit = False
+
+        self.logger.info("Using ACDC blockSplit:\n%s", json.dumps(acdcBlockSplit, indent=2))
+        if acdcBlockSplit:
+            block = acdc.getChunkInfo(acdcInfo['collection'],
+                                      acdcBlockSplit['TaskName'],
+                                      acdcBlockSplit['Offset'],
+                                      acdcBlockSplit['NumOfFiles'])
+            dbsBlock = {}
+            dbsBlock['Name'] = next(iter(self.data))
+
+            dbsBlock['NumberOfFiles'] = block['files']
+            dbsBlock['NumberOfEvents'] = block['events']
+            dbsBlock['NumberOfLumis'] = block['lumis']
+            dbsBlock['ACDC'] = acdcInfo
+            if task.getTrustSitelists().get('trustlists'):
+                dbsBlock["Sites"] = self.sites
+            else:
+                dbsBlock["Sites"] = self.cric.PNNstoPSNs(block["locations"])
+            validBlocks.append(dbsBlock)
+        else:
+            if self.args['SplittingAlgo'] in self.unsupportedAlgos:
+                raise WorkQueueWMSpecError(self.wmspec, 'ACDC is not supported for %s' % self.args['SplittingAlgo'])
+            splittingFunc = self.defaultAlgo
+            if self.args['SplittingAlgo'] in self.algoMapping:
+                splittingFunc = self.algoMapping[self.args['SplittingAlgo']]
+            validBlocks = splittingFunc(acdc, acdcInfo, task)
+        self.logger.info("ACDC with the following validBlocks summary:\n%s", json.dumps(validBlocks, indent=2))
+
+        return validBlocks
+
+    def fixedSizeChunk(self, acdc, acdcInfo, task):
+        """Return a set of blocks with a fixed number of ACDC records"""
+        fixedSizeBlocks = []
+        chunkSize = 250
+        acdcBlocks = acdc.chunkFileset(acdcInfo['collection'],
+                                       acdcInfo['fileset'],
+                                       chunkSize)
+        for block in acdcBlocks:
+            dbsBlock = {}
+            dbsBlock['Name'] = ACDCBlock.name(self.wmspec.name(),
+                                              acdcInfo["fileset"],
+                                              block['offset'], block['files'])
+            dbsBlock['NumberOfFiles'] = block['files']
+            dbsBlock['NumberOfEvents'] = block['events']
+            dbsBlock['NumberOfLumis'] = block['lumis']
+            if task.getTrustSitelists().get('trustlists'):
+                dbsBlock["Sites"] = self.sites
+            else:
+                dbsBlock["Sites"] = self.cric.PNNstoPSNs(block["locations"])
+            dbsBlock['ACDC'] = acdcInfo
+            if dbsBlock['NumberOfFiles']:
+                fixedSizeBlocks.append(dbsBlock)
+        return fixedSizeBlocks
+
+    def singleChunk(self, acdc, acdcInfo, task):
+        """Return a single block (inside a list) with all associated ACDC records"""
+        result = []
+        acdcBlock = acdc.singleChunkFileset(acdcInfo['collection'],
+                                            acdcInfo['fileset'])
+        dbsBlock = {}
+        dbsBlock['Name'] = ACDCBlock.name(self.wmspec.name(),
+                                          acdcInfo["fileset"],
+                                          acdcBlock['offset'], acdcBlock['files'])
+        dbsBlock['NumberOfFiles'] = acdcBlock['files']
+        dbsBlock['NumberOfEvents'] = acdcBlock['events']
+        dbsBlock['NumberOfLumis'] = acdcBlock['lumis']
+        if task.getTrustSitelists().get('trustlists'):
+            dbsBlock["Sites"] = self.sites
+        else:
+            dbsBlock["Sites"] = self.cric.PNNstoPSNs(acdcBlock["locations"])
+        dbsBlock['ACDC'] = acdcInfo
+        if dbsBlock['NumberOfFiles']:
+            result.append(dbsBlock)
+
+        return result