"""Model pool for managing pre-loaded models across GPUs.""" from __future__ import annotations import logging import os import subprocess import time from dataclasses import dataclass 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, 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 class ModelInstance: """A running model instance.""" model_id: str mode: ConnectionMode model_path: str base_url: str port: int process: subprocess.Popen gpu_slot: GPUSlot | None 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}" @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.""" if self.process.poll() is not None: return # Already terminated logger.info("Terminating %s (PID %d)", self.key, self.process.pid) # Try graceful shutdown first self.process.terminate() try: self.process.wait(timeout=timeout) except subprocess.TimeoutExpired: logger.warning("%s did not terminate, killing", self.key) self.process.kill() self.process.wait() 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 def startup( self, requirements: list[tuple[str, ConnectionMode]] | None = None, startup_timeout: int = DEFAULT_STARTUP_TIMEOUT, ) -> None: """Start worker processes for the required models. 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. Args: requirements: List of (model_id, mode) tuples specifying what to start. mode is ConnectionMode.HTTP or ConnectionMode.GRPC. 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)] # 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) continue if mode not in LOCAL_MODES: logger.warning("Invalid mode %s for %s, skipping", mode, model_id) continue valid_requirements.append((model_id, mode)) if not valid_requirements: logger.warning("No valid requirements to start") return logger.info("Starting model pool with: %s", 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), } # Allocate GPU slots slots = self.allocator.allocate_slots(allocation_specs) # Track which models got slots launched_keys = set() 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: logger.info( "%d models deferred for on-demand launch: %s", len(deferred_keys), deferred_keys, ) # 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, ) -> 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. 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"]) if ib_device: cmd.extend(["--disaggregation-ib-device", ib_device]) # Build key based on worker type if 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, 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, ) 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 if instance.process.stderr: stderr = instance.process.stderr.read() if stderr: logger.error("Stderr: %s", stderr.decode()[-2000:]) 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, ) # Terminate failed instances for key in pending: self.instances[key].terminate() del self.instances[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, ) -> 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. 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. Returns: ModelInstance for the requested model/mode/worker_type. Raises: RuntimeError: If worker process 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, ) self._ensure_gpu_available(model_id) # 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] # Update last_used timestamp instance.last_used = time.time() # 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 ) -> 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 (test may need multiple modes of the same model). """ 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 # Exclude instances of the same model_id (test may need multiple modes) evictable = [ inst for inst in self.instances.values() if exclude_model_id is None or inst.model_id != exclude_model_id ] evictable.sort(key=lambda x: x.last_used, reverse=True) freed_gpus = len(available) for inst in evictable: if freed_gpus >= required_gpus: break logger.info("Evicting model %s (MRU) to free GPUs", inst.key) self._evict_instance(inst.key) if inst.gpu_slot: freed_gpus += len(inst.gpu_slot.gpu_ids) def _ensure_gpu_available(self, model_id: str) -> None: """Ensure GPU is available for a model, evicting if needed. Args: model_id: Model ID that needs GPU resources. Raises: RuntimeError: If not enough GPUs after eviction. """ spec = get_model_spec(model_id) required_gpus = spec.get("tp", 1) self._evict_for_gpus(required_gpus, exclude_model_id=model_id) available = self.allocator.available_gpus() if len(available) < required_gpus: raise RuntimeError( f"Cannot launch {model_id}: need {required_gpus} GPUs, " f"only {len(available)} available after eviction" ) 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. Args: model_id: The model ID. worker_type: The worker type to filter by. Returns: List of matching ModelInstance objects. """ return [ inst for inst in self.instances.values() if inst.model_id == model_id and inst.worker_type == worker_type ] 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), ) self._evict_for_gpus(required_gpus, exclude_model_id=model_id) 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, ) 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, ) 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": """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.""" 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()