diff --git a/modelscope/cli/clearcache.py b/modelscope/cli/clearcache.py index dcd3d1df..0713db94 100644 --- a/modelscope/cli/clearcache.py +++ b/modelscope/cli/clearcache.py @@ -2,11 +2,13 @@ import os import shutil from argparse import ArgumentParser -from pathlib import Path from modelscope.cli.base import CLICommand from modelscope.hub.constants import TEMPORARY_FOLDER_NAME from modelscope.hub.utils.utils import get_model_masked_directory +from modelscope.utils.file_utils import (get_dataset_cache_root, + get_model_cache_root, + get_modelscope_cache_dir) def subparser_func(args): @@ -20,9 +22,7 @@ class ClearCacheCMD(CLICommand): def __init__(self, args): self.args = args - self.cache_dir = os.getenv( - 'MODELSCOPE_CACHE', - Path.home().joinpath('.cache', 'modelscope')) + self.cache_dir = get_modelscope_cache_dir() @staticmethod def define_args(parsers: ArgumentParser): @@ -76,11 +76,11 @@ class ClearCacheCMD(CLICommand): self._remove_directory(self.cache_dir) print('Cache cleared.') else: - entity_directory = os.path.join( - self.cache_dir, 'hub' if single_model else 'datasets', id) - temp_directory = os.path.join( - self.cache_dir, 'hub' if single_model else 'datasets', - TEMPORARY_FOLDER_NAME, id) + entity_root = get_model_cache_root( + ) if single_model else get_dataset_cache_root() + entity_directory = os.path.join(entity_root, id) + temp_directory = os.path.join(entity_root, + TEMPORARY_FOLDER_NAME, id) entity_removed = self._remove_directory(entity_directory) temp_removed = self._remove_directory(temp_directory) if (not entity_removed) and (not temp_removed): diff --git a/tests/cli/test_scancache_cmd.py b/tests/cli/test_scancache_cmd.py index a5c409b6..bde80db8 100644 --- a/tests/cli/test_scancache_cmd.py +++ b/tests/cli/test_scancache_cmd.py @@ -1,8 +1,12 @@ import os import subprocess import sys +import tempfile import unittest +from argparse import Namespace +from unittest import mock +from modelscope.hub.constants import TEMPORARY_FOLDER_NAME from modelscope.utils.file_utils import get_modelscope_cache_dir @@ -36,5 +40,47 @@ class TestScanCacheCommand(unittest.TestCase): self.assertIn('not found', output) +class TestClearCacheCommand(unittest.TestCase): + """clear-cache must resolve the right cache root in both scenarios: + * MODELSCOPE_CACHE set -> $MODELSCOPE_CACHE/{models,datasets}/ + * MODELSCOPE_CACHE unset -> ~/.cache/modelscope/hub/{models,datasets}/ + """ + + def _clear_and_assert(self, root, kind, entity_id): + from modelscope.cli.clearcache import ClearCacheCMD + sub = {'model': 'models', 'dataset': 'datasets'}[kind] + entity_dir = os.path.join(root, sub, entity_id) + temp_dir = os.path.join(root, sub, TEMPORARY_FOLDER_NAME, entity_id) + os.makedirs(entity_dir) + os.makedirs(temp_dir) + args = Namespace( + model=entity_id if kind == 'model' else None, + dataset=entity_id if kind == 'dataset' else None) + with mock.patch('builtins.input', return_value='Y'): + ClearCacheCMD(args).execute() + self.assertFalse(os.path.exists(entity_dir)) + self.assertFalse(os.path.exists(temp_dir)) + + def test_env_cache(self): + with tempfile.TemporaryDirectory() as tmp, \ + mock.patch.dict(os.environ, + {'MODELSCOPE_CACHE': tmp}, clear=False): + for kind in ('model', 'dataset'): + with self.subTest(kind=kind): + self._clear_and_assert(tmp, kind, f'org/{kind}') + + def test_default_cache(self): + with tempfile.TemporaryDirectory() as tmp: + env = { + k: v + for k, v in os.environ.items() if k != 'MODELSCOPE_CACHE' + } + env['HOME'] = tmp + with mock.patch.dict(os.environ, env, clear=True): + self._clear_and_assert( + os.path.join(tmp, '.cache', 'modelscope', 'hub'), 'model', + 'org/repo') + + if __name__ == '__main__': unittest.main()