[PD] Refactor Disagg Conn and Fix Hang with total_request/total_tokens Balancing (#21299)

Co-authored-by: Weiliangl User <weiliangl@login-node.hosted.internal>

test/registered/distributed/test_disaggregation_dp_attention.py

@@ -110,7 +110,6 @@ class TestDisaggregationDPAttention(PDDisaggregationServerBase):
 
 class TestDisaggregationDPAttentionRoundRobin(TestDisaggregationDPAttention):
     LOAD_BALANCE_METHOD = "round_robin"
-    # TODO: add test for other load balance methods
     # TODO: add a balancedness metric
 
     def test_bench_serving(self):
@@ -130,6 +129,48 @@ class TestDisaggregationDPAttentionRoundRobin(TestDisaggregationDPAttention):
         self.assertEqual(result["completed"], 1000)
 
 
+class TestDisaggregationDPAttentionTotalRequests(TestDisaggregationDPAttention):
+    LOAD_BALANCE_METHOD = "total_requests"
+    test_gsm8k = unittest.skip(
+        "Covered by base class; this class targets total_requests path."
+    )(TestDisaggregationDPAttention.test_gsm8k)
+
+    def test_bench_serving(self):
+        args = get_benchmark_args(
+            base_url=f"http://{self.base_host}:{self.lb_port}",
+            dataset_name="random",
+            tokenizer=self.model,
+            num_prompts=256,
+            random_input_len=2048,
+            random_output_len=512,
+            request_rate=float("inf"),
+            max_concurrency=128,
+        )
+        result = run_benchmark(args)
+        self.assertEqual(result["completed"], 256)
+
+
+class TestDisaggregationDPAttentionTotalTokens(TestDisaggregationDPAttention):
+    LOAD_BALANCE_METHOD = "total_tokens"
+    test_gsm8k = unittest.skip(
+        "Covered by base class; this class targets total_tokens path."
+    )(TestDisaggregationDPAttention.test_gsm8k)
+
+    def test_bench_serving(self):
+        args = get_benchmark_args(
+            base_url=f"http://{self.base_host}:{self.lb_port}",
+            dataset_name="random",
+            tokenizer=self.model,
+            num_prompts=256,
+            random_input_len=2048,
+            random_output_len=512,
+            request_rate=float("inf"),
+            max_concurrency=128,
+        )
+        result = run_benchmark(args)
+        self.assertEqual(result["completed"], 256)
+
+
 @unittest.skip(
     "Skip this test until new testing logic in mini-lb has been updated in docker image."
 )