decentralized queue scheduling

configs/disagg/wan_t2v_disagg_controller.json

@@ -0,0 +1,26 @@
+{
+    "infer_steps": 50,
+    "target_video_length": 81,
+    "text_len": 512,
+    "target_height": 480,
+    "target_width": 832,
+    "self_attn_1_type": "sage_attn2",
+    "cross_attn_1_type": "sage_attn2",
+    "cross_attn_2_type": "sage_attn2",
+    "sample_guide_scale": 5,
+    "sample_shift": 5,
+    "enable_cfg": true,
+    "cpu_offload": false,
+    "fps": 16,
+    "disagg_mode": "controller",
+    "disagg_config": {
+        "bootstrap_addr": "127.0.0.1",
+        "bootstrap_room": 0,
+        "encoder_engine_rank": 0,
+        "transformer_engine_rank": 1,
+        "decoder_engine_rank": 2,
+        "protocol": "rdma",
+        "local_hostname": "localhost",
+        "metadata_server": "P2PHANDSHAKE"
+    }
+}

configs/disagg/wan_t2v_disagg_decoder.json

@@ -0,0 +1,26 @@
+{
+    "infer_steps": 50,
+    "target_video_length": 81,
+    "text_len": 512,
+    "target_height": 480,
+    "target_width": 832,
+    "self_attn_1_type": "sage_attn2",
+    "cross_attn_1_type": "sage_attn2",
+    "cross_attn_2_type": "sage_attn2",
+    "sample_guide_scale": 5,
+    "sample_shift": 5,
+    "enable_cfg": true,
+    "cpu_offload": false,
+    "fps": 16,
+    "disagg_mode": "decoder",
+    "disagg_config": {
+        "bootstrap_addr": "127.0.0.1",
+        "bootstrap_room": 0,
+        "encoder_engine_rank": 0,
+        "transformer_engine_rank": 1,
+        "decoder_engine_rank": 2,
+        "protocol": "rdma",
+        "local_hostname": "localhost",
+        "metadata_server": "P2PHANDSHAKE"
+    }
+}

configs/disagg/wan_t2v_disagg_encoder.json

@@ -0,0 +1,26 @@
+{
+    "infer_steps": 50,
+    "target_video_length": 81,
+    "text_len": 512,
+    "target_height": 480,
+    "target_width": 832,
+    "self_attn_1_type": "sage_attn2",
+    "cross_attn_1_type": "sage_attn2",
+    "cross_attn_2_type": "sage_attn2",
+    "sample_guide_scale": 5,
+    "sample_shift": 5,
+    "enable_cfg": true,
+    "cpu_offload": false,
+    "fps": 16,
+    "disagg_mode": "encoder",
+    "disagg_config": {
+        "bootstrap_addr": "127.0.0.1",
+        "bootstrap_room": 0,
+        "encoder_engine_rank": 0,
+        "transformer_engine_rank": 1,
+        "decoder_engine_rank": 2,
+        "protocol": "rdma",
+        "local_hostname": "localhost",
+        "metadata_server": "P2PHANDSHAKE"
+    }
+}

configs/disagg/wan_t2v_disagg_transformer.json

@@ -0,0 +1,26 @@
+{
+    "infer_steps": 50,
+    "target_video_length": 81,
+    "text_len": 512,
+    "target_height": 480,
+    "target_width": 832,
+    "self_attn_1_type": "sage_attn2",
+    "cross_attn_1_type": "sage_attn2",
+    "cross_attn_2_type": "sage_attn2",
+    "sample_guide_scale": 5,
+    "sample_shift": 5,
+    "enable_cfg": true,
+    "cpu_offload": false,
+    "fps": 16,
+    "disagg_mode": "transformer",
+    "disagg_config": {
+        "bootstrap_addr": "127.0.0.1",
+        "bootstrap_room": 0,
+        "encoder_engine_rank": 0,
+        "transformer_engine_rank": 1,
+        "decoder_engine_rank": 2,
+        "protocol": "rdma",
+        "local_hostname": "localhost",
+        "metadata_server": "P2PHANDSHAKE"
+    }
+}

lightx2v/disagg/conn.py

@@ -75,6 +75,7 @@ class DataPoll:
 
 RequestPoolType = Dict[int, List[int]]
 WaitingPoolType = Dict[int, Tuple[str, list[int]]]
+MONITOR_POLLING_PORT = 7788
 REQUEST_POLLING_PORT = 12788
 DATASENDER_POLLING_PORT = 17788
 DATARECEIVER_POLLING_PORT = 27788

lightx2v/disagg/examples/run_service.py

@@ -0,0 +1,129 @@
+import argparse
+import json
+import logging
+
+from loguru import logger
+
+from lightx2v.disagg.services.controller import ControllerService
+from lightx2v.disagg.services.decoder import DecoderService
+from lightx2v.disagg.services.encoder import EncoderService
+from lightx2v.disagg.services.transformer import TransformerService
+from lightx2v.disagg.utils import set_config
+from lightx2v.utils.utils import seed_all
+
+logging.basicConfig(level=logging.INFO)
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run a disaggregated LightX2V service process")
+    parser.add_argument("--model_cls", type=str, default="wan2.1")
+    parser.add_argument("--task", type=str, default="t2v")
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--config_json", type=str, required=True)
+
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        type=str,
+        default=(
+            "镜头晃动，色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，"
+            "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，"
+            "畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+        ),
+    )
+    parser.add_argument(
+        "--save_result_path",
+        type=str,
+        default="/root/zht/LightX2V/save_results/test_disagg.mp4",
+    )
+
+    parser.add_argument(
+        "--service",
+        type=str,
+        choices=["encoder", "transformer", "decoder", "controller", "auto"],
+        default="auto",
+        help="Service role. auto = infer from config_json.disagg_mode",
+    )
+    return parser
+
+
+def _normalize_disagg_config(config: dict) -> dict:
+    disagg_cfg = config.get("disagg_config")
+    if isinstance(disagg_cfg, dict):
+        mapping = {
+            "bootstrap_addr": "data_bootstrap_addr",
+            "bootstrap_room": "data_bootstrap_room",
+            "encoder_engine_rank": "encoder_engine_rank",
+            "transformer_engine_rank": "transformer_engine_rank",
+            "decoder_engine_rank": "decoder_engine_rank",
+            "protocol": "protocol",
+            "local_hostname": "local_hostname",
+            "metadata_server": "metadata_server",
+        }
+        for src_key, dst_key in mapping.items():
+            if src_key in disagg_cfg:
+                config[dst_key] = disagg_cfg[src_key]
+    return config
+
+
+def _load_raw_json(path: str) -> dict:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _resolve_service_mode(args: argparse.Namespace, raw_cfg: dict) -> str:
+    if args.service != "auto":
+        return args.service
+    mode = raw_cfg.get("disagg_mode")
+    if mode in {"encoder", "transformer", "decoder", "controller"}:
+        return mode
+    raise ValueError("Cannot resolve service mode: use --service or set disagg_mode in config_json")
+
+
+def _build_runtime_config(args: argparse.Namespace) -> tuple[dict, dict]:
+    raw_cfg = _load_raw_json(args.config_json)
+
+    config = set_config(
+        model_path=args.model_path,
+        task=args.task,
+        model_cls=args.model_cls,
+        config_path=args.config_json,
+    )
+
+    config = _normalize_disagg_config(config)
+    raw_cfg = _normalize_disagg_config(raw_cfg)
+
+    config["seed"] = args.seed
+    config["prompt"] = args.prompt
+    config["negative_prompt"] = args.negative_prompt
+    config["save_path"] = args.save_result_path
+    return config, raw_cfg
+
+
+def main():
+    args = _build_parser().parse_args()
+    config, raw_cfg = _build_runtime_config(args)
+    service_mode = _resolve_service_mode(args, raw_cfg)
+
+    seed_all(args.seed)
+    logger.info("Starting disagg service mode={}", service_mode)
+
+    if service_mode == "encoder":
+        EncoderService(config).run()
+    elif service_mode == "transformer":
+        TransformerService(config).run()
+    elif service_mode == "decoder":
+        DecoderService(config).run()
+    elif service_mode == "controller":
+        ControllerService().run(config)
+    else:
+        raise ValueError(f"Unsupported service mode: {service_mode}")
+
+
+if __name__ == "__main__":
+    main()

lightx2v/disagg/examples/wan_i2v_service.py

@@ -84,21 +84,21 @@ def main():
     # 3. Define service threads
     def run_encoder():
         logger.info("Initializing Encoder Service...")
-        encoder_service = EncoderService()
+        encoder_service = EncoderService(config)
         logger.info("Running Encoder Service...")
         encoder_service.run(stop_event=encoder_stop_event)
         logger.info("Encoder Service completed.")
 
     def run_transformer():
         logger.info("Initializing Transformer Service...")
-        transformer_service = TransformerService()
+        transformer_service = TransformerService(config)
         logger.info("Running Transformer Service...")
         transformer_service.run(stop_event=transformer_stop_event)
         logger.info("Transformer Service completed.")
 
     def run_decoder():
         logger.info("Initializing Decoder Service...")
-        decoder_service = DecoderService()
+        decoder_service = DecoderService(config)
         logger.info("Running Decoder Service...")
         decoder_service.run(stop_event=decoder_stop_event)
         logger.info("Video generation completed.")

lightx2v/disagg/examples/wan_t2v_service.py

@@ -67,21 +67,21 @@ def main():
     # 2. Define service threads
     def run_encoder():
         logger.info("Initializing Encoder Service...")
-        encoder_service = EncoderService()
+        encoder_service = EncoderService(config)
         logger.info("Running Encoder Service...")
         encoder_service.run(stop_event=encoder_stop_event)
         logger.info("Encoder Service completed.")
 
     def run_transformer():
         logger.info("Initializing Transformer Service...")
-        transformer_service = TransformerService()
+        transformer_service = TransformerService(config)
         logger.info("Running Transformer Service...")
         transformer_service.run(stop_event=transformer_stop_event)
         logger.info("Transformer Service completed.")
 
     def run_decoder():
         logger.info("Initializing Decoder Service...")
-        decoder_service = DecoderService()
+        decoder_service = DecoderService(config)
         logger.info("Running Decoder Service...")
         decoder_service.run(stop_event=decoder_stop_event)
         logger.info("Video generation completed.")