Add disaggregation feature #893
Conversation
Summary of ChangesHello @zhtshr, 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 introduces a disaggregation architecture to the LightX2V project, transforming the video generation pipeline into a distributed system. By decoupling the encoder and transformer components into independent services, it significantly enhances resource flexibility, scalability, and memory efficiency. This change allows for more versatile deployments, from multi-GPU single machines to cross-machine distributed environments, without altering existing APIs or workflows. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. 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
|
![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 significant disaggregation feature, enabling distributed deployment of the video generation pipeline. The overall architecture is well-designed, separating concerns into services, connection management, and utilities. My review focuses on improving code quality, portability, and fixing potential bugs. Key areas of feedback include correcting ZeroMQ context management, fixing a bug in an example client, improving the portability of example scripts by removing hardcoded paths, and enhancing code style and consistency. Overall, this is a solid contribution, and the suggested changes will help make the new code more robust and maintainable.
| if ret_value != 0: | ||
| print("Mooncake memory registration failed.") | ||
| raise RuntimeError("Mooncake memory registration failed.") |
There was a problem hiding this comment.
The check if ret_value != 0: will raise a NameError if PROTOCOL is not 'rdma', because ret_value would not have been defined. This check and the following lines should be inside the if PROTOCOL == 'rdma': block.
| if ret_value != 0: | |
| print("Mooncake memory registration failed.") | |
| raise RuntimeError("Mooncake memory registration failed.") | |
| if ret_value != 0: | |
| print("Mooncake memory registration failed.") | |
| raise RuntimeError("Mooncake memory registration failed.") |
| buf = torch.empty((nbytes,), dtype=torch.uint8, #device=torch.device(f"cuda:{self.receiver_engine_rank}") | ||
| ) |
There was a problem hiding this comment.
The device argument for torch.empty is commented out, which means the buffer will be allocated on the CPU. If these buffers are intended for RDMA over InfiniBand with GPU Direct, they typically need to be allocated on the GPU. The torch.cuda.set_device call on line 94 is also commented out, which might be related. Please verify if this is the intended behavior.
buf = torch.empty((nbytes,), dtype=torch.uint8, device=torch.device(f"cuda:{self.receiver_engine_rank}"))| buf = torch.empty( | ||
| (nbytes,), | ||
| dtype=torch.uint8, | ||
| # device=torch.device(f"cuda:{self.sender_engine_rank}"), | ||
| ) |
There was a problem hiding this comment.
The device argument for torch.empty is commented out, which means the buffer will be allocated on the CPU. If these buffers are intended for RDMA over InfiniBand with GPU Direct, they typically need to be allocated on the GPU. The torch.cuda.set_device call on line 135 is also commented out, which might be related. Please verify if this is the intended behavior.
| buf = torch.empty( | |
| (nbytes,), | |
| dtype=torch.uint8, | |
| # device=torch.device(f"cuda:{self.sender_engine_rank}"), | |
| ) | |
| buf = torch.empty( | |
| (nbytes,), | |
| dtype=torch.uint8, | |
| device=torch.device(f"cuda:{self.sender_engine_rank}"), | |
| ) |
| @@ -0,0 +1,89 @@ | |||
| import logging | |||
There was a problem hiding this comment.
The script imports and configures the standard logging module (on line 10) but then exclusively uses loguru.logger for logging. To avoid confusion and unused code, it's better to remove the import and configuration for the standard logging module.
|
|
||
| def main(): | ||
| # 1. Configuration | ||
| model_path = "/root/zht/LightX2V/models/Wan-AI/Wan2.1-T2V-1.3B" |
There was a problem hiding this comment.
Hardcoded absolute paths like this one (and on line 18) make the script difficult to run for other users on different systems. It's better to use relative paths, command-line arguments, or environment variables to specify these paths.
| def failure_exception(self): | ||
| raise Exception("Fake DataSender Exception") |
There was a problem hiding this comment.
Raising a generic Exception with a message like 'Fake DataSender Exception' is not ideal for production code. It's better to use a more specific exception type like RuntimeError and provide a more descriptive message. Consider defining a custom exception class for this module for even better error handling.
| def failure_exception(self): | |
| raise Exception("Fake DataSender Exception") | |
| def failure_exception(self): | |
| raise RuntimeError("DataSender failed due to an internal error.") |
| buffer_ptrs = [buf.data_ptr() for buf in self._rdma_buffers] | ||
| self.data_sender.send(buffer_ptrs) | ||
|
|
||
| import time |
There was a problem hiding this comment.
Imports should be at the top of the file, as per PEP 8 guidelines. Moving import time to the top improves readability and consistency.
| seed_all(self.config["seed"]) | ||
|
|
||
| data_bootstrap_addr = self.config.get("data_bootstrap_addr", "127.0.0.1") | ||
| data_bootstrap_room = self.config.get("data_bootstrap_room", 0) |
There was a problem hiding this comment.
For consistency with how sender_engine_rank and receiver_engine_rank are handled, it's better to cast data_bootstrap_room to an integer here. It's currently cast later in the DataReceiver constructor.
| data_bootstrap_room = self.config.get("data_bootstrap_room", 0) | |
| data_bootstrap_room = int(self.config.get("data_bootstrap_room", 0)) |
|
|
||
| @cache | ||
| def _connect(self, endpoint: str): | ||
| socket = zmq.Context().socket(zmq.PUSH) |
There was a problem hiding this comment.
To follow the single ZMQ context per process best practice, this method should use the context created in __init__ instead of creating a new one.
| socket = zmq.Context().socket(zmq.PUSH) | |
| socket = self.context.socket(zmq.PUSH) |
| return hashlib.sha256(data).hexdigest() | ||
|
|
||
| # Poll for data from EncoderService | ||
| import time |
There was a problem hiding this comment.
Imports should be at the top of the file, as per PEP 8 guidelines. Moving import time to the top improves readability and consistency.
2 participants
Enable Disaggregation Feature
Summary
This PR introduces a disaggregation architecture to LightX2V, enabling distributed deployment of the video generation pipeline across multiple devices or machines.
What's New
Core Functionality
New Components
lightx2v/disagg/: Complete disaggregation packageconn.py: Data connection and managementservices/encoder.py: Encoder service implementationservices/transformer.py: Transformer service implementationexamples/: Usage examples for WAN I2V and T2V modelsKey Benefits
Usage Example
See
lightx2v/disagg/examples/for complete working examples.Backward Compatibility
✅ This is an optional feature that doesn't affect existing functionality:
Testing
Files Changed
lightx2v/disagg/package with all disaggregation modulesFuture Enhancements
Type: Feature
Breaking Changes: None
Documentation: Included in
lightx2v/disagg/examples/