diff --git a/modelscope/__init__.py b/modelscope/__init__.py index 14057ef9..3a23fc3e 100644 --- a/modelscope/__init__.py +++ b/modelscope/__init__.py @@ -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 diff --git a/modelscope/utils/hf_util/patcher.py b/modelscope/utils/hf_util/patcher.py index 1396e153..8aadb596 100644 --- a/modelscope/utils/hf_util/patcher.py +++ b/modelscope/utils/hf_util/patcher.py @@ -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.` that scopes the ModelScope + patch on `kernels.utils._get_hf_api` to the call itself, so + `from kernels import ` 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() diff --git a/tests/utils/test_hf_util_kernels.py b/tests/utils/test_hf_util_kernels.py new file mode 100644 index 00000000..3cafb17f --- /dev/null +++ b/tests/utils/test_hf_util_kernels.py @@ -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()