studio: fix P1 issues from PR review comments

1. n_gpu_layers kwarg: accept (and ignore) in load_model signature
   so callers like llm_assist.py don't get TypeError

2. mmproj exclusion: filter out mmproj files in _find_smallest_fitting_variant
   so fallback doesn't pick a tiny vision projection as the "model"

3. Shard preservation after fallback: re-discover shards for the
   fallback variant instead of resetting to empty list, so split
   GGUFs download all shards

4. Orphan cleanup safety: only kill llama-server processes whose
   cmdline contains ".unsloth/", avoiding termination of unrelated
   llama-server instances on the same machine

5. Path expression sanitization: validate repo_id format before using
   it in cache directory lookups
This commit is contained in:
Daniel Han 2026-03-15 07:40:39 +00:00
commit 1e9d19126b
2 changed files with 41 additions and 7 deletions

View file

@ -307,7 +307,10 @@ class LlamaCppBackend:
from huggingface_hub import get_paths_info, list_repo_files
files = list_repo_files(hf_repo, token = hf_token)
gguf_files = [f for f in files if f.endswith(".gguf")]
gguf_files = [
f for f in files
if f.endswith(".gguf") and "mmproj" not in f.lower()
]
if not gguf_files:
return None
@ -387,6 +390,7 @@ class LlamaCppBackend:
is_vision: bool = False,
n_ctx: int = 4096,
n_threads: Optional[int] = None,
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
) -> bool:
"""
Start llama-server with a GGUF model.
@ -516,12 +520,25 @@ class LlamaCppBackend:
hf_token,
)
if smaller:
fallback_file, fallback_size = smaller
logger.info(
f"Selected variant too large ({total_gb:.1f} GB), "
f"falling back to {smaller[0]} ({smaller[1] / (1024**3):.1f} GB)"
f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
)
gguf_filename = smaller[0]
gguf_extra_shards = []
gguf_filename = fallback_file
# Re-discover shards for the fallback variant
import re as _re
_shard_pat = _re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
_m = _shard_pat.match(gguf_filename)
_prefix = _m.group(1) if _m else None
if _prefix:
gguf_extra_shards = sorted(
f for f in all_gguf_files
if f.startswith(_prefix) and f != gguf_filename
and "mmproj" not in f.lower()
)
else:
gguf_extra_shards = []
else:
raise RuntimeError(
f"Not enough disk space to download any variant. "
@ -727,13 +744,18 @@ class LlamaCppBackend:
@staticmethod
def _kill_orphaned_servers():
"""Kill any orphaned llama-server processes from previous studio runs."""
"""Kill orphaned llama-server processes started by studio.
Only kills processes whose binary lives under ~/.unsloth/llama.cpp/
to avoid terminating unrelated llama-server instances on the machine.
"""
import os
import signal
try:
# Use pgrep with full command match to identify studio-managed servers
result = subprocess.run(
["pgrep", "-f", "llama-server"],
["pgrep", "-a", "-f", "llama-server"],
capture_output = True,
text = True,
timeout = 5,
@ -741,9 +763,16 @@ class LlamaCppBackend:
if result.returncode != 0:
return
for line in result.stdout.strip().splitlines():
pid = int(line.strip())
parts = line.strip().split(None, 1)
if len(parts) < 2:
continue
pid = int(parts[0])
cmdline = parts[1]
if pid == os.getpid():
continue
# Only kill if it's a studio-managed server (lives under .unsloth/)
if ".unsloth/" not in cmdline and "unsloth" not in cmdline.lower():
continue
try:
os.kill(pid, signal.SIGKILL)
logger.info(f"Killed orphaned llama-server process (pid={pid})")

View file

@ -561,8 +561,13 @@ async def get_gguf_variants(
# case-insensitive match.
cached_files: set = set()
try:
import re as _re
from huggingface_hub import constants as hf_constants
# Sanitize repo_id: must be "owner/name" with safe chars only
if not _re.fullmatch(r"[A-Za-z0-9._-]+/[A-Za-z0-9._-]+", repo_id):
raise ValueError(f"Invalid repo_id format: {repo_id}")
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
for entry in cache_dir.iterdir():