Feat/plugin update (#241)

* reload pkg resources after installed

* add version test in order to install correct version

* install only uninstalled package from requirements

* support get module name from distribution

* bug fixed

---------

Co-authored-by: Zhicheng Zhang <zhangzhicheng.zzc@alibaba-inc.com>

modelscope/utils/plugins.py

@@ -1,9 +1,12 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 # This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
+# Part of the implementation is borrowed from wimglenn/johnnydep
+
 import copy
 import importlib
 import os
 import pkgutil
+import shutil
 import sys
 import venv
 from contextlib import contextmanager
@@ -263,14 +266,22 @@ def install_module_from_requirements(requirement_path, ):
 
     """
 
-    install_args = ['-r', requirement_path]
-    status_code, _, args = PluginsManager.pip_command(
-        'install',
-        install_args,
-    )
-    if status_code != 0:
-        raise ImportError(
-            f'Failed to install requirements from {requirement_path}')
+    install_list = []
+    with open(requirement_path, 'r', encoding='utf-8') as f:
+        requirements = f.read().splitlines()
+        for req in requirements:
+            installed, _ = PluginsManager.check_plugin_installed(req)
+            if not installed:
+                install_list.append(req)
+
+    if len(install_list) > 0:
+        status_code, _, args = PluginsManager.pip_command(
+            'install',
+            install_list,
+        )
+        if status_code != 0:
+            raise ImportError(
+                f'Failed to install requirements from {requirement_path}')
 
 
 def import_module_from_file(module_name, file_path):
@@ -298,18 +309,6 @@ def import_module_from_model_dir(model_dir):
         import_module_from_file(module_name, file)
 
 
-def install_modelscope_if_need():
-    plugin_installed, version = PluginsManager.check_plugin_installed(
-        'modelscope')
-    if not plugin_installed:
-        status_code, _, args = PluginsManager.pip_command(
-            'install',
-            ['modelscope'],
-        )
-        if status_code != 0:
-            raise ImportError('Failed to install package modelscope')
-
-
 def install_requirements_by_names(plugins: List[str]):
     plugins_manager = PluginsManager()
     uninstalled_plugins = []
@@ -324,20 +323,21 @@ def install_requirements_by_names(plugins: List[str]):
             f'The required packages {",".join(uninstalled_plugins)} are not installed.',
             f'Please run the command `modelscope plugin install {" ".join(uninstalled_plugins)}` to install them.'
         )
-    install_modelscope_if_need()
 
 
 def install_requirements_by_files(requirements: List[str]):
     for requirement in requirements:
         install_module_from_requirements(requirement)
-    install_modelscope_if_need()
 
 
 def register_plugins_repo(plugins: List[str]) -> None:
     """ Try to install and import plugins from repo"""
     if plugins is not None:
         install_requirements_by_names(plugins)
-        import_plugins(plugins)
+        modules = []
+        for plugin in plugins:
+            modules.extend(get_modules_from_package(plugin))
+        import_plugins(modules)
 
 
 def register_modelhub_repo(model_dir, allow_remote=False) -> None:
@@ -351,6 +351,256 @@ def register_modelhub_repo(model_dir, allow_remote=False) -> None:
             pass
 
 
+DEFAULT_INDEX = 'https://pypi.org/simple/'
+
+
+def get_modules_from_package(package):
+    """ to get the modules from a installed package
+
+    Args:
+        package: The distribution name or package name
+
+    Returns:
+
+    """
+    from zipfile import ZipFile
+    from tempfile import mkdtemp
+    from subprocess import check_output, STDOUT
+    from glob import glob
+    import hashlib
+    from urllib.parse import urlparse
+    from urllib import request as urllib2
+    from pip._internal.utils.packaging import get_requirement
+    req = get_requirement(package)
+    package = req.name
+
+    def urlretrieve(url, filename, data=None, auth=None):
+        if auth is not None:
+            # https://docs.python.org/2.7/howto/urllib2.html#id6
+            password_mgr = urllib2.HTTPPasswordMgrWithDefaultRealm()
+
+            # Add the username and password.
+            # If we knew the realm, we could use it instead of None.
+            username, password = auth
+            top_level_url = urlparse(url).netloc
+            password_mgr.add_password(None, top_level_url, username, password)
+
+            handler = urllib2.HTTPBasicAuthHandler(password_mgr)
+
+            # create "opener" (OpenerDirector instance)
+            opener = urllib2.build_opener(handler)
+        else:
+            opener = urllib2.build_opener()
+
+        res = opener.open(url, data=data)
+
+        headers = res.info()
+
+        with open(filename, 'wb') as fp:
+            fp.write(res.read())
+
+        return filename, headers
+
+    def compute_checksum(target, algorithm='sha256', blocksize=2**13):
+        hashtype = getattr(hashlib, algorithm)
+        hash_ = hashtype()
+        logger.debug('computing checksum', target=target, algorithm=algorithm)
+        with open(target, 'rb') as f:
+            for chunk in iter(lambda: f.read(blocksize), b''):
+                hash_.update(chunk)
+        result = hash_.hexdigest()
+        logger.debug('computed checksum', result=result)
+        return result
+
+    def _get_pip_version():
+        # try to get pip version without actually importing pip
+        # setuptools gets upset if you import pip before importing setuptools..
+        try:
+            import importlib.metadata  # Python 3.8+
+            return importlib.metadata.version('pip')
+        except Exception:
+            pass
+        import pip
+        return pip.__version__
+
+    def _download_dist(url, scratch_file, index_url, extra_index_url):
+        auth = None
+        if index_url:
+            parsed = urlparse(index_url)
+            if parsed.username and parsed.password and parsed.hostname == urlparse(
+                    url).hostname:
+                # handling private PyPI credentials in index_url
+                auth = (parsed.username, parsed.password)
+        if extra_index_url:
+            parsed = urlparse(extra_index_url)
+            if parsed.username and parsed.password and parsed.hostname == urlparse(
+                    url).hostname:
+                # handling private PyPI credentials in extra_index_url
+                auth = (parsed.username, parsed.password)
+        target, _headers = urlretrieve(url, scratch_file, auth=auth)
+        return target, _headers
+
+    def _get_wheel_args(index_url, env, extra_index_url):
+        args = [
+            sys.executable,
+            '-m',
+            'pip',
+            'wheel',
+            '-vvv',  # --verbose x3
+            '--no-deps',
+            '--no-cache-dir',
+            '--disable-pip-version-check',
+        ]
+        if index_url is not None:
+            args += ['--index-url', index_url]
+            if index_url != DEFAULT_INDEX:
+                hostname = urlparse(index_url).hostname
+                if hostname:
+                    args += ['--trusted-host', hostname]
+        if extra_index_url is not None:
+            args += [
+                '--extra-index-url', extra_index_url, '--trusted-host',
+                urlparse(extra_index_url).hostname
+            ]
+        if env is None:
+            pip_version = _get_pip_version()
+        else:
+            pip_version = dict(env)['pip_version']
+            args[0] = dict(env)['python_executable']
+        pip_major, pip_minor = pip_version.split('.')[0:2]
+        pip_major = int(pip_major)
+        pip_minor = int(pip_minor)
+        if pip_major >= 10:
+            args.append('--progress-bar=off')
+        if (20, 3) <= (pip_major, pip_minor) < (21, 1):
+            # See https://github.com/pypa/pip/issues/9139#issuecomment-735443177
+            args.append('--use-deprecated=legacy-resolver')
+        return args
+
+    def get(dist_name,
+            index_url=None,
+            env=None,
+            extra_index_url=None,
+            tmpdir=None,
+            ignore_errors=False):
+        args = _get_wheel_args(index_url, env, extra_index_url) + [dist_name]
+        scratch_dir = mkdtemp(dir=tmpdir)
+        logger.debug(
+            'wheeling and dealing',
+            scratch_dir=os.path.abspath(scratch_dir),
+            args=' '.join(args))
+        try:
+            out = check_output(
+                args, stderr=STDOUT, cwd=scratch_dir).decode('utf-8')
+        except ChildProcessError as err:
+            out = getattr(err, 'output', b'').decode('utf-8')
+            logger.warning(out)
+            if not ignore_errors:
+                raise
+        logger.debug('wheel command completed ok', dist_name=dist_name)
+        links = []
+        local_links = []
+        lines = out.splitlines()
+        for i, line in enumerate(lines):
+            line = line.strip()
+            if line.startswith('Downloading from URL '):
+                parts = line.split()
+                link = parts[3]
+                links.append(link)
+            elif line.startswith('Downloading '):
+                parts = line.split()
+                last = parts[-1]
+                if len(parts) == 3 and last.startswith('(') and last.endswith(
+                        ')'):
+                    link = parts[-2]
+                elif len(parts) == 4 and parts[-2].startswith(
+                        '(') and last.endswith(')'):
+                    link = parts[-3]
+                    if not urlparse(link).scheme:
+                        # newest pip versions have changed to not log the full url
+                        # in the download event. it is becoming more and more annoying
+                        # to preserve compatibility across a wide range of pip versions
+                        next_line = lines[i + 1].strip()
+                        if next_line.startswith(
+                                'Added ') and ' to build tracker' in next_line:
+                            link = next_line.split(
+                                ' to build tracker')[0].split()[-1]
+                else:
+                    link = last
+                links.append(link)
+            elif line.startswith(
+                    'Source in ') and 'which satisfies requirement' in line:
+                link = line.split()[-1]
+                links.append(link)
+            elif line.startswith('Added ') and ' from file://' in line:
+                [link] = [x for x in line.split() if x.startswith('file://')]
+                local_links.append(link)
+        if not links:
+            # prefer http scheme over file
+            links += local_links
+        links = list(dict.fromkeys(links))  # order-preserving dedupe
+        if not links:
+            logger.warning('could not find download link', out=out)
+            raise Exception('failed to collect dist')
+        if len(links) == 2:
+            # sometimes we collect the same link, once with a url fragment/checksum and once without
+            first, second = links
+            if first.startswith(second):
+                del links[1]
+            elif second.startswith(first):
+                del links[0]
+        if len(links) > 1:
+            logger.debug('more than 1 link collected', out=out, links=links)
+            # Since PEP 517, maybe an sdist will also need to collect other distributions
+            # for the build system, even with --no-deps specified. pendulum==1.4.4 is one
+            # example, which uses poetry and doesn't publish any python37 wheel to PyPI.
+            # However, the dist itself should still be the first one downloaded.
+        link = links[0]
+        whls = glob(os.path.join(os.path.abspath(scratch_dir), '*.whl'))
+        try:
+            [whl] = whls
+        except ValueError:
+            if ignore_errors:
+                whl = ''
+            else:
+                raise
+        url, _sep, checksum = link.partition('#')
+        url = url.replace(
+            '/%2Bf/', '/+f/'
+        )  # some versions of pip did not unquote this fragment in the log
+        if not checksum.startswith('md5=') and not checksum.startswith(
+                'sha256='):
+            # PyPI gives you the checksum in url fragment, as a convenience. But not all indices are so kind.
+            algorithm = 'md5'
+            if os.path.basename(whl).lower() == url.rsplit('/', 1)[-1].lower():
+                target = whl
+            else:
+                scratch_file = os.path.join(scratch_dir, os.path.basename(url))
+                target, _headers = _download_dist(url, scratch_file, index_url,
+                                                  extra_index_url)
+            checksum = compute_checksum(target=target, algorithm=algorithm)
+            checksum = '='.join([algorithm, checksum])
+        result = {'path': whl, 'url': url, 'checksum': checksum}
+        return result
+
+    def discover_import_names(whl_file):
+        logger.debug('finding import names')
+        zipfile = ZipFile(file=whl_file)
+        namelist = zipfile.namelist()
+        [top_level_fname
+         ] = [x for x in namelist if x.endswith('top_level.txt')]
+        all_names = zipfile.read(top_level_fname).decode(
+            'utf-8').strip().splitlines()
+        public_names = [n for n in all_names if not n.startswith('_')]
+        return public_names
+
+    tmpdir = mkdtemp()
+    data = get(package, tmpdir=tmpdir)
+    import_names = discover_import_names(data['path'])
+    shutil.rmtree(tmpdir)
+    return import_names
+
+
 class PluginsManager(object):
 
     def __init__(self,
@@ -370,11 +620,31 @@ class PluginsManager(object):
 
     @staticmethod
     def check_plugin_installed(package):
+        """ Check if the plugin is installed, and if the version is valid
+
+        Args:
+            package: the package name need to be installed
+
+        Returns:
+
+        """
+        from pip._internal.utils.packaging import get_requirement, specifiers
+        req = get_requirement(package)
+
         try:
             importlib.reload(pkg_resources)
-            package_meta_info = pkg_resources.working_set.by_key[package]
+            package_meta_info = pkg_resources.working_set.by_key[req.name]
             version = package_meta_info.version
+
+            # To test if the package is installed
             installed = True
+
+            # If installed, test if the version is correct
+            for spec in req.specifier:
+                installed_valid_version = spec.contains(version)
+                if not installed_valid_version:
+                    installed = False
+                    break
         except KeyError:
             version = ''
             installed = False
@@ -402,6 +672,10 @@ class PluginsManager(object):
         options, args = command.parse_args(command_args)
 
         status_code = command.main(command_args)
+
+        # reload the pkg_resources in order to get the latest pkgs information
+        importlib.reload(pkg_resources)
+
         return status_code, options, args
 
     def install_plugins(self,
@@ -722,3 +996,4 @@ class EnvsManager(object):
 
 if __name__ == '__main__':
     install_requirements_by_files(['adaseq'])
+    import_name = get_modules_from_package('pai-easycv')

tests/pipelines/plugin_remote_pipelines/configuration.json

@@ -0,0 +1 @@
+{"framework":"pytorch","task":"bilibili","model":{"type":"my-custom-model","scale":2,"weight_path":"weights_v3/up2x-latest-denoise3x.pth","half":true},"pipeline":{"type":"my-custom-pipeline"}}

tests/pipelines/plugin_remote_pipelines/test_plugin_model.py

@@ -15,8 +15,6 @@ class PluginModelTest(unittest.TestCase, DemoCompatibilityCheck):
 
     def tearDown(self):
         # make sure uninstalled after installing
-        uninstall_args = [self.package, '-y']
-        PluginsManager.pip_command('uninstall', uninstall_args)
         super().tearDown()
         import subprocess
         result = subprocess.run(