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

[smg][ci] Add thread safety to ModelPool and GPUAllocator (#16674)

Browse Source
This commit is contained in:
Simo Lin
2026-01-07 13:25:41 -08:00
committed by GitHub
parent 0241e0460f
commit 6037267f5b
10 changed files with 534 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/e2e_test/infra/gpu_allocator.py
+49 -22
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import logging
import os
import socket
import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass
@@ -200,6 +201,7 @@ class GPUAllocator:
self.gpus = gpus if gpus is not None else self._detect_gpus()
self.slots: list[GPUSlot] = []
self._used_gpus: set[int] = set() # Track GPUs used across all allocations
self._lock = threading.RLock() # Protects slots and _used_gpus
def _detect_gpus(self) -> list[GPUInfo]:
"""Auto-detect available GPUs via nvidia-ml-py (NVML)."""
@@ -261,6 +263,8 @@ class GPUAllocator:
Note: This method tracks used GPUs across multiple calls, so subsequent
allocations will use different GPUs than previous ones.
Thread-safe: Protected by internal lock.
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
@@ -269,6 +273,13 @@ class GPUAllocator:
Returns:
List of GPUSlots with assigned models (only the newly allocated slots)
"""
with self._lock:
return self._allocate_slots_unlocked(model_specs, preserve_order)
def _allocate_slots_unlocked(
self, model_specs: dict[str, dict], preserve_order: bool = False
) -> list[GPUSlot]:
"""Internal allocation logic. Caller must hold _lock."""
if not self.gpus:
logger.warning("No GPUs available for allocation")
return []
@@ -362,23 +373,32 @@ class GPUAllocator:
return new_slots
def get_slot_for_model(self, model_id: str) -> GPUSlot | None:
"""Get the slot assigned to a specific model."""
for slot in self.slots:
if slot.assigned_model == model_id:
return slot
return None
"""Get the slot assigned to a specific model.
Thread-safe: Protected by internal lock.
"""
with self._lock:
for slot in self.slots:
if slot.assigned_model == model_id:
return slot
return None
def release_gpus(self, gpu_ids: list[int]) -> None:
"""Release GPUs back to the available pool.
Thread-safe: Protected by internal lock.
Args:
gpu_ids: List of GPU IDs to release.
"""
for gpu_id in gpu_ids:
self._used_gpus.discard(gpu_id)
# Remove slots that used these GPUs
self.slots = [s for s in self.slots if not any(g in gpu_ids for g in s.gpu_ids)]
logger.info("Released GPUs %s, now used: %s", gpu_ids, self._used_gpus)
with self._lock:
for gpu_id in gpu_ids:
self._used_gpus.discard(gpu_id)
# Remove slots that used these GPUs
self.slots = [
s for s in self.slots if not any(g in gpu_ids for g in s.gpu_ids)
]
logger.info("Released GPUs %s, now used: %s", gpu_ids, self._used_gpus)
def release_slot(self, slot: GPUSlot) -> None:
"""Release a GPU slot back to the available pool.
@@ -391,20 +411,27 @@ class GPUAllocator:
def available_gpus(self) -> list[int]:
"""Get list of available (unused) GPU IDs.
Thread-safe: Protected by internal lock.
Returns:
List of GPU IDs that are not currently allocated.
"""
return [g.id for g in self.gpus if g.id not in self._used_gpus]
with self._lock:
return [g.id for g in self.gpus if g.id not in self._used_gpus]
def summary(self) -> str:
"""Return a summary of GPU allocations."""
lines = ["GPU Allocation Summary:"]
lines.append(f" Total GPUs: {len(self.gpus)}")
lines.append(f" Used GPUs: {sorted(self._used_gpus)}")
lines.append(f" Allocated Slots: {len(self.slots)}")
for slot in self.slots:
lines.append(
f" - {slot.assigned_model}: GPUs {slot.gpu_ids} "
f"({slot.total_memory_gb:.1f}GB) port={slot.port}"
)
return "\n".join(lines)
"""Return a summary of GPU allocations.
Thread-safe: Protected by internal lock.
"""
with self._lock:
lines = ["GPU Allocation Summary:"]
lines.append(f" Total GPUs: {len(self.gpus)}")
lines.append(f" Used GPUs: {sorted(self._used_gpus)}")
lines.append(f" Allocated Slots: {len(self.slots)}")
for slot in self.slots:
lines.append(
f" - {slot.assigned_model}: GPUs {slot.gpu_ids} "
f"({slot.total_memory_gb:.1f}GB) port={slot.port}"
)
return "\n".join(lines)
sgl-model-gateway/e2e_test/infra/model_pool.py
+168 -20
View File
@@ -293,6 +293,7 @@ class ModelPool:
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,
@@ -309,11 +310,22 @@ class ModelPool:
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:
@@ -615,22 +627,76 @@ class ModelPool:
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.
ModelInstance for the requested model/mode/worker_type (already acquired).
Raises:
RuntimeError: If worker process died or failed health check.
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):
@@ -649,7 +715,9 @@ class ModelPool:
"Model %s not running, launching on-demand with MRU eviction if needed",
key,
)
self._ensure_gpu_available(model_id)
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)
@@ -752,14 +820,14 @@ class ModelPool:
if inst.gpu_slot:
freed_gpus += len(inst.gpu_slot.gpu_ids)
def _ensure_gpu_available(self, model_id: str) -> None:
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.
Raises:
RuntimeError: If not enough GPUs after eviction.
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)
@@ -774,10 +842,15 @@ class ModelPool:
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"
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.
@@ -831,38 +904,96 @@ class ModelPool:
) -> 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.
List of matching ModelInstance objects (already acquired).
"""
return [
inst
for inst in self.instances.values()
if inst.model_id == model_id and inst.worker_type == worker_type
]
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 []
@@ -899,6 +1030,19 @@ class ModelPool:
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.",
@@ -976,11 +1120,15 @@ class ModelPool:
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()
"""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