feat(cli): add GPU process context filter

Signed-off-by: Xuehai Pan <XuehaiPan@pku.edu.cn>
This commit is contained in:
Xuehai Pan 2022-08-01 19:33:42 +08:00
parent 9d796d889b
commit 23f581217b
2 changed files with 27 additions and 3 deletions

View file

@ -280,7 +280,8 @@ Type `nvitop --help` for more command options:
usage: nvitop [--help] [--version] [--once] [--monitor [{auto,full,compact}]]
[--interval SEC] [--ascii] [--colorful] [--force-color] [--light]
[--gpu-util-thresh th1 th2] [--mem-util-thresh th1 th2]
[--only idx [idx ...]] [--only-visible] [--compute] [--graphics]
[--only idx [idx ...]] [--only-visible]
[--compute] [--no-compute] [--graphics] [--no-graphics]
[--user [USERNAME ...]] [--pid PID [PID ...]]
An interactive NVIDIA-GPU process viewer.
@ -321,7 +322,9 @@ device filtering:
process filtering:
--compute, -c Only show GPU processes with the compute context. (type: 'C' or 'C+G')
--no-compute, -C Exclude GPU processes with the compute context. (type: 'G' only)
--graphics, -g Only show GPU processes with the graphics context. (type: 'G' or 'C+G')
--no-graphics, -G Exclude GPU processes with the graphics context. (type: 'C' only)
--user [USERNAME ...], -u [USERNAME ...]
Only show processes of the given users (or `$USER` for no argument).
--pid PID [PID ...], -p PID [PID ...]

View file

@ -52,10 +52,12 @@ def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements
version='%(prog)s {}'.format(__version__),
help="Show %(prog)s's version number and exit.",
)
parser.add_argument(
mode = parser.add_mutually_exclusive_group()
mode.add_argument(
'--once', '-1', dest='once', action='store_true', help='Report query data only once.'
)
parser.add_argument(
mode.add_argument(
'--monitor',
'-m',
dest='monitor',
@ -69,6 +71,7 @@ def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements
'(default fallback mode: auto)'
),
)
parser.add_argument(
'--interval',
dest='interval',
@ -165,6 +168,13 @@ def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements
action='store_true',
help="Only show GPU processes with the compute context. (type: 'C' or 'C+G')",
)
process_filtering.add_argument(
'--no-compute',
'-C',
dest='no_compute',
action='store_true',
help="Exclude GPU processes with the compute context. (type: 'G' only)",
)
process_filtering.add_argument(
'--graphics',
'-g',
@ -172,6 +182,13 @@ def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements
action='store_true',
help="Only show GPU processes with the graphics context. (type: 'G' or 'C+G')",
)
process_filtering.add_argument(
'--no-graphics',
'-G',
dest='no_graphics',
action='store_true',
help="Exclude GPU processes with the graphics context. (type: 'C' only)",
)
process_filtering.add_argument(
'--user',
'-u',
@ -297,8 +314,12 @@ def main(): # pylint: disable=too-many-branches,too-many-statements,too-many-lo
filters = []
if args.compute:
filters.append(lambda process: 'C' in process.type)
if args.no_compute:
filters.append(lambda process: 'C' not in process.type)
if args.graphics:
filters.append(lambda process: 'G' in process.type)
if args.no_graphics:
filters.append(lambda process: 'G' not in process.type)
if args.user is not None:
users = set(args.user)
filters.append(lambda process: process.username in users)