| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
[BREAKING][fix] Use jagged tensor as default tensor type (#92) ## Background Previously, TransferQueue would try torch.stack() first when merging per-sample tensors into a batched tensordict for user retrieval. As a result, tensors with uniform shapes were returned as regular dense tensors, while jagged data fell back to nested tensors. This inconsistency forced downstream code to handle two distinct data types (torch.Tensor vs. nested tensor), adding unnecessary branching logic. ## Changes This PR changes the default aggregation strategy so that all tensor fields are returned as nested tensors by default, eliminating the torch.stack() fast-path. Specifically: 1. KVStorageManager._merge_tensors_to_tensordict: Removed the torch.stack(chunk) fallback. The new chain is as_nested_tensor(jagged) → nested_tensor(strided) → NonTensorStack. 2. AsyncSimpleStorageManager._pack_field_values: Removed the torch.stack(values) fast-path for uniform-shape tensors. The new in is as_nested_tensor(jagged) → as_nested_tensor(strided) → NonTensorStack, consistent with the KV backend. 3. Unified strided fallback: Added the missing strided layout fallback to KVStorageManager, ensuring both backends behave identically when jagged layout fails (e.g., for zero-dim tensors). 4. Docstring & comment cleanup: Updated all outdated docstrings and comments that referenced the old torch.stack-first behavior. ## Test updates - Adapted test_async_simple_storage_manager.py, test_kv_storage_manager.py, and e2e tests to accept nested tensors as the default return type. - Reworked the test_kv_storage_manager.py fixture to use realistic variable-length fields (input_ids, prompt_ids, response_ids, response_mask) aligned the single_controller_demo.py schema, replacing the oversimplified text/label/mask example. - Replaced all torch.equal(dense, nested) assertions with safe per-component comparisons (unbind(0) + torch.equal) to accommodate the new nested-tensor contract --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 2 个月前 | |
[feat,CI] Improve KV API usability and KVBatchMeta interactions (#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> | 3 个月前 | |
[fix] Fix BatchMeta.union semantics (#95) ## Problem PR https://gitcode.com/Ascend/TransferQueue/pull/28 incorrectly rewrote BatchMeta.union to behave like concat with global_index deduplication. - **Original semantics**: merge fields for samples with identical global_indexes (row-aligned, column-expanded). - **Broken semantics**: append rows from other whose global_indexes are not present in self. This broke the design boundary between union (same rows, merge columns) and concat (same columns, append rows). ## Changes ### 1. Restore BatchMeta.union (transfer_queue/metadata.py) - Validate that both batches have the same size, global_indexes, and partition_ids. - Merge field_schema with other overriding self on name conflicts. - Merge production_status conservatively via np.bitwise_and (both sides must report ready). - Merge extra_info, custom_meta, and _custom_backend_meta per sample. ### 2. Update tutorial (tutorial/03_metadata_concepts.py) - Example now uses overlapping fields (attention_mask present in both batches) to demonstrate the override behavior. - Corrected comments to clearly distinguish concat (append rows) vs union (merge columns). ### 3. Update unit tests (tests/test_metadata.py) --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 2 个月前 | |
[perf] Improve performance for putting jagged tensor (#36) ## Background When users input a TensorDict containing jagged tensors (nested tensors), the put_data process becomes extremely slow. Specifically, the _filter_storage_data function uses itemgetter(*batch_indexes)(data[fname]) to extract individual items from each tensor in the TensorDict. This indexing approach works efficiently for strided tensors but is extremely inefficient for jagged tensors. ## Root Cause For jagged tensors, itemgetter with multiple batch indexes requires repeated indexing operations, which is $\mathcal{O}(n)$ for each access. When extracting multiple samples, this becomes $\mathcal{O}(n²)$ complexity. ## Solution We unbind nested tensor before accessing each sample from it. python3 # unbind nested tensor results: dict = {} for field in sorted(data.keys()): field_data = data[field] if isinstance(field_data, Tensor) and field_data.is_nested: results[field] = field_data.unbind() else: results[field] = field_data ## Simple Reproduction Script python3 import torch import time from operator import itemgetter # Create a jagged tensor with 1000 samples offsets = torch.tensor([0] + list(torch.randint(10, 50, (1001,)).cumsum(0))) values = torch.randn(offsets[-1].item(), 128) jagged = torch.nested.as_nested_tensor( [values[offsets[i]:offsets[i+1]] for i in range(1000)], layout=torch.jagged ) batch_indexes = list(range(0, 1000, 10)) # 100 indexes # Method 1: Direct itemgetter on jagged tensor (SLOW) start = time.perf_counter() result = itemgetter(*batch_indexes)(jagged) print(f"Direct itemgetter: {(time.perf_counter() - start)*1000:.2f} ms") # Method 2: Unbind first, then itemgetter (FAST) start = time.perf_counter() field_list = jagged.unbind() result = itemgetter(*batch_indexes)(field_list) print(f"Unbind + itemgetter: {(time.perf_counter() - start)*1000:.2f} ms") Output: bash Direct itemgetter: 150.94 ms Unbind + itemgetter: 1.80 ms --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 4 个月前 | |
[feat] Introduce high-level key-value (KV) interface (#28) ## Summary This PR introduces a **High-Level Key-Value (KV) Interface** to TransferQueue, offering a Redis-style API that can enjoy most of the advanced features provided by TransferQueue. ## Background In previous versions of TransferQueue, the learning curve was relatively sharp for new users. To perform basic operations, users had to: 1. Understand BatchMeta SampleMeta and FieldMeta design (as illustrated in [tutorial/02_metadat_concepts.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/02_metadata_concepts.py) 2. Navigate the flexible but complex [TransferQueueClient](https://github.com/Ascend/TransferQueue/blob/main/transfer_queue/client.py) API. Although PR https://github.com/Ascend/TransferQueue/pull/26 simplified the initialization process, the core interaction still required exposing low-level details. This PR bridges that gap by providing a familiar, easy-to-use KV abstraction. ## TransferQueue API Architecture With this PR, TransferQueue now supports a two-level API architecture to satisfy different user needs. | Level | Tier | Style | Fine-Grained Access | Streaming | Sampler | Multiple-Backends | |---|---|---|---|---|---|---| | High | **KV Interface** (this PR) | Put/Get/List/Clear | ✓ | ○ | ✗ | ✓ | | High | **StreamingDataLoader** (#23) | PyTorch DataLoader | ✓ |✓ | ✓ | ✓ | | Low | **TransferQueueClient** | Metadata-based | ✓ | ✓ | ✓ | ✓ | ### High-Level API #### Key-Value based API (This PR) **Methods** - **(async_)kv_put**: Insert/Update a multi-column sample by key, with optional metadata tag - **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batch - **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields) - **(async_)kv_list**: List keys and tags (metadata) in a partition - **(async_)kv_clear**: Remove key-value pairs from storage **Key Features** - **Redis-style Semantics**: Familiar KV interface (Put/Get/List) for zero learning curve - **Fine-grained Access**: Update or retrieve specific fields (columns) within a key (row) without full op. - **Partition Isolation**: Logical separation of storage namespaces - **Metadata Tags**: Lightweight metadata for status tracking - **Pluggable Backends**: Supports multiple backends #### StreamingDataLoader API Refer to our [RoadMap](https://github.com/Ascend/TransferQueue/issues/1) and related PRs(https://github.com/Ascend/TransferQueue/pull/23). The usage example can be found in [tutorial/06_streaming_dataloader.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/06_streaming_dataloader.py). ### Low-Level API Directly manipulate the TransferQueueClient. Refer to [tutorial/03_metadata_concepts.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/03_metadata_concepts.py), [tutorial/04_understanding_controller.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/04_understanding_controller.py) and [tutorial/05_custom_sampler.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/05_custom_sampler.py) for details. ## Usage Example Please refer to [tutorial/02_kv_interface.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/02_kv_interface.py) and [tests/e2e/test_kv_interface_e2e.py](https://github.com/Ascend/TransferQueue/blob/main/tests/e2e/test_kv_interface_e2e.py) for details. python3 import torch from tensordict import TensorDict import transfer_queue as tq # initialize TQ tq.init() # prepare data batch_input_ids = torch.tensor( [ [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15], ] ) batch_attention_mask = torch.ones_like(batch_input_ids) data_batch = TensorDict( { "input_ids": batch_input_ids, "attention_mask": batch_attention_mask, }, batch_size=batch_input_ids.size(0), ) keys = ["1_0", "1_1", "1_2", "2_0"] # 4 keys for 4 samples tags = [{"global_steps": 1, "status": "running", "model_version": 1} for _ in range(len(keys))] partition_id = "test" # use kv interface to put into TQ tq.kv_batch_put(keys=keys, partition_id=partition_id, fields=data_batch, tags=tags) # list all keys and tags all_keys, all_tags = tq.kv_list(partition_id=partition_id) for k, t in zip(all_keys, all_tags, strict=False): print(f" - key='{k}' | tag={t}") # retrieve all data retrieved_all = tq.kv_batch_get(keys=all_keys, partition_id=partition_id) print(f" Fields: {list(retrieved_all.keys())}") ## Use Cases & Limitations **Best For**: - Scenarios requiring fine-grained data access (e.g., updating a reward score for a specific prompt). - Integration with external ReplayBuffers or Single-Controller architectures that manage sample dispatching logic. **Limitations (vs. Streaming/Low-level APIs):** - No built-in production/consumption tracking: Users must manually check status via tags or manage logic externally. - No Built-in Sampler: Must implement data dispatch by ReplayBuffer or single-controller externally. - Not Fully Streaming: Consumers typically wait for a controller to dispatch keys before fetching, rather than a continuous stream. --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 5 个月前 | |
[feat] Support cross-job actor discovery via explicit namespace (#115) When multiple Ray Jobs share the same Ray cluster, Named Actors are isolated by namespace. Without an explicit namespace, a TQ Controller created by one job is invisible to workers in another job. This commit adds namespace="transfer_queue" to both: - ray.get_actor() in _init_from_existing() - TransferQueueController.options() in init() This ensures that the TQ Controller is always registered and discovered in the fixed "transfer_queue" namespace, enabling cross-job TQ sharing (e.g., a teacher server job creates TQ, and a trainer job connects to it). This change is backward-compatible: single-job usage is unaffected since the namespace is consistent between creation and discovery. Signed-off-by: huniu20 <huniumail@gmail.com> | 1 个月前 | |
[BREAKING][fix] Use jagged tensor as default tensor type (#92) ## Background Previously, TransferQueue would try torch.stack() first when merging per-sample tensors into a batched tensordict for user retrieval. As a result, tensors with uniform shapes were returned as regular dense tensors, while jagged data fell back to nested tensors. This inconsistency forced downstream code to handle two distinct data types (torch.Tensor vs. nested tensor), adding unnecessary branching logic. ## Changes This PR changes the default aggregation strategy so that all tensor fields are returned as nested tensors by default, eliminating the torch.stack() fast-path. Specifically: 1. KVStorageManager._merge_tensors_to_tensordict: Removed the torch.stack(chunk) fallback. The new chain is as_nested_tensor(jagged) → nested_tensor(strided) → NonTensorStack. 2. AsyncSimpleStorageManager._pack_field_values: Removed the torch.stack(values) fast-path for uniform-shape tensors. The new in is as_nested_tensor(jagged) → as_nested_tensor(strided) → NonTensorStack, consistent with the KV backend. 3. Unified strided fallback: Added the missing strided layout fallback to KVStorageManager, ensuring both backends behave identically when jagged layout fails (e.g., for zero-dim tensors). 4. Docstring & comment cleanup: Updated all outdated docstrings and comments that referenced the old torch.stack-first behavior. ## Test updates - Adapted test_async_simple_storage_manager.py, test_kv_storage_manager.py, and e2e tests to accept nested tensors as the default return type. - Reworked the test_kv_storage_manager.py fixture to use realistic variable-length fields (input_ids, prompt_ids, response_ids, response_mask) aligned the single_controller_demo.py schema, replacing the oversimplified text/label/mask example. - Replaced all torch.equal(dense, nested) assertions with safe per-component comparisons (unbind(0) + torch.equal) to accommodate the new nested-tensor contract --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 2 个月前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 2 个月前 | ||
| 3 个月前 | ||
| 2 个月前 | ||
| 4 个月前 | ||
| 5 个月前 | ||
| 1 个月前 | ||
| 2 个月前 |