diff --git a/src/openllm/clean.py b/src/openllm/clean.py index 0f26a1beb..19c302edf 100644 --- a/src/openllm/clean.py +++ b/src/openllm/clean.py @@ -5,18 +5,36 @@ import questionary from openllm.analytic import OpenLLMTyper -from openllm.common import CONFIG_FILE, REPO_DIR, VENV_DIR, VERBOSE_LEVEL, output +from openllm.common import (CONFIG_FILE, REPO_DIR, VENV_DIR, VERBOSE_LEVEL, + output) app = OpenLLMTyper(help='clean up and release disk space used by OpenLLM') HUGGINGFACE_CACHE = pathlib.Path.home() / '.cache' / 'huggingface' / 'hub' +def _du(path: pathlib.Path) -> int: + seen_paths = set() + used_space = 0 + + for f in path.rglob('*'): + if os.name == 'nt': # Windows system + # On Windows, directly add file sizes without considering hard links + used_space += f.stat().st_size + else: + # On non-Windows systems, use inodes to avoid double counting + stat = f.stat() + if stat.st_ino not in seen_paths: + seen_paths.add(stat.st_ino) + used_space += stat.st_size + return used_space + + @app.command(help='Clean up all the cached models from huggingface') def model_cache(verbose: bool = False): if verbose: VERBOSE_LEVEL.set(20) - used_space = sum(f.stat().st_size for f in HUGGINGFACE_CACHE.rglob('*')) + used_space = _du(HUGGINGFACE_CACHE) sure = questionary.confirm( f'This will remove all models cached by Huggingface (~{used_space / 1024 / 1024:.2f}MB), are you sure?' ).ask() @@ -31,21 +49,7 @@ def venvs(verbose: bool = False): if verbose: VERBOSE_LEVEL.set(20) - # Set to store paths of files to avoid double counting - seen_paths = set() - used_space = 0 - - for f in VENV_DIR.rglob('*'): - if os.name == 'nt': # Windows system - # On Windows, directly add file sizes without considering hard links - used_space += f.stat().st_size - else: - # On non-Windows systems, use inodes to avoid double counting - stat = f.stat() - if stat.st_ino not in seen_paths: - seen_paths.add(stat.st_ino) - used_space += stat.st_size - + used_space = _du(VENV_DIR) sure = questionary.confirm( f'This will remove all virtual environments created by OpenLLM (~{used_space / 1024 / 1024:.2f}MB), are you sure?' ).ask()