Added possibility to filter processes by user(s).

Closes #20.

README.md

@@ -22,6 +22,8 @@ A tool for enriching the output of `nvidia-smi`.
                                        otherwise print first 100 characters.
       -c|--color                       Colorize the output (green - free GPU, yellow -
                                        moderately used GPU, red - fully used GPU)
+      -u|--user USER[,USER]            Limit the list of processes to selected users
+                                       (comma-separated).
 
 Note: for backward compatibility, `nvidia-smi | nvidia-htop.py [-l [length]]` is also supported.
 

nvidia-htop.py

@@ -30,6 +30,7 @@
 parser = argparse.ArgumentParser()
 parser.add_argument('-l', '--command-length', default=20, const=100, type=int, nargs='?')
 parser.add_argument('-c', '--color', action='store_true')
+parser.add_argument('-u', '--user', default='', help="Limit the list of processes to selected users (comma-separated)")
 # only for testing
 parser.add_argument('-p', '--fake-ps', help="The list of processes to use instead of real output of `ps`")
 
@@ -39,6 +40,7 @@
 command_length = args.command_length
 color = args.color
 fake_ps = args.fake_ps
+users = set(args.user.split(',')) if len(args.user) > 0 else None
 
 # for testing, the stdin can be provided in a file
 fake_stdin_path = os.getenv("FAKE_STDIN_PATH", None)
@@ -132,6 +134,17 @@ def colorize(_lines):
 time = []
 command = []
 
+fields = (
+    gpu_num,
+    pid,
+    gpu_mem,
+    user,
+    cpu,
+    mem,
+    time,
+    command,
+)
+
 gpu_num_idx = 1
 pid_idx = 2 if not is_new_format else 4
 gpu_mem_idx = -3
@@ -167,12 +180,20 @@ def colorize(_lines):
         continue
     parts = re.split(r'\s+', line.strip(), 5)
     # idx = pid.index(parts[0])
+    to_delete = []  # If the command is limited to selected users, we need to delete the other lines
     for idx in filter(lambda p: pid[p] == parts[0], range(len(pid))):
+        if users is not None and parts[1] not in users:
+            to_delete.append(idx)
+            continue
         user[idx] = parts[1]
         cpu[idx] = parts[2]
         mem[idx] = parts[3]
         time[idx] = parts[4] if "-" not in parts[4] else parts[4].split("-")[0] + " days"
         command[idx] = parts[5]
+    # Delete lines not corresponding to the selected users (if some are selected)
+    for idx in reversed(sorted(to_delete)):
+        for field in fields:
+            del field[idx]
 
 max_pid_length = max(5, max([len(x) for x in pid]))
 format = ("|  %3s %" + str(max_pid_length) + "s %8s   %8s %5s %5s %9s  %-" + str(command_length) + "." + str(command_length) + "s  |")

test/DESIRED_STDOUT_NEW_FORMAT_USERS

@@ -0,0 +1,24 @@
+Sun Aug  2 13:44:21 2020
++-----------------------------------------------------------------------------+
+| NVIDIA-SMI 450.57       Driver Version: 450.57       CUDA Version: 11.0     |
+|-------------------------------+----------------------+----------------------+
+| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
+|                               |                      |               MIG M. |
+|===============================+======================+======================|
+|   0  GeForce RTX 2070    Off  | 00000000:01:00.0 Off |                  N/A |
+|  0%   50C    P2    27W / 175W |    749MiB /  7981MiB |      2%      Default |
+|                               |                      |                  N/A |
++-------------------------------+----------------------+----------------------+
+|   1  GeForce GTX 750 Ti  Off  | 00000000:04:00.0  On |                  N/A |
+| 33%   30C    P8     1W /  46W |    871MiB /  2002MiB |      0%      Default |
+|                               |                      |                  N/A |
++-------------------------------+----------------------+----------------------+
+
++-----------------------------------------------------------------------------+
+|  GPU   PID     USER    GPU MEM  %CPU  %MEM      TIME  COMMAND               |
+|    0  1032     root     745MiB  10.0   5.0  11:42:17  python 0.py           |
+|    1 25544     test     139MiB   5.0   3.0  13:42:17  python 2.py           |
+|    1 14518     root       1MiB   8.0   1.0  15:42:17  python3 3.py          |
+|    1 13956     root     472MiB   2.0   0.0  16:42:17  python3 4.py          |
++-----------------------------------------------------------------------------+

test/test_main.py

@@ -22,6 +22,9 @@ def test_with_processes(self):
     def test_new_format(self):
         self.do_test('FAKE_STDIN_NEW_FORMAT', 'DESIRED_STDOUT_NEW_FORMAT')
 
+    def test_new_format_users(self):
+        self.do_test('FAKE_STDIN_NEW_FORMAT', 'DESIRED_STDOUT_NEW_FORMAT_USERS', call_args=["-u", "root,test"])
+
     def test_long_pids(self):
         self.do_test('FAKE_STDIN_LONG_PIDS', 'DESIRED_STDOUT_LONG_PIDS', fake_ps='FAKE_PS_LONG_PIDS')