cache.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403
  1. # coding=utf-8
  2. # Copyright 2025-present, the HuggingFace Inc. team.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. """Contains the 'hf cache' command group with 'scan' and 'delete' subcommands."""
  16. import os
  17. import time
  18. from argparse import Namespace, _SubParsersAction
  19. from functools import wraps
  20. from tempfile import mkstemp
  21. from typing import Any, Callable, Iterable, List, Literal, Optional, Union
  22. from ..utils import CachedRepoInfo, CachedRevisionInfo, CacheNotFound, HFCacheInfo, scan_cache_dir
  23. from . import BaseHuggingfaceCLICommand
  24. from ._cli_utils import ANSI, tabulate
  25. # --- DELETE helpers (from delete_cache.py) ---
  26. try:
  27. from InquirerPy import inquirer
  28. from InquirerPy.base.control import Choice
  29. from InquirerPy.separator import Separator
  30. _inquirer_py_available = True
  31. except ImportError:
  32. _inquirer_py_available = False
  33. SortingOption_T = Literal["alphabetical", "lastUpdated", "lastUsed", "size"]
  34. _CANCEL_DELETION_STR = "CANCEL_DELETION"
  35. def require_inquirer_py(fn: Callable) -> Callable:
  36. @wraps(fn)
  37. def _inner(*args, **kwargs):
  38. if not _inquirer_py_available:
  39. raise ImportError(
  40. "The 'cache delete' command requires extra dependencies for the TUI.\n"
  41. "Please run 'pip install \"huggingface_hub[cli]\"' to install them.\n"
  42. "Otherwise, disable TUI using the '--disable-tui' flag."
  43. )
  44. return fn(*args, **kwargs)
  45. return _inner
  46. class CacheCommand(BaseHuggingfaceCLICommand):
  47. @staticmethod
  48. def register_subcommand(parser: _SubParsersAction):
  49. cache_parser = parser.add_parser("cache", help="Manage local cache directory.")
  50. cache_subparsers = cache_parser.add_subparsers(dest="cache_command", help="Cache subcommands")
  51. # Show help if no subcommand is provided
  52. cache_parser.set_defaults(func=lambda args: cache_parser.print_help())
  53. # Scan subcommand
  54. scan_parser = cache_subparsers.add_parser("scan", help="Scan cache directory.")
  55. scan_parser.add_argument(
  56. "--dir",
  57. type=str,
  58. default=None,
  59. help="cache directory to scan (optional). Default to the default HuggingFace cache.",
  60. )
  61. scan_parser.add_argument(
  62. "-v",
  63. "--verbose",
  64. action="count",
  65. default=0,
  66. help="show a more verbose output",
  67. )
  68. scan_parser.set_defaults(func=CacheCommand, cache_command="scan")
  69. # Delete subcommand
  70. delete_parser = cache_subparsers.add_parser("delete", help="Delete revisions from the cache directory.")
  71. delete_parser.add_argument(
  72. "--dir",
  73. type=str,
  74. default=None,
  75. help="cache directory (optional). Default to the default HuggingFace cache.",
  76. )
  77. delete_parser.add_argument(
  78. "--disable-tui",
  79. action="store_true",
  80. help=(
  81. "Disable Terminal User Interface (TUI) mode. Useful if your platform/terminal doesn't support the multiselect menu."
  82. ),
  83. )
  84. delete_parser.add_argument(
  85. "--sort",
  86. nargs="?",
  87. choices=["alphabetical", "lastUpdated", "lastUsed", "size"],
  88. help=(
  89. "Sort repositories by the specified criteria. Options: "
  90. "'alphabetical' (A-Z), "
  91. "'lastUpdated' (newest first), "
  92. "'lastUsed' (most recent first), "
  93. "'size' (largest first)."
  94. ),
  95. )
  96. delete_parser.set_defaults(func=CacheCommand, cache_command="delete")
  97. def __init__(self, args: Namespace) -> None:
  98. self.args = args
  99. self.verbosity: int = getattr(args, "verbose", 0)
  100. self.cache_dir: Optional[str] = getattr(args, "dir", None)
  101. self.disable_tui: bool = getattr(args, "disable_tui", False)
  102. self.sort_by: Optional[SortingOption_T] = getattr(args, "sort", None)
  103. self.cache_command: Optional[str] = getattr(args, "cache_command", None)
  104. def run(self):
  105. if self.cache_command == "scan":
  106. self._run_scan()
  107. elif self.cache_command == "delete":
  108. self._run_delete()
  109. else:
  110. print("Please specify a cache subcommand (scan or delete). Use -h for help.")
  111. def _run_scan(self):
  112. try:
  113. t0 = time.time()
  114. hf_cache_info = scan_cache_dir(self.cache_dir)
  115. t1 = time.time()
  116. except CacheNotFound as exc:
  117. cache_dir = exc.cache_dir
  118. print(f"Cache directory not found: {cache_dir}")
  119. return
  120. print(get_table(hf_cache_info, verbosity=self.verbosity))
  121. print(
  122. f"\nDone in {round(t1 - t0, 1)}s. Scanned {len(hf_cache_info.repos)} repo(s)"
  123. f" for a total of {ANSI.red(hf_cache_info.size_on_disk_str)}."
  124. )
  125. if len(hf_cache_info.warnings) > 0:
  126. message = f"Got {len(hf_cache_info.warnings)} warning(s) while scanning."
  127. if self.verbosity >= 3:
  128. print(ANSI.gray(message))
  129. for warning in hf_cache_info.warnings:
  130. print(ANSI.gray(str(warning)))
  131. else:
  132. print(ANSI.gray(message + " Use -vvv to print details."))
  133. def _run_delete(self):
  134. hf_cache_info = scan_cache_dir(self.cache_dir)
  135. if self.disable_tui:
  136. selected_hashes = _manual_review_no_tui(hf_cache_info, preselected=[], sort_by=self.sort_by)
  137. else:
  138. selected_hashes = _manual_review_tui(hf_cache_info, preselected=[], sort_by=self.sort_by)
  139. if len(selected_hashes) > 0 and _CANCEL_DELETION_STR not in selected_hashes:
  140. confirm_message = _get_expectations_str(hf_cache_info, selected_hashes) + " Confirm deletion ?"
  141. if self.disable_tui:
  142. confirmed = _ask_for_confirmation_no_tui(confirm_message)
  143. else:
  144. confirmed = _ask_for_confirmation_tui(confirm_message)
  145. if confirmed:
  146. strategy = hf_cache_info.delete_revisions(*selected_hashes)
  147. print("Start deletion.")
  148. strategy.execute()
  149. print(
  150. f"Done. Deleted {len(strategy.repos)} repo(s) and"
  151. f" {len(strategy.snapshots)} revision(s) for a total of"
  152. f" {strategy.expected_freed_size_str}."
  153. )
  154. return
  155. print("Deletion is cancelled. Do nothing.")
  156. def get_table(hf_cache_info: HFCacheInfo, *, verbosity: int = 0) -> str:
  157. if verbosity == 0:
  158. return tabulate(
  159. rows=[
  160. [
  161. repo.repo_id,
  162. repo.repo_type,
  163. "{:>12}".format(repo.size_on_disk_str),
  164. repo.nb_files,
  165. repo.last_accessed_str,
  166. repo.last_modified_str,
  167. ", ".join(sorted(repo.refs)),
  168. str(repo.repo_path),
  169. ]
  170. for repo in sorted(hf_cache_info.repos, key=lambda repo: repo.repo_path)
  171. ],
  172. headers=[
  173. "REPO ID",
  174. "REPO TYPE",
  175. "SIZE ON DISK",
  176. "NB FILES",
  177. "LAST_ACCESSED",
  178. "LAST_MODIFIED",
  179. "REFS",
  180. "LOCAL PATH",
  181. ],
  182. )
  183. else:
  184. return tabulate(
  185. rows=[
  186. [
  187. repo.repo_id,
  188. repo.repo_type,
  189. revision.commit_hash,
  190. "{:>12}".format(revision.size_on_disk_str),
  191. revision.nb_files,
  192. revision.last_modified_str,
  193. ", ".join(sorted(revision.refs)),
  194. str(revision.snapshot_path),
  195. ]
  196. for repo in sorted(hf_cache_info.repos, key=lambda repo: repo.repo_path)
  197. for revision in sorted(repo.revisions, key=lambda revision: revision.commit_hash)
  198. ],
  199. headers=[
  200. "REPO ID",
  201. "REPO TYPE",
  202. "REVISION",
  203. "SIZE ON DISK",
  204. "NB FILES",
  205. "LAST_MODIFIED",
  206. "REFS",
  207. "LOCAL PATH",
  208. ],
  209. )
  210. def _get_repo_sorting_key(repo: CachedRepoInfo, sort_by: Optional[SortingOption_T] = None):
  211. if sort_by == "alphabetical":
  212. return (repo.repo_type, repo.repo_id.lower())
  213. elif sort_by == "lastUpdated":
  214. return -max(rev.last_modified for rev in repo.revisions)
  215. elif sort_by == "lastUsed":
  216. return -repo.last_accessed
  217. elif sort_by == "size":
  218. return -repo.size_on_disk
  219. else:
  220. return (repo.repo_type, repo.repo_id)
  221. @require_inquirer_py
  222. def _manual_review_tui(
  223. hf_cache_info: HFCacheInfo, preselected: List[str], sort_by: Optional[SortingOption_T] = None
  224. ) -> List[str]:
  225. choices = _get_tui_choices_from_scan(repos=hf_cache_info.repos, preselected=preselected, sort_by=sort_by)
  226. checkbox = inquirer.checkbox(
  227. message="Select revisions to delete:",
  228. choices=choices,
  229. cycle=False,
  230. height=100,
  231. instruction=_get_expectations_str(
  232. hf_cache_info, selected_hashes=[c.value for c in choices if isinstance(c, Choice) and c.enabled]
  233. ),
  234. long_instruction="Press <space> to select, <enter> to validate and <ctrl+c> to quit without modification.",
  235. transformer=lambda result: f"{len(result)} revision(s) selected.",
  236. )
  237. def _update_expectations(_):
  238. checkbox._instruction = _get_expectations_str(
  239. hf_cache_info,
  240. selected_hashes=[choice["value"] for choice in checkbox.content_control.choices if choice["enabled"]],
  241. )
  242. checkbox.kb_func_lookup["toggle"].append({"func": _update_expectations})
  243. try:
  244. return checkbox.execute()
  245. except KeyboardInterrupt:
  246. return []
  247. @require_inquirer_py
  248. def _ask_for_confirmation_tui(message: str, default: bool = True) -> bool:
  249. return inquirer.confirm(message, default=default).execute()
  250. def _get_tui_choices_from_scan(
  251. repos: Iterable[CachedRepoInfo], preselected: List[str], sort_by: Optional[SortingOption_T] = None
  252. ) -> List:
  253. choices: List[Union["Choice", "Separator"]] = []
  254. choices.append(
  255. Choice(
  256. _CANCEL_DELETION_STR, name="None of the following (if selected, nothing will be deleted).", enabled=False
  257. )
  258. )
  259. sorted_repos = sorted(repos, key=lambda repo: _get_repo_sorting_key(repo, sort_by))
  260. for repo in sorted_repos:
  261. choices.append(
  262. Separator(
  263. f"\n{repo.repo_type.capitalize()} {repo.repo_id} ({repo.size_on_disk_str}, used {repo.last_accessed_str})"
  264. )
  265. )
  266. for revision in sorted(repo.revisions, key=_revision_sorting_order):
  267. choices.append(
  268. Choice(
  269. revision.commit_hash,
  270. name=(
  271. f"{revision.commit_hash[:8]}: {', '.join(sorted(revision.refs)) or '(detached)'} # modified {revision.last_modified_str}"
  272. ),
  273. enabled=revision.commit_hash in preselected,
  274. )
  275. )
  276. return choices
  277. def _manual_review_no_tui(
  278. hf_cache_info: HFCacheInfo, preselected: List[str], sort_by: Optional[SortingOption_T] = None
  279. ) -> List[str]:
  280. fd, tmp_path = mkstemp(suffix=".txt")
  281. os.close(fd)
  282. lines = []
  283. sorted_repos = sorted(hf_cache_info.repos, key=lambda repo: _get_repo_sorting_key(repo, sort_by))
  284. for repo in sorted_repos:
  285. lines.append(
  286. f"\n# {repo.repo_type.capitalize()} {repo.repo_id} ({repo.size_on_disk_str}, used {repo.last_accessed_str})"
  287. )
  288. for revision in sorted(repo.revisions, key=_revision_sorting_order):
  289. lines.append(
  290. f"{'' if revision.commit_hash in preselected else '#'} {revision.commit_hash} # Refs: {', '.join(sorted(revision.refs)) or '(detached)'} # modified {revision.last_modified_str}"
  291. )
  292. with open(tmp_path, "w") as f:
  293. f.write(_MANUAL_REVIEW_NO_TUI_INSTRUCTIONS)
  294. f.write("\n".join(lines))
  295. instructions = f"""
  296. TUI is disabled. In order to select which revisions you want to delete, please edit
  297. the following file using the text editor of your choice. Instructions for manual
  298. editing are located at the beginning of the file. Edit the file, save it and confirm
  299. to continue.
  300. File to edit: {ANSI.bold(tmp_path)}
  301. """
  302. print("\n".join(line.strip() for line in instructions.strip().split("\n")))
  303. while True:
  304. selected_hashes = _read_manual_review_tmp_file(tmp_path)
  305. if _ask_for_confirmation_no_tui(
  306. _get_expectations_str(hf_cache_info, selected_hashes) + " Continue ?", default=False
  307. ):
  308. break
  309. os.remove(tmp_path)
  310. return sorted(selected_hashes)
  311. def _ask_for_confirmation_no_tui(message: str, default: bool = True) -> bool:
  312. YES = ("y", "yes", "1")
  313. NO = ("n", "no", "0")
  314. DEFAULT = ""
  315. ALL = YES + NO + (DEFAULT,)
  316. full_message = message + (" (Y/n) " if default else " (y/N) ")
  317. while True:
  318. answer = input(full_message).lower()
  319. if answer == DEFAULT:
  320. return default
  321. if answer in YES:
  322. return True
  323. if answer in NO:
  324. return False
  325. print(f"Invalid input. Must be one of {ALL}")
  326. def _get_expectations_str(hf_cache_info: HFCacheInfo, selected_hashes: List[str]) -> str:
  327. if _CANCEL_DELETION_STR in selected_hashes:
  328. return "Nothing will be deleted."
  329. strategy = hf_cache_info.delete_revisions(*selected_hashes)
  330. return f"{len(selected_hashes)} revisions selected counting for {strategy.expected_freed_size_str}."
  331. def _read_manual_review_tmp_file(tmp_path: str) -> List[str]:
  332. with open(tmp_path) as f:
  333. content = f.read()
  334. lines = [line.strip() for line in content.split("\n")]
  335. selected_lines = [line for line in lines if not line.startswith("#")]
  336. selected_hashes = [line.split("#")[0].strip() for line in selected_lines]
  337. return [hash for hash in selected_hashes if len(hash) > 0]
  338. _MANUAL_REVIEW_NO_TUI_INSTRUCTIONS = f"""
  339. # INSTRUCTIONS
  340. # ------------
  341. # This is a temporary file created by running `hf cache delete --disable-tui`. It contains a set of revisions that can be deleted from your local cache directory.
  342. #
  343. # Please manually review the revisions you want to delete:
  344. # - Revision hashes can be commented out with '#'.
  345. # - Only non-commented revisions in this file will be deleted.
  346. # - Revision hashes that are removed from this file are ignored as well.
  347. # - If `{_CANCEL_DELETION_STR}` line is uncommented, the all cache deletion is cancelled and no changes will be applied.
  348. #
  349. # Once you've manually reviewed this file, please confirm deletion in the terminal. This file will be automatically removed once done.
  350. # ------------
  351. # KILL SWITCH
  352. # ------------
  353. # Un-comment following line to completely cancel the deletion process
  354. # {_CANCEL_DELETION_STR}
  355. # ------------
  356. # REVISIONS
  357. # ------------
  358. """.strip()
  359. def _revision_sorting_order(revision: CachedRevisionInfo) -> Any:
  360. return revision.last_modified