[feat]support rollout data transfer with mooncake#591
[feat]support rollout data transfer with mooncake#591zxpdemonio wants to merge 5 commits intoradixark:mainfrom
Conversation
zxpdemonio
commented
Feb 12, 2026
zxpdemonio
commented
- Add Mooncake backend for rollout data transfer
- Use zero-copy put_from/get_into for better performance
- Auto cleanup data after retrieval to free storage
Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>
Signed-off-by: Cruz Zhao <CruzZhao@linux.alibaba.com>
Signed-off-by: Cruz Zhao <CruzZhao@linux.alibaba.com>
Signed-off-by: Cruz Zhao <CruzZhao@linux.alibaba.com>
Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Cruz Zhao <CruzZhao@linux.alibaba.com>
zxpdemonio
requested review from
fzyzcjy,
guapisolo,
maocheng23 and
yueming-yuan
as code owners
Summary of ChangesHello @zxpdemonio, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly upgrades the system's data transfer capabilities by integrating Mooncake as an alternative backend. This integration is designed to boost performance through zero-copy data operations and optimize storage usage with automatic cleanup. The changes provide a more flexible and efficient way to manage rollout data, particularly beneficial for large-scale distributed training environments. Highlights
Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
zxpdemonio
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a data transfer abstraction layer with a new Mooncake backend for improved performance using zero-copy transfers and automatic data cleanup. However, it introduces a significant security risk by using the pickle library for data serialization and deserialization, which is known to be insecure and can lead to remote code execution (RCE) if the data source is untrusted or compromised, especially given Mooncake is a distributed store. Additionally, I've identified a potential race condition in the cleanup thread management within the MooncakeDataTransfer class, an opportunity for efficiency improvements in the cleanup worker, and suggestions for adding type hints to enhance clarity.
| serialized_data = bytes(buffer[:bytes_read]) | ||
|
|
||
| # Deserialize and return | ||
| data = pickle.loads(serialized_data) |
There was a problem hiding this comment.
The use of pickle.loads on data retrieved from an external distributed store (MooncakeDistributedStore) is insecure. pickle can be exploited to execute arbitrary code if the input data is maliciously crafted. Since Mooncake is a distributed storage system, an attacker who can write to the store or intercept network traffic could achieve Remote Code Execution (RCE) on the training actors.
To remediate this, consider using a safer serialization format like json for non-tensor data and safetensors for tensor data. If you must use pickle, you should implement a message authentication code (MAC), such as HMAC with a shared secret, to verify the integrity and authenticity of the data before calling pickle.loads.
| if self._cleanup_thread is None or not self._cleanup_thread.is_alive(): | ||
| self._start_cleanup_thread() |
There was a problem hiding this comment.
There is a potential race condition here. The check for the cleanup thread's existence and liveness is not atomic with the call to _start_cleanup_thread(). If multiple threads call get() concurrently and find that the cleanup thread is not running, they might all attempt to start it. This can lead to a RuntimeError: threads can only be started once.
You should use a threading.Lock to protect this block of code to ensure the thread is started only once. You'll need to add self._cleanup_thread_lock = threading.Lock() in the __init__ method.
| if self._cleanup_thread is None or not self._cleanup_thread.is_alive(): | |
| self._start_cleanup_thread() | |
| with self._cleanup_thread_lock: | |
| if self._cleanup_thread is None or not self._cleanup_thread.is_alive(): | |
| self._start_cleanup_thread() |
|
|
||
|
|
||
| def get_rollout_data(args: Namespace, rollout_data_ref: Box, parallel_state: ParallelState) -> RolloutBatch: | ||
| def get_rollout_data(args: Namespace, rollout_data_ref, parallel_state: ParallelState, transfer_backend) -> RolloutBatch: |
There was a problem hiding this comment.
The type hints for rollout_data_ref and transfer_backend are missing. Adding them would improve code clarity and maintainability. Based on the new data_transfer module, transfer_backend should be of type DataTransferBackend. The type of rollout_data_ref now depends on the backend, so Any would be appropriate.
You will need to add the following imports:
from typing import Any
from miles.utils.data_transfer import DataTransferBackend| def get_rollout_data(args: Namespace, rollout_data_ref, parallel_state: ParallelState, transfer_backend) -> RolloutBatch: | |
| def get_rollout_data(args: Namespace, rollout_data_ref: "Any", parallel_state: ParallelState, transfer_backend: "DataTransferBackend") -> RolloutBatch: |
|
|
||
|
|
||
| def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): | ||
| def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size, transfer_backend): |
There was a problem hiding this comment.
This function is missing type hints for its arguments. Adding them would improve readability and help static analysis tools. Using forward references in strings for types that are not imported is a good practice.
You will need to add the following imports:
from argparse import Namespace
from miles.utils.data_transfer import DataTransferBackend| def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size, transfer_backend): | |
| def process_rollout_data(args: "Namespace", rollout_data_ref: list, dp_rank: int, dp_size: int, transfer_backend: "DataTransferBackend"): |
| def _cleanup_worker(self): | ||
| """Background worker that periodically deletes keys.""" | ||
| while not self._cleanup_stop_event.is_set(): | ||
| try: | ||
| # Collect keys ready for deletion | ||
| current_time = time.time() | ||
| keys_to_delete = [] | ||
| keys_not_ready = [] | ||
|
|
||
| # Process pending deletions | ||
| while len(keys_to_delete) < self.cleanup_batch_size: | ||
| try: | ||
| # Non-blocking get with timeout | ||
| key, deletion_time = self._pending_deletion.get(timeout=0.1) | ||
|
|
||
| if current_time >= deletion_time: | ||
| keys_to_delete.append(key) | ||
| else: | ||
| # Keep track of keys not ready yet | ||
| keys_not_ready.append((key, deletion_time)) | ||
| except queue.Empty: | ||
| break | ||
|
|
||
| # Put back keys that are not ready yet | ||
| for key, deletion_time in keys_not_ready: | ||
| self._pending_deletion.put((key, deletion_time)) | ||
|
|
||
| # Batch delete keys | ||
| if keys_to_delete: | ||
| deleted_count = 0 | ||
| for key in keys_to_delete: | ||
| try: | ||
| result = self.store.remove(key) | ||
| if result == 0: | ||
| deleted_count += 1 | ||
| else: | ||
| logger.warning(f"Failed to delete key {key}, error code: {result}") | ||
| except Exception as e: | ||
| logger.warning(f"Exception while deleting key {key}: {e}") | ||
|
|
||
| if deleted_count > 0: | ||
| logger.debug(f"Deleted {deleted_count} keys from Mooncake store") | ||
|
|
||
| # Sleep a bit to avoid busy waiting | ||
| # If we have keys not ready, sleep until the earliest one is ready | ||
| if keys_not_ready: | ||
| earliest_time = min(dt for _, dt in keys_not_ready) | ||
| sleep_time = min(0.5, max(0.1, earliest_time - current_time)) | ||
| time.sleep(sleep_time) | ||
| else: | ||
| time.sleep(0.5) | ||
|
|
||
| except Exception as e: | ||
| logger.error(f"Error in cleanup worker: {e}", exc_info=True) | ||
| time.sleep(1.0) # Sleep longer on error |
There was a problem hiding this comment.
The _cleanup_worker uses a standard queue.Queue with a polling mechanism (get(timeout=0.1)) to find keys ready for deletion. This approach is inefficient as it involves busy-waiting when the queue is empty and re-queueing items that are not yet ready, consuming unnecessary CPU cycles.
A queue.PriorityQueue would be a more efficient data structure for this task. By storing items with their deletion time as the priority, the worker thread can block on get() to wait for the next item due for deletion and then sleep until its deletion time. This avoids polling and improves efficiency.
2 participants
