ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

studio: revert -- always respect CUDA_VISIBLE_DEVICES in GPU memory query

Browse source
This commit is contained in:
Daniel Han 2026-03-15 07:59:50 +00:00
commit bb57236e29
1 changed files with 16 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

20
studio/backend/core/inference/llama_cpp.py
View file

@ -204,11 +204,12 @@ class LlamaCppBackend:
def _get_gpu_free_memory() -> list[tuple[int, int]]:
"""Query free memory per GPU via nvidia-smi.
Returns list of (gpu_index, free_mib) for ALL physical GPUs,
ignoring CUDA_VISIBLE_DEVICES. llama-server manages its own GPU
allocation via _select_gpus which sets CUDA_VISIBLE_DEVICES
explicitly for the subprocess.
Returns list of (gpu_index, free_mib) sorted by index.
Respects CUDA_VISIBLE_DEVICES if set.
Returns empty list if nvidia-smi is not available.
"""
import os
try:
result = subprocess.run(
[
@ -223,12 +224,23 @@ class LlamaCppBackend:
if result.returncode != 0:
return []
# Parse which GPUs are allowed by existing CUDA_VISIBLE_DEVICES
allowed = None
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if cvd is not None and cvd.strip():
try:
allowed = set(int(x.strip()) for x in cvd.split(","))
except ValueError:
pass # Non-numeric (e.g., "GPU-uuid"), ignore filter
gpus = []
for line in result.stdout.strip().splitlines():
parts = line.split(",")
if len(parts) == 2:
idx = int(parts[0].strip())
free_mib = int(parts[1].strip())
if allowed is not None and idx not in allowed:
continue
gpus.append((idx, free_mib))
return gpus
except Exception: