Repository: wookayin/gpustat

Issue:
Show full commands running
Would you support an option to show the full commands being run instead of just, e.g, `python`? These could be rather long, so one option might be to put each of them on a new line below the GPU that they're on (maybe indented a bit). Since `gpustat` seems to focus on being very concise, I'd understand if not. I have my own modification of `nvidia-smi` to show more about the processes running on each GPU, and I was planning to make it more concise when I ran across your project. I just want to check if you'd support this so I don't have to duplicate some of your work here `nividia-smi` if so.

Discussion:
wookayin: Thanks for your contribution. After all CI tests fixed, I'll have a look to merge this. BTW I am debating what would be the best command line flag (rather than `-f`): one option that comes across is `-C` (similar to `-c`), but not fully convinced myself.
bethune-bryant: Thanks @wookayin 

> After all CI tests fixed,

Yeah, the unicode characters aren't behaving nicely with Python 2.7. I'm trying to figure that out.

> BTW I am debating what would be the best command line flag (rather than `-f`): one option that comes across is `-C` (similar to `-c`), but not fully convinced myself.

I don't have strong feelings one way or another on this `-C` makes sense.


bethune-bryant: @wookayin 
All the tests are passing. I had to add a python 2 specific line.

Let me know what you decide about the flag and I can make that change.
Stonesjtu: Just FYI.
 python2.7 is retiring in Jan 1st 2020, so probably the master branch can focus on only python 3.
Ref. https://pythonclock.org/
bethune-bryant: > 
> 
> Just FYI.
> python2.7 is retiring in Jan 1st 2020, so probably the master branch can focus on only python 3.
> Ref. https://pythonclock.org/

That makes sense to me. That would definitely simplify these changes.
Stonesjtu: @wookayin 
 According to https://stackoverflow.com/questions/40114100/uploading-different-versions-python-2-7-vs-3-5-to-pypi, we can simply add `python_requires='>=3.4'` and bump the version. 
Then if you install the package via `pip2`, the out-dated one instead of the newest version will be selected.
bethune-bryant: I opened an issue (#66) and can make the Python 2 changes separately before we merge these changes.
wookayin: The current master has dropped python2 -- would you want to rebase onto master and drop changesets for python2 support if necessary?
bethune-bryant: @wookayin
Yeah, I'll make another commit to remove the python2 specific changes.
bethune-bryant: @wookayin 
I rebased this PR and removed the python 2 specific changes.

Let me know if there's anything else you'd like me to change.
bethune-bryant: > If we can cache and reuse the Process instances then it can be non-blocking via `interval=0` (for `gpustat -i` cases), but not sure for one-time run without loops. Perhaps we can also parallelize using async/await features.

I will spend some time figuring out the best way to do this.
bethune-bryant: @wookayin 
I made changes to improve the performance if the cpu usage checking, and it works really well on python 3.7. Apparently the asyncio stuff has changed alot throughout 3.4-3.7, so I'll need to look into it more.
wookayin: Looks cool. Two comments:

- In light of getting process information, the gpustat object now requires a context (e.g. pid and process) so it doesn't make any more sense to have a static initializer retrieve all the necessary information (e.g. you had to use static global variables which is not good). But I think this is a matter of refactoring I should take care of -- so let's merge it soon and let me work on this. It is also related to #61.

- Can we make sure asyncio features work well throughout python 3.4-3.7? Or how can we test it? As it requires a bit complicated mocking so writing a test seems to be a bit tricky. But it'd be great to have a minimal test; later on I will refactor test cases using smarter mocking libraries (mockito).
bethune-bryant: @wookayin 
What do you think? Is there anything else you'd like me to change before we merge?
wookayin: Looks good to me, thanks!

Patch:
diff --git a/README.md b/README.md
index fd99f79..df75f98 100644
--- a/README.md
+++ b/README.md
@@ -27,6 +27,7 @@ Options:
 * `--no-color`         : Suppress colored output
 * `-u`, `--show-user`  : Display username of the process owner
 * `-c`, `--show-cmd`   : Display the process name
+* `-f`, `--show-full-cmd`   : Display full command and cpu stats of running process
 * `-p`, `--show-pid`   : Display PID of the process
 * `-F`, `--show-fan`   : Display GPU fan speed
 * `-P`, `--show-power` : Display GPU power usage and/or limit (`draw` or `draw,limit`)
diff --git a/gpustat/__main__.py b/gpustat/__main__.py
index 127973e..2c8e1bf 100644
--- a/gpustat/__main__.py
+++ b/gpustat/__main__.py
@@ -75,6 +75,10 @@ def main(*argv):
 
     parser.add_argument('-c', '--show-cmd', action='store_true',
                         help='Display cmd name of running process')
+    parser.add_argument(
+        '-f', '--show-full-cmd', action='store_true',
+        help='Display full command and cpu stats of running process'
+    )
     parser.add_argument('-u', '--show-user', action='store_true',
                         help='Display username of running process')
     parser.add_argument('-p', '--show-pid', action='store_true',
diff --git a/gpustat/core.py b/gpustat/core.py
index a36656c..6a8ca1a 100644
--- a/gpustat/core.py
+++ b/gpustat/core.py
@@ -18,6 +18,7 @@ import locale
 import os.path
 import platform
 import sys
+import time
 from datetime import datetime
 
 from six.moves import cStringIO as StringIO
@@ -155,6 +156,7 @@ class GPUStat(object):
     def print_to(self, fp,
                  with_colors=True,    # deprecated arg
                  show_cmd=False,
+                 show_full_cmd=False,
                  show_user=False,
                  show_pid=False,
                  show_power=None,
@@ -182,9 +184,11 @@ class GPUStat(object):
         colors['CMemU'] = term.bold_yellow
         colors['CMemT'] = term.yellow
         colors['CMemP'] = term.yellow
+        colors['CCPUMemU'] = term.yellow
         colors['CUser'] = term.bold_black   # gray
         colors['CUtil'] = _conditional(lambda: self.utilization < 30,
                                        term.green, term.bold_green)
+        colors['CCPUUtil'] = term.green
         colors['CPowU'] = _conditional(
             lambda: float(self.power_draw) / self.power_limit < 0.4,
             term.magenta, term.bold_magenta
@@ -245,14 +249,40 @@ class GPUStat(object):
             )
             return r
 
+        def bytes2human(in_bytes):
+            suffixes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
+            suffix = 0
+            result = in_bytes
+            while result > 9999 and suffix < len(suffixes):
+                result = result >> 10
+                suffix += 1
+            return "%d%s" % (result, suffixes[suffix])
+
+        def full_process_info(p):
+            r = "{C0} ├─ {:>6} ".format(
+                    _repr(p['pid'], '--'), **colors
+                )
+            r += "{C0}({CCPUUtil}{:4.0f}%{C0}, {CCPUMemU}{:>6}{C0})".format(
+                    _repr(p['cpu_percent'], '--'),
+                    bytes2human(_repr(p['cpu_memory_usage'], 0)), **colors
+                )
+            r += "{C0}: {C1}{}{C0}".format(
+                    _repr(p['full_command'], '?'), **colors
+                )
+            return r
+
         processes = self.entry['processes']
+        full_processes = ''
         if processes is None:
             # None (not available)
             reps += ' ({})'.format(NOT_SUPPORTED)
         else:
             for p in processes:
                 reps += ' ' + process_repr(p)
-
+                if show_full_cmd:
+                    full_processes += '\n' + full_process_info(p)
+        if show_full_cmd:
+            reps += full_processes[::-1].replace('├', '└', 1)[::-1]
         fp.write(reps)
         return fp
 
@@ -268,6 +298,8 @@ class GPUStat(object):
 
 class GPUStatCollection(object):
 
+    global_processes = {}
+
     def __init__(self, gpu_list, driver_version=None):
         self.gpus = gpu_list
 
@@ -276,6 +308,12 @@ class GPUStatCollection(object):
         self.query_time = datetime.now()
         self.driver_version = driver_version
 
+    @staticmethod
+    def clean_processes():
+        for pid in list(GPUStatCollection.global_processes.keys()):
+            if not psutil.pid_exists(pid):
+                del GPUStatCollection.global_processes[pid]
+
     @staticmethod
     def new_query():
         """Query the information of all the GPUs on local machine"""
@@ -293,7 +331,10 @@ class GPUStatCollection(object):
             def get_process_info(nv_process):
                 """Get the process information of specific pid"""
                 process = {}
-                ps_process = psutil.Process(pid=nv_process.pid)
+                if nv_process.pid not in GPUStatCollection.global_processes:
+                    GPUStatCollection.global_processes[nv_process.pid] = \
+                        psutil.Process(pid=nv_process.pid)
+                ps_process = GPUStatCollection.global_processes[nv_process.pid]
                 process['username'] = ps_process.username()
                 # cmdline returns full path;
                 # as in `ps -o comm`, get short cmdnames.
@@ -301,10 +342,16 @@ class GPUStatCollection(object):
                 if not _cmdline:
                     # sometimes, zombie or unknown (e.g. [kworker/8:2H])
                     process['command'] = '?'
+                    process['full_command'] = '?'
                 else:
                     process['command'] = os.path.basename(_cmdline[0])
+                    process['full_command'] = " ".join(_cmdline)
                 # Bytes to MBytes
                 process['gpu_memory_usage'] = nv_process.usedGpuMemory // MB
+                process['cpu_percent'] = ps_process.cpu_percent()
+                process['cpu_memory_usage'] = \
+                    round((ps_process.memory_percent() / 100.0) *
+                          psutil.virtual_memory().total)
                 process['pid'] = nv_process.pid
                 return process
 
@@ -361,8 +408,6 @@ class GPUStatCollection(object):
                 nv_comp_processes = nv_comp_processes or []
                 nv_graphics_processes = nv_graphics_processes or []
                 for nv_process in nv_comp_processes + nv_graphics_processes:
-                    # TODO: could be more information such as system memory
-                    # usage, CPU percentage, create time etc.
                     try:
                         process = get_process_info(nv_process)
                         processes.append(process)
@@ -370,6 +415,11 @@ class GPUStatCollection(object):
                         # TODO: add some reminder for NVML broken context
                         # e.g. nvidia-smi reset  or  reboot the system
                         pass
+                time.sleep(0.1)
+                for process in processes:
+                    pid = process['pid']
+                    cache_process = GPUStatCollection.global_processes[pid]
+                    process['cpu_percent'] = cache_process.cpu_percent()
 
             index = N.nvmlDeviceGetIndex(handle)
             gpu_info = {
@@ -387,6 +437,7 @@ class GPUStatCollection(object):
                 'memory.total': memory.total // MB if memory else None,
                 'processes': processes,
             }
+            GPUStatCollection.clean_processes()
             return gpu_info
 
         # 1. get the list of gpu and status
@@ -426,9 +477,9 @@ class GPUStatCollection(object):
     # --- Printing Functions ---
 
     def print_formatted(self, fp=sys.stdout, force_color=False, no_color=False,
-                        show_cmd=False, show_user=False, show_pid=False,
-                        show_power=None, show_fan_speed=None, gpuname_width=16,
-                        show_header=True,
+                        show_cmd=False, show_full_cmd=False, show_user=False,
+                        show_pid=False, show_power=None, show_fan_speed=None,
+                        gpuname_width=16, show_header=True,
                         eol_char=os.linesep,
                         ):
         # ANSI color configuration
@@ -473,6 +524,7 @@ class GPUStatCollection(object):
         for g in self:
             g.print_to(fp,
                        show_cmd=show_cmd,
+                       show_full_cmd=show_full_cmd,
                        show_user=show_user,
                        show_pid=show_pid,
                        show_power=show_power,


Test patch:
diff --git a/gpustat/test_gpustat.py b/gpustat/test_gpustat.py
index 0f3bfb8..8081598 100644
--- a/gpustat/test_gpustat.py
+++ b/gpustat/test_gpustat.py
@@ -1,5 +1,3 @@
-# -*- coding: utf-8 -*-
-
 """
 Unit or integration tests for gpustat
 """
@@ -26,7 +24,7 @@ except ImportError:
 MagicMock = mock.MagicMock
 
 
-def _configure_mock(N, Process,
+def _configure_mock(N, Process, virtual_memory,
                     scenario_nonexistent_pid=False):
     """
     Define mock behaviour for N: the pynvml module, and psutil.Process,
@@ -136,24 +134,30 @@ def _configure_mock(N, Process,
     }.get(handle, RuntimeError))
 
     mock_pid_map = {   # mock information for psutil...
-        48448:  ('user1', 'python'),
-        154213: ('user1', 'caffe'),
-        38310:  ('user3', 'python'),
-        153223: ('user2', 'python'),
-        194826: ('user3', 'caffe'),
-        192453: ('user1', 'torch'),
+        48448:  ('user1', 'python', 85.25, 3.1415),
+        154213: ('user1', 'caffe', 16.89, 100.00),
+        38310:  ('user3', 'python', 26.23, 99.9653),
+        153223: ('user2', 'python', 15.25, 0.0000),
+        194826: ('user3', 'caffe', 0.0, 12.5236),
+        192453: ('user1', 'torch', 123.2, 0.7312),
     }
 
     def _MockedProcess(pid):
         if pid not in mock_pid_map:
             raise psutil.NoSuchProcess(pid=pid)
-        username, cmdline = mock_pid_map[pid]
+        username, cmdline, cpuutil, memutil = mock_pid_map[pid]
         p = MagicMock()  # mocked process
         p.username.return_value = username
         p.cmdline.return_value = [cmdline]
+        p.cpu_percent.return_value = cpuutil
+        p.memory_percent.return_value = memutil
         return p
     Process.side_effect = _MockedProcess
 
+    def _MockedMem():
+        return mock_memory_t(total=8589934592, used=0)
+    virtual_memory.side_effect = _MockedMem
+
 
 MOCK_EXPECTED_OUTPUT_DEFAULT = """\
 [0] GeForce GTX TITAN 0 | 80°C,  76 % |  8000 / 12287 MB | user1(4000M) user2(4000M)
@@ -163,7 +167,11 @@ MOCK_EXPECTED_OUTPUT_DEFAULT = """\
 
 MOCK_EXPECTED_OUTPUT_FULL = """\
 [0] GeForce GTX TITAN 0 | 80°C,  16 %,  76 %,  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)
+ ├─ 192453 ( 123%,   59MB): torch
+ └─ 194826 (   0%, 1025MB): caffe
 [2] GeForce GTX TITAN 2 | 71°C, 100 %,  ?? %,  250 /  ?? W |     0 / 12189 MB | (Not Supported)
 """  # noqa: E501
 
@@ -180,30 +188,33 @@ def remove_ansi_codes(s):
 
 class TestGPUStat(unittest.TestCase):
 
+    @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
     @mock.patch('gpustat.core.N')
-    def test_main(self, N, Process):
+    def test_main(self, N, Process, virtual_memory):
         """
         Test whether gpustat.main() works well. The behavior is mocked
         exactly as in test_new_query_mocked().
         """
-        _configure_mock(N, Process)
+        _configure_mock(N, Process, virtual_memory)
         sys.argv = ['gpustat']
         gpustat.main()
 
+    @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
     @mock.patch('gpustat.core.N')
-    def test_new_query_mocked(self, N, Process):
+    def test_new_query_mocked(self, N, Process, virtual_memory):
         """
         A basic functionality test, in a case where everything is just normal.
         """
-        _configure_mock(N, Process)
+        _configure_mock(N, Process, virtual_memory)
 
         gpustats = gpustat.new_query()
         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_cmd=True, show_pid=True, show_power=True, show_fan_speed=True,
+            show_full_cmd=True
         )
 
         result = fp.getvalue()
@@ -216,24 +227,28 @@ class TestGPUStat(unittest.TestCase):
         self.maxDiff = 4096
         self.assertEqual(unescaped, MOCK_EXPECTED_OUTPUT_FULL)
 
+    @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
     @mock.patch('gpustat.core.N')
-    def test_new_query_mocked_nonexistent_pid(self, N, Process):
+    def test_new_query_mocked_nonexistent_pid(self, N, Process,
+                                              virtual_memory):
         """
         Test a case where nvidia query returns non-existent pids (see #16, #18)
         """
-        _configure_mock(N, Process, scenario_nonexistent_pid=True)
+        _configure_mock(N, Process, virtual_memory,
+                        scenario_nonexistent_pid=True)
 
         gpustats = gpustat.new_query()
         gpustats.print_formatted(fp=sys.stdout)
 
+    @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
     @mock.patch('gpustat.core.N')
-    def test_attributes_and_items(self, N, Process):
+    def test_attributes_and_items(self, N, Process, virtual_memory):
         """
         Test whether each property of `GPUStat` instance is well-defined.
         """
-        _configure_mock(N, Process)
+        _configure_mock(N, Process, virtual_memory)
 
         g = gpustat.new_query()[1]  # includes N/A
         print("(keys) : %s" % str(g.keys()))
@@ -253,13 +268,14 @@ class TestGPUStat(unittest.TestCase):
         print("utilization : %s" % (g.utilization))
 
     @unittest.skipIf(sys.version_info < (3, 4), "Only in Python 3.4+")
+    @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
     @mock.patch('gpustat.core.N')
-    def test_args_endtoend(self, N, Process):
+    def test_args_endtoend(self, N, Process, virtual_memory):
         """
         End-to-end testing given command line args.
         """
-        _configure_mock(N, Process)
+        _configure_mock(N, Process, virtual_memory)
 
         def capture_output(*args):
             f = StringIO()
@@ -284,10 +300,11 @@ class TestGPUStat(unittest.TestCase):
         s = capture_output('gpustat', '--no-header')
         self.assertIn("[0]", s.split('\n')[0])
 
+    @mock.patch('psutil.virtual_memory')
     @mock.patch('psutil.Process')
     @mock.patch('gpustat.core.N')
-    def test_json_mocked(self, N, Process):
-        _configure_mock(N, Process)
+    def test_json_mocked(self, N, Process, virtual_memory):
+        _configure_mock(N, Process, virtual_memory)
         gpustats = gpustat.new_query()
 
         fp = StringIO()