findtrouble

#!/usr/bin/env python

import sys
import os
import subprocess

# Covers 2.4-2.5 (where elementtree is in elementtree) to 2.6+ (where it is not)
try:
  import lxml.etree as ET
except:
  try: 
    import elementtree.ElementTree as ET
  except:
    try:
      import xml.etree.ElementTree as ET
    except:
      sys.stderr.write("Could not get an elementtree implementation.\n")
      sys.exit(8)


def parse_error(qstat_cmd, stdout, stderr):
  """ Handles and prints information on XML parsing failures. """
  sys.stderr.write("No information or failed to parse output.\n")
  sys.stderr.write("Command executed was: \"%s\"\n" % qstat_cmd)
  sys.stderr.write("Error is:\n%s\n" % stderr)
  sys.stderr.write("\n")
  sys.exit(1)


# Output example for one host
"""
 <host name='node-u32.data.legion.ucl.ac.uk'>
   <hostvalue name='arch_string'>lx26-amd64</hostvalue>
   <hostvalue name='num_proc'>16</hostvalue>
   <hostvalue name='load_avg'>10.43</hostvalue>
   <hostvalue name='mem_total'>63.0G</hostvalue>
   <hostvalue name='mem_used'>9.2G</hostvalue>
   <hostvalue name='swap_total'>1851.6G</hostvalue>
   <hostvalue name='swap_used'>909.7M</hostvalue>
   <job name='5924228'>
     <jobvalue jobid='5924228' name='priority'>'2.093750'</jobvalue>
     <jobvalue jobid='5924228' name='qinstance_name'>Ulm@node-u32.data.legion.ucl.ac.uk</jobvalue>
     <jobvalue jobid='5924228' name='job_name'>CePydPym_INIT</jobvalue>
     <jobvalue jobid='5924228' name='job_owner'>uccahu0</jobvalue>
     <jobvalue jobid='5924228' name='job_state'>r</jobvalue>
     <jobvalue jobid='5924228' name='start_time'>1428476907</jobvalue>
     <jobvalue jobid='5924228' name='queue_name'>Ulm@node-u32.data.legion.ucl.ac.uk</jobvalue>
     <jobvalue jobid='5924228' name='pe_master'>SLAVE</jobvalue>
   </job>
 </host>
"""

def deSI(value):
  """ Converts an SGE MEMORY type (number + SI-ish suffix) to a float """
  SI_prefixes = { "K": 1000,
                  "M": 1000*1000,
                  "G": 1000*1000*1000,
                  "T": 1000*1000*1000*1000,
                  }
  if type(value) == int:
    return value
  if len(value) == 0:
    raise TypeError

  if len(value) < 2:
    if value.isdigit():
      return int(value)
    else:
      if value == "-":
        return float("nan")
      raise TypeError(value)

  if value[-1] in SI_prefixes:
    return float(value[:-2]) * SI_prefixes[value[-1]]
  else:
    if value[-1].isdigit() or value[-1] == ".":
      return float(value)
    else:
      raise TypeError

def dequote(value):
  """ Just removes surrounding single-quote characters, 
  because for some reason, the priority comes wrapped in single-quotes
  """
  if len(value) > 2:
    if value[0] == "'":
      if value[-1] == "'":
        return value[1:-2]
      else:
        raise TypeError(value)
    else: 
      raise TypeError(value)
  else:
    raise TypeError(value)

class Node:
  """ Storage class for per-node load values. """
  def __init__(self):
    self.hostname   = ""
    self.num_proc   = 0
    self.load_avg   = 0
    self.mem_total  = 0
    self.mem_used   = 0
    self.swap_total = 0
    self.swap_used  = 0
    self.jobs       = list()

  def fix_types(self):
    self.hostname   = str(self.hostname)
    self.num_proc   = int(self.num_proc)
    self.load_avg   = deSI(self.load_avg)
    self.mem_total  = deSI(self.mem_total)
    self.mem_used   = deSI(self.mem_used)
    self.swap_total = deSI(self.swap_total)
    self.swap_used  = deSI(self.swap_used)

  def __repr__(self):
    return repr(self.__dict__)

class Job:
  """ Storage class for per-job information. """
  def __init__(self):
    self.name           ="(no name)"
    self.id             =0
    self.owner          = "(no owner)"
    self.state          = "!"
    self.start_time     = 0
    self.priority       = 0
    self.pe_master      = "NO"
    self.qinstance_name = "(no queue)"

  def fix_types(self):
    self.name           = str(self.name)
    self.id             = str(self.id)
    self.owner          = str(self.owner)
    self.state          = str(self.state)
    self.start_time     = int(self.start_time)
    self.priority       = float(dequote(self.priority))
    self.pe_master      = str(self.pe_master)
    self.qinstance_name = str(self.qinstance_name)

  def __repr__(self):
    return repr(self.__dict__)


def get_nodes(root):
  """ Convert XML Element Tree into a list of Node objects
  (which each have a list of job objects).
  """

  job_property_translations = { 
      'job_name': 'name',
      'job_owner': 'owner',
      'job_state': 'state',
      }

  node_list = list()
  
  for element in root:
    if element.get('name')[0:5] == "node-":
      new_node = Node()
      new_node.hostname = element.get('name')
      
      for node_value in element.findall("hostvalue"):
        new_node.__dict__[node_value.get('name')] = node_value.text

      # Skip if the new is offline
      if '-' in new_node.__dict__.values():
        continue

      for job_element in element.findall("job"):
        new_job = Job()
        new_job.id = job_element.get('name')

        for job_value in job_element.findall("jobvalue"):
          if job_value.get('name') in job_property_translations:
            new_job.__dict__[job_property_translations[job_value.get('name')]] = job_value.text
          else:
            new_job.__dict__[job_value.get('name')] = job_value.text

        new_job.fix_types()
        new_node.jobs.append(new_job)

      new_node.fix_types()
      node_list.append(new_node)

  return node_list


def find_problematic_usage(node_list):
  """ Applies some simple criteria to find users who have jobs on problematic
  nodes. Could probably be more elegantly designed to allow more easily adding
  new checks.
  """
  mem_users = dict()
  load_users = dict()
  underload_users = dict()
  mem_jobs = dict()
  load_jobs = dict()
  underload_jobs = dict()
  arbitrary_memory_threshold = 0.9
  arbitrary_load_threshold = 1.2 
  arbitrary_underload_threshold = 0.3
  arbitrary_probability_threshold = 0.5
  for node in node_list:
    node_load = (node.load_avg/node.num_proc) 

    for job in node.jobs:
      if (node.mem_used/node.mem_total) > arbitrary_memory_threshold:
        is_out_of_mem = 1
      else:
        is_out_of_mem = 0

      if node_load > arbitrary_load_threshold:
        is_overloaded = 1
        is_underloaded = 0
      elif node_load < arbitrary_underload_threshold:
        is_overloaded = 0
        is_underloaded = 1
      else:
        is_overloaded = 0
        is_underloaded = 0

      mem_users[job.owner] = mem_users.get(job.owner,list())
      mem_users[job.owner].append(is_out_of_mem)
      load_users[job.owner] = load_users.get(job.owner,list())
      load_users[job.owner].append(is_overloaded)
      underload_users[job.owner] = underload_users.get(job.owner,list())
      underload_users[job.owner].append(is_underloaded)
      
      mem_jobs[job.id] = mem_jobs.get(job.id,list())
      mem_jobs[job.id].append(is_out_of_mem)
      load_jobs[job.id] = load_jobs.get(job.id,list())
      load_jobs[job.id].append(is_overloaded)
      underload_jobs[job.id] = underload_jobs.get(job.id,list())
      underload_jobs[job.id].append(is_underloaded)

  print ("--No Mem--")

  for user in mem_users.keys():
    user_data = mem_users[user]
    user_no_mem_fraction = sum(user_data)/len(user_data)
   
    if user_no_mem_fraction > arbitrary_probability_threshold:
      print("%s's %3d jobs have a %6.2f%% chance of being on a node that has run out of memory" %
              (user, len(user_data), 100*user_no_mem_fraction))

  print ("\n--High Load--")
  for user in load_users.keys():
    user_data = load_users[user]
    user_load_fraction = float(sum(user_data))/float(len(user_data))
  
#    print ("%s: %d, %f" % (user, len(user_data), user_load_fraction))
    if user_load_fraction > arbitrary_probability_threshold:
      print("%s's %3d jobs have a %6.2f%% chance of being on a node that has a load/proc > %f" % 
              (user, len(user_data), 100*user_load_fraction, arbitrary_load_threshold))

  print ("\n--Low Load--")
  for user in underload_users.keys():
    user_data = underload_users[user]
    user_underload_fraction = float(sum(user_data))/float(len(user_data))
   
    if user_underload_fraction > arbitrary_probability_threshold:
      print("%s's %3d jobs have a %6.2f%% chance of being on a node that has a load/proc < %.2f" % 
              (user, len(user_data), 100*user_underload_fraction, arbitrary_underload_threshold))



def get_input_from_command():
  """ Runs qhost and returns XML output as an XML Element Tree. """
  qhost_cmd = "qhost -xml -j"

  # Run qhost with a pipe
  qhost_proc = subprocess.Popen(qhost_cmd,
    shell=True,
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE,
    env=os.environ)
  stdout, stderr = qhost_proc.communicate()
  # Parse XML
  try:
    root = ET.fromstring(stdout)
  except:
    parse_error(qhost_cmd, stdout, stderr)
  return root

def get_input_from_file(path):
  """ Intended to simulate get_input_from_command for debugging. """
  try:
    tree = ET.parse(path)
    root = tree.getroot()
  except:
    parse_error("","","")  
  return root


def main(source="command", source_file=None):
  # If a user name is provided on the command line, use that
  # otherwise get the USER variable.
  user = os.getenv("USER")
  if len(sys.argv) == 2:
    user = sys.argv[1]
  if len(sys.argv) > 2:
    sys.stderr.write("Too many arguments! Either no arguments or a username, please.\n")
    exit(5)
 
  if source == "file":
    xml_root = get_input_from_file(source_file)
  elif source=="command":
    xml_root = get_input_from_command()

  nodes = get_nodes(xml_root)

  find_problematic_usage(nodes)

  #print(nodes)

if __name__ == "__main__":
  main(source="command")

  # For testing
  #main("file", "test.xml")