cvgpu_console/python/gpu_status.py

import argparse
import requests

from renderer import Renderer

parser = argparse.ArgumentParser(description="Prints the status of server nodes and their GPUs.")


parser.add_argument(
    "watch", 
    type=str,
    nargs="?",
    choices=("watch",),
    help="If the word 'watch' is supplied as first argument, the script will be run in watch mode. Supplying this argument in any other position will have no effect."
)
parser.add_argument(
    "-s", 
    "--server", 
    type=str, 
    default="deimos", 
    help="The name of the server to query."
)
parser.add_argument(
    "-p",
    "--port",
    type=str,
    default="8091",
    help="The port of the server to query.",
)
parser.add_argument(
    "-t", 
    "--protocol", 
    type=str, 
    default="http", 
    choices=("http", "https"), 
    help="The protocol to use for the query.",
)
parser.add_argument(
    "-c", 
    "--compact", 
    action="store_true", 
    help="Flag; if set, a compact overview will be displayed."
)
parser.add_argument(
    "--columns", 
    type=int,
    default=-1,
    help="The number of columns the node information is formatted in. By default, this is set by the compactness style (default, --compact, --ultracompact). Changing this value might break formatting."
)
parser.add_argument(
    "--pbar", 
    type=int,
    default=-1,
    help="The length of the progress bar. By default, this is set by the compactness style (default, --compact, --ultracompact). Changing this value might break formatting."
)
parser.add_argument(
    "-uc", 
    "--ultracompact", 
    action="store_true", 
    help="Flag; if set, a very compact overview will be displayed."
)
parser.add_argument(
    "-u", 
    "--users", 
    action="store_true", 
    help="Flag; if set, the current users of each GPU will be displayed."
)
parser.add_argument(
    "-f",
    "--filter",
    type=str, 
    nargs="*",
    default=None,
    help="The node names of nodes to filter for."
)
parser.add_argument(
    "-a",
    "--available",
    action="store_true",
    help="Flag; if set, displays only available GPUs."
)

args = parser.parse_args()

response = requests.get(
    f"{args.protocol}://{args.server}:{args.port}/api/clients/"
)

if args.ultracompact:
    renderer = Renderer(
        columns=3 if args.columns == -1 else args.columns,
        progress_bar_width=20 if args.pbar == -1 else args.pbar,
        gpu_inner_spacings=False,
        gpu_outer_spacings=False,
        node_names=args.filter,
        display_power=False,
        display_users=args.users,
        dict_type="shortened",
        available_gpus_only=args.available,
    )
elif args.compact:
    renderer = Renderer(
        columns=2 if args.columns == -1 else args.columns,
        progress_bar_width=40 if args.pbar == -1 else args.pbar,
        gpu_inner_spacings=True,
        gpu_outer_spacings=False,
        node_names=args.filter,
        display_power=True,
        display_users=args.users,
        available_gpus_only=args.available,
    )
else:
    renderer = Renderer(
        columns=1 if args.columns == -1 else args.columns,
        progress_bar_width=50 if args.pbar == -1 else args.pbar,
        gpu_inner_spacings=True,
        gpu_outer_spacings=True,
        node_names=args.filter,
        display_users=args.users,
        available_gpus_only=args.available,
    )

print(renderer.render_info_dict(response.json()))