Python CouchWorkQueueElement.CouchWorkQueueElement类代码示例

 def testIdSaved(self):
     """Generated id used as db id"""
     ele = CouchWorkQueueElement(self.couch_db, elementParams = {'RequestName' : 'test'})
     ele.save()
     self.couch_db.commit(timestamp = True)
     self.assertTrue(self.couch_db.documentExists(ele.id))
     self.assertEqual(self.couch_db.info()['doc_count'], 1)

    def insertElements(self, units, parent = None):
        """
        Insert element to database

        @param parent is the parent WorkQueueObject these element's belong to.
                                            i.e. a workflow which has been split
        """
        if not units:
            return
        # store spec file separately - assume all elements share same spec
        self.insertWMSpec(units[0]['WMSpec'])
        for unit in units:

            # cast to couch
            if not isinstance(unit, CouchWorkQueueElement):
                unit = CouchWorkQueueElement(self.db, elementParams = dict(unit))

            if parent:
                unit['ParentQueueId'] = parent.id
                unit['TeamName'] = parent['TeamName']
                unit['WMBSUrl'] = parent['WMBSUrl']

            if unit._couch.documentExists(unit.id):
                self.logger.info('Element "%s" already exists, skip insertion.' % unit.id)
                continue
            unit.save()

        unit._couch.commit(all_or_nothing = True)
        return

    def createWork(self, spec, **kwargs):
        """Return the Inbox element for this spec.

        This does not persist it to the database.
        """
        kwargs.update({'WMSpec' : spec,
                       'RequestName' : spec.name(),
                       'StartPolicy' : spec.startPolicyParameters(),
                       'EndPolicy' : spec.endPolicyParameters(),
                       'OpenForNewData' : True
                      })
        unit = CouchWorkQueueElement(self.inbox, elementParams = kwargs)
        unit.id = spec.name()
        return unit

    def fixConflicts(self):
        """Fix elements in conflict

        Each local queue runs this to resolve its conflicts with global,
        resolution propagates up to global.

        Conflicting elements are merged into one element with others deleted.

        This will fail if elements are modified during the resolution -
        if this happens rerun.
        """
        ordered_states = ['Available', 'Negotiating', 'Acquired', 'Running',
                          'Done', 'Failed', 'CancelRequested', 'Canceled']
        allowed_keys = ['Status', 'EventsWritten', 'FilesProcessed', 'PercentComplete', 'PercentSuccess']
        for db in [self.inbox, self.db]:
            conflicts = db.loadView('WorkQueue', 'conflicts')
            queue = []
            for row in conflicts['rows']:
                previous_value = None
                element_id = row['id']
                for rev in row['value']: # loop over conflicting revisions
                    ele = CouchWorkQueueElement.fromDocument(db, db.document(element_id, rev))
                    if not previous_value: # 1st will contain merged result and become winner
                        previous_value = ele
                        continue
    
                    for key in previous_value:
                        if previous_value[key] == ele.get(key):
                            continue
                        # we need to merge: Take elements from both that seem most advanced, e.g. status & progress stats
                        if key not in allowed_keys:
                            msg = 'Unable to merge conflicting element keys: field "%s" value 1 "%s" value2 "%s"'
                            raise RuntimeError, msg % (key, previous_value.get(key), ele.get(key))
                        if key == 'Status':
                            if ordered_states.index(ele[key]) > ordered_states.index(previous_value[key]):
                                previous_value[key] = ele[key]
                        elif ele[key] > previous_value[key]:
                            previous_value[key] = ele[key]
                    # once losing element has been merged - queue for deletion
                    queue.append(ele)
                # conflict resolved - save element and delete losers
                msg = 'Resolving conflict for wf "%s", id "%s": Losing rev(s): %s'
                self.logger.info(msg % (str(previous_value['RequestName']),
                                         str(previous_value.id),
                                         ", ".join([x._document['_rev'] for x in queue])))
                if self.saveElements(previous_value):
                    for i in queue:
                        i.delete() # delete others (if merged value update accepted)
                    self.saveElements(*queue)

 def testIdFromDbImmutable(self):
     """Modifying element id algorithm doesn't change existing id's"""
     ele = CouchWorkQueueElement(self.couch_db, elementParams = {'RequestName' : 'test'})
     ele.save()
     self.couch_db.commit(timestamp = True)
     ele2 = CouchWorkQueueElement(self.couch_db, id = ele.id).load()
     ele2['RequestName'] = 'ThisWouldCauseIdToChange'
     # id should not change
     self.assertEqual(ele.id, ele2.id)
     # save should modify existing element
     ele2.save()
     self.couch_db.commit(timestamp = True)
     self.assertEqual(self.couch_db.info()['doc_count'], 1)

    def availableWork(self, conditions, teams = None, wfs = None):
        """Get work which is available to be run"""
        elements = []
        for site in conditions.keys():
            if not conditions[site] > 0:
                del conditions[site]
        if not conditions:
            return elements, conditions

        options = {}
        options['include_docs'] = True
        options['descending'] = True
        options['resources'] = conditions
        if teams:
            options['teams'] = teams
        if wfs:
            result = []
            for i in xrange(0, len(wfs), 20):
                options['wfs'] = wfs[i:i+20]
                data = self.db.loadList('WorkQueue', 'workRestrictions', 'availableByPriority', options)
                result.extend(json.loads(data))
            # sort final list
            result.sort(key = lambda x: x['WMCore.WorkQueue.DataStructs.WorkQueueElement.WorkQueueElement']['Priority'])
        else:
            result = self.db.loadList('WorkQueue', 'workRestrictions', 'availableByPriority', options)
            result = json.loads(result)
        for i in result:
            element = CouchWorkQueueElement.fromDocument(self.db, i)
            elements.append(element)

            # Remove 1st random site that can run work
            names = conditions.keys()
            random.shuffle(names)
            for site in names:
                if element.passesSiteRestriction(site):
                    slots_left = conditions[site] - element['Jobs']
                    if slots_left > 0:
                        conditions[site] = slots_left
                    else:
                        conditions.pop(site, None)
                    break
            if not conditions:
                break
        return elements, conditions

    def fixConflicts(self):
        """Fix elements in conflict

        Each local queue runs this to resolve its conflicts with global,
        resolution propagates up to global.

        Conflicting elements are merged into one element with others deleted.

        This will fail if elements are modified during the resolution -
        if this happens rerun.
        """
        for db in [self.inbox, self.db]:
            for row in db.loadView('WorkQueue', 'conflicts')['rows']:
                element_id = row['id']
                try:
                    conflicting_elements = [CouchWorkQueueElement.fromDocument(db, db.document(element_id, rev)) \
                                                                                for rev in row['value']]
                    fixed_elements = fixElementConflicts(*conflicting_elements)
                    if self.saveElements(fixed_elements[0]):
                        self.saveElements(*fixed_elements[1:]) # delete others (if merged value update accepted)
                except Exception, ex:
                    self.logger.error("Error resolving conflict for %s: %s" % (element_id, str(ex)))

    def availableWork(self, thresholds, siteJobCounts, teams = None, wfs = None):
        """
        Get work which is available to be run

        Assume thresholds is a dictionary; keys are the site name, values are
        the maximum number of running jobs at that site.

        Assumes site_job_counts is a dictionary-of-dictionaries; keys are the site
        name and task priorities.  The value is the number of jobs running at that
        priority.
        """
        elements = []

        # We used to pre-filter sites, looking to see if there are idle job slots
        # We don't do this anymore, as we may over-allocate
        # jobs to sites if the new jobs have a higher priority.

        # If there are no sites, punt early.
        if not thresholds:
            return elements, thresholds, siteJobCounts

        options = {}
        options['include_docs'] = True
        options['descending'] = True
        options['resources'] = thresholds
        if teams:
            options['teams'] = teams
        if wfs:
            result = []
            for i in xrange(0, len(wfs), 20):
                options['wfs'] = wfs[i:i+20]
                data = self.db.loadList('WorkQueue', 'workRestrictions', 'availableByPriority', options)
                result.extend(json.loads(data))
            # sort final list
            result.sort(key = lambda x: x['WMCore.WorkQueue.DataStructs.WorkQueueElement.WorkQueueElement']['Priority'])
        else:
            result = self.db.loadList('WorkQueue', 'workRestrictions', 'availableByPriority', options)
            result = json.loads(result)

        # Iterate through the results; apply whitelist / blacklist / data
        # locality restrictions.  Only assign jobs if they are high enough
        # priority.
        for i in result:
            element = CouchWorkQueueElement.fromDocument(self.db, i)
            prio = element['Priority']

            possibleSite = None
            sites = thresholds.keys()
            random.shuffle(sites)
            for site in sites:
                if element.passesSiteRestriction(site):
                    # Count the number of jobs currently running of greater priority
                    prio = element['Priority']
                    curJobCount = sum(map(lambda x : x[1] if x[0] >= prio else 0, siteJobCounts.get(site, {}).items()))
                    if curJobCount < thresholds[site]:
                        possibleSite = site
                        break

            if possibleSite:
                elements.append(element)
                if site not in siteJobCounts:
                    siteJobCounts[site] = {}
                siteJobCounts[site][prio] = siteJobCounts[site].setdefault(prio, 0) + element['Jobs']

        # sort elements to get them in timestamp order
        elements = sorted(elements, key=lambda element: element['CreationTime'])
        return elements, thresholds, siteJobCounts

 def getElementsForParent(self, parent):
     """Get elements with the given parent"""
     elements = self.db.loadView("WorkQueue", "elementsByParent", {"key": parent.id, "include_docs": True})
     return [CouchWorkQueueElement.fromDocument(self.db, x["doc"]) for x in elements.get("rows", [])]

 def getElementsForParentData(self, data):
     """Get active elements for this data """
     elements = self.db.loadView('WorkQueue', 'elementsByParentData', {'key' : data, 'include_docs' : True})
     return [CouchWorkQueueElement.fromDocument(self.db,
                                                x['doc'])
             for x in elements.get('rows', [])]

 def getElementsForParent(self, parent):
     """Get elements with the given parent"""
     elements = self.db.loadView('WorkQueue', 'elementsByParent', {'key' : parent.id, 'include_docs' : True})
     return [CouchWorkQueueElement.fromDocument(self.db,
                                                x['doc'])
             for x in elements.get('rows', [])]

 def getElementsForPileupData(self, data):
     """Get active elements for this data """
     elements = self.db.loadView("WorkQueue", "elementsByPileupData", {"key": data, "include_docs": True})
     return [CouchWorkQueueElement.fromDocument(self.db, x["doc"]) for x in elements.get("rows", [])]

    def availableWork(self, thresholds, siteJobCounts, teams = None, wfs = None):
        """
        Get work which is available to be run

        Assume thresholds is a dictionary; keys are the site name, values are
        the maximum number of running jobs at that site.

        Assumes site_job_counts is a dictionary-of-dictionaries; keys are the site
        name and task priorities.  The value is the number of jobs running at that
        priority.
        """
        self.logger.info("Getting available work from %s/%s" % 
                         (sanitizeURL(self.server.url)['url'], self.db.name))
        elements = []

        # We used to pre-filter sites, looking to see if there are idle job slots
        # We don't do this anymore, as we may over-allocate
        # jobs to sites if the new jobs have a higher priority.

        # If there are no sites, punt early.
        if not thresholds:
            self.logger.error("No thresholds is set: Please check")
            return elements, thresholds, siteJobCounts

        options = {}
        options['include_docs'] = True
        options['descending'] = True
        options['resources'] = thresholds
        if teams:
            options['teams'] = teams
            self.logger.info("setting teams %s" % teams)
        if wfs:
            result = []
            for i in xrange(0, len(wfs), 20):
                options['wfs'] = wfs[i:i+20]
                data = self.db.loadList('WorkQueue', 'workRestrictions', 'availableByPriority', options)
                result.extend(json.loads(data))
            # sort final list
            result.sort(key = lambda x: x['WMCore.WorkQueue.DataStructs.WorkQueueElement.WorkQueueElement']['Priority'])
        else:
            result = self.db.loadList('WorkQueue', 'workRestrictions', 'availableByPriority', options)
            result = json.loads(result)
            if len(result) == 0:
                self.logger.info("""No available work in WQ or didn't pass workqueue restriction 
                                    - check Pileup, site white list, etc""")
            self.logger.debug("Available Work:\n %s \n for resources\n %s" % (result, thresholds))
        # Iterate through the results; apply whitelist / blacklist / data
        # locality restrictions.  Only assign jobs if they are high enough
        # priority.
        for i in result:
            element = CouchWorkQueueElement.fromDocument(self.db, i)
            prio = element['Priority']

            possibleSite = None
            sites = thresholds.keys()
            random.shuffle(sites)
            for site in sites:
                if element.passesSiteRestriction(site):
                    # Count the number of jobs currently running of greater priority
                    prio = element['Priority']
                    curJobCount = sum(map(lambda x : x[1] if x[0] >= prio else 0, siteJobCounts.get(site, {}).items()))
                    self.logger.debug("Job Count: %s, site: %s threshods: %s" % (curJobCount, site, thresholds[site]))
                    if curJobCount < thresholds[site]:
                        possibleSite = site
                        break

            if possibleSite:
                self.logger.debug("Possible site exists %s" % str(possibleSite))
                elements.append(element)
                if site not in siteJobCounts:
                    siteJobCounts[site] = {}
                siteJobCounts[site][prio] = siteJobCounts[site].setdefault(prio, 0) + element['Jobs']*element.get('blowupFactor', 1.0)
            else:
                self.logger.info("No possible site for %s" % element['RequestName'])
        # sort elements to get them in priority first and timestamp order
        elements.sort(key=lambda element: element['CreationTime'])
        elements.sort(key = lambda x: x['Priority'], reverse = True)
        
        return elements, thresholds, siteJobCounts

 def getElementsForWorkflow(self, workflow):
     """Get elements for a workflow"""
     elements = self.db.loadView(
         "WorkQueue", "elementsByWorkflow", {"key": workflow, "include_docs": True, "reduce": False}
     )
     return [CouchWorkQueueElement.fromDocument(self.db, x["doc"]) for x in elements.get("rows", [])]