1232 lines
45 KiB
Python
1232 lines
45 KiB
Python
"""Model pool for managing pre-loaded models across GPUs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING
|
|
|
|
import httpx
|
|
|
|
if TYPE_CHECKING:
|
|
import openai
|
|
|
|
from .constants import (
|
|
DEFAULT_HOST,
|
|
DEFAULT_MODEL,
|
|
DEFAULT_STARTUP_TIMEOUT,
|
|
ENV_SHOW_WORKER_LOGS,
|
|
HEALTH_CHECK_INTERVAL,
|
|
INITIAL_GRACE_PERIOD,
|
|
LAUNCH_STAGGER_DELAY,
|
|
LOCAL_MODES,
|
|
ConnectionMode,
|
|
WorkerType,
|
|
)
|
|
from .gpu_allocator import GPUAllocator, GPUSlot, get_open_port
|
|
from .model_specs import MODEL_SPECS, get_model_spec
|
|
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.
|
|
|
|
Contains both identity (model_id, mode, worker_type) and runtime state
|
|
(process, port, gpu_slot, etc.).
|
|
"""
|
|
|
|
model_id: str
|
|
mode: ConnectionMode
|
|
model_path: str
|
|
base_url: str
|
|
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
|
|
|
|
# Reference counting for safe parallel test execution
|
|
_ref_count: int = 0
|
|
_ref_lock: threading.Lock = field(default_factory=threading.Lock)
|
|
|
|
@property
|
|
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 is_in_use(self) -> bool:
|
|
"""Check if this instance has active references (tests using it)."""
|
|
with self._ref_lock:
|
|
return self._ref_count > 0
|
|
|
|
def acquire(self) -> None:
|
|
"""Acquire a reference to this instance.
|
|
|
|
Call this before using the instance in a test to prevent eviction.
|
|
Must be paired with a release() call when done.
|
|
Also updates last_used timestamp atomically with ref count.
|
|
"""
|
|
with self._ref_lock:
|
|
self._ref_count += 1
|
|
self.last_used = time.time()
|
|
logger.debug(
|
|
"Acquired reference to %s (ref_count=%d)", self.key, self._ref_count
|
|
)
|
|
|
|
def release(self) -> None:
|
|
"""Release a reference to this instance.
|
|
|
|
Call this when done using the instance in a test.
|
|
"""
|
|
with self._ref_lock:
|
|
if self._ref_count > 0:
|
|
self._ref_count -= 1
|
|
logger.debug(
|
|
"Released reference to %s (ref_count=%d)",
|
|
self.key,
|
|
self._ref_count,
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"Attempted to release reference to %s with ref_count=0", self.key
|
|
)
|
|
|
|
@property
|
|
def worker_url(self) -> str:
|
|
"""URL to use when connecting router to this worker."""
|
|
if self.mode == ConnectionMode.GRPC:
|
|
return f"grpc://{DEFAULT_HOST}:{self.port}"
|
|
return self.base_url
|
|
|
|
def is_alive(self) -> bool:
|
|
"""Check if the process is still running."""
|
|
return self.process.poll() is None
|
|
|
|
def health_check(self, timeout: float = 5.0) -> bool:
|
|
"""Check if the model server is healthy.
|
|
|
|
Uses HTTP /health endpoint for HTTP workers, gRPC health check for gRPC workers.
|
|
"""
|
|
if self.mode == ConnectionMode.GRPC:
|
|
return self._grpc_health_check(timeout)
|
|
return self._http_health_check(timeout)
|
|
|
|
def _http_health_check(self, timeout: float = 5.0) -> bool:
|
|
"""Check health via HTTP /health endpoint."""
|
|
try:
|
|
resp = httpx.get(f"{self.base_url}/health", timeout=timeout)
|
|
return resp.status_code == 200
|
|
except (httpx.RequestError, httpx.TimeoutException):
|
|
return False
|
|
|
|
def deep_health_check(self, timeout: float = 30.0) -> bool:
|
|
"""Deep health check that verifies the model can actually generate.
|
|
|
|
Uses /health_generate for HTTP workers (runs actual inference).
|
|
For gRPC workers, falls back to standard health check.
|
|
"""
|
|
if self.mode == ConnectionMode.GRPC:
|
|
# For gRPC, use standard health check (no /health_generate equivalent)
|
|
return self._grpc_health_check(timeout)
|
|
|
|
try:
|
|
resp = httpx.get(f"{self.base_url}/health_generate", timeout=timeout)
|
|
return resp.status_code == 200
|
|
except (httpx.RequestError, httpx.TimeoutException):
|
|
return False
|
|
|
|
def _grpc_health_check(self, timeout: float = 5.0) -> bool:
|
|
"""Check health via gRPC health check protocol."""
|
|
try:
|
|
import grpc
|
|
from grpc_health.v1 import health_pb2, health_pb2_grpc
|
|
except ImportError as e:
|
|
logger.debug("gRPC libraries not available: %s", e)
|
|
return False
|
|
|
|
try:
|
|
channel = grpc.insecure_channel(f"{DEFAULT_HOST}:{self.port}")
|
|
try:
|
|
stub = health_pb2_grpc.HealthStub(channel)
|
|
request = health_pb2.HealthCheckRequest(service="")
|
|
response = stub.Check(request, timeout=timeout)
|
|
is_serving = response.status == health_pb2.HealthCheckResponse.SERVING
|
|
if is_serving:
|
|
logger.debug(
|
|
"gRPC health check passed for port %d (status: SERVING)",
|
|
self.port,
|
|
)
|
|
return is_serving
|
|
finally:
|
|
channel.close()
|
|
except grpc.RpcError as e:
|
|
# gRPC-specific errors (connection refused, deadline exceeded, etc.)
|
|
logger.debug(
|
|
"gRPC health check failed for port %d: %s",
|
|
self.port,
|
|
e.code() if hasattr(e, "code") else str(e),
|
|
)
|
|
return False
|
|
except Exception as e:
|
|
# Other errors
|
|
logger.debug(
|
|
"gRPC health check error for port %d: %s",
|
|
self.port,
|
|
str(e),
|
|
)
|
|
return False
|
|
|
|
def terminate(self, timeout: float = 10.0) -> None:
|
|
"""Terminate the model server process and all child processes.
|
|
|
|
Since workers are started with start_new_session=True, they run in their
|
|
own process group. We must kill the entire process group to ensure child
|
|
processes (e.g., TP workers) are also terminated and GPU memory is freed.
|
|
"""
|
|
if self.process.poll() is not None:
|
|
return # Already terminated
|
|
|
|
pid = self.process.pid
|
|
logger.info("Terminating %s (PID %d)", self.key, pid)
|
|
|
|
# Try graceful shutdown of the entire process group first
|
|
try:
|
|
pgid = os.getpgid(pid)
|
|
os.killpg(pgid, signal.SIGTERM)
|
|
except (ProcessLookupError, OSError) as e:
|
|
logger.debug("Could not send SIGTERM to process group: %s", e)
|
|
# Fall back to terminating just the main process
|
|
self.process.terminate()
|
|
|
|
try:
|
|
self.process.wait(timeout=timeout)
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning("%s did not terminate, killing process group", self.key)
|
|
# Force kill the entire process group
|
|
try:
|
|
pgid = os.getpgid(pid)
|
|
os.killpg(pgid, signal.SIGKILL)
|
|
except (ProcessLookupError, OSError) as e:
|
|
logger.debug("Could not send SIGKILL to process group: %s", e)
|
|
self.process.kill()
|
|
|
|
try:
|
|
self.process.wait(timeout=5) # Brief timeout after kill
|
|
except subprocess.TimeoutExpired:
|
|
logger.error("%s did not die after SIGKILL, abandoning", self.key)
|
|
|
|
|
|
class ModelPool:
|
|
"""Manages long-running SGLang worker processes across GPUs.
|
|
|
|
Workers are expensive to start (~30-60s due to model loading), so this pool
|
|
keeps them running and allows reuse across multiple tests. Routers can then
|
|
be launched cheaply (~1-2s) pointing to these workers.
|
|
|
|
Startup behavior:
|
|
- Workers are pre-launched at startup until GPUs are full
|
|
- When a test needs a model that isn't running, MRU model is evicted
|
|
(models just used are likely done, models not yet used are waiting)
|
|
- The needed model is then launched on-demand
|
|
|
|
Instance keys:
|
|
- Regular workers: "model_id:mode" (e.g., "llama-8b:http")
|
|
- PD workers: "model_id:mode:worker_type" (e.g., "llama-8b:http:prefill")
|
|
|
|
Limitations:
|
|
- Currently one worker instance per (model_id, mode) combination
|
|
- @pytest.mark.workers(count=n) duplicates URLs to router, not distinct workers
|
|
- For true multi-worker LB testing, extend to support multiple instances
|
|
|
|
Usage:
|
|
pool = ModelPool()
|
|
pool.startup(requirements=[("llama-8b", ConnectionMode.HTTP)])
|
|
instance = pool.get("llama-8b", "http") # Pre-launched or on-demand
|
|
"""
|
|
|
|
def __init__(self, allocator: GPUAllocator | None = None):
|
|
"""Initialize the model pool.
|
|
|
|
Args:
|
|
allocator: GPU allocator to use. If None, creates a new one.
|
|
"""
|
|
self.allocator = allocator or GPUAllocator()
|
|
self.instances: dict[str, ModelInstance] = {} # key = "model_id:mode"
|
|
self._startup_timeout = DEFAULT_STARTUP_TIMEOUT
|
|
self._lock = threading.RLock() # Protects instances dict
|
|
|
|
def startup(
|
|
self,
|
|
requirements: list[WorkerIdentity] | None = None,
|
|
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
|
|
) -> None:
|
|
"""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).
|
|
|
|
Thread-safe: Protected by internal lock.
|
|
|
|
Args:
|
|
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.
|
|
"""
|
|
with self._lock:
|
|
self._startup_unlocked(requirements, startup_timeout)
|
|
|
|
def _startup_unlocked(
|
|
self,
|
|
requirements: list[WorkerIdentity] | None = None,
|
|
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
|
|
) -> None:
|
|
"""Internal startup logic. Caller must hold _lock."""
|
|
self._startup_timeout = startup_timeout
|
|
|
|
if requirements is None:
|
|
requirements = [WorkerIdentity(DEFAULT_MODEL, ConnectionMode.HTTP)]
|
|
|
|
# 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 identity.mode not in LOCAL_MODES:
|
|
logger.warning(
|
|
"Invalid mode %s for %s, skipping", identity.mode, identity.model_id
|
|
)
|
|
continue
|
|
valid_requirements.append(identity)
|
|
|
|
if not valid_requirements:
|
|
logger.warning("No valid requirements to start")
|
|
return
|
|
|
|
logger.info(
|
|
"Starting model pool with %d workers: %s",
|
|
len(valid_requirements),
|
|
[str(r) for r in valid_requirements],
|
|
)
|
|
|
|
# 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)
|
|
|
|
deferred: list[str] = []
|
|
launched_count = 0
|
|
|
|
# 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
|
|
|
|
# Each prefill worker needs its own bootstrap port for PD communication
|
|
bootstrap_port = get_open_port() if identity.is_prefill else None
|
|
|
|
# Stagger launches to avoid resource contention during model loading
|
|
if launched_count > 0 and LAUNCH_STAGGER_DELAY > 0:
|
|
logger.info(
|
|
"Staggering launch by %ds to reduce resource contention",
|
|
LAUNCH_STAGGER_DELAY,
|
|
)
|
|
time.sleep(LAUNCH_STAGGER_DELAY)
|
|
|
|
# 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,
|
|
ib_device=(
|
|
ib_device if (identity.is_prefill or identity.is_decode) else None
|
|
),
|
|
instance_key=identity.key,
|
|
)
|
|
launched_count += 1
|
|
|
|
# Log deferred workers
|
|
if deferred:
|
|
logger.info(
|
|
"%d workers deferred for on-demand launch: %s",
|
|
len(deferred),
|
|
deferred,
|
|
)
|
|
|
|
# Wait for all launched models to be healthy
|
|
self._wait_all_healthy()
|
|
|
|
def _launch_model(
|
|
self,
|
|
model_id: str,
|
|
mode: ConnectionMode,
|
|
gpu_slot: GPUSlot | None = None,
|
|
worker_type: WorkerType = WorkerType.REGULAR,
|
|
bootstrap_port: int | None = None,
|
|
ib_device: str | None = None,
|
|
instance_key: str | None = None,
|
|
) -> ModelInstance:
|
|
"""Launch a model instance.
|
|
|
|
Args:
|
|
model_id: Model identifier from MODEL_SPECS.
|
|
mode: Connection mode (HTTP or GRPC).
|
|
gpu_slot: GPU slot assignment, or None for auto.
|
|
worker_type: Worker type (REGULAR, PREFILL, or DECODE).
|
|
bootstrap_port: Bootstrap port for prefill workers in PD mode.
|
|
ib_device: InfiniBand device for PD disaggregation.
|
|
instance_key: Custom instance key, or None to auto-generate.
|
|
|
|
Returns:
|
|
The launched ModelInstance.
|
|
"""
|
|
spec = get_model_spec(model_id)
|
|
model_path = spec["model"]
|
|
tp_size = spec.get("tp", 1)
|
|
features = spec.get("features", [])
|
|
|
|
# Get port - use slot's port if available, otherwise find open port
|
|
port = gpu_slot.port if gpu_slot else get_open_port()
|
|
|
|
# Build environment
|
|
env = os.environ.copy()
|
|
if gpu_slot:
|
|
env["CUDA_VISIBLE_DEVICES"] = gpu_slot.cuda_visible_devices()
|
|
|
|
# Build command
|
|
cmd = [
|
|
"python3",
|
|
"-m",
|
|
"sglang.launch_server",
|
|
"--model-path",
|
|
model_path,
|
|
"--host",
|
|
DEFAULT_HOST,
|
|
"--port",
|
|
str(port),
|
|
"--tp-size",
|
|
str(tp_size),
|
|
"--log-level",
|
|
"warning",
|
|
]
|
|
|
|
if mode == ConnectionMode.GRPC:
|
|
cmd.append("--grpc-mode")
|
|
|
|
# Embedding model flag
|
|
if "embedding" in features:
|
|
cmd.append("--is-embedding")
|
|
|
|
# PD disaggregation arguments
|
|
if worker_type == WorkerType.PREFILL:
|
|
cmd.extend(["--disaggregation-mode", "prefill"])
|
|
if bootstrap_port:
|
|
cmd.extend(["--disaggregation-bootstrap-port", str(bootstrap_port)])
|
|
if ib_device:
|
|
cmd.extend(["--disaggregation-ib-device", ib_device])
|
|
elif worker_type == WorkerType.DECODE:
|
|
cmd.extend(["--disaggregation-mode", "decode"])
|
|
# Base GPU ID 0 since CUDA_VISIBLE_DEVICES remaps the GPU
|
|
cmd.extend(["--base-gpu-id", "0"])
|
|
if ib_device:
|
|
cmd.extend(["--disaggregation-ib-device", ib_device])
|
|
|
|
# Additional worker args from model spec (e.g., --context-length)
|
|
worker_args = spec.get("worker_args", [])
|
|
if worker_args:
|
|
cmd.extend(worker_args)
|
|
|
|
# Build key based on worker type (or use custom key)
|
|
if instance_key:
|
|
key = instance_key
|
|
elif worker_type == WorkerType.REGULAR:
|
|
key = f"{model_id}:{mode.value}"
|
|
else:
|
|
key = f"{model_id}:{mode.value}:{worker_type.value}"
|
|
|
|
gpu_info = gpu_slot.gpu_ids if gpu_slot else "auto"
|
|
logger.info("Launching %s on GPUs %s port %d", key, gpu_info, port)
|
|
|
|
show_output = os.environ.get(ENV_SHOW_WORKER_LOGS, "0") == "1"
|
|
|
|
# Start the process
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
env=env,
|
|
stdout=None if show_output else subprocess.PIPE,
|
|
stderr=None if show_output else subprocess.PIPE,
|
|
start_new_session=True,
|
|
)
|
|
|
|
base_url = f"http://{DEFAULT_HOST}:{port}"
|
|
instance = ModelInstance(
|
|
model_id=model_id,
|
|
mode=mode,
|
|
model_path=model_path,
|
|
base_url=base_url,
|
|
port=port,
|
|
process=proc,
|
|
gpu_slot=gpu_slot,
|
|
key=key,
|
|
worker_type=worker_type,
|
|
bootstrap_port=bootstrap_port,
|
|
last_used=time.time(),
|
|
)
|
|
self.instances[key] = instance
|
|
return instance
|
|
|
|
def _wait_all_healthy(self) -> None:
|
|
"""Wait for all model instances to become healthy.
|
|
|
|
Only checks workers that haven't been marked healthy yet,
|
|
avoiding redundant health checks on already-verified workers.
|
|
"""
|
|
start_time = time.time()
|
|
# Only wait for workers that haven't been verified healthy yet
|
|
pending = {key for key, inst in self.instances.items() if not inst._healthy}
|
|
check_count = 0
|
|
|
|
if not pending:
|
|
logger.info("All workers already healthy, skipping health check")
|
|
return
|
|
|
|
logger.info(
|
|
"Waiting for %d workers to become healthy (timeout: %ds)...",
|
|
len(pending),
|
|
self._startup_timeout,
|
|
)
|
|
|
|
# Initial grace period to allow models to load before health checks
|
|
if INITIAL_GRACE_PERIOD > 0:
|
|
logger.info(
|
|
"Waiting %ds for initial model loading before health checks...",
|
|
INITIAL_GRACE_PERIOD,
|
|
)
|
|
time.sleep(INITIAL_GRACE_PERIOD)
|
|
|
|
while pending and (time.time() - start_time) < self._startup_timeout:
|
|
check_count += 1
|
|
elapsed = time.time() - start_time
|
|
|
|
for key in list(pending):
|
|
instance = self.instances[key]
|
|
|
|
# Check if process died
|
|
if not instance.is_alive():
|
|
logger.error(
|
|
"[%.1fs] %s (PID %d) died during startup",
|
|
elapsed,
|
|
key,
|
|
instance.process.pid,
|
|
)
|
|
# Read stderr for debugging (non-blocking to avoid hangs)
|
|
if instance.process.stderr:
|
|
try:
|
|
import os
|
|
import select
|
|
|
|
# Use select for non-blocking read with short timeout
|
|
# to avoid hanging if child processes keep stderr open
|
|
ready, _, _ = select.select(
|
|
[instance.process.stderr], [], [], 0.5
|
|
)
|
|
if ready:
|
|
# Use os.read with limited size instead of .read()
|
|
# which reads until EOF and can block if pipe stays open
|
|
fd = instance.process.stderr.fileno()
|
|
stderr = os.read(fd, 65536) # Read up to 64KB
|
|
if stderr:
|
|
logger.error(
|
|
"Stderr: %s",
|
|
stderr.decode(errors="replace")[-2000:],
|
|
)
|
|
except Exception as e:
|
|
logger.warning("Could not read stderr: %s", e)
|
|
# Evict dead instance and release GPUs
|
|
self._evict_instance(key)
|
|
pending.discard(key)
|
|
continue
|
|
|
|
# Check health
|
|
if instance.health_check():
|
|
logger.info(
|
|
"[%.1fs] %s is healthy at %s (router url: %s) (check #%d)",
|
|
elapsed,
|
|
key,
|
|
instance.base_url,
|
|
instance.worker_url,
|
|
check_count,
|
|
)
|
|
instance._healthy = True
|
|
pending.discard(key)
|
|
|
|
if pending:
|
|
# Log progress every 30 seconds
|
|
if check_count % 15 == 0: # ~30s at 2s interval
|
|
logger.info(
|
|
"[%.1fs] Still waiting for %d workers: %s",
|
|
elapsed,
|
|
len(pending),
|
|
list(pending),
|
|
)
|
|
time.sleep(HEALTH_CHECK_INTERVAL)
|
|
|
|
if pending:
|
|
elapsed = time.time() - start_time
|
|
logger.error(
|
|
"[%.1fs] Models failed to start within %ds: %s",
|
|
elapsed,
|
|
self._startup_timeout,
|
|
pending,
|
|
)
|
|
# Log stderr from failed workers for debugging
|
|
for key in pending:
|
|
instance = self.instances.get(key)
|
|
if instance and instance.process.stderr:
|
|
try:
|
|
import os
|
|
import select
|
|
|
|
# Use select for non-blocking read with short timeout
|
|
# to avoid hanging if worker is unresponsive
|
|
ready, _, _ = select.select(
|
|
[instance.process.stderr], [], [], 0.1
|
|
)
|
|
if ready:
|
|
# Use os.read with limited size instead of .read()
|
|
# which reads until EOF and can block if pipe stays open
|
|
fd = instance.process.stderr.fileno()
|
|
stderr = os.read(fd, 65536) # Read up to 64KB
|
|
if stderr:
|
|
logger.error(
|
|
"[%s] Last stderr output:\n%s",
|
|
key,
|
|
stderr.decode(errors="replace")[-3000:],
|
|
)
|
|
except Exception as e:
|
|
logger.error("[%s] Could not read stderr: %s", key, e)
|
|
# Terminate failed instances and release their GPUs
|
|
for key in pending:
|
|
self._evict_instance(key)
|
|
else:
|
|
elapsed = time.time() - start_time
|
|
logger.info(
|
|
"[%.1fs] All %d workers healthy after %d health checks",
|
|
elapsed,
|
|
len(self.instances),
|
|
check_count,
|
|
)
|
|
|
|
def get(
|
|
self,
|
|
model_id: str,
|
|
mode: ConnectionMode | str,
|
|
worker_type: WorkerType | str = WorkerType.REGULAR,
|
|
wait_for_gpus: bool = True,
|
|
gpu_wait_timeout: int = 300,
|
|
) -> ModelInstance:
|
|
"""Get a model instance by model_id, mode, and worker_type.
|
|
|
|
If the model is not running, it will be launched on-demand with MRU
|
|
eviction if GPU resources are constrained.
|
|
|
|
Thread-safe: Protected by internal lock. The returned instance has its
|
|
reference count incremented (via acquire()) to prevent eviction.
|
|
Caller MUST call release() on the instance when done.
|
|
|
|
Args:
|
|
model_id: The model ID (e.g., "llama-8b")
|
|
mode: The mode (ConnectionMode.HTTP or ConnectionMode.GRPC, or string)
|
|
worker_type: The worker type (REGULAR, PREFILL, DECODE). Defaults to REGULAR.
|
|
wait_for_gpus: If True, wait for GPUs to become available when all
|
|
are in use by other tests. Defaults to True.
|
|
gpu_wait_timeout: Max seconds to wait for GPUs (default 5 min).
|
|
|
|
Returns:
|
|
ModelInstance for the requested model/mode/worker_type (already acquired).
|
|
|
|
Raises:
|
|
RuntimeError: If worker process died, failed health check, or
|
|
timeout waiting for GPUs.
|
|
"""
|
|
deadline = time.time() + gpu_wait_timeout
|
|
poll_interval = 2.0 # seconds
|
|
|
|
while True:
|
|
with self._lock:
|
|
instance = self._get_unlocked(model_id, mode, worker_type)
|
|
if instance is not None:
|
|
# Acquire while holding lock to prevent race with eviction
|
|
instance.acquire()
|
|
return instance
|
|
|
|
# _get_unlocked returns None when GPUs unavailable after eviction
|
|
if not wait_for_gpus:
|
|
raise RuntimeError(
|
|
f"Cannot get {model_id}: GPUs unavailable and waiting disabled"
|
|
)
|
|
|
|
if time.time() >= deadline:
|
|
raise RuntimeError(
|
|
f"Timeout waiting for GPUs for {model_id} after {gpu_wait_timeout}s"
|
|
)
|
|
|
|
# Release lock while waiting so other tests can release workers
|
|
logger.info(
|
|
"All GPUs in use by other tests, waiting %.1fs for %s...",
|
|
poll_interval,
|
|
model_id,
|
|
)
|
|
time.sleep(poll_interval)
|
|
|
|
def _get_unlocked(
|
|
self,
|
|
model_id: str,
|
|
mode: ConnectionMode | str,
|
|
worker_type: WorkerType | str = WorkerType.REGULAR,
|
|
) -> ModelInstance | None:
|
|
"""Internal get logic. Caller must hold _lock.
|
|
|
|
Returns:
|
|
ModelInstance if successful, None if GPUs unavailable (signals retry).
|
|
|
|
Raises:
|
|
RuntimeError: If worker died or failed health check.
|
|
"""
|
|
# Accept both enum and string for convenience
|
|
if isinstance(mode, str):
|
|
mode = ConnectionMode(mode)
|
|
if isinstance(worker_type, str):
|
|
worker_type = WorkerType(worker_type)
|
|
|
|
if worker_type == WorkerType.REGULAR:
|
|
key = f"{model_id}:{mode.value}"
|
|
else:
|
|
key = f"{model_id}:{mode.value}:{worker_type.value}"
|
|
|
|
# Check if instance exists - if not, launch on-demand with eviction
|
|
if key not in self.instances:
|
|
logger.info(
|
|
"Model %s not running, launching on-demand with MRU eviction if needed",
|
|
key,
|
|
)
|
|
if not self._ensure_gpu_available(model_id):
|
|
# GPUs not available after eviction - signal retry
|
|
return None
|
|
|
|
# Allocate GPU slot for this model
|
|
spec = get_model_spec(model_id)
|
|
allocation_specs = {
|
|
key: {
|
|
"model": spec["model"],
|
|
"memory_gb": spec.get("memory_gb", 16),
|
|
"tp": spec.get("tp", 1),
|
|
}
|
|
}
|
|
slots = self.allocator.allocate_slots(allocation_specs)
|
|
if not slots:
|
|
raise RuntimeError(
|
|
f"Failed to allocate GPU slot for {model_id} after eviction"
|
|
)
|
|
gpu_slot = slots[0]
|
|
|
|
self._launch_model(model_id, mode, gpu_slot=gpu_slot)
|
|
self._wait_for_instance(key)
|
|
|
|
instance = self.instances[key]
|
|
|
|
# Note: last_used is updated in acquire() which should be called by fixtures
|
|
# to prevent eviction during test execution
|
|
|
|
# Verify worker is still alive and healthy
|
|
if not instance.is_alive():
|
|
raise RuntimeError(f"Worker {key} process died (was healthy at startup)")
|
|
|
|
if not instance.deep_health_check(timeout=30.0):
|
|
raise RuntimeError(
|
|
f"Worker {key} failed deep health check (health_generate) - "
|
|
"model may be stuck or crashed"
|
|
)
|
|
|
|
logger.info("Worker %s passed deep health check", key)
|
|
return instance
|
|
|
|
def _evict_for_gpus(
|
|
self,
|
|
required_gpus: int,
|
|
exclude_model_id: str | None = None,
|
|
exclude_mode: ConnectionMode | None = None,
|
|
exclude_worker_types: set[WorkerType] | None = None,
|
|
) -> None:
|
|
"""Evict models until we have enough GPUs available.
|
|
|
|
Uses MRU (most recently used) eviction strategy - evicts models that
|
|
were just used first, keeping models that haven't been used yet
|
|
(which are likely waiting for upcoming tests).
|
|
|
|
Args:
|
|
required_gpus: Number of GPUs needed.
|
|
exclude_model_id: Model ID to exclude from eviction.
|
|
exclude_mode: Connection mode to exclude from eviction (optional).
|
|
exclude_worker_types: Worker types to exclude from eviction.
|
|
If None, falls back to excluding by model_id only (backward compatible).
|
|
"""
|
|
available = self.allocator.available_gpus()
|
|
if len(available) >= required_gpus:
|
|
return # Already have enough
|
|
|
|
# Sort by last_used descending (MRU eviction) - evict most recently used first
|
|
# Store (dict_key, instance) tuples to preserve the actual key for eviction
|
|
# Note: Make a copy of items to avoid RuntimeError if dict is modified during iteration
|
|
evictable: list[tuple[str, ModelInstance]] = []
|
|
for dict_key, inst in list(self.instances.items()):
|
|
# Skip instances with active references (tests using them)
|
|
if inst.is_in_use:
|
|
logger.debug(
|
|
"Skipping eviction of %s - has active references", dict_key
|
|
)
|
|
continue
|
|
if exclude_worker_types is not None:
|
|
# Precise matching with worker types
|
|
# Must match model_id AND worker_type, mode is optional
|
|
if (
|
|
exclude_model_id is not None
|
|
and inst.model_id == exclude_model_id
|
|
and inst.worker_type in exclude_worker_types
|
|
):
|
|
# If mode is specified, also require mode match
|
|
if exclude_mode is None or inst.mode == exclude_mode:
|
|
continue
|
|
else:
|
|
# Backward compatible: exclude by model_id only
|
|
if exclude_model_id is not None and inst.model_id == exclude_model_id:
|
|
continue
|
|
evictable.append((dict_key, inst))
|
|
|
|
evictable.sort(key=lambda x: x[1].last_used, reverse=True)
|
|
|
|
freed_gpus = len(available)
|
|
for dict_key, inst in evictable:
|
|
if freed_gpus >= required_gpus:
|
|
break
|
|
|
|
logger.info("Evicting model %s (MRU) to free GPUs", dict_key)
|
|
self._evict_instance(dict_key)
|
|
if inst.gpu_slot:
|
|
freed_gpus += len(inst.gpu_slot.gpu_ids)
|
|
|
|
def _ensure_gpu_available(self, model_id: str) -> bool:
|
|
"""Ensure GPU is available for a model, evicting if needed.
|
|
|
|
Args:
|
|
model_id: Model ID that needs GPU resources.
|
|
|
|
Returns:
|
|
True if GPUs are available, False if not (all in use by other tests).
|
|
"""
|
|
spec = get_model_spec(model_id)
|
|
required_gpus = spec.get("tp", 1)
|
|
|
|
# Exclude REGULAR workers of same model from eviction (keep them)
|
|
# but allow evicting PD workers (PREFILL/DECODE) to free GPUs
|
|
self._evict_for_gpus(
|
|
required_gpus,
|
|
exclude_model_id=model_id,
|
|
exclude_worker_types={WorkerType.REGULAR},
|
|
)
|
|
|
|
available = self.allocator.available_gpus()
|
|
if len(available) < required_gpus:
|
|
logger.info(
|
|
"Cannot launch %s: need %d GPUs, only %d available after eviction "
|
|
"(all workers in use by other tests)",
|
|
model_id,
|
|
required_gpus,
|
|
len(available),
|
|
)
|
|
return False
|
|
return True
|
|
|
|
def _evict_instance(self, key: str) -> None:
|
|
"""Evict a model instance and free its resources.
|
|
|
|
Args:
|
|
key: Instance key to evict.
|
|
"""
|
|
if key not in self.instances:
|
|
return
|
|
|
|
instance = self.instances[key]
|
|
instance.terminate()
|
|
|
|
# Release GPU slot back to allocator
|
|
if instance.gpu_slot:
|
|
self.allocator.release_slot(instance.gpu_slot)
|
|
|
|
del self.instances[key]
|
|
logger.info("Evicted instance %s", key)
|
|
|
|
def _wait_for_instance(self, key: str, timeout: float | None = None) -> None:
|
|
"""Wait for a specific instance to become healthy.
|
|
|
|
Args:
|
|
key: Instance key to wait for.
|
|
timeout: Timeout in seconds. Defaults to _startup_timeout.
|
|
"""
|
|
if timeout is None:
|
|
timeout = self._startup_timeout
|
|
|
|
start_time = time.time()
|
|
instance = self.instances.get(key)
|
|
if not instance:
|
|
raise KeyError(f"Instance {key} not found")
|
|
|
|
while (time.time() - start_time) < timeout:
|
|
if not instance.is_alive():
|
|
raise RuntimeError(f"Worker {key} died during startup")
|
|
|
|
if instance.health_check():
|
|
logger.info("Instance %s is healthy", key)
|
|
instance._healthy = True
|
|
return
|
|
|
|
time.sleep(HEALTH_CHECK_INTERVAL)
|
|
|
|
raise TimeoutError(f"Instance {key} did not become healthy within {timeout}s")
|
|
|
|
def get_workers_by_type(
|
|
self, model_id: str, worker_type: WorkerType
|
|
) -> list[ModelInstance]:
|
|
"""Get all workers of a specific type for a model.
|
|
|
|
Thread-safe: Protected by internal lock. All returned instances have their
|
|
reference count incremented (via acquire()) to prevent eviction.
|
|
Caller MUST call release() on each instance when done.
|
|
|
|
Args:
|
|
model_id: The model ID.
|
|
worker_type: The worker type to filter by.
|
|
|
|
Returns:
|
|
List of matching ModelInstance objects (already acquired).
|
|
"""
|
|
with self._lock:
|
|
workers = [
|
|
inst
|
|
for inst in self.instances.values()
|
|
if inst.model_id == model_id and inst.worker_type == worker_type
|
|
]
|
|
# Acquire all while holding lock to prevent race with eviction
|
|
for worker in workers:
|
|
worker.acquire()
|
|
return workers
|
|
|
|
def launch_workers(
|
|
self,
|
|
workers: list[WorkerIdentity],
|
|
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
|
|
allow_eviction: bool = True,
|
|
wait_for_gpus: bool = True,
|
|
gpu_wait_timeout: int = 300,
|
|
) -> list[ModelInstance]:
|
|
"""Launch workers of any type.
|
|
|
|
This is the unified method for launching workers. It handles all worker
|
|
types (regular, prefill, decode) uniformly.
|
|
|
|
Thread-safe: Protected by internal lock.
|
|
|
|
Args:
|
|
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.
|
|
wait_for_gpus: If True, wait for GPUs to become available when all
|
|
are in use by other tests (with eviction enabled).
|
|
gpu_wait_timeout: Max seconds to wait for GPUs (default 5 min).
|
|
|
|
Returns:
|
|
List of launched ModelInstance objects.
|
|
"""
|
|
deadline = time.time() + gpu_wait_timeout
|
|
poll_interval = 2.0 # seconds
|
|
|
|
while True:
|
|
with self._lock:
|
|
result = self._launch_workers_unlocked(
|
|
workers, startup_timeout, allow_eviction
|
|
)
|
|
if result is not None:
|
|
return result
|
|
|
|
# _launch_workers_unlocked returns None when GPUs unavailable
|
|
# after eviction attempt (all workers in use by other tests)
|
|
if not wait_for_gpus or not allow_eviction:
|
|
return []
|
|
|
|
if time.time() >= deadline:
|
|
logger.warning(
|
|
"Timeout waiting for GPUs after %ds, giving up",
|
|
gpu_wait_timeout,
|
|
)
|
|
return []
|
|
|
|
# Release lock while waiting so other tests can release workers
|
|
logger.info(
|
|
"All GPUs in use by other tests, waiting %.1fs for availability...",
|
|
poll_interval,
|
|
)
|
|
time.sleep(poll_interval)
|
|
|
|
def _launch_workers_unlocked(
|
|
self,
|
|
workers: list[WorkerIdentity],
|
|
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
|
|
allow_eviction: bool = True,
|
|
) -> list[ModelInstance] | None:
|
|
"""Internal launch logic. Caller must hold _lock.
|
|
|
|
Returns:
|
|
List of launched instances, empty list if no valid workers,
|
|
or None if GPUs unavailable (signals caller to wait and retry).
|
|
"""
|
|
if not workers:
|
|
return []
|
|
|
|
self._startup_timeout = startup_timeout
|
|
|
|
# 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)
|
|
|
|
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) < total_gpus:
|
|
if allow_eviction:
|
|
logger.info(
|
|
"Need %d GPUs for %d workers, only %d available. Evicting...",
|
|
total_gpus,
|
|
len(valid_workers),
|
|
len(available),
|
|
)
|
|
self._evict_for_gpus(total_gpus)
|
|
|
|
# Check again after eviction
|
|
available = self.allocator.available_gpus()
|
|
if len(available) < total_gpus:
|
|
# Still not enough - all workers are in use by other tests
|
|
# Return None to signal caller to wait and retry
|
|
logger.info(
|
|
"Still need %d GPUs, only %d available after eviction. "
|
|
"All workers in use by other tests.",
|
|
total_gpus,
|
|
len(available),
|
|
)
|
|
return None
|
|
else:
|
|
logger.warning(
|
|
"Need %d GPUs, only %d available. Skipping launch.",
|
|
total_gpus,
|
|
len(available),
|
|
)
|
|
return []
|
|
|
|
# Build allocation specs
|
|
allocation_specs = {}
|
|
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": spec.get("tp", 1),
|
|
}
|
|
|
|
# Allocate GPU 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 {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
|
|
|
|
instances: list[ModelInstance] = []
|
|
for w in valid_workers:
|
|
# 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,
|
|
ib_device=ib_device if (w.is_prefill or w.is_decode) else None,
|
|
instance_key=w.key,
|
|
)
|
|
instances.append(instance)
|
|
|
|
self._wait_all_healthy()
|
|
return instances
|
|
|
|
def get_client(
|
|
self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP
|
|
) -> "openai.OpenAI":
|
|
"""Get OpenAI client for a specific model.
|
|
|
|
Args:
|
|
model_id: The model ID to get a client for.
|
|
mode: The mode (ConnectionMode.HTTP or ConnectionMode.GRPC). Defaults to HTTP.
|
|
|
|
Returns:
|
|
OpenAI client configured for this model.
|
|
"""
|
|
import openai
|
|
|
|
instance = self.get(model_id, mode)
|
|
return openai.OpenAI(
|
|
base_url=f"{instance.base_url}/v1",
|
|
api_key="not-used",
|
|
)
|
|
|
|
def get_base_url(
|
|
self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP
|
|
) -> str:
|
|
"""Get the base URL for a specific model."""
|
|
return self.get(model_id, mode).base_url
|
|
|
|
def shutdown(self) -> None:
|
|
"""Tear down all models.
|
|
|
|
Thread-safe: Protected by internal lock.
|
|
"""
|
|
with self._lock:
|
|
logger.info("Shutting down model pool (%d instances)", len(self.instances))
|
|
for instance in self.instances.values():
|
|
instance.terminate()
|
|
self.instances.clear()
|
|
|
|
def __enter__(self) -> "ModelPool":
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
self.shutdown()
|