Add command-line support for enc/dec gpu utilization

wookayin · Mar 20, 2020 · dfc2ecb · dfc2ecb
1 parent f20fd98
commit dfc2ecb
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 15 deletions.
diff --git a/gpustat/cli.py b/gpustat/cli.py
@@ -93,15 +93,18 @@ def main(*argv):
                         help='Display PID of running process')
     parser.add_argument('-F', '--show-fan-speed', '--show-fan',
                         action='store_true', help='Display GPU fan speed')
-    parser.add_argument('--json', action='store_true', default=False,
-                        help='Print all the information in JSON format')
-    parser.add_argument('-v', '--version', action='version',
-                        version=('gpustat %s' % __version__))
+    parser.add_argument(
+        '-e', '--show-codec', nargs='?', const='enc,dec', default='',
+        choices=['enc', 'dec', 'enc,dec'],
+        help='Show encoder/decoder utilization'
+    )
     parser.add_argument(
         '-P', '--show-power', nargs='?', const='draw,limit',
         choices=['', 'draw', 'limit', 'draw,limit', 'limit,draw'],
         help='Show GPU power usage or draw (and/or limit)'
     )
+    parser.add_argument('--json', action='store_true', default=False,
+                        help='Print all the information in JSON format')
     parser.add_argument(
         '-i', '--interval', '--watch', nargs='?', type=float, default=0,
         help='Use watch mode if given; seconds to wait between updates'
@@ -118,12 +121,15 @@ def main(*argv):
         '--debug', action='store_true', default=False,
         help='Allow to print additional informations for debugging.'
     )
+    parser.add_argument('-v', '--version', action='version',
+                        version=('gpustat %s' % __version__))
     args = parser.parse_args(argv[1:])
     if args.show_all:
         args.show_cmd = True
         args.show_user = True
         args.show_pid = True
         args.show_fan_speed = True
+        args.show_codec = 'enc,dec'
         args.show_power = 'draw,limit'
     del args.show_all
 

diff --git a/gpustat/core.py b/gpustat/core.py
@@ -181,8 +181,9 @@ def print_to(self, fp,
                  show_full_cmd=False,
                  show_user=False,
                  show_pid=False,
-                 show_power=None,
                  show_fan_speed=None,
+                 show_codec="",
+                 show_power=None,
                  gpuname_width=16,
                  term=None,
                  ):
@@ -199,8 +200,11 @@ def _conditional(cond_fn, true_value, false_value,
             except Exception:
                 return error_value
 
+        _ENC_THRESHOLD = 50
+
         colors['C0'] = term.normal
         colors['C1'] = term.cyan
+        colors['CBold'] = term.bold
         colors['CName'] = term.blue
         colors['CTemp'] = _conditional(lambda: self.temperature < 50,
                                        term.red, term.bold_red)
@@ -213,6 +217,12 @@ def _conditional(cond_fn, true_value, false_value,
         colors['CUser'] = term.bold_black   # gray
         colors['CUtil'] = _conditional(lambda: self.utilization < 30,
                                        term.green, term.bold_green)
+        colors['CUtilEnc'] = _conditional(
+            lambda: self.utilization_enc < _ENC_THRESHOLD,
+            term.green, term.bold_green)
+        colors['CUtilDec'] = _conditional(
+            lambda: self.utilization_dec < _ENC_THRESHOLD,
+            term.green, term.bold_green)
         colors['CCPUUtil'] = term.green
         colors['CPowU'] = _conditional(
             lambda: float(self.power_draw) / self.power_limit < 0.4,
@@ -239,6 +249,14 @@ def _repr(v, none_value='??'):
             reps += "%(FSpeed)s{entry[fan.speed]:>3} %%%(C0)s, "
 
         reps += "%(CUtil)s{entry[utilization.gpu]:>3} %%%(C0)s"
+        if "enc" in show_codec:
+            reps += (
+                "  %(CBold)sE: %(C0)s"
+                "%(CUtilEnc)s{entry[utilization.enc]:>3} %%%(C0)s")
+        if "dec" in show_codec:
+            reps += (
+                "  %(CBold)sD: %(C0)s"
+                "%(CUtilDec)s{entry[utilization.dec]:>3} %%%(C0)s")
 
         if show_power:
             reps += ",  %(CPowU)s{entry[power.draw]:>3}%(C0)s "
@@ -518,7 +536,8 @@ def __repr__(self):
 
     def print_formatted(self, fp=sys.stdout, force_color=False, no_color=False,
                         show_cmd=False, show_full_cmd=False, show_user=False,
-                        show_pid=False, show_power=None, show_fan_speed=None,
+                        show_pid=False, show_fan_speed=None,
+                        show_codec="", show_power=None,
                         gpuname_width=16, show_header=True,
                         eol_char=os.linesep,
                         ):
@@ -572,8 +591,9 @@ def print_formatted(self, fp=sys.stdout, force_color=False, no_color=False,
                        show_full_cmd=show_full_cmd,
                        show_user=show_user,
                        show_pid=show_pid,
-                       show_power=show_power,
                        show_fan_speed=show_fan_speed,
+                       show_codec=show_codec,
+                       show_power=show_power,
                        gpuname_width=gpuname_width,
                        term=t_color)
             fp.write(eol_char)

diff --git a/gpustat/test_gpustat.py b/gpustat/test_gpustat.py
@@ -179,19 +179,19 @@ def _MockedMem():
 """.splitlines())  # noqa: E501
 
 MOCK_EXPECTED_OUTPUT_FULL = os.linesep.join("""\
-[0] GeForce GTX TITAN 0 | 80°C,  16 %,  76 %,  125 / 250 W |  8000 / 12287 MB | user1:python/48448(4000M) user2:python/153223(4000M)
-[1] GeForce GTX TITAN 1 | 36°C,  53 %,   0 %,   ?? / 250 W |  9000 / 12189 MB | user1:torch/192453(3000M) user3:caffe/194826(6000M)
-[2] GeForce GTX TITAN 2 | 71°C, 100 %,  ?? %,  250 /  ?? W |     0 / 12189 MB | (Not Supported)
+[0] GeForce GTX TITAN 0 | 80°C,  16 %,  76 %  E:  88 %  D:  67 %,  125 / 250 W |  8000 / 12287 MB | user1:python/48448(4000M) user2:python/153223(4000M)
+[1] GeForce GTX TITAN 1 | 36°C,  53 %,   0 %  E:   0 %  D:   0 %,   ?? / 250 W |  9000 / 12189 MB | user1:torch/192453(3000M) user3:caffe/194826(6000M)
+[2] GeForce GTX TITAN 2 | 71°C, 100 %,  ?? %  E:  ?? %  D:  ?? %,  250 /  ?? W |     0 / 12189 MB | (Not Supported)
 """.splitlines())  # noqa: E501
 
 MOCK_EXPECTED_OUTPUT_FULL_PROCESS = os.linesep.join("""\
-[0] GeForce GTX TITAN 0 | 80°C,  16 %,  76 %,  125 / 250 W |  8000 / 12287 MB | user1:python/48448(4000M) user2:python/153223(4000M)
+[0] GeForce GTX TITAN 0 | 80°C,  16 %,  76 %  E:  88 %  D:  67 %,  125 / 250 W |  8000 / 12287 MB | user1:python/48448(4000M) user2:python/153223(4000M)
  ├─  48448 (  85%,  257MB): python
  └─ 153223 (  15%,     0B): python
-[1] GeForce GTX TITAN 1 | 36°C,  53 %,   0 %,   ?? / 250 W |  9000 / 12189 MB | user1:torch/192453(3000M) user3:caffe/194826(6000M)
+[1] GeForce GTX TITAN 1 | 36°C,  53 %,   0 %  E:   0 %  D:   0 %,   ?? / 250 W |  9000 / 12189 MB | user1:torch/192453(3000M) user3:caffe/194826(6000M)
  ├─ 192453 ( 123%,   59MB): torch
  └─ 194826 (   0%, 1025MB): caffe
-[2] GeForce GTX TITAN 2 | 71°C, 100 %,  ?? %,  250 /  ?? W |     0 / 12189 MB | (Not Supported)
+[2] GeForce GTX TITAN 2 | 71°C, 100 %,  ?? %  E:  ?? %  D:  ?? %,  250 /  ?? W |     0 / 12189 MB | (Not Supported)
 """.splitlines())  # noqa: E501
 
 
@@ -232,8 +232,8 @@ def test_new_query_mocked(self, N, Process, virtual_memory):
         fp = StringIO()
         gpustats.print_formatted(
             fp=fp, no_color=False, show_user=True,
-            show_cmd=True, show_pid=True, show_power=True, show_fan_speed=True,
-            show_full_cmd=True
+            show_cmd=True, show_full_cmd=True, show_pid=True,
+            show_fan_speed=True, show_codec="enc,dec", show_power=True,
         )
 
         result = fp.getvalue()