source: trunk/src/PBSQuery.py @ 227

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

forgot to add the old_data_structure function
  • Property svn:keywords set to Id
File size: 10.7 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 227 2010-04-15 08:46:13Z 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                try:
354                        return self.data[name]
355                except KeyError:
356                        error = 'Attribute key error: %s' %(name)
357                        raise PBSError(error)
358
359        def __iter__(self):
360                return iter(self.data.keys())
361
362        def uniq(self, list):
363                """Filter out unique items of a list"""
364                uniq_items = {}
365                for item in list:
366                        uniq_items[item] = 1
367                return uniq_items.keys()
368
369        def return_value(self, key):
370                """Function that returns a value independent of new or old data structure"""
371                if isinstance(self[key], types.ListType):
372                        return self[key][0]
373                else:
374                        return self[key] 
375
376class job(_PBSobject):
377        """PBS job class""" 
378        def is_running(self):
379
380                value = self.return_value('job_state')
381                if value == 'Q':
382                        return self.TRUE
383                else:
384                        return self.FALSE
385
386        def get_nodes(self, unique=None):
387                """Returns a list of the nodes which run this job"""
388                nodes = self.get_value('exec_host')
389                if nodes:
390                        nodelist = string.split(nodes,'+')
391                        if not unique:
392                                return nodelist
393                        else:
394                                return self.uniq(nodelist)
395                return list()
396
397
398class node(_PBSobject):
399        """PBS node class"""
400       
401        def is_free(self):
402                """Check if node is free"""
403
404                value = self.return_value('state')
405                if value == 'free':
406                        return self.TRUE
407                else: 
408                        return self.FALSE
409
410        def has_job(self):
411                """Does the node run a job"""
412                try:
413                        a = self['jobs']
414                        return self.TRUE
415                except KeyError, detail:
416                        return self.FALSE
417       
418        def get_jobs(self, unique=None):
419                """Returns a list of the currently running job-id('s) on the node"""
420                jobstring = self.get_value('jobs')
421                if jobstring:
422                        joblist = re.compile('[^\\ /]\\d+[^/.]').findall( jobstring )
423                        if not unique:
424                                return joblist
425                        else:
426                                return self.uniq(joblist)
427                return list()
428
429
430class queue(_PBSobject):
431        """PBS queue class"""
432        def is_enabled(self):
433
434                value = self.return_value('enabled')
435                if value == 'True':
436                        return self.TRUE
437                else:
438                        return self.FALSE
439
440        def is_execution(self):
441
442                value = self.return_value('queue_type')
443                if value == 'Execution':
444                        return self.TRUE
445                else:
446                        return self.FALSE
447
448class server(_PBSobject):
449        """PBS server class"""
450
451        def get_version(self):
452                return self.get_value('pbs_version')
453
454def main():
455        p = PBSQuery() 
456        serverinfo = p.get_serverinfo()
457        for server in serverinfo.keys():
458                print server, ' version: ', serverinfo[server].get_version()
459        for resource in serverinfo[server].keys():
460                print '\t ', resource, ' = ', serverinfo[server][resource]
461
462        queues = p.getqueues()
463        for queue in queues.keys():
464                print queue
465                if queues[queue].is_execution():
466                        print '\t ', queues[queue]
467                if queues[queue].has_key('acl_groups'):
468                        print '\t acl_groups: yes'
469                else:
470                        print '\t acl_groups: no'
471
472        jobs = p.getjobs()
473        for name,job in jobs.items():
474                if job.is_running():
475                        print job
476
477        l = ['state']
478        nodes = p.getnodes(l)
479        for name,node in nodes.items():
480                if node.is_free(): 
481                        print node
482
483if __name__ == "__main__":
484        main()
Note: See TracBrowser for help on using the repository browser.