[DataOriented] Fix ndarrays on data oriented (reland of #704)#723
Merged
Conversation
…_oriented PR #704 introduced two ndarray-walkers (the hotpath ``_build_struct_nd_paths`` in ``_template_mapper_hotpath.py`` and the compile-time ``_walk_obj`` in ``function_def_transformer.py``) which together caused Genesis unit tests to ``RecursionError`` and crash on first kernel launch. Two independent bugs: 1. **No cycle protection.** Real Genesis containers form attribute graphs with back-pointers (e.g. ``sim.rigid_solver.sim is sim``). Both walkers recursed infinitely on the back-edge. Fixed by threading an ``id(obj)``-keyed seen-set through each walk. 2. **Pydantic metaclass blew the stack.** Both walkers used ``dataclasses.is_dataclass(child) and not isinstance(child, type)`` (which calls ``hasattr(type(obj), '__dataclass_fields__')``) and ``is_data_oriented`` (which called ``getattr(type(obj), '_data_oriented', False)``). Pydantic's ``ModelMetaclass.__getattr__`` recurses infinitely on arbitrary missing attribute names, so probing a ``RigidOptions`` instance for either flag blew the stack. Fixed by adding ``is_dataclass_instance`` in ``util.py`` and rewriting ``is_data_oriented`` to walk ``type(obj).__mro__`` and probe each class's ``__dict__`` directly, bypassing any descriptor / ``__getattr__`` machinery. Both walkers now route through these MRO-safe helpers.
The original #704 cached ndarray-attribute paths per CLASS, which is incorrect under polymorphic instance structure: ``@qd.data_oriented`` and ``qd.Tensor`` members can have different leaf-types across instances of the same class — Genesis ``DataManager`` only allocates ``*_adjoint_cache`` ndarrays when ``requires_grad=True``, and a ``qd.Tensor`` field can wrap an ``Ndarray`` on one instance and a ``MatrixField`` on another (e.g. ndarray-vs-field backend). ``_collect_struct_nd_descriptors`` then crashed reading ``element_type`` / ``shape`` off the ``MatrixField`` (``'MatrixField' object has no attribute 'element_type'``), affecting ~60 Genesis unit tests under the ndarray backend. Fix: cache the path list on the *instance* via ``object.__setattr__`` (so frozen dataclasses work too); ``__slots__`` classes without a ``__dict__`` silently fall back to the class-level cache (Genesis containers don't use slots). Also add defensive ``isinstance(v, Ndarray)`` and ``shape is None`` guards in ``_collect_struct_nd_descriptors`` for the case where a leaf changes type *within* an instance's lifetime.
…-class The args_hash data_oriented walker that #704 added ran unconditionally for every arg of every kernel call. Even with the per-class attribute-path cache, the per- call ``is_data_oriented(arg)`` + ``type(arg).__dict__.get`` chain still cost ~15% FPS on small-step CPU benches (anymal_zero CPU bs=0: -17.6% vs the pre-#704 reference) and several other Genesis CPU/GPU workloads were also significantly regressed (anymal_uniform GPU -11.5%, franka_random CPU -17.5%, …). Two coordinated optimisations: 1. Only iterate ``self.template_slot_locations`` instead of all args. A data_oriented class is never a dataclass and typed-dataclass args carry a specific dataclass type by construction, so the only positions where a data_oriented container can appear are the ``qd.template()`` annotated ones — already tracked by the kernel decorator. Genesis main ``kernel_step_1`` has 4 template positions of 16 args; reduces the per-call work proportionally. 2. Per-``type(arg)`` precomputed dispatch: ``_arg_nd_paths_or_none`` maps each seen type to either the cached path list to walk, or ``None`` (skip — covers primitive templates, non-data_oriented composites, and data_oriented with zero ndarrays). One ``dict.get`` per candidate per call after warmup, replacing the previous chain. Brings the post-#704 Genesis bench numbers from ~ -8 to -18% on CPU and ~ -10% on GPU anymal_uniform back to within noise margin (target: matching pre-#704).
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 42a31e213f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ann = arg_meta.annotation | ||
| if ann is not template and type(ann) is not template: | ||
| continue | ||
| if dataclasses.is_dataclass(val) and not isinstance(val, type) and not is_data_oriented(val): |
There was a problem hiding this comment.
For a qd.template() argument whose class uses a metaclass with a problematic __getattr__ (the same Genesis/Pydantic case this commit fixes elsewhere), this warning path calls dataclasses.is_dataclass(val) before the safe is_data_oriented check, so the first compile can still recurse or crash instead of reaching the new MRO-safe helpers. This should use the new is_dataclass_instance() helper here as well, otherwise non-dataclass template objects can reintroduce the RecursionError on first materialization.
Useful? React with 👍 / 👎.
Comment on lines
+123
to
+124
| _arg_nd_paths_or_none[cls] = paths | ||
| if paths is None: |
There was a problem hiding this comment.
Caching None for the whole class means the first observed @qd.data_oriented instance with no ndarray members disables the ndarray-id walk for every later instance of that type. If a later same-type instance has ndarray members and then reassigns one to a different dtype/ndim, args_hash remains just id(arg) and lookup() can return the stale spec key without calling extract(), reusing a kernel specialized for the old ndarray. This conflicts with the per-instance polymorphism handled in _struct_nd_paths_for; the skip decision needs to be per instance or rechecked when paths may appear.
Useful? React with 👍 / 👎.
The fix in fbfa65d ([Fix] Per-instance ndarray-path cache from PR #704) changes the granularity of the template-mapper's ndarray-attribute-path cache from per-class to per-instance. Update the "Restrictions" bullet under the ``@qd.data_oriented`` section to reflect the new behaviour: - Per-instance caching: two instances of the same class can carry different ndarray attribute sets (e.g. ``DataManager`` with vs without ``requires_grad`` and its optional ``*_adjoint_cache`` members). - The polymorphic-instance / stale-kernel caveat that the original #704 shipped no longer applies — the per-instance cache disambiguates them. - The remaining edge cases (deleting an attribute on an instance after its first launch; adding then reassigning a new attribute on an instance post-first-launch) are now phrased per-instance rather than per-class.
Collaborator
Author
|
Genesis unit testes result: 
|
Collaborator
Author
|
Gensis benchmarks result: 
|
…on test Three CI complaints on #723: 1. **pyright**: ``_template_mapper.py:125`` ``"object" is not iterable``. The ``_arg_nd_paths_or_none.get(cls, _UNCLASSIFIED)`` sentinel pattern widened ``paths`` to ``list[tuple] | None | object``, and the ``if paths is None`` narrowing didn't drop the ``object`` branch. Replaced with a ``try / except KeyError`` lookup so the cached value is unambiguously ``list[tuple] | None``. 2. **Check line wrapping**: prose lines wrapped well below the project's 120c target in ``test_ad_dataclass.py`` (lines 8, 10, 11, 226-227, 266-268), ``test_data_oriented_mixed_combos.py`` (lines 3-4), ``test_data_oriented_ndarray.py`` (lines 510-511, 536-537, 721-722), and ``test_py_dataclass.py`` (lines 2884-2887). Reflowed to ~120c per the ``find_underwrapped`` skill. 3. **Check test coverage for changes**: no regression test exercised the ``_seen`` cycle-detection guard in ``_build_struct_nd_paths`` / ``_walk_obj``. Added ``test_data_oriented_attribute_cycle_does_not_recurse _infinitely`` covering the canonical ``sim.solver.sim is sim`` shape. The fourth failure (arm64 wheel test ``test_no_premature_shutdown[none]``) is unrelated to this PR — it tests pytest's xdist worker retirement and is flaky on arm64.
Contributor
|
ok to merge |
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
2 participants
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.
Relands #704 (reverted in #719 due to Genesis
RecursionError+MatrixFieldcrashes) with only the minimum fixes needed to make Genesis unit tests and benchmarks green. Built directly on top of the original #704 squash commit; no unrelated changes.What's in the branch
Doc changes compared to #704
one stale claim in compound_types.md → "Restrictions" section. The original #704 said the ndarray-path cache is per class and shipped a documented caveat about polymorphic-instance silent-stale-kernel reuse. My per-instance cache fix (commit fbfa65d) makes it per instance and removes that caveat. The bullet now correctly describes per-instance caching and reformulates the remaining edge cases (delete-on-instance, add-and-reassign-after-first-launch) in per-instance terms.
Details
Four commits on top of
main:_build_struct_nd_paths/_walk_obj(Genesis attribute graphs have back-pointers) and MRO-safeis_data_oriented/is_dataclass_instance(Pydantic'sModelMetaclass.__getattr__recursed infinitely on probe-stylegetattrofRigidOptionsinstances).'MatrixField' object has no attribute 'element_type'crash on ~60 Genesis ndarray-backend tests. The cache is stored off-instance in aweakref.finalize-cleaned dict so it isn't visible to the fastcache args walker (which rejectslistchildren).TemplateMapper.lookupiterates onlytemplate_slot_locationsand uses a per-class precomputed dispatch (_arg_nd_paths_or_none), recovering the ~15% FPS that [DataOriented] Fix ndarrays on data oriented #704's unconditional per-call walk cost on small-step CPU benches.Verification
Cluster, 8-GPU
rtx-mid, containergenesis-v1_22.sqsh, Genesis main,pytest -v -ra --backend gpu --dev --forked ./tests:GS_ENABLE_NDARRAY=1: 742 passed, 0 failed, 4 skipped, 6 xfailedGS_ENABLE_NDARRAY=0: 743 passed, 0 failed, 2 skipped, 7 xfailedGenesis speed benchmarks (
bench_cluster_wandb.py, branch vs main from WandB CI):All within typical noise margin. No regression worse than -5.3%.
Supersedes #721.
Made with Cursor