Skip to content

perf(task): lazy task registry backed by a static AST index - #32

Open
geng-haoran wants to merge 1 commit into
mainfrom
perf/lazy-task-registry
Open

perf(task): lazy task registry backed by a static AST index#32
geng-haoran wants to merge 1 commit into
mainfrom
perf/lazy-task-registry

Conversation

@geng-haoran

Copy link
Copy Markdown
Contributor

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

🤖 Generated with Claude Code

https://claude.ai/code/session_017i6VtKoovBNed815mWFqxw

`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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant