chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[smg][ci] preserve model launch order with test collected (#16618)

Browse Source
This commit is contained in:
Simo Lin
2026-01-07 06:16:59 -08:00
committed by GitHub
parent 4d902c8211
commit e432057381
6 changed files with 570 additions and 344 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-rust.yml
+1 -1
View File
@@ -186,7 +186,7 @@ jobs:
extra_deps: ""
env_vars: "SHOW_ROUTER_LOGS=1"
reruns: "--reruns 3 --reruns-delay 2"
- name: router-embeddings
- name: e2e
timeout: 45
test_dirs: "e2e_test/router e2e_test/embeddings"
extra_deps: ""
sgl-model-gateway/e2e_test/conftest.py
+327 -91
View File
@@ -175,15 +175,25 @@ from infra import (
PARAM_MODEL,
PARAM_SETUP_BACKEND,
ConnectionMode,
WorkerIdentity,
WorkerType,
)
# ---------------------------------------------------------------------------
# Test collection: scan for required backends
# Test collection: scan for required workers
# ---------------------------------------------------------------------------
# Global storage for scanned requirements
_scanned_backends: set[str] = set() # {"grpc", "http", "openai", ...}
_scanned_models: set[str] = set() # Models needed by tests
# Track max worker counts: (model_id, mode, worker_type) -> max_count
# This unified approach handles regular, prefill, and decode workers the same way
_worker_counts: dict[tuple[str, ConnectionMode, WorkerType], int] = {}
# Track first-seen order to preserve test collection order
_first_seen_order: list[tuple[str, ConnectionMode, WorkerType]] = []
# Track max GPU requirement for any single test (for validation)
_max_test_gpu_requirement: int = 0
_max_test_name: str = ""
_needs_default_model: bool = False # True if any e2e test lacks explicit model marker
@@ -192,93 +202,259 @@ def pytest_collection_modifyitems(
config: pytest.Config,
items: list[pytest.Item],
) -> None:
"""Scan collected tests to determine required backends and models.
"""Scan collected tests to determine required workers.
This runs after test collection but before tests execute.
It extracts backend requirements from @pytest.mark.parametrize markers.
It extracts worker requirements from markers in test collection order,
tracking the max count needed for each (model, mode, worker_type) combination.
Also tracks the max GPU requirement for any single test for validation.
"""
global _scanned_backends, _scanned_models, _needs_default_model
global _worker_counts, _first_seen_order, _needs_default_model
global _max_test_gpu_requirement, _max_test_name
from infra import MODEL_SPECS
def track_worker(
model_id: str, mode: ConnectionMode, worker_type: WorkerType, count: int
) -> None:
"""Track a worker requirement, updating max count if needed."""
key = (model_id, mode, worker_type)
if key not in _worker_counts:
_first_seen_order.append(key)
_worker_counts[key] = count
else:
_worker_counts[key] = max(_worker_counts[key], count)
def calculate_test_gpus(
model_id: str, prefill: int, decode: int, regular: int
) -> int:
"""Calculate GPU requirement for a single test."""
if model_id not in MODEL_SPECS:
return 0
tp = MODEL_SPECS[model_id].get("tp", 1)
return tp * (prefill + decode + regular)
for item in items:
# Track if this test has an explicit model marker
has_model_marker = False
# Extract model from marker or use default
model_marker = item.get_closest_marker(PARAM_MODEL)
model_id = model_marker.args[0] if model_marker and model_marker.args else None
# Scan parametrize markers for setup_backend
# Check parametrize for model
if model_id is None:
for marker in item.iter_markers("parametrize"):
if marker.args and len(marker.args) >= 2:
param_name = marker.args[0]
if param_name == PARAM_MODEL or PARAM_MODEL in param_name:
param_values = marker.args[1]
if isinstance(param_values, (list, tuple)) and param_values:
model_id = param_values[0] # First model in parametrize
break
# Extract backends from parametrize
backends: list[str] = []
for marker in item.iter_markers("parametrize"):
if marker.args and len(marker.args) >= 2:
param_name = marker.args[0]
param_values = marker.args[1]
if param_name == PARAM_SETUP_BACKEND:
# Extract backend names from parametrize values
if isinstance(param_values, (list, tuple)):
_scanned_backends.update(param_values)
backends.extend(param_values)
elif param_name == PARAM_MODEL or PARAM_MODEL in param_name:
# Extract model names from parametrize
if isinstance(param_values, (list, tuple)):
_scanned_models.update(param_values)
has_model_marker = True
# Check for workers marker (@pytest.mark.workers(...))
workers_marker = item.get_closest_marker("workers")
prefill_count = 0
decode_count = 0
regular_count = 1 # Default to 1 regular worker
if workers_marker:
prefill_count = workers_marker.kwargs.get("prefill") or 0
decode_count = workers_marker.kwargs.get("decode") or 0
regular_count = workers_marker.kwargs.get("count") or 1
# Check for @pytest.mark.model("name") markers
model_marker = item.get_closest_marker(PARAM_MODEL)
if model_marker and model_marker.args:
model_name = model_marker.args[0]
_scanned_models.add(model_name)
has_model_marker = True
# Check if this is an e2e test without an explicit model marker
# Such tests need the DEFAULT_MODEL
if not has_model_marker and item.get_closest_marker("e2e"):
# Track if this test needs default model
is_e2e = item.get_closest_marker("e2e") is not None
if model_id is None and is_e2e:
_needs_default_model = True
model_id = DEFAULT_MODEL
logger.info(
"Scanned test requirements - backends: %s, models: %s, needs default: %s",
_scanned_backends or {"(none)"},
_scanned_models or {"(none)"},
_needs_default_model,
)
# Track worker requirements and calculate this test's GPU requirement
test_gpus = 0
if model_id and backends:
for backend in backends:
# "pd" backend means PD workers
if backend == "pd":
mode = ConnectionMode.HTTP # PD uses HTTP mode
# Default to 1 prefill + 1 decode if not specified
p_count = prefill_count if prefill_count > 0 else 1
d_count = decode_count if decode_count > 0 else 1
track_worker(model_id, mode, WorkerType.PREFILL, p_count)
track_worker(model_id, mode, WorkerType.DECODE, d_count)
test_gpus = max(
test_gpus, calculate_test_gpus(model_id, p_count, d_count, 0)
)
else:
try:
mode = ConnectionMode(backend)
except ValueError:
# Cloud backend (openai, xai, etc.) - skip
continue
# Check if this backend also has PD workers
if prefill_count > 0 or decode_count > 0:
track_worker(model_id, mode, WorkerType.PREFILL, prefill_count)
track_worker(model_id, mode, WorkerType.DECODE, decode_count)
test_gpus = max(
test_gpus,
calculate_test_gpus(
model_id, prefill_count, decode_count, 0
),
)
else:
# Regular worker
track_worker(model_id, mode, WorkerType.REGULAR, regular_count)
test_gpus = max(
test_gpus,
calculate_test_gpus(model_id, 0, 0, regular_count),
)
elif model_id and is_e2e:
# E2E test without explicit backend - will use HTTP by default
track_worker(model_id, ConnectionMode.HTTP, WorkerType.REGULAR, 1)
test_gpus = calculate_test_gpus(model_id, 0, 0, 1)
# Track max GPU requirement across all tests
if test_gpus > _max_test_gpu_requirement:
_max_test_gpu_requirement = test_gpus
_max_test_name = item.nodeid
# Log results
if _worker_counts:
summary = []
for key in _first_seen_order:
model_id, mode, worker_type = key
count = _worker_counts[key]
if worker_type == WorkerType.REGULAR:
summary.append(f"{model_id}:{mode.value}x{count}")
else:
summary.append(f"{model_id}:{mode.value}:{worker_type.value}x{count}")
logger.info("Scanned worker requirements (in test order): %s", summary)
logger.info(
"Max GPU requirement for single test: %d (%s)",
_max_test_gpu_requirement,
_max_test_name,
)
else:
logger.info("Scanned worker requirements: (none)")
def get_pool_requirements() -> list[tuple[str, ConnectionMode]]:
def get_pool_requirements() -> list[WorkerIdentity]:
"""Build pool requirements from scanned test markers.
Returns:
List of (model_id, ConnectionMode) tuples to try to pre-launch.
Models that don't fit will be launched on-demand.
List of WorkerIdentity objects to pre-launch.
Each WorkerIdentity has (model_id, mode, worker_type, index).
Requirements are ordered by first appearance in test collection order,
so workers needed by earlier tests are launched first.
Note:
If a model's first test needs PD workers (prefill/decode), we skip
pre-launching regular workers for that model (they'd be evicted
immediately when PD workers are launched).
"""
models = set(_scanned_models)
# Track which models have PD workers as their first requirement
# These models shouldn't have regular workers pre-launched
models_with_pd_first: set[str] = set()
first_worker_type_per_model: dict[str, WorkerType] = {}
# Add DEFAULT_MODEL if any e2e test lacks an explicit model marker,
# or if no models were specified at all
if _needs_default_model or not models:
models.add(DEFAULT_MODEL)
for model_id, mode, worker_type in _first_seen_order:
if model_id not in first_worker_type_per_model:
first_worker_type_per_model[model_id] = worker_type
if worker_type in (WorkerType.PREFILL, WorkerType.DECODE):
models_with_pd_first.add(model_id)
logger.info(
"Model %s has PD test first - skipping regular worker pre-launch",
model_id,
)
# Convert scanned string backends to ConnectionMode enums
# Filter to local backends only (grpc, http) - cloud backends don't need workers
local_modes: set[ConnectionMode] = set()
for backend in _scanned_backends:
try:
mode = ConnectionMode(backend)
if mode in LOCAL_MODES:
local_modes.add(mode)
except ValueError:
# Not a ConnectionMode (e.g., "openai", "xai", "pd") - skip
pass
# Generate individual WorkerIdentity objects in first-seen order
requirements: list[WorkerIdentity] = []
for model_id, mode, worker_type in _first_seen_order:
# Skip regular workers for models that have PD first
if model_id in models_with_pd_first and worker_type == WorkerType.REGULAR:
continue
# Default to HTTP if no local backends specified
if not local_modes:
local_modes = {ConnectionMode.HTTP}
count = _worker_counts.get((model_id, mode, worker_type), 1)
for i in range(count):
requirements.append(WorkerIdentity(model_id, mode, worker_type, i))
# Build requirements: each model needs each mode
requirements: list[tuple[str, ConnectionMode]] = []
for model in models:
for mode in local_modes:
requirements.append((model, mode))
# Add default if no requirements
if not requirements:
requirements.append(WorkerIdentity(DEFAULT_MODEL, ConnectionMode.HTTP))
return requirements
def validate_gpu_requirements() -> tuple[int, int]:
"""Check if there are enough GPUs for any single test.
Returns:
Tuple of (max_required_gpus, available_gpus).
Note:
We check the max requirement for any single test, not the sum.
Workers can be evicted between tests, so we only need enough GPUs
for the most demanding test.
"""
# Count available GPUs
available_gpus = 0
try:
import torch
if torch.cuda.is_available():
available_gpus = torch.cuda.device_count()
except ImportError:
pass
return _max_test_gpu_requirement, available_gpus
def pytest_collection_finish(session: pytest.Session) -> None:
"""Validate GPU requirements after test collection.
This runs after all tests are collected but before any tests execute.
Fails fast if any single test requires more GPUs than available.
"""
if not _worker_counts:
return
# Skip validation if model pool is disabled
if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
return
max_required, available_gpus = validate_gpu_requirements()
if max_required > available_gpus:
raise pytest.UsageError(
f"\n{'='*60}\n"
f"GPU REQUIREMENTS EXCEEDED\n"
f"{'='*60}\n"
f"Test '{_max_test_name}' requires {max_required} GPUs\n"
f"Available: {available_gpus} GPUs\n"
f"\nOptions:\n"
f" 1. Run tests that fit: pytest -k 'not {_max_test_name.split('::')[0]}'\n"
f" 2. Reduce workers: @pytest.mark.workers(prefill=1, decode=1)\n"
f" 3. Skip GPU tests: SKIP_MODEL_POOL=1 pytest\n"
f"{'='*60}"
)
logger.info(
"GPU validation passed: max %d required (by %s), %d available",
max_required,
_max_test_name,
available_gpus,
)
# ---------------------------------------------------------------------------
# Custom pytest markers
# ---------------------------------------------------------------------------
@@ -333,13 +509,15 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
routers (~1-2s) pointing to these workers.
Startup behavior:
- Scans test markers to determine required (model, mode) combinations
- Launches workers sequentially, but they boot up concurrently
- Scans test markers to determine required workers (model, mode, type, count)
- Launches workers in test collection order
- Waits for all workers to become healthy before returning
Test requirements are auto-detected from:
- @pytest.mark.parametrize("setup_backend", ["grpc", "http"])
- @pytest.mark.parametrize("setup_backend", ["grpc", "http", "pd"])
- @pytest.mark.model("model-name")
- @pytest.mark.workers(count=N) for regular workers
- @pytest.mark.workers(prefill=N, decode=N) for PD workers
Environment variable overrides:
- E2E_MODELS: Comma-separated model IDs (e.g., "llama-8b,qwen-7b")
@@ -388,15 +566,20 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
if not backend_modes:
backend_modes = {ConnectionMode.HTTP}
requirements = [(m, b) for m in models for b in backend_modes]
logger.info("Using env var requirements: %s", requirements)
# Create WorkerIdentity objects (regular workers only from env vars)
requirements = [
WorkerIdentity(m, b, WorkerType.REGULAR, 0)
for m in models
for b in backend_modes
]
logger.info("Using env var requirements: %s", [str(r) for r in requirements])
else:
# Use scanned requirements from test markers
requirements = get_pool_requirements()
logger.info("Using scanned requirements: %s", requirements)
logger.info("Using scanned requirements: %s", [str(r) for r in requirements])
# Filter to valid models
requirements = [(m, b) for m, b in requirements if m in MODEL_SPECS]
requirements = [r for r in requirements if r.model_id in MODEL_SPECS]
if not requirements:
logger.warning("No valid requirements, model pool will be empty")
@@ -408,7 +591,10 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
_model_pool = ModelPool(allocator)
startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300"))
_model_pool.startup(requirements=requirements, startup_timeout=startup_timeout)
_model_pool.startup(
requirements=requirements,
startup_timeout=startup_timeout,
)
# Log final GPU allocation summary
logger.info(_model_pool.allocator.summary())
@@ -624,15 +810,16 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
f"({num_prefill} prefill + {num_decode} decode), found {gpu_count}"
)
# Try to use pre-launched PD workers, or launch new ones if needed
# Try to use pre-launched PD workers, or launch additional ones if needed
existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL)
existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE)
if (
len(existing_prefills) >= num_prefill
and len(existing_decodes) >= num_decode
):
# Use pre-launched workers
# Calculate how many more we need (if any)
missing_prefill = max(0, num_prefill - len(existing_prefills))
missing_decode = max(0, num_decode - len(existing_decodes))
if missing_prefill == 0 and missing_decode == 0:
# Use pre-launched workers (we have enough)
prefills = existing_prefills[:num_prefill]
decodes = existing_decodes[:num_decode]
logger.info(
@@ -641,13 +828,48 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
len(decodes),
)
else:
# Launch new PD workers (custom config or not pre-launched)
prefills, decodes = model_pool.launch_pd_workers(
model_id=model_id,
num_prefill=num_prefill,
num_decode=num_decode,
startup_timeout=300,
# Build WorkerIdentity list for missing workers
workers_to_launch: list[WorkerIdentity] = []
for i in range(missing_prefill):
workers_to_launch.append(
WorkerIdentity(
model_id,
ConnectionMode.HTTP,
WorkerType.PREFILL,
len(existing_prefills) + i,
)
)
for i in range(missing_decode):
workers_to_launch.append(
WorkerIdentity(
model_id,
ConnectionMode.HTTP,
WorkerType.DECODE,
len(existing_decodes) + i,
)
)
logger.info(
"Have %d/%d prefill, %d/%d decode. Launching %d more workers",
len(existing_prefills),
num_prefill,
len(existing_decodes),
num_decode,
len(workers_to_launch),
)
new_instances = model_pool.launch_workers(
workers_to_launch, startup_timeout=300
)
# Combine existing + newly launched
new_prefills = [
w for w in new_instances if w.worker_type == WorkerType.PREFILL
]
new_decodes = [
w for w in new_instances if w.worker_type == WorkerType.DECODE
]
prefills = existing_prefills + new_prefills
decodes = existing_decodes + new_decodes
model_path = prefills[0].model_path if prefills else None
@@ -698,17 +920,31 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
try:
if num_workers > 1:
# Launch multiple workers on separate GPUs
instances = model_pool.launch_regular_workers(
model_id=model_id,
num_workers=num_workers,
mode=connection_mode,
startup_timeout=300,
)
if not instances:
pytest.fail(
f"Failed to launch {num_workers} workers for {model_id}"
# Check existing workers
existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR)
existing_for_mode = [w for w in existing if w.mode == connection_mode]
if len(existing_for_mode) >= num_workers:
instances = existing_for_mode[:num_workers]
else:
# Launch missing workers
missing = num_workers - len(existing_for_mode)
workers_to_launch = [
WorkerIdentity(
model_id,
connection_mode,
WorkerType.REGULAR,
len(existing_for_mode) + i,
)
for i in range(missing)
]
new_instances = model_pool.launch_workers(
workers_to_launch, startup_timeout=300
)
instances = existing_for_mode + new_instances
if not instances:
pytest.fail(f"Failed to get {num_workers} workers for {model_id}")
worker_urls = [inst.worker_url for inst in instances]
model_path = instances[0].model_path
else:
sgl-model-gateway/e2e_test/infra/__init__.py
+3 -2
View File
@@ -37,7 +37,7 @@ from .gpu_allocator import (
)
from .gpu_monitor import GPUMonitor
from .gpu_monitor import should_monitor as should_monitor_gpu
from .model_pool import ModelInstance, ModelPool
from .model_pool import ModelInstance, ModelPool, WorkerIdentity
from .model_specs import ( # Default model paths; Model groups
CHAT_MODELS,
DEFAULT_EMBEDDING_MODEL_PATH,
@@ -63,10 +63,11 @@ from .process_utils import (
from .run_eval import run_eval
__all__ = [
# Enums
# Enums and Identity
"ConnectionMode",
"WorkerType",
"Runtime",
"WorkerIdentity",
# Convenience sets
"LOCAL_MODES",
"LOCAL_RUNTIMES",
sgl-model-gateway/e2e_test/infra/constants.py
+1 -1
View File
@@ -1,6 +1,6 @@
"""Constants and enums for E2E test infrastructure."""
from enum import Enum, auto
from enum import Enum
class ConnectionMode(str, Enum):
sgl-model-gateway/e2e_test/infra/gpu_allocator.py
+21 -9
View File
@@ -244,19 +244,27 @@ class GPUAllocator:
logger.warning("Failed to detect GPUs: %s", e)
return []
def allocate_slots(self, model_specs: dict[str, dict]) -> list[GPUSlot]:
def allocate_slots(
self, model_specs: dict[str, dict], preserve_order: bool = False
) -> list[GPUSlot]:
"""Allocate GPU slots based on model memory requirements.
Uses a first-fit decreasing bin-packing algorithm:
Uses a first-fit decreasing bin-packing algorithm by default:
1. Sort models by memory requirement (largest first)
2. For each model, find the first GPU(s) that can fit it
3. For multi-GPU models, find consecutive GPUs
When preserve_order=True, processes models in dict insertion order
(test collection order) instead of sorting by memory. This ensures
models needed by earlier tests are allocated first.
Note: This method tracks used GPUs across multiple calls, so subsequent
allocations will use different GPUs than previous ones.
Args:
model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys
preserve_order: If True, allocate in dict order (test order) instead
of sorting by memory size. Default False.
Returns:
List of GPUSlots with assigned models (only the newly allocated slots)
@@ -265,17 +273,21 @@ class GPUAllocator:
logger.warning("No GPUs available for allocation")
return []
# Sort models by memory requirement (largest first for better packing)
sorted_models = sorted(
model_specs.items(),
key=lambda x: x[1].get("memory_gb", 0),
reverse=True,
)
if preserve_order:
# Process in dict insertion order (test collection order)
ordered_models = list(model_specs.items())
else:
# Sort models by memory requirement (largest first for better packing)
ordered_models = sorted(
model_specs.items(),
key=lambda x: x[1].get("memory_gb", 0),
reverse=True,
)
# Track new slots allocated in this call
new_slots: list[GPUSlot] = []
for model_id, spec in sorted_models:
for model_id, spec in ordered_models:
memory_gb = spec.get("memory_gb", 16)
tp_size = spec.get("tp", 1)
sgl-model-gateway/e2e_test/infra/model_pool.py
+217 -240
View File
@@ -31,9 +31,63 @@ from .process_utils import detect_ib_device
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class WorkerIdentity:
"""Unique identity for a single worker instance.
Each worker is uniquely identified by (model_id, mode, worker_type, index).
For example:
- llama-8b:http (regular worker, index 0)
- llama-8b:http:prefill_0 (first prefill worker)
- llama-8b:http:prefill_1 (second prefill worker)
- llama-8b:http:decode_0 (first decode worker)
Frozen/hashable so it can be used in sets and as dict keys for deduplication.
"""
model_id: str
mode: ConnectionMode = ConnectionMode.HTTP
worker_type: WorkerType = WorkerType.REGULAR
index: int = 0
@property
def is_prefill(self) -> bool:
"""Check if this is a prefill worker."""
return self.worker_type == WorkerType.PREFILL
@property
def is_decode(self) -> bool:
"""Check if this is a decode worker."""
return self.worker_type == WorkerType.DECODE
@property
def is_regular(self) -> bool:
"""Check if this is a regular worker."""
return self.worker_type == WorkerType.REGULAR
@property
def key(self) -> str:
"""Unique key for this worker instance."""
if self.worker_type == WorkerType.REGULAR:
if self.index == 0:
return f"{self.model_id}:{self.mode.value}"
return f"{self.model_id}:{self.mode.value}:{self.index}"
return (
f"{self.model_id}:{self.mode.value}:{self.worker_type.value}_{self.index}"
)
def __str__(self) -> str:
"""String representation for logging."""
return self.key
@dataclass
class ModelInstance:
"""A running model instance."""
"""A running model instance.
Contains both identity (model_id, mode, worker_type) and runtime state
(process, port, gpu_slot, etc.).
"""
model_id: str
mode: ConnectionMode
@@ -42,21 +96,20 @@ class ModelInstance:
port: int
process: subprocess.Popen
gpu_slot: GPUSlot | None
key: str # Unique instance key (e.g., "llama-8b:http:prefill_0")
worker_type: WorkerType = WorkerType.REGULAR
bootstrap_port: int | None = None # For prefill workers in PD mode
last_used: float = 0.0 # Timestamp for MRU eviction
_healthy: bool = False # Track if initial health check passed
@property
def key(self) -> str:
"""Unique key for this instance.
Regular: 'model_id:mode' (e.g., 'llama-8b:http')
PD workers: 'model_id:mode:worker_type' (e.g., 'llama-8b:http:prefill')
"""
if self.worker_type == WorkerType.REGULAR:
return f"{self.model_id}:{self.mode.value}"
return f"{self.model_id}:{self.mode.value}:{self.worker_type.value}"
def identity(self) -> WorkerIdentity:
"""Get the identity (model_id, mode, worker_type) of this instance."""
return WorkerIdentity(
model_id=self.model_id,
mode=self.mode,
worker_type=self.worker_type,
)
@property
def worker_url(self) -> str:
@@ -200,86 +253,120 @@ class ModelPool:
def startup(
self,
requirements: list[tuple[str, ConnectionMode]] | None = None,
requirements: list[WorkerIdentity] | None = None,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
) -> None:
"""Start worker processes for the required models.
"""Start worker processes for the required workers in order.
Workers are launched sequentially (one Popen at a time) but boot up
concurrently since model loading happens in parallel across processes.
This method blocks until all workers pass health checks.
All worker types (regular, prefill, decode) are handled uniformly.
Each WorkerIdentity uniquely identifies a worker by (model_id, mode,
worker_type, index).
Args:
requirements: List of (model_id, mode) tuples specifying what to start.
mode is ConnectionMode.HTTP or ConnectionMode.GRPC.
requirements: List of WorkerIdentity specifying what to start.
If None, starts default model in HTTP mode.
startup_timeout: Timeout in seconds for all models to become healthy.
"""
self._startup_timeout = startup_timeout
if requirements is None:
requirements = [(DEFAULT_MODEL, ConnectionMode.HTTP)]
requirements = [WorkerIdentity(DEFAULT_MODEL, ConnectionMode.HTTP)]
# Deduplicate and validate
requirements = list(set(requirements))
valid_requirements = []
for model_id, mode in requirements:
if model_id not in MODEL_SPECS:
logger.warning("Unknown model %s, skipping", model_id)
# Validate requirements
valid_requirements: list[WorkerIdentity] = []
for identity in requirements:
if identity.model_id not in MODEL_SPECS:
logger.warning("Unknown model %s, skipping", identity.model_id)
continue
if mode not in LOCAL_MODES:
logger.warning("Invalid mode %s for %s, skipping", mode, model_id)
if identity.mode not in LOCAL_MODES:
logger.warning(
"Invalid mode %s for %s, skipping", identity.mode, identity.model_id
)
continue
valid_requirements.append((model_id, mode))
valid_requirements.append(identity)
if not valid_requirements:
logger.warning("No valid requirements to start")
return
logger.info("Starting model pool with: %s", valid_requirements)
logger.info(
"Starting model pool with %d workers: %s",
len(valid_requirements),
[str(r) for r in valid_requirements],
)
# Build allocation specs - each (model, mode) combo needs its own slot
# Use "model_id:mode" as the allocation key
allocation_specs = {}
for model_id, mode in valid_requirements:
spec = MODEL_SPECS[model_id]
key = f"{model_id}:{mode.value}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": spec.get("tp", 1),
# Detect IB device once for PD workers
has_pd = any(r.is_prefill or r.is_decode for r in valid_requirements)
ib_device = detect_ib_device() if has_pd else None
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
for identity in valid_requirements:
spec = get_model_spec(identity.model_id)
tp = spec.get("tp", 1)
# Check if we have enough GPUs
available_gpus = self.allocator.available_gpus()
if len(available_gpus) < tp:
logger.info(
"Not enough GPUs for %s (need %d, have %d), deferring",
identity,
tp,
len(available_gpus),
)
deferred.append(str(identity))
continue
# Allocate GPU slot
allocation_specs = {
identity.key: {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": tp,
}
}
slots = self.allocator.allocate_slots(allocation_specs, preserve_order=True)
if not slots:
deferred.append(str(identity))
continue
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
# 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]
# Track which models got slots
launched_keys = set()
# Launch the worker
self._launch_model(
model_id=identity.model_id,
mode=identity.mode,
gpu_slot=slots[0],
worker_type=identity.worker_type,
bootstrap_port=bootstrap_port if identity.is_prefill else None,
ib_device=(
ib_device if (identity.is_prefill or identity.is_decode) else None
),
instance_key=identity.key,
)
if not slots:
logger.warning("No GPU slots allocated, launching without GPU assignment")
# Fallback: launch without specific GPU assignment
for model_id, mode in valid_requirements:
self._launch_model(model_id, mode, gpu_slot=None)
launched_keys.add(f"{model_id}:{mode.value}")
else:
# Launch on allocated slots
for slot in slots:
if slot.assigned_model:
# Parse "model_id:mode" back
model_id, mode_str = slot.assigned_model.rsplit(":", 1)
mode = ConnectionMode(mode_str)
self._launch_model(model_id, mode, gpu_slot=slot)
launched_keys.add(slot.assigned_model)
# Log models that will be launched on-demand (not enough GPUs to pre-launch)
all_keys = set(allocation_specs.keys())
deferred_keys = all_keys - launched_keys
if deferred_keys:
# Log deferred workers
if deferred:
logger.info(
"%d models deferred for on-demand launch: %s",
len(deferred_keys),
deferred_keys,
"%d workers deferred for on-demand launch: %s",
len(deferred),
deferred,
)
# Wait for all launched models to be healthy
@@ -391,6 +478,7 @@ class ModelPool:
port=port,
process=proc,
gpu_slot=gpu_slot,
key=key,
worker_type=worker_type,
bootstrap_port=bootstrap_port,
last_used=time.time(),
@@ -714,230 +802,119 @@ class ModelPool:
if inst.model_id == model_id and inst.worker_type == worker_type
]
def launch_regular_workers(
def launch_workers(
self,
model_id: str,
num_workers: int,
mode: ConnectionMode = ConnectionMode.HTTP,
workers: list[WorkerIdentity],
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
allow_eviction: bool = True,
) -> list[ModelInstance]:
"""Launch multiple regular workers for load balancing.
"""Launch workers of any type.
This is the unified method for launching workers. It handles all worker
types (regular, prefill, decode) uniformly.
Args:
model_id: Model identifier from MODEL_SPECS.
num_workers: Number of workers to launch.
mode: Connection mode (HTTP or GRPC).
workers: List of WorkerIdentity objects specifying workers to launch.
startup_timeout: Timeout for workers to become healthy.
allow_eviction: If True, evict MRU models to free GPUs.
Returns:
List of ModelInstance objects.
List of launched ModelInstance objects.
"""
if not workers:
return []
self._startup_timeout = startup_timeout
if model_id not in MODEL_SPECS:
raise ValueError(f"Unknown model: {model_id}")
# Validate all workers
valid_workers: list[WorkerIdentity] = []
for w in workers:
if w.model_id not in MODEL_SPECS:
logger.warning("Unknown model %s, skipping", w.model_id)
continue
if w.mode not in LOCAL_MODES:
logger.warning("Invalid mode %s, skipping", w.mode)
continue
valid_workers.append(w)
spec = get_model_spec(model_id)
tp = spec.get("tp", 1)
required_gpus = num_workers * tp
if not valid_workers:
return []
# Calculate total GPUs needed
total_gpus = 0
for w in valid_workers:
spec = get_model_spec(w.model_id)
total_gpus += spec.get("tp", 1)
# Check if we have enough GPUs
available = self.allocator.available_gpus()
if len(available) < required_gpus:
if len(available) < total_gpus:
if allow_eviction:
logger.info(
"Need %d GPUs for %d workers, only %d available. Evicting MRU models...",
required_gpus,
num_workers,
"Need %d GPUs for %d workers, only %d available. Evicting...",
total_gpus,
len(valid_workers),
len(available),
)
# Exclude REGULAR workers of same model/mode from eviction
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_mode=mode,
exclude_worker_types={WorkerType.REGULAR},
)
self._evict_for_gpus(total_gpus)
else:
logger.info(
"Need %d GPUs for %d workers, only %d available. "
"Skipping (eviction not allowed).",
required_gpus,
num_workers,
logger.warning(
"Need %d GPUs, only %d available. Skipping launch.",
total_gpus,
len(available),
)
return []
# Build allocation specs for all workers
# Build allocation specs
allocation_specs = {}
for i in range(num_workers):
key = f"{model_id}:{mode.value}:{i}"
allocation_specs[key] = {
for w in valid_workers:
spec = get_model_spec(w.model_id)
allocation_specs[w.key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": tp,
"tp": spec.get("tp", 1),
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
slot_map = {slot.assigned_model: slot for slot in slots}
slots = self.allocator.allocate_slots(allocation_specs, preserve_order=True)
slot_map = {s.assigned_model: s for s in slots}
if not slots:
raise RuntimeError(
f"Failed to allocate GPU slots for {num_workers} workers after eviction. "
f"Need {required_gpus} GPUs."
f"Failed to allocate GPU slots for {len(valid_workers)} workers"
)
# Detect IB device for PD workers
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]
# Launch workers
for i in range(num_workers):
key = f"{model_id}:{mode.value}:{i}"
gpu_slot = slot_map.get(key)
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.REGULAR,
instance_key=key,
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,
ib_device=ib_device if (w.is_prefill or w.is_decode) else None,
instance_key=w.key,
)
instances.append(instance)
# Wait for all to be healthy
self._wait_all_healthy()
return instances
def launch_pd_workers(
self,
model_id: str,
num_prefill: int = 1,
num_decode: int = 1,
mode: ConnectionMode = ConnectionMode.HTTP,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
allow_eviction: bool = True,
) -> tuple[list[ModelInstance], list[ModelInstance]]:
"""Launch prefill and decode workers for PD disaggregation.
Args:
model_id: Model identifier from MODEL_SPECS.
num_prefill: Number of prefill workers to launch. Defaults to 1.
num_decode: Number of decode workers to launch. Defaults to 1.
mode: Connection mode (HTTP or GRPC).
startup_timeout: Timeout for workers to become healthy.
allow_eviction: If True, evict MRU models to free GPUs. If False,
return empty lists when not enough GPUs available.
Returns:
Tuple of (prefill_instances, decode_instances).
"""
self._startup_timeout = startup_timeout
if model_id not in MODEL_SPECS:
raise ValueError(f"Unknown model: {model_id}")
spec = get_model_spec(model_id)
ib_device = detect_ib_device()
if ib_device:
logger.info("Detected InfiniBand device: %s", ib_device)
# Calculate total GPUs needed for PD workers
tp = spec.get("tp", 1)
required_gpus = (num_prefill + num_decode) * tp
# Check if we have enough GPUs
available = self.allocator.available_gpus()
if len(available) < required_gpus:
if allow_eviction:
logger.info(
"Need %d GPUs for PD workers, only %d available. Evicting MRU models...",
required_gpus,
len(available),
)
# Exclude PD workers of same model/mode, but evict REGULAR workers
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_mode=mode,
exclude_worker_types={WorkerType.PREFILL, WorkerType.DECODE},
)
else:
logger.info(
"Need %d GPUs for PD workers, only %d available. "
"Skipping pre-launch (eviction not allowed).",
required_gpus,
len(available),
)
return [], []
# Build allocation specs for all PD workers
# Each worker needs its own GPU slot
allocation_specs = {}
for i in range(num_prefill):
key = f"{model_id}:{mode.value}:prefill_{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": tp,
}
for i in range(num_decode):
key = f"{model_id}:{mode.value}:decode_{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": tp,
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
slot_map = {slot.assigned_model: slot for slot in slots}
if not slots:
raise RuntimeError(
f"Failed to allocate GPU slots for PD workers after eviction. "
f"Need {required_gpus} GPUs."
)
prefill_instances: list[ModelInstance] = []
decode_instances: list[ModelInstance] = []
# Launch prefill workers
for i in range(num_prefill):
key = f"{model_id}:{mode.value}:prefill_{i}"
gpu_slot = slot_map.get(key)
bootstrap_port = get_open_port()
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.PREFILL,
bootstrap_port=bootstrap_port,
ib_device=ib_device,
instance_key=key,
)
prefill_instances.append(instance)
# Launch decode workers
for i in range(num_decode):
key = f"{model_id}:{mode.value}:decode_{i}"
gpu_slot = slot_map.get(key)
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.DECODE,
ib_device=ib_device,
instance_key=key,
)
decode_instances.append(instance)
# Wait for all to be healthy
self._wait_all_healthy()
return prefill_instances, decode_instances
def get_client(
self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP
) -> "openai.OpenAI":