[feat] Support data_parser in kv_batch_get for SimpleStorage

[OutstanderWang]

Summary

In RL training with complex parallel strategies (e.g. CP / context parallel), a rank often needs only part of a stored value:

• a slice of a tensor field, e.g. input_ids trimmed to the sample length or the rank's shard;
• some items of a list field, e.g. a few images from a multi-image sample.

TransferQueue can store many columns, and each one may need different handling. This PR adds an optional data_parser to kv_batch_get / async_kv_batch_get, the get-side counterpart of kv_batch_put's data_parser. Each SimpleStorage unit runs it on its share of the batch before replying, so only the parsed values cross the network. It is an advanced feature, and its constraints are documented.

The parser reduces network traffic, not disk I/O. SSD-offloaded values are still read and decoded in full before parsing.

Usage

import transfer_queue as tq

def pick(data, args):
    # data: {field: [stored value of each sample on this unit]}; args: the matching parser_args
    data["input_ids"] = [ids[:n] for ids, (n, _) in zip(data["input_ids"], args)]
    data["images"] = [[imgs[i] for i in idx] for imgs, (_, idx) in zip(data["images"], args)]
    return data

data = tq.kv_batch_get(
    keys=["s0", "s1"],
    partition_id="train",
    select_fields=["input_ids", "images"],
    data_parser=pick,
    parser_args=[(5, [0, 2]), (3, [1])],
)

API

• data_parser(data, args) -> dict. data maps each selected field to a list with one stored value per sample on that unit. args holds the matching parser_args entries, or None without parser_args. Each field may be parsed differently. The result is checked the same way as on put: same keys, and the same number of values per field.
• parser_args are optional per-key arguments aligned with keys. They require data_parser, a matching length, and every key to exist, because missing keys are dropped from the meta and would shift the arguments onto other keys. TransferQueueClient.(async_)get_data checks the first two, so direct client callers get the same errors. kv_batch_get checks that every key was found. All of this runs before any storage request is sent. They are serialized like request data, so tuples arrive as lists.
• No retry with a parser, as on put (max_attempts=1). A storage-unit timeout raises instead of running the parser again.
• SimpleStorage only. KVStorageManager (MooncakeStore, Yuanrong, RayStore) raises NotImplementedError before fetching anything, as it does for put.

Constraints documented in the kv_batch_get docstring and the README:

• Do not modify stored values in place (e.g. tensor.add_(), value.append()). data holds the stored objects, so the change would persist and reach later readers. Build new values instead; replacing entries of data is safe.
• Keep the parser free of side effects. It is not retried internally, but a caller that retries the call runs it again.
• The parser is sent with cloudpickle, so a module-level function must be importable on the storage units.

Design

An earlier revision replaced the parser with a declarative select_items, after the review. That proved too narrow: it applies one dim-0 selection to every selected field, while real layouts need per-column handling (trim one field, pick items from another, decode a third). This revision goes back to a user callable, but with the same contract as put's data_parser instead of a new per-value one. It handles the review's points as follows:

• In-place mutation of stored values is a documented constraint of this advanced API.
• Replayed parsers: gets with a parser are not retried, matching put_data.
• Tuples over msgpack: there are no built-in parsers that depend on tuples anymore. parser_args documents that tuples arrive as lists.
• Non-encodable results: a parser result that fails to serialize now returns GET_ERROR. Before, the exception escaped and the unit's only worker stopped answering.
• Simplifications: each check runs once: parser_args against the parser and the sample count in the client, and every-key-found in async_kv_batch_get. There is no parsed reply flag, no exported parser type, and no built-in parsers. The put and get output checks share one helper, _apply_data_parser.

Compatibility

• kv_batch_get without data_parser behaves as before.
• TransferQueueClient.(async_)get_data passes data_parser / parser_args to StorageManager.get_data only when a parser is set, so custom managers written as get_data(self, metadata) still serve plain gets.

Changes

• transfer_queue/storage/simple_storage.py: _apply_data_parser shared by put and get (put's inline checks moved here unchanged), applied in _handle_get. Response serialization failures become error replies.
• transfer_queue/storage/managers/simple_storage_manager.py: routes parser_args to each unit by batch position. One attempt with a parser.
• transfer_queue/storage/managers/base.py: the new get_data arguments; KV backends reject data_parser.
• transfer_queue/interface.py, transfer_queue/client.py: the public parameters. The client validates parser_args; kv_batch_get requires every key to exist.
• README.md: StorageManager.get_data signature, and an "Advanced: parsing on get" paragraph.

Tests

• tests/test_simple_storage_unit.py: per-field parsing with per-key args leaves the stored values unchanged; a changed element count fails; an unencodable result returns GET_ERROR and the unit keeps serving; a parser returning a view sends only the view's bytes.
• tests/test_async_simple_storage_manager.py: each unit receives the parser_args of its own samples.
• tests/test_storage_request_retry.py: a get is retried only without a parser.
• tests/test_client.py: an old-style get_data(self, metadata) manager still serves plain gets, and parser_args without data_parser raises.
• tests/test_kv_storage_manager.py: KV backends raise before fetching.
• tests/e2e/test_kv_interface_e2e.py, sync and async: parsing across two storage units with per-key args, stored values unchanged afterwards, and the three parser_args validation errors.

pre-commit and python -m compileall -q transfer_queue tutorial tests pass. python -m pytest -q passes, with the Yuanrong e2e module excluded locally because of an environment-specific TORCH_LIBRARY registration error.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-13T13:13:45.410764Z	0058571	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0058571ee3

ℹ️ 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".

[chatgpt-codex-connector]

Preserve existing custom StorageManager gets

For a custom StorageManager implementing the documented get_data(self, metadata) extension contract, an ordinary get with no parser now still passes both new keywords and raises TypeError: unexpected keyword argument 'parser'. Such subclasses remain concrete and registerable because ABC does not check override signatures, and README.md still advertises the old signature, so only pass these keywords when parsing is requested or provide a compatibility path.

AGENTS.md reference: AGENTS.md:L17-L18

Useful? React with 👍 / 👎.

[OutstanderWang]

Option A: done in f15dd10 (woa cherry-pick 1b61e9b).

async_get_data only passes parser / parser_args when a parser is set, so custom managers with get_data(self, metadata) keep working for plain gets. README notes the optional kwargs; test_async_get_data_omits_parser_kwargs_for_old_managers covers the regression.

[OutstanderWang]

Still handled after 82ea7f9. async_get_data passes select_items only when it is set, so managers written as get_data(self, metadata) keep serving plain gets (test_async_get_data_passes_select_items_only_when_set).

[OutstanderWang]

Update: still handled after 61b4f9e, which replaced select_items with data_parser / parser_args. async_get_data passes them to StorageManager.get_data only when data_parser is set, so managers written as get_data(self, metadata) keep serving plain gets. The test is now test_async_get_data_passes_data_parser_only_when_set.

[chatgpt-codex-connector]

Catch parser response serialization failures

When a custom parser returns an object that neither msgpack nor pickle can serialize (for example, a generator or the coroutine produced by an accidentally async parser), this handler builds a success response containing it. _process_one_worker_request calls response_msg.serialize() after its exception guard, so the pickle fallback raises out of _worker_routine and permanently terminates the storage unit's sole worker thread instead of returning a parser error; response serialization failures need to be contained and converted to a basic error reply.

AGENTS.md reference: AGENTS.md:L17-L18

Useful? React with 👍 / 👎.

[OutstanderWang]

Option A: done in f15dd10 (woa cherry-pick 1b61e9b).

Response serialize() is now wrapped in try/except in the worker path: failures become GET_ERROR / PUT_GET_ERROR instead of killing the sole worker. Covered by test_unserializable_get_response_answered_and_worker_survives.

[OutstanderWang]

Superseded by 82ea7f9. The storage unit no longer runs user code: select_items only slices or indexes values that were already encodable. The response-serialization fallback could only be reached through parser output, so it was removed along with parsers.

[OutstanderWang]

Update: my earlier reply is outdated. 61b4f9e brings back user code on the unit (data_parser), so the containment is back as well. If response_msg.serialize() fails, the worker now replies GET_ERROR (or PUT_GET_ERROR for other operations) and keeps running. I checked that without it, the unit stops answering and the request times out. Covered by test_storage_unit_get_data_parser_parses_each_field_without_changing_stored_values: a parser returning generators gets GET_ERROR, and a plain get on the same unit still succeeds.

[chatgpt-codex-connector]

Preserve tuple indexes over the parser RPC

When the advertised multidimensional form is used via kv_batch_get, such as parser_args=[(slice(None), slice(0, 2))], the untyped msgpack transport encodes the tuple as an array and decodes it as a list. The storage unit therefore reaches the explicit list-rejection branch rather than this tuple branch, so tuple indexing works only in the direct unit tests and always fails over the actual RPC; preserve a tuple marker during serialization rather than conflating it with intentionally rejected list indexing.

AGENTS.md reference: AGENTS.md:L17-L18

Useful? React with 👍 / 👎.

[OutstanderWang]

Option A: done in f15dd10 (woa cherry-pick 1b61e9b).

slice_tensor now restores msgpack-decoded list-with-slice shapes back to an indexing tuple so multi-dim RPC parser_args still slice correctly, while pure int lists stay rejected. Covered by test_slice_tensor_restores_indexing_tuple_after_zmq_roundtrip.

[OutstanderWang]

Superseded by 82ea7f9. slice_tensor and _indexing_tuple_from_list are gone. select_items entries are a slice or a list of indices, and both round-trip through msgpack unchanged, so there's no tuple to restore.

[OutstanderWang]

Update for 61b4f9e: there are still no built-in parsers, so no tuple marker is needed. data_parser itself goes through cloudpickle. parser_args are encoded like other request data, so tuples arrive as lists; the kv_batch_get docstring says so.

[chatgpt-codex-connector]

Reject parser_args without a parser on KV backends

When the low-level client is used with Mooncake, Yuanrong, or Ray and supplies parser_args while leaving parser=None, the new argument reaches this method but the guard checks only parser, so the backend fetches and returns the full values while silently ignoring the arguments. The high-level and SimpleStorage paths explicitly reject this combination through validate_get_parser; apply the same validation here so a caller typo cannot look like a successful parsed read.

AGENTS.md reference: AGENTS.md:L17-L18

Useful? React with 👍 / 👎.

[OutstanderWang]

Option A: done in f15dd10 (woa cherry-pick 1b61e9b).

KVStorageManager.get_data now calls validate_get_parser before the NotImplemented check, so parser_args without a parser raise ValueError instead of being silently ignored. Covered by test_get_data_parser_args_without_parser_raises.

[OutstanderWang]

Superseded by 82ea7f9. There is now a single select_items argument, so nothing can be passed without its counterpart. KVStorageManager.get_data raises NotImplementedError for any non-None select_items before fetching (test_get_data_with_select_items_raises_before_fetching).

[OutstanderWang]

Update: my earlier reply is outdated. 61b4f9e reintroduced a separate parser_args, and this case came back with it. Fixed in e7bc18d: TransferQueueClient.(async_)get_data raises ValueError("parser_args requires data_parser"), and also checks that parser_args matches the sample count. It does this before calling the manager, so KV backends and SimpleStorage behave the same way. KVStorageManager.get_data still raises NotImplementedError for any data_parser. Covered in test_async_get_data_passes_data_parser_only_when_set and the KV test test_get_data_rejects_data_parser.

[0oshowero0]

[AI Review]

Thanks for this. The placement is right: trimming runs on the storage unit (data plane), the manager only shards parser_args by batch_positions, and the controller is untouched. Failing loudly on KV backends instead of fetching whole values is the right call, and the all-keys-found check is needed because kv_retrieve_meta drops missing keys. Three issues should be fixed before merging, and the implementation could be roughly half its current size.

Verification

• python -m compileall -q transfer_queue tutorial tests passes.
• pytest tests/test_get_parser.py tests/test_async_simple_storage_manager.py tests/test_client.py tests/test_kv_storage_manager.py tests/test_simple_storage_unit.py: 166 passed.
• pytest tests/e2e/test_kv_interface_e2e.py -k Parser on SimpleStorage: 16 passed. MooncakeStore / Yuanrong were not run.
• The issues below were reproduced against a real SimpleStorage unit (tq.init + kv_batch_put / kv_batch_get) and against HybridStorageUnitData directly.

Must fix

1. An in-place parser permanently corrupts stored data, and only for some samples.
   For memory-resident values, StorageUnitData.get_data returns the stored objects themselves, and apply_get_parser hands them straight to the parser, so any in-place operation is written back to the unit and seen by every later reader. Repro: store zeros(2, 4), then call kv_batch_get(..., parser=lambda v, n: (v.add_(1), v[:n])[1], parser_args=[2, 2]). The reply is the trimmed [[1, 1], [1, 1]], but a plain get afterwards returns [[1, 1, 1, 1], [1, 1, 1, 1]]: all four columns changed, not just the two that were returned.

   Hybrid storage makes this inconsistent rather than safe. Each sample lives either in memory or on SSD (decided per sample at put time by encoded size), and SSD samples are decoded into a fresh object on every read. In one field holding a 4-float sample (memory) and a 64-float sample (SSD), the same parser changed the first and left the second untouched, so tests that only use large SSD-backed values will not catch this. A failed request can also leave changes behind: a parser that mutates and then raises on the second key returned an error, but both keys had already been modified.

   In-place forms that are easy to miss: v.div_(255), v[...] = x, x = v[:n]; x += 1 (x is a view), v.pop(), v.sort(), random.shuffle(v).

   Options: restrict parser to built-in parsers, which removes this issue and item 2 entirely; or keep arbitrary callables and state in the kv_batch_get docstring that the parser must not modify value in place and must return new objects or views. A shallow list(value) before calling the parser is cheap and blocks list-structure mutations, but nothing cheap protects tensors. Which caller this targets matters here: Megatron-style load-balanced CP gives each rank two non-adjacent chunks, which int | slice | tuple cannot express, so built-in-only may need one more built-in.

2. Timeouts replay the parser, and the storage unit runs every attempt.
   Parser gets go through _request_with_retry (by default 3 attempts × TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = 200s). A timeout only discards the manager's socket; the unit's proxy has no deadline or cancellation, so every delivered attempt still runs to completion on the unit's single worker thread and its reply is dropped. Repro with the timeout set to 2s and a parser doing v.add_(10); time.sleep(1.5) over two keys: the caller got StorageUnitTimeout after 9s, and a plain get then read 31 (1 + 3 × 10). The parser ran three times even though the caller only saw a failure. With a pure parser the attempts still queue on the single worker and block every other put/get/clear on that unit; with the defaults, a parser that needs more than 200s for a batch holds the unit for at least 600s.

   put_data already handles the same situation with max_attempts = 1 if data_parser is not None ("the public API does not constrain parser side effects"). Parser gets should do the same; a replayed get-side side effect is worse, because it changes data that other consumers read.

3. slice_tensor indexing tuples behave differently over RPC.
   msgspec decodes tuples as lists, and _indexing_tuple_from_list restores only lists that contain a slice. (0, 1), (0,) and (Ellipsis, 0) work when calling slice_tensor locally but raise TypeError: slice_tensor does not take a list of indices on the unit through kv_batch_get; only tuples like (slice(None), 0) survive. test_slice_tensor_restores_indexing_tuple_after_zmq_roundtrip only covers the slice case. Suggest narrowing slice_tensor to a dim-0 int | slice, which covers prefix trimming and sequence-dim sharding, and dropping _indexing_tuple_from_list together with the list-rejection branch; multi-dim indexing can live in a custom parser.

Suggested

4. Require select_fields whenever parser is set. check_parser_select_fields allows omitting it only while the keys have a single field, so a call that works today starts failing once someone adds a field to the partition. Checking up front in validate_get_parser removes this function and the "omit only when single field" text repeated across four docstrings.

5. Make tensor-arg coercion parser-agnostic. Coercion is gated on parser is slice_tensor, so with functools.partial(slice_tensor) the tensor args reach the unit and fail there with TypeError. _coerce_int_like and _coerce_parser_args (~65 lines) can be replaced by converting tensor/ndarray parser_args for any parser, plus operator.index(index) inside slice_tensor (accepts numpy ints and 0-d integer tensors):

   if isinstance(parser_args, torch.Tensor | np.ndarray):
       parser_args = parser_args.tolist()

6. Validate once per layer. validate_get_parser runs in kv_batch_get, again in async_kv_batch_get, and again in both managers (SimpleStorage, KV); kv_batch_get also repeats the keys normalization. The async public entry plus the manager is enough, since TransferQueueClient.get_data is public too.

7. Naming and coverage. kv_batch_put already has data_parser, a batch-level dict -> dict hook that must keep lengths. Calling a per-value (value, arg) -> value hook with a different contract parser on the paired get API is confusing; a name that says select/trim would read better. Also, kv_batch_get_by_meta delegates to async_kv_batch_get but does not forward the parser.

8. Probably removable:

   • The parsed reply flag (manager, unit): no other request does version negotiation, and units and clients come from the same deployment. If version safety is wanted, it belongs in the handshake, once.
   • _unwrap_non_tensor: NonTensorStack fields written by kv_batch_put reach the unit as plain lists (checked by round-tripping a put body), so this branch and its test look unreachable.
   • _check_all_keys_found is used once; three inline lines read more directly.
   • Exporting GetParser at the top level adds little; it is just Callable[[Any, Any], Any].

9. Docs and comments. The "SSD may still read the full blob" caveat appears 9 times in the added lines, and the kv_batch_get parser doc runs ~12 lines. One copy in kv_batch_get, referenced elsewhere, is enough. NOTE(neowywang) author tags are not used elsewhere in the repo.

Reliability notes

Checked and fine: kv_retrieve_meta keeps key order, so with the all-keys check parser_args stay aligned (duplicate keys keep their positions); KV backends raise before any fetch; the serialize-failure containment keeps the single worker alive when a parser returns a non-encodable value, and since that path is only reachable through parsers it belongs in this PR. For pure parsers the SSD path is correct; the cost is that the full blob is still read and, until the reply is sent, the trimmed view keeps the whole decoded buffer alive.

[0oshowero0]

We can also align the put_data related parser in README

[OutstanderWang]

Done in a5c6205. The README now lists put_data(self, data, metadata, data_parser=None) and get_data(self, metadata, select_items=None), matching StorageManager.

[OutstanderWang]

Update for 61b4f9e: the README now lists get_data(self, metadata: BatchMeta, data_parser=None, parser_args=None), next to put_data(self, data, metadata, data_parser=None).

[0oshowero0]

[AI Review] Follow-up to the review above: one design point (how get_parser.py is organized, and whether this is really a get-specific capability) and the test volume.

get_parser.py mixes code that runs in three places

• Built-in parsers that run on the storage unit and are exported publicly: select_list_items, slice_tensor with _indexing_tuple_from_list, and GetParser (L39-L115).
• Request validation that runs in the client and manager: validate_get_parser, _coerce_*, check_parser_select_fields (L118-L216).
• The storage unit's execution loop: apply_get_parser and _unwrap_non_tensor (L219-L256).

apply_get_parser has a single caller, SimpleStorageUnit._handle_get. Apart from the call loop, it re-checks the arg count that the client and manager already validated, unwraps NonTensor* values that the put path never stores, and adds field/index context to errors. The sibling hooks are applied where they are used: data_parser inline in _handle_put and merge_fn in one line of apply_update. If callables stay, the loop belongs in _handle_get (or a private method of the unit), and the public helpers should not share a utils/ module with storage-unit internals.

The mechanism is not get-specific

Shrinking the reply is a read-side goal, but "run a user callable on stored values inside the storage unit" is a general storage-side hook. This PR adds a third such hook, with its own signature and a different retry policy:

Hook	Signature	Replayed on timeout
put data_parser	dict -> dict, whole batch, lengths unchanged	no (L501-L502)
update merge_fn	(old, new) -> value	no (L581-L583)
get parser (this PR)	(value, arg) -> value	yes (L812-L827)

What the capability really depends on is the backend, not get: only SimpleStorage units can run code next to the data, which is why KV backends have to raise.

Suggested direction

Which design fits depends on the actual need:

1. Return part of a value, which is what both built-ins do. Make it a declarative option of the get request, next to select_fields, which already projects columns: for example a per-key int | slice for tensor fields and an index list for list fields. Naming and exact shape are your call; load-balanced CP, where each rank takes two non-adjacent chunks, still fits as a short list of slices per key. The request then carries data instead of code:

   • it cannot mutate storage, and a replayed attempt is harmless;
   • the storage unit needs no cloudpickle or importable module;
   • plain int / slice / index lists round-trip through the existing encoder without the tuple-vs-list problem.

   Items 1, 2, 3 and 5 of the earlier review go away, get_parser.py shrinks to argument validation, and the storage-unit side becomes a few lines of slicing in _handle_get.

2. Run arbitrary user code next to the data. This is a general capability. It should share one contract with data_parser and merge_fn: where it runs, no in-place mutation, no replay, and how errors are wrapped. Get would be one call site, and the hook should not be named get_parser. This deserves its own design discussion.

3. The consumer layout is fixed (for example, a CP rank always reads the same shard). Split at put time into per-shard keys or fields; a plain get then fetches exactly what is needed, on every backend, with no storage-side compute. The cost is that producers must know the consumers' parallel layout.

For this PR, option 1 looks like the best fit. It covers exactly the two built-in parsers and is a stronger form of the "built-in parsers only" suggestion in item 1 of the earlier review. Arbitrary storage-side code can follow later, on top of an agreed contract, if a caller needs it.

Tests can shrink to under half

The PR adds about 630 test lines for about 460 source lines. tests/test_get_parser.py alone has 322 lines and 31 test functions, and the e2e class has 8 tests, each run in sync and async mode. For a feature that mostly adds an interface option, most of this volume tests machinery that the earlier items suggest removing.

Goes away with the code it tests:

• Tensor-arg coercion: 7 tests (keeps_plain_int_list_and_slice_args, coerces_scalar_length_for_slice_tensor with 5 cases, coerces_1d_lengths_to_list_int with 4 cases, rejects_bad_tensor_shapes_early, coerces_0d_tensor_as_whole_args_for_one_key, does_not_coerce_select_list_items_args, custom_callable_skips_length_coerce). A single test of the .tolist() conversion replaces them.
• Tuple indexing over RPC: restores_indexing_tuple_after_zmq_roundtrip, rejects_index_list_unlike_select_list_items, noncontiguous_slice_encodes_correct_values, and the tuple half of accepts_slice_and_tuple_indexes.
• The single-field select_fields rule: the two check_parser_select_fields tests, plus the e2e requires_select_fields_for_multi_field. Once the rule no longer depends on the partition's fields, this becomes one validation case that runs without Ray.
• The parsed flag: test_get_with_parser_rejects_reply_without_parsed_flag.
• _unwrap_non_tensor: test_apply_get_parser_unwraps_non_tensor_values.
• Repeated validation: the KV test_get_data_parser_args_without_parser_raises duplicates test_validate_get_parser_rejects_args_without_parser.

Duplicates regardless of design:

• e2e selects_video_frames_mixed_with_images runs the same path as selects_list_items_per_key: select_list_items over a list of tensors, with each row packed as a non-tensor list either way. empty_selection can be one more parser_args case of the latter.
• e2e out_of_range_index_raises and exception_reaches_caller both cover error propagation; keeping one is enough.
• prefix_returns_contiguous_view, empty_prefix_is_valid_view and apply_get_parser_slice_tensor_keeps_per_key_views all assert the same view/data_ptr property.
• Once the loop moves into _handle_get, the four apply_get_parser unit tests can become one storage-unit test through the existing MockStorageClient.send_get.

Worth keeping: the per-unit parser_args routing test in the manager, the KV NotImplementedError test, the e2e missing-key and tensor-trim tests, the serialize-failure test on the storage unit, the old-manager test in the client, and slice_tensor_prefix_encode_sends_only_sliced_nbytes (the only test that shows the bandwidth gain).

Missing:

• A no-replay test for parser gets. It can mirror test_put_retries_only_without_a_parser with one more parametrized case.
• If callables stay and a mitigation such as a shallow copy is adopted, a memory-path test showing that a parsed get leaves stored values unchanged. SSD-backed values cannot catch this, because they are decoded fresh on every read.

[OutstanderWang]

Thanks. I reproduced the three must-fix items on a real SimpleStorage setup before changing anything:

• an in-place parser persisted its change into storage, and a failing request still mutated both keys;
• with a 2s timeout, the parser ran 3 times and a plain get then read 31;
• (0, 1), (0,) and (Ellipsis, 0) failed over RPC.

I went with option 1 from the follow-up. In 82ea7f9, the request now carries data instead of code: kv_batch_get(..., select_fields=..., select_items=[...]). Each key gets a slice or a list of indices along dim 0, for tensor and list fields alike. A list of indices covers the two non-adjacent CP chunks.

Your review, item by item:

1. In-place mutation. Fixed: the unit only slices or indexes, so stored values can't change. Covered by test_get_select_items_returns_selected_items_and_keeps_stored_values.
2. Replay on timeout. A retried selection has no side effects, so gets keep the default retry policy.
3. Tuple indexes. Fixed: slices and int lists round-trip through msgpack unchanged, and _indexing_tuple_from_list is gone.
4. select_fields rule. It is now required whenever select_items is set, checked before any RPC.
5. Coercion. _normalize_select_items converts tensors, ndarrays and numpy ints for every entry, via .tolist() / operator.index. The ~65-line coerce helpers are gone.
6. Repeated validation. Validation happens once, in AsyncSimpleStorageManager.get_data. The interface only checks select_fields and that every key exists.
7. Naming and _by_meta. The parameter is now select_items, so it no longer clashes with data_parser. kv_batch_get_by_meta / async_kv_batch_get_by_meta forward it.
8. Removable pieces. Removed: the parsed flag, _unwrap_non_tensor, _check_all_keys_found (now three inline lines), and the GetParser / select_list_items / slice_tensor exports. With no user code on the unit, the response-serialization fallback is unreachable, so it went too.
9. Docs and comments. The SSD note appears once, in kv_batch_get; the other docstrings point there. The NOTE(neowywang) tags are gone.

On the follow-up: utils/get_parser.py is deleted. The unit side is _select_items plus five lines in _handle_get. Tests shrank from ~630 to ~200 lines and follow your keep list:

• per-unit routing;
• KV backends raise NotImplementedError before fetching;
• missing key, tensor trimming, and old-manager compatibility;
• one test that shows only the selected bytes are sent.

Against main, the PR is now +334/−25 across 11 files, down from +1093/−34. Running arbitrary code next to the data stays out of scope, for a separate design shared with data_parser and merge_fn. I've updated the PR description to match.

[0oshowero0]

[AI Review] Re-review of 47826f6, following up on your reply.

Thanks, this is a clear improvement. I checked every item in your reply against the code and they hold: the request now carries data instead of code, so in-place mutation, parser replay and the tuple-over-RPC problem are gone by construction, and keeping the default retry policy is right because a selection has no side effects. Validation runs once in the manager, select_fields is required up front, the _by_meta variants forward select_items, the removable pieces are gone, and the tests follow the keep list. One correctness issue remains; the rest are nits.

Verification

• python -m compileall -q transfer_queue tutorial tests passes.
• pytest tests/test_async_simple_storage_manager.py tests/test_client.py tests/test_kv_storage_manager.py tests/test_simple_storage_unit.py tests/test_storage_request_retry.py: 154 passed.
• pytest tests/e2e/test_kv_interface_e2e.py -k "SelectItems or ByMeta or KVGet" on SimpleStorage: 30 passed. MooncakeStore / Yuanrong were not run.
• The branch merges cleanly into current main, which now includes [chore] Report completed storage requests from the latency histogram #192.

Should fix

1. Boolean masks are silently read as indices. In _normalize_select_items, a bool tensor becomes Python bools through .tolist(), and operator.index(True) == 1, so a mask turns into a list of 0/1 indices. Repro on a real SimpleStorage unit: store x = [0, 1, 2, 3] under one key, then call kv_batch_get(..., select_fields="x", select_items=[torch.tensor([True, False, True, False])]). It returns [1, 0, 1, 0] without an error, while x[mask] is [0, 2]. Numpy bool arrays and plain Python bool lists behave the same way. Masks are a common way to select in PyTorch, so this is easy to hit. Rejecting bools is enough:

   elif isinstance(item, Sequence) and not isinstance(item, str | bytes):
       if any(isinstance(i, bool) for i in item):
           raise TypeError("select_items takes indices, not a boolean mask")
       normalized.append([operator.index(i) for i in item])

   Plus one ([[True, False], [0]], TypeError) case in test_normalize_select_items_rejects_bad_entries.

Nits

2. Slice steps behave differently for tensor and list fields. slice(None, None, -1) fails on the unit with step must be greater than zero for a tensor field but reverses a list field, and step=0 passes normalization and only fails on the unit. Checking step is None or step > 0 in the slice branch makes both field types behave the same and fails before any RPC.
3. _select_items changes the type of other values. Any value that is not a tensor goes through the list comprehension, so an ndarray with an index list comes back as a list of numpy scalars (with a slice it stays an ndarray), and a string comes back as a list of characters. Using the comprehension only for list values and value[items] for everything else keeps ndarrays native and makes unsupported types raise.
4. Test: test_selects_list_items_per_key zips with strict=False, so a reply missing rows would still pass. The three sequences have the same length, so strict=True works.

With item 1 fixed, this looks ready to merge from my side.

[OutstanderWang]

@0oshowero0 A design update. It supersedes my previous reply and makes most of the re-review moot.

Why a parser again. select_items applies one dim-0 selection to every selected field. TransferQueue rows often hold many columns, and each needs its own handling: trim one, pick items from another, decode a third. So 61b4f9e reverts select_items and adds kv_batch_get(..., data_parser=..., parser_args=...) as an advanced feature. It follows your option 2 and shares one contract with put's data_parser:

• it is the same dict -> dict call on the unit's share of the batch;
• the same helper, _apply_data_parser, checks keys and element counts on both put and get;
• it is never replayed;
• errors come back as error replies.

Your first review, item by item:

1. In-place mutation. This is now a documented constraint in the kv_batch_get docstring and the README: data holds the stored objects, so build new values instead of mutating them. A shallow copy would not protect tensors, so I didn't add one. test_storage_unit_get_data_parser_parses_each_field_without_changing_stored_values runs on memory-resident values and checks that a parser replacing entries leaves storage unchanged.
2. Replay. Gets with a parser use max_attempts=1, like put_data. Covered by test_get_retries_only_without_a_parser.
3. Tuples. There are no built-in parsers, so no index tuple has to survive msgpack. The docs say parser_args tuples arrive as lists.
4. select_fields. No longer required. The parser receives a dict and only touches the fields it names, so adding a column to the partition doesn't break an existing parser.
5. Coercion. None. parser_args reach the parser as given.
6. Validation once per layer.
   • The client checks parser_args against data_parser and the sample count (e7bc18d).
   • kv_batch_get checks that every key was found.
   • The unit checks the parser output.
7. Naming. The parameter is data_parser, the same name and contract as on put. The _by_meta variants don't take one in this PR.
8. Removable pieces. The parsed flag, _unwrap_non_tensor, _check_all_keys_found and the GetParser export stay removed. The response-serialization containment is back, because user code can reach it again.
9. Docs. The SSD note appears once, in kv_batch_get. There are no author tags.

Follow-up. utils/get_parser.py stays deleted. _handle_get calls the helper it now shares with _handle_put. Tests follow your keep list:

• per-unit parser_args routing;
• KV backends raise;
• missing key;
• serialize failure;
• old-style manager;
• a parsed view sends only its bytes.

They also add the two you listed as missing: a no-replay test, and the memory-path test that stored values stay unchanged.

Re-review. Boolean masks, slice steps, _select_items value types, and the strict=False zip all went away with select_items, because indexing now happens in the caller's parser. The new e2e test zips with strict=True.

Against main, the PR is +396/−68 across 12 files. The PR description is updated.