source: trunk/src/PBSQuery.py @ 243

Last change on this file since 243 was 243, checked in by bas, 14 years ago

Added some demo code to overload setattr for UserDict? Module
  • Property svn:keywords set to Id
File size: 11.3 KB
Line 
1#
2# Authors: Roy Dragseth (roy.dragseth@cc.uit.no)
3#          Bas van der Vlies (basv@sara.nl)
4#
5# SVN INFO:
6#       $Id: PBSQuery.py 243 2010-07-16 12:25:41Z bas $
7#
8"""
9Usage: from PBSQuery import PBSQuery
10
11This class gets the info from the pbs_server via the pbs.py module
12for the several batch objects. All get..() functions return an dictionary
13with id as key and batch object as value
14
15There are four batch objects:
16 - server
17 - queue
18 - job
19 - node
20
21Each object can be handled as an dictionary and has several member
22functions. The second parameter is an python list and can be used if you
23are only interested in certain resources, see example
24
25There are the following functions for PBSQuery:
26  job -
27        getjob(job_id, attributes=<default is all>)
28        getjobs(attributes=<default is all>)
29 
30  node -
31        getnode(node_id, attributes=<default is all>)
32        getnodes(attributes=<default is all>)
33
34  queue -
35        getqueue(queue_id, attributes=<default is all>)
36        getqueues(attributes=<default is all>)
37
38  server -
39        get_serverinfo(attributes=<default is all>)
40
41Here is an example how to use the module:
42        from PBSQuery import PBSQuery
43        p = PBSQuery()
44        nodes = p.getnodes()
45        for name,node in nodes.items():
46            print name
47            if node.is_free():
48               print node, node['state']
49
50        l = [ 'state', 'np' ]
51        nodes = p.getnodes(l)
52        for name,node in nodes.items():
53               print node, node['state']
54
55The parameter 'attributes' is an python list of resources that
56you are interested in, eg: only show state of nodes
57        l = list()
58        l.append('state')
59        nodes = p.getnodes(l)
60"""
61import pbs
62import UserDict
63import string
64import sys
65import re
66import types
67
68class PBSError(Exception):
69        def __init__(self, msg=''):
70                self.msg = msg
71                Exception.__init__(self, msg)
72               
73        def __repr__(self):
74                return self.msg
75
76        __str__ = __repr__
77
78
79class PBSQuery:
80
81        # a[key] = value, key and value are data type string
82        #
83        OLD_DATA_STRUCTURE = False
84
85        def __init__(self, server=None):
86                if not server:
87                        self.server = pbs.pbs_default()
88                else:
89                        self.server = server
90
91        def _connect(self):
92                """Connect to the PBS/Torque server"""
93                self.con = pbs.pbs_connect(self.server)
94                if self.con < 0:
95                        str = "Could not make a connection with %s\n" %(self.server)
96                        raise PBSError(str)
97
98        def _disconnect(self):
99                """Close the PBS/Torque connection"""
100                pbs.pbs_disconnect(self.con)
101                self.attribs = 'NULL'
102
103        def _list_2_attrib(self, list):
104                """Convert a python list to an attrib list suitable for pbs"""
105                self.attribs = pbs.new_attrl( len(list) )
106                i = 0 
107                for attrib in list:
108                        # So we can user Resource
109                        attrib = attrib.split('.')
110                        self.attribs[i].name = attrib[0]
111                        i = i + 1
112
113        def _pbsstr_2_list(self, str, delimiter):
114                """Convert a string to a python list and use delimiter as spit char"""
115                l = sting.splitfields(str, delimiter)
116                if len(l) > 1:
117                        return l
118
119        def _list_2_dict(self, l, class_func):
120                """
121                Convert a pbsstat function list to a class dictionary, The
122                data structure depends on the function new_data_structure().
123               
124                Default data structure is:
125                        class[key] = value, Where key and value are of type string
126
127                Future release, can be set by new_data_structure():
128                        - class[key] = value where value can be:
129                          1. a list of values of type string
130                          2. a dictionary with as list of values of type string. If
131                             values contain a '=' character
132
133                  eg:
134                            print node['np']
135                                >> [ '2' ]
136
137                                print node['status']['arch']
138                                >> [ 'x86_64' ]
139                """
140                self.d = {}
141                for item in l:
142                        new = class_func()
143
144                        self.d[item.name] = new
145                       
146                        new.name = item.name
147
148                        for a in item.attribs:
149
150                                if self.OLD_DATA_STRUCTURE:
151
152                                        if a.resource:
153                                                key = '%s.%s' %(a.name, a.resource)
154                                        else:
155                                                key = '%s' %(a.name)
156
157                                        new[key] = a.value
158
159                                else:
160                                        values = string.split(a.value, ',') 
161                                        sub_dict = string.split(a.value, '=')
162
163
164                                        # We must creat sub dict, only for specified
165                                        # key values
166                                        #
167                                        if a.name in ['status', 'Variable_List']:
168
169                                                for v in values:
170
171                                                        tmp_l = v.split('=')
172
173                                                        # Check if we already added the key
174                                                        #
175                                                        if new.has_key(a.name):
176                                                                new[a.name][ tmp_l[0] ] = tmp_l[1:]
177
178                                                        else:
179                                                                tmp_d  = dict()
180                                                                tmp_d[ tmp_l[0] ] = tmp_l[1:]
181                                                                new[a.name] = class_func(tmp_d) 
182
183                                        else: 
184                                                # Check if it is a resource type variable, eg:
185                                                #  - Resource_List.(nodes, walltime, ..)
186                                                #
187                                                if a.resource:
188
189                                                        if new.has_key(a.name):
190                                                                new[a.name][a.resource] = values
191
192                                                        else:
193                                                                tmp_d = dict()
194                                                                tmp_d[a.resource] = values
195                                                                new[a.name] = class_func(tmp_d) 
196                                                else:
197                                                        # Simple value
198                                                        #
199                                                        new[a.name] = values
200
201                self._free(l)
202               
203        def _free(self, memory):
204                """
205                freeing up used memmory
206
207                """
208                pbs.pbs_statfree(memory)
209
210        def _statserver(self, attrib_list=None):
211                """Get the server config from the pbs server"""
212                if attrib_list:
213                        self._list_2_attrib(attrib_list)
214                else:
215                        self.attribs = 'NULL' 
216                       
217                self._connect()
218                serverinfo = pbs.pbs_statserver(self.con, self.attribs, 'NULL')
219                self._disconnect() 
220               
221                self._list_2_dict(serverinfo, server)
222
223        def get_serverinfo(self, attrib_list=None):
224                self._statserver(attrib_list)
225                return self.d
226
227        def _statqueue(self, queue_name='', attrib_list=None):
228                """Get the queue config from the pbs server"""
229                if attrib_list:
230                        self._list_2_attrib(attrib_list)
231                else:
232                        self.attribs = 'NULL' 
233                       
234                self._connect()
235                queues = pbs.pbs_statque(self.con, queue_name, self.attribs, 'NULL')
236                self._disconnect() 
237               
238                self._list_2_dict(queues, queue)
239
240        def getqueue(self, name, attrib_list=None):
241                self._statqueue(name, attrib_list)
242                try:
243                        return self.d[name]
244                except KeyError, detail:
245                        return self.d
246       
247        def getqueues(self, attrib_list=None):
248                self._statqueue('', attrib_list)
249                return self.d
250
251        def _statnode(self, select='', attrib_list=None, property=None):
252                """Get the node config from the pbs server"""
253                if attrib_list:
254                        self._list_2_attrib(attrib_list)
255                else:
256                        self.attribs = 'NULL' 
257                       
258                if property:
259                        select = ':%s' %(property)
260
261                self._connect()
262                nodes = pbs.pbs_statnode(self.con, select, self.attribs, 'NULL')
263                self._disconnect() 
264               
265                self._list_2_dict(nodes, node)
266
267        def getnode(self, name, attrib_list=None):
268                self._statnode(name, attrib_list)
269                try:
270                        return self.d[name]
271                except KeyError, detail:
272                        return self.d
273       
274        def getnodes(self, attrib_list=None):
275                self._statnode('', attrib_list)
276                return self.d
277
278        def getnodes_with_property(self, property, attrib_list=None):
279                self._statnode('', attrib_list, property)
280                return self.d
281
282        def _statjob(self, job_name='', attrib_list=None):
283                """Get the job config from the pbs server"""
284                if attrib_list:
285                        self._list_2_attrib(attrib_list)
286                else:
287                        self.attribs = 'NULL' 
288                       
289                self._connect()
290                jobs = pbs.pbs_statjob(self.con, job_name, self.attribs, 'NULL')
291                self._disconnect() 
292               
293                self._list_2_dict(jobs, job)
294
295        def getjob(self, name, attrib_list=None):
296                # To make sure we use the full name of a job; Changes a name
297                # like 1234567 into 1234567.server.name
298                name = name.split('.')[0] + "." + self.get_server_name()
299
300                self._statjob(name, attrib_list)
301                try:
302                        return self.d[name]
303                except KeyError, detail:
304                        return self.d
305       
306        def getjobs(self, attrib_list=None):
307                self._statjob('', attrib_list)
308                return self.d
309
310        def get_server_name(self):
311                return self.server
312
313        def new_data_structure(self): 
314                """
315                Use the new data structure. Is now the default
316                """
317                self.OLD_DATA_STRUCTURE = False
318
319        def old_data_structure(self): 
320                """
321                Use the old data structure. This function is obselete and
322                will be removed in a future release
323                """
324                self.OLD_DATA_STRUCTURE = True
325
326class _PBSobject(UserDict.UserDict):
327        TRUE  = 1
328        FALSE = 0
329
330        def __init__(self, dictin = None):
331                UserDict.UserDict.__init__(self)
332                self.name = None
333
334                if dictin:
335                        if dictin.has_key('name'):
336                                self.name = dictin['name']
337                                del dictin['name']
338                        self.data = dictin
339
340        def get_value(self, key):
341                if self.has_key(key):
342                        return self[key]
343                else:
344                        return None
345
346        def __repr__(self):
347                return repr(self.data)
348
349        def __str__(self):
350                return str(self.data)
351
352        def __getattr__(self, name):
353                """
354                override the class attribute get method. Return the value
355                from the Userdict
356                """
357                try:
358                        return self.data[name]
359                except KeyError:
360                        error = 'Attribute key error: %s' %(name)
361                        raise PBSError(error)
362
363        ## Disabled for this moment, BvdV 16 July 2010
364        #
365        #def __setattr__(self, name, value):
366        #       """
367        #       override the class attribute set method only when the UserDict
368        #       has set its class attribute
369        #       """
370        #       if self.__dict__.has_key('data'):
371        #               self.data[name] = value
372        #       else:
373        #               self.__dict__[name] = value
374
375        def __iter__(self):
376                return iter(self.data.keys())
377
378        def uniq(self, list):
379                """Filter out unique items of a list"""
380                uniq_items = {}
381                for item in list:
382                        uniq_items[item] = 1
383                return uniq_items.keys()
384
385        def return_value(self, key):
386                """Function that returns a value independent of new or old data structure"""
387                if isinstance(self[key], types.ListType):
388                        return self[key][0]
389                else:
390                        return self[key] 
391
392class job(_PBSobject):
393        """PBS job class""" 
394        def is_running(self):
395
396                value = self.return_value('job_state')
397                if value == 'Q':
398                        return self.TRUE
399                else:
400                        return self.FALSE
401
402        def get_nodes(self, unique=None):
403                """Returns a list of the nodes which run this job"""
404                nodes = self.get_value('exec_host')
405                if nodes:
406                        nodelist = string.split(nodes,'+')
407                        if not unique:
408                                return nodelist
409                        else:
410                                return self.uniq(nodelist)
411                return list()
412
413
414class node(_PBSobject):
415        """PBS node class"""
416       
417        def is_free(self):
418                """Check if node is free"""
419
420                value = self.return_value('state')
421                if value == 'free':
422                        return self.TRUE
423                else: 
424                        return self.FALSE
425
426        def has_job(self):
427                """Does the node run a job"""
428                try:
429                        a = self['jobs']
430                        return self.TRUE
431                except KeyError, detail:
432                        return self.FALSE
433       
434        def get_jobs(self, unique=None):
435                """Returns a list of the currently running job-id('s) on the node"""
436
437                jobs = self.get_value('jobs')
438                if jobs:       
439                        if type(jobs) == type('string') :
440                                jobs = re.compile('[^\\ /]\\d+[^/.]').findall( jobs )
441                       
442                        if not unique:
443                                return jobs
444                        else:
445                                job_re = re.compile('^(?:\d+/)?(.+)')
446                                l = list()
447
448                                for j in jobs:
449                                    jobstr = job_re.findall(j.strip())[0]
450                                    if jobstr not in l: 
451                                        l.append(jobstr)           
452                                return l
453
454                return list()
455
456
457class queue(_PBSobject):
458        """PBS queue class"""
459        def is_enabled(self):
460
461                value = self.return_value('enabled')
462                if value == 'True':
463                        return self.TRUE
464                else:
465                        return self.FALSE
466
467        def is_execution(self):
468
469                value = self.return_value('queue_type')
470                if value == 'Execution':
471                        return self.TRUE
472                else:
473                        return self.FALSE
474
475class server(_PBSobject):
476        """PBS server class"""
477
478        def get_version(self):
479                return self.get_value('pbs_version')
480
481def main():
482        p = PBSQuery() 
483        serverinfo = p.get_serverinfo()
484        for server in serverinfo.keys():
485                print server, ' version: ', serverinfo[server].get_version()
486        for resource in serverinfo[server].keys():
487                print '\t ', resource, ' = ', serverinfo[server][resource]
488
489        queues = p.getqueues()
490        for queue in queues.keys():
491                print queue
492                if queues[queue].is_execution():
493                        print '\t ', queues[queue]
494                if queues[queue].has_key('acl_groups'):
495                        print '\t acl_groups: yes'
496                else:
497                        print '\t acl_groups: no'
498
499        jobs = p.getjobs()
500        for name,job in jobs.items():
501                if job.is_running():
502                        print job
503
504        l = ['state']
505        nodes = p.getnodes(l)
506        for name,node in nodes.items():
507                if node.is_free(): 
508                        print node
509
510if __name__ == "__main__":
511        main()
Note: See TracBrowser for help on using the repository browser.