[zxpdemonio]

Motivation

PR2050 added a broad DataProto-oriented transfer stack for RL batch exchange. That PR is now being split into smaller, reviewable pieces so reviewers can validate the structured-object foundation before lazy batch, planner, sharding, and DataProto compatibility layers are added.

This PR adds the minimal flat structured-object object API and tensor fast path on top of the existing structured_object_store.py bundle/manifest infrastructure that has already landed on main.

Scenario

This targets the first step of remote structured-object transfer in RL pipelines:

1. A producer stores a flat mapping or a single wrapped object through MooncakeBundleTransfer.
2. NumPy arrays are stored as structured ndarray members with dtype/shape metadata.
3. Torch tensors are stored through the Mooncake Store tensor path when put_tensor/get_tensor are available.
4. A consumer materializes the object back with get_object() or the existing read_spec() / materialize() path.

Later split PRs will layer nested mappings, TensorDict support, row/list codecs, partial materialization, lazy remote batch, planner, and DataProto compatibility on this foundation.

Description

This PR adds the flat structured-object MVP split from PR2050.

Structured object convenience API

• MooncakeBundleTransfer.put_object(obj, ...)
  • stores a flat mapping directly as structured object members;
  • wraps a single non-mapping object under an internal value member.
• MooncakeBundleTransfer.get_object(ref)
  • materializes objects written by put_object();
  • unwraps single-object payloads automatically.

Torch tensor fast path

• Adds an internal _TensorPayload marker for structured tensor fields.
• Encodes torch tensors as torch_tensor structured members with dtype/shape metadata.
• Uses Store put_tensor() when available.
• Uses Store get_tensor() when materializing tensor payloads.
• Falls back to raw bytes payload storage/read when the tensor APIs are unavailable.

Tests

• Extends the in-memory test store with put_tensor() / get_tensor() hooks.
• Verifies flat NumPy + torch tensor object roundtrip.
• Verifies wrapped NumPy + torch tensor value roundtrip.
• Asserts tensor fields use the direct tensor fast path.

Basic Usage

transfer = MooncakeBundleTransfer(store, key_prefix="rl/structured")                                                                                                               
                                                                                                                                                                                   
ref = transfer.put_object(                                                                                                                                                         
    {                                                                                                                                                                              
        "tokens": torch.arange(8, dtype=torch.int64).reshape(2, 4),                                                                                                                
        "scores": np.arange(2, dtype=np.float32),                                                                                                                                  
    }                                                                                                                                                                              
)                                                                                                                                                                                  
                                                                                                                                                                                   
obj = transfer.get_object(ref)                                                                                                                                                     
                                                                                                                                                                                   
assert torch.equal(obj["tokens"], torch.arange(8, dtype=torch.int64).reshape(2, 4))                                                                                                
assert np.array_equal(obj["scores"], np.arange(2, dtype=np.float32))                                                                                                               
                                                                                                                                                                                   
Single wrapped value:                                                                                                                                                              
                                                                                                                                                                                   
ref = transfer.put_object(torch.arange(4))                                                                                                                                         
tensor = transfer.get_object(ref)     

Architecture

This PR keeps the abstraction generic and framework-neutral.

caller object / flat mapping
│
▼
MooncakeBundleTransfer.put_object()
│
├── numpy fields ─────► structured ndarray payload + dtype/shape metadata
├── torch fields ─────► Mooncake Store put_tensor/get_tensor fast path
└── bytes fields ─────► existing bundle payload path
│
▼
existing structured_object_store bundle manifest

This is intentionally only the first layer from PR2050. It does not introduce DataProto-specific APIs or a parallel manifest format.

Robustness and Safety

• Keeps structured metadata JSON-serializable.
• Rejects user-provided reserved structured metadata keys.
• Keeps tensor payload cleanup under the existing bundle removal path.
• Avoids pickle or general Python object deserialization.
• Preserves the existing bytes/ndarray structured-object behavior.

Performance

The main performance improvement in this split is avoiding raw byte serialization for torch tensor fields when Store tensor APIs are available:

• tensor writes use put_tensor();
• tensor reads use get_tensor();
• non-tensor flat members continue using the existing bundle payload path.

Planner, lazy selection, range reads, sharding, and DataProto adaptive policy are intentionally left to later split PRs.

How Has This Been Tested?

• python -m py_compile mooncake-wheel/mooncake/structured_object_store.py mooncake-wheel/tests/test_structured_object_store.py
• PYTHONPATH=/tmp/mooncake-pr2050-rebase-20260624/mooncake-wheel pytest -q mooncake-wheel/tests/test_structured_object_store.py
  • 38 passed
• PYTHONPATH=/tmp/mooncake-pr2050-rebase-20260624/mooncake-wheel pytest -q mooncake-wheel/tests/test_structured_object_store.py -k 'put_object or structured_object_roundtrip or metadata_defaults or invalid_metadata'
  • 5 passed, 33 deselected
• ./scripts/code_format.sh --check
  • no C/C++ files changed

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit run --all-files and all hooks pass
• I have updated the documentation (if applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specify below)

Claude Opus 4.6 assisted with code generation, code review, benchmark analysis, and PR description drafting. All changes were reviewed and validated by me。

[gemini-code-assist]

Code Review

This pull request introduces support for storing and retrieving PyTorch tensors directly in the structured object store, adding helper methods put_object and get_object along with corresponding serialization and deserialization logic. The review feedback focuses on optimizing the tensor fast path for GPU-enabled stores by avoiding eager CPU and NumPy conversions during encoding. Specifically, it is recommended to redefine _TensorPayload to lazily evaluate the CPU/NumPy array and its size (nbytes), and to raise a clear RuntimeError if a tensor stored via put_tensor cannot be retrieved due to a lack of get_tensor support in the underlying store.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[zxpdemonio]

@gemini-code-assist please review again

[gemini-code-assist]

Code Review

This pull request introduces support for PyTorch tensors in the structured object store, adding new put_object and get_object methods to handle mapping objects or single tensors/arrays, along with corresponding encoding and decoding helper functions. The review feedback suggests two key improvements: first, optimizing _read_torch_tensor_member to avoid double copying of tensor data by reading directly into an allocated numpy array buffer; second, using memoryview.cast("B") in _tensor_payload_bytes_view for a more idiomatic and robust way to obtain a 1D bytes view of the tensor.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[zxpdemonio]

@gemini-code-assist please review again

[gemini-code-assist]

Code Review

This pull request adds support for storing and materializing PyTorch tensors in the structured object store, introducing a fast-path optimization when the underlying store supports tensor-specific operations, as well as a raw fallback path. The review feedback highlights critical edge cases concerning 0-dimensional (scalar) and boolean PyTorch tensors, which currently cause runtime errors during serialization and materialization. Actionable suggestions are provided to resolve these issues by flattening and reshaping the arrays before casting, alongside recommendations to expand the test suite to cover these scenarios.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!