[refactor] Refactor yuanrong_client by dpj135 · Pull Request #18 · Ascend/TransferQueue

dpj135 · 2026-01-28T09:36:00Z

Background:

Currently, the code structure of yuanrong_client.py is complex. The put/get operations of npu_ds_client and cpu_ds_client depend on different tool functions, global constants, and a large number of external dependencies.
In addition, yuanrong_client.py may be optimized in the future, or the interfaces of the data system may be updated or adjusted. As a result, the yuanrong_client.YuanrongStorageClient will be modified in a shotgun manner and changes will be divergent.

Description

Refactor using the Adapter and Strategy patterns.

The interface StorageStrategy is added, which provides a series of abstract methods for encapsulating the yr.datasystem interface.
The two storage paths of the original YuanrongStorageClient are extracted into two new adapter classes: DsTensorClientAdapter and KVClientAdapter.
YuanrongStorageClient is now responsible for dynamically routes data and schedules DsTensorClientAdapter and KVClientAdapter using the strategy pattern.

Todo:

Implement interface StorageStartegy.
Implement adapter class DsTensorClientAdapter and KVClientAdapter.
Add a parameter 'custom_meta' for All TransferQueueStorageClient's clear.
Implement parallelism of methods YuanrongStorageClient::put , YuanrongStorageClient::get and YuanrongStorageClient::clear.
Add a unit test.

Copilot

Pull request overview

This PR refactors the yuanrong_client.py module using the Adapter and Strategy design patterns to improve code organization, maintainability, and enable parallel operations across different storage backends (NPU tensor storage and general KV storage).

Changes:

Introduced StorageStrategy abstract base class with two concrete implementations: DsTensorClientAdapter for high-performance NPU tensor storage and KVClientAdapter for general-purpose serialized object storage
Added custom_meta parameter to all storage client clear() methods to enable proper routing of delete operations based on storage backend type
Implemented parallel execution of put/get/clear operations when both NPU and CPU strategies are active, using ThreadPoolExecutor

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 20 comments.

Show a summary per file

File	Description
transfer_queue/storage/clients/yuanrong_client.py	Major refactor introducing Strategy pattern with adapters for DsTensorClient and KVClient; consolidated batching logic; added parallel dispatch mechanism
transfer_queue/storage/clients/base.py	Updated `get()` and `clear()` method signatures to include `custom_meta` parameter with Optional type annotation
transfer_queue/storage/managers/base.py	Modified `clear_data()` to extract and pass custom_meta to storage clients; added comment clarifying put operation
transfer_queue/storage/clients/ray_storage_client.py	Updated `clear()` signature to accept custom_meta parameter (for interface compliance)
transfer_queue/storage/clients/mooncake_client.py	Updated `get()` and `clear()` signatures to accept custom_meta parameter (for interface compliance)
tests/test_yuanrong_storage_client_e2e.py	New comprehensive E2E test suite with mocked backends testing CPU-only, NPU-only, and mixed data flows
tests/test_yuanrong_client_zero_copy.py	Updated to test KVClientAdapter directly instead of the full YuanrongStorageClient

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

transfer_queue/storage/clients/yuanrong_client.py

tests/test_yuanrong_storage_client_e2e.py

Copilot · 2026-01-30T10:40:58Z

transfer_queue/storage/clients/yuanrong_client.py

-    def _batch_put(self, keys: list[str], values: list[Any]):
-        """Stores a batch of key-value pairs to remote storage, splitting by device type.
+@StorageClientFactory.register("YuanrongStorageClient")
+class YuanrongStorageClient(TransferQueueStorageKVClient):


This class does not call TransferQueueStorageKVClient.init during initialization. (YuanrongStorageClient.init may be missing a call to a base class init)

TransferQueueStorgaeKVClient is an abstract class, it has no Actual data. So ignore it here

Better add this for future compatibility

it's a python convention

tests/test_yuanrong_storage_client_e2e.py

tianyi-ge · 2026-01-31T02:12:57Z

transfer_queue/storage/clients/yuanrong_client.py

-    Args:
-        items: List of memoryview objects to be packed.
+    @abstractmethod
+    def custom_meta(self) -> Any:


maybe strategy_tag?

Now I'm working on the refactor of custom_meta. After the refactor #21 , custom_meta refer to the sample-level information that explicitly provided by the users; while the custom_backend_meta refer to field-level metadata that automatically set by storage backends.

tianyi-ge · 2026-01-31T02:15:22Z

transfer_queue/storage/clients/yuanrong_client.py

+        for meta_str, indexes in self._dispatch_tasks(routed_indexes, put_task):
+            for i in indexes:
+                custom_meta[i] = meta_str
+        return custom_meta


shorter tag string improves performance

in my opinion, custom_meta is a general concept from the perspective of metadata. here we can focus on what it really means in yuanrong client (strategy tag or device type, etc)

transfer_queue/storage/clients/yuanrong_client.py

tianyi-ge · 2026-01-31T02:28:41Z

transfer_queue/storage/clients/yuanrong_client.py

-        """Check if NPU client is available."""
-        return self._npu_ds_client is not None
+            torch_npu_imported = False
+        enable = config.get("enable_yr_npu_optimization", True)


this option is to allow users to disable ds tensor client even if in npu environment? which users will want this?

It's convenient for some users who dont want to configure a complex environment, and for developper to test.

I would recommend enable_yr_npu_transport

0oshowero0 · 2026-01-31T03:10:55Z

transfer_queue/storage/clients/yuanrong_client.py

+
+        return DsTensorClientAdapter(config)
+
+    def custom_meta(self) -> Any:


See above comments

0oshowero0 · 2026-01-31T03:14:14Z

transfer_queue/storage/clients/yuanrong_client.py


-    def mset_zcopy(self, keys: list[str], objs: list[Any]):
+
+class KVClientAdapter(StorageStrategy):


From user side, it's a little bit hard to understand the difference between these two Adapters. Maybe GeneralKVClientAdapter and NPUTensorKVClientAdapter?

0oshowero0 · 2026-01-31T03:18:11Z

transfer_queue/storage/clients/yuanrong_client.py

+        for i, item in enumerate(items):
+            for strategy in self._strategies:
+                if selector(strategy, item):
+                    routed_indexes[strategy].append(i)


It seems that a single sample can have multiple routed_indexes? Need to make sure only one strategy is actually executed according to some priority setting.

'if ... break' proves one-to-one for sample_index and storage_strategy

0oshowero0 · 2026-01-31T03:19:43Z

transfer_queue/storage/clients/yuanrong_client.py

        """
-        self._batch_clear(keys)
+        routed_indexes: dict[StorageStrategy, list[int]] = {s: [] for s in self._strategies}
+        for i, item in enumerate(items):


Try to prevent $O(m\times n)$ loop?

len(_strategies) is 2 now. In fact, this loop is O(m)

len(self._strategies) is a constant.

0oshowero0 · 2026-01-31T03:22:18Z

transfer_queue/storage/clients/yuanrong_client.py

+        Used to route keys/values/custom_meta to storage backends by grouped indexes.
+
+        Args:
+            items: A list of items (e.g., values for put, or custom_meta strings for get/clear).


Better to explicitly distinguish the use case of custom_meta input and ordinary value input

dpj135 · 2026-02-02T06:40:47Z

@0oshowero0 @tianyi-ge Comments are addressed. Please take a look. Than you!

Copilot

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 19 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

transfer_queue/storage/clients/yuanrong_client.py

tests/test_yuanrong_storage_client_e2e.py

transfer_queue/storage/clients/yuanrong_client.py

tests/test_yuanrong_storage_client_e2e.py

Copilot · 2026-02-02T09:01:12Z

transfer_queue/storage/clients/yuanrong_client.py

-                    results[idx] = obj
-            return results
+        # Dispatch tasks and map metadata back to original positions
+        custom_meta: list[str] = [""] * len(keys)


The type annotation list[str] is inconsistent with the actual data being stored. The custom_meta list stores bytes values (from strategy.strategy_tag() which returns bytes), not strings. The type annotation should be list[bytes] to match the actual type being stored.

Copilot · 2026-02-02T09:01:13Z

transfer_queue/storage/clients/yuanrong_client.py

        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")


The put method should validate that the keys list is not empty. Currently, it checks if keys and values are lists and if their lengths match, but doesn't check for empty lists. Calling put with empty lists would result in returning an empty custom_meta list, which may not be the intended behavior. Consider adding validation or documenting that empty lists are allowed.

Suggested change

raise ValueError("Number of keys must match number of values")

raise ValueError("Number of keys must match number of values")

if not keys:

raise ValueError("keys and values must be non-empty lists")

transfer_queue/storage/clients/yuanrong_client.py

tests/test_yuanrong_storage_client_e2e.py

tianyi-ge · 2026-02-02T14:56:43Z

transfer_queue/storage/clients/yuanrong_client.py

-        offset, length = struct.unpack_from(ENTRY_FMT, mv, HEADER_SIZE + i * ENTRY_SIZE)
-        offsets.append((offset, length))
-    return [mv[offset : offset + length] for offset, length in offsets]
+    KEYS_LIMIT: int = 10_000


according to datasystem doc, key>64 can lead to significant performance degradation. please test it for both cpu/npu

Based on current simple test, I find that the performance of DsTensorClient is minimally affected by KEYS_LIMIT, and KVClient is greatly affected by KEYS_LIMIT. From the results, it generally seems that the bigger, the better.

I doubt that the performance degradation of KVClient(adapter class) may be due to thread overhead.

def mset_zero_copy(self, keys: list[str], objs: list[Any]): """Store multiple objects in zero-copy mode using parallel serialization and buffer packing. Args: keys (list[str]): List of string keys under which the objects will be stored. objs (list[Any]): List of Python objects to store (e.g., tensors, strings). """ items_list = [[memoryview(b) for b in _encoder.encode(obj)] for obj in objs] packed_sizes = [self.calc_packed_size(items) for items in items_list] buffers = self._ds_client.mcreate(keys, packed_sizes) tasks = [(target.MutableData(), item) for target, item in zip(buffers, items_list, strict=True)] with ThreadPoolExecutor(max_workers=self.DS_MAX_WORKERS) as executor: list(executor.map(lambda p: self.pack_into(*p), tasks)) self._ds_client.mset_buffer(buffers)

I think we could consider first using multi-threading to execute pack_into, and then set them in batches. @Evelynn-V

tianyi-ge · 2026-02-02T14:58:19Z

transfer_queue/storage/clients/yuanrong_client.py

-        """Check if NPU client is available."""
-        return self._npu_ds_client is not None
+            torch_npu_imported = False
+        enable = config.get("enable_yr_npu_optimization", True)


I would recommend enable_yr_npu_transport

tianyi-ge · 2026-02-02T15:00:30Z

transfer_queue/storage/clients/yuanrong_client.py

+
+    def supports_clear(self, custom_meta: str) -> bool:
+        """Matches 'DsTensorClient' strategy tag."""
+        return isinstance(custom_meta, bytes) and custom_meta == self.strategy_tag()


since you need to compare strategy tag with custom meta, just use "0" and "1" instead of b"\x01" and b"\x02"

tianyi-ge · 2026-02-02T15:01:38Z

transfer_queue/storage/clients/yuanrong_client.py

-    def _batch_put(self, keys: list[str], values: list[Any]):
-        """Stores a batch of key-value pairs to remote storage, splitting by device type.
+@StorageClientFactory.register("YuanrongStorageClient")
+class YuanrongStorageClient(TransferQueueStorageKVClient):


it's a python convention

tianyi-ge · 2026-02-02T15:05:25Z

transfer_queue/storage/clients/yuanrong_client.py

-                    results[idx] = obj
-            return results
+        # Dispatch tasks and map metadata back to original positions
+        custom_meta: list[str] = [""] * len(keys)


#21 is renaming custom_meta to custom_backend_meta, and custom_meta has other meanings. I recommend to rename all the custom_meta here to strategy_tags so that 1. we focus on what it really means 2. other developers won't be confused

I renamed most of the custom names. However, considering the responsibility of this PR is to refactor yuanrong_client, I did not change the function signatures of get and clear in YuanrongStorageClient to remain consistent with the abstract interface.

…client.py' Signed-off-by: dpj135 <958208521@qq.com>

Signed-off-by: dpj135 <958208521@qq.com>

…he order of classes Signed-off-by: dpj135 <958208521@qq.com>

…ong_client can execute correctly) Signed-off-by: dpj135 <958208521@qq.com>

Signed-off-by: dpj135 <958208521@qq.com>

…geClient' Signed-off-by: dpj135 <958208521@qq.com>

Signed-off-by: dpj135 <958208521@qq.com>

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: dpj135 <958208521@qq.com>

…& adjusted annotation related to 'custom_meta()' Signed-off-by: dpj135 <958208521@qq.com>

Signed-off-by: dpj135 <958208521@qq.com>

…tom_name, adjusted annotations ...) Signed-off-by: dpj135 <958208521@qq.com>

Signed-off-by: dpj135 <958208521@qq.com>

tianyi-ge · 2026-02-03T05:24:53Z

look good to merge

Signed-off-by: dpj135 <958208521@qq.com>

dpj135 force-pushed the refactor_yuanrongclient branch from 1533adb to d3720fd Compare January 29, 2026 06:42

dpj135 marked this pull request as ready for review January 30, 2026 05:03

dpj135 force-pushed the refactor_yuanrongclient branch from 143251e to 4db6344 Compare January 30, 2026 06:16

0oshowero0 requested a review from Copilot January 30, 2026 10:34

Copilot started reviewing on behalf of 0oshowero0 January 30, 2026 10:34 View session

Copilot AI reviewed Jan 30, 2026

View reviewed changes

tianyi-ge reviewed Jan 31, 2026

View reviewed changes

0oshowero0 reviewed Jan 31, 2026

View reviewed changes

dpj135 force-pushed the refactor_yuanrongclient branch from 3901b02 to 2aa3955 Compare February 2, 2026 04:59

dpj135 requested review from 0oshowero0 and tianyi-ge February 2, 2026 07:24

0oshowero0 requested a review from Copilot February 2, 2026 08:53

Copilot started reviewing on behalf of 0oshowero0 February 2, 2026 08:53 View session

Copilot AI reviewed Feb 2, 2026

View reviewed changes

tianyi-ge reviewed Feb 2, 2026

View reviewed changes

dpj135 added 10 commits February 3, 2026 11:07

Renamed 'test_yuanrong_storage_manager.py' to 'test_yuanrong_storage_…

3338e37

…client.py' Signed-off-by: dpj135 <958208521@qq.com>

Added abstract interface 'StorageStrategy'

146cbd2

Signed-off-by: dpj135 <958208521@qq.com>

Added DsTensorClient

4d184b8

Signed-off-by: dpj135 <958208521@qq.com>

Added 'KVClientAdapter'

d5ec724

Signed-off-by: dpj135 <958208521@qq.com>

Refactored 'YuanrongStorageClient.put&get'

28320da

Signed-off-by: dpj135 <958208521@qq.com>

Added 'route_to_strategy' to class 'YuanrongStorageClient' & Adjust t…

8c8e9a2

…he order of classes Signed-off-by: dpj135 <958208521@qq.com>

Fixed the order about '@staticmethod' and 'abstractmethod' (Now yuanr…

4d8ae64

…ong_client can execute correctly) Signed-off-by: dpj135 <958208521@qq.com>

Added custom_meta to clear for all TransferQueueKVStorageClient

a9235e6

Signed-off-by: dpj135 <958208521@qq.com>

Added multi-threads optimization to 'put/get/clear' of 'YuanrongStora…

ffa5970

…geClient' Signed-off-by: dpj135 <958208521@qq.com>

Added more annotation for methods

39c72bb

Signed-off-by: dpj135 <958208521@qq.com>

dpj135 and others added 14 commits February 3, 2026 11:19

Added end-to-end test(generated by AI) for 'YuanrongStorageClient'

ed07b61

Signed-off-by: dpj135 <958208521@qq.com>

Fixed tests about yuanrong_clint

91f2532

Signed-off-by: dpj135 <958208521@qq.com>

Added license to test_yuanrong_client

2faa9eb

Signed-off-by: dpj135 <958208521@qq.com>

Added method 'test_mock_can_work' to test_yuanrong_client

7531714

Signed-off-by: dpj135 <958208521@qq.com>

Added an annotation to class 'StorageStrategy'

a8293ef

Signed-off-by: dpj135 <958208521@qq.com>

Fixed test_yuanrong_client_zero_copy

e120eb2

Signed-off-by: dpj135 <958208521@qq.com>

Apply suggestions from code review

c0fc536

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: dpj135 <958208521@qq.com>

Renamed adapter classes & rename 'custom_meta()' to 'strategy_tag()' …

c196940

…& adjusted annotation related to 'custom_meta()' Signed-off-by: dpj135 <958208521@qq.com>

Fixed 'KVClientAdapter' imported error

50c5284

Signed-off-by: dpj135 <958208521@qq.com>

Modified docstrings

a7303cf

Signed-off-by: dpj135 <958208521@qq.com>

Fixed 'test_yuanrong_storage_client_e2e.py' about strategy_tag

1fd53f4

Signed-off-by: dpj135 <958208521@qq.com>

Adjusted annotations of test_yuanrong_storage_client_e2e.py

7bb83a8

Signed-off-by: dpj135 <958208521@qq.com>

Fixed reviews about yuanrong_client(modified strategy_tag, rename cus…

9fac913

…tom_name, adjusted annotations ...) Signed-off-by: dpj135 <958208521@qq.com>

Rename custom_meta to custom_backend_meta

20ae39b

Signed-off-by: dpj135 <958208521@qq.com>

dpj135 force-pushed the refactor_yuanrongclient branch from ddcae5e to 20ae39b Compare February 3, 2026 03:26

Modified annotations about clients

8f3417c

Signed-off-by: dpj135 <958208521@qq.com>

dpj135 force-pushed the refactor_yuanrongclient branch from b499c07 to 8f3417c Compare February 3, 2026 07:16

Adjusted expression of annotations and renamed one variable

4a3be0f

Signed-off-by: dpj135 <958208521@qq.com>

0oshowero0 approved these changes Feb 3, 2026

View reviewed changes

0oshowero0 merged commit fbdb58e into Ascend:main Feb 3, 2026
5 checks passed


		return DsTensorClientAdapter(config)

		def custom_meta(self) -> Any:


		def mset_zcopy(self, keys: list[str], objs: list[Any]):

		class KVClientAdapter(StorageStrategy):

Conversation

dpj135 commented Jan 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Background:

Description

Todo:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dpj135 commented Feb 2, 2026

Uh oh!

Copilot AI left a comment

dpj135 commented Jan 28, 2026 •

edited

Loading