[feat] Support user-defined data parser for SimpleStorage backend (Ascend#82)

## Background

Users often need to store lightweight **references** (e.g., URLs, file
paths, etc.) rather than the full data into TransferQueue to avoid the
expensive loading and decoding processes happening within user code,
which reduces a data copy.

<img width="1032" height="745" alt="image"
src="https://github.com/user-attachments/assets/4a62b350-b6f1-441d-8c05-b35f8cf2e7de"
/>


## Solution

This PR introduces support for a user-defined data parser in the
`SimpleStorage` backend.

The `kv_put` and `kv_batch_put` methods now accept an optional
`data_parser` callable. This parser is executed **inside** each
`SimpleStorageUnit` at put time. It receives the raw `field_data`
dictionary during the put request and should return a dictionary with
the same structure, replacing reference values with the actual parsed
data.

## Limitations & Future Work

- **Synchronous Execution:** In the current design, the data parser
execution is synchronous and part of the `put` operation. This means the
put request is only completed when the data parser finishes execution.
- **Backend Support:** `data_parser` is currently only supported by the
**SimpleStorage** backend.
- **Incorrect Metadata:** Allowing user-provided functions to modify
data in `SimpleStorageUnit` may lead to incorrect `shape` & `dtype`
metadata, which is captured when the data is still in
`TransferQueueClient`. This can lead to problems for RDMA transport,
which leverages these metadata collected by TQ to restore tensor during
`get`.





## Demo Script
```python3
"""Demo: concurrent data_parser with separated single-sample logic.

This demo shows how to structure a data_parser so that:
1. The **core parser** only handles a **single sample**.
2. The **batch wrapper** uses asyncio to process all samples in parallel.
3. The wrapper is **synchronous to the outside**: it blocks until every
   sample finishes, so ``data_parser`` returning means data is ready.

Scenario:
- Users pass URL-like strings in a column.
- The parser sleeps 1 s per sample (simulating I/O / decode) and then
  creates a random tensor of the requested dtype & shape.
- Because the sleeps run concurrently via asyncio, a batch of N samples
  finishes in ~1 s instead of ~N s.
"""

import asyncio
import time

import ray
import torch
from tensordict import TensorDict, NonTensorStack

import transfer_queue as tq


# ---------------------------------------------------------------------------
# Core single-sample parser
# ---------------------------------------------------------------------------
def parse_url(url: str) -> torch.Tensor:
    """Parse a URL-like descriptor 'dtype:HxW' into a random tensor."""
    dtype_str, shape_str = url.split(":")
    dtype = getattr(torch, dtype_str)
    shape = [int(dim) for dim in shape_str.split("x")]
    return torch.randn(shape, dtype=dtype)


# ---------------------------------------------------------------------------
# Batch-level parser
# ---------------------------------------------------------------------------
def concurrent_batch_url_parser(field_data: dict) -> dict:
    """Batch-level data_parser executed inside SimpleStorageUnit.

    It receives a ``dict`` (not a TensorDict) where each value is a
    batched column.  For columns created from ``NonTensorStack`` the
    value is a plain ``list`` of Python objects.

    Workflow:
    1. Spawns one async task per list element.
    2. Waits until *all* tasks finish (``asyncio.gather``).
    3. Replaces the list with the list of results.

    Because ``asyncio.run`` blocks until the loop finishes, this function
    is **synchronous** to its caller: when it returns, every sample has
    been processed.

    Args:
        field_data: Mapping ``field_name -> batched_values``.  The dict
            keys must stay exactly the same; only values may be
            transformed in-place.

    Returns:
        The same dict with parsed values substituted.
    """
    if "data_to_be_parsed" not in field_data:
        return field_data

    urls:list[str] = field_data["data_to_be_parsed"]

    async def _async_parse_single(url: str) -> torch.Tensor:
        await asyncio.sleep(1.0)  # Add fixed delay per sample
        return parse_url(url)

    async def _process_all():
        tasks = [asyncio.create_task(_async_parse_single(url)) for url in urls]
        return await asyncio.gather(*tasks)

    start = time.perf_counter()
    field_data["data_to_be_parsed"] = asyncio.run(_process_all())
    elapsed = time.perf_counter() - start

    print(
        f"[data_parser] Processed {len(urls)} samples in {elapsed:.2f}s "
        f"(serial would be ~{len(urls)}.0s)"
    )
    return field_data


# ---------------------------------------------------------------------------
# Main demo flow
# ---------------------------------------------------------------------------
def main():
    ray.init(ignore_reinit_error=True)
    try:
        tq.init()

        batch_size = 32

        # Column that stays untouched
        normal_data = torch.randn(batch_size, 2)

        # Column to be parsed: URL-like strings describing dtype & shape.
        shapes = [(i % 4 + 1, i % 3 + 2) for i in range(batch_size)]
        urls = [f"float32:{h}x{w}" for h, w in shapes]
        data_to_be_parsed = NonTensorStack(*urls)

        data = TensorDict({
            "normal_data": normal_data,
            "data_to_be_parsed": data_to_be_parsed,
        }, batch_size=batch_size)

        keys = [f"sample_{i}" for i in range(batch_size)]

        # -------------------------------------------------------------------
        # Put with data_parser
        # -------------------------------------------------------------------
        put_start_time = time.perf_counter()
        meta = tq.kv_batch_put(
            keys=keys,
            partition_id="train",
            fields=data,
            data_parser=concurrent_batch_url_parser,
        )
        put_elapsed = time.perf_counter() - put_start_time
        print(f"Put succeeded. Fields: {meta.fields}")
        print(
            f"Total kv_batch_put time: {put_elapsed:.2f}s "
            f"(concurrency keeps it ~1s, not {batch_size}s)\n"
        )

        # -------------------------------------------------------------------
        # Fetch back and verify
        # -------------------------------------------------------------------
        result = tq.kv_batch_get(keys=keys, partition_id="train")

        # 1) normal_data unchanged
        torch.testing.assert_close(result["normal_data"], normal_data)
        print("[PASS] normal_data is unchanged.")

        # 2) Parsed tensors have correct dtype & shape
        expected_shapes = [(i % 4 + 1, i % 3 + 2) for i in range(batch_size)]
        for i, exp_shape in enumerate(expected_shapes):
            tensor = result["data_to_be_parsed"][i]
            assert tensor.dtype == torch.float32, (
                f"dtype mismatch at index {i}: expected torch.float32, got {tensor.dtype}"
            )
            assert tuple(tensor.shape) == exp_shape, (
                f"shape mismatch at index {i}: expected {exp_shape}, got {tuple(tensor.shape)}"
            )
        print(f"[PASS] All {batch_size} parsed tensors have correct dtype & shape.")

        # 3) Timing sanity check
        #    Serial execution would be ~batch_size seconds.
        #    Because asyncio tasks run in parallel, it should be ~1 s.
        #    We allow generous headroom for TQ network / serialization overhead.
        assert put_elapsed < 2.0, (
            f"Expected concurrent execution (~1s), but took {put_elapsed:.2f}s. "
            "Are the asyncio tasks actually running in parallel?"
        )
        print(f"[PASS] Timing looks concurrent: {put_elapsed:.2f}s < 2.0s")

        print("\n=== All verifications passed! ===")

        # wait for Ray log collect
        time.sleep(2)

    except Exception as e:
        print(f"Error: {type(e).__name__}: {e}")
        import traceback

        traceback.print_exc()
    finally:
        tq.close()
        ray.shutdown()


if __name__ == "__main__":
    main()


```

---------

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

tests/test_client.py

@@ -418,7 +418,7 @@ def client_setup(mock_controller, mock_storage):
         client.initialize_storage_manager(manager_type="SimpleStorage", config=config)
 
         # Mock all storage manager methods to avoid real ZMQ operations
-        async def mock_put_data(data, metadata):
+        async def mock_put_data(data, metadata, data_parser=None):
             pass  # Just pretend to store the data
 
         async def mock_get_data(metadata):
@@ -511,7 +511,7 @@ def test_single_controller_multiple_storages():
             client.initialize_storage_manager(manager_type="SimpleStorage", config=config)
 
             # Mock all storage manager methods to avoid real ZMQ operations
-            async def mock_put_data(data, metadata):
+            async def mock_put_data(data, metadata, data_parser=None):
                 pass  # Just pretend to store the data
 
             async def mock_get_data(metadata):

tests/test_simple_storage_unit.py

@@ -34,11 +34,14 @@ def __init__(self, storage_put_get_address):
         self.socket.setsockopt(zmq.RCVTIMEO, 5000)  # 5 second timeout
         self.socket.connect(storage_put_get_address)
 
-    def send_put(self, client_id, global_indexes, field_data):
+    def send_put(self, client_id, global_indexes, field_data, data_parser=None):
+        body = {"global_indexes": global_indexes, "data": field_data}
+        if data_parser is not None:
+            body["data_parser"] = data_parser
         msg = ZMQMessage.create(
             request_type=ZMQRequestType.PUT_DATA,
             sender_id=f"mock_client_{client_id}",
-            body={"global_indexes": global_indexes, "data": field_data},
+            body=body,
         )
         self.socket.send_multipart(msg.serialize())
         return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False))
@@ -434,3 +437,177 @@ def test_storage_unit_data_capacity_uses_active_keys():
     assert len(storage._active_keys) == 2
     storage.put_data({"f": [4]}, global_indexes=[3])
     assert storage._active_keys == {0, 1, 3}
+
+
+def test_storage_unit_data_parser(storage_setup):
+    """Test data_parser functionality in SimpleStorageUnit.
+
+    Writes two columns:
+    - normal_data: regular tensors, should remain unchanged
+    - data_to_be_parsed: list of shape descriptors (list of ints)
+
+    data_parser converts shape descriptors into random tensors of those shapes.
+    """
+    _, put_get_address = storage_setup
+    client = MockStorageClient(put_get_address)
+
+    def create_data_by_shape_parser(field_data):
+        if "data_to_be_parsed" in field_data:
+            shapes = field_data["data_to_be_parsed"]
+            field_data["data_to_be_parsed"] = [torch.randn(shape) for shape in shapes]
+        return field_data
+
+    # Prepare data: normal_data is a batch tensor, data_to_be_parsed is a list of shape lists
+    field_data = {
+        "normal_data": torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]),
+        "data_to_be_parsed": [[2, 3], [1, 4], [3, 2]],
+    }
+    global_indexes = [0, 1, 2]
+
+    # Put with data_parser
+    response = client.send_put(0, global_indexes, field_data, data_parser=create_data_by_shape_parser)
+    assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE, f"Put failed: {response.body}"
+
+    # Get back
+    response = client.send_get(0, global_indexes, ["normal_data", "data_to_be_parsed"])
+    assert response.request_type == ZMQRequestType.GET_DATA_RESPONSE
+
+    result = response.body["data"]
+
+    # Verify normal_data is unchanged
+    torch.testing.assert_close(result["normal_data"][0], torch.tensor([1.0, 2.0]))
+    torch.testing.assert_close(result["normal_data"][1], torch.tensor([3.0, 4.0]))
+    torch.testing.assert_close(result["normal_data"][2], torch.tensor([5.0, 6.0]))
+
+    # Verify data_to_be_parsed shapes match the input shape descriptors
+    expected_shapes = [(2, 3), (1, 4), (3, 2)]
+    for i, expected_shape in enumerate(expected_shapes):
+        actual_shape = tuple(result["data_to_be_parsed"][i].shape)
+        assert actual_shape == expected_shape, (
+            f"Shape mismatch at index {i}: expected {expected_shape}, got {actual_shape}"
+        )
+
+    client.close()
+
+
+def test_storage_unit_data_parser_callable_types(storage_setup):
+    """Test that various callable types (partial, callable class) work as data_parser."""
+    _, put_get_address = storage_setup
+    client = MockStorageClient(put_get_address)
+
+    from functools import partial
+
+    # 1. Test functools.partial
+    def _partial_parser(field_data, prefix):
+        if "text" in field_data:
+            field_data["text"] = [f"{prefix}{t}" for t in field_data["text"]]
+        return field_data
+
+    partial_parser = partial(_partial_parser, prefix="parsed_")
+
+    response = client.send_put(
+        0,
+        [0, 1],
+        {"text": ["a", "b"]},
+        data_parser=partial_parser,
+    )
+    assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE, f"partial parser failed: {response.body}"
+
+    response = client.send_get(0, [0, 1], ["text"])
+    assert response.request_type == ZMQRequestType.GET_DATA_RESPONSE
+    assert response.body["data"]["text"] == ["parsed_a", "parsed_b"]
+
+    # 2. Test callable class instance
+    class CallableParser:
+        def __call__(self, field_data):
+            if "value" in field_data:
+                field_data["value"] = [v * 2 for v in field_data["value"]]
+            return field_data
+
+    callable_parser = CallableParser()
+    response = client.send_put(
+        0,
+        [2, 3],
+        {"value": [1, 2]},
+        data_parser=callable_parser,
+    )
+    assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE, f"callable class parser failed: {response.body}"
+
+    response = client.send_get(0, [2, 3], ["value"])
+    assert response.request_type == ZMQRequestType.GET_DATA_RESPONSE
+    assert response.body["data"]["value"] == [2, 4]
+
+    client.close()
+
+
+def test_storage_unit_data_parser_validation(storage_setup):
+    """Test that invalid data_parser inputs produce clear error messages."""
+    _, put_get_address = storage_setup
+    client = MockStorageClient(put_get_address)
+
+    # 1. Non-callable data_parser should return a clear TypeError
+    response = client.send_put(
+        0,
+        [0],
+        {"data": [1]},
+        data_parser="not_callable",
+    )
+    assert response.request_type == ZMQRequestType.PUT_ERROR
+    assert "data_parser must be callable" in response.body["message"]
+
+    # 2. data_parser returning non-dict should return a clear TypeError
+    def bad_parser(field_data):
+        return "not_a_dict"
+
+    response = client.send_put(
+        0,
+        [1],
+        {"data": [1]},
+        data_parser=bad_parser,
+    )
+    assert response.request_type == ZMQRequestType.PUT_ERROR
+    assert "data_parser must return a dict" in response.body["message"]
+
+    # 3. data_parser deleting a key should return a clear ValueError
+    def delete_key_parser(field_data):
+        del field_data["data"]
+        return field_data
+
+    response = client.send_put(
+        0,
+        [2],
+        {"data": [1], "extra": [2]},
+        data_parser=delete_key_parser,
+    )
+    assert response.request_type == ZMQRequestType.PUT_ERROR
+    assert "data_parser must not change dict keys" in response.body["message"]
+
+    # 4. data_parser adding a key should return a clear ValueError
+    def add_key_parser(field_data):
+        field_data["new_key"] = [999]
+        return field_data
+
+    response = client.send_put(
+        0,
+        [3],
+        {"data": [1]},
+        data_parser=add_key_parser,
+    )
+    assert response.request_type == ZMQRequestType.PUT_ERROR
+    assert "data_parser must not change dict keys" in response.body["message"]
+
+    # 5. data_parser changing element count should return a clear ValueError
+    def wrong_len_parser(field_data):
+        field_data["data"] = field_data["data"][:-1]
+        return field_data
+
+    response = client.send_put(
+        0,
+        [4, 5],
+        {"data": [1, 2]},
+        data_parser=wrong_len_parser,
+    )
+    assert response.request_type == ZMQRequestType.PUT_ERROR
+    assert "data_parser changed the number of elements" in response.body["message"]
+
+    client.close()

transfer_queue/client.py

@@ -324,6 +324,7 @@ async def async_put(
         data: TensorDict,
         metadata: Optional[BatchMeta] = None,
         partition_id: Optional[str] = None,
+        data_parser: Optional[Callable[[Any], Any]] = None,
     ) -> BatchMeta:
         """Asynchronously write data to storage units based on metadata.
 
@@ -342,6 +343,16 @@ async def async_put(
             metadata: Records the metadata of a batch of data samples, containing index and
                       storage unit information. If None, metadata will be auto-generated.
             partition_id: Target data partition id (required if metadata is not provided)
+            data_parser: Optional callable to parse reference data (e.g., URLs) into real
+                         content. The input is a slice of the `data` parameter, in plain
+                         dict format (not TensorDict), mapping field_name -> batched values.
+                         For a regular tensor column the value is a batched tensor; for
+                         nested tensors (jagged or strided) and NonTensorStack columns
+                         the values are extracted into a list. It must modify values
+                         in-place based on the original keys; do not add or remove keys.
+                         The number of elements per column must also remain unchanged.
+                         Do not change the inner order of values within each column.
+                         Only supported by SimpleStorage.
 
         Returns:
             BatchMeta: The metadata used for the put operation (currently returns the input metadata or auto-retrieved
@@ -411,7 +422,7 @@ async def async_put(
         with limit_pytorch_auto_parallel_threads(
             target_num_threads=TQ_NUM_THREADS, info=f"[{self.client_id}] async_put"
         ):
-            await self.storage_manager.put_data(data, metadata)
+            await self.storage_manager.put_data(data, metadata, data_parser=data_parser)
 
         await self.async_set_custom_meta(metadata)
 
@@ -1279,7 +1290,11 @@ def set_custom_meta(self, metadata: BatchMeta) -> None:
         return self._set_custom_meta(metadata=metadata)
 
     def put(
-        self, data: TensorDict, metadata: Optional[BatchMeta] = None, partition_id: Optional[str] = None
+        self,
+        data: TensorDict,
+        metadata: Optional[BatchMeta] = None,
+        partition_id: Optional[str] = None,
+        data_parser: Optional[Callable[[Any], Any]] = None,
     ) -> BatchMeta:
         """Synchronously write data to storage units based on metadata.
 
@@ -1298,6 +1313,16 @@ def put(
             metadata: Records the metadata of a batch of data samples, containing index and
                       storage unit information. If None, metadata will be auto-generated.
             partition_id: Target data partition id (required if metadata is not provided)
+            data_parser: Optional callable to parse reference data (e.g., URLs) into real
+                         content. The input is a slice of the `data` parameter, in plain
+                         dict format (not TensorDict), mapping field_name -> batched values.
+                         For a regular tensor column the value is a batched tensor; for
+                         nested tensors (jagged or strided) and NonTensorStack columns
+                         the values are extracted into a list. It must modify values
+                         in-place based on the original keys; do not add or remove keys.
+                         The number of elements per column must also remain unchanged.
+                         Do not change the inner order of values within each column.
+                         Only supported by SimpleStorage.
 
         Returns:
             BatchMeta: The metadata used for the put operation (currently returns the input metadata or auto-retrieved
@@ -1336,7 +1361,7 @@ def put(
             >>> # This will create metadata in "insert" mode internally.
             >>> metadata = client.put(data=prompts_repeated_batch, partition_id=current_partition_id)
         """
-        return self._put(data=data, metadata=metadata, partition_id=partition_id)
+        return self._put(data=data, metadata=metadata, partition_id=partition_id, data_parser=data_parser)
 
     def get_data(self, metadata: BatchMeta) -> TensorDict:
         """Synchronously fetch data from storage units and organize into TensorDict.