Skip to content

Commit 8f3417c

Browse files
committed
Modified annotations about clients
Signed-off-by: dpj135 <958208521@qq.com>
1 parent 20ae39b commit 8f3417c

2 files changed

Lines changed: 13 additions & 15 deletions

File tree

transfer_queue/storage/clients/mooncake_client.py

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -135,18 +135,14 @@ def _batch_put_bytes(self, keys: list[str], values: list[bytes]):
135135
def get(self, keys: list[str], shapes=None, dtypes=None, custom_backend_meta=None) -> list[Any]:
136136
"""Get multiple key-value pairs from MooncakeStore.
137137
138-
Args:
139-
keys (List[str]): Keys to fetch.
140-
shapes (List[List[int]]): Expected tensor shapes (use [] for scalars).
141-
dtypes (List[Optional[torch.dtype]]): Expected dtypes; use None for non-tensor data.
142-
<<<<<<< HEAD
143-
custom_backend_meta (List[str], optional): Device type (npu/cpu) for each key
144-
=======
145-
custom_meta (List[Any], optional): ...
146-
>>>>>>> acd7686 (Added custom_meta to clear for all TransferQueueKVStorageClient)
147-
148-
Returns:
149-
List[Any]: Retrieved values in the same order as input keys.
138+
Args:
139+
keys (List[str]): Keys to fetch.
140+
shapes (List[List[int]]): Expected tensor shapes (use [] for scalars).
141+
dtypes (List[Optional[torch.dtype]]): Expected dtypes; use None for non-tensor data.
142+
custom_backend_meta (List[str], optional): ...
143+
144+
Returns:
145+
List[Any]: Retrieved values in the same order as input keys.
150146
"""
151147

152148
if shapes is None or dtypes is None:

transfer_queue/storage/clients/yuanrong_client.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,9 @@ def clear(self, keys: list[str]):
7676

7777

7878
class NPUTensorKVClientAdapter(StorageStrategy):
79-
"""Adapter for YuanRong's high-performance NPU tensor storage."""
79+
"""Adapter for YuanRong's high-performance NPU tensor storage.
80+
Using yr.datasystem.DsTensorClient to connect datasystem backends.
81+
"""
8082

8183
KEYS_LIMIT: int = 10_000
8284

@@ -179,8 +181,8 @@ def _create_empty_npu_tensorlist(self, shapes, dtypes):
179181

180182

181183
class GeneralKVClientAdapter(StorageStrategy):
182-
"""
183-
Adapter for general-purpose KV storage with serialization.
184+
"""Adapter for general-purpose KV storage with serialization.
185+
Using yr.datasystem.KVClient to connect datasystem backends.
184186
The serialization method uses '_decoder' and '_encoder' from 'transfer_queue.utils.serial_utils'.
185187
"""
186188

0 commit comments

Comments
 (0)