Python WorkQueue.WorkQueue类代码示例

def main():
    """
    _main_
    """
    if 'WMAGENT_CONFIG' not in os.environ:
        os.environ['WMAGENT_CONFIG'] = '/data/srv/wmagent/current/config/wmagent/config.py'

    config = loadConfigurationFile(os.environ["WMAGENT_CONFIG"])

    # Instantiating central reqmgr and local workqueue
    print "ReqMgr2 URL  : %s" % sanitizeURL(config.JobUpdater.reqMgr2Url)['url']
    print "WorkQueue URL: %s and dbname %s" % (sanitizeURL(config.WorkQueueManager.couchurl)['url'],
                                               config.WorkQueueManager.dbname)

    reqmgr2 = ReqMgr(config.JobUpdater.reqMgr2Url)
    workqueue = WorkQueue(config.WorkQueueManager.couchurl, config.WorkQueueManager.dbname)

    print "\nFirst attempt to update prio of wfs that are not in WMBS and only in local queue"
    priorityCache = {}
    workflowsToUpdate = {}
    workflowsToCheck = [x for x in workqueue.getAvailableWorkflows()]
    print "Retrieved %d workflows from workqueue" % len(workflowsToCheck)

    for workflow, priority in workflowsToCheck:
        if workflow not in priorityCache:
            try:
                priorityCache[workflow] = reqmgr2.getRequestByNames(workflow)[workflow]['RequestPriority']
            except Exception, ex:
                print "Couldn't retrieve the priority of request %s" % workflow
                print "Error: %s" % ex
                continue
        if priority != priorityCache[workflow]:
            workflowsToUpdate[workflow] = priorityCache[workflow]

def main():
    # FIXME update the workflow name here
    wf = "mcremone_task_EXO-RunIISummer15wmLHEGS-04802__v1_T_170811_181808_305"
    print("Looking for problematic inbox elements...")

    wq = WorkQueue("https://cmsweb.cern.ch/couchdb/workqueue")
    print("Workqueue config: server %s and db %s" % (wq.server.url, wq.db.name))

    nonCancelableElements = ['Done', 'Canceled', 'Failed']
    data = wq.db.loadView('WorkQueue', 'elementsDetailByWorkflowAndStatus',
                          {'startkey': [wf], 'endkey': [wf, {}], 'reduce': False})

    elements = [x['id'] for x in data.get('rows', []) if x['key'][1] not in nonCancelableElements]
    print("Found %d elements for wf %s" % (len(elements), wf))
    total = 0
    for eleSlice in grouper(elements, 100):
        try:
            wq.updateElements(*eleSlice, Status='CancelRequested')
        except Exception as ex:
            print("Exception happened, but keep going: %s" % str(ex))
        else:
            total += 100
            print("Elements updated: %s" % total)

    print("Done!")

    sys.exit(0)

def getSiteInfoFromGlobalQueue(serviceURL):

    url, dbName = splitCouchServiceURL(serviceURL)
    globalQ = WorkQueue(url, dbName)
    try:
        queues = globalQ.getChildQueues()
    except Exception, ex:
        errorInfo = {}
        errorInfo['site_name'] = serviceURL
        return [errorInfo]

    def testCompletedWorkflow(self):
        # test getWork
        specName = "RerecoSpec"
        specUrl = self.specGenerator.createReRecoSpec(specName, "file")
        globalQ = globalQueue(DbName='workqueue_t',
                              QueueURL=self.testInit.couchUrl,
                              UnittestFlag=True)
        self.assertTrue(globalQ.queueWork(specUrl, specName, "teamA") > 0)

        wqApi = WorkQueueDS(self.testInit.couchUrl, 'workqueue_t')
        # overwrite default - can't test with stale view
        wqApi.defaultOptions = {'reduce': True, 'group': True}
        # This only checks minimum client call not exactly correctness of return
        # values.
        self.assertEqual(wqApi.getTopLevelJobsByRequest(),
                         [{'total_jobs': 339, 'request_name': specName}])
        results = wqApi.getJobsByStatusAndPriority()
        self.assertEqual(results.keys(), ['Available'])
        self.assertEqual(results['Available'].keys(), [8000])
        self.assertTrue(results['Available'][8000]['sum'], 339)
        result = wqApi.getElementsCountAndJobsByWorkflow()
        self.assertEqual(len(result), 1)
        self.assertEqual(result[specName]['Available']['Jobs'], 339)
        data = wqApi.db.loadView('WorkQueue', 'elementsDetailByWorkflowAndStatus',
                                 {'startkey': [specName], 'endkey': [specName, {}],
                                  'reduce': False})
        elements = [x['id'] for x in data.get('rows', [])]
        wqApi.updateElements(*elements, Status='Canceled')
        # load this view once again to make sure it will be updated in the next assert..
        data = wqApi.db.loadView('WorkQueue', 'elementsDetailByWorkflowAndStatus',
                                 {'startkey': [specName], 'endkey': [specName, {}],
                                  'reduce': False})
        self.assertEqual(len(wqApi.getCompletedWorkflow(stale=False)), 1)
        self.assertEqual(wqApi.getJobsByStatusAndPriority().keys(), ['Canceled'])

def getSiteInfoFromLocalQueue(serviceURL):
    """ get agent status from local agent """

    url, dbName = splitCouchServiceURL(serviceURL)
    wqService = WorkQueue(url, dbName)
    try:
        wmbsUrls = wqService.getWMBSUrl()
    except Exception, ex:
        errorInfo = {}
        errorInfo['site_name'] = serviceURL
        return [errorInfo]

def getRequestInfoFromLocalQueue(serviceURL):
    """ get the request info from local queue """

    url, dbName = splitCouchServiceURL(serviceURL)
    service = WorkQueue(url, dbName)
    try:
        wmbsUrls = service.getWMBSUrl()
        jobStatusInfo = service.getJobInjectStatusByRequest()
    except Exception, ex:
        logging.error("%s: %s" % (serviceURL, str(ex)))
        return DFormatter.errorFormatter(serviceURL, "LocalQueue Down")

def getAgentInfoFromGlobalQueue(serviceURL):

    url, dbName = splitCouchServiceURL(serviceURL)
    globalQ = WorkQueue(url, dbName)
    try:
        childQueues = globalQ.getChildQueues()
    except Exception, ex:
        errorInfo = {}
        errorInfo['url'] = serviceURL
        errorInfo['status'] = "Global Queue down: %s" % serviceURL
        errorInfo['acdc'] = 'N/A'
        return [errorInfo]

def getAgentInfoFromLocalQueue(serviceURL):
    """ get agent status from local agent """
    url, dbName = splitCouchServiceURL(serviceURL)
    localQ = WorkQueue(url, dbName)

    try:
        wmbsUrl = localQ.getWMBSUrl()
    except Exception, ex:
        errorInfo = {}
        errorInfo['url'] = serviceURL
        errorInfo['status'] = "Local Queue down: %s" % serviceURL
        errorInfo['acdc'] = 'N/A'
        return errorInfo

    def testWorkQueueService(self):
        # test getWork
        specName = "RerecoSpec"
        specUrl = self.specGenerator.createReRecoSpec(specName, "file",
                                                      assignKwargs={'SiteWhitelist': ['T2_XX_SiteA']})
        globalQ = globalQueue(DbName='workqueue_t',
                              QueueURL=self.testInit.couchUrl,
                              UnittestFlag=True)
        self.assertTrue(globalQ.queueWork(specUrl, specName, "teamA") > 0)

        wqApi = WorkQueueDS(self.testInit.couchUrl, 'workqueue_t')
        # overwrite default - can't test with stale view
        wqApi.defaultOptions = {'reduce': True, 'group': True}
        # This only checks minimum client call not exactly correctness of return
        # values.
        self.assertEqual(wqApi.getTopLevelJobsByRequest(),
                         [{'total_jobs': 339, 'request_name': specName}])
        # work still available, so no childQueue
        results = wqApi.getChildQueuesAndStatus()
        self.assertItemsEqual(set([item['agent_name'] for item in results]), ["AgentNotDefined"])
        result = wqApi.getElementsCountAndJobsByWorkflow()
        self.assertEqual(len(result), 1)
        self.assertEqual(result[specName]['Available']['Jobs'], 339)

        results = wqApi.getChildQueuesAndPriority()
        resultsPrio = set([item['priority'] for item in results if item['agent_name'] == "AgentNotDefined"])
        self.assertItemsEqual(resultsPrio, [8000])
        self.assertEqual(wqApi.getWMBSUrl(), [])
        self.assertEqual(wqApi.getWMBSUrlByRequest(), [])

def getRequestInfoFromGlobalQueue(serviceURL):
    """ get the request info from global queue """
    url, dbName = splitCouchServiceURL(serviceURL)
    service = WorkQueue(url, dbName)
    try:
        jobInfo = service.getTopLevelJobsByRequest()
        qInfo = service.getChildQueuesByRequest()
        siteWhitelists = service.getSiteWhitelistByRequest()
        childQueueURLs = set()
        for item in qInfo:
            childQueueURLs.add(item['local_queue'])

    except Exception, ex:
        logging.error("%s: %s" % (serviceURL, str(ex)))
        return DFormatter.errorFormatter(serviceURL, "GlobalQueue Down")

    def __init__(self, config):
        """
        initialize properties specified from config
        """
        BaseWorkerThread.__init__(self)
        # set the workqueue service for REST call
        self.config = config
        # need to get campaign, user, owner info
        self.agentInfo = initAgentInfo(self.config)
        self.summaryLevel = config.AnalyticsDataCollector.summaryLevel

        proxyArgs = {'logger': logging.getLogger()}
        self.proxy = Proxy(proxyArgs)
        self.proxyFile = self.proxy.getProxyFilename()  # X509_USER_PROXY
        self.userCertFile = self.proxy.getUserCertFilename()  # X509_USER_CERT
        # credential lifetime warning/error thresholds, in days
        self.credThresholds = {'proxy': {'error': 3, 'warning': 5},
                               'certificate': {'error': 10, 'warning': 20}}

        # Monitoring setup
        self.userAMQ = getattr(config.AgentStatusWatcher, "userAMQ", None)
        self.passAMQ = getattr(config.AgentStatusWatcher, "passAMQ", None)
        self.postToAMQ = getattr(config.AgentStatusWatcher, "enableAMQ", False)
        self.topicAMQ = getattr(config.AgentStatusWatcher, "topicAMQ", None)
        self.hostPortAMQ = getattr(config.AgentStatusWatcher, "hostPortAMQ", [('cms-mb.cern.ch', 61313)])

        # T0 doesn't have WorkQueue, so some monitoring/replication code has to be skipped here
        if hasattr(self.config, "Tier0Feeder"):
            self.isT0agent = True
            self.producer = "tier0wmagent"
        else:
            self.isT0agent = False
            self.producer = "wmagent"
            localWQUrl = config.AnalyticsDataCollector.localQueueURL
            self.workqueueDS = WorkQueueDS(localWQUrl)

    def advanceStatus(self, config):
        """
        gather active data statistics
        """

        reqDBWriter = RequestDBWriter(config.reqmgrdb_url)
        gqService = WorkQueue(config.workqueue_url)
        
        self.logger.info("Getting GQ data for status check")
        wfStatusDict = gqService.getWorkflowStatusFromWQE()
        
        self.logger.info("Advancing status")
        moveForwardStatus(reqDBWriter, wfStatusDict, self.logger)
        moveToArchivedForNoJobs(reqDBWriter, wfStatusDict, self.logger)

        return

    def setup(self, parameters):
        """
        set db connection(couchdb, wmbs) to prepare to gather information
        """
        # set the connection to local queue
        if not hasattr(self.config, "Tier0Feeder"):
            self.localQueue = WorkQueueService(self.config.AnalyticsDataCollector.localQueueURL)

        # set the connection for local couchDB call
        self.localCouchDB = LocalCouchDBData(self.config.AnalyticsDataCollector.localCouchURL, 
                                             self.config.JobStateMachine.summaryStatsDBName,
                                             self.summaryLevel)

        # interface to WMBS/BossAir db
        myThread = threading.currentThread()
        # set wmagent db data
        self.wmagentDB = WMAgentDBData(self.summaryLevel, myThread.dbi, myThread.logger)
        # set the connection for local couchDB call
        self.localSummaryCouchDB = WMStatsWriter(self.config.AnalyticsDataCollector.localWMStatsURL, 
                                                 appName="WMStatsAgent")
        
        if hasattr(self.config, "Tier0Feeder"):
            #use local db for tier0
            centralRequestCouchDBURL = self.config.AnalyticsDataCollector.localT0RequestDBURL
        else:
            centralRequestCouchDBURL = self.config.AnalyticsDataCollector.centralRequestDBURL
        
        self.centralRequestCouchDB = RequestDBWriter(centralRequestCouchDBURL, 
                                                   couchapp = self.config.AnalyticsDataCollector.RequestCouchApp)
        #TODO: change the config to hold couch url
        self.localCouchServer = CouchMonitor(self.config.JobStateMachine.couchurl)
        
        if self.pluginName != None:
            pluginFactory = WMFactory("plugins", "WMComponent.AnalyticsDataCollector.Plugins")
            self.plugin = pluginFactory.loadObject(classname = self.pluginName)

    def setup(self, parameters):
        """
        set db connection(couchdb, wmbs) to prepare to gather information
        """

        #
        self.localQueue = WorkQueueService(self.config.AnalyticsDataCollector.localQueueURL)

        # set the connection for local couchDB call
        self.localCouchDB = LocalCouchDBData(self.config.AnalyticsDataCollector.localCouchURL, self.summaryLevel)

        # interface to WMBS/BossAir db
        myThread = threading.currentThread()
        # set wmagent db data
        self.wmagentDB = WMAgentDBData(self.summaryLevel, myThread.dbi, myThread.logger)
        # set the connection for local couchDB call
        self.localSummaryCouchDB = WMStatsWriter(self.config.AnalyticsDataCollector.localWMStatsURL)
        logging.info("Setting the replication to central monitor ...")
        self.localSummaryCouchDB.replicate(self.config.AnalyticsDataCollector.centralWMStatsURL)
        
        self.centralWMStatsCouchDB = WMStatsWriter(self.config.AnalyticsDataCollector.centralWMStatsURL)
        
        if self.pluginName != None:
            pluginFactory = WMFactory("plugins", "WMComponent.AnalyticsDataCollector.Plugins")
            self.plugin = pluginFactory.loadObject(classname = self.pluginName)

 def __init__(self, app, api, config, mount):
     # main CouchDB database where requests/workloads are stored
     RESTEntity.__init__(self, app, api, config, mount)
     self.reqmgr_db = api.db_handler.get_db(config.couch_reqmgr_db)
     self.reqmgr_db_service = RequestDBWriter(self.reqmgr_db, couchapp="ReqMgr")
     # this need for the post validtiaon
     self.reqmgr_aux_db = api.db_handler.get_db(config.couch_reqmgr_aux_db)
     self.gq_service = WorkQueue(config.couch_host, config.couch_workqueue_db)