diff --git a/llama_stack/cli/model/list.py b/llama_stack/cli/model/list.py index a89b9c143..2f62cb9ce 100644 --- a/llama_stack/cli/model/list.py +++ b/llama_stack/cli/model/list.py @@ -7,6 +7,7 @@ import argparse import os import time +from pathlib import Path from llama_stack.cli.subcommand import Subcommand from llama_stack.cli.table import print_table @@ -14,17 +15,8 @@ from llama_stack.distribution.utils.config_dirs import DEFAULT_CHECKPOINT_DIR from llama_stack.models.llama.sku_list import all_registered_models -def _convert_size_to_gb(size): - return size / (1024**3) - - def _get_model_size(model_dir): - model_size = 0 - for dirpath, _, files in os.walk(model_dir): - for file in files: - file_path = os.path.join(dirpath, file) - model_size += os.path.getsize(file_path) - return model_size + return sum(f.stat().st_size for f in Path(model_dir).rglob("*") if f.is_file()) def _run_model_list_downloaded_cmd() -> None: @@ -34,7 +26,7 @@ def _run_model_list_downloaded_cmd() -> None: for model in os.listdir(DEFAULT_CHECKPOINT_DIR): abs_path = os.path.join(DEFAULT_CHECKPOINT_DIR, model) space_usage = _get_model_size(abs_path) - model_size = f"{_convert_size_to_gb(space_usage):.2f} GB" + model_size = f"{space_usage / (1024**3):.2f} GB" modified_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(os.path.getmtime(abs_path))) rows.append( [