Merge pull request #22

Added the possibility to filter by GPU IDs

Author: peci1

README.md

@@ -24,6 +24,7 @@ A tool for enriching the output of `nvidia-smi`.
                                        moderately used GPU, red - fully used GPU)
       -u|--user USER[,USER]            Limit the list of processes to selected users
                                        (comma-separated).
+      -i|--id ID[,ID[,ID]]             Limit the command to selected GPU IDs (comma-separated).
 
 Note: for backward compatibility, `nvidia-smi | nvidia-htop.py [-l [length]]` is also supported.
 

nvidia-htop.py

@@ -3,14 +3,15 @@
 #######
 # USAGE
 #
-# [nvidia-smi | ] nvidia-htop.py [-l [length]]
+# [nvidia-smi | ] nvidia-htop.py [-l [length]] [-i ID]
 #   print GPU utilization with usernames and CPU stats for each GPU-utilizing process
 #
 #   -l|--command-length [length]     Print longer part of the commandline. If `length'
 #                                    is provided, use it as the commandline length,
 #                                    otherwise print first 100 characters.
 #   -c|--color                       Colorize the output (green - free GPU, yellow -
-#                                    moderately used GPU, red - fully used GPU)
+#                                    moderately used GPU, red - fully used GPU).
+#   -i|--id ID[,ID[,ID...]]          Limit the command to selected GPU IDs (comma-separated).
 ######
 
 import sys
@@ -31,6 +32,7 @@
 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)")
+parser.add_argument('-i', '--id', default='', help="Limit the command to selected GPU IDs (comma-separated)")
 # only for testing
 parser.add_argument('-p', '--fake-ps', help="The list of processes to use instead of real output of `ps`")
 
@@ -53,12 +55,18 @@
         lines = f.readlines()
 elif stdin_lines:
     lines = stdin_lines
+    if len(args.id) > 0:
+        print('nvidia-htop argument -i/--id cannot be used when nvidia-smi output is being piped into it. To filter the'
+              ' shown GPUs, pass the -i argument to the nvidia-smi call instead.', file=sys.stderr)
 else:
-    ps_call = subprocess.run('nvidia-smi', stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    nvidiasmi_args = []
+    if len(args.id) > 0:
+        nvidiasmi_args = ['-i', args.id]
+    ps_call = subprocess.run(['nvidia-smi'] + nvidiasmi_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     if ps_call.returncode != 0:
-        print('nvidia-smi exited with error code {}:'.format(ps_call.returncode))
-        print(ps_call.stdout.decode() + ps_call.stderr.decode())
-        sys.exit()
+        print('nvidia-smi exited with error code {}:'.format(ps_call.returncode), file=sys.stderr)
+        print(ps_call.stdout.decode() + ps_call.stderr.decode(), file=sys.stderr)
+        sys.exit(ps_call.returncode)
     lines_proc = ps_call.stdout.decode().split("\n")
     lines = [line + '\n' for line in lines_proc[:-1]]
     lines += lines_proc[-1]

test/test_main.py

@@ -25,6 +25,10 @@ def test_new_format(self):
     def test_new_format_users(self):
         self.do_test('FAKE_STDIN_NEW_FORMAT', 'DESIRED_STDOUT_NEW_FORMAT_USERS', call_args=["-u", "root,test"])
 
+    # The --id option cannot be tested in this way. So we just check that the option is considered valid.
+    def test_new_format_filter_ids(self):
+        self.do_test('FAKE_STDIN_NEW_FORMAT', 'DESIRED_STDOUT_NEW_FORMAT', call_args=["-i", "1,2"])
+
     def test_long_pids(self):
         self.do_test('FAKE_STDIN_LONG_PIDS', 'DESIRED_STDOUT_LONG_PIDS', fake_ps='FAKE_PS_LONG_PIDS')