[feat] Provide user-defined `custom_meta` methods by 0oshowero0 · Pull Request #21 · Ascend/TransferQueue

0oshowero0 · 2026-01-30T11:59:32Z

Background

When integrating with upstream RL frameworks, it is often necessary to maintain some sample-level metadata alongside the batch data for computation. Typical examples include:

uid
session_id
trajectory_id
start_model_version/end_model_version
current_status
...

This PR introduces a flexible User-Defined Metadata mechanism. Users can now attach arbitrary key-value pairs (custom_meta) to BatchMeta and synchronize them with the TransferQueueController.

Note: Different from custom_meta that maintains sample-level info, the extra_info stores batch-level info and will not be stored into the controller.

Key Changes

BatchMeta Enhancement: Added interfaces to BatchMeta for setting and retrieving user-defined custom_meta.
Put custom_meta:
- Explicit Sync: Metadata can be sent to the controller via TransferQueueClient.set_custom_meta.
- Implicit Sync: Metadata is automatically stored when calling TransferQueueClient.put with given batch_meta input.
Get custom_meta: Users can query the stored custom_meta from the controller using get_meta with mode="force_fetch".

We refactor the previous meta for storage backend as _custom_backend_meta, which is a field-level information.

Related API

BatchMeta

set_custom_meta(global_index:int, value:Any)
Insert or update a single key-value pair for a specific index.
update_custom_meta(new_meta:dict[int, Any])
Batch update custom_meta using a dictionary.
get_all_custom_meta() -> dict[int, Any]
Retrieve all stored custom_meta.

TransferQueueClient

(async_)set_custom_meta(metadata: BatchMeta)
Explicitly send custom_meta to the TransferQueueController
(async_)put(data:TensorDict, metadata: BatchMeta)
Implicitly send custom_meta to the controller while putting data.
(async_)get_meta(partition_id:str, mode="force_fetch")
Query and retrieve all BatchMeta alone with custom_meta stored in the controller (requires mode="force_fetch").

Usage Example

Refer to tutorial/02_metadata_concepts.py.

Other Changes

Add CI for tutorial scripts.
Improve docstring for TransferQueueClient

TODO

Move custom_meta -> SampleMeta, _custom_backend_meta -> FieldMeta

CC: @wuxibin89 @tianyi-ge

Copilot

Pull request overview

This PR introduces a user-defined metadata mechanism allowing users to attach arbitrary key-value pairs (custom_meta) to samples in BatchMeta. The implementation distinguishes between:

custom_meta: User-defined per-sample metadata
_custom_backend_meta: Backend-specific per-sample per-field metadata (renamed from field_custom_metas)

Changes:

Added custom_meta field and related methods (set_custom_meta, update_custom_meta, get_all_custom_meta, clear_custom_meta) to BatchMeta
Renamed backend metadata fields from _custom_meta/field_custom_metas to _custom_backend_meta/field_custom_backend_meta throughout the codebase
Added set_custom_meta and get_custom_meta methods to DataPartitionStatus and TransferQueueController
Added chunk_by_partition method to split batches by partition_id
Added SET_CUSTOM_META request type and client methods (async_set_custom_meta, set_custom_meta)
Implicit sync of custom_meta when calling client.put()

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
transfer_queue/utils/zmq_utils.py	Added SET_CUSTOM_META and SET_CUSTOM_META_RESPONSE request types
transfer_queue/metadata.py	Added custom_meta field and methods to BatchMeta, renamed _custom_backend_meta, added chunk_by_partition method, updated post_init to filter metadata
transfer_queue/controller.py	Renamed field_custom_metas to field_custom_backend_meta, added custom_meta storage and methods to DataPartitionStatus and controller, updated generate_batch_meta to include both metadata types
transfer_queue/client.py	Added async_set_custom_meta and set_custom_meta methods, implicit sync in async_put, updated documentation for force_fetch mode
tests/test_metadata.py	Added comprehensive tests for custom_meta methods, reorganized test structure
tests/test_controller_data_partitions.py	Updated tests to use new field_custom_backend_meta naming, added tests for custom_meta methods
tests/test_controller.py	Added integration tests for custom_meta and custom_backend_meta functionality
tests/test_client.py	Added tests for set_custom_meta client interface methods

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

transfer_queue/controller.py

transfer_queue/metadata.py

transfer_queue/controller.py

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 10 comments.

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

Copilot · 2026-02-02T01:27:46Z

transfer_queue/metadata.py

+        for partition_id, global_index in zip(self.partition_ids, self.global_indexes, strict=False):
+            grouped_global_indexes[partition_id].append(global_index)
+
+        chunk_list = [self.select_samples(samples) for samples in grouped_global_indexes.values()]


The chunk_by_partition method uses select_samples which does not preserve custom_meta and _custom_backend_meta. When samples are selected, the custom metadata is lost. The implementation should pass custom_meta and _custom_backend_meta when creating new BatchMeta instances.

CC @tianyi-ge

transfer_queue/metadata.py

transfer_queue/controller.py

transfer_queue/client.py

tutorial/02_metadata_concepts.py

tests/test_controller.py

transfer_queue/client.py

tutorial/02_metadata_concepts.py

.github/workflows/tutorial-check.yml

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 8 comments.

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

transfer_queue/controller.py

tutorial/02_metadata_concepts.py

transfer_queue/metadata.py

transfer_queue/client.py

transfer_queue/metadata.py

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 5 comments.

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

transfer_queue/metadata.py

Copilot · 2026-02-02T06:47:21Z

transfer_queue/metadata.py

    def __getitem__(self, item):
        if isinstance(item, int | np.integer):
            sample_meta = self.samples[item] if self.samples else []
-            return BatchMeta(samples=[sample_meta], extra_info=self.extra_info)
+            global_idx = self.global_indexes[item]
+            return BatchMeta(
+                samples=[sample_meta],
+                extra_info=self.extra_info,
+                custom_meta={global_idx: self.custom_meta[global_idx]},
+                _custom_backend_meta={global_idx: self._custom_backend_meta[global_idx]},
+            )


Missing test coverage for the __getitem__ method with custom_meta and _custom_backend_meta. Since this method was modified to handle these new fields, tests should verify that indexing works correctly when these fields are present or absent.

tianyi-ge · 2026-02-02T07:27:23Z

transfer_queue/metadata.py


        # construct new SampleMeta instance
-        # TODO(tianyi): move custom_meta to FieldMeta level
+        # TODO(tianyi): (maybe) move _custom_backend_meta and custom_meta to FieldMeta level?


If _custom_backend_meta is in FieldMeta, some details need to be modified. For example, if users call BatchMeta.add_fields after put, it _extract_field_metas and then add_fields -> _union_fields, which overwrites the _custom_backend_meta and cannot be recovered

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 8 comments.

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

transfer_queue/client.py

tests/test_metadata.py

transfer_queue/metadata.py

transfer_queue/controller.py

transfer_queue/metadata.py

transfer_queue/controller.py

transfer_queue/client.py

wuxibin89 · 2026-02-02T08:59:13Z

We should support put data with metadata?

tq.put(data={"prompt_ids": torch.Tensor(), "response_ids": torch.Tensor(), ...}, metadata={"uid": "xxx", "session_id": "xxx", "status": "RUNNING"})

wuxibin89 · 2026-02-02T09:19:56Z

transfer_queue/metadata.py

-    _custom_meta: dict[int, dict[str, Any]] = dataclasses.field(default_factory=dict)
+
+    # user-defined meta for each sample
+    custom_meta: dict[int, dict[str, Any]] = dataclasses.field(default_factory=dict)


Better move custom_meta into SampleMeta?

class SampleMeta: """Records the metadata of a single data sample (stored as a row in the data system).""" partition_id: str # Partition id, used for data versioning uid: str # 原始prompt的uid session_id: int = 0 # 单个prompt n多次采样，每个采样对应一个AgentLoop trajectory_id: int = 0 # 每次AgentLoop可能有多个输出：prefix切换 extra_metadata: dict[str, Any] = dataclasses.field(default_factory=dict) # global_index: int # Global row index, uniquely identifies a data sample fields: dict[str, FieldMeta] # Fields of interest for this sample

Great point. Actually this is planned in following PR. https://github.com/Ascend/TransferQueue/pull/21/changes#diff-a1fde8d325ec343f7422450e654b20e60cdd34497c772d2a8d76d0874e3cf2edR140

To keep SampleMeta universal and avoid hardcoding specific fields, we are introducing a generic dictionary to store these information.

0oshowero0 · 2026-02-02T12:16:38Z

We should support put data with metadata?

tq.put(data={"prompt_ids": torch.Tensor(), "response_ids": torch.Tensor(), ...}, metadata={"uid": "xxx", "session_id": "xxx", "status": "RUNNING"})

That's right. Now this can be done by the following usage example:

    custom_meta = {
        global_index: {"uid": uuid.uuid4().hex[:4], "session_id": uuid.uuid4().hex[:4], "model_version": 0}
        for global_index in batch_meta.global_indexes
    }
    batch_meta.update_custom_meta(custom_meta)

    tq_client.put(data, batch_meta)

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> # Conflicts: # transfer_queue/client.py

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Copilot AI review requested due to automatic review settings January 30, 2026 11:59

Copilot started reviewing on behalf of 0oshowero0 January 30, 2026 11:59 View session

Copilot AI reviewed Jan 30, 2026

View reviewed changes

0oshowero0 force-pushed the custom_meta branch from d955b85 to 7405c1b Compare January 30, 2026 12:30

0oshowero0 mentioned this pull request Jan 31, 2026

[refactor] Refactor yuanrong_client #18

Merged

5 tasks

0oshowero0 requested a review from Copilot February 2, 2026 01:22

Copilot started reviewing on behalf of 0oshowero0 February 2, 2026 01:22 View session

Copilot AI reviewed Feb 2, 2026

View reviewed changes

0oshowero0 requested a review from Copilot February 2, 2026 03:12

Copilot started reviewing on behalf of 0oshowero0 February 2, 2026 03:12 View session

Copilot AI reviewed Feb 2, 2026

View reviewed changes

0oshowero0 requested a review from Copilot February 2, 2026 06:42

Copilot started reviewing on behalf of 0oshowero0 February 2, 2026 06:43 View session

Copilot AI reviewed Feb 2, 2026

View reviewed changes

tianyi-ge reviewed Feb 2, 2026

View reviewed changes

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

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

Copilot AI reviewed Feb 2, 2026

View reviewed changes

wuxibin89 reviewed Feb 2, 2026

View reviewed changes

0oshowero0 added 9 commits February 3, 2026 09:13

try

30fdc31

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

try

d2d8c3e

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

support custom_meta

9a1e75b

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

add CI

dfb07e5

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

update batchmeta & ci

b6b0d00

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

098b7d1

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

add tutorial & update CI

02e1b48

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

2b5eb43

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

update var name

177fe5b

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

0oshowero0 added 12 commits February 3, 2026 09:13

fix

b3b1abe

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix ut

f5c95ac

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix metadata chunk & concat & union

4ed69ef

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> # Conflicts: # transfer_queue/client.py

modify select_samples by global_index

a01e6bd

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

bdadf4c

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix param name

339d9ed

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

optimize docstring

44fcea0

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix ci

e66561e

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

bc34a6a

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

f422aeb

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

minor update

abf2112

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

c2d75ce

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

0oshowero0 force-pushed the custom_meta branch from 07ed406 to c2d75ce Compare February 3, 2026 01:19

0oshowero0 added 2 commits February 3, 2026 09:22

update

c5ea293

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

45394fd

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

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

0oshowero0 deleted the custom_meta branch February 3, 2026 01:55

Conversation

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

Uh oh!

Background

Key Changes

Related API

BatchMeta

TransferQueueClient

Usage Example

Other Changes

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!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

0oshowero0 Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

tianyi-ge Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wuxibin89 commented Feb 2, 2026

Uh oh!

0oshowero0 commented Jan 30, 2026 •

edited

Loading