Zero-copy joint access from Python - #836
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe changes enable FlatBuffers field-name reflection and add shared-storage tracker wrappers. Full-body updates now publish a newly allocated pose per frame. Python bindings expose owner-retaining, strided NumPy views for hand and body joint fields. Retargeting sources consume these views directly, while tests verify shapes, aliasing, write-through behavior, lifetime retention, and copying. Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Tracker
participant PybindBindings
participant SchemaViews
participant RetargetingSource
Tracker->>PybindBindings: publish shared tracker data
PybindBindings->>SchemaViews: expose joint storage
SchemaViews->>RetargetingSource: provide strided NumPy views
RetargetingSource->>RetargetingSource: use positions, orientations, and validity fields
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/core/deviceio_trackers/python/tracker_bindings.cpp`:
- Around line 251-259: Update the lifetime documentation attached to
FullBodyTracker.get_body_pose and TRACKED_LIFETIME_DOC to state that returned
wrappers are frame snapshots and are not refilled on subsequent updates. Clarify
that callers must request a new pose to observe new body data, matching
LiveFullBodyTrackerPicoImpl::update() replacing tracked_.data.
In `@src/core/retargeting_engine/python/deviceio_source_nodes/hands_source.py`:
- Around line 136-141: Update the hand joint forwarding logic in the source-node
method containing the HandInputIndex assignments to snapshot each joints buffer
before storing it in group. Copy positions, orientations, radii, and is_valid so
downstream asynchronous consumers cannot observe later in-place mutations of the
reused HandJoints object.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 64d2d779-cc79-40f9-8d82-04898aef9570
📒 Files selected for processing (11)
cmake/GenerateFlatBuffers.cmakesrc/core/deviceio_trackers/python/tracker_bindings.cppsrc/core/live_trackers/cpp/live_full_body_tracker_pico_impl.cppsrc/core/retargeting_engine/python/deviceio_source_nodes/full_body_source.pysrc/core/retargeting_engine/python/deviceio_source_nodes/hands_source.pysrc/core/schema/python/CMakeLists.txtsrc/core/schema/python/full_body_bindings.hsrc/core/schema/python/hand_bindings.hsrc/core/schema/python/schema_array_views.hsrc/core/schema_tests/python/test_full_body.pysrc/core/schema_tests/python/test_hand.py
HandJoints and BodyJoints gain per-field accessors (positions, orientations, radii, is_valid) returning strided NumPy arrays that alias the FlatBuffers joint storage. Reading all joints through poses(i)/joints(i) costs one Python->C++ round trip per field per joint; the views cost one call and no copy. HandsSource and FullBodySource now feed the retargeting graph directly from them. The joints are an interleaved array of structs, so each field is strided rather than packed. NumPy and DLPack both carry strides natively, so the views satisfy the NDArrayType tensor contract as-is. They alias schema storage and are writable, so callers that modify data must .copy() first; read-only would be the safer contract but NumPy cannot export a read-only array over DLPack before 2.1 and the wheel declares numpy>=1.23. Field locations are derived, not restated. Nested struct fields are reached through the generated accessors, so a schema rename fails the build. The two scalar fields have no addressable accessor (flatc returns scalars by value through EndianScalar) and instead resolve an offset from flatc's reflection table once at binding time, with the name compile-time checked against the accessor via FBS_FIELD_OFFSET. flatc was passed --reflect-names --reflect-types; both assign the same option, so names were silently dropped and the reflection tables carried none. It now passes --reflect-names alone. The generator command also depends on the file holding the flag list, so a flag change retriggers regeneration under generators that compare timestamps. Signed-off-by: Andrei Aristarkhov <aaristarkhov@nvidia.com>
Every tracker accessor returned the generated -T wrapper by value. flatc gives each -T a deep-copying copy constructor even though the payload sits behind a shared_ptr, so get_left_hand() cloned 936 bytes of joints plus two allocations on every call, twice per frame, and get_messages() cloned every drained payload. The accessors now copy the data pointer into a fresh wrapper: one small allocation and a refcount bump, which is what the .data property on these wrappers has always done. The Python-visible type is unchanged, so call sites and isinstance checks are untouched. The result is therefore a view into live tracker storage rather than a snapshot. Head, controller, full-body and the schema trackers refill that storage in place on the next session.update(), so every accessor docstring now carries the lifetime note. TeleopSession already runs update, tracker polling and graph execution together inside one step, deliberately, to avoid passing raw DeviceIO state across threads; that is the invariant this relies on. MessageChannelSource is the one consumer that retains the wrapper in a member across frames, and it stays correct: the live impl clears its vector and pushes freshly allocated messages, so the copied vector keeps them alive and unmutated. The hand tracker likewise allocates a fresh pose each frame, but that is an implementation detail and the docstring does not promise it. Signed-off-by: Andrei Aristarkhov <aaristarkhov@nvidia.com>
Sharing tracker snapshots with Python turned the strided NumPy joint views into views over live tracker storage. The PICO full-body impl refilled its BodyJoints array in place, so an application holding JOINT_POSITIONS across a step() would see it change: TeleopSession._step_sync returns graph outputs without a snapshot, and the next step() updates DeviceIO before recomputing. The impl now builds a fresh FullBodyPoseT per frame and installs it, so each frame's joints are independent. One 768-byte allocation per frame, well under the per-call copy this series removed. The hand tracker and both replay impls already worked this way, so full body was the only accessor whose views aliased reused storage. Also from review of that change: share_tracked() carries only .data, and a Tracked wrapper that grew a second field would have been dropped silently -- the parallel Record wrappers already carry a timestamp, so that is a plausible schema edit. A static_assert on the wrapper's width now fails the build instead; NativeTable is empty, so a one-field wrapper is exactly as wide as its member. get_messages() no longer carries the shared-storage lifetime note. Its payload is a vector, so share_tracked() copies the list, and the live and replay impls both push freshly allocated messages. That accessor really does return a snapshot, and the note would have prompted a needless defensive copy. Signed-off-by: Andrei Aristarkhov <aaristarkhov@nvidia.com>
The full-body impls install freshly allocated joint storage each frame rather than refilling the previous frame's, so the wrapper handed to Python is a snapshot of the frame it was read on. The shared lifetime note claimed the opposite and would have prompted a defensive copy that buys nothing. The note stays on the other accessors, which do still alias: the head impl mutates its pose in place, the controller impl reuses its snapshot and swaps sub-pointers, and the schema trackers UnPackTo into the same object. Signed-off-by: Andrei Aristarkhov <aaristarkhov@nvidia.com>
The live hand impl allocated its pose only when the target was empty, and got fresh storage per frame purely because update_hand() passes a fresh local candidate. Nothing stated that, so folding the candidate away to avoid the move would have silently turned the accessor into an in-place refill -- and HandsSource hands out strided NumPy views over exactly that joint array, so callers holding an earlier frame's views would have seen them change. try_update_hand() now builds the pose unconditionally and installs it, matching the full-body impl and making the invariant explicit rather than emergent. The replay impl already moves in a freshly read record, so it has no reuse pattern that could regress. get_left_hand / get_right_hand consequently drop the shared lifetime note: like full body, they return a snapshot of the frame they were read on. Signed-off-by: Andrei Aristarkhov <aaristarkhov@nvidia.com>
a536dec to
d5d4b8b
Compare
Reading hand and body joints from Python cost a Python→C++ round trip per field per joint, plus a full deep copy of the payload on every tracker call. This removes both.
What changed
Joint arrays reach Python as zero-copy NumPy views.
HandJointsandBodyJointsgain per-field accessors —positions,orientations,radii,is_valid— returning strided arrays that alias the FlatBuffers joint storage. Reading all joints throughposes(i)/joints(i)cost one round trip per field per joint; the views cost one call and no copy.HandsSourceandFullBodySourcenow feed the retargeting graph directly from them.The joints are an interleaved array of structs, so each field is strided rather than packed. NumPy and DLPack both carry strides natively, so the views satisfy the
NDArrayTypetensor contract as-is.Field locations are derived, not restated: nested struct fields are reached through the generated accessors, so a schema rename fails the build. The two scalar fields have no addressable accessor — flatc returns scalars by value through
EndianScalar— and instead resolve an offset from flatc's reflection table once at binding time, with the name compile-time checked viaFBS_FIELD_OFFSET. This also fixes a latent flatc flag bug:--reflect-namesand--reflect-typesassign the same option, so passing both silently dropped the names and left the reflection tables empty.Tracker accessors share their payload instead of deep-copying it. Every accessor returned the generated
-Twrapper by value, and flatc gives each-Ta deep-copying copy constructor even though the payload sits behind ashared_ptr.get_left_hand()cloned 936 bytes of joints plus two allocations on every call, twice per frame;get_messages()cloned every drained payload. The accessors now copy the data pointer into a fresh wrapper: one small allocation and a refcount bump, which is what the.dataproperty has always done. The Python-visible type is unchanged, so call sites andisinstancechecks are untouched.The joint trackers publish fresh storage each frame. Sharing payloads turns the views into views over live tracker storage, so anything emitting them must not refill in place. The PICO full-body impl did, and now builds a fresh
FullBodyPoseTper frame. The hand impl was already fresh, but only becauseupdate_handhappened to pass an empty local — folding that away would have silently reintroduced the aliasing, so the allocation is now unconditional and the invariant is stated in both impls.Lifetime contract
Accessors fall into two groups, and each docstring says which:
session.update()— head, controller, and the six schema-backed trackers. Read before that call, or copy.TeleopSessionruns update, tracker polling, and graph execution together inside one step — deliberately, to avoid passing raw DeviceIO state across threads — so the in-repo consumer path reads before any refill.Known trade-offs
.copy()first. Read-only would be safer, but NumPy cannot export a read-only array over DLPack before 2.1 and the wheel declaresnumpy>=1.23. In-repo transform nodes copy before mutating; nothing enforces it for external nodes.hands_source.pydoes not None-check the optionaljointsstruct, whilefull_body_source.pydoes. Replaying aHandPoseserialized withoutjointsraisesAttributeError. Pre-existing rather than introduced here — the priorjoints.poses(i)path had the same gap — but the asymmetry is worth closing separately.Testing
Full suite green (297/297) at each commit. Schema tests cover the view mechanics: layout and offsets, stride correctness, aliasing, and DLPack export on all seven properties.
Not covered: the per-frame-fresh invariant for the live trackers needs real XR runtimes (
XR_BD_body_tracking,XR_EXT_hand_tracking), and the existing joint coverage runs through the replay path, which allocates fresh regardless. That invariant rests on the comments in the two live impls.🤖 Generated with Claude Code