diff --git a/llama_stack/cli/info.py b/llama_stack/cli/info.py new file mode 100644 index 0000000000..2ba545e8f9 --- /dev/null +++ b/llama_stack/cli/info.py @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import argparse +import importlib +import platform +import subprocess +import sys + +import psutil + +from llama_stack.cli.subcommand import Subcommand + + +def _system_info(): + sys_info = { + "sys.version": sys.version, + "sys.platform": sys.platform, + "platform.machine": platform.machine(), + "platform.node": platform.node(), + "platform.python_version": platform.python_version(), + } + if sys.platform == "linux": + os_release = platform.freedesktop_os_release() + for key in ["ID", "VERSION_ID", "PRETTY_NAME"]: + value = os_release.get(key) + if value: + sys_info[f"os-release.{key}"] = value + elif sys.platform == "darwin": + try: + cpu_info = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"]).decode("utf-8").strip() + sys_info["platform.cpu_brand"] = cpu_info + except subprocess.CalledProcessError: + sys_info["platform.cpu_brand"] = "Unknown" + sys_memory_info = psutil.virtual_memory() + sys_info["memory.used"] = f"{(sys_memory_info.used / 1024**3):.2f} GB" + sys_info["memory.available"] = f"{(sys_memory_info.available / 1024**3):.2f} GB" + sys_info["memory.total"] = f"{(sys_memory_info.total / 1024**3):.2f} GB" + return sys_info + + +def _list_pkg_version(pkg_name): + packages = sorted( + (distribution.name, distribution.version) + for distribution in importlib.metadata.distributions() + if distribution.name.startswith(pkg_name) + ) + return {f"{name}.version": version for name, version in packages} + + +def _cuda_info(): + import torch + + if torch.cuda.is_available(): + cuda_info = { + "cuda.version": torch.version.cuda, + "gpu.count": torch.cuda.device_count(), + "cuda.bf16": torch.cuda.is_bf16_supported(), + "cuda.current.device": torch.cuda.current_device(), + } + for idx in range(torch.cuda.device_count()): + device = torch.device("cuda", idx) + free, total = torch.cuda.mem_get_info(device) + cuda_info[f"cuda.{idx}.name"] = torch.cuda.get_device_name(device) + cuda_info[f"cuda.{idx}.free"] = f"{(free / 1024**3):.1f} GB" + cuda_info[f"cuda.{idx}.total"] = f"{(total / 1024**3):.1f} GB" + return cuda_info + + +def _add_to_group(groups, item_name, info): + if item_name not in groups: + groups[item_name] = [] + groups[item_name].extend(info.items()) + + +def _get_env_info_by_group(): + groups = {} + _add_to_group(groups, "system", _system_info()) + _add_to_group(groups, "llama packages", _list_pkg_version("llama")) + _add_to_group(groups, "torch packages", _list_pkg_version("torch")) + + # only cuda exist + cuda_info = _cuda_info() + if cuda_info: + _add_to_group(groups, "cuda info", _cuda_info()) + + return groups + + +class LlamaStackInfo(Subcommand): + """Llama cli for showing llama stack environment information""" + + def __init__(self, subparsers: argparse._SubParsersAction): + super().__init__() + self.parser = subparsers.add_parser( + "info", + prog="llama info", + description="Show llama stack environment information", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + self._add_arguments() + self.parser.set_defaults(func=self._run_llama_stack_info_cmd) + + def _run_llama_stack_info_cmd(self, args: argparse.Namespace) -> None: + groups = _get_env_info_by_group() + for idx, (group, items) in enumerate(groups.items()): + if idx > 0: + print() + print(f"{group}:") + for k, v in items: + print(f" {k}: {v}") diff --git a/llama_stack/cli/llama.py b/llama_stack/cli/llama.py index 8ff5800298..0890b6ac78 100644 --- a/llama_stack/cli/llama.py +++ b/llama_stack/cli/llama.py @@ -7,6 +7,7 @@ import argparse from .download import Download +from .info import LlamaStackInfo from .model import ModelParser from .stack import StackParser from .stack.utils import print_subcommand_description @@ -34,6 +35,7 @@ def __init__(self): StackParser.create(subparsers) Download.create(subparsers) VerifyDownload.create(subparsers) + LlamaStackInfo.create(subparsers) print_subcommand_description(self.parser, subparsers)