[feat] Support data_parser in kv_batch_get for SimpleStorage

README.md

@@ -72,8 +72,8 @@ In the data plane, we utilize a pluggable design, enabling TransferQueue to inte
 
 Specifically, we provide a `StorageManager` abstraction class that defines the core APIs as follows:
 
-- `async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None`
-- `async def get_data(self, metadata: BatchMeta) -> TensorDict`
+- `async def put_data(self, data: TensorDict, metadata: BatchMeta, data_parser=None) -> None`
+- `async def get_data(self, metadata: BatchMeta, data_parser=None, parser_args=None) -> TensorDict`
 - `async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) -> None`
 
 This class encapsulates the core interaction logic within the TransferQueue system. You only need to write a simple subclass to integrate your custom storage backend. Refer to the [Customize](#customize) section for details.
@@ -115,7 +115,7 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API
 
 - **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag.
 - **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches.
-- **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields).
+- **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields) and, on SimpleStorage, an optional `data_parser` run on the storage units.
 - **(async_)kv_update / (async_)kv_batch_update**: Merge new values into already-produced fields with `merge_fn(old, new)`. SimpleStorage only.
 - **(async_)kv_list**: List keys and tags (metadata) in a partition.
 - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
@@ -126,6 +126,18 @@ caller stops waiting, so non-idempotent merge operations must not be retried
 blindly. Updates are atomic within one storage-unit request; a batch spanning
 multiple units is not a distributed transaction.
 
+**Advanced: parsing on get.** Like `kv_batch_put`, `kv_batch_get` takes a
+`data_parser`, here run by each SimpleStorage unit before it replies, so only the
+parsed values cross the network (e.g. trimming padding or picking a few frames).
+It is called as `data_parser(data, args)`: `data` maps each selected field to the
+unit's stored values, one per sample, and `args` holds the matching entries of the
+optional per-key `parser_args`. Each field may be parsed differently, but the
+result must keep the keys and the number and order of values. `data` holds the
+stored objects themselves, so never modify them in place; build new values
+instead. As with put, a get with a parser is not retried on a storage-unit
+timeout; the call raises, and retrying it yourself may run the parser again, so
+keep it free of side effects. See the `kv_batch_get` docstring for details.
+
 **Key Features**
 
 - **Redis-style Semantics**: Familiar KV interface (Put/Get/List) for a zero learning curve.

tests/e2e/test_kv_interface_e2e.py

@@ -727,6 +727,42 @@ def test_kv_batch_get_nonexistent_key(self, controller, tq_api):
             # Or it might raise an error about keys not found
             assert "not found" in str(e).lower() or "empty" in str(e).lower()
 
+    def test_kv_batch_get_data_parser_parses_each_field_with_per_key_args(self, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("kv_batch_get data_parser is implemented only for SimpleStorage")
+
+        partition_id = "test_partition"
+        keys = ["parse_0", "parse_1", "parse_2"]
+        input_ids = torch.tensor([[1, 2, 0, 0], [3, 4, 5, 0], [6, 0, 0, 0]])
+        scores = torch.tensor([1.0, 2.0, 3.0])
+        tq_api.kv_batch_put(
+            keys=keys, partition_id=partition_id, fields=TensorDict({"input_ids": input_ids, "scores": scores}, 3)
+        )
+
+        def parse(data, args):
+            data["input_ids"] = [ids[:n] for ids, n in zip(data["input_ids"], args, strict=True)]
+            data["scores"] = [score * 10 for score in data["scores"]]
+            return data
+
+        retrieved = tq_api.kv_batch_get(keys=keys, partition_id=partition_id, data_parser=parse, parser_args=[2, 3, 1])
+        assert_nested_tensor_equal(retrieved["input_ids"], [input_ids[0, :2], input_ids[1, :3], input_ids[2, :1]])
+        assert_tensor_equal(retrieved["scores"], scores * 10)
+
+        retrieved = tq_api.kv_batch_get(keys=keys, partition_id=partition_id)
+        assert_tensor_equal(retrieved["input_ids"], input_ids)
+        assert_tensor_equal(retrieved["scores"], scores)
+
+        with pytest.raises(ValueError, match="2 entries for 3 samples"):
+            tq_api.kv_batch_get(keys=keys, partition_id=partition_id, data_parser=parse, parser_args=[2, 3])
+        with pytest.raises(ValueError, match="requires data_parser"):
+            tq_api.kv_batch_get(keys=keys, partition_id=partition_id, parser_args=[2, 3, 1])
+        with pytest.raises(ValueError, match="parser_args needs every key"):
+            tq_api.kv_batch_get(
+                keys=["parse_0", "absent"], partition_id=partition_id, data_parser=parse, parser_args=[1, 1]
+            )
+
+        tq_api.kv_clear(keys=keys, partition_id=partition_id)
+
 
 class TestKVBatchGetByMetaE2E:
     """End-to-end tests for kv_batch_get_by_meta functionality."""

tests/test_async_simple_storage_manager.py

@@ -286,6 +286,36 @@ async def fake_get(global_indexes, fields, target_storage_unit=None, **kwargs):
     assert set(called_with["storage_1"]) == {1, 3}
 
 
+@pytest.mark.asyncio
+async def test_get_data_sends_each_unit_the_parser_args_of_its_samples():
+    manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
+    manager.storage_manager_id = "test_get_parser"
+    manager.storage_unit_infos = {
+        su: ZMQServerInfo(role=Role.STORAGE, id=su, ip="127.0.0.1", ports={"put_get_socket": port})
+        for su, port in (("storage_0", 19012), ("storage_1", 19013))
+    }
+    batch_meta = BatchMeta(
+        global_indexes=[3, 0, 1, 2],
+        partition_ids=["p0"] * 4,
+        field_schema={"f": {"dtype": torch.float32, "shape": (2,), "is_nested": False, "is_non_tensor": False}},
+        production_status=np.ones(4, dtype=np.int8),
+    )
+    called_with: dict[str, tuple] = {}
+
+    async def fake_get(global_indexes, fields, target_storage_unit=None, data_parser=None, parser_args=None, **_):
+        called_with[target_storage_unit] = (list(global_indexes), data_parser, parser_args)
+        return fields, {"f": [torch.zeros(2) for _ in global_indexes]}
+
+    manager._get_from_single_storage_unit = fake_get
+
+    def parser(data, args):
+        return data
+
+    await manager.get_data(batch_meta, data_parser=parser, parser_args=["a", "b", "c", "d"])
+
+    assert called_with == {"storage_0": ([0, 2], parser, ["b", "d"]), "storage_1": ([3, 1], parser, ["a", "c"])}
+
+
 @pytest.mark.asyncio
 async def test_clear_data_routes_from_hash():
     """clear_data should route using global_idx % num_su (hash routing)."""

tests/test_client.py

@@ -23,6 +23,7 @@
 from tensordict import NonTensorStack, TensorDict
 
 from transfer_queue import TransferQueueClient
+from transfer_queue.client import AsyncTransferQueueClient
 from transfer_queue.metadata import BatchMeta
 from transfer_queue.utils.enum_utils import Role
 from transfer_queue.utils.zmq_utils import (
@@ -1472,3 +1473,27 @@ def test_load_storage_checkpoint_without_manager(self, mock_controller, tmp_path
 
         with pytest.raises(RuntimeError, match="Storage manager not initialized"):
             client.load_storage_checkpoint(ckpt_dir)
+
+
+@pytest.mark.asyncio
+async def test_async_get_data_passes_data_parser_only_when_set():
+    """Custom managers written as get_data(self, metadata) keep serving plain gets."""
+
+    class OldStyleManager:
+        async def get_data(self, metadata):
+            return TensorDict({"x": torch.zeros(metadata.size)}, batch_size=[metadata.size])
+
+    client = AsyncTransferQueueClient.__new__(AsyncTransferQueueClient)
+    client.client_id = "old_manager_client"
+    client.storage_manager = OldStyleManager()
+    metadata = BatchMeta(
+        global_indexes=[0],
+        partition_ids=["0"],
+        field_schema={"x": {"dtype": torch.float32, "shape": (1,), "is_nested": False, "is_non_tensor": False}},
+    )
+
+    assert "x" in await client.async_get_data(metadata)
+    with pytest.raises(TypeError):
+        await client.async_get_data(metadata, data_parser=lambda data, args: data)
+    with pytest.raises(ValueError, match="parser_args requires data_parser"):
+        await client.async_get_data(metadata, parser_args=[1])

tests/test_kv_storage_manager.py

@@ -516,3 +516,15 @@ def test_put_data_custom_backend_meta_length_mismatch_raises_error(test_data_for
         asyncio.run(manager.put_data(test_data_for_put_data["data"], test_data_for_put_data["metadata"]))
 
     assert "does not match" in str(exc_info.value)
+
+
+@patch.object(KVStorageManager, "_connect_to_controller", lambda self: None)
+def test_get_data_rejects_data_parser(test_data_for_put_data):
+    mock_storage_client = MagicMock()
+    config = {"controller_info": MagicMock(), "client_name": "MockClient"}
+    with patch(f"{STORAGE_CLIENT_FACTORY_PATH}.create", return_value=mock_storage_client):
+        manager = KVStorageManager(controller_info=MagicMock(), config=config)
+
+    with pytest.raises(NotImplementedError, match="data_parser is not supported"):
+        asyncio.run(manager.get_data(test_data_for_put_data["metadata"], data_parser=lambda data, args: data))
+    mock_storage_client.get.assert_not_called()

tests/test_simple_storage_unit.py

@@ -46,8 +46,10 @@
     SimpleStorageUnit,
     SSDFileStore,
     StorageKeyNotFoundError,
+    _apply_data_parser,
     _SSDValueRef,
 )
+from transfer_queue.utils.serial_utils import MsgpackDecoder, MsgpackEncoder
 from transfer_queue.utils.zmq_utils import (
     STORAGE_MANAGER_IDENTITY_PREFIX,
     ZMQMessage,
@@ -81,11 +83,14 @@ def send_put(self, client_id, global_indexes, field_data, data_parser=None):
         self.socket.send_multipart(msg.serialize())
         return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False))
 
-    def send_get(self, client_id, global_indexes, fields):
+    def send_get(self, client_id, global_indexes, fields, data_parser=None, parser_args=None):
+        body = {"global_indexes": global_indexes, "fields": fields}
+        if data_parser is not None:
+            body.update(data_parser=data_parser, parser_args=parser_args)
         msg = ZMQMessage.create(
             request_type=ZMQRequestType.GET_DATA,
             sender_id=f"mock_client_{client_id}",
-            body={"global_indexes": global_indexes, "fields": fields},
+            body=body,
         )
         self.socket.send_multipart(msg.serialize())
         return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False))
@@ -960,6 +965,59 @@ def __call__(self, field_data):
     client.close()
 
 
+def test_storage_unit_get_data_parser_parses_each_field_without_changing_stored_values(storage_setup):
+    _, put_get_address, storage_ip = storage_setup
+    client = MockStorageClient(put_get_address, storage_ip)
+    tokens = torch.arange(8).reshape(2, 4)
+    images = [["i0", "i1", "i2"], ["j0", "j1"]]
+    response = client.send_put(0, [0, 1], {"tokens": tokens, "images": images})
+    assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE, response.body
+
+    def pick(data, args):
+        data["tokens"] = [value[:n] for value, (n, _) in zip(data["tokens"], args, strict=True)]
+        data["images"] = [[value[i] for i in idx] for value, (_, idx) in zip(data["images"], args, strict=True)]
+        return data
+
+    response = client.send_get(0, [1, 0], ["tokens", "images"], data_parser=pick, parser_args=[(1, [1]), (3, [2, 0])])
+    assert response.request_type == ZMQRequestType.GET_DATA_RESPONSE, response.body
+    torch.testing.assert_close(response.body["data"]["tokens"], [tokens[1, :1], tokens[0, :3]])
+    assert response.body["data"]["images"] == [["j1"], ["i2", "i0"]]
+
+    response = client.send_get(0, [0, 1], ["tokens", "images"])
+    torch.testing.assert_close(response.body["data"]["tokens"], [tokens[0], tokens[1]])
+    assert response.body["data"]["images"] == images
+
+    def drop_one(data, _args):
+        data["tokens"] = data["tokens"][:-1]
+        return data
+
+    response = client.send_get(0, [0, 1], ["tokens"], data_parser=drop_one)
+    assert response.request_type == ZMQRequestType.GET_ERROR
+    assert "data_parser changed the number of elements" in response.body["message"]
+
+    def unencodable(data, _args):
+        data["tokens"] = [(x for x in ()) for _ in data["tokens"]]
+        return data
+
+    response = client.send_get(0, [0, 1], ["tokens"], data_parser=unencodable)
+    assert response.request_type == ZMQRequestType.GET_ERROR
+    assert "failed to serialize response" in response.body["message"]
+    response = client.send_get(0, [0], ["tokens"])
+    torch.testing.assert_close(response.body["data"]["tokens"], [tokens[0]])
+
+    client.close()
+
+
+def test_get_data_parser_view_sends_only_the_parsed_bytes():
+    value = torch.arange(12, dtype=torch.int64).view(3, 4)
+    parsed = _apply_data_parser({"x": [value]}, lambda data, _args: {"x": [v[:2] for v in data["x"]]}, None)
+
+    buffers = MsgpackEncoder().encode(parsed)
+    assert parsed["x"][0].data_ptr() == value.data_ptr()
+    assert buffers[1].nbytes == value[:2].nbytes < value.nbytes
+    assert torch.equal(MsgpackDecoder().decode(buffers)["x"][0], value[:2])
+
+
 def test_storage_unit_data_parser_validation(storage_setup):
     """Test that invalid data_parser inputs produce clear error messages."""
     _, put_get_address, storage_ip = storage_setup

tests/test_storage_request_retry.py

@@ -107,6 +107,23 @@ async def test_put_retries_only_without_a_parser(data_parser, expected_attempts)
     manager.notify_data_update.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("data_parser, expected_attempts", [(None, 3), (lambda data, args: data, 1)])
+async def test_get_retries_only_without_a_parser(data_parser, expected_attempts):
+    manager = _manager()
+    manager._get_from_single_storage_unit = AsyncMock(side_effect=StorageUnitTimeout("no answer"))
+    _, metadata = _single_sample_batch()
+
+    with (
+        patch.object(ssm, "TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3),
+        patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"),
+        pytest.raises(StorageUnitTimeout),
+    ):
+        await manager.get_data(metadata, data_parser=data_parser)
+
+    assert manager._get_from_single_storage_unit.await_count == expected_attempts
+
+
 @pytest.mark.asyncio
 async def test_update_reports_controller_rejection_after_storage_commit():
     manager = _manager()

transfer_queue/client.py

@@ -522,11 +522,19 @@ async def async_update(
         field_schema = await self.storage_manager.update_data(metadata, values, merge_fn)
         return metadata.apply_field_schema(field_schema)
 
-    async def async_get_data(self, metadata: BatchMeta) -> TensorDict:
+    async def async_get_data(
+        self,
+        metadata: BatchMeta,
+        data_parser: Callable[[Any, Any], Any] | None = None,
+        parser_args: list | None = None,
+    ) -> TensorDict:
         """Asynchronously fetch data from storage units and organize into TensorDict.
 
         Args:
             metadata: Batch metadata containing data location information and global indexes
+            data_parser: Optional ``data_parser(data, args)`` the storage units apply before replying.
+                See ``transfer_queue.kv_batch_get``. Only supported by SimpleStorage.
+            parser_args: Optional per-sample arguments aligned with ``metadata.global_indexes``.
 
         Returns:
             TensorDict containing:
@@ -556,11 +564,22 @@ async def async_get_data(self, metadata: BatchMeta) -> TensorDict:
         if not metadata or metadata.size == 0 or len(metadata.field_names) == 0:
             logger.warning(f"[{self.client_id}]: Empty BatchMeta provided to get_data. Returning empty TensorDict.")
             return TensorDict({}, batch_size=0)
+        if parser_args is not None:
+            if data_parser is None:
+                raise ValueError("parser_args requires data_parser")
+            if len(parser_args) != metadata.size:
+                raise ValueError(f"parser_args has {len(parser_args)} entries for {metadata.size} samples")
 
         with limit_pytorch_auto_parallel_threads(
             target_num_threads=TQ_NUM_THREADS, info=f"[{self.client_id}] async_get_data"
         ):
-            results = await self.storage_manager.get_data(metadata)
+            # Passed only when set, so custom managers with get_data(self, metadata) keep working.
+            if data_parser is None:
+                results = await self.storage_manager.get_data(metadata)
+            else:
+                results = await self.storage_manager.get_data(
+                    metadata, data_parser=data_parser, parser_args=parser_args
+                )
 
         logger.debug(f"[{self.client_id}]: get_data with {metadata.size} samples successfully.")
 
@@ -1645,11 +1664,19 @@ def update(
         """
         return self._update(metadata, values, merge_fn)
 
-    def get_data(self, metadata: BatchMeta) -> TensorDict:
+    def get_data(
+        self,
+        metadata: BatchMeta,
+        data_parser: Callable[[Any, Any], Any] | None = None,
+        parser_args: list | None = None,
+    ) -> TensorDict:
         """Synchronously fetch data from storage units and organize into TensorDict.
 
         Args:
             metadata: Batch metadata containing data location information and global indexes
+            data_parser: Optional ``data_parser(data, args)`` the storage units apply before replying.
+                See ``transfer_queue.kv_batch_get``. Only supported by SimpleStorage.
+            parser_args: Optional per-sample arguments aligned with ``metadata.global_indexes``.
 
         Returns:
             TensorDict containing:
@@ -1667,7 +1694,7 @@ def get_data(self, metadata: BatchMeta) -> TensorDict:
             >>> print(batch)
             >>> # TensorDict with fields "prompts", "attention_mask", and sample order matching metadata global_indexes
         """
-        return self._get_data(metadata=metadata)
+        return self._get_data(metadata=metadata, data_parser=data_parser, parser_args=parser_args)
 
     def clear_partition(self, partition_id: str):
         """Synchronously clear the whole partition from all storage units and the controller.