test: add manual init test for mooncake transfer engine

test/manual/kv_transfer/test_mooncake_transfer_engine_init.py

@@ -0,0 +1,394 @@
+#!/usr/bin/env python3
+"""
+Test script for verifying init_shared_mooncake_transfer_engine functionality.
+Tests the code at model_runner.py lines 911-952.
+
+This test verifies:
+1. MooncakeTransferEngine initialization conditions
+2. Different server argument combinations that trigger mooncake TE
+3. Mooncake transfer engine initialization with hostname, gpu_id, and ib_device
+
+Usage:
+    # Run on 2 GPUs from project root
+    CUDA_VISIBLE_DEVICES=0,1 python test/manual/kv_transfer/test_mooncake_transfer_engine_init.py
+
+    # With strict mode (will fail if mooncake is not available)
+    CUDA_VISIBLE_DEVICES=0,1 python test/manual/kv_transfer/test_mooncake_transfer_engine_init.py --strict
+"""
+
+import argparse
+import multiprocessing
+import os
+import sys
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class ServerArgs:
+    """Mock ServerArgs for testing."""
+    disaggregation_mode: str = "null"
+    disaggregation_transfer_backend: str = "mooncake"
+    enable_hierarchical_cache: bool = False
+    hicache_storage_backend: str = "mooncake"
+    encoder_only: bool = False
+    language_only: bool = False
+    encoder_transfer_backend: str = "mooncake"
+    enable_elastic_expert_backup: bool = False
+    elastic_ep_backend: Optional[str] = None
+    disaggregation_ib_device: Optional[str] = None
+    mooncake_ib_device: Optional[str] = None
+
+
+def test_mooncake_te_condition(server_args: ServerArgs) -> bool:
+    """
+    Test the condition logic for using MooncakeTransferEngine.
+    This mirrors the logic in model_runner.py lines 917-935.
+    """
+    # Use os.environ directly to avoid import issues
+    hicache_mooncake_reuse_te = os.environ.get("SGLANG_HICACHE_MOONCAKE_REUSE_TE", "1").lower() in ("1", "true", "yes")
+
+    use_mooncake_te = (
+        (
+            server_args.disaggregation_mode != "null"
+            and server_args.disaggregation_transfer_backend == "mooncake"
+        )
+        or (
+            server_args.enable_hierarchical_cache
+            and server_args.hicache_storage_backend == "mooncake"
+            and hicache_mooncake_reuse_te
+        )
+        or (
+            server_args.encoder_only
+            and server_args.encoder_transfer_backend == "mooncake"
+        )
+        or (
+            server_args.language_only
+            and server_args.encoder_transfer_backend == "mooncake"
+        )
+        or (
+            server_args.enable_elastic_expert_backup
+            and server_args.elastic_ep_backend is not None
+        )
+    )
+
+    return use_mooncake_te
+
+
+def run_mooncake_init(rank: int, world_size: int, master_port: int, args: argparse.Namespace,
+                      server_args: ServerArgs) -> bool:
+    """Worker function for testing mooncake transfer engine initialization."""
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(master_port)
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["LOCAL_RANK"] = str(rank)
+
+    import torch
+    import torch.distributed as dist
+
+    try:
+        # Initialize distributed environment
+        print(f"[Rank {rank}] Initializing distributed environment...")
+        dist.init_process_group(
+            backend="nccl",
+            world_size=world_size,
+            rank=rank,
+            init_method=f"tcp://127.0.0.1:{master_port}",
+        )
+
+        # Set device
+        torch.cuda.set_device(rank)
+
+        # Sync to ensure all ranks are ready
+        dist.barrier()
+        print(f"[Rank {rank}] Distributed initialization complete.")
+
+        # Test the condition logic
+        use_mooncake_te = test_mooncake_te_condition(server_args)
+        print(f"[Rank {rank}] use_mooncake_te = {use_mooncake_te}")
+
+        if use_mooncake_te:
+            print(f"[Rank {rank}] Attempting to initialize MooncakeTransferEngine...")
+
+            # Import and initialize mooncake transfer engine
+            from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+                init_mooncake_transfer_engine,
+            )
+            from sglang.srt.utils import get_local_ip_auto
+
+            ib_device = (
+                server_args.disaggregation_ib_device
+                or server_args.mooncake_ib_device
+            )
+
+            print(f"[Rank {rank}] IB device: {ib_device}")
+
+            if args.strict:
+                # Strict mode: actually initialize mooncake
+                init_mooncake_transfer_engine(
+                    hostname=get_local_ip_auto(),
+                    gpu_id=rank,
+                    ib_device=ib_device,
+                )
+                print(f"[Rank {rank}] MooncakeTransferEngine initialized successfully!")
+            else:
+                # Mock mode: just verify the import works
+                print(f"[Rank {rank}] MooncakeTransferEngine import successful (mock mode)")
+
+            dist.barrier()
+
+        print(f"[Rank {rank}] Test completed successfully!")
+        return True
+
+    except ImportError as e:
+        print(f"[Rank {rank}] Mooncake not available (ImportError): {e}")
+        if args.strict:
+            raise
+        return True  # Non-strict mode: this is OK
+
+    except Exception as e:
+        print(f"[Rank {rank}] Test failed with error: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+    finally:
+        # Cleanup
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        print(f"[Rank {rank}] Process group destroyed.")
+
+
+def run_test(args: argparse.Namespace, server_args: ServerArgs) -> bool:
+    """Run the mooncake transfer engine test."""
+    # Set CUDA visible devices
+    cuda_devices = args.cuda_visible_devices.split(",")
+    world_size = len(cuda_devices)
+
+    if world_size < 2:
+        print("ERROR: This test requires at least 2 GPUs.")
+        print("Usage: CUDA_VISIBLE_DEVICES=0,1 python test/manual/kv_transfer/test_mooncake_transfer_engine_init.py")
+        sys.exit(1)
+
+    # Check GPU availability
+    import torch
+    if not torch.cuda.is_available():
+        print("ERROR: CUDA is not available")
+        sys.exit(1)
+
+    available_gpus = torch.cuda.device_count()
+    if world_size > available_gpus:
+        print(f"ERROR: Requested {world_size} GPUs but only {available_gpus} available")
+        sys.exit(1)
+
+    print(f"Testing with {world_size} GPUs: {cuda_devices}")
+    print(f"Strict mode: {args.strict}")
+    print()
+
+    # Print server args configuration
+    print("ServerArgs configuration:")
+    for key, value in vars(server_args).items():
+        print(f"  {key}: {value}")
+    print()
+
+    # Check if mooncake should be used
+    use_mooncake_te = test_mooncake_te_condition(server_args)
+    print(f"use_mooncake_te = {use_mooncake_te}")
+    print()
+
+    # Find a free port
+    import socket
+    with socket.socket() as s:
+        s.bind(("", 0))
+        master_port = s.getsockname()[1]
+
+    print(f"Using master port: {master_port}")
+
+    # Spawn worker processes
+    ctx = multiprocessing.get_context("spawn")
+    processes = []
+
+    for rank in range(world_size):
+        p = ctx.Process(
+            target=run_mooncake_init,
+            args=(rank, world_size, master_port, args, server_args),
+        )
+        p.start()
+        processes.append(p)
+
+    # Wait for all processes to complete
+    success = True
+    for i, p in enumerate(processes):
+        p.join(timeout=60)
+        if p.exitcode != 0:
+            print(f"Process {i} failed with exit code: {p.exitcode}")
+            success = False
+
+    # Cleanup any remaining processes
+    for p in processes:
+        if p.is_alive():
+            print(f"Process {p.pid} is still alive, terminating...")
+            p.terminate()
+            p.join(timeout=5)
+
+    return success
+
+
+def test_condition_logic():
+    """Test the condition logic for different server argument combinations."""
+    print("=" * 60)
+    print("Testing condition logic for use_mooncake_te")
+    print("=" * 60)
+    print()
+
+    test_cases = [
+        # (name, server_args, expected_result)
+        (
+            "PD disaggregation with mooncake",
+            ServerArgs(disaggregation_mode="prefill", disaggregation_transfer_backend="mooncake"),
+            True,
+        ),
+        (
+            "PD disaggregation without mooncake",
+            ServerArgs(disaggregation_mode="prefill", disaggregation_transfer_backend="other"),
+            False,
+        ),
+        (
+            "No disaggregation",
+            ServerArgs(),
+            False,
+        ),
+        (
+            "HiCache with mooncake (env=False)",
+            ServerArgs(enable_hierarchical_cache=True, hicache_storage_backend="mooncake"),
+            False,  # env var is False by default
+        ),
+        (
+            "Encoder only with mooncake",
+            ServerArgs(encoder_only=True, encoder_transfer_backend="mooncake"),
+            True,
+        ),
+        (
+            "Language only with mooncake",
+            ServerArgs(language_only=True, encoder_transfer_backend="mooncake"),
+            True,
+        ),
+        (
+            "Elastic expert backup with backend",
+            ServerArgs(enable_elastic_expert_backup=True, elastic_ep_backend="mooncake"),
+            True,
+        ),
+        (
+            "Elastic expert backup without backend",
+            ServerArgs(enable_elastic_expert_backup=True, elastic_ep_backend=None),
+            False,
+        ),
+    ]
+
+    # Set env var for HiCache test
+    os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = "1"
+    test_cases[3] = (
+        "HiCache with mooncake (env=True)",
+        ServerArgs(enable_hierarchical_cache=True, hicache_storage_backend="mooncake"),
+        True,
+    )
+
+    passed = 0
+    failed = 0
+
+    for name, server_args, expected in test_cases:
+        result = test_mooncake_te_condition(server_args)
+        status = "✓ PASS" if result == expected else "✗ FAIL"
+
+        if result == expected:
+            passed += 1
+        else:
+            failed += 1
+
+        print(f"{status}: {name}")
+        print(f"       Expected: {expected}, Got: {result}")
+        print()
+
+    print(f"Condition logic tests: {passed} passed, {failed} failed")
+    print()
+
+    return failed == 0
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Test init_shared_mooncake_transfer_engine functionality"
+    )
+    parser.add_argument(
+        "--cuda-visible-devices",
+        type=str,
+        default="0,1",
+        help="CUDA visible devices (default: 0,1)",
+    )
+    parser.add_argument(
+        "--strict",
+        action="store_true",
+        help="Strict mode: actually initialize mooncake (will fail if not available)",
+    )
+    parser.add_argument(
+        "--test-case",
+        type=str,
+        choices=["pd_disaggregation", "hicache", "encoder_only", "language_only", "elastic_ep"],
+        default="pd_disaggregation",
+        help="Test case to run",
+    )
+
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("Mooncake Transfer Engine Init Test")
+    print("=" * 60)
+    print()
+
+    # First run condition logic tests
+    condition_passed = test_condition_logic()
+
+    if not condition_passed:
+        print("Condition logic tests failed, skipping distributed test.")
+        sys.exit(1)
+
+    # Configure server args based on test case
+    server_args = ServerArgs()
+
+    if args.test_case == "pd_disaggregation":
+        server_args.disaggregation_mode = "prefill"
+        server_args.disaggregation_transfer_backend = "mooncake"
+    elif args.test_case == "hicache":
+        server_args.enable_hierarchical_cache = True
+        server_args.hicache_storage_backend = "mooncake"
+        os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = "1"
+    elif args.test_case == "encoder_only":
+        server_args.encoder_only = True
+        server_args.encoder_transfer_backend = "mooncake"
+    elif args.test_case == "language_only":
+        server_args.language_only = True
+        server_args.encoder_transfer_backend = "mooncake"
+    elif args.test_case == "elastic_ep":
+        server_args.enable_elastic_expert_backup = True
+        server_args.elastic_ep_backend = "mooncake"
+
+    start_time = time.time()
+    success = run_test(args, server_args)
+    elapsed_time = time.time() - start_time
+
+    print()
+    print("=" * 60)
+    if success:
+        print(f"TEST PASSED (elapsed: {elapsed_time:.2f}s)")
+    else:
+        print(f"TEST FAILED (elapsed: {elapsed_time:.2f}s)")
+    print("=" * 60)
+
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()