TransferQueue 0.1.1.dev1__py3-none-any.whl → 0.1.2.dev0__py3-none-any.whl

transfer_queue/metadata.py

@@ -14,15 +14,22 @@
 
 import dataclasses
 import itertools
-from collections import ChainMap
+import logging
+import os
+from collections import defaultdict
 from dataclasses import dataclass
 from typing import Any, Optional
 
 import numpy as np
+import torch
 from tensordict import TensorDict
+from tensordict.tensorclass import NonTensorData, NonTensorStack
 
 from transfer_queue.utils.utils import ProductionStatus
 
+logger = logging.getLogger(__name__)
+logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
+
 
 @dataclass
 class FieldMeta:
@@ -296,8 +303,33 @@ class BatchMeta:
 
         # Combine all samples
         all_samples = list(itertools.chain.from_iterable(chunk.samples for chunk in data))
+
         # Merge all extra_info dictionaries from the chunks
-        merged_extra_info = dict(ChainMap(*(chunk.extra_info for chunk in data)))
+        merged_extra_info = dict()
+
+        values_by_key = defaultdict(list)
+        for chunk in data:
+            for key, value in chunk.extra_info.items():
+                values_by_key[key].append(value)
+        for key, values in values_by_key.items():
+            if all(isinstance(v, torch.Tensor) for v in values):
+                try:
+                    if all(v.dim() == 0 for v in values):
+                        merged_extra_info[key] = torch.cat([v.unsqueeze(0) for v in values], dim=0)
+                    else:
+                        merged_extra_info[key] = torch.cat(values, dim=0)
+                except RuntimeError as e:
+                    logger.warning(
+                        f"BatchMeta.concat try to use torch.cat(dim=0) to merge extra_info key '{key}'"
+                        f" fails, with RuntimeError {e}. Falling back to use list."
+                    )
+                    merged_extra_info[key] = values
+            elif all(isinstance(v, NonTensorStack | NonTensorData) for v in values):
+                merged_extra_info[key] = torch.stack(values)
+            elif all(isinstance(v, list) for v in values):
+                merged_extra_info[key] = list(itertools.chain.from_iterable(values))
+            else:
+                merged_extra_info[key] = values[-1]
 
         return BatchMeta(samples=all_samples, extra_info=merged_extra_info)
 

transfer_queue/storage/managers/simple_backend_manager.py

@@ -34,6 +34,9 @@ from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServer
 logger = logging.getLogger(__name__)
 logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
 
+TQ_SIMPLE_STORAGE_MANAGER_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_MANAGER_RECV_TIMEOUT", 200))  # seconds
+TQ_SIMPLE_STORAGE_MANAGER_SEND_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_MANAGER_SEND_TIMEOUT", 200))  # seconds
+
 
 @TransferQueueStorageManagerFactory.register("AsyncSimpleStorageManager")
 class AsyncSimpleStorageManager(TransferQueueStorageManager):
@@ -132,8 +135,8 @@ class AsyncSimpleStorageManager(TransferQueueStorageManager):
                 try:
                     sock.connect(address)
                     # Timeouts to avoid indefinite await on recv/send
-                    sock.setsockopt(zmq.RCVTIMEO, 10_000)  # 10s
-                    sock.setsockopt(zmq.SNDTIMEO, 10_000)  # 10s
+                    sock.setsockopt(zmq.RCVTIMEO, TQ_SIMPLE_STORAGE_MANAGER_RECV_TIMEOUT * 1000)
+                    sock.setsockopt(zmq.SNDTIMEO, TQ_SIMPLE_STORAGE_MANAGER_SEND_TIMEOUT * 1000)
                     logger.info(
                         f"[{self.storage_manager_id}]: Connected to StorageUnit {server_info.id} at {address} "
                         f"with identity {identity.decode()}"

transfer_queue/version/version

@@ -1 +1 @@
-0.1.1.dev1
+0.1.2.dev0

{transferqueue-0.1.1.dev1.dist-info → transferqueue-0.1.2.dev0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: TransferQueue
-Version: 0.1.1.dev1
+Version: 0.1.2.dev0
 Summary: TransferQueue: An Asynchronous Streaming Data Management Module
 Author-email: The TransferQueue Team <hanzy19@tsinghua.org.cn>
 License: Apache-2.0
@@ -40,13 +40,13 @@ Dynamic: license-file
 TransferQueue is a high-performance data storage and transfer module with panoramic data visibility and streaming scheduling capabilities, optimized for efficient dataflow in post-training workflows.
 
 <p align="center">
-  <img src="https://cdn.nlark.com/yuque/0/2025/png/23208217/1761356010763-b05751d3-f975-4890-ba59-c8d753cf95f2.png" width="70%">
+  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/tq_arch.png?raw=true" width="70%">
 </p>
 
 TransferQueue offers **fine-grained, sample-level** data management and **load-balancing** (on the way) capabilities, serving as a data gateway that decouples explicit data dependencies across computational tasks. This enables a divide-and-conquer approach, significantly simplifies the algorithm controller design.
 
 <p align="center">
-  <img src="https://cdn.nlark.com/yuque/0/2025/png/23208217/1758696791245-fa7baf96-46af-4c19-8606-28ffadc4556c.png" width="70%">
+  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/main_func.png?raw=true" width="70%">
 </p>
 
 <h2 id="updates">🔄 Updates</h2>
@@ -69,7 +69,7 @@ In the control plane, `TransferQueueController` tracks the **production status**
 For consumption status, we record the consumption records for each computational task (e.g., `generate_sequences`, `compute_log_prob`, etc.). Therefore, even when different computation tasks require the same data field, they can consume the data independently without interfering with each other.
 
 <p align="center">
-  <img src="https://cdn.nlark.com/yuque/0/2025/png/23208217/1758696820173-456c1784-42ba-40c8-a292-2ff1401f49c5.png" width="70%">
+  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/control_plane.png?raw=true" width="70%">
 </p>
 
 To make the data retrieval process more customizable, we provide a `Sampler` class that allows users to define their own data retrieval and consumption logic. Refer to the [Customize](#customize) section for details.
@@ -91,9 +91,9 @@ This class encapsulates the core interaction logic within the TransferQueue syst
 Currently, we support the following storage backends:
 
 - SimpleStorageUnit: A basic CPU memory storage with minimal data format constraints and easy usability.
-- [MoonCakeStore](https://github.com/kvcache-ai/Mooncake): A high-performance, KV-based hierarchical storage that supports RDMA transport between GPU and DRAM.
 - [Yuanrong](https://gitee.com/openeuler/yuanrong-datasystem): An Ascend native data system that provides hierarchical storage interfaces including HBM/DRAM/SSD.
-- [Ray Direct Transport](https://docs.ray.io/en/master/ray-core/direct-transport.html): Ray's new feature that allows Ray to store and pass objects directly between Ray actors.
+- [MoonCakeStore](https://github.com/kvcache-ai/Mooncake) (WIP): A high-performance, KV-based hierarchical storage that supports RDMA transport between GPU and DRAM.
+- [Ray Direct Transport](https://docs.ray.io/en/master/ray-core/direct-transport.html) ([WIP](https://github.com/TransferQueue/TransferQueue/pull/108)): Ray's new feature that allows Ray to store and pass objects directly between Ray actors.
 
 Among them, `SimpleStorageUnit` serves as our default storage backend, coordinated by the `AsyncSimpleStorageManager` class. Each storage unit can be deployed on a separate node, allowing for distributed data management.
 
@@ -105,7 +105,7 @@ Among them, `SimpleStorageUnit` serves as our default storage backend, coordinat
 This data structure design is motivated by the computational characteristics of the post-training process, where each training sample is generated in a relayed manner across task pipelines. It provides an accurate addressing capability, which allows fine-grained, concurrent data read/write operations in a streaming manner.
 
 <p align="center">
-  <img src="https://cdn.nlark.com/yuque/0/2025/png/23208217/1758696805154-3817011f-84e6-40d0-a80c-58b7e3e5f6a7.png" width="70%">
+  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/data_plane.png?raw=true" width="70%">
 </p>
 
 ### User Interface: Asynchronous & Synchronous Client
@@ -140,7 +140,7 @@ We will soon release a detailed tutorial and API documentation.
 #### verl
 The primary motivation for integrating TransferQueue to verl now is to **alleviate the data transfer bottleneck of the single controller `RayPPOTrainer`**. Currently, all `DataProto` objects must be routed through `RayPPOTrainer`, resulting in a single point bottleneck of the whole post-training system.
 
-![verl_dataflow_DataProto](https://cdn.nlark.com/yuque/0/2025/jpeg/23208217/1758704289414-bcc54228-716b-4d4a-ad3b-f9ace6d10fcf.jpeg)
+![verl_dataflow_DataProto](https://github.com/TransferQueue/community_doc/blob/main/docs/verl_workflow.jpeg?raw=true)
 
 Leveraging TransferQueue, we separate experience data transfer from metadata dispatch by
 
@@ -148,7 +148,7 @@ Leveraging TransferQueue, we separate experience data transfer from metadata dis
 - Preserving verl's original Dispatch/Collect logic via BatchMeta (maintaining single-controller debuggability)
 - Accelerating data transfer by TransferQueue's distributed storage units
 
-![verl_dataflow_TransferQueue](https://cdn.nlark.com/yuque/0/2025/jpeg/23208217/1758704301666-0807dc06-766c-4a2d-9cde-889a6bb56b34.jpeg)
+![verl_dataflow_TransferQueue](https://github.com/TransferQueue/community_doc/blob/main/docs/verl_workflow_with_tq.jpeg?raw=true)
 
 You may refer to the [recipe](https://github.com/TransferQueue/TransferQueue/tree/dev/recipe/simple_use_case), where we mimic the verl usage in both async & sync scenarios. Official integration to verl is also available now at [verl/pulls/3649](https://github.com/volcengine/verl/pull/3649) (with subsequent PRs to further optimize the integration).
 
@@ -157,7 +157,7 @@ You may refer to the [recipe](https://github.com/TransferQueue/TransferQueue/tre
 Work in progress :)
 
 <p align="center">
-  <img src="https://cdn.nlark.com/yuque/0/2025/png/23208217/1758696840817-14ba4c3b-b96e-4390-ac7c-4ecf7b8c0ac3.png" width="70%">
+  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/tq_streaming_dataloader.png?raw=true" width="70%">
 </p>
 
 <h2 id="quick-start">🚀 Quick Start</h2>
@@ -188,12 +188,12 @@ Follow these steps to build and install:
 <h2 id="performance">📊 Performance</h2>
 
 <p align="center">
-  <img src="https://cdn.nlark.com/yuque/0/2025/png/23208217/1761294403612-76ca20a7-9108-42fc-b3f5-60f84d70f39b.png" width="100%">
+  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/performance_0.1.1.dev2.png?raw=true" width="100%">
 </p>
 
 > Note: The above benchmark for TransferQueue is based on our naive `SimpleStorageUnit` backend. By introducing high-performance storage backends and optimizing serialization/deserialization, we expect to achieve even better performance. Warmly welcome contributions from the community!
 
-For detailed performance benchmarks, please refer to [this blog](https://www.yuque.com/haomingzi-lfse7/hlx5g0/obi4ovmy9wf08zz3?singleDoc#).
+For detailed performance benchmarks, please refer to [this blog](https://www.yuque.com/haomingzi-lfse7/hlx5g0/tml8ke0zkgn6roey?singleDoc#).
 
 <h2 id="customize"> 🛠️ Customize TransferQueue</h2>
 

{transferqueue-0.1.1.dev1.dist-info → transferqueue-0.1.2.dev0.dist-info}/RECORD

@@ -5,7 +5,6 @@ tests/test_client.py,sha256=74Pm1D4SI_GCg0Kxwm5Wqa4ppSfc57mpHJGuIdtNUrs,15325
 tests/test_controller.py,sha256=ZcvFCC3jSnNN_fEerjA37RQv0SSO0Xh8vjcL2mvF03o,11084
 tests/test_controller_data_partitions.py,sha256=qZxMHerMwKIwyRmT8FZke8TEd80Z9vAkBzU_k5Jz1bY,19185
 tests/test_kv_storage_manager.py,sha256=Eh6xykdhLBMpxikXfRHxn1crhLsQjn9QGa0O7TLO-5o,3582
-tests/test_put.py,sha256=WnRKCGPXmRITAvbD8KWlZor4jpvv6sdg2gg3NCw9gyQ,10453
 tests/test_samplers.py,sha256=CvYqfmbHEWWa1RyymztCAn0GcitAPOBbfJ4ud1VvO2o,19168
 tests/test_serial_utils_on_cpu.py,sha256=Hju5yAV52JP-cPw-PT_jsut-y6J_7lUX_SbG5EO1lNQ,7379
 tests/test_simple_storage_unit.py,sha256=Sczhw2bdCfTfa7RwnpW-aKCV1of8mO6z3l3q2PDzZCA,16535
@@ -13,7 +12,7 @@ tests/test_storage_client_factory.py,sha256=U0gS_l4bc_bP7K_uPhy8UlJ186HyTrA10TSJ
 transfer_queue/__init__.py,sha256=68c0sBfqHPqTa7OdzO4sAZB52XvwtjpwLqP9BWAh4fA,1535
 transfer_queue/client.py,sha256=zDlH1beWwRbjz0a7S8QH9IOJ1wp5yEQ36XYwFwJTXOY,24949
 transfer_queue/controller.py,sha256=uc_MQAlG_QmJ8szxc5yPLaZvCeP2CRULnAADtcf41-8,49702
-transfer_queue/metadata.py,sha256=LHQhM7vixw5QcA05nZcqWbxSh7mhJ2ETJOo_1gFb3eI,17645
+transfer_queue/metadata.py,sha256=mH3nev_lK1HiIqZ6przEgCd0BiHkO7jI6ji3x6VSBgY,19076
 transfer_queue/sampler/__init__.py,sha256=1oauDy2Dwb5GXhKi7tl5DWAHv8i4t2MQK1S4U36Sy4g,788
 transfer_queue/sampler/base.py,sha256=wFti4dNJb3YArYpGzxA_YDfyUTdTG8wVz6HclPDyZPw,3299
 transfer_queue/sampler/grpo_group_n_sampler.py,sha256=Kq3hGAz8mboBNvw4Dj0P8lP6Qs8TDojx81fxSh57w28,6566
@@ -27,15 +26,15 @@ transfer_queue/storage/clients/yuanrong_client.py,sha256=rWOPQPLgHLav7cFdGCr8CeV
 transfer_queue/storage/managers/__init__.py,sha256=y5x5OzZwK_YormdHdzc-smnJNew3niuqU2g_kkaiXIk,876
 transfer_queue/storage/managers/base.py,sha256=ntlo6sLWIbTiAOUxgJTUyjLl5m4HBuAUUNCYZxq9BFM,20352
 transfer_queue/storage/managers/factory.py,sha256=58kp2mCKz1K8Ea7RWMsWxdDhN3y4ZhgE-G647AKq7-I,1752
-transfer_queue/storage/managers/simple_backend_manager.py,sha256=fCj-0BnS3RbzPn03KEcxmUWbxZVWK11cfxpj8tSf-yY,27381
+transfer_queue/storage/managers/simple_backend_manager.py,sha256=zn0RrRng2zo0kLHl3eKQ9Ksf9O1W9DCEzpuZes8uSjM,27684
 transfer_queue/storage/managers/yuanrong_manager.py,sha256=NjHC3LBW0fQwm30Oq_qEKoCQEq8oWO0D-AobcpQNPNg,777
 transfer_queue/utils/__init__.py,sha256=vki-5RVaRBKxVc6Q7XPQox3VNPio2DvJYvRz0SZtu-w,586
 transfer_queue/utils/serial_utils.py,sha256=9ZgsytTp-441YKtIRFqyH5NhNifSqRKa2h0FI44ltcc,10200
 transfer_queue/utils/utils.py,sha256=Pno4h3WjX_eT7q4xiVV6Jkhquc39Fp-Ycsg1cv0qNKQ,4544
 transfer_queue/utils/zmq_utils.py,sha256=jCg2pQfvy_IYdGyZq4nvL4CAwjJc7Li0Trp3T3GDBMg,5118
-transfer_queue/version/version,sha256=IBGGJN0Ii7hsCVlxnB5t6XS0CWaPTLXLjPL0jyY3vKE,11
-transferqueue-0.1.1.dev1.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-transferqueue-0.1.1.dev1.dist-info/METADATA,sha256=nYLoCMpxUeP1Go3Pnai9W7lVRc4wRAUocqrzyP4K5ds,19531
-transferqueue-0.1.1.dev1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-transferqueue-0.1.1.dev1.dist-info/top_level.txt,sha256=BiBclu7jWJ0AZ35vUr3hN9_cg8JL9EiH_hjFxquMxtw,33
-transferqueue-0.1.1.dev1.dist-info/RECORD,,
+transfer_queue/version/version,sha256=Q11DRFEP_n2U67FVBUV-pmge1VJtg5UP4Tj8Ski4SkU,11
+transferqueue-0.1.2.dev0.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+transferqueue-0.1.2.dev0.dist-info/METADATA,sha256=JnPLADNXjifw9xZxuj6vvNahLyT8k0bDc1gDuJxosOk,19502
+transferqueue-0.1.2.dev0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+transferqueue-0.1.2.dev0.dist-info/top_level.txt,sha256=BiBclu7jWJ0AZ35vUr3hN9_cg8JL9EiH_hjFxquMxtw,33
+transferqueue-0.1.2.dev0.dist-info/RECORD,,

tests/test_put.py

@@ -1,327 +0,0 @@
-# Copyright 2025 The TransferQueue Team
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import asyncio
-import logging
-import sys
-from pathlib import Path
-
-import pytest
-import ray
-import torch
-from tensordict import TensorDict
-
-parent_dir = Path(__file__).resolve().parent.parent
-sys.path.append(str(parent_dir))
-
-from transfer_queue import (  # noqa: E402
-    AsyncTransferQueueClient,
-    SimpleStorageUnit,
-    TransferQueueController,
-    process_zmq_server_info,
-)
-from transfer_queue.utils.utils import get_placement_group  # noqa: E402
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-@pytest.fixture(scope="module")
-def ray_setup():
-    """Initialize Ray for testing."""
-    ray.init(ignore_reinit_error=True)
-    yield
-    ray.shutdown()
-
-
-@pytest.fixture(scope="module")
-def data_system_setup(ray_setup):
-    """Set up data system for testing."""
-    # Initialize storage units
-    num_storage_units = 2
-    storage_size = 10000
-    storage_units = {}
-
-    storage_placement_group = get_placement_group(num_storage_units, num_cpus_per_actor=1)
-    for storage_unit_rank in range(num_storage_units):
-        storage_node = SimpleStorageUnit.options(
-            placement_group=storage_placement_group, placement_group_bundle_index=storage_unit_rank
-        ).remote(storage_unit_size=storage_size)
-        storage_units[storage_unit_rank] = storage_node
-        logger.info(f"SimpleStorageUnit #{storage_unit_rank} has been created.")
-
-    # Initialize controller
-    controller = TransferQueueController.remote()
-    logger.info("TransferQueueController has been created.")
-
-    # Prepare connection info
-    controller_info = process_zmq_server_info(controller)
-    storage_unit_infos = process_zmq_server_info(storage_units)
-
-    # Create config
-    config = {
-        "controller_info": controller_info,
-        "storage_unit_infos": storage_unit_infos,
-    }
-
-    yield controller, storage_units, config
-
-    # Cleanup
-    ray.kill(controller)
-    for storage_unit in storage_units.values():
-        ray.kill(storage_unit)
-
-
-@pytest.fixture
-async def client_setup(data_system_setup):
-    """Set up client for testing."""
-    controller, storage_units, config = data_system_setup
-
-    client = AsyncTransferQueueClient(
-        client_id="TestClient",
-        controller_info=config["controller_info"],
-    )
-
-    client.initialize_storage_manager(manager_type="AsyncSimpleStorageManager", config=config)
-
-    # Wait a bit for connections to establish
-    await asyncio.sleep(1)
-
-    return client
-
-
-class TestMultipleAsyncPut:
-    """Test class for multiple async_put operations."""
-
-    def __init__(self):
-        self.controller = None
-        self.storage_units = None
-        self.config = None
-
-    async def setup(self):
-        """Setup for the test class."""
-        if not ray.is_initialized():
-            ray.init(ignore_reinit_error=True)
-
-        # Initialize data system
-        num_storage_units = 2
-        self.storage_units = {}
-        storage_placement_group = get_placement_group(num_storage_units, num_cpus_per_actor=1)
-
-        for i in range(num_storage_units):
-            self.storage_units[i] = SimpleStorageUnit.options(
-                placement_group=storage_placement_group, placement_group_bundle_index=i
-            ).remote(storage_unit_size=10000)
-
-        self.controller = TransferQueueController.remote()
-
-        # Wait for initialization
-        await asyncio.sleep(2)
-
-        controller_info = process_zmq_server_info(self.controller)
-        storage_unit_infos = process_zmq_server_info(self.storage_units)
-
-        self.config = {
-            "controller_info": controller_info,
-            "storage_unit_infos": storage_unit_infos,
-        }
-        logger.info("TestMultipleAsyncPut setup completed")
-
-    async def teardown(self):
-        """Teardown for the test class."""
-        if self.controller:
-            ray.kill(self.controller)
-        if self.storage_units:
-            for storage in self.storage_units.values():
-                ray.kill(storage)
-        if ray.is_initialized():
-            ray.shutdown()
-
-    async def create_client(self, client_id="TestClient"):
-        """Create a new client instance."""
-        if self.config is None:
-            await self.setup()
-
-        client = AsyncTransferQueueClient(
-            client_id=client_id,
-            controller_info=self.config["controller_info"],
-        )
-        client.initialize_storage_manager(manager_type="AsyncSimpleStorageManager", config=self.config)
-        await asyncio.sleep(1)  # Wait for connections
-        return client
-
-
-@pytest.mark.asyncio
-async def test_concurrent_async_put():
-    """Test concurrent async_put operations."""
-    test_instance = TestMultipleAsyncPut()
-    await test_instance.setup()
-
-    try:
-        client = await test_instance.create_client("ConcurrentTestClient")
-
-        async def put_operation(partition_id, data):
-            await client.async_put(data=data, partition_id=partition_id)
-            return partition_id
-
-        # Create multiple put operations
-        tasks = []
-        for i in range(5):
-            data = TensorDict({f"concurrent_data_{i}": torch.randn(4, 8)}, batch_size=[4])
-            task = put_operation(f"concurrent_partition_{i}", data)
-            tasks.append(task)
-
-        # Execute concurrently
-        results = await asyncio.gather(*tasks)
-
-        assert len(results) == 5
-        logger.info(f"Completed {len(results)} concurrent put operations")
-
-        # Cleanup
-        for partition in results:
-            try:
-                await client.async_clear(partition)
-            except Exception as e:
-                logger.warning(f"Failed to clear partition {partition}: {e}")
-
-        client.close()
-    finally:
-        await test_instance.teardown()
-
-
-@pytest.mark.asyncio
-async def test_sequential_async_put_with_verification():
-    """Test sequential async_put operations with data verification."""
-    test_instance = TestMultipleAsyncPut()
-    await test_instance.setup()
-
-    try:
-        client = await test_instance.create_client("SequentialTestClient")
-
-        # Test data
-        test_cases = [
-            ("sequential_1", torch.randn(4, 5)),
-            ("sequential_2", torch.randn(4, 10)),
-            ("sequential_3", torch.randn(4, 15)),
-        ]
-
-        for partition_id, tensor_data in test_cases:
-            data = TensorDict({"sequential_data": tensor_data}, batch_size=[4])
-
-            # Put data
-            await client.async_put(data=data, partition_id=partition_id)
-            logger.info(f"Put data to {partition_id}")
-
-            # Verify by reading back
-            metadata = await client.async_get_meta(
-                data_fields=["sequential_data"],
-                batch_size=4,
-                partition_id=partition_id,
-                mode="fetch",
-                task_name="verification_task",
-            )
-
-            retrieved_data = await client.async_get_data(metadata)
-
-            # Verify shape and content
-            assert retrieved_data["sequential_data"].shape == tensor_data.shape
-            logger.info(f"Verified data in {partition_id}")
-
-            # Cleanup
-            await client.async_clear(partition_id)
-
-        client.close()
-    finally:
-        await test_instance.teardown()
-
-
-@pytest.mark.asyncio
-async def test_multiple_puts_same_partition():
-    """Test multiple puts to the same partition."""
-    test_instance = TestMultipleAsyncPut()
-    await test_instance.setup()
-
-    try:
-        client = await test_instance.create_client("SamePartitionTestClient")
-
-        partition_id = "same_partition_test"
-        batch_size = 4
-
-        # First put
-        data1 = TensorDict({"first_batch": torch.randn(batch_size, 8)}, batch_size=[batch_size])
-
-        await client.async_put(data=data1, partition_id=partition_id)
-        logger.info("First put completed")
-
-        # Second put to same partition
-        data2 = TensorDict({"second_batch": torch.randn(batch_size, 8) * 2}, batch_size=[batch_size])
-
-        await client.async_put(data=data2, partition_id=partition_id)
-        logger.info("Second put completed")
-
-        # Third put to same partition
-        data3 = TensorDict({"third_batch": torch.randn(batch_size, 8) * 3}, batch_size=[batch_size])
-
-        await client.async_put(data=data3, partition_id=partition_id)
-        logger.info("Third put completed")
-
-        # Verify the last data is what we get
-        metadata = await client.async_get_meta(
-            data_fields=["third_batch"],
-            batch_size=batch_size,
-            partition_id=partition_id,
-            mode="fetch",
-            task_name="verification_task",
-        )
-
-        retrieved_data = await client.async_get_data(metadata)
-
-        assert "third_batch" in retrieved_data
-        assert retrieved_data["third_batch"].shape == (batch_size, 8)
-        logger.info("Verified multiple puts to same partition")
-
-        # Cleanup
-        await client.async_clear(partition_id)
-
-        client.close()
-    finally:
-        await test_instance.teardown()
-
-
-@pytest.mark.asyncio
-async def test_simple_multiple_async_put():
-    """Simple test using the existing client_setup fixture."""
-    test_instance = TestMultipleAsyncPut()
-    await test_instance.setup()
-
-    try:
-        client = await test_instance.create_client("SimpleTestClient")
-
-        # Test basic multiple puts
-        for i in range(3):
-            data = TensorDict({f"simple_data_{i}": torch.randn(2, 4)}, batch_size=[2])
-
-            partition_id = f"simple_partition_{i}"
-            await client.async_put(data=data, partition_id=partition_id)
-            logger.info(f"Put data to {partition_id}")
-
-            # Cleanup each partition
-            await client.async_clear(partition_id)
-
-        logger.info("Simple multiple async_put test completed")
-
-        client.close()
-    finally:
-        await test_instance.teardown()