diff --git a/README.md b/README.md index 4137911..e2a3d29 100644 --- a/README.md +++ b/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. diff --git a/nvidia-htop.py b/nvidia-htop.py index 3d9a5dc..af53528 100755 --- a/nvidia-htop.py +++ b/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] diff --git a/test/test_main.py b/test/test_main.py index d19a0d2..c6dea28 100644 --- a/test/test_main.py +++ b/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')