438 lines
15 KiB
Python
438 lines
15 KiB
Python
"""GPU detection and slot allocation for parallel test execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import socket
|
|
import threading
|
|
import time
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Try to import nvidia-ml-py for GPU detection
|
|
try:
|
|
import pynvml
|
|
|
|
NVML_AVAILABLE = True
|
|
except ImportError:
|
|
NVML_AVAILABLE = False
|
|
logger.debug("nvidia-ml-py not available, GPU detection will be limited")
|
|
|
|
|
|
@contextmanager
|
|
def nvml_context():
|
|
"""Context manager for NVML initialization/shutdown.
|
|
|
|
Usage:
|
|
with nvml_context():
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
|
...
|
|
"""
|
|
if not NVML_AVAILABLE:
|
|
yield
|
|
return
|
|
|
|
try:
|
|
pynvml.nvmlInit()
|
|
yield
|
|
finally:
|
|
pynvml.nvmlShutdown()
|
|
|
|
|
|
@dataclass
|
|
class GPUInfo:
|
|
"""Information about a single GPU."""
|
|
|
|
id: int
|
|
name: str
|
|
memory_mb: int
|
|
|
|
@property
|
|
def memory_gb(self) -> float:
|
|
return self.memory_mb / 1024
|
|
|
|
|
|
@dataclass
|
|
class GPUSlot:
|
|
"""A slot representing one or more GPUs allocated for a model."""
|
|
|
|
gpu_ids: list[int]
|
|
total_memory_mb: int
|
|
assigned_model: str | None = None
|
|
port: int | None = None
|
|
|
|
@property
|
|
def total_memory_gb(self) -> float:
|
|
return self.total_memory_mb / 1024
|
|
|
|
def cuda_visible_devices(self) -> str:
|
|
"""Return CUDA_VISIBLE_DEVICES string for this slot."""
|
|
return ",".join(str(g) for g in self.gpu_ids)
|
|
|
|
|
|
def get_open_port() -> int:
|
|
"""Get an available port by binding to port 0 and reading the assigned port."""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind(("", 0))
|
|
s.listen(1)
|
|
port = s.getsockname()[1]
|
|
return port
|
|
|
|
|
|
def get_physical_device_indices(devices: list[int]) -> list[int]:
|
|
"""Map logical device indices to physical indices based on CUDA_VISIBLE_DEVICES."""
|
|
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
if visible_devices is None:
|
|
return devices
|
|
|
|
visible_indices = [int(x) for x in visible_devices.split(",")]
|
|
index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
|
|
return [index_mapping[i] for i in devices if i in index_mapping]
|
|
|
|
|
|
def get_gpu_memory_usage(device_id: int) -> tuple[float, float]:
|
|
"""Get GPU memory usage in GB (used, total).
|
|
|
|
Args:
|
|
device_id: Physical GPU device ID
|
|
|
|
Returns:
|
|
Tuple of (used_gb, total_gb)
|
|
"""
|
|
if not NVML_AVAILABLE:
|
|
return (0.0, 0.0)
|
|
|
|
with nvml_context():
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
|
|
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
return (mem_info.used / (1024**3), mem_info.total / (1024**3))
|
|
|
|
|
|
def wait_for_gpu_memory_to_clear(
|
|
*,
|
|
devices: list[int],
|
|
threshold_bytes: int | None = None,
|
|
threshold_ratio: float | None = None,
|
|
timeout_s: float = 120,
|
|
) -> None:
|
|
"""Wait for GPU memory to be freed below a threshold.
|
|
|
|
Args:
|
|
devices: List of logical GPU device IDs to check
|
|
threshold_bytes: Memory threshold in bytes (used <= threshold)
|
|
threshold_ratio: Memory threshold as ratio (used/total <= ratio)
|
|
timeout_s: Timeout in seconds
|
|
|
|
Raises:
|
|
ValueError: If memory doesn't clear within timeout
|
|
"""
|
|
if not NVML_AVAILABLE:
|
|
logger.warning("nvidia-ml-py not available, skipping memory wait")
|
|
return
|
|
|
|
if threshold_bytes is None and threshold_ratio is None:
|
|
raise ValueError("Must specify threshold_bytes or threshold_ratio")
|
|
|
|
physical_devices = get_physical_device_indices(devices)
|
|
start_time = time.time()
|
|
|
|
with nvml_context():
|
|
while True:
|
|
output: dict[int, str] = {}
|
|
output_raw: dict[int, tuple[float, float]] = {}
|
|
|
|
for device in physical_devices:
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(device)
|
|
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
gb_used = mem_info.used / (1024**3)
|
|
gb_total = mem_info.total / (1024**3)
|
|
output_raw[device] = (gb_used, gb_total)
|
|
output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
|
|
|
|
logger.debug(
|
|
"GPU memory used/total (GiB): %s",
|
|
" ".join(f"{k}={v}" for k, v in output.items()),
|
|
)
|
|
|
|
if threshold_bytes is not None:
|
|
|
|
def is_free(used: float, total: float) -> bool:
|
|
return used <= threshold_bytes / (1024**3)
|
|
|
|
threshold_desc = f"{threshold_bytes / (1024**3):.1f} GiB"
|
|
else:
|
|
|
|
def is_free(used: float, total: float) -> bool:
|
|
return used / total <= threshold_ratio # type: ignore[operator]
|
|
|
|
threshold_desc = f"{threshold_ratio:.2%}" # type: ignore[str-format]
|
|
|
|
dur_s = time.time() - start_time
|
|
if all(is_free(used, total) for used, total in output_raw.values()):
|
|
logger.info(
|
|
"GPU memory cleared on devices %s (threshold=%s) in %.1fs",
|
|
devices,
|
|
threshold_desc,
|
|
dur_s,
|
|
)
|
|
return
|
|
|
|
if dur_s >= timeout_s:
|
|
raise ValueError(
|
|
f"GPU memory on devices {devices} not freed after {dur_s:.1f}s "
|
|
f"(threshold={threshold_desc})"
|
|
)
|
|
|
|
time.sleep(5)
|
|
|
|
|
|
class GPUAllocator:
|
|
"""Detects GPUs and assigns them to model slots using bin-packing."""
|
|
|
|
def __init__(self, gpus: list[GPUInfo] | None = None):
|
|
"""Initialize the allocator.
|
|
|
|
Args:
|
|
gpus: Optional list of GPUs. If None, auto-detects via nvidia-ml-py.
|
|
"""
|
|
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)."""
|
|
if not NVML_AVAILABLE:
|
|
logger.warning("nvidia-ml-py not available - no GPUs detected")
|
|
return []
|
|
|
|
# Check for CUDA_VISIBLE_DEVICES restriction
|
|
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
allowed_ids: set[int] | None = None
|
|
if visible_devices:
|
|
allowed_ids = set(int(x) for x in visible_devices.split(",") if x.strip())
|
|
|
|
try:
|
|
with nvml_context():
|
|
device_count = pynvml.nvmlDeviceGetCount()
|
|
|
|
gpus = []
|
|
for idx in range(device_count):
|
|
# Skip GPUs not in CUDA_VISIBLE_DEVICES if set
|
|
if allowed_ids is not None and idx not in allowed_ids:
|
|
continue
|
|
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
|
|
name = pynvml.nvmlDeviceGetName(handle)
|
|
# Handle bytes vs string return type (varies by pynvml version)
|
|
if isinstance(name, bytes):
|
|
name = name.decode("utf-8", errors="replace")
|
|
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
# Convert bytes to MB
|
|
memory_mb = mem_info.total // (1024 * 1024)
|
|
|
|
gpus.append(GPUInfo(idx, name, memory_mb))
|
|
|
|
logger.info("Detected %d GPUs: %s", len(gpus), [g.name for g in gpus])
|
|
return gpus
|
|
|
|
except pynvml.NVMLError as e:
|
|
logger.warning("NVML error during GPU detection: %s", e)
|
|
return []
|
|
except Exception as e:
|
|
logger.warning("Failed to detect GPUs: %s", e)
|
|
return []
|
|
|
|
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 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.
|
|
|
|
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
|
|
of sorting by memory size. Default False.
|
|
|
|
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 []
|
|
|
|
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 ordered_models:
|
|
memory_gb = spec.get("memory_gb", 16)
|
|
tp_size = spec.get("tp", 1)
|
|
|
|
# Find available GPUs (not used by any previous allocation)
|
|
available = [g for g in self.gpus if g.id not in self._used_gpus]
|
|
|
|
if tp_size == 1:
|
|
# Single GPU - find one with enough memory
|
|
for gpu in available:
|
|
if gpu.memory_gb >= memory_gb:
|
|
slot = GPUSlot(
|
|
gpu_ids=[gpu.id],
|
|
total_memory_mb=gpu.memory_mb,
|
|
assigned_model=model_id,
|
|
port=get_open_port(),
|
|
)
|
|
new_slots.append(slot)
|
|
self._used_gpus.add(gpu.id)
|
|
logger.info(
|
|
"Allocated GPU %d (%s, %.1fGB) for %s",
|
|
gpu.id,
|
|
gpu.name,
|
|
gpu.memory_gb,
|
|
model_id,
|
|
)
|
|
break
|
|
else:
|
|
logger.warning(
|
|
"No GPU with %.1fGB available for %s (used: %s)",
|
|
memory_gb,
|
|
model_id,
|
|
self._used_gpus,
|
|
)
|
|
else:
|
|
# Multi-GPU - find consecutive GPUs with enough total memory
|
|
# Sort available by ID for consecutive allocation
|
|
available_sorted = sorted(available, key=lambda g: g.id)
|
|
|
|
for i in range(len(available_sorted) - tp_size + 1):
|
|
candidate_gpus = available_sorted[i : i + tp_size]
|
|
total_mem = sum(g.memory_mb for g in candidate_gpus)
|
|
|
|
if total_mem >= memory_gb * 1024:
|
|
gpu_ids = [g.id for g in candidate_gpus]
|
|
slot = GPUSlot(
|
|
gpu_ids=gpu_ids,
|
|
total_memory_mb=total_mem,
|
|
assigned_model=model_id,
|
|
port=get_open_port(),
|
|
)
|
|
new_slots.append(slot)
|
|
self._used_gpus.update(gpu_ids)
|
|
logger.info(
|
|
"Allocated GPUs %s (%.1fGB total) for %s (tp=%d)",
|
|
gpu_ids,
|
|
total_mem / 1024,
|
|
model_id,
|
|
tp_size,
|
|
)
|
|
break
|
|
else:
|
|
logger.warning(
|
|
"No %d consecutive GPUs with %.1fGB available for %s (used: %s)",
|
|
tp_size,
|
|
memory_gb,
|
|
model_id,
|
|
self._used_gpus,
|
|
)
|
|
|
|
# Add new slots to existing slots list
|
|
self.slots.extend(new_slots)
|
|
return new_slots
|
|
|
|
def get_slot_for_model(self, model_id: str) -> GPUSlot | 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.
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
slot: The GPUSlot to release.
|
|
"""
|
|
self.release_gpus(slot.gpu_ids)
|
|
|
|
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.
|
|
"""
|
|
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.
|
|
|
|
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)
|