Skip to content

Commit f048354

Browse files
author
ascend-robot
committed
[feat] Introduce zero-copy to use YuanrongStorageClient for transmitting CPU Tensors in NPU environment
Co-authored-by: liwenlin<liwenlin8@huawei.com> # message auto-generated for no-merge-commit merge: !24 merge npu_zcopy into main [feat] Introduce zero-copy to use `YuanrongStorageClient` for transmitting CPU Tensors in NPU environment Created-by: Lexie-7 Commit-by: liwenlin Merged-by: ascend-robot Description: ### Summary Transmitting CPU tensor in the NPU environment and achieving zero-copy. ### Change Modified the `transfer_queue/storage/clients/yuanrong_client.py` , enable zero-copy transmission of CPU data In the NPU branch. See merge request: Ascend/TransferQueue!24
1 parent 4df106b commit f048354

1 file changed

Lines changed: 4 additions & 4 deletions

File tree

transfer_queue/storage/clients/yuanrong_client.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,7 @@ def _batch_put(self, keys: list[str], values: list[Any]):
260260
for i in range(0, len(cpu_keys), CPU_DS_CLIENT_KEYS_LIMIT):
261261
batch_keys = cpu_keys[i : i + CPU_DS_CLIENT_KEYS_LIMIT]
262262
batch_values = cpu_values[i : i + CPU_DS_CLIENT_KEYS_LIMIT]
263-
self._cpu_ds_client.mset(batch_keys, batch_values)
263+
self.mset_zcopy(batch_keys, batch_values)
264264

265265
else:
266266
# All data goes through CPU path
@@ -356,9 +356,9 @@ def _batch_get(self, keys: list[str], shapes: list, dtypes: list) -> list[Any]:
356356
for i in range(0, len(cpu_keys), CPU_DS_CLIENT_KEYS_LIMIT):
357357
batch_keys = cpu_keys[i : i + CPU_DS_CLIENT_KEYS_LIMIT]
358358
batch_indices = cpu_indices[i : i + CPU_DS_CLIENT_KEYS_LIMIT]
359-
raw_values = self._cpu_ds_client.get(batch_keys)
360-
for idx, raw_val in zip(batch_indices, raw_values, strict=False):
361-
results[idx] = pickle.loads(raw_val)
359+
objects = self.mget_zcopy(batch_keys)
360+
for idx, obj in zip(batch_indices, objects, strict=False):
361+
results[idx] = obj
362362

363363
return results
364364

0 commit comments

Comments
 (0)