From e156dfb51ab001f941523c794ce182c5d2361c24 Mon Sep 17 00:00:00 2001 From: Haoran Geng <71596067+geng-haoran@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:44:44 +0000 Subject: [PATCH] perf(task): lazy task registry backed by a static AST index MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `list_tasks()` / `get_task_class()` imported every module under every task package (`pkgutil.walk_packages` + `import_module`): with RoboVerse that is ~730 modules, two 20k-line generated files, and every import-time side effect (asset downloads, optional-dependency warnings) before the first lookup — 7.3 s, and a failure blast radius that grows with each task family. - `metasim/task/_static_index.py`: walks the package directories (no imports — `pkgutil.walk_packages` would import each sub-package `__init__`, some of which eagerly import their whole family), parses each file with `ast`, and records every `register_task(...)` call with string-literal arguments as name -> module. Calls with non-literal arguments (loops) flag the module as dynamic. Results are cached per file (mtime + size) in `$METASIM_CACHE_DIR/task_index.json`. - `registry.get_task_class` imports only the registering module; unknown names fall back to the old import-everything discovery so the KeyError still lists import failures. `registry.list_tasks` returns the index plus dynamic modules' names. The cached index is keyed by the package configuration and rebuilt when it changes (env vars, cwd); an emptied registry re-runs discovery. `METASIM_TASK_DISCOVERY=eager` restores the previous behaviour. Measured on RoboVerse main (6051 tasks): index cold 1.08 s / warm 0.04 s; `get_task_class` 0.01-0.1 s (mjlab 0.9 s = its own imports); `list_tasks` 4.9 s (3.5 s of it is importing `maniskill.native_tasks`, which registers 5360 names in a loop) versus 7.3 s before. All 2898 decorators in RoboVerse are string literals. Tests: `metasim/test/test_task_registry_lazy_general.py` (6 tests on a synthetic package: single module imported per lookup, list without imports, eager fallback + failure reporting, env var, cache reuse/invalidation). `pytest -k general` -> 492 passed; `-k mujoco` unchanged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017i6VtKoovBNed815mWFqxw --- CHANGELOG.md | 9 + metasim/task/_static_index.py | 199 ++++++++++++++++++ metasim/task/registry.py | 87 +++++++- .../test/test_task_registry_lazy_general.py | 172 +++++++++++++++ 4 files changed, 456 insertions(+), 11 deletions(-) create mode 100644 metasim/task/_static_index.py create mode 100644 metasim/test/test_task_registry_lazy_general.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b507a6..a4243ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- **Task registry is lazy.** `get_task_class(name)` resolves the name through a static AST index of + every `@register_task(...)` literal (`metasim/task/_static_index.py`, per-file cache under + `$METASIM_CACHE_DIR`, default `/metasim_cache`) and imports only the module that registers it; + `list_tasks()` no longer imports every task module (modules that register names in loops are + imported to learn them). RoboVerse: cold index 1.1 s, warm 0.04 s, a lookup ~0.1 s, versus 7.3 s + and every import-time side effect before. `METASIM_TASK_DISCOVERY=eager` restores the old behaviour. + +### Changed + - **Distribution renamed to `roboverse-metasim`** (`metasim` on PyPI is an unrelated project); the import name stays `metasim`. Downstream requirements must say `roboverse-metasim @ git+...`. `metasim.__version__` now reads the installed version from package metadata. diff --git a/metasim/task/_static_index.py b/metasim/task/_static_index.py new file mode 100644 index 0000000..d6d7e3f --- /dev/null +++ b/metasim/task/_static_index.py @@ -0,0 +1,199 @@ +"""Static task index: find ``@register_task("name")`` decorators without importing the modules. + +``list_tasks()`` and ``get_task_class()`` used to import *every* module of every task package +(``pkgutil.walk_packages`` + ``import_module``). With RoboVerse that is ~500 modules, two of them +20k-line generated files, plus import-time side effects (asset downloads, optional-dependency +warnings) — about 7 s before the first task can be looked up, and a failure blast radius that grows +with every task family. + +This module scans the same packages with :mod:`ast` instead: every ``register_task(...)`` decorator +whose arguments are string literals is recorded as ``name -> module``. Lookups then import only the +one module that registers the requested task. Decorators with non-literal arguments are reported so +the registry can fall back to the full import for those packages. +""" + +from __future__ import annotations + +import ast +import json +import os +import tempfile +from dataclasses import dataclass, field +from importlib.util import find_spec + + +@dataclass +class StaticIndex: + """Result of :func:`build_static_index`.""" + + names: dict[str, str] = field(default_factory=dict) + """lower-cased task name -> fully qualified module that registers it.""" + dynamic_modules: set[str] = field(default_factory=set) + """Modules with a ``register_task`` call whose arguments are not all string literals.""" + scanned_modules: int = 0 + parse_failures: dict[str, str] = field(default_factory=dict) + + +def _decorator_names(node: ast.AST) -> list[str] | None: + """String-literal arguments of a ``register_task(...)`` call, or None if not one / not literal.""" + if not isinstance(node, ast.Call): + return None + func = node.func + fname = func.id if isinstance(func, ast.Name) else func.attr if isinstance(func, ast.Attribute) else None + if fname != "register_task": + return None + names = [] + for arg in node.args: + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + names.append(arg.value) + else: + return [] # dynamic: caller must import the module to learn the names + if node.keywords: + return [] + return names + + +def _parse_registrations(path: str) -> tuple[list[str], bool]: + """(literal task names, has_dynamic_registration) for one source file. + + Every ``register_task(...)`` call is inspected — decorators and plain calls such as + ``register_task(f"family.{n}")(cls)`` in a loop — so a module that registers names the parser + cannot see is flagged dynamic and imported when the full name list is needed. + """ + with open(path, encoding="utf-8") as f: + tree = ast.parse(f.read(), filename=path) + names: list[str] = [] + dynamic = False + for node in ast.walk(tree): + found = _decorator_names(node) + if found is None: + continue + if not found: + dynamic = True + names.extend(found) + return names, dynamic + + +def default_cache_path() -> str: + """Per-file index cache (JSON). ``$METASIM_CACHE_DIR`` overrides the temp-dir default.""" + return os.path.join( + os.environ.get("METASIM_CACHE_DIR", os.path.join(tempfile.gettempdir(), "metasim_cache")), "task_index.json" + ) + + +def _load_cache(path: str) -> dict: + try: + with open(path, encoding="utf-8") as f: + data = json.load(f) + return data if isinstance(data, dict) and data.get("version") == 1 else {} + except (OSError, ValueError): + return {} + + +def _save_cache(path: str, entries: dict) -> None: + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = f"{path}.tmp-{os.getpid()}" + with open(tmp, "w", encoding="utf-8") as f: + json.dump({"version": 1, "files": entries}, f) + os.replace(tmp, path) + except OSError: # a read-only or full cache location must not break lookups + pass + + +def _scan_source(module_name: str, path: str, index: StaticIndex, cache: dict | None = None) -> None: + try: + st = os.stat(path) + stamp = [st.st_mtime, st.st_size] + entry = cache.get(path) if cache is not None else None + if entry is not None and entry.get("stamp") == stamp: + names, dynamic = entry["names"], entry["dynamic"] + else: + names, dynamic = _parse_registrations(path) + if cache is not None: + cache[path] = {"stamp": stamp, "names": names, "dynamic": dynamic} + except (OSError, SyntaxError, ValueError) as exc: + index.parse_failures[module_name] = f"{type(exc).__name__}: {exc}" + return + index.scanned_modules += 1 + if dynamic: + index.dynamic_modules.add(module_name) + for raw in names: + key = raw.strip().lower() + if key: + index.names.setdefault(key, module_name) + + +def _module_source_path(module_name: str) -> str | None: + try: + spec = find_spec(module_name) + except (ImportError, ValueError): + return None + if spec is None or not spec.origin or not spec.origin.endswith(".py"): + return None + return spec.origin + + +def _iter_package_sources(pkg_name: str, pkg_paths) -> list[tuple[str, str]]: + """(module_name, source_path) for every ``.py`` under the package directories, without importing. + + ``pkgutil.walk_packages`` imports each sub-package to descend into it, which runs every + ``__init__`` side effect (some task families import all their modules there) — the very cost this + index exists to avoid. Walking the directories is enough because a module's name is its path. + """ + out: list[tuple[str, str]] = [] + for root_dir in pkg_paths: + root_dir = os.path.abspath(root_dir) + for dirpath, dirnames, filenames in os.walk(root_dir): + dirnames[:] = sorted(d for d in dirnames if d != "__pycache__" and not d.startswith(".")) + rel = os.path.relpath(dirpath, root_dir) + prefix = pkg_name if rel == "." else pkg_name + "." + rel.replace(os.sep, ".") + for fname in sorted(filenames): + if not fname.endswith(".py"): + continue + mod = prefix if fname == "__init__.py" else f"{prefix}.{fname[:-3]}" + out.append((mod, os.path.join(dirpath, fname))) + return out + + +def build_static_index( + package_names: list[str], *, local_modules: list[str] = (), cache_path: str | None = None +) -> StaticIndex: + """Scan ``package_names`` (importable package roots) and ``local_modules`` for task registrations. + + Nothing is imported: package locations come from ``importlib.util.find_spec`` and every source + file is parsed with ``ast``. Parse results are cached per file (keyed by mtime + size) in + ``cache_path`` (default :func:`default_cache_path`; ``""`` disables the cache). + """ + index = StaticIndex() + cache_path = default_cache_path() if cache_path is None else cache_path + cache: dict | None = _load_cache(cache_path).get("files", {}) if cache_path else None + seen_paths: set[str] = set() + for pkg_name in package_names: + if "." not in pkg_name and pkg_name in local_modules: + continue + try: + spec = find_spec(pkg_name) + except (ImportError, ValueError) as exc: + index.parse_failures[pkg_name] = f"{type(exc).__name__}: {exc}" + continue + if spec is None: + index.parse_failures[pkg_name] = "ModuleNotFoundError: no such package" + continue + if spec.submodule_search_locations: + for module_name, src in _iter_package_sources(pkg_name, list(spec.submodule_search_locations)): + if src in seen_paths: + continue + seen_paths.add(src) + _scan_source(module_name, src, index, cache) + elif spec.origin and spec.origin.endswith(".py"): + _scan_source(pkg_name, spec.origin, index, cache) + for mod in local_modules: + src = _module_source_path(mod) + if src is None and os.path.isfile(mod + ".py"): + src = mod + ".py" + if src: + _scan_source(mod, src, index, cache) + if cache is not None: + _save_cache(cache_path, cache) + return index diff --git a/metasim/task/registry.py b/metasim/task/registry.py index 27f3488..6b8a147 100644 --- a/metasim/task/registry.py +++ b/metasim/task/registry.py @@ -7,6 +7,7 @@ from loguru import logger as log +from metasim.task._static_index import StaticIndex, build_static_index from metasim.task.base import BaseTaskEnv from metasim.utils.package_discovery import get_package_candidates @@ -18,6 +19,14 @@ # they look up a missing task, rather than having to enable DEBUG logging. _DISCOVERY_FAILURES: dict[str, str] = {} +# Static (AST) index of ``@register_task`` names -> module, built once per process; see +# ``metasim/task/_static_index.py``. ``METASIM_TASK_DISCOVERY=eager`` restores the old +# import-everything behaviour (useful when debugging import-time side effects). +_STATIC_INDEX: StaticIndex | None = None +_STATIC_INDEX_KEY: tuple[str, ...] | None = None +"""Package list the cached index was built for; a different configuration (env vars, cwd) rebuilds it.""" +_EAGER_DONE = False + def register_task(*names): """Class decorator to register a task under one or more names. @@ -46,8 +55,8 @@ def _decorator(cls): return _decorator -def _discover_task_modules() -> None: - """Import configured task packages so @register_task decorators run.""" +def _task_packages() -> list[str]: + """Configured task packages (plus ``_*task.py`` modules in the CWD), in discovery order.""" cwd = os.getcwd() if cwd not in sys.path: sys.path.insert(0, cwd) @@ -55,13 +64,51 @@ def _discover_task_modules() -> None: local_task_modules = [ os.path.splitext(fname)[0] for fname in os.listdir(cwd) if fname.endswith("task.py") and fname.startswith("_") ] - packages_to_scan = get_package_candidates( + return get_package_candidates( "tasks", defaults=["metasim.example.example_pack.tasks"], local_modules=local_task_modules, cwd=cwd, ) + +def _eager_discovery_requested() -> bool: + return os.environ.get("METASIM_TASK_DISCOVERY", "").strip().lower() == "eager" + + +def _static_index() -> StaticIndex: + """Build (once) the AST index of task names -> module for the configured packages.""" + global _STATIC_INDEX, _STATIC_INDEX_KEY + packages = _task_packages() + key = tuple(packages) + if _STATIC_INDEX is None or _STATIC_INDEX_KEY != key: + local = [m for m in packages if "." not in m] # cwd ``_*task.py`` modules come through as bare names + _STATIC_INDEX = build_static_index(packages, local_modules=local) + _STATIC_INDEX_KEY = key + for mod, err in _STATIC_INDEX.parse_failures.items(): + _DISCOVERY_FAILURES.setdefault(mod, err) + log.warning(f"Task discovery: could not scan '{mod}': {err}") + return _STATIC_INDEX + + +def _import_registering_module(module_name: str) -> None: + """Import one task module so its ``@register_task`` decorators run; record failures.""" + try: + import_module(module_name) + except Exception as e: + err_str = f"{type(e).__name__}: {e}" + _DISCOVERY_FAILURES[module_name] = err_str + log.warning(f"Task discovery: skip module '{module_name}': {err_str}") + + +def _discover_task_modules() -> None: + """Import configured task packages so @register_task decorators run (eager fallback).""" + global _EAGER_DONE + if _EAGER_DONE and TASK_REGISTRY: # an emptied registry (tests, reloads) re-runs discovery + return + _EAGER_DONE = True + packages_to_scan = _task_packages() + for pkg_name in packages_to_scan: try: # Import the root package @@ -98,13 +145,23 @@ def _discover_task_modules() -> None: def get_task_class(name: str) -> type[BaseTaskEnv]: """Return the task wrapper class registered under the given name. - Name lookup is case-insensitive. + Name lookup is case-insensitive. Only the module that registers ``name`` is imported (found + through the static index); the eager import of every task module happens only when the name is + unknown to the index or ``METASIM_TASK_DISCOVERY=eager`` is set. """ - # ensure modules are imported so registry is populated - if not TASK_REGISTRY: - _discover_task_modules() - key = name.strip().lower() + if key in TASK_REGISTRY: + return TASK_REGISTRY[key] + if _eager_discovery_requested(): + _discover_task_modules() + else: + index = _static_index() + module_name = index.names.get(key) + if module_name is not None: + _import_registering_module(module_name) + if key not in TASK_REGISTRY: + # dynamic registrations or a name the index could not see: last resort, import everything + _discover_task_modules() try: return TASK_REGISTRY[key] except KeyError as exc: @@ -119,7 +176,15 @@ def get_task_class(name: str) -> type[BaseTaskEnv]: def list_tasks(): - """List all registered task names (sorted).""" - if not TASK_REGISTRY: + """List all task names (sorted) without importing the task modules. + + Names come from the static index plus anything registered so far; modules whose + ``register_task`` arguments are not string literals are imported to learn their names. + """ + if _eager_discovery_requested(): _discover_task_modules() - return sorted(TASK_REGISTRY.keys()) + return sorted(TASK_REGISTRY.keys()) + index = _static_index() + for module_name in sorted(index.dynamic_modules): + _import_registering_module(module_name) + return sorted(set(index.names) | set(TASK_REGISTRY.keys())) diff --git a/metasim/test/test_task_registry_lazy_general.py b/metasim/test/test_task_registry_lazy_general.py new file mode 100644 index 0000000..232d6cc --- /dev/null +++ b/metasim/test/test_task_registry_lazy_general.py @@ -0,0 +1,172 @@ +"""The task registry resolves names through a static (AST) index and imports one module per lookup. + +A synthetic task package is written to ``tmp_path``: two modules with literal ``@register_task`` +decorators, one that registers names in a loop (dynamic), and one whose import fails. Each module +appends to a marker file when imported, so the tests can assert *which* modules a call imported. +""" + +from __future__ import annotations + +import json +import os +import sys +import textwrap + +import pytest + +from metasim.task import _static_index, registry + + +@pytest.fixture +def task_pkg(tmp_path, monkeypatch): + pkg = tmp_path / "lazypack" + (pkg / "tasks").mkdir(parents=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "tasks" / "__init__.py").write_text("", encoding="utf-8") + marker = tmp_path / "imported.txt" + common = textwrap.dedent( + f""" + from metasim.task.base import BaseTaskEnv + from metasim.task.registry import register_task + + with open({str(marker)!r}, "a") as _f: + _f.write(__name__ + "\\n") + """ + ) + (pkg / "tasks" / "alpha.py").write_text( + common + + textwrap.dedent( + """ + @register_task("lazy.alpha", "Alpha") + class AlphaTask(BaseTaskEnv): + pass + """ + ), + encoding="utf-8", + ) + (pkg / "tasks" / "beta.py").write_text( + common + + textwrap.dedent( + """ + @register_task("lazy.beta") + class BetaTask(BaseTaskEnv): + pass + """ + ), + encoding="utf-8", + ) + (pkg / "tasks" / "looped.py").write_text( + common + + textwrap.dedent( + """ + for _i in range(2): + register_task(f"lazy.looped_{_i}")(type(f"Looped{_i}", (BaseTaskEnv,), {})) + """ + ), + encoding="utf-8", + ) + (pkg / "tasks" / "broken.py").write_text( + common + + textwrap.dedent( + """ + import module_that_does_not_exist_anywhere # noqa: F401 + + @register_task("lazy.broken") + class BrokenTask(BaseTaskEnv): + pass + """ + ), + encoding="utf-8", + ) + monkeypatch.syspath_prepend(str(tmp_path)) + monkeypatch.setenv("METASIM_TASK_PACKAGES", "lazypack.tasks") + monkeypatch.setenv("METASIM_CACHE_DIR", str(tmp_path / "cache")) + monkeypatch.delenv("METASIM_TASK_DISCOVERY", raising=False) + # fresh registry state for every test + monkeypatch.setattr(registry, "TASK_REGISTRY", {}) + monkeypatch.setattr(registry, "_DISCOVERY_FAILURES", {}) + monkeypatch.setattr(registry, "_STATIC_INDEX", None) + monkeypatch.setattr(registry, "_EAGER_DONE", False) + for name in list(sys.modules): + if name.startswith("lazypack"): + del sys.modules[name] + + def imported() -> list[str]: + return marker.read_text(encoding="utf-8").split() if marker.exists() else [] + + return imported + + +@pytest.mark.general +def test_get_task_class_imports_only_the_registering_module(task_pkg): + cls = registry.get_task_class("Lazy.Alpha") + assert cls.__name__ == "AlphaTask" + assert task_pkg() == ["lazypack.tasks.alpha"], "beta / looped / broken must not have been imported" + assert registry._EAGER_DONE is False + + +@pytest.mark.general +def test_list_tasks_sees_static_names_without_importing_them(task_pkg): + names = registry.list_tasks() + assert {"lazy.alpha", "alpha", "lazy.beta", "lazy.looped_0", "lazy.looped_1", "lazy.broken"} <= set(names) + # only the dynamic module had to be imported to learn its names + assert task_pkg() == ["lazypack.tasks.looped"] + + +@pytest.mark.general +def test_unknown_name_falls_back_to_eager_discovery_and_reports_failures(task_pkg): + with pytest.raises(KeyError) as excinfo: + registry.get_task_class("lazy.missing") + msg = str(excinfo.value) + assert "lazypack.tasks.broken" in msg and "module_that_does_not_exist_anywhere" in msg + assert registry._EAGER_DONE is True + + +@pytest.mark.general +def test_broken_module_error_is_surfaced_on_lookup(task_pkg): + with pytest.raises(KeyError, match="module_that_does_not_exist_anywhere"): + registry.get_task_class("lazy.broken") + + +@pytest.mark.general +def test_eager_env_var_restores_import_everything(task_pkg, monkeypatch): + monkeypatch.setenv("METASIM_TASK_DISCOVERY", "eager") + registry.get_task_class("lazy.beta") + assert set(task_pkg()) == { + "lazypack.tasks.alpha", + "lazypack.tasks.beta", + "lazypack.tasks.looped", + "lazypack.tasks.broken", + } + + +@pytest.mark.general +def test_index_cache_is_written_and_reused(task_pkg, tmp_path): + registry.list_tasks() + cache_file = tmp_path / "cache" / "task_index.json" + assert cache_file.is_file() + data = json.loads(cache_file.read_text(encoding="utf-8")) + alpha_path = str(tmp_path / "lazypack" / "tasks" / "alpha.py") + assert data["files"][alpha_path]["names"] == ["lazy.alpha", "Alpha"] + assert data["files"][str(tmp_path / "lazypack" / "tasks" / "looped.py")]["dynamic"] is True + # a second build with an untouched tree reads the cache: parse nothing, same names + parsed_before = _static_index._parse_registrations + calls = [] + + def _counting(path): + calls.append(path) + return parsed_before(path) + + _static_index._parse_registrations = _counting + try: + idx = _static_index.build_static_index(["lazypack.tasks"]) + finally: + _static_index._parse_registrations = parsed_before + assert calls == [] + assert idx.names["lazy.alpha"] == "lazypack.tasks.alpha" + # editing a file invalidates only that file + beta = tmp_path / "lazypack" / "tasks" / "beta.py" + beta.write_text(beta.read_text(encoding="utf-8").replace('"lazy.beta"', '"lazy.beta2"'), encoding="utf-8") + os.utime(beta, (os.stat(beta).st_atime + 5, os.stat(beta).st_mtime + 5)) + idx = _static_index.build_static_index(["lazypack.tasks"]) + assert "lazy.beta2" in idx.names and "lazy.beta" not in idx.names