update display style, add customlize support

gpuutil/__main__.py

@@ -1,18 +1,67 @@
 from gpuutil import GPUStat
 import sys
+import json
+import argparse
+import os
+
+def csv2list(csv):
+    l = [col.strip() for col in csv.split(',')]
+    return [col for col in l if col != '']
+def str2bool(s):
+    if s.lower() in ['t', 'yes', 'y', 'aye', 'positive', 'true']:
+        return True
+    else:
+        return False
+
+def load_config():
+    home_dir = os.path.expanduser('~')
+    configpath = os.path.join(home_dir, '.gpuutil.conf')
+    if not os.path.isfile(configpath):
+        return {}
+    with open(configpath, 'r', encoding='utf-8') as f:
+        return json.load(f)
+def save_config(config):
+    home_dir = os.path.expanduser('~')
+    configdir = os.path.join(home_dir, '.gpuutil.conf')
+    with open(configdir, 'w+', encoding='utf-8') as f:
+        json.dump(config, f, ensure_ascii=False, indent=4)
 
 if __name__ == '__main__':
     stat = GPUStat()
-    show_types = ['brief', 'detail']
-    default_type = 'brief'
-    show_type = default_type
-    if len(sys.argv) > 1:
-        show_type = str(sys.argv[1])
-    if show_type in show_types:
-        stat.show(disp_type=show_type)
-    else:
-        print('The given type is \"{0}\" not understood, and it should be choosen from {1}\nUsing default type \"{2}\".'.format(show_type, show_types, default_type))
-        show_type = default_type
-        stat.show(disp_type=show_type)
 
-    # auto_set(1, ask=True, blacklist=[], show=True)
+    avaliable_cols = ['ID', 'Fan', 'Temp', 'TempMax', 'Pwr', 'PwrMax', 'Freq', 'FreqMax', 'Util', 'Vmem', 'UsedMem', 'TotalMem', 'FreeMem', 'Users']
+    recommended_cols = ['ID', 'Fan', 'Temp', 'Pwr', 'Freq', 'Util', 'Vmem']
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--profile', '-p', default=None, type=str, help='profile keyword, corresponding configuration are saved in ~/.gpuutil.conf')
+    parser.add_argument('--cols', '-c', type=csv2list, help='colums to show')
+    parser.add_argument('--show-process', '-sp', default=True, type=str2bool, help='whether show process or not')
+    parser.add_argument('--save', default=False, action="store_true", 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
+    unexpected_cols = []
+    for col in cols:
+        if col not in avaliable_cols:
+            unexpected_cols.append(col)
+    if len(unexpected_cols) > 0:
+        raise ValueError('Unexpected cols {0} occured. Cols must be chosen from {1}'.format(unexpected_cols, ','.join(avaliable_cols)))
+
+    if args.save:
+        params = {
+            "cols": cols,
+            "show-process": show_process
+        }
+        profile = args.profile if args.profile is not None else input('Please input your profile name:\n>>> ')
+        config = load_config()
+        config[profile] = params
+        save_config(config)
+    elif args.profile is not None:
+        config = load_config()
+        if args.profile in config:
+            params = config[args.profile]
+            cols = params["cols"]
+            show_process = params["show-process"]
+        else:
+            raise ValueError('Profile do not exist.\nAvaliable Profiles:{0}'.format(','.join(list(config.keys()))))
+    stat.show(enabled_cols = cols, show_command=show_process)

gpuutil/gpuutil.py

@@ -133,32 +133,6 @@ def short_gpu_info(stat, disp_type='brief'):
         info += process_info
     return info
     
-def show_gpu_info_v2(stat, disp_type='brief'):
-    # for gpu information.
-    stats = {
-        "id": stat['id'],
-        "fan": stat['fan_speed'].split(' ')[0].strip(),
-        "temp_cur": stat['temperature']['current'].split(' ')[0].strip(),
-        "temp_max": stat['temperature']['max'].split(' ')[0].strip(),
-        "power_cur": stat['power']['current'].split(' ')[0].strip(),
-        "power_max": stat['power']['max'].split(' ')[0].strip(),
-        "clock_cur": stat['clocks']['current'].split(' ')[0].strip(),
-        "clock_max": stat['clocks']['max'].split(' ')[0].strip(),
-        "util": stat['utilization'],
-        "mem_used": stat['memory']['used'].split(' ')[0].strip(),
-        "mem_total": stat['memory']['total'].split(' ')[0].strip(),
-        "mem_free": stat['memory']['free'].split(' ')[0].strip()
-    }
-
-    gpu_stat_fmt = "{id} | {fan} | {temp_cur} | {power_cur} | {clock_cur} | {util} | {mem_used}/{mem_all}"
-    info = gpu_stat_fmt
-    for key in stats:
-        placeholder = '{%s}'%key
-        if info.find(placeholder) != -1:
-            info.replace(placeholder, stats[key])
-    return info
-    
-    
 
 def get_basic_process_info_linux():
     pipe = os.popen('ps axo user:20,pid,args:1024')
@@ -194,6 +168,51 @@ def get_basic_process_info_windows():
         }
     return processes
 
+def draw_table(table, header_line = 0, c_align = 'r', h_align='c', delemeter = ' | ', joint_delemeter = '-+-'):
+    # calculate max lengths.
+    num_columns = len(table[0])
+    def cvt_align(align, num_columns):
+        if type(align) is str:
+            if len(align) == 1:
+                return [align] * num_columns
+            elif len(align) == num_columns:
+                return list(align)
+            else:
+                raise ValueError('align flag length mismatch')
+        else:
+            return align
+    c_align = cvt_align(c_align, num_columns)
+    h_align = cvt_align(h_align, num_columns)
+    max_lengths = [0] * num_columns
+    for row in table:
+        for i, col in enumerate(row):
+            if len(col) > max_lengths[i]:
+                max_lengths[i] = len(col)
+    width = sum(max_lengths) + num_columns * len(delemeter) + 1
+    hline = '+'
+    hline += joint_delemeter.join(['-' * length for length in max_lengths])
+    hline += '+\n'
+    info =  hline
+    for i, row in enumerate(table):
+        info += '|'
+        row_just = []
+        align = h_align if i <= header_line else c_align
+        for w, col, a in zip(max_lengths, row, align):
+            if a == 'c':
+                row_just.append(col.center(w))
+            elif a == 'l':
+                row_just.append(col.ljust(w))
+            elif a == 'r':
+                row_just.append(col.rjust(w))
+        info += delemeter.join(row_just)
+        info += '|\n'
+        if i == header_line:
+            info += hline
+    info +=  hline
+    return info
+    
+
+
 class GPUStat():
     def __init__(self):
         self.gpus = []
@@ -236,31 +255,83 @@ class GPUStat():
             gpu['id'] = i
             self.gpus.append(gpu)
 
-    def show(self, disp_type='brief', command=True):
+    def show(self, enabled_cols = ['ID', 'Fan', 'Temp', 'Pwr', 'Freq', 'Util', 'Vmem', 'Users'], show_command=True):
         self.parse()
-        lines = [short_gpu_info(info, disp_type=disp_type) for info in self.gpus]
+        gpu_infos = []
-        print('================== GPU INFO ==================')
+        # stats = {
-        print('\n'.join(lines))
+        #     "id": stat['id'],
-        if command:
+        #     "fan": stat['fan_speed'].split(' ')[0].strip(),
-            print('================ PROCESS INFO ================')
+        #     "temp_cur":  stat['temperature']['current'].split(' ')[0].strip(),
+        #     "temp_max":  stat['temperature']['max'].split(' ')[0].strip(),
+        #     "power_cur": stat['power']['current'].split(' ')[0].strip(),
+        #     "power_max": stat['power']['max'].split(' ')[0].strip(),
+        #     "clock_cur": stat['clocks']['current'].split(' ')[0].strip(),
+        #     "clock_max": stat['clocks']['max'].split(' ')[0].strip(),
+        #     "util":      stat['utilization'],
+        #     "mem_used":  stat['memory']['used'].split(' ')[0].strip(),
+        #     "mem_total": stat['memory']['total'].split(' ')[0].strip(),
+        #     "mem_free":  stat['memory']['free'].split(' ')[0].strip()
+        # }
+        for gpu in self.gpus:
+            process_fmt = '{user}({pid})'
+            process_info = ','.join([process_fmt.format(
+                user = proc['user'],
+                pid = proc['pid']
+            ) for proc in gpu['processes']])
+            info_gpu = {
+                'ID': '{0}'.format(str(gpu['id'])),
+                'Fan': '{0} %'.format(gpu['fan_speed'].split(' ')[0].strip()),
+                'Temp': '{0} C'.format(gpu['temperature']['current'].split(' ')[0].strip()),
+                'TempMax': '{0} C'.format(gpu['temperature']['max'].split(' ')[0].strip()),
+                'Pwr': '{0} W'.format(gpu['power']['current'].split(' ')[0].strip()),
+                'PwrMax': '{0} W'.format(gpu['power']['max'].split(' ')[0].strip()),
+                'Freq': '{0} MHz'.format(gpu['clocks']['current'].split(' ')[0].strip()),
+                'FreqMax': '{0} MHz'.format(gpu['clocks']['max'].split(' ')[0].strip()),
+                'Util': '{0} %'.format(gpu['utilization'].split(' ')[0]),
+                'Vmem': '{0}/{1} MiB'.format(
+                    gpu['memory']['used'].split(' ')[0].strip(),
+                    gpu['memory']['total'].split(' ')[0].strip(),
+                ),
+                'UsedMem': '{0} MiB'.format(gpu['memory']['used'].split(' ')[0].strip()),
+                'TotalMem': '{0} MiB'.format(gpu['memory']['total'].split(' ')[0].strip()),
+                'FreeMem': '{0} MiB'.format(gpu['memory']['free'].split(' ')[0].strip()),
+                'Users': process_info
+            }
+            gpu_infos.append(info_gpu)
+        align_methods = {key:'r' for key in gpu_infos[0]}
+        align_methods['Users'] = 'l'
+        if enabled_cols is None:
+            enabled_cols = list(align_methods.keys())
+        c_align = [align_methods[col] for col in enabled_cols]
+        info_table = [enabled_cols]
+        for info in gpu_infos:
+            this_row = [info[key] for key in enabled_cols]
+            info_table.append(this_row)
+        info = draw_table(info_table, header_line=0, delemeter=' | ', joint_delemeter='-+-', c_align=c_align)
+        if show_command:
             procs = {}
             for gpu in self.gpus:
                 for proc in gpu['processes']:
                     pid = proc['pid']
+                    proc['gpu'] = [str(gpu['id'])]
                     if pid not in procs:
                         procs[pid] = proc
-            proc_fmt = '[{pid}] {user}({vmem} MiB) {cmd}'
+                    else:
+                        procs[pid].append(gpu['id'])
+            proc_fmt = '[{pid}|{gpus}] {user}({vmem} MiB) {cmd}'
             proc_strs = []
             for pid in procs:
                 this_proc_str = proc_fmt.format(
                     user = procs[pid]['user'],
                     vmem = procs[pid]['vmem'].split(' ')[0],
-                    pid = procs[pid]['pid'], 
+                    pid = procs[pid]['pid'].rjust(5), 
-                    cmd = procs[pid]['command']
+                    cmd = procs[pid]['command'],
+                    gpus = ','.join(procs[pid]['gpu'])
                 )
                 proc_strs.append(this_proc_str)
             proc_info = '\n'.join(proc_strs)
-            print(proc_info)
+            info += proc_info
+        print(info)
 
 class MoreGPUNeededError(Exception):
     def __init__(self):