Support hardware metadata queries (#152)

* feat: Include agent hardware metadata query in "admin agent" command
* feat: Translate empty/null values using cli.utils.format_nested_dict() and cli.utils.format_value()
* feat: Add "admin storage", "admin storage-list" commands
* feat: Include performance metric in "admin storage" result

Author: achimnol

changes/152.feature

@@ -0,0 +1,3 @@
+Add support for hardware metadata queries for super-admins
+- Add `admin storages` and `admin storage {vfolder-host}` commands which provides the hardware metadata and performance metric
+- Include `hardware_metadata` field in the `Agent` graph object type.

src/ai/backend/client/cli/admin/__init__.py

@@ -10,9 +10,20 @@ def admin():
 
 def _attach_command():
     from . import (  # noqa
-        agents, domains, etcd, groups, images, keypairs, resources, resource_policies,
-        scaling_groups, sessions, users, vfolders,
+        agents,
+        domains,
+        etcd,
+        groups,
+        images,
+        keypairs,
         license,
+        resources,
+        resource_policies,
+        scaling_groups,
+        sessions,
+        storage,
+        users,
+        vfolders,
     )
 
 

src/ai/backend/client/cli/admin/agents.py

@@ -13,6 +13,7 @@
     echo_via_pager,
     tabulate_items,
 )
+from ..utils import format_nested_dicts
 from ...exceptions import NoItems
 
 
@@ -85,6 +86,7 @@ def agent(agent_id):
         ('CPU Usage (%)', 'cpu_cur_pct'),
         ('Total slots', 'available_slots'),
         ('Occupied slots', 'occupied_slots'),
+        ('Hardware Metadata', 'hardware_metadata'),
         ('Live Stat', 'live_stat'),
     ]
     if is_legacy_server():
@@ -104,6 +106,8 @@ def agent(agent_id):
             if key in resp:
                 if key == 'live_stat' and resp[key] is not None:
                     rows.append((name, format_stats(json.loads(resp[key]))))
+                elif key == 'hardware_metadata':
+                    rows.append((name, format_nested_dicts(json.loads(resp[key]))))
                 else:
                     rows.append((name, resp[key]))
         print(tabulate(rows, headers=('Field', 'Value')))

src/ai/backend/client/cli/admin/storage.py

@@ -0,0 +1,79 @@
+import json
+import sys
+
+import click
+from tabulate import tabulate
+
+from . import admin
+from ...session import Session
+from ..pretty import print_error
+from ..pagination import (
+    get_preferred_page_size,
+    echo_via_pager,
+    tabulate_items,
+)
+from ..utils import format_nested_dicts, format_value
+from ...exceptions import NoItems
+
+
+@admin.command()
+@click.argument('vfolder_host')
+def storage(vfolder_host):
+    """
+    Show the information about the given storage volume.
+    (super-admin privilege required)
+    """
+    fields = [
+        ('ID', 'id'),
+        ('Backend', 'backend'),
+        ('Capabilities', 'capabilities'),
+        ('Path', 'path'),
+        ('FS prefix', 'fsprefix'),
+        ('Hardware Metadata', 'hardware_metadata'),
+        ('Perf. Metric', 'performance_metric'),
+    ]
+    with Session() as session:
+        try:
+            resp = session.Storage.detail(
+                vfolder_host=vfolder_host,
+                fields=(item[1] for item in fields))
+        except Exception as e:
+            print_error(e)
+            sys.exit(1)
+        rows = []
+        for name, key in fields:
+            if key in resp:
+                if key in ('hardware_metadata', 'performance_metric'):
+                    rows.append((name, format_nested_dicts(json.loads(resp[key]))))
+                else:
+                    rows.append((name, format_value(resp[key])))
+        print(tabulate(rows, headers=('Field', 'Value')))
+
+
+@admin.command()
+def storage_list():
+    """
+    List storage volumes.
+    (super-admin privilege required)
+    """
+    fields = [
+        ('ID', 'id'),
+        ('Backend', 'backend'),
+        ('Capabilities', 'capabilities'),
+    ]
+    try:
+        with Session() as session:
+            page_size = get_preferred_page_size()
+            try:
+                items = session.Storage.paginated_list(
+                    fields=[f[1] for f in fields],
+                    page_size=page_size,
+                )
+                echo_via_pager(
+                    tabulate_items(items, fields)
+                )
+            except NoItems:
+                print("There are no storage volumes.")
+    except Exception as e:
+        print_error(e)
+        sys.exit(1)

src/ai/backend/client/cli/utils.py

@@ -1,4 +1,6 @@
 import re
+import textwrap
+from typing import Any, Mapping
 
 import click
 
@@ -37,3 +39,34 @@ def convert(self, value, param, ctx):
         if m is None:
             self.fail(f"{value!r} is not a valid byte-size expression", param, ctx)
         return value
+
+
+def format_nested_dicts(value: Mapping[str, Mapping[str, Any]]) -> str:
+    """
+    Format a mapping from string keys to sub-mappings.
+    """
+    rows = []
+    if not value:
+        rows.append("(empty)")
+    for outer_key, outer_value in value.items():
+        if isinstance(outer_value, dict):
+            if outer_value:
+                rows.append(f"+ {outer_key}")
+                inner_rows = format_nested_dicts(outer_value)
+                rows.append(textwrap.indent(inner_rows, prefix="  "))
+            else:
+                rows.append(f"+ {outer_key}: (empty)")
+        else:
+            if outer_value is None:
+                rows.append(f"- {outer_key}: (null)")
+            else:
+                rows.append(f"- {outer_key}: {outer_value}")
+    return "\n".join(rows)
+
+
+def format_value(value: Any) -> str:
+    if value is None:
+        return "(null)"
+    if isinstance(value, (dict, list, set)) and not value:
+        return "(empty)"
+    return str(value)

src/ai/backend/client/func/storage.py

@@ -0,0 +1,85 @@
+import textwrap
+from typing import (
+    AsyncIterator,
+    Sequence,
+)
+
+from .base import api_function, BaseFunction
+from ..request import Request
+from ..pagination import generate_paginated_results
+
+__all__ = (
+    'Storage',
+)
+
+_default_list_fields = (
+    'id',
+    'backend',
+    'capabilities',
+)
+
+_default_detail_fields = (
+    'id',
+    'backend',
+    'path',
+    'fsprefix',
+    'capabilities',
+    'hardware_metadata',
+)
+
+
+class Storage(BaseFunction):
+    """
+    Provides a shortcut of :func:`Admin.query()
+    <ai.backend.client.admin.Admin.query>` that fetches various straoge volume
+    information keyed by vfolder hosts.
+
+    .. note::
+
+      All methods in this function class require your API access key to
+      have the *super-admin* privilege.
+    """
+
+    @api_function
+    @classmethod
+    async def paginated_list(
+        cls,
+        status: str = 'ALIVE',
+        *,
+        fields: Sequence[str] = _default_list_fields,
+        page_size: int = 20,
+    ) -> AsyncIterator[dict]:
+        """
+        Lists the keypairs.
+        You need an admin privilege for this operation.
+        """
+        async for item in generate_paginated_results(
+            'storage_volume_list',
+            {},
+            fields,
+            page_size=page_size,
+        ):
+            yield item
+
+    @api_function
+    @classmethod
+    async def detail(
+        cls,
+        vfolder_host: str,
+        fields: Sequence[str] = _default_detail_fields,
+    ) -> Sequence[dict]:
+        query = textwrap.dedent("""\
+            query($vfolder_host: String!) {
+                storage_volume(id: $vfolder_host) {$fields}
+            }
+        """)
+        query = query.replace('$fields', ' '.join(fields))
+        variables = {'vfolder_host': vfolder_host}
+        rqst = Request('POST', '/admin/graphql')
+        rqst.set_json({
+            'query': query,
+            'variables': variables,
+        })
+        async with rqst.fetch() as resp:
+            data = await resp.json()
+            return data['storage_volume']

src/ai/backend/client/session.py

@@ -243,7 +243,7 @@ class BaseSession(metaclass=abc.ABCMeta):
         '_config', '_closed', '_context_token', '_proxy_mode',
         'aiohttp_session', 'api_version',
         'System', 'Manager', 'Admin',
-        'Agent', 'AgentWatcher', 'ScalingGroup',
+        'Agent', 'AgentWatcher', 'ScalingGroup', 'Storage',
         'Image', 'ComputeSession', 'SessionTemplate',
         'Domain', 'Group', 'Auth', 'User', 'KeyPair',
         'BackgroundTask',
@@ -273,6 +273,7 @@ def __init__(
         from .func.system import System
         from .func.admin import Admin
         from .func.agent import Agent, AgentWatcher
+        from .func.storage import Storage
         from .func.auth import Auth
         from .func.bgtask import BackgroundTask
         from .func.domain import Domain
@@ -295,6 +296,7 @@ def __init__(
         self.Admin = Admin
         self.Agent = Agent
         self.AgentWatcher = AgentWatcher
+        self.Storage = Storage
         self.Auth = Auth
         self.BackgroundTask = BackgroundTask
         self.EtcdConfig = EtcdConfig