[4/N] (Elastic EP) Back up Expert Weights in DRAM (#17374)

Co-authored-by: UNIDY2002 <unidy2002@outlook.com>
2026-02-27 15:59:13 +08:00
parent eef44ec916
commit 43fade5f69
11 changed files with 553 additions and 2 deletions
--- a/test/manual/ep/test_mooncake_expert_backup.py
+++ b/test/manual/ep/test_mooncake_expert_backup.py
@@ -0,0 +1,142 @@
+import time
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    CustomTestCase,
+    popen_launch_pd_server,
+)
+
+ib_devices = get_rdma_devices_args()
+
+
+class TestBackup(CustomTestCase):
+    extra_args = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+        cls.base_port = 20000
+        cls.num_processes = 2
+        # TODO (stage 100): in the future, implement a specified multiprocess launcher
+        cls.processes = [
+            popen_launch_pd_server(
+                cls.model,
+                f"http://127.0.0.1:{cls.base_port + i}",
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--trust-remote-code",
+                    "--tp",
+                    "4",
+                    "--enable-dp-attention",
+                    "--dp",
+                    "4",
+                    "--elastic-ep-backend",
+                    "mooncake",
+                    "--mooncake-ib-device",
+                    ib_devices,
+                    "--moe-a2a-backend",
+                    "mooncake",
+                    "--deepep-mode",
+                    "low_latency",
+                    "--moe-dense-tp-size",
+                    "1",
+                    "--enable-dp-lm-head",
+                    "--enable-two-batch-overlap",
+                    "--disable-custom-all-reduce",
+                    "--enable-elastic-expert-backup",
+                    "--enable-eplb",
+                    "--eplb-rebalance-num-iterations",
+                    "50",
+                    "--chunked-prefill-size",
+                    "512",
+                    "--cuda-graph-max-bs",
+                    "128",
+                    "--max-running-requests",
+                    "512",
+                    "--mem-fraction-static",
+                    "0.5",
+                    "--dist-init-addr",
+                    "127.0.0.1:5000",
+                    "--nnodes",
+                    f"{cls.num_processes}",
+                    "--node-rank",
+                    f"{i}",
+                    "--base-gpu-id",
+                    f"{i * 2}",
+                ],
+            )
+            for i in range(cls.num_processes)
+        ]
+
+        server_ready = [False] * cls.num_processes
+        start_time = time.perf_counter()
+        with requests.Session() as session:
+            while (
+                time.perf_counter() - start_time < DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+                and not all(server_ready)
+            ):
+                for i, process in enumerate(cls.processes):
+                    return_code = process.poll()
+                    if return_code is not None:
+                        # Server failed to start (non-zero exit code) or crashed
+                        raise Exception(
+                            f"Server process exited with code {return_code}. "
+                            "Check server logs for errors."
+                        )
+
+                    try:
+                        headers = {
+                            "Content-Type": "application/json; charset=utf-8",
+                        }
+                        response = session.get(
+                            f"http://127.0.0.1:{cls.base_port + i}/health_generate",
+                            headers=headers,
+                        )
+                        if response.status_code == 200:
+                            server_ready[i] = True
+                    except requests.RequestException:
+                        pass
+
+                    return_code = process.poll()
+                    if return_code is not None:
+                        raise Exception(
+                            f"Server unexpectedly exits ({return_code=}). Usually there will be error logs describing the cause far above this line."
+                        )
+
+                    time.sleep(10)
+        if not all(server_ready):
+            for process in cls.processes:
+                kill_process_tree(process.pid)
+            raise TimeoutError("Server failed to start within the timeout period.")
+
+    @classmethod
+    def tearDownClass(cls):
+        for process in cls.processes:
+            kill_process_tree(process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=self.base_port,
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(metrics)
+
+        self.assertGreater(metrics["accuracy"], 0.60)
+
+
+if __name__ == "__main__":
+    unittest.main()