perf(task): lazy task registry backed by a static AST index - #32
Open
geng-haoran wants to merge 1 commit into
Open
perf(task): lazy task registry backed by a static AST index#32geng-haoran wants to merge 1 commit into
geng-haoran wants to merge 1 commit into
Conversation
`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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017i6VtKoovBNed815mWFqxw
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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-linegenerated 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_packageswould import each sub-package
__init__, some of which eagerly import their whole family), parseseach file with
ast, and records everyregister_task(...)call with string-literal arguments asname -> 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_classimports only the registering module; unknown names fall back to theold import-everything discovery so the KeyError still lists import failures.
registry.list_tasksreturns the index plus dynamic modules' names. The cached index is keyed bythe package configuration and rebuilt when it changes (env vars, cwd); an emptied registry re-runs
discovery.
METASIM_TASK_DISCOVERY=eagerrestores the previous behaviour.Measured on RoboVerse main (6051 tasks): index cold 1.08 s / warm 0.04 s;
get_task_class0.01-0.1 s(mjlab 0.9 s = its own imports);
list_tasks4.9 s (3.5 s of it is importingmaniskill.native_tasks, which registers 5360 names in a loop) versus 7.3 s before. All 2898decorators in RoboVerse are string literals.
Tests:
metasim/test/test_task_registry_lazy_general.py(6 tests on a synthetic package: singlemodule imported per lookup, list without imports, eager fallback + failure reporting, env var,
cache reuse/invalidation).
pytest -k general-> 492 passed;-k mujocounchanged.Co-Authored-By: Claude Fable 5.1 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_017i6VtKoovBNed815mWFqxw
🤖 Generated with Claude Code
https://claude.ai/code/session_017i6VtKoovBNed815mWFqxw