add watch support

gpuutil/__main__.py

@@ -1,3 +1,4 @@
+import curses
 from gpuutil import GPUStat, loaddict, savedict
 import sys
 import argparse
@@ -64,6 +65,7 @@ if __name__ == '__main__':
     parser.add_argument('--show-process', '-sp', default=True, type=str2bool, help='whether show process or not')
     parser.add_argument('--vertical', '-v', default=False, type=str2bool, help='whether show each user in different lines. (show user vertically)')
     parser.add_argument('--save', default=False, action="store_true", help='save config to profile')
+    parser.add_argument('--watch', '-w', default=-1, type=float, help='save config to profile')
     args = parser.parse_args()
     cols = args.cols if args.cols is not None else recommended_cols
     show_process = args.show_process
@@ -106,4 +108,33 @@ if __name__ == '__main__':
 
         else:
             raise ValueError('Profile do not exist.\nAvaliable Profiles:{0}'.format(','.join(list(config.keys()))))
-    stat.show(enabled_cols = cols, colsty=style, colsz=limit, vertical=vertical, show_command=show_process)
+    info = stat.show(enabled_cols = cols, colsty=style, colsz=limit, vertical=vertical, show_command=show_process, tostdout=False)
+    if args.watch < 0:
+        print(info)
+    else:
+        from curses import wrapper
+        import time
+        def continuous_watch(stdscr, info):
+            curses.curs_set(0)
+            stdscr.clear()
+            stdscr.nodelay(True)
+            lasttime = time.time()
+            try:
+                while True:
+                    c = stdscr.getch()
+                    if c in [ord('q'), ord('Q')]:
+                        break
+                    curses.flushinp()
+                    hint = "Interval: {0} S | CurrentTime: {1}".format(args.watch, time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
+                    stdscr.clear()
+                    stdscr.addstr(0, 0, hint + '\n' + info)
+                    stdscr.refresh()
+                    passed_time = time.time() - lasttime
+                    if passed_time < args.watch:
+                        time.sleep(args.watch - passed_time)
+                    lasttime = time.time()
+                    info = stat.show(enabled_cols = cols, colsty=style, colsz=limit, vertical=vertical, show_command=show_process, tostdout=False)
+            except KeyboardInterrupt:
+                curses.flushinp()
+                pass
+        wrapper(continuous_watch, info)

gpuutil/gpuutil.py

@@ -378,7 +378,7 @@ class GPUStat():
             gpu['id'] = i
             self.gpus.append(gpu)
 
-    def show(self, enabled_cols = ['ID', 'Fan', 'Temp', 'Pwr', 'Freq', 'Util', 'Vmem', 'Users'], colsty=None, colsz=None, show_command=True, vertical=False):
+    def show(self, enabled_cols = ['ID', 'Fan', 'Temp', 'Pwr', 'Freq', 'Util', 'Vmem', 'Users'], colsty=None, colsz=None, show_command=True, vertical=False, tostdout=True):
         self.parse()
         gpu_infos = []
         # stats = {
@@ -450,7 +450,10 @@ class GPUStat():
                     pid = proc['pid']
                     proc['gpu'] = [str(gpu['id'])]
                     if type(proc['vmem']) is str:
+                        try:
                             proc['vmem'] = int(proc['vmem'].split(' ')[0])
+                        except:
+                            proc['vmem'] = 0
                     if pid not in procs:
                         procs[pid] = proc
                     else:
@@ -471,7 +474,9 @@ class GPUStat():
             table_width = info.find('\n')
             proc_info = draw_table([['Process Info'.center(table_width-4)], [proc_info]], rowsty="c|c|", colsty="|l|", colsz=[table_width-4])
             info += proc_info
+        if tostdout:
             print(info)
+        return info
 
 class MoreGPUNeededError(Exception):
     def __init__(self):

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name = 'gpuutil',
-    version = '0.0.4',
+    version = '0.0.5',
     keywords='gpu utils',
     description = 'A tool for observing gpu stat and auto set visible gpu in python code.',
     license = 'MIT License',