<commit_before>import re
import datetime
import ProcessGroups
from threading import Lock
from DBProxyPool import DBProxyPool
from brokerage.SiteMapper import SiteMapper
from dataservice.Setupper import Setupper
from dataservice.Closer import Closer
from dataservice.TaLauncher import TaLauncher


class TaskBuffer:
    """
    task queue
    
    """

    # constructor 
    def __init__(self):
        self.proxyPool = None
        self.lock = Lock()


    # initialize
    def init(self,dbname,dbpass,nDBConnection=10,useTimeout=False):
        # lock
        self.lock.acquire()
        # create Proxy Pool
        if self.proxyPool == None:
            self.proxyPool = DBProxyPool(dbname,dbpass,nDBConnection,useTimeout)
        # release
        self.lock.release()


    # check production role
    def checkProdRole(self,fqans):
        for fqan in fqans:
            # check production role
            match = re.search('/([^/]+)/Role=production',fqan)
            if match != None:
                return True,match.group(1)
        return False,None

                                                        
    # store Jobs into DB
    def storeJobs(self,jobs,user,joinThr=False,forkSetupper=False,fqans=[],hostname='',resetLocInSetupper=False,
                  checkSpecialHandling=True):
        # check quota for priority calculation
        weight         = 0.0
        userJobID      = -1
        userJobsetID   = -1
        userStatus     = True
        priorityOffset = 0
        userVO         = 'atlas'
        userCountry    = None
        nExpressJobs   = 0
        if len(jobs) > 0 and (jobs[0].prodSourceLabel in ['user','panda','ptest','rc_test']) \
               and (not jobs[0].processingType in ['merge','unmerge']):
            # get DB proxy
            proxy = self.proxyPool.getProxy()
            # check quota
            weight = proxy.checkQuota(user)
            # get JobID and status
            userJobID,userJobsetID,userStatus = proxy.getUserParameter(user,jobs[0].jobDefinitionID,jobs[0].jobsetID)
            # get site access
            userSiteAccess = proxy.checkSiteAccess(jobs[0].computingSite,user)
            # check quota for express jobs
            if 'express' in jobs[0].specialHandling:
                expressQuota = proxy.getExpressJobs(user)
                if expressQuota != None and expressQuota['status'] and expressQuota['quota'] > 0:
                    nExpressJobs = expressQuota['quota']
            # release proxy
            self.proxyPool.putProxy(proxy)
            # get site spec
            siteMapper  = SiteMapper(self)
            tmpSiteSpec = siteMapper.getSite(jobs[0].computingSite)
            # check allowed groups
            if userStatus and hasattr(tmpSiteSpec,'allowedgroups') and (not tmpSiteSpec.allowedgroups in ['',None]):
                # set status to False when allowedgroups is defined
                userStatus = False
                # loop over all groups
                for tmpGroup in tmpSiteSpec.allowedgroups.split(','):
                    if tmpGroup == '':
                        continue
                    # loop over all FQANs
                    for tmpFQAN in fqans:
                        if re.search('^%s' % tmpGroup,tmpFQAN) != None:
                            userStatus = True
                            break
                    # escape
                    if userStatus:
                        break
            # get priority offset
            if hasattr(tmpSiteSpec,'priorityoffset') and (not tmpSiteSpec.priorityoffset in ['',None]):
                # loop over all groups
                for tmpGP in tmpSiteSpec.priorityoffset.split(','):
                    if tmpGP == '':
                        continue
                    # get group and offset
                    tmpGroup = tmpGP.split(':')[0]
                    try:
                        tmpOffset = int(tmpGP.split(':')[-1])
                    except:
                        tmpOffset = 0
                    # loop over all FQANs
                    for tmpFQAN in fqans:
                        if re.search('^%s/' % tmpGroup,tmpFQAN) != None:
                            # use the largest offset
                            if tmpOffset > priorityOffset:
                                priorityOffset = tmpOffset
                            break
            # check site access
            if hasattr(tmpSiteSpec,'accesscontrol') and tmpSiteSpec.accesscontrol == 'grouplist':
                if userSiteAccess == {} or userSiteAccess['status'] != 'approved':
                    # user is not allowed
                    userStatus = False
            # set priority offset
            if userStatus:        
                if userSiteAccess.has_key('poffset') and userSiteAccess['poffset'] > priorityOffset: 
                    priorityOffset = userSiteAccess['poffset']
            # extract country group
            for tmpFQAN in fqans:
                match = re.search('^/atlas/([^/]+)/',tmpFQAN)
                if match != None:
                    tmpCountry = match.group(1)
                    # use country code or usatlas
                    if len(tmpCountry) == 2:
                        userCountry = tmpCountry
                        break
                    # usatlas
                    if tmpCountry in ['usatlas']:
                        userCountry = 'us'
                        break
        # check production role
        withProdRole,workingGroup = self.checkProdRole(fqans)
        if withProdRole:
            # check dataset name
            for tmpFile in jobs[-1].Files:
                if tmpFile.type in ['output','log'] and not tmpFile.lfn.startswith('group'):
                    # reset
                    withProdRole,workingGroup = False,None
                    break
        # return if DN is blocked
        if not userStatus:
            return []
        # extract VO
        for tmpFQAN in fqans:
            match = re.search('^/([^/]+)/',tmpFQAN)
            if match != None:
                userVO = match.group(1)
                break
        # get DB proxy
        proxy = self.proxyPool.getProxy()
        # get number of jobs currently in PandaDB
        serNum = 0
        userDefinedWG = False
        validWorkingGroup = False
        usingBuild = False
        if len(jobs) > 0 and (jobs[0].prodSourceLabel in ['user','panda']) \
               and (not jobs[0].processingType in ['merge','unmerge']):
            # set high prioryty for production role
            if withProdRole:
                serNum = 0
                weight = 0.0                
                priorityOffset = 2000
            # check workingGroup
            if not jobs[0].workingGroup in ['',None,'NULL']:
                userDefinedWG = True
                if userSiteAccess != {}:
                    if userSiteAccess['status'] == 'approved' and jobs[0].workingGroup in userSiteAccess['workingGroups']:
                        # valid workingGroup
                        validWorkingGroup = True
            # get nJob
            if userDefinedWG and validWorkingGroup:
                serNum = proxy.getNumberJobsUser(user,workingGroup=jobs[0].workingGroup)
            else:
                serNum = proxy.getNumberJobsUser(user,workingGroup=None)
            # using build for analysis
            if jobs[0].prodSourceLabel == 'panda':
                usingBuild = True
            # reset nJob/weight for HC   
            if jobs[0].processingType in ['hammercloud','gangarobot'] \
                   or jobs[0].processingType.startswith('gangarobot-'):
                serNum = 0
                weight = 0.0
                priorityOffset = 3000                
        # get group job serial number
        groupJobSerialNum = 0
        if len(jobs) > 0 and (jobs[0].prodSourceLabel in ['user','panda']) \
               and (not jobs[0].processingType in ['merge','unmerge']):
            for tmpFile in jobs[-1].Files:
                if tmpFile.type in ['output','log'] and '$GROUPJOBSN' in tmpFile.lfn:<commit_msg>put priority boost only for gangarobot<commit_after>import re
import datetime
import ProcessGroups
from threading import Lock
from DBProxyPool import DBProxyPool
from brokerage.SiteMapper import SiteMapper
from dataservice.Setupper import Setupper
from dataservice.Closer import Closer
from dataservice.TaLauncher import TaLauncher


class TaskBuffer:
    """
    task queue
    
    """

    # constructor 
    def __init__(self):
        self.proxyPool = None
        self.lock = Lock()


    # initialize
    def init(self,dbname,dbpass,nDBConnection=10,useTimeout=False):
        # lock
        self.lock.acquire()
        # create Proxy Pool
        if self.proxyPool == None:
            self.proxyPool = DBProxyPool(dbname,dbpass,nDBConnection,useTimeout)
        # release
        self.lock.release()


    # check production role
    def checkProdRole(self,fqans):
        for fqan in fqans:
            # check production role
            match = re.search('/([^/]+)/Role=production',fqan)
            if match != None:
                return True,match.group(1)
        return False,None

                                                        
    # store Jobs into DB
    def storeJobs(self,jobs,user,joinThr=False,forkSetupper=False,fqans=[],hostname='',resetLocInSetupper=False,
                  checkSpecialHandling=True):
        # check quota for priority calculation
        weight         = 0.0
        userJobID      = -1
        userJobsetID   = -1
        userStatus     = True
        priorityOffset = 0
        userVO         = 'atlas'
        userCountry    = None
        nExpressJobs   = 0
        if len(jobs) > 0 and (jobs[0].prodSourceLabel in ['user','panda','ptest','rc_test']) \
               and (not jobs[0].processingType in ['merge','unmerge']):
            # get DB proxy
            proxy = self.proxyPool.getProxy()
            # check quota
            weight = proxy.checkQuota(user)
            # get JobID and status
            userJobID,userJobsetID,userStatus = proxy.getUserParameter(user,jobs[0].jobDefinitionID,jobs[0].jobsetID)
            # get site access
            userSiteAccess = proxy.checkSiteAccess(jobs[0].computingSite,user)
            # check quota for express jobs
            if 'express' in jobs[0].specialHandling:
                expressQuota = proxy.getExpressJobs(user)
                if expressQuota != None and expressQuota['status'] and expressQuota['quota'] > 0:
                    nExpressJobs = expressQuota['quota']
            # release proxy
            self.proxyPool.putProxy(proxy)
            # get site spec
            siteMapper  = SiteMapper(self)
            tmpSiteSpec = siteMapper.getSite(jobs[0].computingSite)
            # check allowed groups
            if userStatus and hasattr(tmpSiteSpec,'allowedgroups') and (not tmpSiteSpec.allowedgroups in ['',None]):
                # set status to False when allowedgroups is defined
                userStatus = False
                # loop over all groups
                for tmpGroup in tmpSiteSpec.allowedgroups.split(','):
                    if tmpGroup == '':
                        continue
                    # loop over all FQANs
                    for tmpFQAN in fqans:
                        if re.search('^%s' % tmpGroup,tmpFQAN) != None:
                            userStatus = True
                            break
                    # escape
                    if userStatus:
                        break
            # get priority offset
            if hasattr(tmpSiteSpec,'priorityoffset') and (not tmpSiteSpec.priorityoffset in ['',None]):
                # loop over all groups
                for tmpGP in tmpSiteSpec.priorityoffset.split(','):
                    if tmpGP == '':
                        continue
                    # get group and offset
                    tmpGroup = tmpGP.split(':')[0]
                    try:
                        tmpOffset = int(tmpGP.split(':')[-1])
                    except:
                        tmpOffset = 0
                    # loop over all FQANs
                    for tmpFQAN in fqans:
                        if re.search('^%s/' % tmpGroup,tmpFQAN) != None:
                            # use the largest offset
                            if tmpOffset > priorityOffset:
                                priorityOffset = tmpOffset
                            break
            # check site access
            if hasattr(tmpSiteSpec,'accesscontrol') and tmpSiteSpec.accesscontrol == 'grouplist':
                if userSiteAccess == {} or userSiteAccess['status'] != 'approved':
                    # user is not allowed
                    userStatus = False
            # set priority offset
            if userStatus:        
                if userSiteAccess.has_key('poffset') and userSiteAccess['poffset'] > priorityOffset: 
                    priorityOffset = userSiteAccess['poffset']
            # extract country group
            for tmpFQAN in fqans:
                match = re.search('^/atlas/([^/]+)/',tmpFQAN)
                if match != None:
                    tmpCountry = match.group(1)
                    # use country code or usatlas
                    if len(tmpCountry) == 2:
                        userCountry = tmpCountry
                        break
                    # usatlas
                    if tmpCountry in ['usatlas']:
                        userCountry = 'us'
                        break
        # check production role
        withProdRole,workingGroup = self.checkProdRole(fqans)
        if withProdRole:
            # check dataset name
            for tmpFile in jobs[-1].Files:
                if tmpFile.type in ['output','log'] and not tmpFile.lfn.startswith('group'):
                    # reset
                    withProdRole,workingGroup = False,None
                    break
        # return if DN is blocked
        if not userStatus:
            return []
        # extract VO
        for tmpFQAN in fqans:
            match = re.search('^/([^/]+)/',tmpFQAN)
            if match != None:
                userVO = match.group(1)
                break
        # get DB proxy
        proxy = self.proxyPool.getProxy()
        # get number of jobs currently in PandaDB
        serNum = 0
        userDefinedWG = False
        validWorkingGroup = False
        usingBuild = False
        if len(jobs) > 0 and (jobs[0].prodSourceLabel in ['user','panda']) \
               and (not jobs[0].processingType in ['merge','unmerge']):
            # set high prioryty for production role
            if withProdRole:
                serNum = 0
                weight = 0.0                
                priorityOffset = 2000
            # check workingGroup
            if not jobs[0].workingGroup in ['',None,'NULL']:
                userDefinedWG = True
                if userSiteAccess != {}:
                    if userSiteAccess['status'] == 'approved' and jobs[0].workingGroup in userSiteAccess['workingGroups']:
                        # valid workingGroup
                        validWorkingGroup = True
            # get nJob
            if userDefinedWG and validWorkingGroup:
                serNum = proxy.getNumberJobsUser(user,workingGroup=jobs[0].workingGroup)
            else:
                serNum = proxy.getNumberJobsUser(user,workingGroup=None)
            # using build for analysis
            if jobs[0].prodSourceLabel == 'panda':
                usingBuild = True
            # reset nJob/weight for HC   
            if jobs[0].processingType in ['hammercloud','gangarobot'] \
                   or jobs[0].processingType.startswith('gangarobot-'):
                serNum = 0
                weight = 0.0
            if jobs[0].processingType == 'gangarobot':    
                priorityOffset = 3000                
        # get group job serial number
        groupJobSerialNum = 0
        if len(jobs) > 0 and (jobs[0].prodSourceLabel in ['user','panda']) \
               and (not jobs[0].processingType in ['merge','unmerge']):
            for tmpFile in jobs[-1].Files:
                if tmpFile.type in ['output','log'] and '$GROUPJOBSN' in tmpFile.lfn:
{'fraction_non_alphanumeric': 0.08600565179997542, 'fraction_numerical': 0.006696154318712373, 'mean_word_length': 1.4483380959542789, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4875071', 'n_tokens_mistral': 4260, 'n_tokens_neox': 3981, 'n_words': 1051}