[Tutorial] Provide tutorials for TransferQueue usage#141
Conversation
Signed-off-by: 0oshowero0 <[email protected]>
Signed-off-by: 0oshowero0 <[email protected]>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughA comprehensive tutorial suite is added to TransferQueue, consisting of four new Python modules demonstrating core components, metadata concepts, controller functionality, and custom sampling strategies. The README is updated to note the official tutorial release on December 20, 2025. Each tutorial includes setup helpers, demonstration functions, and end-to-end workflow examples with Ray, TensorDict, and OmegaConf. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Poem
Note 🎁 Summarized by CodeRabbit FreeYour organization is on the Free plan. CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please upgrade your subscription to CodeRabbit Pro by visiting https://app.coderabbit.ai/login. Comment |
There was a problem hiding this comment.
Pull request overview
This PR adds two comprehensive tutorials for the TransferQueue library, demonstrating core components and metadata system usage. The tutorials provide hands-on examples with detailed explanations to help users understand how to use TransferQueue for distributed data transfer workflows.
Key Changes:
- New tutorial introducing TransferQueue's three core components (Controller, Storage Backend, Client)
- New tutorial demonstrating the metadata system (FieldMeta, SampleMeta, BatchMeta)
- Both tutorials include practical examples with actual TransferQueue API usage
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| tutorial/01_core_components.py | Introduces TransferQueue architecture and demonstrates basic put/get/clear workflow with SimpleStorageUnit backend |
| tutorial/02_metadata_concepts.py | Demonstrates metadata classes (FieldMeta, SampleMeta, BatchMeta) with operations like chunk, concat, union, and select |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| parent_dir = Path(__file__).resolve().parent.parent | ||
| sys.path.append(str(parent_dir)) | ||
|
|
||
| from transfer_queue import ( # noqa: E402 # noqa: E402 |
There was a problem hiding this comment.
Duplicate noqa comment. The second "noqa: E402" comment is redundant.
| from transfer_queue import ( # noqa: E402 # noqa: E402 | |
| from transfer_queue import ( # noqa: E402 |
| ray.init() | ||
|
|
||
| # Setup TransferQueue | ||
| config = OmegaConf.create( | ||
| { | ||
| "global_batch_size": 8, | ||
| "num_data_storage_units": 2, | ||
| } | ||
| ) | ||
|
|
||
| storage_units = {} | ||
| for i in range(2): | ||
| storage_units[i] = SimpleStorageUnit.remote(storage_unit_size=100) | ||
|
|
||
| controller = TransferQueueController.remote() | ||
| controller_info = process_zmq_server_info(controller) | ||
| storage_unit_infos = process_zmq_server_info(storage_units) | ||
|
|
||
| client = TransferQueueClient( | ||
| client_id="TutorialClient", | ||
| controller_info=controller_info, | ||
| ) | ||
|
|
||
| tq_config = OmegaConf.create({}, flags={"allow_objects": True}) | ||
| tq_config.controller_info = controller_info | ||
| tq_config.storage_unit_infos = storage_unit_infos | ||
| config = OmegaConf.merge(tq_config, config) | ||
|
|
||
| client.initialize_storage_manager(manager_type="AsyncSimpleStorageManager", config=config) | ||
|
|
||
| print("[Step 1] Putting data into TransferQueue...") | ||
| input_ids = torch.randint(0, 1000, (8, 512)) | ||
| attention_mask = torch.ones(8, 512) | ||
|
|
||
| data_batch = TensorDict( | ||
| { | ||
| "input_ids": input_ids, | ||
| "attention_mask": attention_mask, | ||
| }, | ||
| batch_size=8, | ||
| ) | ||
|
|
||
| partition_id = "demo_partition" | ||
| batch_meta = client.put(data=data_batch, partition_id=partition_id) | ||
| print(f"✓ Put {data_batch.batch_size[0]} samples into partition '{partition_id}', got BatchMeta back {batch_meta}.") | ||
|
|
||
| print("[Step 2] Try to get metadata from TransferQueue from other places...") | ||
| batch_meta = client.get_meta( | ||
| data_fields=["input_ids", "attention_mask"], | ||
| batch_size=8, | ||
| partition_id=partition_id, | ||
| task_name="demo_task", # TransferQueueController prevents same task_name from getting data repeatedly | ||
| ) | ||
| print("✓ Got BatchMeta from TransferQueue:") | ||
| print(f" Number of samples: {len(batch_meta)}") | ||
| print(f" Global indexes: {batch_meta.global_indexes}") | ||
| print(f" Field names: {batch_meta.field_names}") | ||
| print(f" Partition ID: {batch_meta.samples[0].partition_id}") | ||
| print(f" Sample structure: {batch_meta.samples[0]}") | ||
|
|
||
| print("[Step 3] Retrieve samples with specific fields..") | ||
| selected_meta = batch_meta.select_fields(["input_ids"]) | ||
| print("✓ Selected 'input_ids' field only:") | ||
| print(f" New field names: {selected_meta.field_names}") | ||
| print(f" Samples still have same global indexes: {selected_meta.global_indexes}") | ||
| retrieved_data = client.get_data(batch_meta) | ||
| print(f" Retrieved data keys: {list(retrieved_data.keys())}") | ||
|
|
||
| print("[Step 4] Select specific samples from the retrieved BatchMeta...") | ||
| partial_meta = batch_meta.select_samples([0, 2, 4, 6]) | ||
| print("✓ Selected samples at indices [0, 2, 4, 6]:") | ||
| print(f" New global indexes: {partial_meta.global_indexes}") | ||
| print(f" Number of samples: {len(partial_meta)}") | ||
| retrieved_data = client.get_data(partial_meta) | ||
| print(f" Retrieved data samples: {retrieved_data}, all the data samples: {data_batch}") | ||
|
|
||
| print("[Step 5] Demonstrate chunk operation...") | ||
| chunks = batch_meta.chunk(2) | ||
| print(f"✓ Chunked into {len(chunks)} parts:") | ||
| for i, chunk in enumerate(chunks): | ||
| print(f" Chunk {i}: {len(chunk)} samples, indexes={chunk.global_indexes}") | ||
| chunk_data = client.get_data(chunk) | ||
| print(f" Chunk {i}: Retrieved chunk data: {chunk_data}") | ||
|
|
||
| # Cleanup | ||
| client.clear(partition_id=partition_id) | ||
| client.close() | ||
| ray.shutdown() |
There was a problem hiding this comment.
Potential resource leak if an exception occurs between ray.init() and ray.shutdown(). The function should use try-finally or a context manager to ensure Ray is properly shutdown even if an exception occurs. Consider wrapping the main logic in a try-finally block with ray.shutdown() in the finally clause.
Signed-off-by: 0oshowero0 <[email protected]>
Signed-off-by: 0oshowero0 <[email protected]>
Signed-off-by: 0oshowero0 <[email protected]>
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 14 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: 0oshowero0 <[email protected]>
Signed-off-by: 0oshowero0 <[email protected]>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 11 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: 0oshowero0 <[email protected]>
Signed-off-by: 0oshowero0 <[email protected]>
Signed-off-by: 0oshowero0 <[email protected]>
This pull request introduces comprehensive tutorial scripts for TransferQueue, significantly enhancing the documentation and onboarding experience for new users. It adds detailed, runnable Python scripts that demonstrate the core components, controller behavior, and key features of TransferQueue. Additionally, a new update is added to the
README.mdto announce the release of the official tutorial.Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.