source: trunk/plugin/PBSQuery.py @ 94

#
# Authors: Roy Dragseth (roy.dragseth@cc.uit.no) 
#          Bas van der Vlies (basv@sara.nl)
#
# SVN INFO:
#       $Id: PBSQuery.py 392 2005-03-02 16:47:25Z bas $
#
"""
Usage: from PBSQuery import PBSQuery

This class gets the info from the pbs_server via the pbs.py module
for the several pbs objects. All get..() functions returns an dictionary
of openpbs/torque objects. There are four objects:
 - server 
 - queue
 - job
 - node

Each object can be handled as an dictionary and has several member
functions.

There are the following functions for PBSQuery:
  job - 
        getjob(name, attributes=<default is all>)
        getjobs(attributes=<default is all>)
 
  node -
        getnode(name, attributes=<default is all>)
        getnodes(attributes=<default is all>)

  queue -
        getqueue(name, attributes=<default is all>)
        getqueues(attributes=<default is all>)

  server -
        get_serverinfo(attributes=<default is all>)

Here is an example how to use the module:
        from PBSQuery import PBSQuery
        p = PBSQuery()
        nodes = p.getnodes()
        for name,node in nodes.items():
            print name
            if node.is_free():
               print node

"""
import pbs


class PBSQuery:
    def __init__(self, server=None):
        if not server:
            self.server = pbs.pbs_default()
        else:
            self.server = server

        self.con = pbs.pbs_connect(self.server)
        self.attribs = 'NULL'

    def __del__(self):
        pbs.pbs_disconnect(self.con)
        self.con = None
        self.attribs = None

    def _list_2_attrib(self, list):
        """Convert an python list to an attrib list suitable for pbs"""
        self.attribs = pbs.new_attrl( len(list) )
        i = 0 
        for attrib in list:
            self.attribs[i].name = attrib
            i = i + 1

    def _list_2_dict(self, l, class_func):
        """Convert an pbsstat function list to an class dictionary"""
        self.d = dict()
        for item in l:
            new = class_func()
            for a in item.attribs:
                new.name = item.name 
                if a.resource:
                   new.attribs[a.name + '.' + a.resource] = a.value
                else:
                   new.attribs[a.name] = a.value
                self.d[item.name] = new
        self._free(l)
                
    def _free(self, memory):
        pbs.pbs_statfree(memory)

    def _statserver(self, attrib_list=None):
        """Get the server config from the pbs server"""
        if attrib_list:
           self._list_2_attrib(attrib_list)
        else:
           self.attribs = 'NULL'

        serverinfo = pbs.pbs_statserver(self.con, self.attribs, 'NULL')
        self.serverinfo = dict()
        self._list_2_dict(serverinfo, server)

    def get_serverinfo(self, attrib_list=None):
        self._statserver(attrib_list)
        return self.d

    def _statqueue(self, queue_name='', attrib_list=None):
        """Get the queue config from the pbs server"""
        if attrib_list:
           self._list_2_attrib(attrib_list)
        else:
           self.attribs = 'NULL'

        queues = pbs.pbs_statque(self.con, queue_name, self.attribs, 'NULL')
        self.queues = dict()
        self._list_2_dict(queues, queue)

    def getqueue(self, name, attrib_list=None):
        self._statqueue(name, attrib_list)
        return self.d
        
    def getqueues(self, attrib_list=None):
        self._statqueue('', attrib_list)
        return self.d

    def _statnode(self, node_name='', attrib_list=None):
        """Get the node config from the pbs server"""
        if attrib_list:
           self._list_2_attrib(attrib_list)
        else:
           self.attribs = 'NULL'

        nodes = pbs.pbs_statnode(self.con, node_name, self.attribs, 'NULL')
        self.nodes = dict()
        self._list_2_dict(nodes, node)

    def getnode(self, name, attrib_list=None):
        self._statnode(name, attrib_list)
        return self.d
        
    def getnodes(self, attrib_list=None):
        self._statnode('', attrib_list)
        return self.d

    def _statjob(self, job_name='', attrib_list=None):
        """Get the job config from the pbs server"""
        if attrib_list:
           self._list_2_attrib(attrib_list)
        else:
           self.attribs = 'NULL'
           
        jobs = pbs.pbs_statjob(self.con, job_name, self.attribs, 'NULL')
        self.jobs = dict()
        self._list_2_dict(jobs, job)

    def getjob(self, name, attrib_list=None):
        self._statjob(name, attrib_list)
        return self.d
        
    def getjobs(self, attrib_list=None):
        self._statjob('', attrib_list)
        return self.d

class _PBSobject:
    TRUE  = 1
    FALSE = 0

    def __init__(self):
        self.attribs = dict()
        self.name = None

    def __len__(self):
        return(len(self.attribs.keys()))

    def __str__(self):
        data = self.name + '\n'
        for name,value in self.attribs.items():
                data = data + '\t%s = %s\n' %(name, value)
        return(data)

    __repr__ = __str__

    def __getitem__(self,key):
        if key == 'name':
            return self.name
        
        return(self.attribs[key])

    def has_key(self, key):
        return self.attribs.has_key(key)

    def get_value(self, key):
        if self.attribs.has_key(key):
           return self.attribs[key]
        else:
           return None

    def keys(self):
        return self.attribs.keys()

    def values(self):
        return self.attribs.values()

    def items(self):
        return self.attribs.items()
    
class job(_PBSobject):
    """PBS job class"""

    def is_running(self):
        if self.get_value('job_state') == 'Q':
           return self.FALSE
        else:
           return self.TRUE 

class node(_PBSobject):
    """PBS node class"""

    def is_free(self):
        """Check if node is free"""
        if self.get_value('state') == 'free':
              return self.TRUE
        else:
              return self.FALSE 

    def has_job(self):
        """Does the node run a job"""
        if self.get_value('jobs'):
              return self.TRUE
        else:
              return self.FALSE

class queue(_PBSobject):
    """PBS queue class"""

    def is_enabled(self):
        if self.get_value('enabled') == 'True':
           return self.TRUE 
        else:
           return self.FALSE

    def is_execution(self):
        if self.get_value('queue_type') == 'Execution':
           return self.TRUE 
        else:
           return self.FALSE


class server(_PBSobject):
    """PBS server class"""

def main():
    p = PBSQuery()

    serverinfo = p.get_serverinfo()
    for name, server in serverinfo.items():
        print server
        for key in server.keys():
                print key, ' = ', server[key]

    queues = p.getqueues()
    for queue in queues.values():
        if queue.is_execution():
           print queue

    jobs = p.getjobs()
    for name,job in jobs.items():
        if job.is_running():
           print job

    l = ['state']
    nodes = p.getnodes(l)
    for name,node in nodes.items():
        if node.is_free(): 
           print node

if __name__ == "__main__":
    main()