From 346af6773fd18fe1cd13030eddd0c933cf1619a1 Mon Sep 17 00:00:00 2001 From: pangda Date: Wed, 11 Jan 2023 10:35:09 +0800 Subject: [PATCH] support plugin mechanism for second-party/third-party modules --- modelscope/pipelines/builder.py | 12 ++ modelscope/utils/plugins.py | 215 +++++++++++++++++++++++ tests/utils/plugins/.modelscope_plugins | 1 + tests/utils/plugins/dummy/__init__.py | 1 + tests/utils/plugins/dummy/dummy_model.py | 8 + tests/utils/test_plugin.py | 41 +++++ 6 files changed, 278 insertions(+) create mode 100644 modelscope/utils/plugins.py create mode 100644 tests/utils/plugins/.modelscope_plugins create mode 100644 tests/utils/plugins/dummy/__init__.py create mode 100644 tests/utils/plugins/dummy/dummy_model.py create mode 100644 tests/utils/test_plugin.py diff --git a/modelscope/pipelines/builder.py b/modelscope/pipelines/builder.py index 2514b367..613b9430 100644 --- a/modelscope/pipelines/builder.py +++ b/modelscope/pipelines/builder.py @@ -316,6 +316,7 @@ def pipeline(task: str = None, framework: str = None, device: str = 'gpu', model_revision: Optional[str] = DEFAULT_MODEL_REVISION, + plugins: List[str] = None, **kwargs) -> Pipeline: """ Factory method to build an obj:`Pipeline`. @@ -349,6 +350,8 @@ def pipeline(task: str = None, if task is None and pipeline_name is None: raise ValueError('task or pipeline_name is required') + try_import_plugins(plugins) + model = normalize_model_input(model, model_revision) pipeline_props = {'type': pipeline_name} if pipeline_name is None: @@ -362,6 +365,7 @@ def pipeline(task: str = None, model, str) else read_config( model[0], revision=model_revision) check_config(cfg) + try_import_plugins(cfg.safe_get('plugins')) pipeline_props = cfg.pipeline elif model is not None: # get pipeline info from Model object @@ -370,6 +374,7 @@ def pipeline(task: str = None, # model is instantiated by user, we should parse config again cfg = read_config(first_model.model_dir) check_config(cfg) + try_import_plugins(cfg.safe_get('plugins')) first_model.pipeline = cfg.pipeline pipeline_props = first_model.pipeline else: @@ -427,3 +432,10 @@ def get_default_pipeline_info(task): else: pipeline_name, default_model = DEFAULT_MODEL_FOR_PIPELINE[task] return pipeline_name, default_model + + +def try_import_plugins(plugins: List[str]) -> None: + """ Try to import plugins """ + if plugins is not None: + from modelscope.utils.plugins import import_plugins + import_plugins(plugins) diff --git a/modelscope/utils/plugins.py b/modelscope/utils/plugins.py new file mode 100644 index 00000000..6c2f2975 --- /dev/null +++ b/modelscope/utils/plugins.py @@ -0,0 +1,215 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +# This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp +import importlib +import os +import pkgutil +import sys +from contextlib import contextmanager +from fnmatch import fnmatch +from pathlib import Path +from typing import Iterable, List, Optional, Set + +from modelscope.utils.logger import get_logger + +logger = get_logger() + +LOCAL_PLUGINS_FILENAME = '.modelscope_plugins' +GLOBAL_PLUGINS_FILENAME = os.path.join(Path.home(), '.modelscope', 'plugins') +DEFAULT_PLUGINS = [] + + +@contextmanager +def pushd(new_dir: str, verbose: bool = False): + """ + Changes the current directory to the given path and prepends it to `sys.path`. + This method is intended to use with `with`, so after its usage, the current + directory will be set to the previous value. + """ + previous_dir = os.getcwd() + if verbose: + logger.info(f'Changing directory to {new_dir}') # type: ignore + os.chdir(new_dir) + try: + yield + finally: + if verbose: + logger.info(f'Changing directory back to {previous_dir}') + os.chdir(previous_dir) + + +@contextmanager +def push_python_path(path: str): + """ + Prepends the given path to `sys.path`. + This method is intended to use with `with`, so after its usage, its value + will be removed from `sys.path`. + """ + path = Path(path).resolve() + path = str(path) + sys.path.insert(0, path) + try: + yield + finally: + sys.path.remove(path) + + +def discover_file_plugins( + filename: str = LOCAL_PLUGINS_FILENAME) -> Iterable[str]: + """ + Discover plugins from file + """ + with open(filename) as f: + for module_name in f: + module_name = module_name.strip() + if module_name: + yield module_name + + +def discover_plugins() -> Iterable[str]: + """ + Discover plugins + """ + plugins: Set[str] = set() + if os.path.isfile(LOCAL_PLUGINS_FILENAME): + with push_python_path('.'): + for plugin in discover_file_plugins(LOCAL_PLUGINS_FILENAME): + if plugin in plugins: + continue + yield plugin + plugins.add(plugin) + if os.path.isfile(GLOBAL_PLUGINS_FILENAME): + for plugin in discover_file_plugins(GLOBAL_PLUGINS_FILENAME): + if plugin in plugins: + continue + yield plugin + plugins.add(plugin) + + +def import_all_plugins(plugins: List[str] = None) -> List[str]: + """ + Imports default plugins, input plugins and file discovered plugins. + """ + import_module_and_submodules( + 'modelscope', + include={ + 'modelscope.metrics.builder', + 'modelscope.models.builder', + 'modelscope.pipelines.builder', + 'modelscope.preprocessors.builder', + 'modelscope.trainers.builder', + }, + exclude={ + 'modelscope.metrics.*', + 'modelscope.models.*', + 'modelscope.pipelines.*', + 'modelscope.preprocessors.*', + 'modelscope.trainers.*', + 'modelscope.msdatasets', + 'modelscope.utils', + 'modelscope.exporters', + }) + + imported_plugins: List[str] = [] + + imported_plugins.extend(import_plugins(DEFAULT_PLUGINS)) + imported_plugins.extend(import_plugins(plugins)) + imported_plugins.extend(import_file_plugins()) + + return imported_plugins + + +def import_plugins(plugins: List[str] = None) -> List[str]: + """ + Imports the plugins listed in the arguments. + """ + imported_plugins: List[str] = [] + if plugins is None or len(plugins) == 0: + return imported_plugins + + # Workaround for a presumed Python issue where spawned processes can't find modules in the current directory. + cwd = os.getcwd() + if cwd not in sys.path: + sys.path.append(cwd) + + for module_name in plugins: + try: + import_module_and_submodules(module_name) + logger.info('Plugin %s available', module_name) + imported_plugins.append(module_name) + except ModuleNotFoundError as e: + logger.error(f'Plugin {module_name} could not be loaded: {e}') + + return imported_plugins + + +def import_file_plugins() -> List[str]: + """ + Imports the plugins found with `discover_plugins()`. + """ + imported_plugins: List[str] = [] + + # Workaround for a presumed Python issue where spawned processes can't find modules in the current directory. + cwd = os.getcwd() + if cwd not in sys.path: + sys.path.append(cwd) + + for module_name in discover_plugins(): + try: + importlib.import_module(module_name) + logger.info('Plugin %s available', module_name) + imported_plugins.append(module_name) + except ModuleNotFoundError as e: + logger.error(f'Plugin {module_name} could not be loaded: {e}') + + return imported_plugins + + +def import_module_and_submodules(package_name: str, + include: Optional[Set[str]] = None, + exclude: Optional[Set[str]] = None) -> None: + """ + Import all public submodules under the given package. + """ + # take care of None + include = include if include else set() + exclude = exclude if exclude else set() + + def fn_in(packge_name: str, pattern_set: Set[str]) -> bool: + for pattern in pattern_set: + if fnmatch(package_name, pattern): + return True + return False + + if not fn_in(package_name, include) and fn_in(package_name, exclude): + return + + importlib.invalidate_caches() + + # For some reason, python doesn't always add this by default to your path, but you pretty much + # always want it when using `--include-package`. And if it's already there, adding it again at + # the end won't hurt anything. + with push_python_path('.'): + # Import at top level + try: + module = importlib.import_module(package_name) + path = getattr(module, '__path__', []) + path_string = '' if not path else path[0] + + # walk_packages only finds immediate children, so need to recurse. + for module_finder, name, _ in pkgutil.walk_packages(path): + # Sometimes when you import third-party libraries that are on your path, + # `pkgutil.walk_packages` returns those too, so we need to skip them. + if path_string and module_finder.path != path_string: # type: ignore[union-attr] + continue + if name.startswith('_'): + # skip directly importing private subpackages + continue + if name.startswith('test'): + # skip tests + continue + subpackage = f'{package_name}.{name}' + import_module_and_submodules(subpackage, exclude=exclude) + except Exception as e: + logger.warning(f'{package_name} not imported: {str(e)}') + if len(package_name.split('.')) == 1: + raise ModuleNotFoundError('Package not installed') diff --git a/tests/utils/plugins/.modelscope_plugins b/tests/utils/plugins/.modelscope_plugins new file mode 100644 index 00000000..421376db --- /dev/null +++ b/tests/utils/plugins/.modelscope_plugins @@ -0,0 +1 @@ +dummy diff --git a/tests/utils/plugins/dummy/__init__.py b/tests/utils/plugins/dummy/__init__.py new file mode 100644 index 00000000..a0d86001 --- /dev/null +++ b/tests/utils/plugins/dummy/__init__.py @@ -0,0 +1 @@ +import dummy.dummy_model diff --git a/tests/utils/plugins/dummy/dummy_model.py b/tests/utils/plugins/dummy/dummy_model.py new file mode 100644 index 00000000..8a89c12e --- /dev/null +++ b/tests/utils/plugins/dummy/dummy_model.py @@ -0,0 +1,8 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from modelscope.models.base import Model +from modelscope.models.builder import MODELS + + +@MODELS.register_module(group_key='dummy-group', module_name='dummy-model') +class DummyModel(Model): + pass diff --git a/tests/utils/test_plugin.py b/tests/utils/test_plugin.py new file mode 100644 index 00000000..40d86f9d --- /dev/null +++ b/tests/utils/test_plugin.py @@ -0,0 +1,41 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +from modelscope.models.builder import MODELS +from modelscope.utils.plugins import (discover_plugins, import_all_plugins, + import_file_plugins, import_plugins, + pushd) + + +class PluginTest(unittest.TestCase): + + def setUp(self): + self.plugins_root = 'tests/utils/plugins/' + + def test_no_plugins(self): + available_plugins = set(discover_plugins()) + assert available_plugins == set() + + def test_file_plugins(self): + with pushd(self.plugins_root): + available_plugins = set(discover_plugins()) + assert available_plugins == {'dummy'} + + import_file_plugins() + assert MODELS.get('dummy-model', 'dummy-group') is not None + + def test_custom_plugins(self): + with pushd(self.plugins_root): + available_plugins = set(discover_plugins()) + assert available_plugins == {'dummy'} + + import_plugins(['dummy']) + assert MODELS.get('dummy-model', 'dummy-group') is not None + + def test_all_plugins(self): + with pushd(self.plugins_root): + available_plugins = set(discover_plugins()) + assert available_plugins == {'dummy'} + + import_all_plugins() + assert MODELS.get('dummy-model', 'dummy-group') is not None