refactor(e2e_test): fix smg ci e2e test code quality (#16664)
This commit is contained in:
@@ -17,6 +17,8 @@ from .constants import ( # Enums; Convenience sets; Fixture parameters; Default
|
||||
HEALTH_CHECK_INTERVAL,
|
||||
LOCAL_MODES,
|
||||
LOCAL_RUNTIMES,
|
||||
LOG_SEPARATOR_WIDTH,
|
||||
MAX_RETRY_ATTEMPTS,
|
||||
PARAM_BACKEND_ROUTER,
|
||||
PARAM_MODEL,
|
||||
PARAM_SETUP_BACKEND,
|
||||
@@ -82,6 +84,8 @@ __all__ = [
|
||||
"DEFAULT_STARTUP_TIMEOUT",
|
||||
"DEFAULT_ROUTER_TIMEOUT",
|
||||
"HEALTH_CHECK_INTERVAL",
|
||||
"MAX_RETRY_ATTEMPTS",
|
||||
"LOG_SEPARATOR_WIDTH",
|
||||
# Env vars
|
||||
"ENV_MODELS",
|
||||
"ENV_BACKENDS",
|
||||
|
||||
@@ -58,3 +58,11 @@ DEFAULT_HOST = "127.0.0.1"
|
||||
DEFAULT_STARTUP_TIMEOUT = 300
|
||||
DEFAULT_ROUTER_TIMEOUT = 60
|
||||
HEALTH_CHECK_INTERVAL = 5
|
||||
|
||||
# Retry configuration
|
||||
MAX_RETRY_ATTEMPTS = (
|
||||
6 # Max retries with exponential backoff (total ~63s: 1+2+4+8+16+32)
|
||||
)
|
||||
|
||||
# Display formatting
|
||||
LOG_SEPARATOR_WIDTH = 60 # Width for log separator lines (e.g., "="*60)
|
||||
|
||||
@@ -328,6 +328,30 @@ class Gateway:
|
||||
# Worker Management APIs
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _worker_from_api_response(self, w: dict) -> WorkerInfo:
|
||||
"""Convert API response dict to WorkerInfo.
|
||||
|
||||
Args:
|
||||
w: Worker dict from API response.
|
||||
|
||||
Returns:
|
||||
WorkerInfo object.
|
||||
"""
|
||||
status = "healthy" if w.get("is_healthy", False) else "unhealthy"
|
||||
return WorkerInfo(
|
||||
id=w.get("id", ""),
|
||||
url=w.get("url", ""),
|
||||
model=w.get("model_id"),
|
||||
status=status,
|
||||
pending_requests=w.get("load", 0),
|
||||
metadata={
|
||||
"worker_type": w.get("worker_type"),
|
||||
"connection_mode": w.get("connection_mode"),
|
||||
"priority": w.get("priority"),
|
||||
"cost": w.get("cost"),
|
||||
},
|
||||
)
|
||||
|
||||
def list_workers(self, timeout: float = 5.0) -> list[WorkerInfo]:
|
||||
"""List all workers connected to the gateway.
|
||||
|
||||
@@ -338,26 +362,9 @@ class Gateway:
|
||||
resp = httpx.get(f"{self.base_url}/workers", timeout=timeout)
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
workers = []
|
||||
for w in data.get("workers", []):
|
||||
# Map API fields to WorkerInfo
|
||||
status = "healthy" if w.get("is_healthy", False) else "unhealthy"
|
||||
workers.append(
|
||||
WorkerInfo(
|
||||
id=w.get("id", ""),
|
||||
url=w.get("url", ""),
|
||||
model=w.get("model_id"),
|
||||
status=status,
|
||||
pending_requests=w.get("load", 0),
|
||||
metadata={
|
||||
"worker_type": w.get("worker_type"),
|
||||
"connection_mode": w.get("connection_mode"),
|
||||
"priority": w.get("priority"),
|
||||
"cost": w.get("cost"),
|
||||
},
|
||||
)
|
||||
)
|
||||
return workers
|
||||
return [
|
||||
self._worker_from_api_response(w) for w in data.get("workers", [])
|
||||
]
|
||||
return []
|
||||
except (httpx.RequestError, httpx.TimeoutException):
|
||||
return []
|
||||
@@ -374,21 +381,7 @@ class Gateway:
|
||||
try:
|
||||
resp = httpx.get(f"{self.base_url}/workers/{worker_id}", timeout=timeout)
|
||||
if resp.status_code == 200:
|
||||
w = resp.json()
|
||||
status = "healthy" if w.get("is_healthy", False) else "unhealthy"
|
||||
return WorkerInfo(
|
||||
id=w.get("id", ""),
|
||||
url=w.get("url", ""),
|
||||
model=w.get("model_id"),
|
||||
status=status,
|
||||
pending_requests=w.get("load", 0),
|
||||
metadata={
|
||||
"worker_type": w.get("worker_type"),
|
||||
"connection_mode": w.get("connection_mode"),
|
||||
"priority": w.get("priority"),
|
||||
"cost": w.get("cost"),
|
||||
},
|
||||
)
|
||||
return self._worker_from_api_response(resp.json())
|
||||
return None
|
||||
except (httpx.RequestError, httpx.TimeoutException):
|
||||
return None
|
||||
|
||||
@@ -227,7 +227,7 @@ class GPUAllocator:
|
||||
name = pynvml.nvmlDeviceGetName(handle)
|
||||
# Handle bytes vs string return type (varies by pynvml version)
|
||||
if isinstance(name, bytes):
|
||||
name = name.decode("utf-8")
|
||||
name = name.decode("utf-8", errors="replace")
|
||||
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
# Convert bytes to MB
|
||||
memory_mb = mem_info.total // (1024 * 1024)
|
||||
|
||||
@@ -305,9 +305,6 @@ class ModelPool:
|
||||
if ib_device:
|
||||
logger.info("Detected InfiniBand device: %s", ib_device)
|
||||
|
||||
# Track bootstrap ports for PD groups (all PD workers of same model/mode share one)
|
||||
pd_bootstrap_ports: dict[tuple[str, ConnectionMode], int] = {}
|
||||
|
||||
deferred: list[str] = []
|
||||
|
||||
# Process requirements in order - all workers treated uniformly
|
||||
@@ -340,13 +337,8 @@ class ModelPool:
|
||||
deferred.append(str(identity))
|
||||
continue
|
||||
|
||||
# Get bootstrap port for PD workers (shared within model/mode group)
|
||||
bootstrap_port = None
|
||||
if identity.is_prefill or identity.is_decode:
|
||||
pd_key = (identity.model_id, identity.mode)
|
||||
if pd_key not in pd_bootstrap_ports:
|
||||
pd_bootstrap_ports[pd_key] = get_open_port()
|
||||
bootstrap_port = pd_bootstrap_ports[pd_key]
|
||||
# Each prefill worker needs its own bootstrap port for PD communication
|
||||
bootstrap_port = get_open_port() if identity.is_prefill else None
|
||||
|
||||
# Launch the worker
|
||||
self._launch_model(
|
||||
@@ -354,7 +346,7 @@ class ModelPool:
|
||||
mode=identity.mode,
|
||||
gpu_slot=slots[0],
|
||||
worker_type=identity.worker_type,
|
||||
bootstrap_port=bootstrap_port if identity.is_prefill else None,
|
||||
bootstrap_port=bootstrap_port,
|
||||
ib_device=(
|
||||
ib_device if (identity.is_prefill or identity.is_decode) else None
|
||||
),
|
||||
@@ -888,25 +880,17 @@ class ModelPool:
|
||||
has_pd = any(w.is_prefill or w.is_decode for w in valid_workers)
|
||||
ib_device = detect_ib_device() if has_pd else None
|
||||
|
||||
# Track bootstrap ports for PD groups (shared within model/mode)
|
||||
pd_bootstrap_ports: dict[tuple[str, ConnectionMode], int] = {}
|
||||
|
||||
instances: list[ModelInstance] = []
|
||||
for w in valid_workers:
|
||||
# Get bootstrap port for PD workers
|
||||
bootstrap_port = None
|
||||
if w.is_prefill or w.is_decode:
|
||||
pd_key = (w.model_id, w.mode)
|
||||
if pd_key not in pd_bootstrap_ports:
|
||||
pd_bootstrap_ports[pd_key] = get_open_port()
|
||||
bootstrap_port = pd_bootstrap_ports[pd_key]
|
||||
# Each prefill worker needs its own bootstrap port for PD communication
|
||||
bootstrap_port = get_open_port() if w.is_prefill else None
|
||||
|
||||
instance = self._launch_model(
|
||||
model_id=w.model_id,
|
||||
mode=w.mode,
|
||||
gpu_slot=slot_map.get(w.key),
|
||||
worker_type=w.worker_type,
|
||||
bootstrap_port=bootstrap_port if w.is_prefill else None,
|
||||
bootstrap_port=bootstrap_port,
|
||||
ib_device=ib_device if (w.is_prefill or w.is_decode) else None,
|
||||
instance_key=w.key,
|
||||
)
|
||||
|
||||
@@ -30,12 +30,10 @@ if TYPE_CHECKING:
|
||||
from .simple_eval_common import Eval
|
||||
|
||||
from .simple_eval_common import ChatCompletionSampler, set_ulimit
|
||||
from .simple_eval_mmlu import MMLU_DATASET_URL
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# MMLU dataset URL
|
||||
MMLU_DATASET_URL = "https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv"
|
||||
|
||||
|
||||
@dataclass
|
||||
class EvalConfig:
|
||||
|
||||
@@ -20,6 +20,8 @@ import requests
|
||||
from openai import OpenAI
|
||||
from tqdm import tqdm
|
||||
|
||||
from .constants import MAX_RETRY_ATTEMPTS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
OPENAI_SYSTEM_MESSAGE_API = "You are a helpful assistant."
|
||||
@@ -119,7 +121,6 @@ class ChatCompletionSampler(SamplerBase):
|
||||
image: str,
|
||||
encoding: str = "base64",
|
||||
format: str = "png",
|
||||
fovea: int = 768,
|
||||
):
|
||||
new_image = {
|
||||
"type": "image_url",
|
||||
@@ -141,7 +142,7 @@ class ChatCompletionSampler(SamplerBase):
|
||||
self._pack_message("system", self.system_message)
|
||||
] + message_list
|
||||
trial = 0
|
||||
while trial < 6: # Max 63 seconds backoff (1+2+4+8+16+32)
|
||||
while trial < MAX_RETRY_ATTEMPTS:
|
||||
try:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
@@ -162,14 +163,15 @@ class ChatCompletionSampler(SamplerBase):
|
||||
log_fn(
|
||||
"Request failed (retry %d/%d, backoff %ds): %s",
|
||||
trial + 1,
|
||||
6,
|
||||
MAX_RETRY_ATTEMPTS,
|
||||
exception_backoff,
|
||||
e,
|
||||
)
|
||||
time.sleep(exception_backoff)
|
||||
trial += 1
|
||||
logger.warning(
|
||||
"All retry attempts exhausted after 6 retries, returning empty response"
|
||||
"All retry attempts exhausted after %d retries, returning empty response",
|
||||
MAX_RETRY_ATTEMPTS,
|
||||
)
|
||||
return ""
|
||||
|
||||
|
||||
@@ -25,6 +25,9 @@ from .simple_eval_common import (
|
||||
if TYPE_CHECKING:
|
||||
from .simple_eval_common import SamplerBase
|
||||
|
||||
# MMLU dataset URL (hosted by OpenAI)
|
||||
MMLU_DATASET_URL = "https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv"
|
||||
|
||||
SUBJECT_TO_CATEGORY = {
|
||||
"abstract_algebra": "stem",
|
||||
"anatomy": "other",
|
||||
|
||||
Reference in New Issue
Block a user