From 52ea10789fa568eb55c07f7ec8edfb8d1b003740 Mon Sep 17 00:00:00 2001 From: Martin Pecka Date: Thu, 4 Jan 2024 23:29:17 +0100 Subject: [PATCH 1/3] Added possibility to filter processes by user(s). Closes #20. --- README.md | 2 ++ nvidia-htop.py | 21 +++++++++++++++++++++ test/DESIRED_STDOUT_NEW_FORMAT_USERS | 24 ++++++++++++++++++++++++ test/test_main.py | 3 +++ 4 files changed, 50 insertions(+) create mode 100644 test/DESIRED_STDOUT_NEW_FORMAT_USERS diff --git a/README.md b/README.md index 5a6f055..4137911 100644 --- a/README.md +++ b/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. diff --git a/nvidia-htop.py b/nvidia-htop.py index a53860a..997c1df 100755 --- a/nvidia-htop.py +++ b/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 |") diff --git a/test/DESIRED_STDOUT_NEW_FORMAT_USERS b/test/DESIRED_STDOUT_NEW_FORMAT_USERS new file mode 100644 index 0000000..85e9252 --- /dev/null +++ b/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 | ++-----------------------------------------------------------------------------+ diff --git a/test/test_main.py b/test/test_main.py index 3b06c54..d19a0d2 100644 --- a/test/test_main.py +++ b/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') From 0ab2f97a8a477bd116c07e539be6f7001ffe3e12 Mon Sep 17 00:00:00 2001 From: Martin Pecka Date: Fri, 5 Jan 2024 00:20:22 +0100 Subject: [PATCH 2/3] Fix case with no processes left after filtering by user. --- nvidia-htop.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nvidia-htop.py b/nvidia-htop.py index 997c1df..3d9a5dc 100755 --- a/nvidia-htop.py +++ b/nvidia-htop.py @@ -195,6 +195,10 @@ def colorize(_lines): for field in fields: del field[idx] +if len(pid) == 0: + print("| " + no_running_process + " " * (73 - len(no_running_process)) + " |") + sys.exit() + 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 |") From 14ab7542f0f5c8ea537909d230288b4885421354 Mon Sep 17 00:00:00 2001 From: Martin Pecka Date: Fri, 5 Jan 2024 00:22:43 +0100 Subject: [PATCH 3/3] 1.0.7 --- CHANGELOG.md | 4 ++++ setup.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 82277dc..cf7f330 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## 1.0.7 + +- Added CLI argument `--user` to filter processes by selected users. + ## 1.0.6 - Fixed coloring output. Thanks @tomix1024 ! diff --git a/setup.py b/setup.py index f926f92..d29f2fe 100644 --- a/setup.py +++ b/setup.py @@ -9,7 +9,7 @@ long_description = (here / 'README.md').read_text(encoding='utf-8') setup(name='nvidia-htop', - version='1.0.6', + version='1.0.7', description='A tool for enriching the output of nvidia-smi', long_description=long_description, long_description_content_type='text/markdown',