diff --git a/gpuutil/__main__.py b/gpuutil/__main__.py index 02ef1e7..a4a3c52 100644 --- a/gpuutil/__main__.py +++ b/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) \ No newline at end of file + 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) \ No newline at end of file diff --git a/gpuutil/gpuutil.py b/gpuutil/gpuutil.py index f27151a..9d18891 100644 --- a/gpuutil/gpuutil.py +++ b/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: - proc['vmem'] = int(proc['vmem'].split(' ')[0]) + 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 - print(info) + if tostdout: + print(info) + return info class MoreGPUNeededError(Exception): def __init__(self): diff --git a/setup.py b/setup.py index f9256cc..7ef1a9e 100644 --- a/setup.py +++ b/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',