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

studio: ignore CUDA_VISIBLE_DEVICES in GPU memory query for llama-server

Browse source 
_get_gpu_free_memory was filtering by CUDA_VISIBLE_DEVICES, so with
CUDA_VISIBLE_DEVICES='0' set by the training env, llama-server only
saw 1 GPU and used --fit for CPU offloading instead of spreading
across all 8 GPUs.

Since llama-server manages its own GPU allocation (the _select_gpus
method picks GPUs and sets CUDA_VISIBLE_DEVICES for the subprocess),
the query must see ALL physical GPUs to make the right decision.
This commit is contained in:
Daniel Han 2026-03-15 07:58:20 +00:00
commit 5603ced75f
1 changed files with 4 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -204,12 +204,11 @@ 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) sorted by index.
Only returns GPUs that are allowed by CUDA_VISIBLE_DEVICES
(if set). Returns empty list if nvidia-smi is not available.
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.
"""
import os
try:
result = subprocess.run(
[
@ -224,23 +223,12 @@ 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: