crabFunctions.py

Go to the documentation of this file.

#!/usr/bin/env python

from __future__ import print_function
import os,sys,glob
import tarfile
import xml.etree.ElementTree as ET
import imp
import json
import optparse
import subprocess
import logging
import datetime
import uuid
import time
from http.client import HTTPException
from multiprocessing import Process, Queue

from CRABAPI.RawCommand import crabCommand
from CRABClient.UserUtilities import getConsoleLogLevel, setConsoleLogLevel
from CRABClient.ClientUtilities import LOGLEVEL_MUTE
from CRABClient.ClientExceptions import CachefileNotFoundException



class CrabController():

    
    def __init__(self, debug=0, logger = None , workingArea = None, voGroup = None, username = None):

        setConsoleLogLevel(LOGLEVEL_MUTE)
        self.debug = debug
        if workingArea is not None:
            self.workingArea = workingArea
        else:
            self.workingArea = os.getcwd()
        self.dry_run = False
        if voGroup is not None:
            self.voGroup = voGroup
        else:
            self.voGroup = "dcms"
        if username is not None:
            self.username = username
        else:
            self.username = None

        if logger is not None:
            self.logger = logger.getChild("CrabController")
        else:
            # add instance logger as logger to root
            self.logger = logging.getLogger("CrabController")
            # check if handlers are present for root logger
            # we assume that the default logging is not configured
            # if handler is present
            if len(logging.getLogger().handlers) < 1 :
                ch = logging.FileHandler('crabController.log', mode='a', encoding=None, delay=False)
                ch.setLevel(logging.DEBUG)
                # create formatter
                formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s' )
                # add formatter to ch
                ch.setFormatter(formatter)
                self.logger.addHandler(ch)

        self.crab_q = Queue()
    
    def checkwrite(self,site='T2_DE_RWTH',path='noPath'):
        if self.username is None: self.checkusername()
        try:
            self.logger.info( "Checking if user can write to /store/user/%s on site %s with voGroup %s"%(self.username,site , self.voGroup) )
            if not 'noPath' in path:
                res = crabCommand('checkwrite','--site',site,'--voGroup',self.voGroup,'--lfn', path)
            else:
                res = crabCommand('checkwrite','--site',site,'--voGroup',self.voGroup)
            if res['status'] == 'SUCCESS':
                self.logger.info("Checkwrite was sucessfully called.")
                return True
            else:
                self.logger.error( "The crab checkwrite command failed for site: %s"%site )
                return False
        except:
            self.logger.error( 'Unable to perform crab checkwrite')
            return False

    
    def submit(self,name):
        if self.dry_run:
            res = self.callCrabCommand(('submit', '--dryrun', name))
            self.logger.info('Dry-run: You may check the created config and sandbox')
        else:
            res = self.callCrabCommand(('submit','--wait' , name))
            self.logger.info("crab sumbit called for task %s"%name)
            if self.debug > 1:
                self.logger.info(str(res))
        return res
    
    def resubmit(self,name,joblist = None):
        if self.dry_run:
            self.logger.info('Dry-run: Created config file. ')
            return {}
        #~ if joblist is not None:
            #~ jobstring ="%s"%','.join(joblist)
            #~ cmd = ('resubmit','--wait', '--jobids=',jobstring, os.path.join(self.workingArea,self._prepareFoldername(name)) )
        if False:
            pass
        else:
            cmd = ('resubmit','--wait', os.path.join(self.workingArea,self._prepareFoldername(name)) )
        res = self.callCrabCommand( cmd )
        self.logger.info("crab resumbit called for task %s"%name)
        return res
    
    def checkusername(self):
        #depreceated string: cmd = 'crab checkHNname --voGroup=dcms'
        #~ cmd = 'crab checkusername --voGroup=dcms'
        try:
            username = os.environ["CERNUSERNAME"]
            return username
        except:pass
        res = crabCommand('checkusername')
        try:
            self.username = res['username']
            return res['username']
        except:
            return "noHNname"

    
    def status(self,name):
        if self.dry_run:
            self.logger.info('Dry-run: Created config file. crab command would have been: %s'%cmd)
        else:
            try:
                if not "crab_" in name:
                    callname = "crab_" + name
                else:
                    callname = name
                res = self.callCrabCommand( ('status', '--long', callname) )
                if 'taskFailureMsg' in res and 'jobs' in res:
                    return res['status'], res['jobs'], res['taskFailureMsg']
                elif 'jobs' in res and 'taskFailureMsg' not in res:
                    return res['status'], res['jobs'],None
                elif 'jobs' not in res and 'taskFailureMsg' in res:
                    return res['status'], {},res['taskFailureMsg']
                else:
                     return res['status'],{},None
            except Exception as e:
                print(e)
                self.logger.error("Can not run crab status request")
                return "NOSTATE",{},None

    
    def callCrabCommand( self, crabArgs ):
        crabCommandProcessArgs = (self.crab_q, crabArgs)
        p = Process(target=crabCommandProcess, args=(crabCommandProcessArgs))
        p.start()
        res = self.crab_q.get()
        p.join()
        return res

    
    def getlog(self, name):
        foldername = self._prepareFoldername( name)
        try:
            #res = crabCommand('--quiet','status', dir = 'crab_%s' % name)
            res = self.callCrabCommand( ('getlog',  '%s' % foldername) )
            return res['success'], res['failed']
        except:
            self.logger.error("Error calling crab getlog for %s" %foldername)
            return {}, {}

    
    def report(self, name):
        foldername = self._prepareFoldername( name)
        try:
            res = self.callCrabCommand( ('report',  '%s' % foldername) )
            return res['analyzedLumis']
        except:
            self.logger.error("Error calling crab report for %s" %foldername)


    
    def readCrabConfig( self, name ):
        try:
            if os.path.exists(name):
                pset = name
            else:
                pset = 'crab_%s_cfg.py' % name
            with open( pset, 'r') as cfgfile:
                cfo = imp.load_source("pycfg", pset, cfgfile )
                config = cfo.config
                del cfo
            return config
        except:
            return False

    
    @property
    def crabFolders(self):
        results = []
        dirlist = [ x for x in os.listdir( self.workingArea ) if (x.startswith('crab_') and os.path.isdir( os.path.join(self.workingArea,x) ) )]
        return dirlist

    
    def _prepareFoldername(self, name):
        if name.startswith("crab_"):
            crabfolder = '%s'%name
        else:
            crabfolder = "crab_%s "%name
        return crabfolder.strip()
    
    def commandlineOptions(self,parser = optparse.OptionParser( 'usage: %prog' )):
        # we first need to call parse_args with a dummy string at the beginning to
        # check for existing options later
        (currentoptions, args ) = parser.parse_args([" "])

        # The following block shows how variables should be added, where
        # conflicts are possible if the option is overridden by this function
        # they raise a value error
        #try:
        #    parser.add_option( '--someOption',metavar='DIR',default=None,
        #                       help='Dummy option for future integration')
        #except OptionConflictError as e:
        #    conditionalLog(crablog,"There are conflicts extending the optparse options object",'error')
        #    conditionalLog(crablog,e.strerror,'error')

        # options where it is checked if they exists and new options are added
        # otherwise
        if not hasattr(currentoptions, 'dry_run'):
            parser.add_option( '--dry-run', action='store_true', default=False,
                            help='Do everything except calling CRAB or registering samples to the database.' )
        if not hasattr(currentoptions, 'workingArea'):
            parser.add_option( '--workingArea',metavar='DIR',default=os.getcwd(),help='The area (full or relative path) where to create the CRAB project directory. '
                                 'If the area doesn\'t exist, CRAB will try to create it using the mkdir command' \
                                ' (without -p option). Defaults to the current working directory.'       )


        # Some options can be added without expected trouble with other parser
        # parts, simply because it is quite fixed what they represent.
        # those variables should be added here and will throw no exception if
        # they already exist in the parser
        #parser.set_conflict_handler('resolve')
        #parser.add_option( '--someOption',metavar='DIR',default=None,
        #                       help='Dummy option for future integration')

        return parser




def crabCommandProcess(q,crabCommandArgs):
    # give crab3 the chance for one server glitch
    i=0
    while True:
        i+=1
        try:
            res = crabCommand(*crabCommandArgs)
            break
        except HTTPException as e:
            print("crab error ---------------")
            print(e)
            print("end error ---------------")
            print("will try again!")
            import time
            time.sleep(5)
        except CachefileNotFoundException as e:
            print("crab error ---------------")
            print(e)
            print("end error ---------------")
            print(crabCommandArgs)
            res={ 'status':"CachefileNotFound",'jobs':{}}
            break
        if i>5:
            res={ 'status':"UnexpectedError",'jobs':{}}
            break
    q.put( res )

class CertInfo:
    def __init__( self ):
        p = subprocess.Popen("voms-proxy-info  --fqan",
                              stdout = subprocess.PIPE,
                              stderr = subprocess.PIPE,
                              shell=True)
        stdout, stderr = p.communicate()
        print(stdout)
        if p.returncode != 0:
            self.vo = ""
            self.voGroup = ""
            self.voRole = ""
        else:
            lines = stdout.split(b"\n")
            splitline = lines[0].split(b"/")
            if len(splitline) < 4:
                splitline = lines[1].split(b"/")
            self.vo = splitline[1]
            self.voGroup = splitline[2]
            try:
                self.voRole = splitline[2].split("=")[1]
                if  "NULL" in self.voRole:
                    self.voGroup = ""
            except:
                self.voRole = ""


class CrabTask:

    
    def __init__(self,
                 taskname="",
                 crab_config="",
                 crabController = None ,
                 initUpdate = True,
                 debuglevel = "ERROR",
                 datasetpath = "",
                 localDir = "",
                 outlfn = "" ,):

        # crab config as a python object should only be used via .config
        self._crabConfig = None

        self._crabFolder = None

        if taskname:
            self.name = taskname
        else:
            if not crab_config:
                raise ValueError("Either taskname or crab_config needs to be set")
            if not os.path.exists( crab_config):
                raise IOError("File %s not found" % crab_config )
            self.name = crab_config
            self.name = self.crabConfig.General.requestName
        self.uuid = uuid.uuid4()
        #~ self.lock = multiprocessing.Lock()
        #setup logging
        self.log = logging.getLogger( 'crabTask' )
        self.log.setLevel(logging.getLevelName(debuglevel))
        self.jobs = {}
        self.localDir = localDir
        self.outlfn = outlfn
        self.isUpdating = False
        self.taskId = -1
        #variables for statistics
        self.nJobs = 0
        self.state = "NOSTATE"
        self.maxjobnumber = 0
        self.nUnsubmitted   = 0
        self.nIdle = 0
        self.nRunning = 0
        self.nTransferring    = 0
        self.nCooloff    = 0
        self.nFailed    = 0
        self.nFinished    = 0
        self.nComplete    = 0
        self.failureReason = None
        self.lastUpdate = datetime.datetime.now().strftime( "%Y-%m-%d_%H.%M.%S" )

        self._isData = None
        self.resubmitCount = 0

        self.debug = False

        self.finalFiles = []
        self.totalEvents = 0


        self._datasetpath_default = datasetpath

        #start with first updates
        if initUpdate:
            self.update()
            self.updateJobStats()

    
    @property
    def isData( self ):
        if self._isData is None:
            try:
                test = self.crabConfig.Data.lumiMask
                self._isData = True
            except:
                if self.name.startswith( "Data_" ):
                    self._isData = True
                else:
                    self._isData = False
        return self._isData


    
    @property
    def crabConfig( self ):
        if self._crabConfig is None:
            crab = CrabController()
            self._crabConfig = crab.readCrabConfig( self.name )
        return self._crabConfig

    @property
    def datasetpath( self ):
        try:
            return self.crabConfig.Data.inputDataset
        except:
            pass
        return self._datasetpath_default

    @property
    def crabFolder( self ):
        if not self._crabFolder is None: return self._crabFolder
        crab = CrabController()
        if os.path.exists( os.path.join( self.crabConfig.General.workArea, crab._prepareFoldername( self.name ) ) ):
            self._crabFolder = os.path.join( self.crabConfig.General.workArea, crab._prepareFoldername( self.name ) )
            return self._crabFolder
        alternative_path = os.path.join(os.path.cwd(), crab._prepareFoldername( self.name ) )
        if os.path.exists( alternative_path ):
            self._crabFolder = alternative_path
            return self._crabFolder
        self.log.error( "Unable to find folder for Task")
        return ""

    
    def resubmit_failed( self ):
        failedJobIds = []
        controller =  CrabController()
        for jobkey in self.jobs.keys():
            job = self.jobs[jobkey]
            if job['State'] == 'failed':
                failedJobIds.append( job['JobIds'][-1] )
        controller.resubmit( self.name, joblist = failedJobIds )
        self.lastUpdate = datetime.datetime.now().strftime( "%Y-%m-%d_%H.%M.%S" )

    @property
    def crab_folder(self):
        return os.path.join( self.crabConfig.General.workArea,
            "crab_" + self.crabConfig.General.requestName)
    
    def update(self):
        #~ self.lock.acquire()
        self.log.debug( "Start update for task %s" % self.name )
        self.isUpdating = True
        controller =  CrabController()
        self.state = "UPDATING"
        # check if we should drop this sample due to missing info

        self.log.debug( "Try to get status for task" )
        self.state , self.jobs,self.failureReason = controller.status(self.crab_folder)
        self.log.debug( "Found state: %s" % self.state )
        if self.state=="FAILED":
            #try it once more
            time.sleep(2)
            self.state , self.jobs,self.failureReason = controller.status(self.crab_folder)
        self.nJobs = len(self.jobs)
        self.updateJobStats()
        if self.state == "NOSTATE":
            self.log.debug( "Trying to resubmit because of NOSTATE" )
            if self.resubmitCount < 3: self.self.handleNoState()
        # add to db if not
        # Final solution inf state not yet found
        self.isUpdating = False
        self.lastUpdate = datetime.datetime.now().strftime( "%Y-%m-%d_%H.%M.%S" )
        #~ self.lock.release()

    
    def handleNoState( self ):
        crab = CrabController()
        if "The CRAB3 server backend could not resubmit your task because the Grid scheduler answered with an error." in task.failureReason:
            # move folder and try it again
            cmd = 'mv %s bak_%s' %(crab._prepareFoldername( self.name ),crab._prepareFoldername( self.name ))
            p = subprocess.Popen(cmd,stdout=subprocess.PIPE, shell=True)#,shell=True,universal_newlines=True)
            (out,err) = p.communicate()
            self.state = "SHEDERR"
            configName = '%s_cfg.py' %(crab._prepareFoldername( self.name ))
            crab.submit( configName )

        elif task.failureReason is not None:
            self.state = "ERRHANDLE"
            crab.resubmit( self.name )
        self.resubmitCount += 1

    def test_print(self):
        return self.uuid
    
    def updateJobStats(self,dCacheFileList = None):
        jobKeys = sorted(self.jobs.keys())
        try:
            intJobkeys = [int(x) for x in jobKeys]
        except:
            print("error parsing job numers to int")

        #maxjobnumber = max(intJobkeys)

        stateDict = {'unsubmitted':0,'idle':0,'running':0,'transferring':0,'cooloff':0,'failed':0,'finished':0}
        nComplete = 0

        # loop through jobs
        for key in jobKeys:
            job = self.jobs[key]
             #check if all completed files are on decache
            for statekey in stateDict.keys():
                if statekey in job['State']:
                    stateDict[statekey]+=1
                    # check if finished fails are found on dCache if dCacheFilelist is given
                    if dCacheFileList is not None:
                        outputFilename = "%s_%s"%( self.name, key)
                        if 'finished' in statekey and any(outputFilename in s for s in dCacheFileList):
                            nComplete +=1

        for state in stateDict:
            attrname = "n" + state.capitalize()
            setattr(self, attrname, stateDict[state])
        self.nComplete = nComplete

    
    def readLogArch(self, logArchName):
        JobNumber = logArchName.split("/")[-1].split("_")[1].split(".")[0]
        log = {'readEvents' : 0}
        with tarfile.open( logArchName, "r") as tar:
            try:
                JobXmlFile = tar.extractfile('FrameworkJobReport-%s.xml' % JobNumber)
                root = ET.fromstring( JobXmlFile.read() )
                for child in root:
                    if child.tag == 'InputFile':
                        for subchild in child:
                            if subchild.tag == 'EventsRead':
                                nEvents = int(subchild.text)
                                log.update({'readEvents' : nEvents})
                                break
                        break
            except:
                print("Can not parse / read %s" % logArchName)
        return log


class TaskStats:

    
    def __init__(self, tasklist = None):
        if tasklist is not None:
            self.updateStats(tasklist)
        else:
            self.clearStats()

    
    def updateStats(self,tasklist):
        self.clearStats()
        self.nTasks = len(tasklist)
        for task in tasklist:
            if not task.isUpdating:
                self.nUnsubmitted   += task.nUnsubmitted
                self.nIdle += task.nIdle
                self.nRunning += task.nRunning
                self.nTransferring    += task.nTransferring
                self.nCooloff    += task.nCooloff
                self.nFailed    += task.nFailed
                self.nFinished    += task.nFinished
                self.nComplete    += task.nComplete

    
    def clearStats(self):
        self.nTasks = 0
        self.nUnsubmitted   = 0
        self.nIdle = 0
        self.nRunning = 0
        self.nTransferring    = 0
        self.nCooloff    = 0
        self.nFailed    = 0
        self.nFinished    = 0
        self.nComplete    = 0