From 27ebee0471269de4ab51fcb126a5605f5cffdbd4 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Sat, 5 Sep 2026 02:58:34 +0800 Subject: [PATCH] feat(registry): discover third-party task packages via the unilab.tasks entry-point group ensure_registries() now also collects bootstrap packages declared through the 'unilab.tasks' entry-point group, in addition to the built-in unilab.tasks package and the UNILAB_EXTRA_REGISTRY_PACKAGES env var (kept as the test-fixture/ad-hoc seam). Entry-point metadata lives in site-packages, so spawn-based collector subprocesses discover the same packages without env-var forwarding; entry-point import failures fail closed, unlike env-var packages. Closes the gap tracked in #1500 (roadmap #1252 item 4). --- .../1-architecture/5-registry.md | 19 ++++- .../3-extending/1-new_task.md | 3 + .../1-architecture/5-registry.md | 15 +++- .../3-extending/1-new_task.md | 3 + src/unilab/base/registry.py | 17 +++++ tests/base/test_registry.py | 70 +++++++++++++++++++ 6 files changed, 122 insertions(+), 5 deletions(-) diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md index 3bf01b4eb..643781a7c 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md @@ -8,9 +8,13 @@ defined by {doc}`/adr/ADR-0004-registry-bootstrap-contract` and implemented in 1. Training entrypoints call `unilab.training.common.ensure_registries()`. 2. That helper delegates to `unilab.base.registry.ensure_registries()`. -3. The registry imports its sole declared bootstrap package, `unilab.tasks`. -4. `unilab.tasks` exposes `__unilab_registry_modules__`, an explicit tuple of - task leaf modules that contain registration side effects. +3. The registry collects bootstrap packages from three sources: the built-in + `unilab.tasks`, third-party packages declaring an entry point under + `[project.entry-points."unilab.tasks"]` (the value is an importable package + name), and the `UNILAB_EXTRA_REGISTRY_PACKAGES` environment variable + (mainly for test fixtures). +4. Each bootstrap package exposes `__unilab_registry_modules__`, an explicit + tuple of task leaf modules that contain registration side effects. 5. Imported modules register configs with `@registry.envcfg(...)` and env implementations with `@registry.env(..., sim_backend=...)` or `registry.register_env(...)`. @@ -22,6 +26,15 @@ defined by {doc}`/adr/ADR-0004-registry-bootstrap-contract` and implemented in - Add new task leaves to `unilab.tasks.__unilab_registry_modules__` when they are not imported by an existing bootstrap entry. +- Third-party task packages living outside this repo self-register by declaring + `[project.entry-points."unilab.tasks"]` in their own `pyproject.toml` (e.g. + `microduck = "microduck_rl_unilab.tasks"`); the declared package exposes the + same `__unilab_registry_modules__` tuple. Entry-point metadata lives in + site-packages, so spawn-based collector subprocesses discover the same + packages without any env-var forwarding; import failures of entry-point + packages fail closed (installing the package is a deliberate act). The + `UNILAB_EXTRA_REGISTRY_PACKAGES` env var remains available for test fixtures + and ad-hoc debugging. - Keep registration cheap. Scene materialization, XML processing, asset access, and backend construction belong after `registry.make(...)`, not in decorator registration. diff --git a/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md b/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md index 3c597ba1d..5cfe95525 100644 --- a/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md +++ b/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md @@ -14,6 +14,9 @@ Start from the contracts: {doc}`../2-contracts/1-env_contract`, `@registry.env("EnvName", sim_backend="motrix")`. 4. If the task lives in a new module, add that module to the package `__unilab_registry_modules__` tuple so `ensure_registries()` imports it. + Third-party packages outside this repo additionally declare + `[project.entry-points."unilab.tasks"]` in their `pyproject.toml` (see + {doc}`../1-architecture/5-registry`). 5. Keep `obs_groups_spec` accurate. It must include `obs` and may include `critic`; wrappers and learners trust these dimensions. 6. Keep reset and step semantics at the env owner layer: diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md index c96662e06..518ec6ce2 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md @@ -8,8 +8,11 @@ Registry bootstrap 是一个针对环境的显式导入契约。它由 1. 训练入口调用 `unilab.training.common.ensure_registries()`。 2. 该 helper 委托给 `unilab.base.registry.ensure_registries()`。 -3. registry 导入唯一声明的 bootstrap 包 `unilab.tasks`。 -4. `unilab.tasks` 暴露 `__unilab_registry_modules__`,即一个包含注册副作用的 +3. registry 依次从三个来源收集 bootstrap 包:内置的 `unilab.tasks`、第三方包 + 通过 `[project.entry-points."unilab.tasks"]` 声明的 entry point(值为可导入 + 包名)、以及 `UNILAB_EXTRA_REGISTRY_PACKAGES` 环境变量(主要供测试 fixture + 使用)。 +4. 每个 bootstrap 包暴露 `__unilab_registry_modules__`,即一个包含注册副作用的 task leaf module 显式元组。 5. 被导入的模块通过 `@registry.envcfg(...)` 注册 config,并通过 `@registry.env(..., sim_backend=...)` 或 `registry.register_env(...)` 注册 @@ -21,6 +24,14 @@ Registry bootstrap 是一个针对环境的显式导入契约。它由 - 如果新的 task leaf 尚未被现有 bootstrap 条目导入,需将其加入 `unilab.tasks.__unilab_registry_modules__`。 +- 仓库外的第三方任务包在自己的 `pyproject.toml` 中声明 + `[project.entry-points."unilab.tasks"]`(例如 + `microduck = "microduck_rl_unilab.tasks"`)完成自注册,被声明的包同样暴露 + `__unilab_registry_modules__`。entry-point 元数据位于 site-packages 中, + spawn 出的 collector 子进程无需转发环境变量即可发现同样的包;entry-point + 包的导入失败按 fail-closed 处理(安装是显式行为)。 + `UNILAB_EXTRA_REGISTRY_PACKAGES` 环境变量仍保留,供测试 fixture 与临时 + 调试使用。 - 保持注册过程轻量。场景 materialization、XML 处理、资源访问以及 backend 构造 应放在 `registry.make(...)` 之后,而不是放在装饰器注册中。 - 重复的 env config 以及重复的 `(env, sim_backend)` 注册会在 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md index 7e43d4c7c..9b61dc5c7 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md @@ -14,6 +14,9 @@ 后端实现。 4. 如果任务位于一个新模块中,请把该模块加入包的 `__unilab_registry_modules__` 元组,以便 `ensure_registries()` 导入它。 + 仓库外的第三方任务包还需在 `pyproject.toml` 声明 + `[project.entry-points."unilab.tasks"]`(见 + {doc}`../1-architecture/5-registry`)。 5. 保持 `obs_groups_spec` 准确。它必须包含 `obs`,并且可以包含 `critic`;wrapper 和 learner 都信任这些维度。 6. 把 reset 与 step 语义保留在 env owner 层: diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index 1943c4927..663a806a2 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -4,6 +4,7 @@ import os from collections.abc import Sequence from dataclasses import dataclass, field +from importlib.metadata import entry_points from typing import ( Any, Callable, @@ -47,6 +48,12 @@ def __call__( # Mainly intended for test setups that need to ship a fixture-only registry into # spawn subprocesses (which do not inherit pytest conftest state). _EXTRA_REGISTRY_PACKAGES_ENV = "UNILAB_EXTRA_REGISTRY_PACKAGES" +# Entry-point group through which third-party task packages self-register. +# A package declares e.g. ``[project.entry-points."unilab.tasks"]`` with +# ``microduck = "microduck_rl_unilab.tasks"``; the value is the importable +# package name, consumed exactly like a default/env-var package (it must +# declare ``__unilab_registry_modules__``). +_REGISTRY_ENTRY_POINT_GROUP = "unilab.tasks" logger = logging.getLogger(__name__) @@ -322,6 +329,16 @@ def ensure_registries( # have the env var leaked from a parent shell. optional.add(extra) + # Third-party task packages discovered through the "unilab.tasks" + # entry-point group. Entry-point metadata lives in site-packages, so + # spawn-based collector subprocesses (fresh interpreters re-running + # ensure_registries) discover the same packages without any env-var + # forwarding. Unlike env-var packages, an installed entry point is a + # deliberate installation choice, so import failures stay strict. + for ep in entry_points(group=_REGISTRY_ENTRY_POINT_GROUP): + if ep.value and ep.value not in package_names: + package_names.append(ep.value) + for package_name in package_names: is_optional = package_name in optional try: diff --git a/tests/base/test_registry.py b/tests/base/test_registry.py index bb0250bca..1b969b8f4 100644 --- a/tests/base/test_registry.py +++ b/tests/base/test_registry.py @@ -2,7 +2,9 @@ from __future__ import annotations +import importlib from dataclasses import dataclass +from types import SimpleNamespace import gymnasium as gym import numpy as np @@ -368,3 +370,71 @@ def test_make_with_invalid_cfg_override_raises(): """make() with a config key that doesn't exist raises ValueError.""" with pytest.raises(ValueError, match="has no attribute"): registry_mod.make(_TEST_ENV_A, sim_backend="mujoco", env_cfg_override={"__bogus_key__": 1}) + + +# --------------------------------------------------------------------------- +# Third-party "unilab.tasks" entry-point discovery (issue #1500) +# --------------------------------------------------------------------------- + + +def _track_imports(monkeypatch): + """Patch importlib.import_module to record every imported module name.""" + imported: list[str] = [] + real_import = importlib.import_module + + def tracking_import(name, package=None): + imported.append(name) + return real_import(name, package) + + monkeypatch.setattr("unilab.base.registry.importlib.import_module", tracking_import) + return imported + + +def _fake_entry_points(*values: str): + eps = [SimpleNamespace(value=v) for v in values] + + def fake_entry_points(*, group: str): + assert group == "unilab.tasks" + return eps + + return fake_entry_points + + +def test_ensure_registries_discovers_entry_point_packages(monkeypatch): + """Packages declared via the "unilab.tasks" entry-point group are imported + together with their declared ``__unilab_registry_modules__`` leaf modules.""" + imported = _track_imports(monkeypatch) + monkeypatch.setattr(registry_mod, "entry_points", _fake_entry_points("tests._test_registry")) + # Scrub the env var so the fixture package can only arrive via the + # entry-point seam under test. + monkeypatch.delenv("UNILAB_EXTRA_REGISTRY_PACKAGES", raising=False) + + registry_mod.ensure_registries(packages=[]) + + assert "tests._test_registry" in imported + assert "tests._test_registry.dummy_flat_env" in imported + + +def test_ensure_registries_entry_point_dedupes_default_packages(monkeypatch): + """An entry point naming an already-listed package must not be imported twice.""" + imported = _track_imports(monkeypatch) + monkeypatch.setattr(registry_mod, "entry_points", _fake_entry_points("unilab.tasks")) + monkeypatch.delenv("UNILAB_EXTRA_REGISTRY_PACKAGES", raising=False) + + registry_mod.ensure_registries() + + assert imported.count("unilab.tasks") == 1 + + +def test_ensure_registries_entry_point_import_failure_is_strict(monkeypatch): + """Unlike env-var packages, an installed entry point is deliberate: a broken + import must fail closed instead of being downgraded to a warning.""" + monkeypatch.setattr( + registry_mod, + "entry_points", + _fake_entry_points("definitely_not_a_real_pkg_xyz"), + ) + monkeypatch.delenv("UNILAB_EXTRA_REGISTRY_PACKAGES", raising=False) + + with pytest.raises(ImportError, match="definitely_not_a_real_pkg_xyz"): + registry_mod.ensure_registries(packages=[])