feat: Add command to get available resources

changes/165.feature

@@ -0,0 +1 @@
+Add `resource check` command to get available resources for given groups and scaling groups

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

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)

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)

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)

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):