cvgpu_console/python/gpu_status.py

import argparse
from typing import Protocol

import requests
from renderer import Renderer

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

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(
    "-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."
)

args = parser.parse_args()

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

if args.ultracompact:
    renderer = Renderer(
        columns=3,
        progress_bar_width=30,
        use_space_lines=False,
        node_names=args.filter,
        display_power=False,
        display_users=args.users,
    )
elif args.compact:
    renderer = Renderer(
        columns=2,
        progress_bar_width=40,
        use_space_lines=False,
        node_names=args.filter,
        display_power=True,
        display_users=args.users,
    )
else:
    renderer = Renderer(
        columns=1,
        progress_bar_width=50,
        use_space_lines=True,
        node_names=args.filter,
        display_users=args.users,
    )

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