chore: du

bentoml · Aug 12, 2024 · cacc7d5 · cacc7d5
1 parent 82dd97f
commit cacc7d5
Showing 1 changed file with 21 additions and 17 deletions.
diff --git 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()