mirror of
https://github.com/modelscope/modelscope.git
synced 2026-09-02 12:09:36 +02:00
114 lines
4.0 KiB
Python
114 lines
4.0 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
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
|
|
|
|
|
|
def subparser_func(args):
|
|
""" Function which will be called for a specific sub parser.
|
|
"""
|
|
return ClearCacheCMD(args)
|
|
|
|
|
|
class ClearCacheCMD(CLICommand):
|
|
name = 'clear-cache'
|
|
|
|
def __init__(self, args):
|
|
self.args = args
|
|
self.cache_dir = os.getenv(
|
|
'MODELSCOPE_CACHE',
|
|
Path.home().joinpath('.cache', 'modelscope'))
|
|
|
|
@staticmethod
|
|
def define_args(parsers: ArgumentParser):
|
|
""" define args for clear-cache command.
|
|
"""
|
|
parser = parsers.add_parser(ClearCacheCMD.name)
|
|
group = parser.add_mutually_exclusive_group()
|
|
group.add_argument(
|
|
'--model',
|
|
type=str,
|
|
help=
|
|
'The id of the model whose cache will be cleared. For clear-cache, '
|
|
'if neither model or dataset id is provided, entire cache will be cleared.'
|
|
)
|
|
group.add_argument(
|
|
'--dataset',
|
|
type=str,
|
|
help=
|
|
'The id of the dataset whose cache will be cleared. For clear-cache, '
|
|
'if neither model or dataset id is provided, entire cache will be cleared.'
|
|
)
|
|
|
|
parser.set_defaults(func=subparser_func)
|
|
|
|
def execute(self):
|
|
self._execute_with_confirmation()
|
|
|
|
def _execute_with_confirmation(self):
|
|
all = False
|
|
single_model = False
|
|
prompt = '\nYou are about to delete '
|
|
|
|
if self.args.model or self.args.dataset:
|
|
if self.args.model:
|
|
id = self.args.model
|
|
single_model = True
|
|
prompt = prompt + f'local cache for model {id}. '
|
|
else:
|
|
id = self.args.dataset
|
|
prompt = prompt + f'local cache for dataset {id}. '
|
|
else:
|
|
prompt = prompt + f'entire ModelScope cache at {self.cache_dir}, including ALL models and dataset.\n'
|
|
all = True
|
|
user_input = input(
|
|
prompt
|
|
+ '\nPlease press Y or y to proceed, any other key to abort.\n'
|
|
).strip().upper()
|
|
|
|
if user_input == 'Y':
|
|
if all:
|
|
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_removed = self._remove_directory(entity_directory)
|
|
temp_removed = self._remove_directory(temp_directory)
|
|
if (not entity_removed) and (not temp_removed):
|
|
if single_model:
|
|
print(
|
|
f'Cache for Model {id} not found. Nothing to do.')
|
|
else:
|
|
print(
|
|
f'Cache for Dataset {id} not found. Nothing to do.'
|
|
)
|
|
else:
|
|
print('Cache cleared.')
|
|
else:
|
|
print('Operation aborted.')
|
|
return
|
|
|
|
def _remove_directory(self, path):
|
|
if os.path.exists(path):
|
|
try:
|
|
if os.path.islink(path):
|
|
shutil.rmtree(os.readlink(path))
|
|
os.remove(path)
|
|
print(f'Cache and link for {path} removed.')
|
|
else:
|
|
shutil.rmtree(path)
|
|
print(f'Cache folder {path} removed.')
|
|
return True
|
|
except Exception as e:
|
|
print(f'An error occurred while clearing cache at {path}: {e}')
|
|
return False
|