diff --git a/README.md b/README.md
index 2b7cd65f..d7785108 100644
--- a/README.md
+++ b/README.md
@@ -54,7 +54,7 @@ TransferQueue offers **fine-grained, sub-sample-level** data management and **lo
### Control Plane: Panoramic Data Management
-In the control plane, `TransferQueueController` tracks the **production status** and **consumption status** of each training sample as metadata. Once all required data fields are ready (i.e., written to the `StorageManager`), the data sample can be consumed by downstream tasks.
+In the control plane, `TransferQueueController` tracks the **production status** and **consumption status** of each training sample as metadata. A sample becomes available to downstream tasks when every required field is ready.
We also track the consumption history for each computational task (e.g., `generate_sequences`, `compute_log_prob`, etc.). Therefore, even when different computational tasks require the same data field, they can consume the data independently without interfering with each other.
@@ -80,12 +80,12 @@ This class encapsulates the core interaction logic within the TransferQueue syst
Currently, we support the following storage backends:
-- SimpleStorage: A basic CPU memory storage with minimal data format constraints and ease of use.
+- SimpleStorage: Distributed CPU storage with optional [local SSD offload](docs/ssd_offload.md), storage-side parsing and updates, and shared fields. It is the default backend and needs no external storage service.
- [Yuanrong](https://gitee.com/openeuler/yuanrong-datasystem) ([usage guide](docs/storage_backends/openyuanrong_datasystem.md), beta, [#PR107](https://github.com/TransferQueue/TransferQueue/pull/107), [#PR96](https://github.com/TransferQueue/TransferQueue/pull/96)): An Ascend native data system that provides hierarchical storage interfaces including HBM/DRAM/SSD.
- [MooncakeStore](https://github.com/kvcache-ai/Mooncake) (beta, [#PR162](https://github.com/TransferQueue/TransferQueue/pull/162)): A high-performance, KV-based hierarchical storage that supports RDMA transport between GPU and DRAM.
- [RayRDT](https://docs.ray.io/en/master/ray-core/direct-transport.html) (alpha, [#PR167](https://github.com/TransferQueue/TransferQueue/pull/167)): Ray's new feature that allows Ray to store and pass objects directly between Ray actors.
-Among them, `SimpleStorageUnit` serves as our default storage backend, coordinated by the `AsyncSimpleStorageManager` class. Each storage unit can be deployed on a separate node, allowing for distributed data management.
+Among them, `SimpleStorageUnit` serves as our default storage backend, coordinated by the `AsyncSimpleStorageManager` class. Storage units are distributed across eligible Ray nodes. Within SimpleStorage, each partition's rows are balanced across units by row count; this does not account for differences in payload size.
`SimpleStorageUnit` employs a 2D data structure as follows:
@@ -100,54 +100,29 @@ This data structure design is motivated by the computational characteristics of
### User Interface: High-Level & Low-Level APIs
-| Level | Tier | Style | Fine-Grained Access | Streaming | Sampler | Multiple-Backends |
-|---|---|---|---|------------------|---|---|
-| High | **KV Interface** ([PR#28](https://github.com/Ascend/TransferQueue/pull/28))| Put/Get/List/Clear | ✓ | ○ | ✗ | ✓ |
-| High | **StreamingDataLoader** ([PR#23](https://github.com/Ascend/TransferQueue/pull/23)) | PyTorch DataLoader | ✓ | ✓ | ✓ | ✓ |
-| Low | **TransferQueueClient** | Metadata-based | ✓ | ✓ | ✓ | ✓ |
-
-
-#### Key-Value based API
-
-To simplify the usage of TransferQueue, we provide a Redis-style high-level API that exposes most of its advanced features ([PR#28](https://github.com/Ascend/TransferQueue/pull/28)).
-
-**Methods**
-
-- **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag.
-- **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches.
-- **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields) and, on SimpleStorage, an optional `data_parser` run on the storage units.
-- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`, adding missing keys and fields. SimpleStorage only.
-- **(async_)kv_list**: List keys and tags (metadata) in a partition.
-- **(async_)kv_describe**: Show each key's stored fields, including those `kv_update(finalize=False)` keeps hidden.
-- **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
-
-`kv_update` merges only the rows the controller records as holding the field
-(including hidden ones); every other row takes the new value as given, without
-calling `merge_fn`, and overwrites anything a storage unit kept after a failed
-clear. Every stored value must keep the field's tensor/non-tensor type and
-tensor dtype in the partition. A timeout has an unknown outcome because a
-storage unit may commit after the caller stops waiting, so non-idempotent merge
-operations must not be retried blindly. Updates are atomic within one
-storage-unit request; a batch spanning multiple units is not a distributed
-transaction.
-
-To build a field over several updates without exposing intermediate values, pass
-`finalize=False` to every `kv_update` except the last one, which uses the default
-`finalize=True`. Until then the field is stored but not ready, so samplers, the
-dataloader and `kv_batch_get` do not see it. Updates to one key must come from
-one writer at a time; use `kv_global_lock` if several processes write the same key.
-
-**Advanced: parsing on get.** Like `kv_batch_put`, `kv_batch_get` takes a
-`data_parser`, here run by each SimpleStorage unit before it replies, so only the
-parsed values cross the network (e.g. trimming padding or picking a few frames).
-It is called as `data_parser(data, args)`: `data` maps each selected field to the
-unit's stored values, one per sample, and `args` holds the matching entries of the
-optional per-key `parser_args`. Each field may be parsed differently, but the
-result must keep the keys and the number and order of values. `data` holds the
-stored objects themselves, so never modify them in place; build new values
-instead. As with put, a get with a parser is not retried on a storage-unit
-timeout; the call raises, and retrying it yourself may run the parser again, so
-keep it free of side effects. See the `kv_batch_get` docstring for details.
+| Interface | How consumers select data | Typical use |
+|---|---|---|
+| **KV Interface** | Explicit string keys and selected fields | A trainer or replay buffer dispatches samples |
+| **StreamingDataLoader** | A sampler selects samples whose required fields are ready | Distributed ranks consume a streaming pipeline |
+| **TransferQueueClient** | `BatchMeta` identifies rows, fields and readiness | Custom scheduling and metadata-based integrations |
+
+#### Key-Value API
+
+Use string keys to read or write selected fields when your application controls
+sample dispatch. Each key identifies a sample within a partition. Start with
+[the interactive notebook](tutorial/basic.ipynb) or [Tutorial 02](tutorial/02_kv_interface.py).
+
+| API | Behavior |
+|---|---|
+| `(async_)kv_put`, `(async_)kv_batch_put` | Store samples or add/replace fields, with optional metadata tags |
+| `(async_)kv_update`, `(async_)kv_batch_update` | Merge new values into stored fields, such as appending response chunks |
+| `(async_)kv_batch_get` | Read selected fields for the requested keys |
+| `(async_)kv_list` | List keys and tags in one or all partitions |
+| `(async_)kv_describe` | Show a key's fields, readiness and shared-field sources |
+| `(async_)kv_exists` | Check key existence and list ready fields |
+| `(async_)kv_link` | Link fields to shared data in `tq.PUBLIC_PARTITION` |
+| `(async_)kv_clear` | Delete keys or selected fields |
+| `(async_)kv_clear_unlinked_public` | Reclaim shared data no longer referenced by trajectories |
**Key Features**
@@ -156,35 +131,10 @@ keep it free of side effects. See the `kv_batch_get` docstring for details.
- **Partition Isolation**: Logical separation of storage namespaces.
- **Metadata Tags**: Lightweight metadata for status tracking.
- **Pluggable Backends**: Supports multiple backends.
-
-Refer to [tutorials/basic.ipynb](https://github.com/Ascend/TransferQueue/blob/main/tutorial/basic.ipynb) and [tutorials/02_kv_interface.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/02_kv_interface.py) for detailed usage examples.
-
-**KV Key Locks**
-
-Locks are **advisory**: KV calls never check them, so every writer that needs mutual exclusion must take the lock. `(async_)kv_local_lock` excludes threads and asyncio tasks of one process and needs no setup.
-
-- **(async_)kv_global_lock(keys, partition_id, timeout=None, lease_s=30)**: Hold an exclusive cluster-wide lock on `(partition_id, key)` pairs under a lease that a background thread renews. Yields a **GlobalLease**: call `check()` right before writes (it raises **LockLostError** once the lease is lost); `names` lists the locked pairs. Exiting the block also raises `LockLostError` if the lease was lost.
-- **kv_lock_list(partition_id=None)**: Return `{"holders": [...], "waiters": n}` for live global locks.
-
-```python
-import os
-os.environ["TQ_ENABLE_KV_GLOBAL_LOCK"] = "1" # only in the process that creates TransferQueue, before tq.init()
-tq.init()
-with tq.kv_global_lock(["a", "b"], "train", timeout=30) as lease:
- lease.check()
- tq.kv_put("a", "train", fields={"v": torch.tensor([1])})
-```
-
-- **Enabling**: set these before `tq.init()` in the process that creates TransferQueue; only that process reads them. Processes and Ray workers that attach find the lock actors by name and need nothing. Without the switch, `kv_global_lock` raises `RuntimeError`.
-
-| Variable | Default | Meaning |
-| --- | --- | --- |
-| `TQ_ENABLE_KV_GLOBAL_LOCK` | off | `1`/`true`/`yes`/`y`/`on` starts the lock actors |
-| `TQ_KV_GLOBAL_LOCK_NUM_SHARDS` | `8` | Number of lock actors (integer >= 1) |
-| `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` | `10000` | Acquire calls each lock actor lets wait at once (integer >= 1) |
-
-- **Shards**: keys are spread by hash over `TQ_KV_GLOBAL_LOCK_NUM_SHARDS` `TransferQueueLockManager` actors, which `tq.init()` pins round-robin to the alive Ray nodes (per-node counts differ by at most 1). Keys on several actors are taken one actor at a time, so a multi-key request may hold some keys while it waits for others.
-- **Capacity**: `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` caps the acquire calls each actor lets wait at once. A request past the cap keeps its queue place but polls once per window (5 s, or `lease_s / 3` if shorter), so it may see its grant up to a window late. About `TQ_KV_GLOBAL_LOCK_NUM_SHARDS` × this value (80,000 by default) requests can wait without that delay. More shards spread the load, but all waiters on one key share one actor, so a hot key is bounded by this value alone. Each actor's Ray `max_concurrency` is this value plus 16 reserved slots shared, without priority, by renew, release, `kv_lock_list` and turned-away acquires. Every waiter still costs one call per window: past an actor's call throughput, renewals are delayed and leases can be lost.
+- **Incremental Updates**: Build responses in chunks and publish complete fields to consumers. See [Tutorial 02](tutorial/02_kv_interface.py).
+- **Shared Fields**: Store a field once for reuse across samples. See [Tutorial 07](tutorial/07_shared_fields.py).
+- **Storage-side Parsing**: Decode references on put or return selected tokens and frames on get. See [Tutorial 08](tutorial/08_data_parser.py) and [the notebook](tutorial/basic.ipynb).
+- **Selective Recovery**: Save unfinished rollout samples and restore them into a running TQ. See [Tutorial 09](tutorial/09_data_dump.py).
#### StreamingDataLoader API
@@ -193,17 +143,34 @@ Designed as a drop-in replacement for the standard PyTorch `DataLoader`, this AP
In this scenario, `TransferQueueController` serves as a side-controller for data dispatching, with a user-defined `Sampler` class to organize the dataflow.
It encapsulates the complex scheduling and data transfer logic required for various parallelism strategies, seamlessly integrating TransferQueue into existing training workflows and simplifying the development of disaggregated frameworks.
-See the [Roadmap](https://github.com/Ascend/TransferQueue/issues/1) and [tutorials/06_streaming_dataloader.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/06_streaming_dataloader.py) for more details.
+See [tutorial/06_streaming_dataloader.py](tutorial/06_streaming_dataloader.py) for more details.
#### Low-Level Native API
The native interfaces of TransferQueue are implemented in `TransferQueueClient`. It offers maximum flexibility through native, atomic operations.
Developers can leverage `TransferQueueClient` directly to implement advanced features that require fine-grained control and fully streamed data scheduling, as illustrated in the following tutorials:
-- [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)
-- [tutorial/05_custom_sampler.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/05_custom_sampler.py)
+- [tutorial/03_metadata_concepts.py](tutorial/03_metadata_concepts.py)
+- [tutorial/04_understanding_controller.py](tutorial/04_understanding_controller.py)
+- [tutorial/05_custom_sampler.py](tutorial/05_custom_sampler.py)
+
+### Tutorials
+
+Start with [the basic notebook](tutorial/basic.ipynb) for an interactive walkthrough.
+The scripts below run from the repository root with CPU tensors and SimpleStorage.
+
+| Tutorial | What you will learn |
+|---|---|
+| [01: Core components](tutorial/01_core_components.py) | Controller, client and storage responsibilities |
+| [02: KV interface](tutorial/02_kv_interface.py) | Put/get, staged updates, readiness probes and field-level clear |
+| [03: Metadata](tutorial/03_metadata_concepts.py) | BatchMeta operations and shared fields without string keys |
+| [04: Controller](tutorial/04_understanding_controller.py) | Partitions, field readiness and consumption tracking |
+| [05: Custom samplers](tutorial/05_custom_sampler.py) | Select and consume ready samples |
+| [06: StreamingDataLoader](tutorial/06_streaming_dataloader.py) | Distributed consumption, DP rank coordination and micro-batches |
+| [07: Shared fields](tutorial/07_shared_fields.py) | Share fields across samples, read shared and private fields, and reclaim shared data |
+| [08: Data parsers](tutorial/08_data_parser.py) | Load files on put and select tensor rows on get using storage-side parsers |
+| [09: Selective recovery](tutorial/09_data_dump.py) | Save and resume unfinished samples with shared fields |
🔥 Showcases
@@ -225,7 +192,7 @@ We have experimentally implemented a **standardized, fully-streamed distributed*
By leveraging the `RankAwareSampler` and `StreamingDataLoader` interfaces, we achieve a **streamlined micro-batch-level producer-consumer pipeline**. This design eliminates the need to manually determine data dispatching logic across varying parallelism strategies—a typical complexity in the single-controller paradigm—thereby greatly simplifying framework design.
-Please refer to our [Roadmap](https://github.com/Ascend/TransferQueue/issues/1) and [tutorials/05_streaming_dataloader.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/05_streaming_dataloader.py) for more details.
+Please refer to [tutorial/06_streaming_dataloader.py](tutorial/06_streaming_dataloader.py) for more details.
@@ -404,4 +371,4 @@ Please kindly cite our paper if you find this repo is useful:
journal={arXiv preprint arXiv:2507.01663},
year={2025}
}
-```
\ No newline at end of file
+```
diff --git a/transfer_queue/version/version b/transfer_queue/version/version
index 260e27af..0ea3a944 100644
--- a/transfer_queue/version/version
+++ b/transfer_queue/version/version
@@ -1 +1 @@
-0.1.12.dev0
+0.2.0
diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py
index 01752b5e..b1d5b915 100644
--- a/tutorial/02_kv_interface.py
+++ b/tutorial/02_kv_interface.py
@@ -146,6 +146,7 @@ def demonstrate_kv_api():
# Step 4: Only update tags through kv_put
print("\n[Step 4] Update existing tags without providing value...")
+ print(" Tags hold application metadata; TQ tracks field readiness separately.")
key_for_update_tags = "0_0"
tag_update = {"global_steps": 0, "status": "finish", "model_version": 0}
print(f" Target Key: {key_for_update_tags}")
@@ -153,18 +154,31 @@ def demonstrate_kv_api():
tq.kv_put(key=key_for_update_tags, partition_id=partition_id, fields=None, tag=tag_update)
print(f" ✓ Update success: Samples '0_0' now has tag as {tag_update}.")
- # Step 5: Concatenate a field with kv_update, then clear only that field
- print("\n[Step 5] Updating a field with kv_update (concat, then clear the field)...")
- tq.kv_put(key=key, partition_id=partition_id, fields={"scratch": torch.tensor([1, 2])})
+ response_key = keys[0]
+ print(f"\n[Step 5] Building a response for '{response_key}' in chunks...")
+ print(" kv_update requires SimpleStorage; use one writer per key or coordinate writers with a KV lock.")
+ print(" A timed-out merge may have committed; retrying it blindly can append the same chunk twice.")
tq.kv_update(
- key=key,
+ key=response_key,
partition_id=partition_id,
- fields={"scratch": torch.tensor([3, 4])},
+ fields={"response": torch.tensor([1, 2])},
merge_fn=lambda old, new: torch.cat([old, new]),
+ finalize=False,
)
- print(" ✓ kv_update concat: scratch of '0_0' is now [1, 2, 3, 4].")
- tq.kv_clear(keys=key, partition_id=partition_id, fields="scratch")
- print(" ✓ tq.kv_clear(fields=...): scratch of '0_0' is removed (key and other fields remain).")
+ print(" Stored fields:", tq.kv_describe(response_key, partition_id)[response_key])
+ print(" Ready fields:", tq.kv_exists(response_key, partition_id)[response_key])
+
+ tq.kv_update(
+ key=response_key,
+ partition_id=partition_id,
+ fields={"response": torch.tensor([3, 4])},
+ merge_fn=lambda old, new: torch.cat([old, new]),
+ finalize=True,
+ )
+ response = tq.kv_batch_get(response_key, partition_id, select_fields="response")
+ assert response["response"][0].tolist() == [1, 2, 3, 4]
+ print(" Published response:", response["response"][0])
+ tq.kv_put(response_key, partition_id, tag={"status": "finish"})
# Step 6: List all keys and tags in a partition
print("\n[Step 6] Listing all keys and tags in partition...")
@@ -183,17 +197,17 @@ def demonstrate_kv_api():
retrieved_input_ids = tq.kv_batch_get(keys=all_keys, partition_id=partition_id, select_fields="input_ids")
print(f" ✓ Successfully retrieved only {list(retrieved_input_ids.keys())} field for all samples.")
- # Step 8: Retrieve all fields using kv_batch_get
- print("\n[Step 8] Retrieving all fields with kv_batch_get...")
+ print("\n[Step 8] Retrieving the fields ready in every requested key...")
retrieved_all = tq.kv_batch_get(keys=all_keys, partition_id=partition_id)
- print(f" Retrieved all fields for {all_keys}:")
- print(f" Fields: {list(retrieved_all.keys())}")
- print(
- f" Note: We cannot retrieve fields {list(response_batch.keys())}, since they only available in {append_keys}"
- )
+ print(f" Requested keys: {all_keys}")
+ print(f" Common ready fields: {list(retrieved_all.keys())}")
+
+ print("\n[Step 9] Clearing selected fields while keeping keys...")
+ tq.kv_clear(all_keys, partition_id, fields="attention_mask")
+ print(" Keys and their other fields remain:", tq.kv_exists(all_keys, partition_id))
- # Step 9: Clear specific keys
- print("\n[Step 9] Clearing keys from partition...")
+ # Step 10: Clear specific keys
+ print("\n[Step 10] Clearing keys from partition...")
keys_to_clear = all_keys[:2] # Delete the first 2 keys
tq.kv_clear(keys=keys_to_clear, partition_id=partition_id)
print(f" ✓ Cleared keys: {keys_to_clear}")
@@ -231,10 +245,6 @@ def main():
- Focusing on fine-grained data access where extreme streaming performance is non-essential
- Integration with external ReplayBuffer/single-controller that manage sample dispatching
- Limitations (vs low-level native APIs):
- - No built-in production/consumption tracking: Users must manually check status via tags externally.
- - No built-in Sampler support: Must implement data dispatch by ReplayBuffer or single-controller externally.
- - Not fully streaming: Consumers must wait for single-controller to dispatch `keys`.
"""
)
)
@@ -252,9 +262,11 @@ def main():
print("=" * 80)
print("\nKey Takeaways:")
print(" 1. KV API simplifies data access with Redis-style semantics")
- print(" 2. Use 'fields' parameter to get/put specific fields only")
- print(" 3. Tags enable custom metadata for production status, scores, etc.")
- print(" 4. Use kv_list to inspect partition contents")
+ print(" 2. Read or write only the fields you need")
+ print(" 3. Tags store custom metadata such as status and scores")
+ print(" 4. Use `kv_list` to inspect keys and tags")
+ print(" 5. Supports incremental updates and publishing completed responses")
+ print(" 6. Supports field-level clearing while keeping keys")
# Cleanup
tq.close()
diff --git a/tutorial/03_metadata_concepts.py b/tutorial/03_metadata_concepts.py
index a7482196..5f5a61b0 100644
--- a/tutorial/03_metadata_concepts.py
+++ b/tutorial/03_metadata_concepts.py
@@ -317,6 +317,30 @@ def demonstrate_real_workflow():
chunk_data = tq_client.get_data(chunk)
print(f" Chunk {i}: Retrieved chunk data: {chunk_data}")
+ print("[Step 7] Selecting rows and fields from BatchMeta with shared fields...")
+ video = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+ public_meta = tq_client.put(
+ TensorDict({"video": video.unsqueeze(0)}, batch_size=1),
+ partition_id=tq.PUBLIC_PARTITION,
+ )
+ linked_meta = tq_client.link(batch_meta.select_samples([0, 1]), public_meta.select_samples([0, 0]), fields="video")
+ print(" Linked the first two samples to the same public video field.")
+ # Payloads stay on the public row; _links holds source indexes aligned with global_indexes.
+ # Selection filters these local references without changing the controller's links.
+ print(f" Rows: {linked_meta.global_indexes}, fields: {linked_meta.field_names}")
+ print(f" Internal links: {linked_meta._links}")
+
+ selected_meta = linked_meta.select_samples([1]).select_fields(["video"])
+ print(" Selected the second sample and kept only its video field.")
+ print(f" Selected rows: {selected_meta.global_indexes}, fields: {selected_meta.field_names}")
+ print(f" Links after selection: {selected_meta._links}")
+ torch.testing.assert_close(tq_client.get_data(selected_meta)["video"][0], video)
+ print(" Reading through the selected metadata returned the original public video.")
+
+ private_meta = linked_meta.select_fields(["input_ids"])
+ print(" Kept only the private input_ids field for both samples; video links are excluded from this metadata.")
+ print(f" Private fields: {private_meta.field_names}, links: {private_meta._links}")
+
# Cleanup
tq_client.clear_partition(partition_id=partition_id)
tq.close()
@@ -341,6 +365,7 @@ def main():
- custom_meta: list[dict] aligned with global_indexes (one dict per sample)
- Metadata operations: chunk, concat, union, select_fields, select_samples, reorder
- batch.samples[i] returns a lazy view with .fields -> field_schema (read-only)
+ - client.link shares public fields across keyless rows on SimpleStorage
"""
)
)
diff --git a/tutorial/06_streaming_dataloader.py b/tutorial/06_streaming_dataloader.py
index 0571e108..284ed408 100644
--- a/tutorial/06_streaming_dataloader.py
+++ b/tutorial/06_streaming_dataloader.py
@@ -14,7 +14,7 @@
# limitations under the License.
"""
-Tutorial 5: Streaming DataLoader for Distributed Training
+Tutorial 6: Streaming DataLoader for Distributed Training
This script demonstrates how to use StreamingDataset and StreamingDataLoader
for efficient streaming data loading in distributed training scenarios.
diff --git a/tutorial/07_shared_fields.py b/tutorial/07_shared_fields.py
new file mode 100644
index 00000000..0a5c9e7e
--- /dev/null
+++ b/tutorial/07_shared_fields.py
@@ -0,0 +1,108 @@
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tutorial 7: Shared Fields
+
+Share stored fields across samples through PUBLIC_PARTITION and kv_link. This
+example uses a tensor field named features and a private score for each sample.
+
+Run: python tutorial/07_shared_fields.py
+
+Requires SimpleStorage. Public keys are shared by applications on the same TQ;
+use an application-specific prefix. Linked fields and referenced public rows
+cannot be overwritten.
+"""
+
+import sys
+from pathlib import Path
+
+import ray
+import torch
+from tensordict import TensorDict
+
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+import transfer_queue as tq # noqa: E402
+
+
+def main():
+ """Demonstrate linking, reading and reclaiming shared fields."""
+ print("=" * 80)
+ print("TransferQueue Tutorial 7: Shared Fields")
+ print("=" * 80)
+ print("A field can be stored once and referenced by multiple samples.")
+ print("Each sample can also have its own private fields.")
+ print("Setting up TransferQueue with the default SimpleStorage backend...")
+ ray.init(namespace="TransferQueueTutorial", include_dashboard=False)
+ try:
+ tq.init()
+ partition_id = "train"
+
+ print("\n[Step 1] Storing shared fields in the public partition...")
+ public_values = {
+ "tutorial07/shared_0": torch.arange(24, dtype=torch.float32).reshape(6, 4),
+ "tutorial07/shared_1": torch.arange(24, 48, dtype=torch.float32).reshape(6, 4),
+ }
+ public_keys = list(public_values)
+ for public_key, value in public_values.items():
+ tq.kv_put(public_key, tq.PUBLIC_PARTITION, fields={"features": value})
+ print(f" Stored {public_key}: features shape={tuple(value.shape)}")
+
+ print("\n[Step 2] Linking sample fields to the shared values...")
+ keys = [f"sample_{i}" for i in range(4)]
+ sources = [public_keys[0], public_keys[0], public_keys[1], public_keys[1]]
+ tq.kv_link(keys, partition_id, public_key=sources, fields="features")
+ descriptions = tq.kv_describe(keys, partition_id)
+ for key in keys:
+ print(f" {key} reads features from {descriptions[key]['links']['features']}")
+ print(" The links reference the public rows; no features payload is copied into the sample rows.")
+
+ print("\n[Step 3] Adding a private field to each sample...")
+ scores = torch.tensor([0.1, 0.2, 0.3, 0.4])
+ tq.kv_batch_put(keys, partition_id, fields=TensorDict({"score": scores}, batch_size=len(keys)))
+ print(" Added score values:", scores.tolist())
+ print(" Each sample now has a shared features field and its own score field.")
+
+ print("\n[Step 4] Reading shared and private fields together...")
+ batch = tq.kv_batch_get(keys, partition_id)
+ print(" Returned fields:", list(batch.keys()))
+ for i, (key, source) in enumerate(zip(keys, sources, strict=True)):
+ torch.testing.assert_close(batch["features"][i], public_values[source])
+ print(f" {key}: features shape={tuple(batch['features'][i].shape)}, score={batch['score'][i].item():.1f}")
+
+ print("\n[Step 5] Clearing sample keys and their links...")
+ tq.kv_clear(keys, partition_id)
+ print(" Cleared sample keys:", keys)
+ print(" Public keys still available:", list(tq.kv_exists(public_keys, tq.PUBLIC_PARTITION)))
+
+ print("\n[Step 6] Reclaiming shared data with no remaining links...")
+ freed = tq.kv_clear_unlinked_public()
+ assert set(freed) == set(public_keys)
+ print(" Reclaimed public keys:", freed)
+
+ print("\n" + "=" * 80)
+ print("Tutorial Complete!")
+ print("=" * 80)
+ print("\nKey Takeaways:")
+ print(" 1. Multiple samples can share a stored field through links")
+ print(" 2. Reads combine shared fields with each sample's private fields")
+ print(" 3. Shared data can be reclaimed after its last link is removed")
+ finally:
+ tq.close()
+ ray.shutdown()
+ print("\nCleanup complete")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tutorial/08_data_parser.py b/tutorial/08_data_parser.py
new file mode 100644
index 00000000..6a01916b
--- /dev/null
+++ b/tutorial/08_data_parser.py
@@ -0,0 +1,160 @@
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tutorial 8: Data Parsers
+
+Load tensors from filesystem paths on put, then select tensor rows on get.
+The caller sends file paths; the storage units read the files and store the data.
+See https://github.com/Ascend/TransferQueue/pull/82 for the original put-parser design.
+
+Run: python tutorial/08_data_parser.py
+
+This single-node example creates two small .pt files in a temporary directory.
+For your own dataset, use existing file paths. In a multi-node deployment, every
+storage node must be able to read those paths, for example on a shared filesystem.
+Keep the files available and unchanged until the put call returns.
+
+Requires SimpleStorage. Each parser receives the samples assigned to its storage
+unit and must preserve field names, sample count and sample order. A get parser
+must leave stored values unchanged. It reduces reply traffic; SSD-backed values
+are still read and decoded in full before parsing.
+
+Parsers are serialized to storage units; their dependencies must be available
+there. Keep them free of side effects: parser-backed requests are not retried
+automatically, and retrying a failed call can execute the parser again.
+"""
+
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import ray
+import torch
+from tensordict import NonTensorStack, TensorDict
+
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+import transfer_queue as tq # noqa: E402
+
+
+def load_features(data):
+ """Read each tensor file on the storage unit and replace its path with the data."""
+ features = []
+ for file_path in data["features"]:
+ features.append(torch.load(file_path, map_location="cpu", weights_only=True))
+ data["features"] = features
+ return data
+
+
+def select_rows(data, row_indexes):
+ """Return the requested rows of each tensor without modifying stored values."""
+ selected = []
+ for value, indexes in zip(data["features"], row_indexes, strict=True):
+ selected.append(value[indexes])
+ return {"features": selected}
+
+
+def main():
+ """Parse values on put and get, showing the input and output of each operation."""
+ print("=" * 80)
+ print("TransferQueue Tutorial 8: Data Parsers")
+ print("=" * 80)
+ print("A put parser can load data from file paths directly on the storage units.")
+ print("A get parser transforms values before they are returned to the caller.")
+ print("Both run on the storage units and keep field names and sample count/order unchanged.")
+ print("Setting up TransferQueue with the default SimpleStorage backend...")
+ ray.init(namespace="TransferQueueTutorial", include_dashboard=False)
+ try:
+ tq.init()
+ partition_id = "parser"
+
+ keys = ["sample_0", "sample_1"]
+ source_features = [
+ torch.arange(12, dtype=torch.float32).reshape(3, 4),
+ torch.arange(12, 24, dtype=torch.float32).reshape(3, 4),
+ ]
+ with TemporaryDirectory(prefix="tq-parser-") as data_dir:
+ print("\n[Step 1] Preparing tensor files on the filesystem...")
+ print(" This example creates two .pt files; a dataset can supply existing paths.")
+ file_paths = []
+ for key, value in zip(keys, source_features, strict=True):
+ file_path = Path(data_dir) / f"{key}.pt"
+ torch.save(value, file_path)
+ file_paths.append(str(file_path))
+ print(f" {key}: saved shape={tuple(value.shape)} to {file_path}")
+ print(" Each file contains a tensor with 3 rows and 4 columns.")
+
+ print("\n[Step 2] Passing file paths to the put parser...")
+ fields = TensorDict({"features": NonTensorStack(*file_paths)}, batch_size=len(keys))
+ print(" NonTensorStack holds one file path per sample.")
+ print(" The put request carries these paths; load_features reads the files on the storage units.")
+ tq.kv_batch_put(keys, partition_id, fields=fields, data_parser=load_features)
+ print(" The put call returned after the files were loaded into storage.")
+
+ print("\n[Step 3] Reading the stored tensors without a get parser...")
+ stored = tq.kv_batch_get(keys, partition_id, select_fields="features")
+ for i, key in enumerate(keys):
+ value = stored["features"][i]
+ torch.testing.assert_close(value, source_features[i])
+ print(f" {key}: shape={tuple(value.shape)}, dtype={value.dtype}")
+ print(f" Stored values: {value.tolist()}")
+ print(" The stored tensors match the file contents.")
+
+ print("\n[Step 4] Reading selected rows with a get parser...")
+ print(" parser_args supplies one argument per key, in the same order as keys.")
+ print(" select_rows(data, row_indexes) interprets each argument as row indexes starting at 0.")
+ row_indexes = [[0, 2], [1]]
+ for key, indexes in zip(keys, row_indexes, strict=True):
+ print(f" {key} requests rows {indexes}")
+ selected = tq.kv_batch_get(
+ keys,
+ partition_id,
+ select_fields="features",
+ data_parser=select_rows,
+ parser_args=row_indexes,
+ )
+ for i, key in enumerate(keys):
+ value = selected["features"][i]
+ torch.testing.assert_close(value, stored["features"][i][row_indexes[i]])
+ print(f" {key}: returned shape={tuple(value.shape)}")
+ print(f" Returned values: {value.tolist()}")
+ print(" Only the selected rows were sent back from storage.")
+
+ print("\n[Step 5] Reading the full tensors again...")
+ full = tq.kv_batch_get(keys, partition_id, select_fields="features")
+ for i, key in enumerate(keys):
+ torch.testing.assert_close(full["features"][i], stored["features"][i])
+ print(f" {key}: full shape={tuple(full['features'][i].shape)}")
+ print(" Both tensors still contain all 3 rows; the get parser left stored values unchanged.")
+
+ print("\n[Step 6] Clearing the example data...")
+ tq.kv_clear(keys, partition_id)
+ print(" Cleared sample keys:", keys)
+
+ print("\n" + "=" * 80)
+ print("Tutorial Complete!")
+ print("=" * 80)
+ print("\nKey Takeaways:")
+ print(" 1. Put parsers can load actual data from file paths on the storage units")
+ print(" 2. Get parsers transform returned values while leaving stored data unchanged")
+ print(" 3. parser_args provides per-sample options for a get parser")
+ print(" 4. Input files must be accessible from the storage nodes")
+ finally:
+ tq.close()
+ ray.shutdown()
+ print("\nCleanup complete")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tutorial/09_data_dump.py b/tutorial/09_data_dump.py
new file mode 100644
index 00000000..44ff0b6e
--- /dev/null
+++ b/tutorial/09_data_dump.py
@@ -0,0 +1,147 @@
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tutorial 9: Selective Data Dump and Restore
+
+Save selected samples and resume processing after restarting TransferQueue.
+The example saves shared fields separately, restores them before the samples,
+then continues the unfinished responses.
+
+Run: python tutorial/09_data_dump.py
+
+This single-node SimpleStorage example uses a temporary directory. Multi-node
+deployments need a shared filesystem accessible to every storage unit. Pause
+writers and clears for the selected keys during dump/load, and load only trusted
+dumps. A selective restore preserves unrelated data but does not restore sampler
+or consumption state; applications track which keys still need processing.
+
+Restores also support a different number of storage units. For sharing and
+reclaiming public data, see tutorial/07_shared_fields.py.
+"""
+
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import ray
+import torch
+
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+import transfer_queue as tq # noqa: E402
+
+
+def main():
+ """Save selected samples, restore their data and continue updating them."""
+ print("=" * 80)
+ print("TransferQueue Tutorial 9: Selective Data Dump and Restore")
+ print("=" * 80)
+ print("dump_data_by_key saves selected samples; load_data_by_key restores them.")
+ print("This example saves two unfinished samples and resumes them after a restart.")
+ print("Setting up TransferQueue with the default SimpleStorage backend...")
+ ray.init(namespace="TransferQueueTutorial", include_dashboard=False)
+ try:
+ tq.init()
+ partition_id = "resume"
+
+ print("\n[Step 1] Preparing samples with shared fields and unfinished responses...")
+ public_key = "tutorial09/shared_0"
+ features = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+ tq.kv_put(public_key, tq.PUBLIC_PARTITION, fields={"features": features})
+ print(f" Stored {public_key}: features shape={tuple(features.shape)}")
+ pending_keys = ["pending_0", "pending_1"]
+ tq.kv_link(pending_keys, partition_id, public_key=public_key, fields="features")
+ print(f" Linked {pending_keys} to the same public features field.")
+ for i, key in enumerate(pending_keys):
+ first_chunk = torch.tensor([i, i + 1])
+ tq.kv_update(
+ key,
+ partition_id,
+ fields={"response": first_chunk},
+ merge_fn=lambda old, new: torch.cat([old, new]),
+ finalize=False,
+ )
+ tq.kv_put(key, partition_id, tag={"status": "pending", "chunks": 1})
+ print(f" {key}: stored response chunk={first_chunk.tolist()}, finalize=False")
+ print(" The features field is readable; each response remains unfinished.")
+
+ with TemporaryDirectory(prefix="tq-data-dump-") as dump_root:
+ print("\n[Step 2] Saving the selected samples and their shared data...")
+ public_dump = Path(dump_root) / "public"
+ pending_dump = Path(dump_root) / "pending"
+ print(" Sample dumps record links; the referenced public data needs its own dump.")
+ print(" All writes have finished for now; the selected keys stay unchanged during the dump.")
+ tq.dump_data_by_key(public_dump, [public_key], tq.PUBLIC_PARTITION)
+ print(f" Saved public key {public_key} to {public_dump}")
+ report = tq.dump_data_by_key(pending_dump, pending_keys, partition_id)
+ print(f" Saved selected keys {pending_keys} to {pending_dump}")
+ print(" Sample dump report:", report)
+
+ print("\n[Step 3] Restarting TransferQueue before restoring the samples...")
+ tq.close()
+ tq.init()
+ print(" Started an empty TransferQueue; both dumps remain on disk.")
+
+ print("\n[Step 4] Restoring public data first, then the samples...")
+ tq.load_data_by_key(public_dump)
+ print(" Restored public fields:", tq.kv_exists(public_key, tq.PUBLIC_PARTITION))
+ print(" Sample links resolve by public key name, so the expected public data must already be ready.")
+
+ tags = tq.load_data_by_key(pending_dump)
+ descriptions = tq.kv_describe(pending_keys, partition_id)
+ for key in pending_keys:
+ row = descriptions[key]
+ assert row["unfinalized"] == ["response"]
+ print(f" {key}: tags={tags[key]}")
+ print(f" Stored fields={row['fields']}, unfinalized={row['unfinalized']}")
+ print(f" Shared-field links={row['links']}")
+ print(" Restoring data keeps unfinished responses unfinalized.")
+ print(" Sampler and consumption state are not restored; the application chooses which keys to resume.")
+
+ print("\n[Step 5] Continuing the responses and publishing the completed fields...")
+ for i, key in enumerate(pending_keys):
+ final_chunk = torch.tensor([i + 2])
+ print(f" {key}: appending {final_chunk.tolist()} with finalize=True")
+ tq.kv_update(
+ key,
+ partition_id,
+ fields={"response": final_chunk},
+ merge_fn=lambda old, new: torch.cat([old, new]),
+ finalize=True,
+ )
+ tq.kv_put(key, partition_id, tag={"status": "complete", "chunks": 2})
+ batch = tq.kv_batch_get(key, partition_id, select_fields=["response", "features"])
+ assert batch["response"][0].tolist() == [i, i + 1, i + 2]
+ torch.testing.assert_close(batch["features"][0], features)
+ print(f" Readable response={batch['response'][0].tolist()}")
+ print(f" Shared features shape={tuple(batch['features'][0].shape)}")
+ print(" Updated tags:", tq.kv_list(partition_id=partition_id)[partition_id])
+
+ print("\n" + "=" * 80)
+ print("Tutorial Complete!")
+ print("=" * 80)
+ print("\nKey Takeaways:")
+ print(" 1. Selected samples can be saved and restored by key")
+ print(" 2. Dumps preserve fields, tags, links and unfinished field state")
+ print(" 3. Shared data is saved separately and must be ready before restoring linked samples")
+ print(" 4. Restored unfinished responses can be updated and published later")
+ print("\nSee tutorial/07_shared_fields.py for sharing and reclaiming public data.")
+ finally:
+ tq.close()
+ ray.shutdown()
+ print("\nCleanup complete")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tutorial/10_kv_locks.py b/tutorial/10_kv_locks.py
new file mode 100644
index 00000000..eaa6cf7f
--- /dev/null
+++ b/tutorial/10_kv_locks.py
@@ -0,0 +1,134 @@
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tutorial 10: KV Locks
+
+Coordinate operations on keys inside a with block.
+Local locks coordinate threads and async tasks within one process. Global locks
+coordinate callers across processes and nodes in the same TQ cluster.
+
+Run: python tutorial/10_kv_locks.py
+
+Locks work with every backend. Callers that need coordination must use the same
+kind of lock for the same keys; KV APIs do not check these locks automatically.
+A key lock does not automatically cover linked public keys or an entire partition.
+Exceptions do not roll back completed operations.
+
+Global locks renew a lease automatically. Call lease.check() before writing;
+if LockLostError occurs, verify the data before retrying a write that may already
+have completed.
+
+Set TQ_ENABLE_KV_GLOBAL_LOCK=1 before the owning tq.init(); global locks are
+disabled by default. Attached workers discover the lock actors automatically.
+
+Pass multiple keys in one lock call. Nesting the same kind of lock is rejected;
+when combining kinds, acquire global before local. In a running event loop, use
+the async context managers instead; see the KV lock examples in basic.ipynb.
+"""
+
+import os
+import sys
+from pathlib import Path
+
+import ray
+import torch
+
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+import transfer_queue as tq # noqa: E402
+
+
+def main():
+ """Read, modify and write a value while holding a local or global lock."""
+ print("=" * 80)
+ print("TransferQueue Tutorial 10: KV Locks")
+ print("=" * 80)
+ print("The with block can contain several operations on the locked keys.")
+ print("Local locks coordinate threads and async tasks within one process.")
+ print("Global locks coordinate callers across processes and nodes in the same TQ cluster.")
+ print("Callers that need coordination must use the same kind of lock for the same keys.")
+ print("\nOperations that can be grouped inside a lock:")
+ print(" Read or inspect: kv_batch_get, kv_exists, kv_describe, kv_list")
+ print(" Write fields or tags: kv_put, kv_batch_put")
+ print(" Merge or finalize fields: kv_update, kv_batch_update")
+ print(" Clear keys or fields: kv_clear")
+ print(" Link shared fields: kv_link")
+ print(" Save or restore selected data: dump_data_by_key, load_data_by_key")
+ print(" Application logic: validation, computation and decisions between calls")
+ print("A key lock does not automatically lock public dependencies or the whole partition.")
+
+ print("\nSetting up TransferQueue with global locks enabled...")
+ os.environ["TQ_ENABLE_KV_GLOBAL_LOCK"] = "1"
+ ray.init(namespace="TransferQueueTutorial", include_dashboard=False)
+ try:
+ tq.init()
+ partition_id = "locks"
+ print(" TQ_ENABLE_KV_GLOBAL_LOCK=1 was set before tq.init().")
+
+ print("\n[Step 1] Creating counters for the local and global lock examples...")
+ tq.kv_put("local", partition_id, fields={"count": torch.tensor(0)})
+ tq.kv_put("global", partition_id, fields={"count": torch.tensor(0)})
+ print(" Created keys 'local' and 'global', each with count=0.")
+
+ print("\n[Step 2] Updating a counter under a local lock...")
+ print(" Other threads using this local lock must wait until the block ends.")
+ with tq.kv_local_lock("local", partition_id, timeout=10):
+ batch = tq.kv_batch_get("local", partition_id)
+ count = int(batch["count"][0])
+ print(f" Acquired the local lock; read count={count}")
+
+ new_count = count + 1
+ tq.kv_put("local", partition_id, fields={"count": torch.tensor(new_count)})
+ print(f" Wrote count={new_count} while still holding the lock.")
+ print(" Leaving the with block released the local lock.")
+ local_count = int(tq.kv_batch_get("local", partition_id)["count"][0])
+ print(f" Stored count after the update: {local_count}")
+
+ print("\n[Step 3] Updating a counter under a global lock...")
+ print(" Processes on the same or different nodes use this global lock to coordinate writes.")
+ with tq.kv_global_lock("global", partition_id, timeout=10) as lease:
+ batch = tq.kv_batch_get("global", partition_id)
+ count = int(batch["count"][0])
+ print(f" Acquired the global lock; read count={count}")
+
+ new_count = count + 1
+ lease.check()
+ tq.kv_put("global", partition_id, fields={"count": torch.tensor(new_count)})
+ print(f" Checked the lease and wrote count={new_count} while holding the lock.")
+ print(" Leaving the with block released the global lock.")
+ global_count = int(tq.kv_batch_get("global", partition_id)["count"][0])
+ print(f" Stored count after the update: {global_count}")
+
+ print("\n[Step 4] Clearing the example data...")
+ tq.kv_clear(["local", "global"], partition_id)
+ print(" Cleared both counter keys.")
+
+ print("\n" + "=" * 80)
+ print("Tutorial Complete!")
+ print("=" * 80)
+ print("\nKey Takeaways:")
+ print(" 1. Local locks coordinate threads and async tasks within one process")
+ print(" 2. Global locks coordinate callers across processes and nodes in the same TQ cluster")
+ print(" 3. A lock can coordinate several operations within one with block")
+ print(" 4. Global locks use a lease, checked before writing")
+ print(" 5. Locks are released when the with block ends")
+ print("\nSee tutorial/basic.ipynb for async with examples.")
+ finally:
+ tq.close()
+ ray.shutdown()
+ print("\nCleanup complete")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tutorial/basic.ipynb b/tutorial/basic.ipynb
index 29976e61..bedd6f9b 100644
--- a/tutorial/basic.ipynb
+++ b/tutorial/basic.ipynb
@@ -18,29 +18,52 @@
"4. Retrieve data — `kv_batch_get`\n",
"5. List stored keys & tags — `kv_list`\n",
"6. Partial-key and partial-field retrieval\n",
- "7. Updating fields incrementally\n",
+ "7. Add fields, merge response chunks and clear selected fields\n",
"8. Working with nested (variable-length) tensors\n",
"9. Storing variable-size image data\n",
"10. Storing non-tensor data (`NonTensorStack`)\n",
- "11. Lazy Data Parsing with `data_parser`\n",
+ "11. Load files on put and select tensor rows on get\n",
"12. Multiple partitions\n",
- "13. Clean up — `kv_clear` / `close`\n",
+ "13. Share fields across samples\n",
+ "14. Coordinate operations with local and global KV locks\n",
+ "15. Save and restore selected samples\n",
+ "16. Reclaim data and close TransferQueue\n",
"\n",
- "> **Prerequisites:** `pip install TransferQueue` (or install from source). \n",
- "> Ray will be started automatically in this notebook."
+ "**Prerequisites:** Install this checkout with `pip install -e .` from the repository\n",
+ "root and select a notebook kernel using that environment. The examples use CPU\n",
+ "PyTorch and SimpleStorage.\n",
+ "\n",
+ "Restart the kernel and run the cells in order. The initialization cell enables\n",
+ "global locks before starting TQ so the later lock examples can use them."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 1. Initialization\n\nTransferQueue runs on top of [Ray](https://www.ray.io/). \nWe start Ray, then call `tq.init()` with a minimal configuration that uses the\nbuilt-in **SimpleStorage** backend."
+ "## 1. Initialization\n",
+ "\n",
+ "TransferQueue runs on top of [Ray](https://www.ray.io/). \n",
+ "We start Ray, then call `tq.init()` with a minimal configuration that uses the\n",
+ "built-in **SimpleStorage** backend."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "TransferQueue is ready!\n"
+ ]
+ }
+ ],
"source": [
+ "import os\n",
+ "\n",
"import ray\n",
"import torch\n",
"from omegaconf import OmegaConf\n",
@@ -49,7 +72,8 @@
"\n",
"import transfer_queue as tq\n",
"\n",
- "ray.init(ignore_reinit_error=True)\n",
+ "os.environ[\"TQ_ENABLE_KV_GLOBAL_LOCK\"] = \"1\"\n",
+ "ray.init(ignore_reinit_error=True, include_dashboard=False)\n",
"\n",
"config = OmegaConf.create(\n",
" {\n",
@@ -66,65 +90,65 @@
"\n",
"tq.init(config)\n",
"print(\"TransferQueue is ready!\")"
- ],
- "outputs": [
- {
- "name": "stderr",
- "output_type": "stream",
- "text": [
- ""
- ]
- },
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "TransferQueue is ready!\n"
- ]
- }
- ],
- "execution_count": 1
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 2. Store a Single Sample — `kv_put`\n\n`kv_put` stores **one** key-value pair. \n- `key` — a unique string identifier for the sample \n- `partition_id` — a logical namespace (like a table name) \n- `fields` — a `dict` of tensors **or** a `TensorDict` \n- `tag` — optional metadata dict attached to the key"
+ "## 2. Store a Single Sample — `kv_put`\n",
+ "\n",
+ "`kv_put` stores **one** key-value pair. \n",
+ "- `key` — a unique string identifier for the sample \n",
+ "- `partition_id` — a logical namespace (like a table name) \n",
+ "- `fields` — a `dict` of tensors **or** a `TensorDict` \n",
+ "- `tag` — optional metadata dict attached to the key"
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "tq.kv_put(\n",
- " key=\"sample_0\",\n",
- " partition_id=\"train\",\n",
- " fields={\"input_ids\": torch.tensor([1, 2, 3, 4])},\n",
- " tag={\"source\": \"wikipedia\", \"score\": 0.95},\n",
- ")\n",
- "print(\"Stored sample_0\")"
- ],
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
"Stored sample_0\n"
]
}
],
- "execution_count": 2
+ "source": [
+ "tq.kv_put(\n",
+ " key=\"sample_0\",\n",
+ " partition_id=\"train\",\n",
+ " fields={\"input_ids\": torch.tensor([1, 2, 3, 4])},\n",
+ " tag={\"source\": \"wikipedia\", \"score\": 0.95},\n",
+ ")\n",
+ "print(\"Stored sample_0\")"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "You can also pass a pre-built `TensorDict` directly (the batch dimension\nmust be 1 for `kv_put`):"
+ "You can also pass a pre-built `TensorDict` directly (the batch dimension\n",
+ "must be 1 for `kv_put`):"
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Stored sample_1\n"
+ ]
+ }
+ ],
"source": [
"fields_td = TensorDict(\n",
" {\n",
@@ -141,28 +165,32 @@
" tag={\"source\": \"books\", \"score\": 0.88},\n",
")\n",
"print(\"Stored sample_1\")"
- ],
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Stored sample_1\n"
- ]
- }
- ],
- "execution_count": 3
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 3. Store a Batch of Samples — `kv_batch_put`\n\nWhen you have multiple samples, `kv_batch_put` is more efficient than\ncalling `kv_put` in a loop. The `fields` TensorDict must have\n`batch_size == len(keys)`."
+ "## 3. Store a Batch of Samples — `kv_batch_put`\n",
+ "\n",
+ "When you have multiple samples, `kv_batch_put` is more efficient than\n",
+ "calling `kv_put` in a loop. The `fields` TensorDict must have\n",
+ "`batch_size == len(keys)`."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Stored 3 samples in one call\n"
+ ]
+ }
+ ],
"source": [
"keys = [\"batch_0\", \"batch_1\", \"batch_2\"]\n",
"\n",
@@ -182,73 +210,48 @@
"\n",
"tq.kv_batch_put(keys=keys, partition_id=\"train\", fields=fields, tags=tags)\n",
"print(f\"Stored {len(keys)} samples in one call\")"
- ],
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Stored 3 samples in one call\n"
- ]
- }
- ],
- "execution_count": 4
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 4. Retrieve Data — `kv_batch_get`\n\nRetrieve samples by key(s). The result is always a `TensorDict`."
+ "## 4. Retrieve Data — `kv_batch_get`\n",
+ "\n",
+ "Retrieve samples by key(s). The result is always a `TensorDict`."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "# Retrieve a single key (pass a string)\n",
- "result = tq.kv_batch_get(keys=\"sample_0\", partition_id=\"train\")\n",
- "print(\"sample_0 →\", result)\n",
- "print(\"input_ids:\", result[\"input_ids\"])"
- ],
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "sample_0 → TensorDict(\n",
- " fields={\n",
- " input_ids: NestedTensor(shape=torch.Size([1, j1]), device=cpu, dtype=torch.int64, is_shared=False)},\n",
- " batch_size=torch.Size([1]),\n",
- " device=None,\n",
- " is_shared=False)\n",
+ "sample_0 → TensorDict(batch_size=[1], fields={input_ids: i64[j1]})\n",
"input_ids: NestedTensor(size=(1, j1), offsets=tensor([0, 4]), contiguous=True)\n"
]
}
],
- "execution_count": 5
+ "source": [
+ "# Retrieve a single key (pass a string)\n",
+ "result = tq.kv_batch_get(keys=\"sample_0\", partition_id=\"train\")\n",
+ "print(\"sample_0 →\", result)\n",
+ "print(\"input_ids:\", result[\"input_ids\"])"
+ ]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "# Retrieve multiple keys at once\n",
- "result = tq.kv_batch_get(keys=keys, partition_id=\"train\")\n",
- "print(\"batch result →\", result)\n",
- "print(\"input_ids:\\n\", result[\"input_ids\"])\n",
- "print(\"attention_mask:\\n\", result[\"attention_mask\"])"
- ],
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "batch result → TensorDict(\n",
- " fields={\n",
- " attention_mask: NestedTensor(shape=torch.Size([3, j2]), device=cpu, dtype=torch.int64, is_shared=False),\n",
- " input_ids: NestedTensor(shape=torch.Size([3, j3]), device=cpu, dtype=torch.int64, is_shared=False)},\n",
- " batch_size=torch.Size([3]),\n",
- " device=None,\n",
- " is_shared=False)\n",
+ "batch result → TensorDict(batch_size=[3], fields={attention_mask: i64[j2], input_ids: i64[j3]})\n",
"input_ids:\n",
" NestedTensor(size=(3, j3), offsets=tensor([0, 2, 4, 6]), contiguous=True)\n",
"attention_mask:\n",
@@ -256,30 +259,35 @@
]
}
],
- "execution_count": 6
+ "source": [
+ "# Retrieve multiple keys at once\n",
+ "result = tq.kv_batch_get(keys=keys, partition_id=\"train\")\n",
+ "print(\"batch result →\", result)\n",
+ "print(\"input_ids:\\n\", result[\"input_ids\"])\n",
+ "print(\"attention_mask:\\n\", result[\"attention_mask\"])"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 5. List Keys & Tags — `kv_list`\n\n`kv_list` returns a nested dict:\n```\n{ partition_id: { key: tag_dict, ... }, ... }\n```\nPass `partition_id` to filter, or omit it to see everything."
+ "## 5. List Keys & Tags — `kv_list`\n",
+ "\n",
+ "`kv_list` returns a nested dict:\n",
+ "```\n",
+ "{ partition_id: { key: tag_dict, ... }, ... }\n",
+ "```\n",
+ "Pass `partition_id` to filter, or omit it to see everything."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "info = tq.kv_list(partition_id=\"train\")\n",
- "\n",
- "for partition, key_tags in info.items():\n",
- " print(f\"\\nPartition: {partition}\")\n",
- " for key, tag in key_tags.items():\n",
- " print(f\" {key}: {tag}\")"
- ],
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
"\n",
"Partition: train\n",
@@ -291,52 +299,79 @@
]
}
],
- "execution_count": 7
+ "source": [
+ "info = tq.kv_list(partition_id=\"train\")\n",
+ "\n",
+ "for partition, key_tags in info.items():\n",
+ " print(f\"\\nPartition: {partition}\")\n",
+ " for key, tag in key_tags.items():\n",
+ " print(f\" {key}: {tag}\")"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 6. Partial-Key and Partial-Field Retrieval\n\nYou don't have to retrieve *all* keys or *all* fields at once."
+ "## 6. Partial-Key and Partial-Field Retrieval\n",
+ "\n",
+ "You don't have to retrieve *all* keys or *all* fields at once."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "### 6a. Partial Keys\n\nJust pass a subset of the keys you stored."
+ "### 6a. Partial Keys\n",
+ "\n",
+ "Just pass a subset of the keys you stored."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "partial = tq.kv_batch_get(keys=[\"batch_0\", \"batch_2\"], partition_id=\"train\")\n",
- "print(\"Partial-key input_ids:\\n\", partial[\"input_ids\"])\n",
- "assert partial[\"input_ids\"].shape[0] == 2 # only 2 rows"
- ],
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
"Partial-key input_ids:\n",
" NestedTensor(size=(2, j5), offsets=tensor([0, 2, 4]), contiguous=True)\n"
]
}
],
- "execution_count": 8
+ "source": [
+ "partial = tq.kv_batch_get(keys=[\"batch_0\", \"batch_2\"], partition_id=\"train\")\n",
+ "print(\"Partial-key input_ids:\\n\", partial[\"input_ids\"])\n",
+ "assert partial[\"input_ids\"].shape[0] == 2 # only 2 rows"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "### 6b. Partial Fields\n\nUse the `fields` argument to select specific columns."
+ "### 6b. Partial Fields\n",
+ "\n",
+ "Use `select_fields` to select columns. Every requested key must exist, and every\n",
+ "explicitly selected field must be ready in every key; otherwise the read raises\n",
+ "`ValueError`. Without `select_fields`, a batch read returns only fields ready in\n",
+ "all requested keys. Use `kv_exists` to probe optional keys without an error."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Fields returned: ['input_ids']\n",
+ "Fields returned: ['attention_mask', 'input_ids']\n"
+ ]
+ }
+ ],
"source": [
"# Retrieve only input_ids (single field)\n",
"result = tq.kv_batch_get(keys=\"sample_1\", partition_id=\"train\", select_fields=\"input_ids\")\n",
@@ -351,29 +386,36 @@
" select_fields=[\"input_ids\", \"attention_mask\"],\n",
")\n",
"print(\"Fields returned:\", list(result.keys()))"
- ],
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Fields returned: ['input_ids']\n",
- "Fields returned: ['attention_mask', 'input_ids']\n"
- ]
- }
- ],
- "execution_count": 9
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 7. Updating Fields Incrementally\n\nTransferQueue tracks each field (column) independently per key (row). \nYou can **add new fields** to existing keys with another `kv_put` /\n`kv_batch_put` call — the earlier fields are preserved."
+ "## 7. Adding and Updating Fields\n",
+ "\n",
+ "### 7a. Add a new field\n",
+ "\n",
+ "`kv_put` / `kv_batch_put` add or replace the fields supplied in the call, preserving\n",
+ "the other fields. Tags hold application metadata; TQ tracks field readiness\n",
+ "separately from those tags."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "All fields now: ['attention_mask', 'input_ids', 'response']\n",
+ "response:\n",
+ " NestedTensor(size=(3, j11), offsets=tensor([0, 2, 4, 6]), contiguous=True)\n"
+ ]
+ }
+ ],
"source": [
"# Add a \"response\" field to the batch keys\n",
"response_fields = TensorDict(\n",
@@ -386,30 +428,146 @@
"result = tq.kv_batch_get(keys=keys, partition_id=\"train\")\n",
"print(\"All fields now:\", list(result.keys()))\n",
"print(\"response:\\n\", result[\"response\"])"
- ],
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### 7b. Build a response in chunks\n",
+ "\n",
+ "`kv_update` merges new values into a stored field with `merge_fn(old, new)`.\n",
+ "A missing field takes the first value directly. With `finalize=False`, the field\n",
+ "stays unready until the final update publishes it with `finalize=True`.\n",
+ "\n",
+ "`kv_describe` shows stored and unfinalized fields; `kv_exists` reports ready fields.\n",
+ "This example uses two fixed token chunks to illustrate the updates. See\n",
+ "[Tutorial 02](02_kv_interface.py) for the KV interface walkthrough."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
"outputs": [
{
+ "output_type": "stream",
"name": "stdout",
+ "text": [
+ "[Step 1] Storing the first response chunk with finalize=False...\n",
+ " Stored fields: {'fields': ['response'], 'unfinalized': ['response'], 'links': {}}\n",
+ " Ready fields: []\n",
+ "\n",
+ "[Step 2] Appending the final chunk and publishing the response...\n",
+ " Published response: [10, 11, 12, 13]\n"
+ ]
+ },
+ {
+ "output_type": "execute_result",
+ "metadata": {},
+ "data": {
+ "text/plain": [
+ "KVBatchMeta(keys=['streaming_sample'], tags=[{'status': 'finish'}], partition_id='train', fields=['response'], extra_info={})"
+ ]
+ },
+ "execution_count": null
+ }
+ ],
+ "source": [
+ "streaming_key = \"streaming_sample\"\n",
+ "tq.kv_put(streaming_key, \"train\", tag={\"status\": \"running\"})\n",
+ "\n",
+ "print(\"[Step 1] Storing the first response chunk with finalize=False...\")\n",
+ "tq.kv_update(\n",
+ " streaming_key,\n",
+ " \"train\",\n",
+ " fields={\"response\": torch.tensor([10, 11])},\n",
+ " merge_fn=lambda old, new: torch.cat([old, new]),\n",
+ " finalize=False,\n",
+ ")\n",
+ "print(\" Stored fields:\", tq.kv_describe(streaming_key, \"train\")[streaming_key])\n",
+ "print(\" Ready fields:\", tq.kv_exists(streaming_key, \"train\")[streaming_key])\n",
+ "\n",
+ "print(\"\\n[Step 2] Appending the final chunk and publishing the response...\")\n",
+ "tq.kv_update(\n",
+ " streaming_key,\n",
+ " \"train\",\n",
+ " fields={\"response\": torch.tensor([12, 13])},\n",
+ " merge_fn=lambda old, new: torch.cat([old, new]),\n",
+ " finalize=True,\n",
+ ")\n",
+ "complete = tq.kv_batch_get(streaming_key, \"train\", select_fields=\"response\")\n",
+ "assert complete[\"response\"][0].tolist() == [10, 11, 12, 13]\n",
+ "print(\" Published response:\", complete[\"response\"][0].tolist())\n",
+ "tq.kv_put(streaming_key, \"train\", tag={\"status\": \"finish\"})"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Incremental updates require SimpleStorage. Use one writer per key or coordinate\n",
+ "writers with a KV lock. A timed-out merge may have committed, so retrying it\n",
+ "blindly can append the same chunk twice.\n",
+ "\n",
+ "### 7c. Clear selected fields while keeping keys\n",
+ "\n",
+ "The batch samples from section 3 still have `attention_mask`. Clearing that field\n",
+ "keeps their other fields, tags and consumption state."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
"output_type": "stream",
+ "name": "stdout",
"text": [
- "All fields now: ['attention_mask', 'input_ids', 'response']\n",
- "response:\n",
- " NestedTensor(size=(3, j11), offsets=tensor([0, 2, 4, 6]), contiguous=True)\n"
+ "Clearing attention_mask from the batch samples...\n",
+ " Keys and remaining fields: {'batch_0': ['input_ids', 'response'], 'batch_1': ['input_ids', 'response'], 'batch_2': ['input_ids', 'response']}\n",
+ " Tags: {'batch_0': {'split': 'train', 'idx': 0}, 'batch_1': {'split': 'train', 'idx': 1}, 'batch_2': {'split': 'train', 'idx': 2}}\n"
]
}
],
- "execution_count": 10
+ "source": [
+ "print(\"Clearing attention_mask from the batch samples...\")\n",
+ "tq.kv_clear(keys, \"train\", fields=\"attention_mask\")\n",
+ "print(\" Keys and remaining fields:\", tq.kv_exists(keys, \"train\"))\n",
+ "tags_by_key = tq.kv_list(\"train\")[\"train\"]\n",
+ "print(\" Tags:\", {key: tags_by_key[key] for key in keys})"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 8. Working with Nested (Variable-Length) Tensors\n\nIn many NLP and RL workloads each sample has a **different sequence length**\n(e.g. generated responses). PyTorch represents these as\n[nested tensors](https://pytorch.org/docs/stable/nested.html) with the\n**jagged layout** (`layout=torch.jagged`), and TransferQueue handles them\nnatively.\n\n> **Note:** Because individual samples have different shapes, you must use\n> `kv_batch_put` (not `kv_put`) to store nested tensors."
+ "## 8. Working with Nested (Variable-Length) Tensors\n",
+ "\n",
+ "In many NLP and RL workloads each sample has a **different sequence length**\n",
+ "(e.g. generated responses). PyTorch represents these as\n",
+ "[nested tensors](https://pytorch.org/docs/stable/nested.html) with the\n",
+ "**jagged layout** (`layout=torch.jagged`), and TransferQueue handles them\n",
+ "natively.\n",
+ "\n",
+ "> **Note:** Because individual samples have different shapes, you must use\n",
+ "> `kv_batch_put` (not `kv_put`) to store nested tensors."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Stored 3 samples with variable-length responses\n"
+ ]
+ }
+ ],
"source": [
"nested_keys = [\"nested_0\", \"nested_1\", \"nested_2\"]\n",
"\n",
@@ -438,21 +596,27 @@
" tags=[{\"len\": 3}, {\"len\": 1}, {\"len\": 2}],\n",
")\n",
"print(\"Stored 3 samples with variable-length responses\")"
- ],
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "Stored 3 samples with variable-length responses\n"
+ "input_ids:\n",
+ " NestedTensor(size=(3, j14), offsets=tensor([0, 2, 4, 6]), contiguous=True)\n",
+ "\n",
+ "response (nested tensor):\n",
+ " sample 0: tensor([10, 11, 12]) (length 3)\n",
+ " sample 1: tensor([20]) (length 1)\n",
+ " sample 2: tensor([30, 31]) (length 2)\n"
]
}
],
- "execution_count": 11
- },
- {
- "cell_type": "code",
- "metadata": {},
"source": [
"# Retrieve all nested samples\n",
"result = tq.kv_batch_get(keys=nested_keys, partition_id=\"train\")\n",
@@ -460,54 +624,27 @@
"print(\"\\nresponse (nested tensor):\")\n",
"for i, sample in enumerate(result[\"response\"]):\n",
" print(f\" sample {i}: {sample} (length {sample.shape[0]})\")"
- ],
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "input_ids:\n",
- " NestedTensor(size=(3, j13), offsets=tensor([0, 2, 4, 6]), contiguous=True)\n",
- "\n",
- "response (nested tensor):\n",
- " sample 0: tensor([10, 11, 12]) (length 3)\n",
- " sample 1: tensor([20]) (length 1)\n",
- " sample 2: tensor([30, 31]) (length 2)\n"
- ]
- }
- ],
- "execution_count": 12
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "Partial-key retrieval works the same way with nested tensors — only the\nrequested samples are returned:"
+ "Partial-key retrieval works the same way with nested tensors — only the\n",
+ "requested samples are returned:"
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "# Retrieve only the first and last sample\n",
- "partial = tq.kv_batch_get(keys=[\"nested_0\", \"nested_2\"], partition_id=\"train\")\n",
- "print(\"Partial-key input_ids:\\n\", partial[\"input_ids\"])\n",
- "print(\"\\nPartial-key responses:\")\n",
- "for i, sample in enumerate(partial[\"response\"]):\n",
- " print(f\" sample {i}: {sample}\")\n",
- "\n",
- "# Verify correctness\n",
- "assert torch.equal(partial[\"response\"][0], torch.tensor([10, 11, 12]))\n",
- "assert torch.equal(partial[\"response\"][1], torch.tensor([30, 31]))\n",
- "print(\"\\nAssertions passed!\")"
- ],
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
"Partial-key input_ids:\n",
- " NestedTensor(size=(2, j15), offsets=tensor([0, 2, 4]), contiguous=True)\n",
+ " NestedTensor(size=(2, j16), offsets=tensor([0, 2, 4]), contiguous=True)\n",
"\n",
"Partial-key responses:\n",
" sample 0: tensor([10, 11, 12])\n",
@@ -517,18 +654,44 @@
]
}
],
- "execution_count": 13
+ "source": [
+ "# Retrieve only the first and last sample\n",
+ "partial = tq.kv_batch_get(keys=[\"nested_0\", \"nested_2\"], partition_id=\"train\")\n",
+ "print(\"Partial-key input_ids:\\n\", partial[\"input_ids\"])\n",
+ "print(\"\\nPartial-key responses:\")\n",
+ "for i, sample in enumerate(partial[\"response\"]):\n",
+ " print(f\" sample {i}: {sample}\")\n",
+ "\n",
+ "# Verify correctness\n",
+ "assert torch.equal(partial[\"response\"][0], torch.tensor([10, 11, 12]))\n",
+ "assert torch.equal(partial[\"response\"][1], torch.tensor([30, 31]))\n",
+ "print(\"\\nAssertions passed!\")"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "Higher-dimensional nested tensors work too. Here each sample is a 3D\ntensor with a variable first dimension (e.g. a different number of\nattention heads or generated candidates):"
+ "Higher-dimensional nested tensors work too. Here each sample is a 3D\n",
+ "tensor with a variable first dimension (e.g. a different number of\n",
+ "attention heads or generated candidates):"
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "sample 0 hidden_states shape: torch.Size([2, 3, 4])\n",
+ "sample 1 hidden_states shape: torch.Size([5, 3, 4])\n",
+ "sample 2 hidden_states shape: torch.Size([1, 3, 4])\n"
+ ]
+ }
+ ],
"source": [
"nested_3d_keys = [\"nd3d_0\", \"nd3d_1\", \"nd3d_2\"]\n",
"\n",
@@ -558,40 +721,48 @@
"# Clean up nested-tensor keys\n",
"tq.kv_clear(keys=nested_keys, partition_id=\"train\")\n",
"tq.kv_clear(keys=nested_3d_keys, partition_id=\"train\")"
- ],
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "sample 0 hidden_states shape: torch.Size([2, 3, 4])\n",
- "sample 1 hidden_states shape: torch.Size([5, 3, 4])\n",
- "sample 2 hidden_states shape: torch.Size([1, 3, 4])\n"
- ]
- }
- ],
- "execution_count": 14
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 9. Storing Variable-Size Image Data\n\nA common multimodal scenario: each sample in a batch contains a\n**different number of images**, and each image has a **different\nresolution**. We can model this with a list of nested tensors — one\nnested tensor per sample — wrapped inside a `TensorDict`.\n\nSince the data is doubly ragged (variable count *and* variable size),\nwe store each image as a flattened 1-D tensor and pack all images per\nsample into a single jagged nested tensor. This way every sample is\none element of the batch, yet images retain their individual sizes."
+ "## 9. Storing Variable-Size Image Data\n",
+ "\n",
+ "A common multimodal scenario: each sample in a batch contains a\n",
+ "**different number of images**, and each image has a **different\n",
+ "resolution**. We can model this with a list of nested tensors — one\n",
+ "nested tensor per sample — wrapped inside a `TensorDict`.\n",
+ "\n",
+ "Since the data is doubly ragged (variable count *and* variable size),\n",
+ "we store each image as a flattened 1-D tensor and pack all images per\n",
+ "sample into a single jagged nested tensor. This way every sample is\n",
+ "one element of the batch, yet images retain their individual sizes."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Stored 3 samples with variable numbers of variable-size images\n"
+ ]
+ }
+ ],
"source": [
"image_keys = [\"img_sample_0\", \"img_sample_1\", \"img_sample_2\"]\n",
"\n",
- "# Sample 0: 2 images — 3\u00d732\u00d732 (RGB, 32\u00d732) and 3\u00d764\u00d764\n",
+ "# Sample 0: 2 images — 3×32×32 (RGB, 32×32) and 3×64×64\n",
"sample_0_images = [torch.randn(3, 32, 32), torch.randn(3, 64, 64)]\n",
"\n",
- "# Sample 1: 1 image — 3\u00d748\u00d748\n",
+ "# Sample 1: 1 image — 3×48×48\n",
"sample_1_images = [torch.randn(3, 48, 48)]\n",
"\n",
- "# Sample 2: 3 images — 3\u00d716\u00d716, 3\u00d724\u00d724, 3\u00d732\u00d764\n",
+ "# Sample 2: 3 images — 3×16×16, 3×24×24, 3×32×64\n",
"sample_2_images = [torch.randn(3, 16, 16), torch.randn(3, 24, 24), torch.randn(3, 32, 64)]\n",
"\n",
"\n",
@@ -636,21 +807,23 @@
"\n",
"tq.kv_batch_put(keys=image_keys, partition_id=\"train\", fields=fields_img, tags=tags_img)\n",
"print(\"Stored 3 samples with variable numbers of variable-size images\")"
- ],
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "Stored 3 samples with variable numbers of variable-size images\n"
+ "Sample 0: 2 image(s) → [(3, 32, 32), (3, 64, 64)]\n",
+ "Sample 1: 1 image(s) → [(3, 48, 48)]\n",
+ "Sample 2: 3 image(s) → [(3, 16, 16), (3, 24, 24), (3, 32, 64)]\n"
]
}
],
- "execution_count": 15
- },
- {
- "cell_type": "code",
- "metadata": {},
"source": [
"# Retrieve and reconstruct the images\n",
"result_img = tq.kv_batch_get(keys=image_keys, partition_id=\"train\")\n",
@@ -673,19 +846,7 @@
"\n",
"# Clean up\n",
"tq.kv_clear(keys=image_keys, partition_id=\"train\")"
- ],
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Sample 0: 2 image(s) → [(3, 32, 32), (3, 64, 64)]\n",
- "Sample 1: 1 image(s) → [(3, 48, 48)]\n",
- "Sample 2: 3 image(s) → [(3, 16, 16), (3, 24, 24), (3, 32, 64)]\n"
- ]
- }
- ],
- "execution_count": 16
+ ]
},
{
"cell_type": "markdown",
@@ -704,7 +865,17 @@
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Stored 3 samples with non-tensor fields\n"
+ ]
+ }
+ ],
"source": [
"nt_keys = [\"nt_sample_0\", \"nt_sample_1\", \"nt_sample_2\"]\n",
"\n",
@@ -718,7 +889,7 @@
" \"Translate to French:\",\n",
" \"Write a poem about rain.\",\n",
" ),\n",
- " # You can also store richer Python objects (dicts, lists, \u2026)\n",
+ " # You can also store richer Python objects (dicts, lists, …)\n",
" \"metadata\": NonTensorStack(\n",
" {\"source\": \"wiki\", \"lang\": \"en\"},\n",
" {\"source\": \"books\", \"lang\": \"fr\"},\n",
@@ -730,21 +901,34 @@
"\n",
"tq.kv_batch_put(keys=nt_keys, partition_id=\"train\", fields=nt_fields)\n",
"print(\"Stored 3 samples with non-tensor fields\")"
- ],
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "Stored 3 samples with non-tensor fields\n"
+ "input_ids:\n",
+ " NestedTensor(size=(3, j26), offsets=tensor([0, 3, 6, 9]), contiguous=True)\n",
+ "\n",
+ "prompt_text (NonTensorStack):\n",
+ " [0] 'Summarise the following article.'\n",
+ " [1] 'Translate to French:'\n",
+ " [2] 'Write a poem about rain.'\n",
+ "\n",
+ "metadata (NonTensorStack of dicts):\n",
+ " [0] {'source': 'wiki', 'lang': 'en'}\n",
+ " [1] {'source': 'books', 'lang': 'fr'}\n",
+ " [2] {'source': 'user', 'lang': 'en'}\n",
+ "\n",
+ "Single-key non-tensor field: label=[['positive']]\n"
]
}
],
- "execution_count": 17
- },
- {
- "cell_type": "code",
- "metadata": {},
"source": [
"# Retrieve and inspect non-tensor fields\n",
"result_nt = tq.kv_batch_get(keys=nt_keys, partition_id=\"train\")\n",
@@ -771,300 +955,457 @@
"# Clean up\n",
"tq.kv_clear(keys=nt_keys, partition_id=\"train\")\n",
"tq.kv_clear(keys=\"single_nt\", partition_id=\"train\")"
- ],
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 11. Parsing on the Storage Unit\n",
+ "\n",
+ "### 11a. Load a file on put\n",
+ "\n",
+ "A `data_parser` runs on SimpleStorage: the caller sends a file path, and storage\n",
+ "loads its tensor. The file must be accessible to storage until the put returns."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "input_ids:\n",
- " NestedTensor(size=(3, j25), offsets=tensor([0, 3, 6, 9]), contiguous=True)\n",
- "\n",
- "prompt_text (NonTensorStack):\n",
- " [0] 'Summarise the following article.'\n",
- " [1] 'Translate to French:'\n",
- " [2] 'Write a poem about rain.'\n",
- "\n",
- "metadata (NonTensorStack of dicts):\n",
- " [0] {'source': 'wiki', 'lang': 'en'}\n",
- " [1] {'source': 'books', 'lang': 'fr'}\n",
- " [2] {'source': 'user', 'lang': 'en'}\n",
- "\n",
- "Single-key non-tensor field: label=[['positive']]\n"
+ "Loading features.pt on the storage unit...\n",
+ "Stored features:\n",
+ " tensor([[ 0., 1., 2., 3.],\n",
+ " [ 4., 5., 6., 7.],\n",
+ " [ 8., 9., 10., 11.]])\n"
]
}
],
- "execution_count": 18
+ "source": [
+ "from pathlib import Path\n",
+ "from tempfile import TemporaryDirectory\n",
+ "\n",
+ "\n",
+ "def load_features(data):\n",
+ " data[\"features\"] = [torch.load(path, map_location=\"cpu\", weights_only=True) for path in data[\"features\"]]\n",
+ " return data\n",
+ "\n",
+ "\n",
+ "data_parser_keys = [\"data_parser_sample_0\"]\n",
+ "parser_partition = \"parser\"\n",
+ "with TemporaryDirectory(prefix=\"tq-notebook-parser-\") as data_dir:\n",
+ " file_path = Path(data_dir) / \"features.pt\"\n",
+ " torch.save(torch.arange(12, dtype=torch.float32).reshape(3, 4), file_path)\n",
+ " print(\"Loading features.pt on the storage unit...\")\n",
+ " tq.kv_batch_put(\n",
+ " data_parser_keys,\n",
+ " parser_partition,\n",
+ " fields=TensorDict({\"features\": NonTensorStack(str(file_path))}, batch_size=1),\n",
+ " data_parser=load_features,\n",
+ " )\n",
+ "\n",
+ "stored = tq.kv_batch_get(data_parser_keys, parser_partition)\n",
+ "print(\"Stored features:\\n\", stored[\"features\"][0])"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 11. Lazy Data Parsing with `data_parser`\n",
- "\n",
- "Sometimes you want to store lightweight **references** (e.g. URLs or file paths)\n",
- "and defer the expensive loading / decoding until the data reaches the storage unit.\n",
- "\n",
- "`kv_put` / `kv_batch_put` accept an optional `data_parser` callable. It is executed **inside**\n",
- "each `SimpleStorageUnit` at put time. The callable receives a plain `dict` (not a TensorDict)\n",
- "mapping `field_name -> batched_values`. For a regular tensor column the value is a batched tensor;\n",
- "for nested tensors (jagged or strided) and `NonTensorStack` columns the values are extracted into\n",
- "a `list`. It must modify values in-place based on the original keys; do not add or remove keys.\n",
- "The number of elements per column must also remain unchanged. Do not change the inner order of\n",
- "values within each column.\n",
- "\n",
- "> **Design tip:** Separate the **core single-sample parser** from the **batch concurrency wrapper**.\n",
- "> The wrapper can use `asyncio` to process all samples in parallel while the parser function itself\n",
- "> remains synchronous to the caller.\n",
+ "### 11b. Select rows on get\n",
"\n",
- "> **Note:** `data_parser` is only supported by the **SimpleStorage** backend."
+ "A get parser can select tensor rows before returning data to the caller.\n",
+ "`parser_args` supplies one argument per requested key. Parsers preserve field\n",
+ "names and sample order; get parsers must leave stored values unchanged.\n",
+ "See [Tutorial 08](08_data_parser.py) for batch examples and usage constraints."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Rows 0 and 2 returned from storage:\n",
+ " tensor([[ 0., 1., 2., 3.],\n",
+ " [ 8., 9., 10., 11.]])\n"
+ ]
+ }
+ ],
"source": [
- "import asyncio\n",
- "import time\n",
- "\n",
- "\n",
- "# 1. Define core single-sample parser (pure business logic, no asyncio, no batch)\n",
- "def parse_url(url: str) -> torch.Tensor:\n",
- " \"\"\"Parse a URL-like descriptor 'dtype:HxW' into a random tensor.\"\"\"\n",
- " dtype_str, shape_str = url.split(\":\")\n",
- " dtype = getattr(torch, dtype_str)\n",
- " shape = [int(dim) for dim in shape_str.split(\"x\")]\n",
- " return torch.randn(shape, dtype=dtype)\n",
- "\n",
- "\n",
- "# 2. Define Batch-level parser: sync on the outside, async-parallel on the inside\n",
- "def concurrent_batch_url_parser(field_data: dict) -> dict:\n",
- " \"\"\"Batch-level data_parser executed inside SimpleStorageUnit.\n",
- "\n",
- " It receives a ``dict`` (not a TensorDict) where each value is a\n",
- " batched column. For columns created from ``NonTensorStack`` the\n",
- " value is a plain ``list`` of Python objects.\n",
- "\n",
- " Workflow:\n",
- " 1. Spawns one async task per list element.\n",
- " 2. Waits until *all* tasks finish (``asyncio.gather``).\n",
- " 3. Replaces the list with the list of results.\n",
- "\n",
- " Because ``asyncio.run`` blocks until the loop finishes, this function\n",
- " is **synchronous** to its caller: when it returns, every sample has\n",
- " been processed.\n",
- "\n",
- " Args:\n",
- " field_data: Mapping ``field_name -> batched_values``. The dict\n",
- " keys must stay exactly the same; only values may be\n",
- " transformed in-place.\n",
- "\n",
- " Returns:\n",
- " The same dict with parsed values substituted.\n",
- " \"\"\"\n",
- " if \"data_to_be_parsed\" not in field_data:\n",
- " return field_data\n",
+ "def select_rows(data, row_indexes):\n",
+ " return {\"features\": [value[indexes] for value, indexes in zip(data[\"features\"], row_indexes, strict=True)]}\n",
"\n",
- " urls: list[str] = field_data[\"data_to_be_parsed\"]\n",
"\n",
- " async def _async_parse_single(url: str) -> torch.Tensor:\n",
- " await asyncio.sleep(1.0) # Add fixed delay per sample\n",
- " return parse_url(url)\n",
- "\n",
- " async def _process_all():\n",
- " tasks = [asyncio.create_task(_async_parse_single(url)) for url in urls]\n",
- " return await asyncio.gather(*tasks)\n",
- "\n",
- " start = time.perf_counter()\n",
- " field_data[\"data_to_be_parsed\"] = asyncio.run(_process_all())\n",
- " elapsed = time.perf_counter() - start\n",
- "\n",
- " print(f\"[data_parser] Processed {len(urls)} samples in {elapsed:.2f}s (serial would be ~{len(urls)}.0s)\")\n",
- " return field_data\n",
- "\n",
- "\n",
- "# ---------------------------------------------------------------------------\n",
- "# Build the batch\n",
- "# ---------------------------------------------------------------------------\n",
- "batch_size = 32\n",
- "\n",
- "normal_data = torch.randn(batch_size, 2)\n",
- "\n",
- "# URL-like strings: all use the same dtype so TQ can pack them on get\n",
- "shapes = [(i % 4 + 1, i % 3 + 2) for i in range(batch_size)]\n",
- "urls = [f\"float32:{h}x{w}\" for h, w in shapes]\n",
- "\n",
- "parser_fields = TensorDict(\n",
- " {\n",
- " \"normal_data\": normal_data,\n",
- " \"data_to_be_parsed\": NonTensorStack(*urls),\n",
- " },\n",
- " batch_size=batch_size,\n",
+ "selected = tq.kv_batch_get(\n",
+ " data_parser_keys,\n",
+ " parser_partition,\n",
+ " select_fields=\"features\",\n",
+ " data_parser=select_rows,\n",
+ " parser_args=[[0, 2]],\n",
")\n",
+ "print(\"Rows 0 and 2 returned from storage:\\n\", selected[\"features\"][0])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 12. Multiple Partitions\n",
"\n",
- "data_parser_keys = [f\"data_parser_sample_{i}\" for i in range(batch_size)]\n",
- "\n",
- "put_start_time = time.perf_counter()\n",
- "meta = tq.kv_batch_put(\n",
- " keys=data_parser_keys,\n",
- " partition_id=\"train\",\n",
- " fields=parser_fields,\n",
- " data_parser=concurrent_batch_url_parser,\n",
- ")\n",
- "put_elapsed = time.perf_counter() - put_start_time\n",
- "print(f\"Put succeeded. Fields: {meta.fields}\")\n",
- "print(f\"Total kv_batch_put time: {put_elapsed:.2f}s (concurrency keeps it ~1s, not {batch_size}s)\\n\")"
- ],
+ "Partitions provide logical isolation — the same key name can exist in\n",
+ "different partitions without conflict."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "Put succeeded. Fields: ['data_to_be_parsed', 'normal_data']\n",
- "Total kv_batch_put time: 1.02s (concurrency keeps it ~1s, not 32s)\n",
- "\n"
+ "All partitions: ['train', 'parser', 'validation']\n"
]
}
],
- "execution_count": 19
+ "source": [
+ "tq.kv_put(\n",
+ " key=\"val_sample_0\",\n",
+ " partition_id=\"validation\",\n",
+ " fields={\"input_ids\": torch.tensor([99, 98, 97])},\n",
+ " tag={\"split\": \"val\"},\n",
+ ")\n",
+ "\n",
+ "all_info = tq.kv_list() # no partition_id → list everything\n",
+ "print(\"All partitions:\", list(all_info.keys()))"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
- "source": "When `kv_batch_put` returns, the user-defined data parser has also finished executing. We can then safely call `kv_batch_get` to retrieve the parsed data."
+ "source": [
+ "## 13. Share Fields Across Samples\n",
+ "\n",
+ "Store a field once in `tq.PUBLIC_PARTITION` and link several samples to it with\n",
+ "`kv_link`. This example shares a `features` tensor; each sample also gets its own\n",
+ "private `score` field. Public key names need an application-specific prefix\n",
+ "because applications on the same TQ share the public partition.\n",
+ "\n",
+ "Shared fields require SimpleStorage. Linked fields and referenced public rows\n",
+ "cannot be overwritten. See [Tutorial 07](07_shared_fields.py) for sharing and\n",
+ "reclaiming public data, and [Tutorial 03](03_metadata_concepts.py) for `BatchMeta` links."
+ ]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "result = tq.kv_batch_get(keys=data_parser_keys, partition_id=\"train\")\n\n# normal_data should be unchanged\n# normal_data is packed as a nested tensor on retrieval; compare per-sample.\nfor t1, t2 in zip(result[\"normal_data\"], normal_data, strict=True):\n torch.testing.assert_close(t1, t2)\nprint(\"[PASS] normal_data is unchanged.\")\n\n# data_to_be_parsed should now be tensors with the requested shapes\nexpected_shapes = [(i % 4 + 1, i % 3 + 2) for i in range(batch_size)]\nfor i, expected in enumerate(expected_shapes):\n tensor = result[\"data_to_be_parsed\"][i]\n assert tensor.dtype == torch.float32\n actual = tuple(tensor.shape)\n assert actual == expected, f\"Mismatch at {i}: {actual} != {expected}\"\nprint(f\"[PASS] All {batch_size} parsed tensors have correct dtype & shape.\")\n\n# Timing sanity check: serial would be ~batch_size seconds.\n# Because asyncio tasks run in parallel inside the parser, it should be ~1 s.\nassert put_elapsed < 2.0, f\"Expected concurrent execution (~1s), but took {put_elapsed:.2f}s.\"\nprint(f\"[PASS] Timing looks concurrent: {put_elapsed:.2f}s < 2.0s\")\n\n# wait for Ray log collect\ntime.sleep(2)"
- ],
"outputs": [
{
- "name": "stderr",
"output_type": "stream",
+ "name": "stdout",
"text": [
- ""
+ "[Step 1] Storing a shared field in the public partition...\n",
+ " Stored notebook/shared_0: features shape=(6, 4)\n",
+ "\n",
+ "[Step 2] Linking samples to the shared field...\n",
+ " shared_sample_0 reads features from notebook/shared_0\n",
+ " shared_sample_1 reads features from notebook/shared_0\n",
+ " shared_sample_2 reads features from notebook/shared_0\n",
+ " The links reference the public row; the tensor is stored once.\n"
]
- },
+ }
+ ],
+ "source": [
+ "print(\"[Step 1] Storing a shared field in the public partition...\")\n",
+ "public_key = \"notebook/shared_0\"\n",
+ "shared_features = torch.arange(24, dtype=torch.float32).reshape(6, 4)\n",
+ "tq.kv_put(public_key, tq.PUBLIC_PARTITION, fields={\"features\": shared_features})\n",
+ "print(f\" Stored {public_key}: features shape={tuple(shared_features.shape)}\")\n",
+ "\n",
+ "print(\"\\n[Step 2] Linking samples to the shared field...\")\n",
+ "shared_keys = [f\"shared_sample_{i}\" for i in range(3)]\n",
+ "tq.kv_link(shared_keys, \"shared\", public_key=public_key, fields=\"features\")\n",
+ "descriptions = tq.kv_describe(shared_keys, \"shared\")\n",
+ "for key in shared_keys:\n",
+ " print(f\" {key} reads features from {descriptions[key]['links']['features']}\")\n",
+ "print(\" The links reference the public row; the tensor is stored once.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Reads combine linked fields with each sample's private fields. Here the three\n",
+ "samples share `features` and keep separate `score` values."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
- "[PASS] normal_data is unchanged.\n",
- "[PASS] All 32 parsed tensors have correct dtype & shape.\n",
- "[PASS] Timing looks concurrent: 1.02s < 2.0s\n"
+ "[Step 3] Adding a private score to each sample...\n",
+ " Each sample now has shared features and its own score.\n",
+ "\n",
+ "[Step 4] Reading shared and private fields together...\n"
]
},
{
- "name": "stderr",
"output_type": "stream",
+ "name": "stdout",
"text": [
- ""
+ " Returned fields: ['features', 'score']\n",
+ " shared_sample_0: features shape=(6, 4), score=0.1\n",
+ " shared_sample_1: features shape=(6, 4), score=0.2\n",
+ " shared_sample_2: features shape=(6, 4), score=0.3\n"
]
}
],
- "execution_count": 20
+ "source": [
+ "print(\"[Step 3] Adding a private score to each sample...\")\n",
+ "scores = torch.tensor([0.1, 0.2, 0.3])\n",
+ "tq.kv_batch_put(\n",
+ " shared_keys,\n",
+ " \"shared\",\n",
+ " fields=TensorDict({\"score\": scores}, batch_size=len(shared_keys)),\n",
+ ")\n",
+ "print(\" Each sample now has shared features and its own score.\")\n",
+ "\n",
+ "print(\"\\n[Step 4] Reading shared and private fields together...\")\n",
+ "shared = tq.kv_batch_get(shared_keys, \"shared\")\n",
+ "print(\" Returned fields:\", list(shared.keys()))\n",
+ "for i, key in enumerate(shared_keys):\n",
+ " torch.testing.assert_close(shared[\"features\"][i], shared_features)\n",
+ " print(f\" {key}: features shape={tuple(shared['features'][i].shape)}, score={shared['score'][i].item():.1f}\")"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 12. Multiple Partitions\n\nPartitions provide logical isolation — the same key name can exist in\ndifferent partitions without conflict."
+ "## 14. Coordinate Operations with KV Locks\n",
+ "\n",
+ "Local locks coordinate threads and async tasks within one process; global locks\n",
+ "coordinate processes across nodes in the same TQ cluster. Callers must use the\n",
+ "same kind of lock for the same keys; KV operations do not check locks automatically.\n",
+ "\n",
+ "Use `async with` in a notebook. Global leases renew automatically; `lease.check()`\n",
+ "checks local lease state before writing and raises `LockLostError` if the lease\n",
+ "has expired or been lost. Global locks were enabled during initialization.\n",
+ "See [Tutorial 10](10_kv_locks.py) for supported operations and further details."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "tq.kv_put(\n",
- " key=\"val_sample_0\",\n",
- " partition_id=\"validation\",\n",
- " fields={\"input_ids\": torch.tensor([99, 98, 97])},\n",
- " tag={\"split\": \"val\"},\n",
- ")\n",
- "\n",
- "all_info = tq.kv_list() # no partition_id → list everything\n",
- "print(\"All partitions:\", list(all_info.keys()))"
- ],
"outputs": [
{
+ "output_type": "stream",
"name": "stdout",
+ "text": [
+ "Updating count under a local lock...\n",
+ "Local lock released; count=1\n"
+ ]
+ }
+ ],
+ "source": [
+ "tq.kv_put(\"local\", \"locks\", fields={\"count\": torch.tensor(0)})\n",
+ "tq.kv_put(\"global\", \"locks\", fields={\"count\": torch.tensor(0)})\n",
+ "\n",
+ "print(\"Updating count under a local lock...\")\n",
+ "async with tq.async_kv_local_lock(\"local\", \"locks\", timeout=10):\n",
+ " batch = await tq.async_kv_batch_get(\"local\", \"locks\")\n",
+ " count = int(batch[\"count\"][0]) + 1\n",
+ " await tq.async_kv_put(\"local\", \"locks\", fields={\"count\": torch.tensor(count)})\n",
+ "print(f\"Local lock released; count={count}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
"output_type": "stream",
+ "name": "stdout",
"text": [
- "All partitions: ['train', 'validation']\n"
+ "Updating count under a global lock...\n",
+ "Global lock released; count=1\n"
]
}
],
- "execution_count": 21
+ "source": [
+ "print(\"Updating count under a global lock...\")\n",
+ "async with tq.async_kv_global_lock(\"global\", \"locks\", timeout=10) as lease:\n",
+ " batch = await tq.async_kv_batch_get(\"global\", \"locks\")\n",
+ " count = int(batch[\"count\"][0]) + 1\n",
+ " lease.check()\n",
+ " await tq.async_kv_put(\"global\", \"locks\", fields={\"count\": torch.tensor(count)})\n",
+ "print(f\"Global lock released; count={count}\")"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## 13. Clean Up — `kv_clear` and `close`\n\nRemove specific keys with `kv_clear`, then shut down the system with `tq.close()`."
+ "## 15. Save and Restore Selected Samples\n",
+ "\n",
+ "Save the existing `sample_0`, clear it, then restore its fields and tags.\n",
+ "For unfinished responses, shared fields and recovery after a restart, see\n",
+ "[Tutorial 09](09_data_dump.py)."
]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Saving sample_0...\n",
+ "Cleared sample_0; restoring it from the dump...\n",
+ "Restored input_ids: [1, 2, 3, 4]\n",
+ "Restored tags: {'source': 'wikipedia', 'score': 0.95}\n"
+ ]
+ }
+ ],
"source": [
- "# Clear individual keys\n",
- "tq.kv_clear(keys=\"sample_0\", partition_id=\"train\")\n",
- "tq.kv_clear(keys=\"sample_1\", partition_id=\"train\")\n",
- "tq.kv_clear(keys=keys, partition_id=\"train\")\n",
- "tq.kv_clear(keys=data_parser_keys, partition_id=\"train\")\n",
- "tq.kv_clear(keys=\"val_sample_0\", partition_id=\"validation\")\n",
- "print(\"All keys cleared.\")\n",
+ "with TemporaryDirectory(prefix=\"tq-notebook-dump-\") as dump_dir:\n",
+ " dump_path = Path(dump_dir) / \"sample\"\n",
+ " print(\"Saving sample_0...\")\n",
+ " tq.dump_data_by_key(dump_path, [\"sample_0\"], \"train\")\n",
+ " tq.kv_clear(\"sample_0\", \"train\")\n",
+ " print(\"Cleared sample_0; restoring it from the dump...\")\n",
+ " restored_tags = tq.load_data_by_key(dump_path)\n",
+ "\n",
+ "restored = tq.kv_batch_get(\"sample_0\", \"train\")\n",
+ "print(\"Restored input_ids:\", restored[\"input_ids\"][0].tolist())\n",
+ "print(\"Restored tags:\", restored_tags[\"sample_0\"])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 16. Reclaim Data and Close\n",
"\n",
- "remaining = tq.kv_list()\n",
- "total_keys = sum(len(v) for v in remaining.values())\n",
- "print(f\"Remaining keys across all partitions: {total_keys}\")"
- ],
+ "Clearing sample keys removes their links. Their public data remains until\n",
+ "`kv_clear_unlinked_public` reclaims the rows with no remaining links.\n",
+ "\n",
+ "Each process closes its own client with `tq.close()`. The process that created TQ\n",
+ "also shuts down the shared services. `ray.shutdown()` closes the Ray connection."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
"outputs": [
{
+ "output_type": "stream",
"name": "stdout",
+ "text": [
+ "[Step 1] Clearing shared samples and reclaiming public data...\n",
+ " Cleared sample keys: ['shared_sample_0', 'shared_sample_1', 'shared_sample_2']\n",
+ " Public fields still available: {'notebook/shared_0': ['features']}\n",
+ " Reclaimed public keys: ['notebook/shared_0']\n",
+ "\n",
+ "[Step 2] Clearing the other example keys...\n"
+ ]
+ },
+ {
"output_type": "stream",
+ "name": "stdout",
"text": [
- "All keys cleared.\n",
- "Remaining keys across all partitions: 0\n"
+ " Remaining keys across all partitions: 0\n"
]
}
],
- "execution_count": 22
+ "source": [
+ "print(\"[Step 1] Clearing shared samples and reclaiming public data...\")\n",
+ "tq.kv_clear(shared_keys, \"shared\")\n",
+ "print(\" Cleared sample keys:\", shared_keys)\n",
+ "print(\" Public fields still available:\", tq.kv_exists(public_key, tq.PUBLIC_PARTITION))\n",
+ "freed = tq.kv_clear_unlinked_public()\n",
+ "print(\" Reclaimed public keys:\", freed)\n",
+ "\n",
+ "print(\"\\n[Step 2] Clearing the other example keys...\")\n",
+ "tq.kv_clear([\"local\", \"global\"], \"locks\")\n",
+ "tq.kv_clear(streaming_key, \"train\")\n",
+ "tq.kv_clear(keys=\"sample_0\", partition_id=\"train\")\n",
+ "tq.kv_clear(keys=\"sample_1\", partition_id=\"train\")\n",
+ "tq.kv_clear(keys=keys, partition_id=\"train\")\n",
+ "tq.kv_clear(keys=data_parser_keys, partition_id=parser_partition)\n",
+ "tq.kv_clear(keys=\"val_sample_0\", partition_id=\"validation\")\n",
+ "\n",
+ "remaining = tq.kv_list()\n",
+ "total_keys = sum(len(values) for values in remaining.values())\n",
+ "assert total_keys == 0\n",
+ "print(f\" Remaining keys across all partitions: {total_keys}\")"
+ ]
},
{
"cell_type": "code",
+ "execution_count": null,
"metadata": {},
- "source": [
- "tq.close()\n",
- "ray.shutdown()\n",
- "print(\"TransferQueue and Ray shut down.\")"
- ],
"outputs": [
{
- "name": "stdout",
"output_type": "stream",
+ "name": "stdout",
"text": [
"TransferQueue and Ray shut down.\n"
]
}
],
- "execution_count": 23
+ "source": [
+ "tq.close()\n",
+ "ray.shutdown()\n",
+ "print(\"TransferQueue and Ray shut down.\")"
+ ]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "---\n\n## Summary\n\n| Operation | Function | Notes |\n|---|---|---|\n| Init | `tq.init(config)` | Call once; subsequent processes auto-connect |\n| Put single | `tq.kv_put(key, partition_id, fields, tag)` | `fields` can be a plain dict |\n| Put batch | `tq.kv_batch_put(keys, partition_id, fields, tags)` | `fields` must be a `TensorDict` |\n| Put with parser | `tq.kv_batch_put(..., data_parser=fn)` | Only for **SimpleStorage**; receives dict, can use asyncio inside |\n| Get | `tq.kv_batch_get(keys, partition_id, select_fields=None)` | Returns a `TensorDict` |\n| List | `tq.kv_list(partition_id=None)` | Returns `{partition: {key: tag}}` |\n| Clear | `tq.kv_clear(keys, partition_id)` | Removes keys + data |\n| Close | `tq.close()` | Tears down controller & storage |\n\nFor **async** variants, use `async_kv_put`, `async_kv_batch_put`,\n`async_kv_batch_get`, `async_kv_list`, and `async_kv_clear`.\n\nFor low-level, metadata-based access, see `tq.get_client()` and the\n[official tutorials](https://github.com/Ascend/TransferQueue/tree/main/tutorial)."
+ "## API Reference for This Walkthrough\n",
+ "\n",
+ "| Operation | Function | Notes |\n",
+ "|---|---|---|\n",
+ "| Initialize | `tq.init(config)` | Attached processes discover the existing TQ |\n",
+ "| Write | `kv_put`, `kv_batch_put` | Add or replace supplied fields and optionally merge tags |\n",
+ "| Merge | `kv_update`, `kv_batch_update` | SimpleStorage; keep chunks hidden with `finalize=False` |\n",
+ "| Read | `kv_batch_get(..., select_fields=...)` | Selected fields must be ready in every requested key |\n",
+ "| Parse on put/get | `data_parser`, get-side `parser_args` | SimpleStorage; preserve sample order and leave stored values unchanged on get |\n",
+ "| List / inspect | `kv_list`, `kv_describe`, `kv_exists` | Inspect tags, stored fields, links and readiness |\n",
+ "| Share | `kv_link`, `client.link` | SimpleStorage; link to fields in `PUBLIC_PARTITION` |\n",
+ "| Clear fields/keys | `kv_clear(..., fields=...)` | Omit `fields` to remove keys as well |\n",
+ "| Reclaim public rows | `kv_clear_unlinked_public` | Clear referencing samples first |\n",
+ "| Lock | `(async_)kv_local_lock`, `(async_)kv_global_lock` | Advisory; global locks must be enabled by the owner |\n",
+ "| Selective recovery | `dump_data_by_key`, `load_data_by_key` | SimpleStorage; restore public sources before linked rows |\n",
+ "| Close | `tq.close()` | Every process releases its own resources; only the owner shuts down TQ |\n",
+ "\n",
+ "KV operations also have `async_` variants. Selective dump/load are synchronous.\n",
+ "For metadata-based access, see [Tutorial 03](03_metadata_concepts.py); for\n",
+ "sampler-driven consumption, see [Tutorial 06](06_streaming_dataloader.py)."
]
- },
- {
- "cell_type": "code",
- "metadata": {},
- "source": [],
- "outputs": [],
- "execution_count": 23
}
],
"metadata": {
@@ -1088,4 +1429,4 @@
},
"nbformat": 4,
"nbformat_minor": 4
-}
\ No newline at end of file
+}