GBSUtilities.py

Go to the documentation of this file.

#  Utilities that are either classless or serve two or more classes.

import os
import re
import signal
import time

def RemoveChildProcesses(timeout=0):
    """Delete all child processes starting with those without children and working up.
    Allow child 'timeout' seconds before killing."""
    killed_process = 1
    my_pid = os.getpid()
    print "Debug: Removing child processes after %d secs. Parent process is %d" % (timeout,my_pid)
    while killed_process:  #Continue until fail to kill any process
        #  Collect a list of children and their parents
        killed_process = 0
        parents   = {}  # Hash: pid -> parent pid
        children  = {}  # Hash: pid -> list of child pids
        time_cmds = {}  # Hash: pid -> stime,cmd
        ps_cmd = "ps -o pid,ppid,stime,cmd"
        inp = os.popen(ps_cmd,"r")
        for line in inp:
            if re.search(ps_cmd,line): continue #Skip process running ps!
            print "Debug PS line: ",line,
            mo = re.search(r"(\d+)\s+(\d+)\s+(.*)",line)
            if mo:
                (pid,ppid,cmd)      = mo.groups()
                parents[int(pid)]   = int(ppid)
                time_cmds[int(pid)] = cmd
        inp.close()
        for child,parent in parents.iteritems():
            antecedent = parent
            while parents.has_key(antecedent):
                if children.has_key(antecedent): children[antecedent].append(child)
                else: children[antecedent] = [child]
                antecedent = parents[antecedent]
        # For the current process, attempt to kill any children without children
        if not children.has_key(my_pid): break
        for child in children[my_pid]:
            if children.has_key(child): continue # Skip children with children
            #  Allow child process 'timeout' seconds before killing.
            mo = re.search(r"(\d\d):(\d\d)",time_cmds[child])
            if mo:
                start_hour = int(mo.group(1))
                start_min  = int(mo.group(2))
                now_hour = int(time.strftime("%H"))
                now_min  = int(time.strftime("%M"))
                time_run = (now_hour*24 + now_min - start_hour*24 -start_min)*60
                if time_run >= 0 and time_run < timeout:
                     print "Debug: Child %d has run %d secs, delay killing" % (child,time_run)
                     time.sleep(timeout - time_run)
                     killed_process = 1 #Claim to have killed just to keep loop going.
                     continue
            #  Attempt to kill twice.  Use /proc/pid/status to check if successful
            if os.path.isfile('/proc/' + str(child) + '/status'):
                print "Attempting to kill " + str(child) + " " + time_cmds[child]
                os.kill(child,signal.SIGTERM)
                time.sleep(1)
                if os.path.isfile('/proc/' + str(child) + '/status'): os.kill(child,signal.SIGKILL)
                time.sleep(1)
                if os.path.isfile('/proc/' + str(child) + '/status'): print "Failed to kill " + str(child)
                else: killed_process = 1


def BuildEnvStr(d):

    """Convert dictionary into an alphabetically sorted environment string: 'var1=val1,var2=val2...'"""

    s = ""
    var_list = d.keys()
    var_list.sort()
    for var in var_list:
        val = d[var]
        val = re.compile(',').sub('\\,',val)
        if s: s += ','
        s += var + '=' + val
    return s

def ListCommaSepList(s,prefix=""):
    """List a comma separated list string s with each entry on a separate list starting with prefix."""
    l = []
    ParseCommaSepList(s,l)
    for element in l:
        print prefix + element
        
def ParseCommaSepList(s,l):

    """Parse comma separated list string s (with possibly embedded \,) and append into a list l"""
    if not s: return l
    part_arg = ""
    for arg in s.split(','):
        if part_arg: arg = part_arg + arg
        part_arg = ""
        if arg[-1] == "\\": part_arg = arg[:-1] + ','
        else:               l.append(arg)
    if part_arg: l.append(part_arg)
    return l

def ParseEnvStr(s):

    """Parse environment string: 'var1=val1,var2=val2...' into dictionary and return.

    If parse error, dictionary will be empty"""

    d = {}
    l = []
    ParseCommaSepList(s,l)
    for entry in l:
        mo = re.search(r"^\s*(\S+)\s*=(.*)$",entry)
        if not mo: return None
        d[mo.group(1)] = mo.group(2)
    return d

def ProcessSandboxSetup(obj,in_str,sb_type):

    """Support setting of input and output sandbox lists for Task, Jobs and Protojobs objects.

    Args are:-
      obj         Object whose sandbox is being setup
      in_str     Comma separated list of sandbox files
      sb_type     Sandbox type: either "input" or "output"

    Returns:
      ok          Flag, True if O.K.
      out_str     The string to be stored ( but only if ok is True)
      """

    #  Form file list and validate it.

    l = []
    ParseCommaSepList(in_str,l)
    out_str = ""
    err_msg = ""
    for file in l:
        if sb_type == "input":
            if not os.path.isfile(file):
                if not err_msg: err_msg = "Cannot find input sandbox file(s):"
                err_msg += " '" + file + "'"
            else:
                if out_str: out_str += ","
                out_str += os.path.split(file)[1]
        else:
            if re.search(r"/",file):
                if not err_msg: err_msg = "Output sandbox file(s) contains directory:"
                err_msg += " '" + file + "'"
            else:
                if out_str: out_str += ","
                out_str += file
    if err_msg:
        print err_msg
        return (False,"")

    #  For input sandbox, except for GBSProtoJob objects, create or wipe
    #  sandbox directory and copy files.

    if sb_type != "input" or obj.GetType() == "GBSProtoJob": return (True,out_str)
    sbox_dir = obj.GetStoreLocation("child_dir") + "/InputSandbox"
    if not os.path.isdir(sbox_dir):
        try:
            os.makedirs(sbox_dir,0755)
        except OSError:
                print "Unable to create directory: " + sbox_dir
                return (False,"")
    else:
        os.system("rm -f " + sbox_dir + "/*")

    for file in l:
        print "Copying " + file + " -> " + sbox_dir
        if os.system("cp " + file + " " + sbox_dir):
            print "Failed to copy file!"
            return (False,"")
    
    return (True,out_str)
            

---