reformat

Signed-off-by: ZhengWG <zwg0606@gmail.com>

Author: ZhengWG

tests/e2e/eplb/test_eplb_e2e.py

@@ -1,12 +1,13 @@
+import json
 import os
-import time
-import subprocess
+import random
 import signal
+import subprocess
+import time
+
 import psutil
-import requests
 import pytest
-import random
-import json
+import requests
 
 
 def kill_process_and_children(pid):
@@ -36,13 +37,17 @@ def kill_all_vllm_related():
             continue
 
 
-def build_expert_map(expert_map_path, num_redundant_expert=0,
-                     num_layer=58, num_device=16, num_original_expert=256,
+def build_expert_map(expert_map_path,
+                     num_redundant_expert=0,
+                     num_layer=58,
+                     num_device=16,
+                     num_original_expert=256,
                      random_seed=42):
     expert_num_list = list(range(num_original_expert))
     random.seed(random_seed)
     if num_redundant_expert > 0:
-        expert_num_list = expert_num_list + random.choices(expert_num_list, k=num_redundant_expert)
+        expert_num_list = expert_num_list + random.choices(
+            expert_num_list, k=num_redundant_expert)
     local_num_expert = len(expert_num_list) // num_device
 
     expert_map = {
@@ -52,16 +57,21 @@ def build_expert_map(expert_map_path, num_redundant_expert=0,
     }
     for layer_id in range(num_layer):
         random.shuffle(expert_num_list)
-        current_expert_distribution = [expert_num_list[i*local_num_expert:(i+1)*local_num_expert] for i in range(num_device)]
+        current_expert_distribution = [
+            expert_num_list[i * local_num_expert:(i + 1) * local_num_expert]
+            for i in range(num_device)
+        ]
         layer_info = {
             "layer_id": layer_id,
             "device_count": num_device,
             "device_list": []
         }
         for device_id in range(num_device):
             layer_info["device_list"].append({
-                "device_id": device_id,
-                "device_expert": current_expert_distribution[device_id]
+                "device_id":
+                device_id,
+                "device_expert":
+                current_expert_distribution[device_id]
             })
         expert_map["layer_list"].append(layer_info)
     with open(expert_map_path, "w") as f:
@@ -77,7 +87,6 @@ def is_port_in_use(port):
 
 def ensure_port_available(port, timeout=30):
     """Wait for a port to become available."""
-    import socket
     start = time.time()
     while time.time() - start < timeout:
         if not is_port_in_use(port):
@@ -107,12 +116,16 @@ def wait_for_port(port, timeout=30):
 def test_eplb_with_redundant_expert(num_redundant_expert):
     # Ensure port is available before starting the test
     if is_port_in_use(PROXY_PORT):
-        print(f"Port {PROXY_PORT} is still in use from previous test, waiting for it to become available...")
+        print(
+            f"Port {PROXY_PORT} is still in use from previous test, waiting for it to become available..."
+        )
         if not ensure_port_available(PROXY_PORT, timeout=300):
-            pytest.skip(f"Port {PROXY_PORT} is still in use after waiting 60 seconds")
-    
+            pytest.skip(
+                f"Port {PROXY_PORT} is still in use after waiting 60 seconds")
+
     print("Launching bash script to run eplb setup...")
-    build_expert_map(EXPERT_MAP_PATH, num_redundant_expert=num_redundant_expert)
+    build_expert_map(EXPERT_MAP_PATH,
+                     num_redundant_expert=num_redundant_expert)
     proc = subprocess.Popen(["bash", SCRIPT_PATH, str(num_redundant_expert)])
     try:
         print("Waiting for proxy port to be available...")
@@ -135,7 +148,7 @@ def test_eplb_with_redundant_expert(num_redundant_expert):
         print("Response:", result)
         assert "text" in result["choices"][0]
         assert len(result["choices"][0]["text"].strip()) > 0
-    
+
     finally:
         # clean up subprocesses
         print("Cleaning up subprocess...")
@@ -147,7 +160,7 @@ def test_eplb_with_redundant_expert(num_redundant_expert):
         if os.path.exists(EXPERT_MAP_PATH):
             os.remove(EXPERT_MAP_PATH)
         kill_all_vllm_related()
-        
+
         # Wait for port to be fully released
         print("Waiting for port to be fully released...")
-        time.sleep(3)
+        time.sleep(3)