refactor(e2e): unify RouterInstance into Gateway class, split conftest.py into modular fixtures (#16671)
This commit is contained in:
Simo Lin
@@ -40,10 +40,11 @@ class Gateway:
|
||||
- Worker management (list, add, remove)
|
||||
- Health and metrics endpoints
|
||||
|
||||
Three startup modes:
|
||||
Four startup modes:
|
||||
1. Regular mode: Start with worker URLs
|
||||
2. PD mode: Start with prefill/decode workers
|
||||
3. IGW mode: Start empty, add workers via API
|
||||
4. Cloud mode: Start with cloud backend (OpenAI, xAI)
|
||||
|
||||
Example (regular mode):
|
||||
gateway = Gateway()
|
||||
@@ -71,6 +72,11 @@ class Gateway:
|
||||
|
||||
# Cleanup
|
||||
gateway.shutdown()
|
||||
|
||||
Example (cloud mode):
|
||||
gateway = Gateway()
|
||||
gateway.start(cloud_backend="openai") # or "xai"
|
||||
# Requires OPENAI_API_KEY or XAI_API_KEY env var
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -97,7 +103,10 @@ class Gateway:
|
||||
self.policy: str = "round_robin"
|
||||
self.pd_mode: bool = False
|
||||
self.igw_mode: bool = False
|
||||
self.cloud_mode: bool = False
|
||||
self.cloud_backend: str | None = None
|
||||
self._started: bool = False
|
||||
self._env: dict[str, str] | None = None # Custom env for subprocess
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
@@ -115,6 +124,9 @@ class Gateway:
|
||||
decode_workers: list["ModelInstance"] | None = None,
|
||||
# IGW mode arguments
|
||||
igw_mode: bool = False,
|
||||
# Cloud mode arguments
|
||||
cloud_backend: str | None = None,
|
||||
history_backend: str = "memory",
|
||||
# Common arguments
|
||||
policy: str = "round_robin",
|
||||
timeout: float = DEFAULT_ROUTER_TIMEOUT,
|
||||
@@ -123,10 +135,11 @@ class Gateway:
|
||||
) -> None:
|
||||
"""Start the gateway.
|
||||
|
||||
Can be started in three modes:
|
||||
Can be started in four modes:
|
||||
1. Regular mode: Provide worker_urls and model_path
|
||||
2. PD mode: Provide prefill_workers and decode_workers
|
||||
3. IGW mode: Set igw_mode=True, add workers later via add_worker()
|
||||
4. Cloud mode: Provide cloud_backend ("openai" or "xai")
|
||||
|
||||
Args:
|
||||
worker_urls: List of worker URLs for regular mode.
|
||||
@@ -134,6 +147,8 @@ class Gateway:
|
||||
prefill_workers: List of prefill ModelInstance objects for PD mode.
|
||||
decode_workers: List of decode ModelInstance objects for PD mode.
|
||||
igw_mode: Start in IGW mode (no workers, add via API).
|
||||
cloud_backend: Cloud backend type ("openai" or "xai").
|
||||
history_backend: History backend for cloud mode ("memory" or "oracle").
|
||||
policy: Routing policy (round_robin, random, etc.)
|
||||
timeout: Startup timeout in seconds.
|
||||
show_output: Show subprocess output (env var override).
|
||||
@@ -150,19 +165,21 @@ class Gateway:
|
||||
is_pd_mode = prefill_workers is not None or decode_workers is not None
|
||||
is_regular_mode = worker_urls is not None
|
||||
is_igw_mode = igw_mode
|
||||
is_cloud_mode = cloud_backend is not None
|
||||
|
||||
# Validate mode exclusivity
|
||||
modes_specified = sum([is_pd_mode, is_regular_mode, is_igw_mode])
|
||||
modes_specified = sum([is_pd_mode, is_regular_mode, is_igw_mode, is_cloud_mode])
|
||||
if modes_specified > 1:
|
||||
raise ValueError(
|
||||
"Cannot specify multiple modes. Choose one of: "
|
||||
"worker_urls (regular), prefill/decode_workers (PD), or igw_mode"
|
||||
"worker_urls (regular), prefill/decode_workers (PD), "
|
||||
"igw_mode, or cloud_backend"
|
||||
)
|
||||
|
||||
if modes_specified == 0:
|
||||
raise ValueError(
|
||||
"Must specify one mode: worker_urls (regular), "
|
||||
"prefill/decode_workers (PD), or igw_mode=True"
|
||||
"prefill/decode_workers (PD), igw_mode=True, or cloud_backend"
|
||||
)
|
||||
|
||||
if show_output is None:
|
||||
@@ -201,6 +218,47 @@ class Gateway:
|
||||
extra_args=extra_args,
|
||||
log_msg=f"PD gateway ({len(prefills)} prefill, {len(decodes)} decode)",
|
||||
)
|
||||
elif is_cloud_mode:
|
||||
# Cloud mode: OpenAI/xAI backend
|
||||
self.pd_mode = False
|
||||
self.igw_mode = False
|
||||
self.cloud_mode = True
|
||||
self.cloud_backend = cloud_backend
|
||||
|
||||
# Get worker URL and API key based on backend
|
||||
if cloud_backend == "openai":
|
||||
worker_url = "https://api.openai.com"
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("OPENAI_API_KEY environment variable required")
|
||||
self._env = os.environ.copy()
|
||||
self._env["OPENAI_API_KEY"] = api_key
|
||||
elif cloud_backend == "xai":
|
||||
worker_url = "https://api.x.ai"
|
||||
api_key = os.environ.get("XAI_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("XAI_API_KEY environment variable required")
|
||||
self._env = os.environ.copy()
|
||||
self._env["XAI_API_KEY"] = api_key
|
||||
else:
|
||||
raise ValueError(f"Unsupported cloud backend: {cloud_backend}")
|
||||
|
||||
mode_args = [
|
||||
"--backend",
|
||||
"openai", # Both OpenAI and xAI use openai backend type
|
||||
"--worker-urls",
|
||||
worker_url,
|
||||
"--history-backend",
|
||||
history_backend,
|
||||
]
|
||||
|
||||
self._launch(
|
||||
mode_args=mode_args,
|
||||
timeout=timeout,
|
||||
show_output=show_output,
|
||||
extra_args=extra_args,
|
||||
log_msg=f"{cloud_backend} cloud gateway",
|
||||
)
|
||||
else:
|
||||
# Regular mode: worker URLs
|
||||
if model_path is None:
|
||||
@@ -249,6 +307,7 @@ class Gateway:
|
||||
|
||||
self.process = subprocess.Popen(
|
||||
cmd,
|
||||
env=self._env, # Use custom env if set (e.g., for cloud mode API keys)
|
||||
stdout=None if show_output else subprocess.PIPE,
|
||||
stderr=None if show_output else subprocess.PIPE,
|
||||
start_new_session=True,
|
||||
|
||||
Reference in New Issue
Block a user