diff --git a/changes/165.feature b/changes/165.feature new file mode 100644 index 00000000..4131b3f7 --- /dev/null +++ b/changes/165.feature @@ -0,0 +1 @@ +Add `resource check` command to get available resources for given groups and scaling groups diff --git a/src/ai/backend/client/cli/__init__.py b/src/ai/backend/client/cli/__init__.py index a0bdcee0..7a0d699c 100644 --- a/src/ai/backend/client/cli/__init__.py +++ b/src/ai/backend/client/cli/__init__.py @@ -39,6 +39,7 @@ def _attach_command(): from . import session_template # noqa from . import dotfile # noqa from . import server_log # noqa + from . import resource # noqa _attach_command() diff --git a/src/ai/backend/client/cli/manager.py b/src/ai/backend/client/cli/manager.py index 8b81c4b5..c1cf1228 100644 --- a/src/ai/backend/client/cli/manager.py +++ b/src/ai/backend/client/cli/manager.py @@ -9,8 +9,9 @@ from . import main from .interaction import ask_yn -from .pretty import print_done, print_error, print_fail, print_info, print_wait +from .pretty import print_done, print_error, print_fail, print_info, print_wait, print_warn from ..session import Session +from .config import get_config @main.group() @@ -203,4 +204,4 @@ def exclude_agents(agent_ids): print_done('The given agents will no longer start new sessions.') except Exception as e: print_error(e) - sys.exit(1) + sys.exit(1) \ No newline at end of file diff --git a/src/ai/backend/client/cli/pretty.py b/src/ai/backend/client/cli/pretty.py index 3074ceca..76a7bc08 100644 --- a/src/ai/backend/client/cli/pretty.py +++ b/src/ai/backend/client/cli/pretty.py @@ -1,8 +1,11 @@ +from decimal import Decimal import enum import functools +import math import sys import textwrap import traceback +from typing import TextIO from click import echo, style @@ -175,3 +178,23 @@ def show_warning(message, category, filename, lineno, file=None, line=None): style(str(category.__name__), fg='yellow', bold=True), style(str(message), fg='yellow'), ), file=file) + + +def print_resource( + slots: dict, + slot_types: dict, + *, + prefix: str = "", + file: TextIO = sys.stdout, + nan_as_infinite: bool = False +) -> None: + for key in slot_types.keys(): + value = Decimal(slots[key]) + value_expr: str + if math.isnan(value): + value_expr = "Unlimited" if nan_as_infinite else "Unavailable" + elif math.isinf(value): + value_expr = "Unlimited" + else: + value_expr = str(value) + print(f"{prefix}{key}: {value_expr}", file=file) \ No newline at end of file diff --git a/src/ai/backend/client/cli/resource.py b/src/ai/backend/client/cli/resource.py new file mode 100644 index 00000000..f59b7134 --- /dev/null +++ b/src/ai/backend/client/cli/resource.py @@ -0,0 +1,62 @@ +import sys + +import click + +from . import main +from .pretty import print_error, print_resource +from ..session import Session + + +@main.group() +def resource(): + """ + Provides resource check operations + """ + + +@resource.command() +@click.argument('group', metavar='NAME', default="default") +def check_group(group: str) -> None: + """ + Display the available resources from all allowed scaling groups of the given user group (project). + """ + try: + with Session() as session: + ret = session.Resource.check_group(group) + slot_types = session.Resource.get_resource_slots() + print("Limits:") + prefix = "- " + print_resource(ret["limits"], slot_types, prefix=prefix, nan_as_infinite=True) + print("Occupied:") + print_resource(ret["occupied"], slot_types, prefix=prefix) + print("Remaining:") + print_resource(ret["remaining"], slot_types, prefix=prefix) + print("Scaling groups allowed for this group:") + for sgroup_name in ret["scaling_groups"].keys(): + print(f" [{sgroup_name}]") + print(" Occupied:") + prefix = " - " + print_resource(ret["scaling_groups"][sgroup_name]["occupied"], slot_types, prefix=prefix) + print(" Remaining:") + print_resource(ret["scaling_groups"][sgroup_name]["remaining"], slot_types, prefix=prefix) + except Exception as e: + print_error(e) + sys.exit(1) + + +@resource.command() +@click.argument('scaling_group', metavar='NAME', default="default") +def check_scaling_group(scaling_group: str) -> None: + """ + Display the available resource from the given scaling group. + """ + try: + with Session() as session: + ret = session.Resource.check_scaling_group(scaling_group) + slot_types = session.Resource.get_resource_slots() + print(f"Total remaining resources of scaling group [{scaling_group}]:") + prefix = "- " + print_resource(ret["remaining"], slot_types, prefix=prefix) + except Exception as e: + print_error(e) + sys.exit(1) \ No newline at end of file diff --git a/src/ai/backend/client/func/resource.py b/src/ai/backend/client/func/resource.py index 1b7c7a8f..e1c44408 100644 --- a/src/ai/backend/client/func/resource.py +++ b/src/ai/backend/client/func/resource.py @@ -34,6 +34,26 @@ async def check_presets(cls): async with rqst.fetch() as resp: return await resp.json() + @api_function + @classmethod + async def check_group(cls, group_name: str): + """ + Lists available resources from the scaling groups. + """ + rqst = Request('GET', '/resource/group', params={"name": group_name}) + async with rqst.fetch() as resp: + return await resp.json() + + @api_function + @classmethod + async def check_scaling_group(cls, sgroup_name: str): + """ + Lists available resources from the scaling groups. + """ + rqst = Request('GET', '/resource/scaling-group', params={"name": sgroup_name}) + async with rqst.fetch() as resp: + return await resp.json() + @api_function @classmethod async def get_docker_registries(cls):