mirror of
https://github.com/modelscope/modelscope.git
synced 2026-08-29 10:08:40 +02:00
Support kernels downloading (#1697)
This commit is contained in:
tastelikefeet
@@ -119,21 +119,30 @@ else:
|
||||
extra_objects[_import] = getattr(hf_util, _import)
|
||||
|
||||
def try_import_from_hf(name):
|
||||
hf_pkgs = ['transformers', 'peft', 'diffusers']
|
||||
hf_pkgs = ['transformers', 'peft', 'diffusers', 'kernels']
|
||||
module = None
|
||||
matched_pkg = None
|
||||
for pkg in hf_pkgs:
|
||||
try:
|
||||
module = getattr(importlib.import_module(pkg), name)
|
||||
matched_pkg = pkg
|
||||
break
|
||||
except Exception: # noqa
|
||||
pass
|
||||
|
||||
if module is not None:
|
||||
module = _patch_pretrained_class([module], wrap=True)
|
||||
else:
|
||||
if module is None:
|
||||
raise AttributeError(
|
||||
f'Cannot import available module of {name} in modelscope,'
|
||||
f' or related packages({hf_pkgs})')
|
||||
|
||||
if matched_pkg == 'kernels':
|
||||
if callable(module):
|
||||
from modelscope.utils.hf_util.patcher import \
|
||||
_wrap_kernels_callable
|
||||
return _wrap_kernels_callable(name)
|
||||
return module
|
||||
|
||||
module = _patch_pretrained_class([module], wrap=True)
|
||||
return module[0]
|
||||
|
||||
import sys
|
||||
|
||||
@@ -8,7 +8,7 @@ import sys
|
||||
from asyncio import Future
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from types import MethodType
|
||||
from types import MethodType, SimpleNamespace
|
||||
from typing import BinaryIO, Dict, Iterable, List, Optional, Union
|
||||
|
||||
from modelscope.hub.constants import DEFAULT_MODELSCOPE_DATA_ENDPOINT
|
||||
@@ -515,6 +515,147 @@ def _unpatch_pretrained_class(all_imported_modules):
|
||||
delattr(dynamic_module_utils, 'origin_get_class_from_dynamic_module')
|
||||
|
||||
|
||||
def _patch_kernels():
|
||||
"""Monkey-patch the `kernels` library to route HF API calls to ModelScope.
|
||||
|
||||
Only `kernels.utils._get_hf_api` is replaced; every download, file check
|
||||
and ref listing performed by `kernels` goes through `_MsKernelApi`, so
|
||||
the kernel loading/variant/lock logic stays untouched.
|
||||
"""
|
||||
try:
|
||||
from kernels import utils as kernels_utils
|
||||
except ImportError:
|
||||
return
|
||||
if hasattr(kernels_utils, '_get_hf_api_origin'):
|
||||
return
|
||||
kernels_utils._get_hf_api_origin = kernels_utils._get_hf_api
|
||||
kernels_utils._get_hf_api = lambda user_agent=None: _MsKernelApi()
|
||||
|
||||
|
||||
def _unpatch_kernels():
|
||||
try:
|
||||
from kernels import utils as kernels_utils
|
||||
except ImportError:
|
||||
return
|
||||
origin = getattr(kernels_utils, '_get_hf_api_origin', None)
|
||||
if origin is not None:
|
||||
kernels_utils._get_hf_api = origin
|
||||
del kernels_utils._get_hf_api_origin
|
||||
|
||||
|
||||
def _ms_revision(revision):
|
||||
"""Translate an HF revision string into one ModelScope accepts."""
|
||||
return 'master' if revision in (None, 'main') else revision
|
||||
|
||||
|
||||
class _MsKernelApi:
|
||||
"""Minimal `HfApi` look-alike that forwards to ModelScope. Only the
|
||||
handful of methods that `kernels` actually calls are implemented.
|
||||
"""
|
||||
|
||||
def snapshot_download(self,
|
||||
repo_id,
|
||||
*,
|
||||
allow_patterns=None,
|
||||
ignore_patterns=None,
|
||||
cache_dir=None,
|
||||
revision=None,
|
||||
local_files_only=False,
|
||||
**kwargs):
|
||||
from modelscope import snapshot_download as ms_snapshot_download
|
||||
return ms_snapshot_download(
|
||||
repo_id,
|
||||
revision=_ms_revision(revision),
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
allow_patterns=allow_patterns,
|
||||
ignore_patterns=ignore_patterns,
|
||||
**kwargs)
|
||||
|
||||
def list_repo_tree(self,
|
||||
repo_id,
|
||||
*,
|
||||
path_in_repo=None,
|
||||
revision=None,
|
||||
**kwargs):
|
||||
from huggingface_hub.hf_api import RepoFolder
|
||||
from modelscope.hub.api import HubApi
|
||||
entries = HubApi().get_model_files(
|
||||
repo_id,
|
||||
revision=_ms_revision(revision),
|
||||
root=path_in_repo,
|
||||
recursive=False)
|
||||
folders = []
|
||||
for entry in entries:
|
||||
if entry.get('Type') != 'tree':
|
||||
continue
|
||||
path = entry.get('Path') or entry.get('Name')
|
||||
folders.append(RepoFolder(path=path, oid='', last_commit=None))
|
||||
return folders
|
||||
|
||||
def file_exists(self, repo_id, filename, *, revision=None, **kwargs):
|
||||
from modelscope.hub.api import HubApi
|
||||
return HubApi().file_exists(
|
||||
repo_id, filename, revision=_ms_revision(revision))
|
||||
|
||||
def list_repo_refs(self, repo_id, **kwargs):
|
||||
from huggingface_hub.hf_api import GitRefInfo
|
||||
from modelscope.hub.api import HubApi
|
||||
branches, tags = HubApi().get_model_branches_and_tags(repo_id)
|
||||
# `target_commit` doubles as the revision in later calls, so reuse
|
||||
# the branch/tag name (ModelScope accepts it as a revision).
|
||||
return SimpleNamespace(
|
||||
branches=[
|
||||
GitRefInfo(name=n, ref=f'refs/heads/{n}', target_commit=n)
|
||||
for n in (branches or [])
|
||||
],
|
||||
tags=[
|
||||
GitRefInfo(name=n, ref=f'refs/tags/{n}', target_commit=n)
|
||||
for n in (tags or [])
|
||||
],
|
||||
converts=[])
|
||||
|
||||
def hf_hub_download(self, *args, **kwargs):
|
||||
# Only called for the optional `kernel-status.toml`. Raising
|
||||
# `EntryNotFoundError` makes kernels treat the repo as having no
|
||||
# redirect status.
|
||||
from huggingface_hub.errors import EntryNotFoundError
|
||||
raise EntryNotFoundError(
|
||||
'kernel-status.toml lookup is skipped on ModelScope')
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _kernels_patch_scope():
|
||||
"""Apply `_patch_kernels` for the duration of the `with` block, unless an
|
||||
outer patch (e.g. `patch_hub()`) is already in effect.
|
||||
"""
|
||||
from kernels import utils as kernels_utils
|
||||
if hasattr(kernels_utils, '_get_hf_api_origin'):
|
||||
yield
|
||||
return
|
||||
_patch_kernels()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_unpatch_kernels()
|
||||
|
||||
|
||||
def _wrap_kernels_callable(attr_name):
|
||||
"""Return a wrapper around `kernels.<attr_name>` that scopes the ModelScope
|
||||
patch on `kernels.utils._get_hf_api` to the call itself, so
|
||||
`from kernels import <attr_name>` stays on HuggingFace unless the user
|
||||
explicitly calls `patch_hub()` / `patch_context()`.
|
||||
"""
|
||||
|
||||
def _wrapped(*args, **kwargs):
|
||||
import kernels
|
||||
with _kernels_patch_scope():
|
||||
return getattr(kernels, attr_name)(*args, **kwargs)
|
||||
|
||||
_wrapped.__name__ = _wrapped.__qualname__ = attr_name
|
||||
return _wrapped
|
||||
|
||||
|
||||
def _patch_hub():
|
||||
import huggingface_hub
|
||||
from huggingface_hub import hf_api
|
||||
@@ -854,11 +995,13 @@ def _unpatch_hub():
|
||||
|
||||
def patch_hub():
|
||||
_patch_hub()
|
||||
_patch_kernels()
|
||||
_patch_pretrained_class(get_all_imported_modules())
|
||||
|
||||
|
||||
def unpatch_hub():
|
||||
_unpatch_pretrained_class(get_all_imported_modules())
|
||||
_unpatch_kernels()
|
||||
_unpatch_hub()
|
||||
|
||||
|
||||
|
||||
274
tests/utils/test_hf_util_kernels.py
Normal file
274
tests/utils/test_hf_util_kernels.py
Normal file
@@ -0,0 +1,274 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
"""Unit tests for the `kernels` library monkey-patch in modelscope.
|
||||
|
||||
Requirements verified:
|
||||
|
||||
1. `from modelscope import get_kernel` works without `patch_hub()` first and
|
||||
routes downloads through ModelScope, without leaking the
|
||||
`kernels.utils._get_hf_api` patch to anyone else.
|
||||
2. `patch_hub()` / `patch_context()` makes `from kernels import get_kernel`
|
||||
also route downloads through ModelScope; `unpatch_hub()` restores it.
|
||||
"""
|
||||
import contextlib
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from modelscope.utils.hf_util.patcher import (_patch_kernels, _unpatch_kernels,
|
||||
patch_context, patch_hub,
|
||||
unpatch_hub)
|
||||
|
||||
|
||||
def _ensure_kernels_installed():
|
||||
try:
|
||||
from kernels import get_kernel # noqa: F401
|
||||
from kernels.utils import _get_hf_api # noqa: F401
|
||||
except ImportError:
|
||||
subprocess.check_call(
|
||||
[sys.executable, '-m', 'pip', 'install', '-q', 'kernels'])
|
||||
for mod in list(sys.modules):
|
||||
if mod == 'kernels' or mod.startswith('kernels.'):
|
||||
sys.modules.pop(mod, None)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _isolate_hub_patches():
|
||||
"""Neutralize the non-kernels parts of `patch_hub()` so tests focus on
|
||||
the kernels monkey-patch behaviour only.
|
||||
"""
|
||||
targets = [
|
||||
'modelscope.utils.hf_util.patcher._patch_hub',
|
||||
'modelscope.utils.hf_util.patcher._unpatch_hub',
|
||||
'modelscope.utils.hf_util.patcher._patch_pretrained_class',
|
||||
'modelscope.utils.hf_util.patcher._unpatch_pretrained_class',
|
||||
]
|
||||
with contextlib.ExitStack() as stack:
|
||||
for t in targets:
|
||||
stack.enter_context(patch(t))
|
||||
stack.enter_context(
|
||||
patch(
|
||||
'modelscope.utils.hf_util.patcher.get_all_imported_modules',
|
||||
return_value=[]))
|
||||
yield
|
||||
|
||||
|
||||
class _KernelsTestBase(unittest.TestCase):
|
||||
"""Installs `kernels`, captures the original `_get_hf_api`, and keeps the
|
||||
state clean between tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
_ensure_kernels_installed()
|
||||
|
||||
def setUp(self):
|
||||
_unpatch_kernels()
|
||||
# Instance attribute, so the descriptor protocol is not triggered.
|
||||
from kernels.utils import _get_hf_api
|
||||
from kernels import utils as kernels_utils
|
||||
self.original_get_hf_api = _get_hf_api
|
||||
self.kernels_utils = kernels_utils
|
||||
|
||||
def tearDown(self):
|
||||
_unpatch_kernels()
|
||||
import modelscope
|
||||
for name in ('get_kernel', 'has_kernel', 'install_kernel',
|
||||
'load_kernel', 'get_locked_kernel', 'snapshot_download'):
|
||||
try:
|
||||
delattr(modelscope, name)
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
|
||||
class KernelsProxyApiTest(_KernelsTestBase):
|
||||
"""Low-level proxy API behaviour exercised against the real
|
||||
`kernels.utils` module.
|
||||
"""
|
||||
|
||||
def _patched_api(self):
|
||||
_patch_kernels()
|
||||
return self.kernels_utils._get_hf_api()
|
||||
|
||||
def test_patch_replaces_get_hf_api(self):
|
||||
api = self._patched_api()
|
||||
self.assertTrue(hasattr(self.kernels_utils, '_get_hf_api_origin'))
|
||||
for name in ('snapshot_download', 'list_repo_tree', 'file_exists',
|
||||
'list_repo_refs', 'hf_hub_download'):
|
||||
self.assertTrue(callable(getattr(api, name, None)), name)
|
||||
|
||||
def test_patch_is_idempotent(self):
|
||||
_patch_kernels()
|
||||
first = self.kernels_utils._get_hf_api
|
||||
_patch_kernels()
|
||||
self.assertIs(self.kernels_utils._get_hf_api, first)
|
||||
|
||||
def test_unpatch_restores_original(self):
|
||||
_patch_kernels()
|
||||
_unpatch_kernels()
|
||||
self.assertFalse(hasattr(self.kernels_utils, '_get_hf_api_origin'))
|
||||
self.assertIs(self.kernels_utils._get_hf_api, self.original_get_hf_api)
|
||||
|
||||
def test_snapshot_download_routes_to_modelscope(self):
|
||||
api = self._patched_api()
|
||||
with patch(
|
||||
'modelscope.hub.snapshot_download.snapshot_download',
|
||||
return_value='/tmp/fake_path') as mocked:
|
||||
result = api.snapshot_download(
|
||||
'foo/bar',
|
||||
allow_patterns=['build/*'],
|
||||
ignore_patterns=['*.md'],
|
||||
revision='main',
|
||||
max_workers=4)
|
||||
self.assertEqual(result, '/tmp/fake_path')
|
||||
kwargs = mocked.call_args.kwargs
|
||||
# `main` is normalized to `master` for ModelScope.
|
||||
self.assertEqual(kwargs['revision'], 'master')
|
||||
# `allow_patterns` / `ignore_patterns` are forwarded as-is and extra
|
||||
# kwargs (e.g. `max_workers`) are passed through.
|
||||
self.assertEqual(kwargs['allow_patterns'], ['build/*'])
|
||||
self.assertEqual(kwargs['ignore_patterns'], ['*.md'])
|
||||
self.assertEqual(kwargs['max_workers'], 4)
|
||||
|
||||
def test_hf_hub_download_raises_entry_not_found(self):
|
||||
from huggingface_hub.errors import EntryNotFoundError
|
||||
api = self._patched_api()
|
||||
with self.assertRaises(EntryNotFoundError):
|
||||
api.hf_hub_download(repo_id='foo/bar', filename='x.toml')
|
||||
|
||||
def test_file_exists_routes_to_hubapi(self):
|
||||
api = self._patched_api()
|
||||
fake = MagicMock()
|
||||
fake.file_exists.return_value = True
|
||||
with patch('modelscope.hub.api.HubApi', return_value=fake):
|
||||
self.assertTrue(api.file_exists('foo/bar', 'README.md'))
|
||||
fake.file_exists.assert_called_once_with(
|
||||
'foo/bar', 'README.md', revision='master')
|
||||
|
||||
def test_list_repo_refs_routes_to_hubapi(self):
|
||||
api = self._patched_api()
|
||||
fake = MagicMock()
|
||||
fake.get_model_branches_and_tags.return_value = (['master',
|
||||
'v1'], ['r1.0'])
|
||||
with patch('modelscope.hub.api.HubApi', return_value=fake):
|
||||
refs = api.list_repo_refs('foo/bar')
|
||||
self.assertEqual([b.name for b in refs.branches], ['master', 'v1'])
|
||||
self.assertEqual([t.name for t in refs.tags], ['r1.0'])
|
||||
# `target_commit` doubles as the ModelScope revision for later calls.
|
||||
self.assertEqual(refs.branches[1].target_commit, 'v1')
|
||||
|
||||
|
||||
class PatchHubFlowTest(_KernelsTestBase):
|
||||
"""Requirement 2: `patch_hub` / `patch_context` toggle the kernels patch."""
|
||||
|
||||
def test_patch_hub_then_unpatch_hub_round_trip(self):
|
||||
with _isolate_hub_patches():
|
||||
self.assertIs(self.kernels_utils._get_hf_api,
|
||||
self.original_get_hf_api)
|
||||
|
||||
patch_hub()
|
||||
self.assertTrue(hasattr(self.kernels_utils, '_get_hf_api_origin'))
|
||||
self.assertTrue(
|
||||
hasattr(self.kernels_utils._get_hf_api(), 'snapshot_download'))
|
||||
|
||||
unpatch_hub()
|
||||
self.assertFalse(hasattr(self.kernels_utils, '_get_hf_api_origin'))
|
||||
self.assertIs(self.kernels_utils._get_hf_api,
|
||||
self.original_get_hf_api)
|
||||
|
||||
def test_patch_context_round_trip(self):
|
||||
with _isolate_hub_patches():
|
||||
with patch_context():
|
||||
self.assertTrue(
|
||||
hasattr(self.kernels_utils, '_get_hf_api_origin'))
|
||||
self.assertFalse(hasattr(self.kernels_utils, '_get_hf_api_origin'))
|
||||
|
||||
|
||||
class ModelscopeImportTest(_KernelsTestBase):
|
||||
"""Requirement 1: `from modelscope import get_kernel` delegates to the
|
||||
real `kernels.get_kernel`, scoping the patch to the wrapped call only.
|
||||
"""
|
||||
|
||||
def _fake_get_kernel_capturing_api(self, sink):
|
||||
"""Build a fake `kernels.get_kernel` that records `_get_hf_api()` at
|
||||
call time, so the test can observe the patch state mid-call.
|
||||
"""
|
||||
|
||||
def _fake(*args, **kwargs):
|
||||
sink['api'] = self.kernels_utils._get_hf_api()
|
||||
sink['args'] = args
|
||||
sink['kwargs'] = kwargs
|
||||
return 'kernel-module'
|
||||
|
||||
return _fake
|
||||
|
||||
def test_from_modelscope_wraps_kernels_get_kernel(self):
|
||||
import kernels
|
||||
import modelscope
|
||||
|
||||
ms_get_kernel = modelscope.get_kernel
|
||||
self.assertIsNot(ms_get_kernel, kernels.get_kernel)
|
||||
# Before any call, `_get_hf_api` stays the original.
|
||||
self.assertIs(self.kernels_utils._get_hf_api, self.original_get_hf_api)
|
||||
|
||||
captured = {}
|
||||
with patch.object(kernels, 'get_kernel',
|
||||
self._fake_get_kernel_capturing_api(captured)):
|
||||
result = ms_get_kernel('foo/bar', revision='v1')
|
||||
|
||||
self.assertEqual(result, 'kernel-module')
|
||||
self.assertEqual(captured['args'], ('foo/bar', ))
|
||||
self.assertEqual(captured['kwargs'], {'revision': 'v1'})
|
||||
# Mid-call: ModelScope proxy was active.
|
||||
self.assertTrue(hasattr(captured['api'], 'snapshot_download'))
|
||||
# After the call: patch is rolled back.
|
||||
self.assertIs(self.kernels_utils._get_hf_api, self.original_get_hf_api)
|
||||
|
||||
def test_wrapped_call_nests_inside_patch_hub(self):
|
||||
import kernels
|
||||
import modelscope
|
||||
captured = {}
|
||||
with _isolate_hub_patches():
|
||||
patch_hub()
|
||||
with patch.object(kernels, 'get_kernel',
|
||||
self._fake_get_kernel_capturing_api(captured)):
|
||||
modelscope.get_kernel('foo/bar')
|
||||
|
||||
# The outer `patch_hub` patch survives the wrapped call.
|
||||
self.assertTrue(hasattr(self.kernels_utils, '_get_hf_api_origin'))
|
||||
self.assertTrue(hasattr(captured['api'], 'snapshot_download'))
|
||||
|
||||
unpatch_hub()
|
||||
self.assertIs(self.kernels_utils._get_hf_api,
|
||||
self.original_get_hf_api)
|
||||
|
||||
|
||||
class TinyGradRMSIntegrationTest(_KernelsTestBase):
|
||||
"""Real-world check using `kernels-community/tinygrad-rms`, which is
|
||||
published on both HuggingFace and ModelScope. Verifies that the
|
||||
ModelScope download path actually works end to end.
|
||||
"""
|
||||
|
||||
REPO = 'kernels-community/tinygrad-rms'
|
||||
|
||||
def test_from_modelscope_get_kernel(self):
|
||||
import modelscope
|
||||
# Routes through `try_import_from_hf` and scopes the ModelScope
|
||||
# monkey-patch to this single call.
|
||||
module = modelscope.get_kernel(self.REPO)
|
||||
self.assertIsNotNone(module)
|
||||
# Wrapper must leave `_get_hf_api` restored afterwards.
|
||||
self.assertIs(self.kernels_utils._get_hf_api, self.original_get_hf_api)
|
||||
|
||||
def test_patch_hub_then_kernels_get_kernel(self):
|
||||
with _isolate_hub_patches(), patch_context():
|
||||
from kernels import get_kernel
|
||||
module = get_kernel(self.REPO)
|
||||
self.assertIsNotNone(module)
|
||||
# `patch_context` rolled the kernels patch back on exit.
|
||||
self.assertIs(self.kernels_utils._get_hf_api, self.original_get_hf_api)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user