Skip to content

Zero-copy joint access from Python - #836

Open
aristarkhovNV wants to merge 5 commits into
mainfrom
aaristarkhov/zero-copy-joints-access
Open

Zero-copy joint access from Python#836
aristarkhovNV wants to merge 5 commits into
mainfrom
aaristarkhov/zero-copy-joints-access

Conversation

@aristarkhovNV

@aristarkhovNV aristarkhovNV commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

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. HandJoints and BodyJoints gain per-field accessors — positions, orientations, radii, is_valid — returning strided arrays that alias the FlatBuffers joint storage. Reading all joints through poses(i)/joints(i) cost one 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.

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 via FBS_FIELD_OFFSET. This also fixes a latent flatc flag bug: --reflect-names and --reflect-types assign 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 -T wrapper by value, and flatc gives each -T a deep-copying copy constructor even though the payload sits behind a shared_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 .data property has always done. The Python-visible type is unchanged, so call sites and isinstance checks 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 FullBodyPoseT per frame. The hand impl was already fresh, but only because update_hand happened 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:

  • Snapshot — hand, full body, and message channel. Safe to retain; call again for newer data.
  • Aliased, refilled by the next session.update() — head, controller, and the six schema-backed trackers. Read before that call, or copy.

TeleopSession runs 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

  • The views are writable and alias schema storage; callers that modify data must .copy() first. Read-only would be safer, but NumPy cannot export a read-only array over DLPack before 2.1 and the wheel declares numpy>=1.23. In-repo transform nodes copy before mutating; nothing enforces it for external nodes.
  • hands_source.py does not None-check the optional joints struct, while full_body_source.py does. Replaying a HandPose serialized without joints raises AttributeError. Pre-existing rather than introduced here — the prior joints.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

@coderabbitai

coderabbitai Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c41cf3e6-130a-4880-9f77-90f52a18b0bb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

The 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
Loading
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 78.95% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly captures the main change: zero-copy Python access to joint data.
✨ Finishing Touches 💡 1
📝 Generate docstrings 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch aaristarkhov/zero-copy-joints-access

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between a3513b7 and 9950866.

📒 Files selected for processing (11)
  • cmake/GenerateFlatBuffers.cmake
  • src/core/deviceio_trackers/python/tracker_bindings.cpp
  • src/core/live_trackers/cpp/live_full_body_tracker_pico_impl.cpp
  • src/core/retargeting_engine/python/deviceio_source_nodes/full_body_source.py
  • src/core/retargeting_engine/python/deviceio_source_nodes/hands_source.py
  • src/core/schema/python/CMakeLists.txt
  • src/core/schema/python/full_body_bindings.h
  • src/core/schema/python/hand_bindings.h
  • src/core/schema/python/schema_array_views.h
  • src/core/schema_tests/python/test_full_body.py
  • src/core/schema_tests/python/test_hand.py

Comment thread src/core/deviceio_trackers/python/tracker_bindings.cpp Outdated
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>
@aristarkhovNV
aristarkhovNV force-pushed the aaristarkhov/zero-copy-joints-access branch from a536dec to d5d4b8b Compare July 31, 2026 17:24
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