feat: No longer need to re-install plugins after ./pants export :: (#…

…439)

* Remove excessive debug logs for scanning entrypoints
* test: Add health check to the Redis container to prevent races in CI

changes/439.misc.md

@@ -0,0 +1 @@
+Improve auto-detection of plugins in development setups so that we no longer need to reinstall them after running `./pants export`

docs/dev/daily-workflows.rst

@@ -204,25 +204,22 @@ Working with plugins
 
 To develop Backend.AI plugins together, the repository offers a special location
 ``./plugins`` where you can clone plugin repositories and a shortcut script
-``scripts/install-plugin.sh`` to install them inside the exported venv.
+``scripts/install-plugin.sh`` that does this for you.
 
 .. code-block:: console
 
     $ scripts/install-plugin.sh lablup/backend.ai-accelerator-cuda-mock
 
-Manual installation:
+This is equivalent to:
 
 .. code-block:: console
 
-    $ git clone https://github.com/lablup/backend.ai-accelerator-cuda-mock plugins/cuda-mock
-    $ ./py -m pip install -e plugins/cuda-mock
-
-.. warning::
-
-   Whenever you re-export the venv using ``./pants export ::``, you must
-   *reinstall* the plugins.  There is also a shrotcut script that does this
-   for you: ``scripts/reinstall-plugins.sh``.
+    $ git clone \
+    >   https://github.com/lablup/backend.ai-accelerator-cuda-mock \
+    >   plugins/backend.ai-accelerator-cuda-mock
 
+These plugins are auto-detected by scanning ``setup.cfg`` of plugin subdirectories
+by the ``ai.backend.plugin.entrypoint`` module, even without explicit editable installations.
 
 Writing test cases
 ------------------

py

@@ -11,4 +11,8 @@ if [ $? -ne 0 ]; then
 fi
 LOCKSET=${LOCKSET:-python-default/$PYTHON_VERSION}
 source dist/export/python/virtualenvs/$LOCKSET/bin/activate
-PYTHONPATH=src python "$@"
+PYTHONPATH="${PYTHONPATH}"
+for plugin_dir in $(ls -d plugins/*/); do
+  PYTHONPATH="${plugin_dir}/src:${PYTHONPATH}"
+done
+PYTHONPATH="src:${PYTHONPATH}" python "$@"

scripts/install-plugin.sh

@@ -5,4 +5,3 @@ PLUGIN_BRANCH=${PLUGIN_BRANCH:-main}
 PLUGIN_OWNER=$(echo $1 | cut -d / -f 1)
 PLUGIN_REPO=$(echo $1 | cut -d / -f 2)
 git clone "https://github.com/$1" "./plugins/$PLUGIN_REPO"
-./py -m pip install -e "./plugins/$PLUGIN_REPO"

src/ai/backend/common/plugin/__init__.py

@@ -3,7 +3,6 @@
 from abc import ABCMeta, abstractmethod
 import asyncio
 import logging
-import pkg_resources
 import re
 from typing import (
     Any,
@@ -20,6 +19,7 @@
 from weakref import WeakSet
 
 from ai.backend.common.asyncio import cancel_tasks
+from ai.backend.plugin.entrypoint import scan_entrypoints
 
 from ..etcd import AsyncEtcd
 from ..logging_utils import BraceStyleAdapter
@@ -130,7 +130,7 @@ def discover_plugins(
     ) -> Iterator[Tuple[str, Type[P]]]:
         if blocklist is None:
             blocklist = set()
-        for entrypoint in pkg_resources.iter_entry_points(plugin_group):
+        for entrypoint in scan_entrypoints(plugin_group):
             if entrypoint.name in blocklist:
                 continue
             log.info('loading plugin (group:{}): {}', plugin_group, entrypoint.name)

src/ai/backend/plugin/entrypoint.py

@@ -1,4 +1,5 @@
 import ast
+import configparser
 import itertools
 import logging
 from importlib.metadata import EntryPoint, entry_points
@@ -17,15 +18,27 @@ def scan_entrypoints(
     existing_names: dict[str, EntryPoint] = {}
     for entrypoint in itertools.chain(
         scan_entrypoint_from_buildscript(group_name),
+        scan_entrypoint_from_plugin_checkouts(group_name),
         scan_entrypoint_from_package_metadata(group_name),
     ):
         if entrypoint.name in blocklist:
             continue
-        if existing_entrypoint := existing_names.get(entrypoint.name):
-            raise RuntimeError(
-                f"Detected a duplicate plugin entrypoint name {entrypoint.name!r} "
-                f"from {existing_entrypoint.value} and {entrypoint.value}",
-            )
+        if existing_entrypoint := existing_names.get(entrypoint.name, None):
+            if existing_entrypoint.value == entrypoint.value:
+                # Allow if the same plugin is scanned multiple times.
+                # This may happen if:
+                # - A plugin is installed via `./py -m pip install -e ...`
+                # - The unified venv is re-exported *without* `./py -m pip uninstall ...`.
+                # - Adding PYTHONPATH with plugin src directories in `./py` results in
+                #   *duplicate* scan results from the remaining `.egg-info` directory (pkg metadata)
+                #   and the `setup.cfg` scan results.
+                # TODO: compare the plugin versions as well? (need to remember version with entrypoints)
+                continue
+            else:
+                raise RuntimeError(
+                    f"Detected a duplicate plugin entrypoint name {entrypoint.name!r} "
+                    f"from {existing_entrypoint.value} and {entrypoint.value}",
+                )
         existing_names[entrypoint.name] = entrypoint
         yield entrypoint
 
@@ -40,7 +53,6 @@ def scan_entrypoint_from_buildscript(group_name: str) -> Iterator[EntryPoint]:
     ai_backend_ns_path = Path(__file__).parent.parent
     log.debug("scan_entrypoint_from_buildscript(%r): Namespace path: %s", group_name, ai_backend_ns_path)
     for buildscript_path in ai_backend_ns_path.glob("**/BUILD"):
-        log.debug("reading entry points [%s] from %s", group_name, buildscript_path)
         for entrypoint in extract_entrypoints_from_buildscript(group_name, buildscript_path):
             entrypoints[entrypoint.name] = entrypoint
     # Override with the entrypoints found in the current source directories,
@@ -50,18 +62,35 @@ def scan_entrypoint_from_buildscript(group_name: str) -> Iterator[EntryPoint]:
         pass
     else:
         src_path = build_root / 'src'
-        plugins_path = build_root / 'plugins'
         log.debug("scan_entrypoint_from_buildscript(%r): current src: %s", group_name, src_path)
         for buildscript_path in src_path.glob("**/BUILD"):
-            if buildscript_path.is_relative_to(plugins_path):
-                # Prevent loading BUILD files in plugin checkouts if they use Pants on their own.
-                continue
-            log.debug("reading entry points [%s] from %s", group_name, buildscript_path)
             for entrypoint in extract_entrypoints_from_buildscript(group_name, buildscript_path):
                 entrypoints[entrypoint.name] = entrypoint
     yield from entrypoints.values()
 
 
+def scan_entrypoint_from_plugin_checkouts(group_name: str) -> Iterator[EntryPoint]:
+    entrypoints = {}
+    try:
+        build_root = find_build_root()
+    except ValueError:
+        pass
+    else:
+        plugins_path = build_root / 'plugins'
+        log.debug("scan_entrypoint_from_plugin_checkouts(%r): plugin parent dir: %s", group_name, plugins_path)
+        # For cases when plugins use Pants
+        for buildscript_path in plugins_path.glob("**/BUILD"):
+            for entrypoint in extract_entrypoints_from_buildscript(group_name, buildscript_path):
+                entrypoints[entrypoint.name] = entrypoint
+        # For cases when plugins use standard setup.cfg
+        for setup_cfg_path in plugins_path.glob("**/setup.cfg"):
+            for entrypoint in extract_entrypoints_from_setup_cfg(group_name, setup_cfg_path):
+                if entrypoint.name not in entrypoints:
+                    entrypoints[entrypoint.name] = entrypoint
+        # TODO: implement pyproject.toml scanner
+    yield from entrypoints.values()
+
+
 def find_build_root(path: Optional[Path] = None) -> Path:
     cwd = Path.cwd() if path is None else path
     while True:
@@ -97,3 +126,24 @@ def extract_entrypoints_from_buildscript(
                                 yield EntryPoint(name=name, value=ref, group=group_name)
                             except ValueError:
                                 pass
+
+
+def extract_entrypoints_from_setup_cfg(
+    group_name: str,
+    setup_cfg_path: Path,
+) -> Iterator[EntryPoint]:
+    cfg = configparser.ConfigParser()
+    cfg.read(setup_cfg_path)
+    raw_data = cfg.get('options.entry_points', group_name, fallback="").strip()
+    if not raw_data:
+        return
+    data = {
+        k.strip(): v.strip()
+        for k, v in (
+            line.split("=", maxsplit=1)
+            for line in
+            raw_data.splitlines()
+        )
+    }
+    for name, ref in data.items():
+        yield EntryPoint(name=name, value=ref, group=group_name)

src/ai/backend/testutils/bootstrap.py

@@ -12,6 +12,21 @@
 from ai.backend.common.types import HostPortPair
 
 
+def wait_health_check(container_id):
+    while True:
+        proc = subprocess.run(
+            [
+                'docker', 'inspect', container_id,
+            ],
+            capture_output=True,
+        )
+        container_info = json.loads(proc.stdout)
+        if container_info[0]['State']['Health']['Status'].lower() != 'healthy':
+            time.sleep(0.2)
+            continue
+        return container_info
+
+
 @pytest.fixture(scope='session')
 def etcd_container() -> Iterator[tuple[str, HostPortPair]]:
     # Spawn a single-node etcd container for a testing session.
@@ -52,6 +67,9 @@ def redis_container() -> Iterator[tuple[str, HostPortPair]]:
         [
             'docker', 'run', '-d',
             '-p', ':6379',
+            '--health-cmd', 'redis-cli ping',
+            '--health-interval', '0.3s',
+            '--health-start-period', '0.2s',
             'redis:6.2-alpine',
         ],
         capture_output=True,
@@ -64,6 +82,7 @@ def redis_container() -> Iterator[tuple[str, HostPortPair]]:
         capture_output=True,
     )
     container_info = json.loads(proc.stdout)
+    container_info = wait_health_check(container_id)
     host_port = int(container_info[0]['NetworkSettings']['Ports']['6379/tcp'][0]['HostPort'])
     yield container_id, HostPortPair('127.0.0.1', host_port)
     subprocess.run(
@@ -91,19 +110,8 @@ def postgres_container() -> Iterator[tuple[str, HostPortPair]]:
         capture_output=True,
     )
     container_id = proc.stdout.decode().strip()
-    host_port = 0
-    while host_port == 0:
-        proc = subprocess.run(
-            [
-                'docker', 'inspect', container_id,
-            ],
-            capture_output=True,
-        )
-        container_info = json.loads(proc.stdout)
-        if container_info[0]['State']['Health']['Status'].lower() != 'healthy':
-            time.sleep(0.2)
-            continue
-        host_port = int(container_info[0]['NetworkSettings']['Ports']['5432/tcp'][0]['HostPort'])
+    container_info = wait_health_check(container_id)
+    host_port = int(container_info[0]['NetworkSettings']['Ports']['5432/tcp'][0]['HostPort'])
     yield container_id, HostPortPair('127.0.0.1', host_port)
     subprocess.run(
         [

tests/common/test_plugin.py

@@ -1,8 +1,13 @@
 import asyncio
 import functools
+from dataclasses import dataclass
 from typing import (
     Any,
+    Callable,
+    Iterator,
     Mapping,
+    Type,
+    overload,
 )
 from unittest.mock import AsyncMock
 
@@ -39,27 +44,66 @@ async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None:
         pass
 
 
+@dataclass
 class DummyEntrypoint:
 
-    def __init__(self, name: str, load_result: Any) -> None:
-        self.name = name
-        self._load_result = load_result
+    name: str
+    load_result: Type[AbstractPlugin] | Callable[..., AbstractPlugin]
+    value: str = "dummy.mod.DummyPlugin"
+    group: str = "dummy_group_v00"
 
-    def load(self) -> Any:
-        return self._load_result
+    def load(self):
+        return self.load_result
 
 
 def mock_entrypoints_with_instance(plugin_group_name: str, *, mocked_plugin):
     # Since mocked_plugin is already an instance constructed via AsyncMock,
     # we emulate the original constructor using a lambda fucntion.
-    yield DummyEntrypoint('dummy', lambda plugin_config, local_config: mocked_plugin)
+    yield DummyEntrypoint(
+        name='dummy',
+        group=plugin_group_name,
+        load_result=lambda plugin_config, local_config: mocked_plugin,
+    )
+
+
+@overload
+def mock_entrypoints_with_class(
+    plugin_group_name: str,
+    *,
+    plugin_cls: list[Type[AbstractPlugin]],
+) -> Iterator[DummyEntrypoint]:
+    ...
+
+
+@overload
+def mock_entrypoints_with_class(
+    plugin_group_name: str,
+    *,
+    plugin_cls: Type[AbstractPlugin],
+) -> DummyEntrypoint:
+    ...
 
 
-def mock_entrypoints_with_class(plugin_group_name: str, *, plugin_cls):
+def mock_entrypoints_with_class(
+    plugin_group_name: str,
+    *,
+    plugin_cls,
+):
     if isinstance(plugin_cls, list):
-        yield from (DummyEntrypoint(getattr(p, '_entrypoint_name', 'dummy'), p) for p in plugin_cls)
+        yield from (
+            DummyEntrypoint(
+                getattr(p, '_entrypoint_name', 'dummy'),
+                group=plugin_group_name,
+                load_result=p,
+            )
+            for p in plugin_cls
+        )
     else:
-        yield DummyEntrypoint('dummy', plugin_cls)
+        yield DummyEntrypoint(
+            'dummy',
+            group=plugin_group_name,
+            load_result=plugin_cls,
+        )
 
 
 @pytest.mark.asyncio
@@ -68,7 +112,7 @@ async def test_plugin_context_init_cleanup(etcd, mocker):
     mocked_plugin = AsyncMock(DummyPlugin)
     mocked_entrypoints = functools.partial(mock_entrypoints_with_instance,
                                            mocked_plugin=mocked_plugin)
-    mocker.patch('ai.backend.common.plugin.pkg_resources.iter_entry_points', mocked_entrypoints)
+    mocker.patch('ai.backend.common.plugin.scan_entrypoints', mocked_entrypoints)
     ctx = BasePluginContext(etcd, {})
     try:
         assert not ctx.plugins
@@ -83,7 +127,7 @@ async def test_plugin_context_init_cleanup(etcd, mocker):
 @pytest.mark.asyncio
 async def test_plugin_context_config(etcd, mocker):
     mocked_entrypoints = functools.partial(mock_entrypoints_with_class, plugin_cls=DummyPlugin)
-    mocker.patch('ai.backend.common.plugin.pkg_resources.iter_entry_points', mocked_entrypoints)
+    mocker.patch('ai.backend.common.plugin.scan_entrypoints', mocked_entrypoints)
     await etcd.put('config/plugins/XXX/dummy/etcd-key', 'etcd-value')
     ctx = BasePluginContext(
         etcd,
@@ -105,7 +149,7 @@ async def test_plugin_context_config_autoupdate(etcd, mocker):
     mocked_plugin = AsyncMock(DummyPlugin)
     mocked_entrypoints = functools.partial(mock_entrypoints_with_instance,
                                            mocked_plugin=mocked_plugin)
-    mocker.patch('ai.backend.common.plugin.pkg_resources.iter_entry_points', mocked_entrypoints)
+    mocker.patch('ai.backend.common.plugin.scan_entrypoints', mocked_entrypoints)
     await etcd.put_prefix('config/plugins/XXX/dummy', {'a': '1', 'b': '2'})
     ctx = BasePluginContext(
         etcd,
@@ -223,7 +267,7 @@ async def test_hook_dispatch(etcd, mocker):
         mock_entrypoints_with_class,
         plugin_cls=[DummyHookPassingPlugin, DummyHookRejectingPlugin, DummyHookErrorPlugin],
     )
-    mocker.patch('ai.backend.common.plugin.pkg_resources.iter_entry_points', mocked_entrypoints)
+    mocker.patch('ai.backend.common.plugin.scan_entrypoints', mocked_entrypoints)
     ctx = HookPluginContext(etcd, {})
     try:
         await ctx.init()
@@ -279,7 +323,7 @@ async def test_hook_notify(etcd, mocker):
         mock_entrypoints_with_class,
         plugin_cls=[DummyHookPassingPlugin, DummyHookRejectingPlugin, DummyHookErrorPlugin],
     )
-    mocker.patch('ai.backend.common.plugin.pkg_resources.iter_entry_points', mocked_entrypoints)
+    mocker.patch('ai.backend.common.plugin.scan_entrypoints', mocked_entrypoints)
     ctx = HookPluginContext(etcd, {})
     try:
         await ctx.init()

tests/manager/test_registry.py

@@ -30,7 +30,7 @@ async def test_handle_heartbeat(
     def mocked_entrypoints(entry_point_group: str):
         return []
 
-    mocker.patch('ai.backend.common.plugin.pkg_resources.iter_entry_points', mocked_entrypoints)
+    mocker.patch('ai.backend.common.plugin.scan_entrypoints', mocked_entrypoints)
 
     registry, mock_dbconn, mock_dbresult, mock_shared_config, _, _ = registry_ctx
     image_data = snappy.compress(msgpack.packb([