[feat,CI] Improve KV API usability and `KVBatchMeta` interactions by 0oshowero0 · Pull Request #57 · Ascend/TransferQueue

0oshowero0 · 2026-03-25T02:03:41Z

📝 Description

This PR aims to improve the usability and flexibility of the KV (Key-Value) API. It optimizes how KVBatchMeta is handled during read and write operations, refactors some internal naming for clarity, and adds unit test coverage for the asynchronous interfaces to ensure stability.

🔄 What's Changed

Enhanced kv_put: The kv_put operation now directly returns a KVBatchMeta object containing all the available fields after the put operation.
Simplified kv_get: Provide an additional API (async_)kv_batch_get_by_meta that supports taking a KVBatchMeta object directly as input, simplifying the calling logic.
Flexible KVBatchMeta Fields: Introduced public interfaces that allow for more flexible and customizable field selection within KVBatchMeta.
Naming Refactor: Refactored the fields variables to select_fields during kv_batch_get to semantically distinguish from the fields (real data TensorDict) in kv_put.
Async UT Coverage: Added UT specifically for the async KV interfaces to ensure robustness.

ascend-robot · 2026-03-25T02:03:51Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Copilot

Pull request overview

This PR improves the high-level key-value (KV) interface to be more user-friendly by returning metadata from put operations, adding “get by meta” helpers, and renaming the “fields selection” parameter for clarity.

Changes:

kv_put / kv_batch_put (and async variants) now return KVBatchMeta containing keys/tags/partition and cumulative field names.
Added kv_batch_get_by_meta / async_kv_batch_get_by_meta convenience helpers.
Renamed kv_batch_get / async_kv_batch_get keyword argument from fields to select_fields, and updated tutorial + e2e tests (including sync/async parametrization).

Reviewed changes

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

File	Description
`tutorial/02_kv_interface.py`	Updates tutorial usage to `select_fields=...` for partial-field retrieval.
`transfer_queue/interface.py`	Updates KV API to return `KVBatchMeta`, adds get-by-meta helpers, and renames `fields` → `select_fields` for get calls.
`tests/e2e/test_kv_interface_e2e.py`	Expands e2e coverage for new return metadata semantics and runs the suite against both sync and async KV APIs via a wrapper fixture.

Comments suppressed due to low confidence (1)

transfer_queue/interface.py:388

In kv_put's docstring, the Raises section mentions “length of keys”, but this API accepts a single key. Updating the wording would keep the documentation accurate and reduce confusion for users.

        ValueError: If nested tensors are provided (use kv_batch_put instead)
        RuntimeError: If retrieved BatchMeta size doesn't match length of `keys`

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

Copilot · 2026-03-25T02:08:28Z

 async def async_kv_batch_get(
-    keys: list[str] | str, partition_id: str, fields: Optional[list[str] | str] = None
+    keys: list[str] | str, partition_id: str, select_fields: Optional[list[str] | str] = None
 ) -> TensorDict:


async_kv_batch_get also renames the public keyword argument from fields to select_fields, which will break callers using fields=.... Consider supporting fields as a deprecated alias here as well (and error if both are provided) for API compatibility.

Copilot · 2026-03-25T02:08:29Z

+def kv_batch_get(
+    keys: list[str] | str, partition_id: str, select_fields: Optional[list[str] | str] = None
+) -> TensorDict:


kv_batch_get renamed the public keyword argument from fields to select_fields. This is a backwards-incompatible API change for callers using fields=.... Consider accepting both (fields as a deprecated alias) and raising if both are provided, to avoid breaking existing users while still promoting select_fields.

ascend-robot · 2026-03-25T02:50:54Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Copilot

Pull request overview

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

Comments suppressed due to low confidence (1)

transfer_queue/interface.py:491

The error message says "fields or tag" but this function parameter is tags (plural). Update the message to avoid confusing users (especially now that kv_batch_put is part of the public KV API).

    if fields is None and tags is None:
        raise ValueError("Please provide at least one parameter of fields or tag.")

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

Copilot · 2026-03-25T02:57:27Z

+def kv_batch_get(
+    keys: list[str] | str, partition_id: str, select_fields: Optional[list[str] | str] = None
+) -> TensorDict:


kv_batch_get renamed the keyword argument from fields to select_fields, which is a breaking change for any callers using fields=. Consider keeping fields as a deprecated alias (e.g., accept both params and error if both are set) to preserve backward compatibility.

ascend-robot · 2026-03-25T03:11:47Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot · 2026-03-25T03:21:14Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

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

Copilot · 2026-03-25T03:21:30Z

        try:
-            retrieved = tq.kv_batch_get(keys="nonexistent_key", partition_id=partition_id)
+            retrieved = tq_api.kv_batch_get(keys="nonexistent_key", partition_id=partition_id)
            # If it returns, it should be empty
            assert retrieved.batch_size[0] == 0
        except RuntimeError as e:
            # Or it might raise an error about keys not found
            assert "not found" in str(e).lower() or "empty" in str(e).lower()


This test still catches RuntimeError, but kv_batch_get was updated to raise ValueError for missing keys/partitions. Update the assertion to expect ValueError (or broaden the exception type) so the test matches the new interface contract.

ascend-robot · 2026-03-25T03:33:04Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot · 2026-03-25T08:06:21Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot · 2026-03-25T08:24:06Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

Comments suppressed due to low confidence (1)

transfer_queue/interface.py:490

The error message says "fields or tag" but this function takes tags (plural). Updating the message to reference tags would make the API clearer for users.

    if fields is None and tags is None:
        raise ValueError("Please provide at least one parameter of fields or tag.")

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

Copilot · 2026-03-25T08:46:12Z

+def kv_batch_get(
+    keys: list[str] | str, partition_id: str, select_fields: Optional[list[str] | str] = None
+) -> TensorDict:


Renaming the public kwarg from fields to select_fields in kv_batch_get is a breaking change for any external callers using keyword arguments. Consider keeping fields as a deprecated alias (e.g., accept both and error if both are set) to preserve backwards compatibility while transitioning to select_fields.

Copilot · 2026-03-25T08:46:12Z

 async def async_kv_batch_get(
-    keys: list[str] | str, partition_id: str, fields: Optional[list[str] | str] = None
+    keys: list[str] | str, partition_id: str, select_fields: Optional[list[str] | str] = None
 ) -> TensorDict:


Same as the sync API: renaming the kwarg from fields to select_fields in async_kv_batch_get breaks callers that pass fields=. Consider supporting fields as a deprecated alias (and/or emitting a DeprecationWarning) to avoid a hard break.

ascend-robot · 2026-03-25T08:52:35Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

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

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> # Conflicts: # tests/e2e/test_kv_interface_e2e.py

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

ascend-robot · 2026-03-25T09:08:55Z

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

…cend#57) ## 📝 Description This PR aims to improve the usability and flexibility of the KV (Key-Value) API. It optimizes how `KVBatchMeta` is handled during read and write operations, refactors some internal naming for clarity, and adds unit test coverage for the asynchronous interfaces to ensure stability. ## 🔄 What's Changed 1. **Enhanced `kv_put`**: The `kv_put` operation now directly returns a `KVBatchMeta` object containing all the available fields after the put operation. 2. **Simplified `kv_get`**: Provide an additional API `(async_)kv_batch_get_by_meta` that supports taking a `KVBatchMeta` object directly as input, simplifying the calling logic. 3. **Flexible `KVBatchMeta` Fields**: Introduced public interfaces that allow for more flexible and customizable field selection within `KVBatchMeta`. 4. **Naming Refactor**: Refactored the `fields` variables to `select_fields` during `kv_batch_get` to semantically distinguish from the `fields` (real data TensorDict) in `kv_put`. 5. **Async UT Coverage**: Added UT specifically for the async KV interfaces to ensure robustness. --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> Signed-off-by: Chi Zhang <czhangseu@gmail.com>

…rollout and add CI (#63) ## Summary - **Rewrite `single_controller_demo.py`** to showcase a realistic agentic RLHF training loop: replaces the previous hardcoded tensor inputs and `AsyncvLLMServer` with a multi-turn `AgentLoop` that interleaves LLM generation with simulated tool calls, an OpenAI-style `MessageDataset` with interleaved image support, a proper `DataLoader`-based training loop, and an explicit `compute_reward` step that writes advantages back through TQ. - **Adopt the new KV APIs** introduced in #57 throughout the recipe — uses `kv_batch_get_by_meta`, the return value of `kv_batch_put` (cumulative `KVBatchMeta`), and removes all manual `kv_meta.fields.append(...)` calls. - **Add structured configuration** via `@dataclass` classes (`TrainerConfig`, `AgentLoopConfig`, `MessageDatasetConfig`) and `argparse` CLI, replacing the flat `OmegaConf.create` dict. Trainer config and TQ config are now cleanly separated. - **Add `recipe-check.yml` CI workflow** that runs the demo end-to-end on every push/PR with reduced parameters (`--num-samples 8 --global-batch-size 4 --rollout-agent-num-workers 1`) to keep CI fast. ## Key changes | Area | Before | After | |------|--------|-------| | Rollout | `AsyncvLLMServer` (Ray actor, single-turn) | `AgentLoop` (multi-turn with tool calls, per-sample async) | | Data | Hardcoded `[[1,2],[3,4],...]` tensors | `MessageDataset` + `DataLoader` with random multi-modal messages | | Reward | Simulated inline (`time.sleep`) | `compute_reward()` producing per-token advantages via TQ | | KV API | `kv_batch_get(keys=..., fields=...)` + manual field tracking | `kv_batch_get_by_meta(meta=...)` + `kv_batch_put` return value | | Config | Flat `OmegaConf` dict | Typed `@dataclass` hierarchy + `argparse` CLI | | CI | None | `recipe-check.yml` workflow | ## Test plan - [ ] Run the recipe locally: `python recipe/simple_use_case/single_controller_demo.py --num-samples 8 --global-batch-size 4` - [ ] Verify the new `recipe-check.yml` workflow passes in CI - [ ] Confirm existing tests (`pytest tests`) still pass --------- Signed-off-by: Chi Zhang <czhangseu@gmail.com>

Copilot AI review requested due to automatic review settings March 25, 2026 02:03

ascend-robot added the ascend-cla/yes label Mar 25, 2026

Copilot started reviewing on behalf of 0oshowero0 March 25, 2026 02:04 View session

0oshowero0 changed the title ~~[feat] User-friendly KV API~~ [feat,CI] Improve KV API usability and KVBatchMeta interactions Mar 25, 2026

0oshowero0 changed the title ~~[feat,CI] Improve KV API usability and KVBatchMeta interactions~~ [feat,CI] Improve KV API usability and KVBatchMeta interactions Mar 25, 2026

Copilot AI reviewed Mar 25, 2026

View reviewed changes

0oshowero0 requested a review from Copilot March 25, 2026 02:51

Copilot started reviewing on behalf of 0oshowero0 March 25, 2026 02:51 View session

Copilot AI reviewed Mar 25, 2026

View reviewed changes

0oshowero0 requested a review from Copilot March 25, 2026 03:16

Copilot started reviewing on behalf of 0oshowero0 March 25, 2026 03:17 View session

Copilot AI reviewed Mar 25, 2026

View reviewed changes

0oshowero0 requested a review from Copilot March 25, 2026 08:34

Copilot started reviewing on behalf of 0oshowero0 March 25, 2026 08:35 View session

Copilot AI reviewed Mar 25, 2026

View reviewed changes

0oshowero0 added 7 commits March 25, 2026 16:55

better API for usage

d7759bd

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

add UT

c00e878

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

provide async kv API test cover

3f725c3

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> # Conflicts: # tests/e2e/test_kv_interface_e2e.py

fix comments

aab02f4

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

fix

44f765b

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

fix import

888dfa7

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

fix

177f156

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

0oshowero0 added 4 commits March 25, 2026 16:57

enhance get by meta function

32b94f4

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

fix pre commit

0725f89

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

fix

0e51004

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

update demo

34fb300

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

0oshowero0 force-pushed the han/api branch from 667e1bb to 34fb300 Compare March 25, 2026 09:08

0oshowero0 merged commit 28b846b into Ascend:main Mar 25, 2026
6 checks passed

vermouth1992 mentioned this pull request Mar 28, 2026

[recipe] feat: Revamp single-controller demo with agentic multi-turn rollout and add CI #63

Merged

3 tasks

0oshowero0 deleted the han/api branch April 2, 2026 07:22

Uh oh!

Conversation

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

Uh oh!

📝 Description

🔄 What's Changed

Uh oh!

ascend-robot commented Mar 25, 2026

CLA Signature Pass

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

ascend-robot commented Mar 25, 2026

CLA Signature Pass

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ascend-robot commented Mar 25, 2026

CLA Signature Pass

Uh oh!

ascend-robot commented Mar 25, 2026

CLA Signature Pass

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ascend-robot commented Mar 25, 2026

CLA Signature Pass

Uh oh!

ascend-robot commented Mar 25, 2026

CLA Signature Pass

Uh oh!

ascend-robot commented Mar 25, 2026

CLA Signature Pass

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ascend-robot commented Mar 25, 2026

0oshowero0 commented Mar 25, 2026 •

edited

Loading