Files
bentoml--openllm/src/openllm/clean.py
T
2026-07-13 12:48:25 +08:00

84 lines
2.5 KiB
Python

from __future__ import annotations
import os, pathlib, shutil
import questionary
from openllm.analytic import OpenLLMTyper
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) -> None:
if verbose:
VERBOSE_LEVEL.set(20)
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()
if not sure:
return
shutil.rmtree(HUGGINGFACE_CACHE, ignore_errors=True)
output('All models cached by Huggingface have been removed', style='green')
@app.command(help='Clean up all the virtual environments created by OpenLLM')
def venvs(verbose: bool = False) -> None:
if verbose:
VERBOSE_LEVEL.set(20)
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()
if not sure:
return
shutil.rmtree(VENV_DIR, ignore_errors=True)
output('All virtual environments have been removed', style='green')
@app.command(help='Clean up all the repositories cloned by OpenLLM')
def repos(verbose: bool = False) -> None:
if verbose:
VERBOSE_LEVEL.set(20)
shutil.rmtree(REPO_DIR, ignore_errors=True)
output('All repositories have been removed', style='green')
@app.command(help='Reset configurations to default')
def configs(verbose: bool = False) -> None:
if verbose:
VERBOSE_LEVEL.set(20)
shutil.rmtree(CONFIG_FILE, ignore_errors=True)
output('All configurations have been reset', style='green')
@app.command(name='all', help='Clean up all above and bring OpenLLM to a fresh start')
def all_cache(verbose: bool = False) -> None:
if verbose:
VERBOSE_LEVEL.set(20)
repos()
venvs()
model_cache()
configs()