Support managing multiple data partitions for Train/Val/Test in controller

[ji-huazhong]

Summary by CodeRabbit

• Documentation

  • Extensive transfer queue documentation updates with new sections on customization, storage backends, and integration patterns; includes code examples and architecture guides.

• Refactor

  • Streamlined transfer queue client creation with simplified parameters and improved initialization flow.
  • Consolidated multi-controller data system management into a unified single-controller approach.

• Chores

  • Updated transfer queue dependency version.

[coderabbitai]

Caution

Review failed

The pull request is closed.

Walkthrough

This PR refactors the TransferQueue integration to consolidate multi-role train/val clients into a unified single-controller approach. Method signatures across agent_loop, worker, and utils modules are updated to accept controller_info and config instead of separate controller_infos, storage_infos, and role parameters. Data operations shift to partition_id-based identifiers. Documentation receives extensive updates covering interfaces, customization patterns, and integration examples.

Changes

Cohort / File(s)	Summary
Documentation docs/data/transfer_queue.md	Extensive updates: metadata refresh, terminology refinements (e.g., "storage manager"), new structural sections (Updates timeline, expanded Components with Control/Data Plane details, Customization guides), code examples, Verl integration showcase, new interfaces (BaseSampler, TransferQueueStorageManager), and directory structure diagrams.
API Refactoring: TransferQueue Client Initialization recipe/transfer_queue/agent_loop.py, verl/single_controller/base/worker.py, verl/utils/transferqueue_utils.py	Method signature consolidation: create_transferqueue_client(controller_infos, storage_infos, role) → create_transferqueue_client(controller_info, config). Removal of VAL-specific client paths and role-based differentiation. Client ID generation simplified from role-prefixed to generic worker_* format.
Ray Trainer Data System Refactoring recipe/transfer_queue/ray_trainer.py	Unified data system initialization via new _initialize_data_system() method; single TransferQueueController instead of per-role controllers. RPC calls (async_put, async_get_meta, async_get_data, async_clear) switch from global_step to partition_id-based identifiers (e.g., "train_X", "val_X"). Replaced TransferQueueStorageSimpleUnit with SimpleStorageUnit. Added keep_minibatch parameter to _balance_batch() signature. OmegaConf flag handling for ZMQ server info preservation.
Dependency Update requirements_transferqueue.txt	Replaced git+https VCS dependency with pinned PyPI pre-release: transferqueue==0.1.1.dev2.

Sequence Diagram(s)

sequenceDiagram
    participant Trainer as RayPPOTrainer
    participant Controller as TransferQueueController
    participant Client as AsyncTransferQueueClient
    participant Storage as AsyncSimpleStorageManager

    Trainer->>Trainer: _initialize_data_system()
    Trainer->>Controller: Initialize (unified, single controller)
    Trainer->>Client: create_transferqueue_client(controller_info, config)
    Client->>Storage: initialize_storage_manager(config)
    
    Trainer->>Client: async_put(partition_id="train_X", data)
    Client->>Storage: Store at partition_id
    
    Trainer->>Client: async_get_meta(partition_id="train_X")
    Client->>Storage: Retrieve metadata
    
    Trainer->>Client: async_get_data(partition_id="train_X")
    Client->>Storage: Retrieve data

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Areas requiring extra attention:

• ray_trainer.py: Dense refactoring affecting data system initialization flow, RPC call migration from global_step to partition_id identifiers, and consolidated controller setup; validate logic correctness for train/val data routing
• transferqueue_utils.py: Removal of VAL client branching logic; ensure all conditional paths previously selecting per-role clients now correctly route through unified client
• API contract consistency: Verify all callers of create_transferqueue_client pass correct controller_info and config structure across agent_loop and worker modules

Possibly related PRs

• [data] feat: TransferQueue - Support managing multiple data partitions for Train/Val/Test in controller #43: Consolidates per-role train/val clients into a single controller/config API, implementing the same signature changes to create_transferqueue_client, removal of VAL client paths, and unified data system initialization.

Suggested reviewers

• baymax591
• 0oshowero0

Poem

🐰 A queue once split in train and val,
Now dances as a single cohort, whole and tall!
With partition_id tags and config in hand,
We unified the client—oh, what a grand land!
Storage sings in harmony, one controller reigns,
Transfer Queue flows sweetly through refactored veins. 🎉

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch main_tq

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5c466c5 and 2ae18a6.

📒 Files selected for processing (6)

• docs/data/transfer_queue.md (2 hunks)
• recipe/transfer_queue/agent_loop.py (1 hunks)
• recipe/transfer_queue/ray_trainer.py (17 hunks)
• requirements_transferqueue.txt (1 hunks)
• verl/single_controller/base/worker.py (1 hunks)
• verl/utils/transferqueue_utils.py (3 hunks)

Tip

📝 Customizable high-level summaries are now available in beta!

You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.

• Provide your own instructions using the high_level_summary_instructions setting.
• Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
• Use high_level_summary_in_walkthrough to move the summary from the description to the walkthrough section.

Example instruction:

"Divide the high-level summary into five sections:

1. 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
2. 📓 References — List relevant issues, discussions, documentation, or related PRs.
3. 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
4. 📊 Contributor Summary — Include a Markdown table showing contributions:
   | Contributor | Lines Added | Lines Removed | Files Changed |
5. ✔️ Additional Notes — Add any extra reviewer context.
   Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."

Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}