Better to provide a base class TransferQueueStorageClient for storage client?

done

param name refer to TransferQueueStorageManagerFactory (name->client_type)

done

As above comment, better to provide a base abstract class

done

cfg -> config so we can align with other parts of the code~

done

better raise an error; how to make sure the address space is continues?

tensor has a method is_contiguous

better raise an error; how to make sure the address space is continues?

Yuanrong data system imposes a limit on the number of keys for a single query; this will be modified later to use multiple queries. Additionally, empty tensors are created by KVStorageManager._create_empty_tensorlist and should therefore be contiguous. If the additional performance overhead is acceptable, a check for tensor contiguity can be added in the future.

extra space; please use
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always to auto fix these issues

it can be a general function. We can just call it _generate_keys. same for following function names

done

Maybe we can use itertools to speed up this process? there may be lots of keys.

And is this key order can get better performance than the other order? Typically rows can be much more than columns

And is this key order can get better performance than the other order? Typically rows can be much more than columns

Generally speaking, for two unrelated lists, I think placing the shorter list in the outer loop is better, as it means fewer conditional checks and fewer pointer jumps.

as above comment

done

Add a TODO here. We need to support complex data types (NonTensorStack/NonTensorData/NestedTensor)

done

_merge_kv_to_tensordict

suggest use from transfer_queue.storage.managers.base import KVStorageManager

had updated

suggest use from transfer_queue.storage.clients.base import TransferQueueStorageClient

Extra space between 'import' and 'torch'. Should be 'import torch'.

done

Missing space after comma in import statement. Should be 'ABC, abstractmethod'.

done

Missing import for 'StorageClientFactory'. The class is used but not imported from 'transfer_queue.storage.clients.factory'.

done

Error messages at lines 90 and 92 are incomplete. They should clearly state what is missing, e.g., 'shapes parameter is required for Yuanrong DataSystem' and 'dtypes parameter is required for Yuanrong DataSystem'.

had upated more prompts

Incomplete TODO comment. Should clarify what is wrong with 'DSTensorClient.dev_mget', e.g., 'TODO: DSTensorClient.dev_mget has wrong behavior' or 'TODO: DSTensorClient.dev_mget has wrong return type'.

done

Variable client is not used.

done

Import of 'YuanrongStorageManager' is not used.

could be ignored

Import of 'TensorDict' is not used.

done

Import of 'torch_npu' is not used.

torch_npu is neccessary for yuanrong

'except' clause does nothing but pass and there is no explanatory comment.

it try to import YRStorageClient for registering

⚠️ Potential issue | 🟠 Major

Replace the assertion with explicit validation.

assert len(keys) <= 10000 is skipped when Python runs with optimizations (python -O), so the critical safeguard on key count silently disappears. That leaves dev_mset exposed to oversized batches, violating the documented constraint. Please convert this to an explicit runtime check that raises a ValueError.

-        # 约束：传入的key的数量不能超过1万。&Tensor的地址空间必须连续。
-        assert len(keys) <= 10000
+        # 约束：传入的key的数量不能超过1万。&Tensor的地址空间必须连续。
+        if len(keys) > 10000:
+            raise ValueError("Number of keys cannot exceed 10_000")

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

In transfer_queue/storage/clients/yuanrong_client.py around lines 63 to 71,
replace the assert len(keys) <= 10000 with an explicit runtime validation that
checks if len(keys) > 10000 and raises a ValueError (including the actual key
count and the 10000 limit in the message); keep the subsequent tensor type and
device checks as-is so oversized batches are rejected reliably even when Python
is run with optimizations.

done

@dpj135 Thank you for addressing this issue! The explicit validation will now properly enforce the 10,000 key limit even when Python runs with optimizations enabled. 🎉

⚠️ Potential issue | 🟠 Major

Validate key/metadata alignment and enforce the limit without assert.

get never checks that the number of keys matches the supplied shapes/dtypes, so callers can provide mismatched metadata and silently feed _ds_client.dev_mget with too few or too many output buffers. That risks read-overwrite errors or misaligned tensors. Additionally, the key-count guard again relies on assert, which is stripped with python -O. Please add explicit validation for both the length alignment and the 10 000-key limit.

         if dtypes is None:
             raise ValueError('Yuanrong DataSystem')
         if len(dtypes) != len(shapes):
             raise ValueError('Length of dtypes must equal length of shapes')
+        if len(keys) != len(shapes):
+            raise ValueError('Number of keys must match provided shapes/dtypes')
+        if len(keys) > 10000:
+            raise ValueError('Number of keys cannot exceed 10_000')
 
         values: list[Tensor] = self._create_empty_tensorlist(shapes=shapes, dtypes=dtypes)
 
-        # 约束：传入的key的数量不能超过1万。&Tensor的地址空间必须连续。
-        # print(f'get_keys: {keys}')
-        assert len(keys) <= 10000
+        # 约束：传入的key的数量不能超过1万。&Tensor的地址空间必须连续。
+        # print(f'get_keys: {keys}')

In transfer_queue/storage/clients/yuanrong_client.py around lines 89-104, the
method currently uses assert for the 10,000-key guard and does not validate that
len(keys) matches the supplied shapes and dtypes, which can cause misaligned
buffers or be skipped under python -O; replace the assert with an explicit check
that raises a ValueError if len(keys) > 10000, and add a validation that
len(keys) == len(shapes) == len(dtypes) (raising ValueError with a clear message
if not) before creating the output tensors and calling _ds_client.dev_mget.

done

@dpj135, thank you for addressing the validation issues! The explicit checks will ensure robust error handling in production environments.

⚠️ Potential issue | 🔴 Critical

Missing import causes NameError.

Line 286 references StorageClientFactory but it's not imported. This will cause a runtime error when instantiating KVStorageManager.

Add the missing import at the top of the file:

+from transfer_queue.storage.clients.factory import StorageClientFactory

Committable suggestion skipped: line range outside the PR's diff.

In transfer_queue/storage/managers/base.py around lines 280 to 287, the
constructor uses StorageClientFactory but it's not imported; add the missing
import at the top of the file (e.g., add a line importing StorageClientFactory —
use a relative import like "from ..clients import StorageClientFactory" or the
full package path "from transfer_queue.storage.clients import
StorageClientFactory") so the name is defined when KVStorageManager is
instantiated.

done

🏁 Script executed:

#!/bin/bash
# Verify the import statement is present in the file
grep -n "from transfer_queue.storage.clients.factory import StorageClientFactory" transfer_queue/storage/managers/base.py

Length of output: 74

@dpj135 Thank you for adding the import! The StorageClientFactory import is now properly included on line 31, which resolves the issue. The KVStorageManager class can now be instantiated without any NameError. ✅

[transfer_queue/storage/managers/factory.py:26-40]

def register(cls, manager_type: str, manager_cls: type[TransferQueueStorageManager]):
if not issubclass(manager_cls, TransferQueueStorageManager):
raise TypeError(
f"manager_cls {getattr(manager_cls, 'name', repr(manager_cls))} must be "
f"a subclass of TransferQueueStorageManager"
)
cls._registry[manager_type] = manager_cls

End snippet

[transfer_queue/storage/managers/factory.py:35-40]

def create(cls, manager_type: str, config: dict[str, Any]) -> TransferQueueStorageManager:
if manager_type not in cls._registry:
raise ValueError(
f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}"
)
return cls._registrymanager_type

End snippet

[transfer_queue/storage/clients/factory.py:35-40]

def register(cls, client_type: str):
"""
Decorator to register a concrete client class with the factory.
Args:
client_type (str): The name used to identify the client
Returns:
Callable: The decorator function that returns the original class
"""
def decorator(client_class: TransferQueueStorageClient) -> TransferQueueStorageClient:
cls._registry[client_type] = client_class
return client_class
return decorator

End snippet

[transfer_queue/storage/clients/factory.py:59-71]

def create(cls, client_type: str, config: dict) -> TransferQueueStorageClient:
"""
Create and return an instance of the storage client by name.
Args:
client_type (str): The registered name of the client
Returns:
StorageClientFactory: An instance of the requested client
Raises:
ValueError: If no client is registered with the given name
"""
if client_type not in cls._registry:
raise ValueError(f"Unknown StorageClient: {client_type}")
return cls._registryclient_type

End snippet

[transfer_queue/utils/zmq_utils.py:31-65]

class ZMQRequestType(ExplicitEnum):
# HANDSHAKE
HANDSHAKE = "HANDSHAKE" # TransferQueueStorageUnit -> TransferQueueController
HANDSHAKE_ACK = "HANDSHAKE_ACK" # TransferQueueController -> TransferQueueStorageUnit

# DATA_OPERATION
GET_DATA = "GET"
PUT_DATA = "PUT"
GET_DATA_RESPONSE = "GET_DATA_RESPONSE"
PUT_DATA_RESPONSE = "PUT_DATA_RESPONSE"
CLEAR_DATA = "CLEAR_DATA"
CLEAR_DATA_RESPONSE = "CLEAR_DATA_RESPONSE"

PUT_GET_OPERATION_ERROR = "PUT_GET_OPERATION_ERROR"
PUT_GET_ERROR = "PUT_GET_ERROR"
PUT_ERROR = "PUT_ERROR"
GET_ERROR = "GET_ERROR"
CLEAR_DATA_ERROR = "CLEAR_DATA_ERROR"

# META_OPERATION
GET_META = "GET_META"
GET_META_RESPONSE = "GET_META_RESPONSE"
GET_CLEAR_META = "GET_CLEAR_META"
GET_CLEAR_META_RESPONSE = "GET_CLEAR_META_RESPONSE"
CLEAR_META = "CLEAR_META"
CLEAR_META_RESPONSE = "CLEAR_META_RESPONSE"

# CHECK_CONSUMPTION
CHECK_CONSUMPTION = "CHECK_CONSUMPTION"
CONSUMPTION_RESPONSE = "CONSUMPTION_RESPONSE"

# NOTIFY_DATA_UPDATE
NOTIFY_DATA_UPDATE = "NOTIFY_DATA_UPDATE"
NOTIFY_DATA_UPDATE_ACK = "NOTIFY_DATA_UPDATE_ACK"
NOTIFY_DATA_UPDATE_ERROR = "NOTIFY_DATA_UPDATE_ERROR"

End snippet

[transfer_queue/utils/zmq_utils.py:68-87]

class ZMQServerInfo:
def init(self, role: TransferQueueRole, id: str, ip: str, ports: dict[str, str]):
self.role = role
self.id = id
self.ip = ip
self.ports = ports

def to_addr(self, port_name: str) -> str:
    return f"tcp://{self.ip}:{self.ports[port_name]}"

def to_dict(self):
    return {
        "role": self.role,
        "id": self.id,
        "ip": self.ip,
        "ports": self.ports,
    }

def __str__(self) -> str:
    return f"ZMQSocketInfo(role={self.role}, id={self.id}, ip={self.ip}, ports={self.ports})"

End snippet

[transfer_queue/utils/zmq_utils.py:100-114]

def create(
cls,
request_type: ZMQRequestType,
sender_id: str,
body: dict[str, Any],
receiver_id: Optional[str] = None,
) -> "ZMQMessage":
return cls(
request_type=request_type,
sender_id=sender_id,
receiver_id=receiver_id,
body=body,
request_id=str(uuid4().hex[:8]),
timestamp=time.time(),
)

End snippet

[transfer_queue/utils/zmq_utils.py:116-118]

def serialize(self) -> bytes:
"""Using pickle to serialize ZMQMessage objects"""
return pickle.dumps(self)

End snippet

[transfer_queue/utils/zmq_utils.py:121-131]

def deserialize(cls, data: bytes | list[bytes]):
"""Using pickle to deserialize ZMQMessage objects"""
if isinstance(data, list):
# Process multiple byte streams by deserializing each in sequence
result = []
for d in data:
result.append(pickle.loads(d))
return result
else:
# Single byte stream case
return pickle.loads(data)

End snippet

[transfer_queue/metadata.py:136-151]

class BatchMeta:
"""Records the metadata of a batch of data samples."""

samples: list[SampleMeta]
extra_info: dict[str, Any] = dataclasses.field(default_factory=dict)

def __post_init__(self):
    """Initialize all computed properties during initialization"""
    # Basic properties
    object.__setattr__(self, "_size", len(self.samples))
    object.__setattr__(self, "_is_ready", all(sample.is_ready for sample in self.samples))

    # Pre-compute all list properties for better performance
    if self.samples:
        for idx, sample in enumerate(self.samples):
            object.__setattr__(sample, "_batch_index", idx)  # Ensure batch_index is set correctly

        object.__setattr__(self, "_global_indexes", [sample.global_index for sample in self.samples])

        # assume all samples have the same fields.
        object.__setattr__(self, "_field_names", sorted(self.samples[0].field_names))
    else:
        object.__setattr__(self, "_global_indexes", [])
        object.__setattr__(self, "_field_names", [])

End snippet

[transfer_queue/metadata.py:167-169]

def global_indexes(self) -> list[int]:
"""Get all global indexes in this batch"""
return getattr(self, "_global_indexes", [])

End snippet

[transfer_queue/metadata.py:74-76]

def get_field_by_name(self, name: str) -> Optional[FieldMeta]:
"""Get FieldMeta by field name"""
return self.fields.get(name)

End snippet

[transfer_queue/storage/clients/yuanrong_client.py:41-63]

def put(self, keys: list[str], values: list[Tensor]):
"""
Store tensors to remote storage.
Args:
keys (list): List of string keys
values (list): List of torch.Tensor on NPU
"""
if not isinstance(keys, list) or not isinstance(values, list):
raise ValueError("keys and values must be lists")
if len(keys) != len(values):
raise ValueError("Number of keys must match number of values")

    # TODO: Support the situation when the number of keys is greater than 10000
    if len(keys) > 10000:
        raise NotImplementedError('We will support the number of keys greater than 10000 int the future')

    for value in values:
        if not isinstance(value, torch.Tensor):
            raise ValueError(f"Expected torch.Tensor, got {type(value)}")
        if value.device.type != 'npu':
            raise ValueError(f"Tensor is on {value.device}, not on NPU")

    self._ds_client.dev_mset(keys, values)

End snippet

[transfer_queue/storage/clients/yuanrong_client.py:65-90]

def get(self, keys: list[str], shapes=None, dtypes=None) -> list[Tensor]:
"""
Retrieve tensors from remote storage.
Args:
keys (list): List of keys to fetch
shapes (list): Expected shapes of returned tensors
dtypes (list): Expected dtypes of returned tensors
Returns:
list: List of retrieved NPU tensors
"""
if shapes is None:
raise ValueError('Yuanrong storage client needs Expected shapes of returned tensors')
if dtypes is None:
raise ValueError('Yuanrong storage client needs Expected dtypes of returned tensors')
if len(dtypes) != len(shapes):
raise ValueError('Length of dtypes must equal length of shapes')

    values: list[Tensor] = self._create_empty_tensorlist(shapes=shapes, dtypes=dtypes)

    # TODO: Support the situation when the number of keys is greater than 10000
    if len(keys) > 10000:
        raise NotImplementedError('We will support the number of keys greater than 10000 int the future')

    # Timeout set to 2000ms
    self._ds_client.dev_mget(keys, values, 2000)
    return values

End snippet

[transfer_queue/storage/clients/yuanrong_client.py:92-98]

def clear(self, keys: list[str]):
"""
Delete entries from storage by keys.
Args:
keys (list): List of keys to delete
"""
self._ds_client.dev_delete(keys)

End snippet

[transfer_queue/storage/metadata.py:136-151]

class BatchMeta:
"""Records the metadata of a batch of data samples."""

samples: list[SampleMeta]
extra_info: dict[str, Any] = dataclasses.field(default_factory=dict)

def __post_init__(self):
    """Initialize all computed properties during initialization"""
    # Basic properties
    object.__setattr__(self, "_size", len(self.samples))
    object.__setattr__(self, "_is_ready", all(sample.is_ready for sample in self.samples))

    # Pre-compute all list properties for better performance
    if self.samples:
        for idx, sample in enumerate(self.samples):
            object.__setattr__(sample, "_batch_index", idx)  # Ensure batch_index is set correctly

        object.__setattr__(self, "_global_indexes", [sample.global_index for sample in self.samples])

        # assume all samples have the same fields.
        object.__setattr__(self, "_field_names", sorted(self.samples[0].field_names))
    else:
        object.__setattr__(self, "_global_indexes", [])
        object.__setattr__(self, "_field_names", [])

End snippet

[transfer_queue/storage/managers/base.py:— end —]

Note: The above snippets are provided as-is from the referenced files to help understand interactions with the file transfer_queue/storage/managers/base.py and its related components.

⚠️ Potential issue | 🟡 Minor

RuntimeError catch may be too broad.

The except RuntimeError clause assumes any RuntimeError during torch.stack is due to shape mismatches, but other runtime errors (e.g., CUDA errors, device mismatches) could also raise RuntimeError. Consider catching more specific exceptions or adding validation.

Consider validating tensor shapes before stacking, or at minimum logging the exception to aid debugging:

         try:
             tensor_data[field] = torch.stack(tensor_list)
-        except RuntimeError as re:
+        except RuntimeError as e:
+            logger.warning(f"Cannot stack tensors for field {field}, using nested tensor. Error: {e}")
             # Fallback to nested tensor if shapes are irregular
             tensor_data[field] = torch.nested.as_nested_tensor(tensor_list)

Committable suggestion skipped: line range outside the PR's diff.

In transfer_queue/storage/managers/base.py around lines 375 to 379, the current
except RuntimeError: fallback to nested tensor is too broad; validate tensors
before stacking (ensure all tensors have identical .shape and .device), only
attempt torch.stack when validation passes, and if torch.stack still fails catch
the exception, log the exception details, and only use
torch.nested.as_nested_tensor when the failure is confirmed to be due to
irregular shapes/devices; otherwise re-raise or propagate the error so true
runtime (e.g., CUDA) issues aren’t silently swallowed.

ignored

@dpj135, understood. I've noted your decision to skip this suggestion.

⚠️ Potential issue | 🔴 Critical

Methods marked async but perform no async operations.

All three methods (put_data, get_data, clear_data) are declared async but call synchronous storage_client methods without using await. This blocks the event loop and defeats the purpose of async/await.

Either:

1. If storage_client methods are synchronous, remove async from these method signatures (but this may break the abstract base class contract).
2. If storage_client methods should be async, update them to be async and use await here.
3. If synchronous calls are unavoidable, run them in an executor to avoid blocking:

async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
    keys = self._generate_yr_keys(metadata)
    values = self._generate_yr_values(data)
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, self.storage_client.put, keys, values)

🛠️ Refactor suggestion | 🟠 Major

Add error handling for storage client operations.

None of the data operations have error handling. If storage_client operations fail (network issues, storage full, etc.), exceptions will propagate unhandled. Add try-except blocks with appropriate logging and error recovery.

Example for put_data:

 async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
-    keys = self._generate_yr_keys(metadata)
-    values = self._generate_yr_values(data)
-    self.storage_client.put(keys=keys, values=values)
+    try:
+        keys = self._generate_yr_keys(metadata)
+        values = self._generate_yr_values(data)
+        self.storage_client.put(keys=keys, values=values)
+    except Exception as e:
+        logger.error(f"[{self.storage_manager_id}]: Failed to put data: {e}")
+        raise

Apply similar error handling to get_data and clear_data.

Committable suggestion skipped: line range outside the PR's diff.

ignored

@dpj135, understood! I've noted that you're choosing not to implement these suggestions for now. Feel free to reach out if you'd like to discuss them in the future.

(\_/)
( •_•) 
/ >💔

⚠️ Potential issue | 🟠 Major

Missing data update notification after put.

After successfully storing data, put_data should call notify_data_update (inherited from the base class) to inform the controller that new data is available. This is essential for the queue coordination workflow.

Add notification after the put operation:

 async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
     keys = self._generate_yr_keys(metadata)
     values = self._generate_yr_values(data)
     self.storage_client.put(keys=keys, values=values)
+    
+    # Notify controller of new data
+    dtypes = {idx: {} for idx in metadata.global_indexes}
+    shapes = {idx: {} for idx in metadata.global_indexes}
+    for idx, sample in enumerate(metadata.samples):
+        for field_name in metadata.field_names:
+            field = sample.get_field_by_name(field_name)
+            dtypes[metadata.global_indexes[idx]][field_name] = field.dtype
+            shapes[metadata.global_indexes[idx]][field_name] = field.shape
+    
+    await self.notify_data_update(
+        fields=list(metadata.field_names),
+        global_indexes=metadata.global_indexes,
+        dtypes=dtypes,
+        shapes=shapes,
+    )

In transfer_queue/storage/managers/base.py around lines 407 to 410, put_data
stores keys/values but does not notify the controller; after successfully
calling self.storage_client.put(keys=keys, values=values) call await
self.notify_data_update(metadata) so the inherited notification is sent to the
controller; ensure the notify call is awaited (and only run after put completes)
so the queue coordination workflow is triggered.

done

@dpj135, thank you for implementing the fix! I can confirm that the put_data method now properly calls notify_data_update after storing the data, which ensures the controller is informed about the new data availability. The implementation looks correct, with the dtypes and shapes being properly extracted and passed to the notification method. Great work! ✨

🐰

This interface seems to fit KV storage client. We should call it TransferQueueStorageKVClient. And this class should be placed in a separated file named base.py

And a general base class is needed

done

Need to translate it into English

done

Unnecessary comments should be removed

done

Why we need these codes?

done

Now this function is a little hard to understand. The input params is metadata and values, which are not keys and values.

Can we propose a better calling logic of its upstream function? Or simply change the name of it?

    async def get_data(self, metadata: BatchMeta) -> TensorDict:
        keys = self._generate_keys(metadata)
        shapes, dtypes = self._get_shape_type_list(metadata)
        values = self.storage_client.get(keys=keys, shapes=shapes, dtypes=dtypes)
        return self._merge_kv_to_tensordict(metadata, values)

could we call it _merge_tensors_to_tensordict?

    async def get_data(self, metadata: BatchMeta) -> TensorDict:
        keys = self._generate_keys(metadata)
        shapes, dtypes = self._get_shape_type_list(metadata)
        values = self.storage_client.get(keys=keys, shapes=shapes, dtypes=dtypes)
        return self._merge_tensors_to_tensordict(metadata, values)