wehub-resource-sync
249 lines
11 KiB
Python
249 lines
11 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import importlib.util
|
|
import os
|
|
import requests
|
|
from modelscope.hub.api import HubApi, ModelScopeConfig
|
|
from modelscope.hub.utils.utils import get_cache_dir
|
|
from pathlib import Path
|
|
from tqdm import tqdm
|
|
from typing import List, Optional
|
|
|
|
from .env import use_hf_hub
|
|
from .logger import get_logger
|
|
from .torch_utils import is_local_master, safe_ddp_context
|
|
from .utils import subprocess_run
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def safe_snapshot_download(model_id_or_path: str,
|
|
revision: Optional[str] = None,
|
|
download_model: bool = True,
|
|
use_hf: Optional[bool] = None,
|
|
hub_token: Optional[str] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
check_local: bool = False,
|
|
**kwargs) -> str:
|
|
"""Download model snapshot safely with DDP context protection.
|
|
|
|
This function attempts to download a model from HuggingFace or ModelScope hub,
|
|
with support for local paths, subfolder specification, and distributed training
|
|
context protection. It handles various path formats and provides flexible
|
|
file filtering options.
|
|
|
|
Args:
|
|
model_id_or_path (str): The model identifier on the hub (e.g., 'Qwen/Qwen2.5-7B-Instruct')
|
|
or a local path to the model directory. Supports subfolder specification
|
|
using colon syntax (e.g., 'model_id:subfolder').
|
|
revision (Optional[str], optional): Specific model version/revision to download
|
|
(branch name, tag, or commit hash). Defaults to None (latest version).
|
|
download_model (bool, optional): Whether to download model weight files
|
|
(.bin, .safetensors). If False, only config and tokenizer files are
|
|
downloaded. Defaults to True.
|
|
use_hf (Optional[bool], optional): Force using HuggingFace Hub (True) or ModelScope (False).
|
|
If None, it is controlled by the environment variable `USE_HF`, which defaults to '0'.
|
|
Default: None.
|
|
hub_token (Optional[str], optional): Authentication token for accessing private
|
|
or gated models. Defaults to None.
|
|
ignore_patterns (Optional[List[str]], optional): List of glob patterns for files
|
|
to exclude from download. If None, uses default patterns to exclude zip,
|
|
gguf, pth, pt, and other auxiliary files. Defaults to None.
|
|
check_local (bool, optional): Whether to check for a local directory matching
|
|
the last component of model_id_or_path before attempting download.
|
|
Defaults to False.
|
|
**kwargs: Additional keyword arguments passed to the underlying hub download function.
|
|
|
|
Returns:
|
|
str: Absolute path to the model directory where files are stored.
|
|
|
|
Raises:
|
|
ValueError: If model_id_or_path starts with '/' (absolute path) and the path
|
|
does not exist.
|
|
Examples:
|
|
>>> # Download from hub
|
|
>>> model_dir = safe_snapshot_download('Qwen/Qwen2.5-7B-Instruct')
|
|
|
|
>>> # Download config only (no weights)
|
|
>>> model_dir = safe_snapshot_download('Qwen/Qwen2.5-7B-Instruct', download_model=False)
|
|
"""
|
|
from swift.hub import get_hub
|
|
if check_local:
|
|
model_suffix = model_id_or_path.rsplit('/', 1)[-1]
|
|
if os.path.exists(model_suffix):
|
|
model_dir = os.path.abspath(os.path.expanduser(model_suffix))
|
|
logger.info(f'Loading the model using local model_dir: {model_dir}')
|
|
return model_dir
|
|
if ignore_patterns is None:
|
|
ignore_patterns = [
|
|
'*.zip', '*.gguf', '*.pth', '*.pt', 'consolidated*', 'onnx/*', '*.safetensors.md', '*.msgpack', '*.onnx',
|
|
'*.ot', '*.h5'
|
|
]
|
|
if not download_model:
|
|
ignore_patterns += ['*.bin', '*.safetensors']
|
|
hub = get_hub(use_hf)
|
|
if model_id_or_path.startswith('~'):
|
|
model_id_or_path = os.path.abspath(os.path.expanduser(model_id_or_path))
|
|
model_path_to_check = '/'.join(model_id_or_path.split(':', 1))
|
|
if os.path.exists(model_id_or_path):
|
|
model_dir = model_id_or_path
|
|
sub_folder = None
|
|
elif os.path.exists(model_path_to_check):
|
|
model_dir = model_path_to_check
|
|
sub_folder = None
|
|
else:
|
|
if model_id_or_path.startswith('/'): # startswith
|
|
raise ValueError(f"path: '{model_id_or_path}' not found")
|
|
model_id_or_path = model_id_or_path.split(':', 1) # get sub_folder
|
|
if len(model_id_or_path) == 1:
|
|
model_id_or_path = [model_id_or_path[0], None]
|
|
model_id_or_path, sub_folder = model_id_or_path
|
|
if sub_folder is not None:
|
|
kwargs['allow_patterns'] = [f"{sub_folder.rstrip('/')}/*"]
|
|
with safe_ddp_context(hash_id=model_id_or_path):
|
|
model_dir = hub.download_model(model_id_or_path, revision, ignore_patterns, token=hub_token, **kwargs)
|
|
|
|
logger.info(f'Loading the model using model_dir: {model_dir}')
|
|
|
|
model_dir = os.path.abspath(os.path.expanduser(model_dir))
|
|
if sub_folder:
|
|
model_dir = os.path.join(model_dir, sub_folder)
|
|
assert os.path.isdir(model_dir), f'model_dir: {model_dir}'
|
|
return model_dir
|
|
|
|
|
|
def git_clone_github(github_url: str,
|
|
*,
|
|
local_repo_name: Optional[str] = None,
|
|
branch: Optional[str] = None,
|
|
commit_hash: Optional[str] = None) -> str:
|
|
if github_url.endswith('.git'):
|
|
github_url = github_url[:-4]
|
|
git_cache_dir = os.path.join(get_cache_dir(), '_github')
|
|
os.makedirs(git_cache_dir, exist_ok=True)
|
|
if local_repo_name is None:
|
|
github_url = github_url.rstrip('/')
|
|
local_repo_name = github_url.rsplit('/', 1)[1]
|
|
github_url = f'{github_url}.git'
|
|
local_repo_path = os.path.join(git_cache_dir, local_repo_name)
|
|
with safe_ddp_context('git_clone', use_barrier=True):
|
|
repo_existed = os.path.exists(local_repo_path)
|
|
if not is_local_master() and repo_existed:
|
|
return local_repo_path
|
|
if repo_existed:
|
|
command = ['git', '-C', local_repo_path, 'fetch']
|
|
subprocess_run(command)
|
|
if branch is not None:
|
|
command = ['git', '-C', local_repo_path, 'checkout', branch]
|
|
subprocess_run(command)
|
|
else:
|
|
command = ['git', '-C', git_cache_dir, 'clone', github_url, local_repo_name]
|
|
if branch is not None:
|
|
command += ['--branch', branch]
|
|
subprocess_run(command)
|
|
|
|
if commit_hash is not None:
|
|
command = ['git', '-C', local_repo_path, 'reset', '--hard', commit_hash]
|
|
subprocess_run(command)
|
|
elif repo_existed:
|
|
command = ['git', '-C', local_repo_path, 'pull']
|
|
subprocess_run(command)
|
|
logger.info(f'local_repo_path: {local_repo_path}')
|
|
return local_repo_path
|
|
|
|
|
|
def download_ms_file(url: str, local_path: str, cookies=None) -> None:
|
|
if cookies is None:
|
|
cookies = ModelScopeConfig.get_cookies()
|
|
resp = requests.get(url, cookies=cookies, stream=True)
|
|
with open(local_path, 'wb') as f:
|
|
for data in tqdm(resp.iter_lines()):
|
|
f.write(data)
|
|
|
|
|
|
def _resolve_kernel_variant_str(repo_id: str) -> Optional[str]:
|
|
"""Resolve the kernel build variant matching the current torch/cuda/platform
|
|
by listing the ``build/`` folder of the ModelScope kernel repository. Returns
|
|
``None`` if listing or parsing fails (caller should fall back to downloading
|
|
the whole repo).
|
|
"""
|
|
try:
|
|
from kernels.variants import parse_variant, resolve_variant
|
|
files = HubApi().get_model_files(repo_id, root='build', recursive=False)
|
|
variants = []
|
|
for f in files:
|
|
name = f.get('Name') or f.get('Path', '').rsplit('/', 1)[-1]
|
|
if not name:
|
|
continue
|
|
try:
|
|
variants.append(parse_variant(name))
|
|
except ValueError:
|
|
continue
|
|
variant = resolve_variant(variants)
|
|
return variant.variant_str if variant else None
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def patch_kernels() -> bool:
|
|
"""Install a process-wide monkey patch on
|
|
``transformers.integrations.hub_kernels.get_kernel`` so that kernel
|
|
repositories are downloaded from ModelScope and loaded via
|
|
``kernels.get_local_kernel``.
|
|
|
|
The runtime behavior is controlled by the ``USE_HF`` env (read on each
|
|
``get_kernel`` call):
|
|
- ``USE_HF=1``: fall back to the original HuggingFace-based loading.
|
|
- otherwise (default): use ModelScope.
|
|
|
|
Returns True if the patch was installed, False if skipped (``kernels`` not
|
|
installed, or the ``transformers`` integration is unavailable). Callers are
|
|
expected to guarantee idempotency (e.g. via a module-level flag).
|
|
"""
|
|
if importlib.util.find_spec('kernels') is None:
|
|
return False
|
|
try:
|
|
from kernels import get_local_kernel
|
|
from transformers.integrations import hub_kernels
|
|
except ImportError:
|
|
return False
|
|
|
|
origin_get_kernel = hub_kernels.get_kernel
|
|
|
|
def patched_get_kernel(repo_id, *args, **kwargs):
|
|
if use_hf_hub():
|
|
return origin_get_kernel(repo_id, *args, **kwargs)
|
|
try:
|
|
variant_str = _resolve_kernel_variant_str(repo_id)
|
|
allow_patterns = [f'build/{variant_str}/*'] if variant_str else None
|
|
model_dir = safe_snapshot_download(repo_id, use_hf=False, allow_patterns=allow_patterns)
|
|
package_name = repo_id.split('/')[-1].replace('-', '_')
|
|
# kernels < 0.14
|
|
kernel = get_local_kernel(Path(model_dir), package_name)
|
|
logger.info(f'Loaded kernel `{repo_id}` from ModelScope: {model_dir}')
|
|
return kernel
|
|
except Exception as e:
|
|
logger.warning(f'Failed to load kernel `{repo_id}` from ModelScope ({e}), fallback to HuggingFace.')
|
|
return origin_get_kernel(repo_id, *args, **kwargs)
|
|
|
|
hub_kernels.get_kernel = patched_get_kernel
|
|
return True
|
|
|
|
|
|
def download_file(url: str) -> str:
|
|
url = url.rstrip('/')
|
|
file_name = url.rsplit('/', 1)[-1]
|
|
cache_dir = os.path.join(get_cache_dir(), 'files')
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
file_path = os.path.join(cache_dir, file_name)
|
|
if os.path.exists(file_path):
|
|
return file_path
|
|
resp = requests.get(url, stream=True)
|
|
resp.raise_for_status()
|
|
total_size = int(resp.headers.get('content-length', 0))
|
|
with open(file_path, 'wb') as f, tqdm(
|
|
total=total_size, unit='B', unit_scale=True, unit_divisor=1024, desc=file_name) as pbar:
|
|
for chunk in resp.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
pbar.update(len(chunk))
|
|
return file_path
|