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

persist llama.cpp KV cache across idle auto-unload (slot save/restore) (#7204)

Browse source 
* Studio: persist llama.cpp KV cache across idle auto-unload (slot save/restore)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: address KV persistence review feedback

* Studio: guard KV restore on launch config

* Studio: fix KV resume purge race, fingerprint requested ctx, purge on disable

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: re-check idle/keep-KV settings after slot save, ns file identity

* Studio: shard-aware KV guard, honor user --no-cache-prompt, early save cap

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: honor LLAMA_ARG_CACHE_PROMPT env in slot-save guard

* Studio: derive prompt-cache state from final argv for slot saves

* Studio: stat LoRA/control-vector sidecars in KV restore fingerprint

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: parse csv and FNAME:SCALE sidecar syntax in KV fingerprint

* Studio: address codex review on idle-unload KV resume

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: harden slot-save cleanup, cap accounting, stale-KV guard, save timeout

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: treat unavailable KV estimate as full-cap for slot-save disk check

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Nilay 2026-07-20 12:42:42 +05:30 • committed by GitHub
commit 95d9970233
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 1644 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -23,6 +23,7 @@ import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path
from typing import (
Callable,
@ -453,6 +454,23 @@ def _hf_offline_if_dns_dead():
os.environ.pop("TRANSFORMERS_OFFLINE", None)
try:
_SLOT_SAVE_MAX_BYTES = int(os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30))
except ValueError:
_SLOT_SAVE_MAX_BYTES = 10 << 30
# The idle loop holds the lifecycle gate across a slot save, so a newly arriving
# request waits on the in-flight save's HTTP call. Bound it (was 120s) so a slow
# or stuck save can't stall the next request for minutes; best-effort save just
# falls back to a plain unload. Override with UNSLOTH_SLOT_SAVE_TIMEOUT (seconds).
try:
_SLOT_SAVE_HTTP_TIMEOUT = float(os.environ.get("UNSLOTH_SLOT_SAVE_TIMEOUT") or 30.0)
except ValueError:
_SLOT_SAVE_HTTP_TIMEOUT = 30.0
if _SLOT_SAVE_HTTP_TIMEOUT <= 0:
_SLOT_SAVE_HTTP_TIMEOUT = 30.0
def _swa_cache_path() -> Path:
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
base = Path(home) if home else Path.home() / ".unsloth" / "studio"
@ -2000,6 +2018,12 @@ class LlamaCppBackend:
self._llama_log_path: Optional[Path] = None
self._cancel_event = threading.Event()
self._api_key: Optional[str] = None
self._slot_save_dir: Optional[str] = None
self._slot_save_binary: Optional[tuple[str, int]] = None
# (gguf_identity, launch_fingerprint) snapshotted at load, so a later slot
# save can tell whether the model files were swapped on disk since load.
self._slot_loaded_identity: Optional[tuple] = None
self._prompt_cache_disabled: bool = False
# True once a probe has completed; cleared on transient failure.
self._is_audio: bool = False
self._audio_type: Optional[str] = None
@ -2638,6 +2662,7 @@ class LlamaCppBackend:
"supports_ctx_checkpoints": False,
"supports_no_cache_prompt": False,
"supports_metrics": False,
"supports_slot_save": False,
}
try:
mtime = int(Path(bin_path).stat().st_mtime)
@ -2658,6 +2683,7 @@ class LlamaCppBackend:
supports_ctx_checkpoints = False
supports_no_cache_prompt = False
supports_metrics = False
supports_slot_save = False
try:
probe_env = cls._llama_server_env_for_binary(bin_path)
result = subprocess.run(
@ -2756,6 +2782,7 @@ class LlamaCppBackend:
supports_ctx_checkpoints = _is_real("--ctx-checkpoints")
supports_no_cache_prompt = _is_real("--no-cache-prompt")
supports_metrics = _is_real("--metrics")
supports_slot_save = _is_real("--slot-save-path")
except (OSError, subprocess.SubprocessError) as exc:
logger.debug(f"llama-server --help probe failed: {exc}")
@ -2773,6 +2800,7 @@ class LlamaCppBackend:
"supports_ctx_checkpoints": supports_ctx_checkpoints,
"supports_no_cache_prompt": supports_no_cache_prompt,
"supports_metrics": supports_metrics,
"supports_slot_save": supports_slot_save,
}
cls._capability_cache[cache_key] = info
return info
@ -7332,6 +7360,26 @@ class LlamaCppBackend:
# when the binary advertises it (older/custom binaries may not).
if server_caps.get("supports_metrics"):
cmd.append("--metrics")
self._slot_save_dir = None
self._slot_save_binary = None
self._prompt_cache_disabled = False
if server_caps.get("supports_slot_save"):
try:
from utils.paths.storage_roots import ( # noqa: WPS433
llama_slot_cache_root,
)
slot_dir = llama_slot_cache_root()
slot_dir.mkdir(parents = True, exist_ok = True)
# Saved KV encodes chat content; keep it from other local users.
with contextlib.suppress(OSError):
os.chmod(slot_dir, 0o700)
cmd.extend(["--slot-save-path", str(slot_dir)])
self._slot_save_dir = str(slot_dir)
self._slot_save_binary = (binary, Path(binary).stat().st_mtime_ns)
except OSError:
self._slot_save_dir = None
self._slot_save_binary = None
cmd.extend(
self._ctx_integrity_flags(
n_parallel,
@ -7529,6 +7577,7 @@ class LlamaCppBackend:
unsupported_cache_flags.append("--ctx-checkpoints")
if server_caps.get("supports_no_cache_prompt"):
cmd.append("--no-cache-prompt")
self._prompt_cache_disabled = True
else:
unsupported_cache_flags.append("--no-cache-prompt")
if unsupported_cache_flags:
@ -8105,6 +8154,15 @@ class LlamaCppBackend:
if not self._healthy:
return False
# Snapshot the files the server actually loaded. If a GGUF shard or a
# LoRA/control-vector sidecar is swapped on disk afterwards while the
# old weights stay mapped, save_slots_for_resume() compares against
# this and refuses to persist KV that a reload could misapply.
if self._slot_save_dir:
self._slot_loaded_identity = (
self._gguf_file_identity(self._gguf_path),
self._slot_launch_fingerprint(),
)
return True
def _build_speculative_flags(
@ -8690,6 +8748,10 @@ class LlamaCppBackend:
self._effective_context_length = None
self._max_context_length = None
self._reset_effective_parallel_slots()
self._slot_save_dir = None
self._slot_save_binary = None
self._slot_loaded_identity = None
self._prompt_cache_disabled = False
self._chat_template = None
self._chat_template_override = None
self._supports_reasoning = False
@ -9216,6 +9278,237 @@ class LlamaCppBackend:
return False
return True
def _slot_launch_fingerprint(self) -> tuple:
# KV validity keys on extra args, stat'd sidecar weights, effective ctx.
sidecars = []
for path in self._sidecar_weight_files():
try:
st = os.stat(path)
sidecars.append((path, st.st_size, st.st_mtime_ns))
except OSError:
sidecars.append((path, None, None))
return (
tuple(self._extra_args or ()),
tuple(sidecars),
self._requested_n_ctx,
self._effective_context_length,
getattr(self, "_cache_type_kv", None),
self.effective_parallel_slots,
)
def _gguf_file_identity(self, path) -> Optional[tuple]:
# (size, mtime_ns) per shard: a split GGUF keys KV validity on every sibling.
p = Path(path)
paths = [p]
m = _SHARD_FULL_RE.match(p.name)
if m:
prefix, _first, total = m.groups()
paths = [
p.with_name(f"{prefix}-{i:05d}-of-{total}{p.suffix}")
for i in range(1, int(total) + 1)
]
try:
return tuple((sp.stat().st_size, sp.stat().st_mtime_ns) for sp in paths)
except OSError:
return None
_SIDECAR_WEIGHT_FLAGS = (
"--lora",
"--lora-scaled",
"--control-vector",
"--control-vector-scaled",
)
def _sidecar_weight_files(self) -> list[str]:
# llama.cpp: comma-separated paths, FNAME:SCALE on -scaled (older builds: FNAME SCALE).
args = [str(a).strip() for a in (self._extra_args or ())]
files: list[str] = []
for i, arg in enumerate(args):
flag, sep, inline = arg.partition("=")
if flag not in self._SIDECAR_WEIGHT_FLAGS:
continue
operand = inline if sep else (args[i + 1] if i + 1 < len(args) else "")
if not operand:
continue
candidates = [operand]
pieces = [p for p in operand.split(",") if p]
if len(pieces) > 1:
candidates.extend(pieces)
if flag.endswith("-scaled"):
for item in list(candidates):
# ":<number>" tail is a scale; rpartition spares drive letters.
head, colon, tail = item.rpartition(":")
if not (colon and head):
continue
try:
float(tail)
except ValueError:
continue
candidates.append(head)
for cand in candidates:
if cand not in files:
files.append(cand)
return files
def _prompt_cache_off(self) -> bool:
# Caching off makes restores useless; last prompt-cache flag wins, env only when unset.
last = None
for arg in self._extra_args or ():
flag = arg.strip().split("=", 1)[0]
if flag in ("--cache-prompt", "--no-cache-prompt"):
last = flag
if last is not None:
return last == "--no-cache-prompt"
if self._prompt_cache_disabled:
return True
if os.environ.get("LLAMA_ARG_NO_CACHE_PROMPT") is not None:
return True
env = (os.environ.get("LLAMA_ARG_CACHE_PROMPT") or "").strip().lower()
return env in {"off", "disabled", "false", "0"}
def save_slots_for_resume(
self, should_abort: Optional[Callable[[], bool]] = None
) -> Optional[dict]:
if (
not self.is_loaded
or not self._slot_save_dir
or not self._gguf_path
or self._prompt_cache_off()
):
return None
save_dir = Path(self._slot_save_dir)
gguf_stat = self._gguf_file_identity(self._gguf_path)
if gguf_stat is None:
return None
launch = self._slot_launch_fingerprint()
# If the GGUF or a sidecar was swapped on disk while the original weights
# stayed mapped, the live KV belongs to the old weights but a reload would
# load the new file. Persisting it would let restore misapply stale KV.
if self._slot_loaded_identity is not None and self._slot_loaded_identity != (
gguf_stat,
launch,
):
logger.debug("Skipping slot save: model files changed on disk since load")
return None
try:
estimate = self._estimate_kv_cache_bytes(
self._effective_context_length or self._context_length or 0,
self._cache_type_kv,
n_parallel = self.effective_parallel_slots,
)
# Skip before writing anything when the estimate alone blows the cap,
# rather than fully writing a slot and discarding it afterwards.
if estimate > _SLOT_SAVE_MAX_BYTES:
logger.debug(
"Skipping slot save: estimated %d bytes exceeds cap %d",
estimate,
_SLOT_SAVE_MAX_BYTES,
)
return None
# A 0 estimate means metadata was insufficient, not a zero-byte cache:
# a slot can still be many GiB, so demand room for the whole cap before
# trusting the post-write check.
required = (estimate if estimate > 0 else _SLOT_SAVE_MAX_BYTES) + (1 << 30)
if shutil.disk_usage(save_dir).free < required:
logger.debug("Skipping slot save: insufficient free disk")
return None
except Exception:
pass
token = uuid.uuid4().hex[:8]
entries: list[dict] = []
total_bytes = 0
for slot in range(self.effective_parallel_slots):
# A request pending mid-save waits on the gate; stop wasting its time.
if should_abort is not None and should_abort():
break
filename = f"resume-{token}-slot{slot}.bin"
path = save_dir / filename
try:
resp = httpx.post(
f"{self.base_url}/slots/{slot}",
params = {"action": "save"},
json = {"filename": filename},
headers = self._auth_headers,
timeout = _SLOT_SAVE_HTTP_TIMEOUT,
trust_env = False,
)
except Exception as e:
logger.debug(f"slot {slot} save failed: {e}")
with contextlib.suppress(OSError):
path.unlink()
break
if resp.status_code != 200:
logger.debug(f"slot {slot} save returned HTTP {resp.status_code}")
with contextlib.suppress(OSError):
path.unlink()
continue
try:
body = resp.json()
if not isinstance(body, dict):
raise ValueError("slot save response was not a JSON object")
n_saved = int(body.get("n_saved") or 0)
except Exception as e:
# A 200 that still wrote a file but returns a malformed body must
# clean up like the transport/HTTP error paths above, or the file
# (which holds chat KV) is orphaned until the next startup sweep.
logger.debug(f"slot {slot} save returned an invalid response: {e}")
with contextlib.suppress(OSError):
path.unlink()
continue
if n_saved <= 0:
with contextlib.suppress(OSError):
path.unlink()
continue
# Account by the bytes actually on disk, not the server-reported
# count, so the cap holds even if a custom binary under-reports.
try:
n_written = path.stat().st_size
except OSError:
n_written = 0
total_bytes += n_written
entries.append({"id": slot, "filename": filename, "n_saved": n_saved})
if total_bytes > _SLOT_SAVE_MAX_BYTES:
break # already over the cap; the discard below cleans up
if not entries:
return None
if total_bytes > _SLOT_SAVE_MAX_BYTES:
logger.debug(
"Discarding slot save: %d bytes exceeds cap %d",
total_bytes,
_SLOT_SAVE_MAX_BYTES,
)
for entry in entries:
with contextlib.suppress(OSError):
(save_dir / entry["filename"]).unlink()
return None
return {
"dir": self._slot_save_dir,
"binary": self._slot_save_binary,
"gguf": str(self._gguf_path),
"gguf_stat": gguf_stat,
"launch": launch,
"slots": entries,
}
def restore_slots_for_resume(self, manifest: dict) -> None:
if not self.is_loaded or not self._slot_save_dir:
return
for entry in manifest.get("slots") or []:
try:
resp = httpx.post(
f"{self.base_url}/slots/{int(entry['id'])}",
params = {"action": "restore"},
json = {"filename": str(entry["filename"])},
headers = self._auth_headers,
timeout = _SLOT_SAVE_HTTP_TIMEOUT,
trust_env = False,
)
except Exception as e:
logger.debug(f"slot restore failed: {e}")
break
if resp.status_code != 200:
logger.debug(f"slot {entry.get('id')} restore returned HTTP {resp.status_code}")
def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool:
"""Schedule one background reload without MTP after a mid-generation death.

138
studio/backend/core/inference/llama_keepwarm.py
View file

@ -15,6 +15,7 @@ import asyncio
import contextlib
import threading
import time
from pathlib import Path
from loggers import get_logger
@ -30,6 +31,8 @@ _last_active = time.monotonic()
# otherwise 503 against an empty backend can reload it (set on unload, cleared on
# reload). Storing the quant means the reload restores the exact freed variant.
_last_unloaded_model = None
# Slot KV manifest saved by the idle unload; whoever pops it owns deleting its files.
_kv_resume = None
# Guards inflight bumps against the idle-check-then-unload race, and blocks new
# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is
# shared across every event loop in the process, so a per-loop gate would let a
@ -161,11 +164,17 @@ def inference_lifecycle_gate():
return _unload_gate()
def note_model_loaded() -> None:
"""Record a successful GGUF load: stamp activity and drop any reload stash so
a manual load clears it synchronously, not only on the next idle poll."""
def note_model_loaded(backend = None) -> None:
"""Stamp activity and synchronously drop any reload stash."""
_note_activity()
resume = take_kv_resume()
_set_last_unloaded(None)
if resume is None:
return
if backend is not None:
restore_kv_resume(backend, resume)
else:
_delete_resume_files(resume)
def note_model_unloaded() -> None:
@ -182,9 +191,81 @@ def get_last_unloaded_model():
def _set_last_unloaded(value) -> None:
global _last_unloaded_model
global _last_unloaded_model, _kv_resume
stale = None
with _lock:
_last_unloaded_model = value
if value is None and _kv_resume is not None:
stale, _kv_resume = _kv_resume, None
if stale:
_delete_resume_files(stale)
def _delete_resume_files(manifest) -> None:
try:
base = Path(manifest.get("dir") or "")
for entry in manifest.get("slots") or []:
with contextlib.suppress(OSError):
(base / str(entry.get("filename"))).unlink()
except Exception:
pass
def _set_kv_resume(value) -> None:
global _kv_resume
stale = None
with _lock:
if _kv_resume is not None and _kv_resume is not value:
stale = _kv_resume
_kv_resume = value
if stale:
_delete_resume_files(stale)
def take_kv_resume():
global _kv_resume
with _lock:
manifest, _kv_resume = _kv_resume, None
return manifest
def purge_kv_resume() -> None:
resume = take_kv_resume()
if resume:
_delete_resume_files(resume)
def restore_kv_resume(backend, manifest) -> None:
try:
gguf = manifest.get("gguf")
binary = manifest.get("binary")
current = getattr(backend, "_gguf_path", None)
same_gguf = bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve()
if same_gguf:
# Same path is not enough: shards may have been rewritten meanwhile.
identity = getattr(backend, "_gguf_file_identity", None)
same_gguf = callable(identity) and identity(current) == manifest.get("gguf_stat")
if same_gguf:
# Nor the same file: launch overrides can invalidate KV numerics.
fingerprint = getattr(backend, "_slot_launch_fingerprint", None)
same_gguf = callable(fingerprint) and manifest.get("launch") == fingerprint()
if same_gguf and binary and binary == getattr(backend, "_slot_save_binary", None):
logger.info("Restoring saved slot KV onto the reloaded model")
backend.restore_slots_for_resume(manifest)
except Exception as exc:
logger.debug("slot restore after reload failed: %s", exc)
finally:
_delete_resume_files(manifest)
def sweep_slot_save_dir() -> None:
try:
from utils.paths.storage_roots import llama_slot_cache_root
for path in llama_slot_cache_root().glob("resume-*.bin"):
with contextlib.suppress(OSError):
path.unlink()
except Exception:
pass
class LlamaKeepWarmMiddleware:
@ -266,7 +347,10 @@ def _loaded_identity(backend):
async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
"""Unload the loaded GGUF once idle past the configured TTL. Inert when off."""
from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds
from utils.openai_auto_switch_settings import (
get_auto_unload_idle_seconds,
get_auto_unload_keep_kv,
)
seen_model = None
while True:
@ -281,17 +365,47 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
# Track by (id, variant): a (re)loaded model -- including the same repo
# at a different quant -- counts as activity so it survives one TTL
# before its first request (loads bypass the activity middleware).
current = _loaded_identity(backend)
if current != seen_model:
seen_model = current
if current is not None:
_note_activity()
_set_last_unloaded(None) # a model is loaded; drop stale stash
async with _unload_gate():
# Purging the stash mid-reload would race the restore.
current = _loaded_identity(backend)
if current != seen_model:
seen_model = current
if current is not None:
_note_activity()
_set_last_unloaded(None) # a model is loaded; drop stale stash
if backend.is_loaded and _is_idle(ttl):
freed = _loaded_identity(backend)
await asyncio.to_thread(backend.unload_model)
manifest = None
if get_auto_unload_keep_kv():
try:
manifest = await asyncio.to_thread(
backend.save_slots_for_resume,
lambda: not _is_idle(ttl),
)
except Exception as exc:
logger.debug("slot save before idle unload failed: %s", exc)
# Re-read settings: the save can outlive a settings change.
ttl = get_auto_unload_idle_seconds()
if ttl <= 0 or not _is_idle(ttl):
if manifest:
_delete_resume_files(manifest)
continue
if manifest and not get_auto_unload_keep_kv():
_delete_resume_files(manifest)
manifest = None
try:
await asyncio.to_thread(backend.unload_model)
except Exception:
# Failed unload means nothing will stash the manifest.
if manifest:
_delete_resume_files(manifest)
raise
_set_last_unloaded(freed) # let an alias request reload it
if manifest and freed:
_set_kv_resume({"identity": freed, **manifest})
logger.info("Idle auto-unload: saved slot KV for restore on reload")
elif manifest:
_delete_resume_files(manifest)
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)
seen_model = None
except Exception as exc:

2
studio/backend/core/inference/llama_server_args.py
View file

@ -70,6 +70,8 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# llama-server's own built-in tools flag would silently stack on top of
# Unsloth's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
# Slot-state dir: Studio owns it for KV persistence across idle unload.
frozenset({"--slot-save-path"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)

3
studio/backend/main.py
View file

@ -547,8 +547,9 @@ async def lifespan(app: FastAPI):
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
# Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set).
from core.inference.llama_keepwarm import idle_unload_loop
from core.inference.llama_keepwarm import idle_unload_loop, sweep_slot_save_dir
sweep_slot_save_dir()
app.state.idle_unload_task = asyncio.create_task(idle_unload_loop())
# Initialize RSA key pair for API key encryption (external providers).

2
studio/backend/routes/inference.py
View file

@ -4710,7 +4710,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Clear any idle-unload reload stash now, not only on the next poll.
from core.inference.llama_keepwarm import note_model_loaded
note_model_loaded()
await asyncio.to_thread(note_model_loaded, llama_backend)
# A plain load advertises its own identifier; auto-switch overwrites
# this with the repo id right after _load_model_impl returns.
llama_backend._openai_advertised_id = None

21
studio/backend/routes/settings.py
View file

@ -36,9 +36,10 @@ from utils.helper_precache_settings import (
)
from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents
from utils.openai_auto_switch_settings import (
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS,
DEFAULT_AUTO_UNLOAD_KEEP_KV,
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
get_auto_unload_idle_seconds,
get_auto_unload_keep_kv,
get_model_overrides,
get_openai_auto_switch_enabled,
get_stored_auto_unload_idle_seconds,
@ -90,7 +91,9 @@ class HelperPrecacheResponse(BaseModel):
class OpenAIAutoSwitchPayload(BaseModel):
enabled: bool
auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0)
# None leaves the stored value untouched (partial updates can't clobber it).
auto_unload_idle_seconds: Optional[int] = Field(default = None, ge = 0)
auto_unload_keep_kv: Optional[bool] = None
class OpenAIAutoSwitchResponse(BaseModel):
@ -101,6 +104,7 @@ class OpenAIAutoSwitchResponse(BaseModel):
# UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled
# is false, so the UI can show idle-unload as active instead of "needs enable".
idle_unload_active: bool = False
auto_unload_keep_kv: bool = DEFAULT_AUTO_UNLOAD_KEEP_KV
class ModelOverridePayload(BaseModel):
@ -198,6 +202,7 @@ def get_openai_auto_switch(
enabled = get_openai_auto_switch_enabled(),
auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(),
idle_unload_active = get_auto_unload_idle_seconds() > 0,
auto_unload_keep_kv = get_auto_unload_keep_kv(),
)
@ -206,8 +211,8 @@ def update_openai_auto_switch(
payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject)
) -> OpenAIAutoSwitchResponse:
try:
enabled, idle_seconds = set_openai_auto_switch(
payload.enabled, payload.auto_unload_idle_seconds
enabled, idle_seconds, keep_kv = set_openai_auto_switch(
payload.enabled, payload.auto_unload_idle_seconds, payload.auto_unload_keep_kv
)
except ValueError as exc:
raise log_and_http_error(
@ -217,10 +222,16 @@ def update_openai_auto_switch(
event = "settings.update_openai_auto_switch_failed",
log = logger,
) from exc
idle_unload_active = get_auto_unload_idle_seconds() > 0
if not keep_kv or not idle_unload_active:
# Keep-KV off or idle unload disabled: drop already-saved chat context too.
from core.inference.llama_keepwarm import purge_kv_resume
purge_kv_resume()
return OpenAIAutoSwitchResponse(
enabled = enabled,
auto_unload_idle_seconds = idle_seconds,
idle_unload_active = get_auto_unload_idle_seconds() > 0,
idle_unload_active = idle_unload_active,
auto_unload_keep_kv = keep_kv,
)

19
studio/backend/tests/test_llama_cpp_mtp_detection.py
View file

@ -741,6 +741,25 @@ def test_probe_reports_windows_cache_flags_absent_for_older_binary(tmp_path):
assert caps["supports_no_cache_prompt"] is False
@_NEEDS_BASH
def test_probe_detects_slot_save_path(tmp_path):
fake = _make_fake_llama_server(
tmp_path / "llama-server",
"--slot-save-path PATH path to save slot kv cache\n--threads N\n",
)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["supports_slot_save"] is True
@_NEEDS_BASH
def test_probe_reports_slot_save_absent_for_older_binary(tmp_path):
fake = _make_fake_llama_server(tmp_path / "llama-server", "--threads N\n")
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["supports_slot_save"] is False
def test_build_ngram_mod_flags_new():
flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"})
assert flags == [

494
studio/backend/tests/test_llama_cpp_slot_resume.py Normal file
View file

@ -0,0 +1,494 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import os
from types import SimpleNamespace
import core.inference.llama_cpp as llama_cpp
from core.inference.llama_cpp import LlamaCppBackend
def _resume_backend(tmp_path, n_slots = 1):
backend = LlamaCppBackend()
backend._healthy = True
# No-op lifecycle methods so the atexit cleanup can kill the fake quietly.
backend._process = SimpleNamespace(
poll = lambda: None,
terminate = lambda: None,
wait = lambda *a, **k: 0,
kill = lambda: None,
pid = 0,
)
backend._port = 8081
backend._slot_save_dir = str(tmp_path)
backend._slot_save_binary = ("/bin/llama-server", 1)
(tmp_path / "model.gguf").write_bytes(b"gguf")
backend._gguf_path = str(tmp_path / "model.gguf")
backend._effective_parallel_slots = n_slots
backend._estimate_kv_cache_bytes = lambda *a, **k: 0
return backend
def _fake_disk(monkeypatch, free = 1 << 40):
monkeypatch.setattr(llama_cpp.shutil, "disk_usage", lambda _p: SimpleNamespace(free = free))
class _Resp:
def __init__(
self,
status_code = 200,
body = None,
):
self.status_code = status_code
self._body = body or {}
def json(self):
return self._body
def test_save_returns_none_when_slot_save_disabled(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._slot_save_dir = None
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_prompt_cache_disabled(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._prompt_cache_disabled = True
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_insufficient_free_disk(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40
_fake_disk(monkeypatch, free = 1 << 20)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_collects_manifest_across_slots(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 2)
_fake_disk(monkeypatch)
calls = []
def fake_post(url, **kwargs):
calls.append((url, kwargs["params"], kwargs["json"]))
return _Resp(200, {"n_saved": 40, "n_written": 100})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
manifest = backend.save_slots_for_resume()
assert manifest is not None
assert manifest["dir"] == str(tmp_path)
assert manifest["binary"] == ("/bin/llama-server", 1)
assert manifest["gguf"] == str(tmp_path / "model.gguf")
st = os.stat(manifest["gguf"])
assert manifest["gguf_stat"] == ((st.st_size, st.st_mtime_ns),)
assert manifest["launch"] == backend._slot_launch_fingerprint()
assert [e["id"] for e in manifest["slots"]] == [0, 1]
assert all(e["n_saved"] == 40 for e in manifest["slots"])
assert [c[1] for c in calls] == [{"action": "save"}] * 2
assert "/slots/0" in calls[0][0] and "/slots/1" in calls[1][0]
def test_save_unlinks_empty_slot_and_returns_none(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"")
return _Resp(200, {"n_saved": 0, "n_written": 0})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == [] # empty-slot file removed
def test_save_cap_breach_discards_all_files(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 2)
_fake_disk(monkeypatch)
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100)
return _Resp(200, {"n_saved": 40, "n_written": 100})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None # 200 bytes > 150 cap
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_transport_error_aborts_remaining_slots(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 3)
_fake_disk(monkeypatch)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
raise OSError("connection refused")
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert len(calls) == 1 # no retries against a dead server
def test_save_transport_error_unlinks_partial_file(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"partial")
raise OSError("timed out")
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == []
def test_fingerprint_tracks_lora_sidecar_rewrite(tmp_path):
backend = _resume_backend(tmp_path)
adapter = tmp_path / "adapter.gguf"
adapter.write_bytes(b"v1")
backend._extra_args = ["--lora", str(adapter)]
before = backend._slot_launch_fingerprint()
adapter.write_bytes(b"v2-different") # re-exported adapter, same path
assert backend._slot_launch_fingerprint() != before
backend._extra_args = [f"--lora={adapter}"]
assert backend._sidecar_weight_files() == [str(adapter)]
backend._extra_args = ["--lora-scaled", str(adapter), "0.5"]
assert backend._sidecar_weight_files() == [str(adapter)]
backend._extra_args = ["--control-vector", str(adapter), "--threads", "4"]
assert backend._sidecar_weight_files() == [str(adapter)]
def test_sidecar_files_parse_csv_and_colon_scale(tmp_path):
backend = _resume_backend(tmp_path)
a, b = tmp_path / "a.gguf", tmp_path / "b.gguf"
backend._extra_args = ["--lora", f"{a},{b}"]
files = backend._sidecar_weight_files()
assert str(a) in files and str(b) in files
backend._extra_args = ["--lora-scaled", f"{a}:0.5"]
assert str(a) in backend._sidecar_weight_files()
backend._extra_args = ["--control-vector-scaled", f"{a}:1.0,{b}:2.0"]
files = backend._sidecar_weight_files()
assert str(a) in files and str(b) in files
# Windows drive letter must not be mistaken for a scale separator.
backend._extra_args = ["--lora-scaled", "C:\\adapters\\a.gguf:0.75"]
assert "C:\\adapters\\a.gguf" in backend._sidecar_weight_files()
backend._extra_args = ["--lora", "C:\\adapters\\a.gguf"]
assert backend._sidecar_weight_files() == ["C:\\adapters\\a.gguf"]
def test_fingerprint_tracks_colon_scaled_adapter_rewrite(tmp_path):
backend = _resume_backend(tmp_path)
adapter = tmp_path / "adapter.gguf"
adapter.write_bytes(b"v1")
backend._extra_args = ["--lora-scaled", f"{adapter}:0.5"]
before = backend._slot_launch_fingerprint()
adapter.write_bytes(b"v2-different") # re-exported adapter, same path
assert backend._slot_launch_fingerprint() != before
def test_fingerprint_tracks_effective_context_length(tmp_path):
backend = _resume_backend(tmp_path)
backend._effective_context_length = 8192
before = backend._slot_launch_fingerprint()
backend._effective_context_length = 4096 # auto-fit landed smaller on reload
assert backend._slot_launch_fingerprint() != before
def test_gguf_file_identity_covers_split_shards(tmp_path):
backend = _resume_backend(tmp_path)
first = tmp_path / "m-00001-of-00002.gguf"
second = tmp_path / "m-00002-of-00002.gguf"
first.write_bytes(b"a")
second.write_bytes(b"bb")
before = backend._gguf_file_identity(str(first))
st1, st2 = os.stat(first), os.stat(second)
assert before == ((st1.st_size, st1.st_mtime_ns), (st2.st_size, st2.st_mtime_ns))
second.write_bytes(b"rewritten") # sibling changes, primary untouched
after = backend._gguf_file_identity(str(first))
assert after is not None and after != before
assert after[0] == before[0] # primary shard unchanged
second.unlink()
assert backend._gguf_file_identity(str(first)) is None # missing shard
def test_save_skipped_when_user_disabled_prompt_cache(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._extra_args = ["--no-cache-prompt"]
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_env_disables_prompt_cache(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0")
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
monkeypatch.delenv("LLAMA_ARG_CACHE_PROMPT")
monkeypatch.setenv("LLAMA_ARG_NO_CACHE_PROMPT", "1") # legacy negative form
assert backend.save_slots_for_resume() is None
def test_explicit_cache_prompt_flag_overrides_env(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0")
backend._extra_args = ["--cache-prompt"] # CLI wins over env in llama.cpp
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
raising = False,
)
assert backend.save_slots_for_resume() is not None
def test_user_cache_prompt_overrides_studio_no_cache_flag(monkeypatch, tmp_path):
# User extras follow Studio's flags, so an explicit --cache-prompt wins.
backend = _resume_backend(tmp_path)
backend._prompt_cache_disabled = True
backend._extra_args = ["--cache-prompt"]
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
raising = False,
)
assert backend.save_slots_for_resume() is not None
# Last flag wins when both appear in extras.
backend._extra_args = ["--cache-prompt", "--no-cache-prompt"]
assert backend.save_slots_for_resume() is None
def test_save_stops_writing_once_cap_exceeded(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 3)
_fake_disk(monkeypatch)
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100)
return _Resp(200, {"n_saved": 1, "n_written": 100})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert len(calls) == 2 # cap blown after slot 1; slot 2 never attempted
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_aborts_between_slots_when_no_longer_idle(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 3)
_fake_disk(monkeypatch)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
return _Resp(200, {"n_saved": 5, "n_written": 10})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
aborts = iter([False, True, True])
manifest = backend.save_slots_for_resume(should_abort = lambda: next(aborts))
assert len(calls) == 1 # slots 1 and 2 skipped
assert manifest is not None
assert [e["id"] for e in manifest["slots"]] == [0]
def test_save_non_200_slot_is_skipped_but_others_kept(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 2)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
if "/slots/0" in url:
return _Resp(500)
return _Resp(200, {"n_saved": 5, "n_written": 10})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
manifest = backend.save_slots_for_resume()
assert manifest is not None
assert [e["id"] for e in manifest["slots"]] == [1]
def test_restore_posts_each_slot_and_tolerates_failures(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
calls = []
def fake_post(url, **kwargs):
calls.append((url, kwargs["params"], kwargs["json"]))
return _Resp(500 if "/slots/0" in url else 200, {"n_restored": 5})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
backend.restore_slots_for_resume(
{
"slots": [
{"id": 0, "filename": "resume-a-slot0.bin", "n_saved": 5},
{"id": 1, "filename": "resume-a-slot1.bin", "n_saved": 5},
]
}
)
assert [c[1] for c in calls] == [{"action": "restore"}] * 2
assert calls[0][2] == {"filename": "resume-a-slot0.bin"}
def test_restore_transport_error_stops_early(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
raise OSError("connection refused")
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
backend.restore_slots_for_resume(
{"slots": [{"id": 0, "filename": "a.bin"}, {"id": 1, "filename": "b.bin"}]}
)
assert len(calls) == 1
def test_save_deletes_orphan_on_malformed_response(monkeypatch, tmp_path):
# A 200 that writes a file but returns a non-numeric counter must be cleaned
# up like any other save failure, not left orphaned holding chat KV.
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv")
return _Resp(200, {"n_saved": "not-an-int"})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_deletes_orphan_on_non_dict_response(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv")
return _Resp(200, ["unexpected", "list"])
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_cap_uses_actual_file_size_not_reported_bytes(monkeypatch, tmp_path):
# A binary under-reporting n_written must not slip past the disk cap: the
# cap is enforced against the bytes actually on disk.
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 200)
return _Resp(200, {"n_saved": 5, "n_written": 1}) # under-reported
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None # 200 real bytes > 150 cap
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path):
# An estimate over the cap skips before writing any slot at all.
backend = _resume_backend(tmp_path)
backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 1 << 20)
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path):
# The GGUF/sidecars were swapped on disk after the server loaded them, so the
# live KV belongs to the old weights: refuse to persist it (no POST at all).
backend = _resume_backend(tmp_path)
backend._slot_loaded_identity = ((("stale", 0),), ()) # != current identity
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_proceeds_when_load_identity_matches(monkeypatch, tmp_path):
# Matching load-time snapshot: the save runs normally.
backend = _resume_backend(tmp_path)
backend._slot_loaded_identity = (
backend._gguf_file_identity(backend._gguf_path),
backend._slot_launch_fingerprint(),
)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"kv")
return _Resp(200, {"n_saved": 5, "n_written": 2})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
manifest = backend.save_slots_for_resume()
assert manifest is not None
assert [e["id"] for e in manifest["slots"]] == [0]
def test_save_skipped_when_estimate_unavailable_and_low_disk(monkeypatch, tmp_path):
# A 0 estimate means metadata was insufficient, not a zero-byte cache: the save
# must demand room for the whole cap, not just 1 GiB, on a low-disk host.
backend = _resume_backend(tmp_path)
backend._estimate_kv_cache_bytes = lambda *a, **k: 0 # metadata unavailable
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 8 << 30) # 8 GiB cap
_fake_disk(monkeypatch, free = 2 << 30) # 2 GiB free < 8 + 1 GiB required
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None

12
studio/backend/tests/test_llama_server_args.py
View file

@ -183,6 +183,8 @@ def test_non_flag_token_passes_through():
"--reranking",
# llama-server's own --tools clashes with Unsloth's tool policy.
"--tools",
# Slot-state dir: Studio owns it for KV persistence across idle unload.
"--slot-save-path",
],
)
def test_denylist_rejects_all_aliases(denied):
@ -224,6 +226,16 @@ def test_denylist_rejects_equals_form():
validate_extra_args(["--port=9000"])
def test_slot_save_path_is_managed_in_all_forms():
for args in (["--slot-save-path", "/tmp/x"], ["--slot-save-path=/tmp/x"], ["--slot-save-path"]):
with pytest.raises(ValueError, match = "--slot-save-path"):
validate_extra_args(args)
assert is_managed_flag("--slot-save-path") is True
assert is_managed_flag("--slot-save-path=/tmp/x") is True
# --slots (read-only diagnostics endpoint) stays a user choice.
assert is_managed_flag("--slots") is False
@pytest.mark.parametrize(
"padded",
[" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"],

588
studio/backend/tests/test_openai_auto_switch.py
View file

@ -8,6 +8,7 @@ tests/test_gguf_completion_usage.py.
"""
import asyncio
import os
import pytest
@ -18,6 +19,10 @@ from utils import openai_auto_switch_settings as settings
class _FakeBackend:
effective_parallel_slots = 1
_slot_save_binary = None
_gguf_path = None
def __init__(
self,
loaded_id = None,
@ -29,6 +34,22 @@ class _FakeBackend:
self.hf_variant = hf_variant
self._openai_advertised_id = advertised_id
def save_slots_for_resume(self, should_abort = None):
return None
def restore_slots_for_resume(self, manifest):
return None
def _slot_launch_fingerprint(self):
return ((), None, None, 1)
def _gguf_file_identity(self, path):
try:
st = os.stat(path)
except OSError:
return None
return ((st.st_size, st.st_mtime_ns),)
class _LoadRecorder:
"""Stand-in for the load route: records calls and simulates a load."""
@ -53,10 +74,15 @@ class _LoadRecorder:
from fastapi import HTTPException
raise HTTPException(status_code = 503, detail = "load failed")
self.backend.model_identifier = request.model_path
self.backend.hf_variant = getattr(request, "gguf_variant", None)
self.backend._gguf_path = request.model_path
self.backend.is_loaded = True
# Mirror _load_model_impl: a load advertises its own id until the
# auto-switch caller overwrites it with the repo id.
self.backend._openai_advertised_id = None
from core.inference import llama_keepwarm as kw
kw.note_model_loaded(self.backend)
return None
@ -446,6 +472,75 @@ def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch):
assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M"
def test_idle_loop_deletes_saved_kv_when_unload_fails(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
saved = tmp_path / "resume-abc-slot0.bin"
backend = _FakeBackend("unsloth/Idle-GGUF")
manifests = []
def _save(should_abort = None):
if manifests:
return None
saved.write_bytes(b"kv")
manifest = {"dir": str(tmp_path), "slots": [{"id": 0, "filename": saved.name}]}
manifests.append(manifest)
return manifest
def _unload():
raise RuntimeError("cuda teardown failed")
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
async def _drive():
task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01))
for _ in range(200):
await asyncio.sleep(0.01)
if manifests and not saved.exists():
break
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
asyncio.run(_drive())
assert manifests and not saved.exists()
assert kw._kv_resume is None
def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path):
# PUT leaves keep-KV on but makes idle unload inactive: saved KV must go too.
import routes.settings as settings_route
from core.inference import llama_keepwarm as kw
saved = tmp_path / "resume-abc-slot0.bin"
saved.write_bytes(b"kv")
kw._kv_resume = {
"identity": ("m", None, "m"),
"dir": str(tmp_path),
"slots": [{"id": 0, "filename": saved.name}],
}
monkeypatch.setattr(settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True))
monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0)
payload = settings_route.OpenAIAutoSwitchPayload(enabled = False)
resp = settings_route.update_openai_auto_switch(payload, "tester")
assert resp.idle_unload_active is False and resp.auto_unload_keep_kv is True
assert kw._kv_resume is None and not saved.exists()
def test_audio_generate_is_tracked_as_inference_path():
# Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so
# the keep-warm middleware must count it as in-flight inference.
@ -2912,8 +3007,10 @@ def test_non_gguf_load_clears_reload_stash():
# A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF
# branch, so it never lingers until the idle poll (or forever, idle-unload off).
import inspect
src = inspect.getsource(inference_route._load_model_impl)
assert src.count("note_model_loaded()") >= 2
assert src.count("note_model_loaded()") >= 1 # non-GGUF branch
assert "to_thread(note_model_loaded, llama_backend)" in src # GGUF branch
def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch):
@ -3121,6 +3218,495 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp
assert "Model auto-switch" in non_gguf_loaded
# ── idle-unload KV persistence (slot save/restore) ──────────────────
def _seed_kv_manifest(
tmp_path,
identity = ("unsloth/A-GGUF", "Q4_K_M", "unsloth/A-GGUF"),
gguf = None,
):
if gguf is None:
gguf_file = tmp_path / "model.gguf"
gguf_file.write_bytes(b"gguf")
gguf = str(gguf_file)
st = os.stat(gguf)
state_file = tmp_path / "resume-abc-slot0.bin"
state_file.write_bytes(b"kv")
return state_file, {
"identity": identity,
"dir": str(tmp_path),
"binary": ("/bin/llama-server", 111),
"gguf": gguf,
"gguf_stat": ((st.st_size, st.st_mtime_ns),),
"launch": ((), None, None, 1),
"slots": [{"id": 0, "filename": state_file.name, "n_saved": 42}],
}
def _drive_idle_loop(
kw,
poll_seconds = 0.02,
run_for = 0.2,
):
async def _drive():
task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = poll_seconds))
await asyncio.sleep(run_for)
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
asyncio.run(_drive())
def test_idle_unload_saves_slots_before_unload_and_stashes_manifest(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
events = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
manifest = {
"dir": str(tmp_path),
"binary": ("bin", 1),
"slots": [{"id": 0, "filename": "f.bin", "n_saved": 42}],
}
def _save(should_abort = None):
events.append("save")
return manifest
def _unload():
events.append("unload")
backend.is_loaded = False
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
# KV must be saved while the server is still alive, then exactly one unload.
assert events == ["save", "unload"]
assert kw.get_last_unloaded_model()[:2] == ("unsloth/Idle-GGUF", "Q4_K_M")
resume = kw.take_kv_resume()
assert resume is not None
assert resume["identity"][:2] == ("unsloth/Idle-GGUF", "Q4_K_M")
assert resume["slots"][0]["filename"] == "f.bin"
def test_idle_save_failure_still_unloads_plain(monkeypatch):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
unloads = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
def _save(should_abort = None):
raise RuntimeError("slot save exploded")
def _unload():
unloads.append(1)
backend.is_loaded = False
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert unloads == [1] # the save failure must not skip the unload
assert kw.get_last_unloaded_model() is not None
assert kw.take_kv_resume() is None
def test_keep_kv_setting_off_skips_save(monkeypatch):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: False)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
saves, unloads = [], []
backend = _FakeBackend("unsloth/Idle-GGUF")
def _unload():
unloads.append(1)
backend.is_loaded = False
backend.save_slots_for_resume = lambda *a, **k: saves.append(1)
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert saves == []
assert unloads == [1]
assert kw.take_kv_resume() is None
def test_keep_kv_disabled_mid_save_discards_manifest(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
keep = {"on": True}
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: keep["on"])
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
unloads = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
state_file = tmp_path / "resume-mid-slot0.bin"
state_file.write_bytes(b"kv")
manifest = {
"dir": str(tmp_path),
"binary": ("bin", 1),
"slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}],
}
def _save(should_abort = None):
keep["on"] = False # user flips the toggle while the save runs
return manifest
def _unload():
unloads.append(1)
backend.is_loaded = False
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert unloads == [1] # still unloads; only the stash is dropped
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_idle_ttl_disabled_mid_save_skips_unload(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
ttl = {"v": 0.005}
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: ttl["v"])
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
unloads = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
state_file = tmp_path / "resume-mid-slot0.bin"
state_file.write_bytes(b"kv")
manifest = {
"dir": str(tmp_path),
"binary": ("bin", 1),
"slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}],
}
def _save(should_abort = None):
ttl["v"] = 0 # user turns idle unload off while the save runs
return manifest
backend.save_slots_for_resume = _save
backend.unload_model = lambda: unloads.append(1)
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert unloads == [] # the unload was cancelled by the setting change
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_alias_reload_restores_slots_and_deletes_files(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
backend = _FakeBackend(None) # idle-unload emptied the backend
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
rec = _LoadRecorder(backend)
_wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec)
monkeypatch.setattr(kw, "_inflight", 0)
state_file, manifest = _seed_kv_manifest(tmp_path)
monkeypatch.setattr(kw, "_last_unloaded_model", (manifest["gguf"], "Q4_K_M"))
monkeypatch.setattr(kw, "_kv_resume", manifest)
_run_hook("gpt-4o-mini")
assert len(rec.calls) == 1
assert len(restored) == 1 # same model + binary: restore ran
assert not state_file.exists() # state file deleted after the restore
assert kw._kv_resume is None
def test_no_restore_when_different_model_loads(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
backend = _FakeBackend(None)
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
rec = _LoadRecorder(backend)
_wire(
monkeypatch,
enabled = True,
resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"),
backend = backend,
recorder = rec,
)
monkeypatch.setattr(kw, "_inflight", 0)
state_file, manifest = _seed_kv_manifest(tmp_path) # manifest is for model A
monkeypatch.setattr(kw, "_kv_resume", manifest)
_run_hook("unsloth/B-GGUF")
assert len(rec.calls) == 1
assert restored == [] # different model: never restored
assert not state_file.exists() # but the stale files are gone
assert kw._kv_resume is None
def test_restore_skipped_when_binary_changed(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
backend._gguf_path = manifest["gguf"]
backend._slot_save_binary = ("/bin/llama-server", 222) # newer mtime
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
kw.restore_kv_resume(backend, manifest)
assert restored == []
assert not state_file.exists()
def test_restore_skipped_when_launch_config_changed(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
backend._gguf_path = manifest["gguf"]
backend._slot_save_binary = ("/bin/llama-server", 111)
backend._slot_launch_fingerprint = lambda: (("--rope-freq-scale", "0.5"), None, None, 1)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
kw.restore_kv_resume(backend, manifest)
assert restored == []
assert not state_file.exists()
def test_restore_skipped_when_gguf_rewritten_in_place(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
with open(manifest["gguf"], "wb") as fh:
fh.write(b"different weights") # same path, new content
backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
backend._gguf_path = manifest["gguf"]
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
kw.restore_kv_resume(backend, manifest)
assert restored == []
assert not state_file.exists()
def test_note_model_unloaded_purges_manifest_and_files(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M"))
kw._set_kv_resume(manifest)
kw.note_model_unloaded()
assert kw.get_last_unloaded_model() is None
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_note_model_loaded_purges_manifest_and_files(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M"))
kw._set_kv_resume(manifest)
kw.note_model_loaded()
assert kw.get_last_unloaded_model() is None
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_new_idle_save_purges_previous_manifest_files(tmp_path):
from core.inference import llama_keepwarm as kw
old_file, old_manifest = _seed_kv_manifest(tmp_path)
kw._set_kv_resume(old_manifest)
new_file = tmp_path / "resume-def-slot0.bin"
new_file.write_bytes(b"kv2")
kw._set_kv_resume(
{
"identity": ("unsloth/B-GGUF", None, "unsloth/B-GGUF"),
"dir": str(tmp_path),
"binary": ("/bin/llama-server", 111),
"slots": [{"id": 0, "filename": new_file.name, "n_saved": 7}],
}
)
assert not old_file.exists() # replaced manifest's files purged
assert new_file.exists()
assert kw.take_kv_resume()["slots"][0]["filename"] == new_file.name
def test_sweep_slot_save_dir_removes_only_resume_files(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
from utils.paths import storage_roots
monkeypatch.setattr(storage_roots, "llama_slot_cache_root", lambda: tmp_path)
stale = tmp_path / "resume-old-slot0.bin"
stale.write_bytes(b"kv")
other = tmp_path / "unrelated.txt"
other.write_text("keep")
kw.sweep_slot_save_dir()
assert not stale.exists()
assert other.exists()
def test_keep_kv_setting_roundtrip_and_default(monkeypatch):
import storage.studio_db as db
store = {}
monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
assert settings.get_auto_unload_keep_kv() is True # default when never stored
assert settings.set_openai_auto_switch(True, 60, False)[2] is False
assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False
assert settings.get_auto_unload_keep_kv() is False
# None leaves the stored value untouched (older clients can't reset it).
assert settings.set_openai_auto_switch(True, 60, None)[2] is False
assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False
with pytest.raises(ValueError, match = "true or false"):
settings.set_openai_auto_switch(True, 60, "garbage")
def test_stale_stash_cleanup_waits_for_lifecycle_gate(monkeypatch, tmp_path):
# The loop's stale-stash purge must wait on the gate a mid-reload holds.
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 3600)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic()
backend = _FakeBackend("unsloth/New-GGUF")
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
state_file, manifest = _seed_kv_manifest(tmp_path)
kw._kv_resume = manifest
kw._last_unloaded_model = ("unsloth/A-GGUF", "Q4_K_M")
assert kw._lifecycle_lock.acquire(blocking = False) # simulate in-flight reload
try:
_drive_idle_loop(kw)
assert kw._kv_resume is manifest # purge deferred while the gate is held
assert state_file.exists()
finally:
kw._lifecycle_lock.release()
_drive_idle_loop(kw)
assert kw._kv_resume is None # gate freed: genuinely stale stash purged
assert not state_file.exists()
def test_put_route_disabling_keep_kv_purges_saved_state(monkeypatch, tmp_path):
import routes.settings as settings_route
import storage.studio_db as db
from core.inference import llama_keepwarm as kw
store = {}
monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
state_file, manifest = _seed_kv_manifest(tmp_path)
monkeypatch.setattr(kw, "_kv_resume", manifest)
payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_keep_kv = False)
resp = settings_route.update_openai_auto_switch(payload, "tester")
assert resp.auto_unload_keep_kv is False
assert kw._kv_resume is None
assert not state_file.exists()
def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch):
# A keep-KV-only update must not materialize the env TTL as a stored value.
import routes.settings as settings_route
import storage.studio_db as db
store = {}
monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600")
assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None
enabled, idle, keep_kv = settings.set_openai_auto_switch(False, None, False)
assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched
assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active
assert (enabled, idle, keep_kv) == (False, 600, False)
def test_load_impl_notes_loaded_with_backend_off_loop():
import inspect
src = inspect.getsource(inference_route._load_model_impl)
assert "to_thread(note_model_loaded, llama_backend)" in src
def test_restore_matches_gguf_realpath_across_naming(tmp_path):
from core.inference import llama_keepwarm as kw
blob = tmp_path / "blob.gguf"
blob.write_bytes(b"gguf")
link = tmp_path / "snapshot.gguf"
try:
link.symlink_to(blob)
except OSError:
pytest.skip("symlinks unsupported on this host")
backend = _FakeBackend("/hf/snapshots/d7f5", hf_variant = None)
backend._gguf_path = str(link) # reload resolved the symlink spelling
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
state_file, manifest = _seed_kv_manifest(
tmp_path, identity = ("unsloth/A-GGUF", None, "unsloth/A-GGUF"), gguf = str(blob)
)
kw.restore_kv_resume(backend, manifest)
assert len(restored) == 1 # names differ, file identical: restore ran
assert not state_file.exists()
def test_setter_rejects_idle_below_floor(monkeypatch):
import storage.studio_db as db

61
studio/backend/utils/openai_auto_switch_settings.py
View file

@ -30,11 +30,13 @@ from typing import Any, Optional
OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model"
AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds"
AUTO_UNLOAD_KEEP_KV_SETTING_KEY = "openai_api_auto_unload_keep_kv"
MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides"
MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
DEFAULT_AUTO_UNLOAD_KEEP_KV = True
MIN_AUTO_UNLOAD_IDLE_SECONDS = 60
_CACHE_TTL_S = 2.0
@ -158,29 +160,54 @@ def get_auto_unload_idle_seconds() -> int:
return env if env is not None else 0
def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]:
"""Set both auto-switch flags in one transaction so a settings PUT can't leave
one key updated and the other stale. Both values are coerced before any write,
so an invalid value raises without persisting either."""
def get_auto_unload_keep_kv() -> bool:
"""Whether the idle unload persists slot KV to disk for restore on reload."""
parsed = _coerce_bool(_cached_setting(AUTO_UNLOAD_KEEP_KV_SETTING_KEY, None))
return parsed if parsed is not None else DEFAULT_AUTO_UNLOAD_KEEP_KV
def set_openai_auto_switch(
enabled: Any,
idle_seconds: Any,
keep_kv: Any = None,
) -> tuple[bool, int, bool]:
"""One-transaction write; ``None`` leaves a stored value untouched."""
parsed_enabled = _coerce_bool(enabled)
if parsed_enabled is None:
raise ValueError("OpenAI auto-switch must be true or false.")
parsed_idle = _coerce_int(idle_seconds)
if parsed_idle is None:
raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS:
raise ValueError(
f"Auto-unload idle seconds must be 0 (off) or at least "
f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}."
)
parsed_idle = None
if idle_seconds is not None:
parsed_idle = _coerce_int(idle_seconds)
if parsed_idle is None:
raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS:
raise ValueError(
f"Auto-unload idle seconds must be 0 (off) or at least "
f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}."
)
parsed_keep_kv = None
if keep_kv is not None:
parsed_keep_kv = _coerce_bool(keep_kv)
if parsed_keep_kv is None:
raise ValueError("Keep KV on idle unload must be true or false.")
from storage.studio_db import upsert_app_settings
upsert_app_settings(
{OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle}
)
updates: dict[str, Any] = {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled}
if parsed_idle is not None:
updates[AUTO_UNLOAD_IDLE_SETTING_KEY] = parsed_idle
if parsed_keep_kv is not None:
updates[AUTO_UNLOAD_KEEP_KV_SETTING_KEY] = parsed_keep_kv
upsert_app_settings(updates)
_invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY)
_invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
return parsed_enabled, parsed_idle
if parsed_idle is not None:
_invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
if parsed_keep_kv is not None:
_invalidate(AUTO_UNLOAD_KEEP_KV_SETTING_KEY)
return (
parsed_enabled,
parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(),
parsed_keep_kv if parsed_keep_kv is not None else get_auto_unload_keep_kv(),
)
def get_model_overrides() -> dict[str, dict]:

5
studio/backend/utils/paths/storage_roots.py
View file

@ -61,6 +61,11 @@ def cache_root() -> Path:
return studio_root() / "cache"
def llama_slot_cache_root() -> Path:
"""Dir llama-server saves/restores slot KV state in across idle unloads."""
return cache_root() / "llama-slots"
def studio_bin_root() -> Path:
"""Dir for Unsloth-managed executables (the `unsloth` shim, downloaded tools like cloudflared)."""
return studio_root() / "bin"

19
studio/frontend/src/features/settings/api/openai-auto-switch.ts
View file

@ -11,6 +11,8 @@ export type OpenAIAutoSwitchSettings = {
// True when the idle-unload loop will actually unload (e.g. enabled via the
// UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off).
idleUnloadActive: boolean;
// Persist the KV cache to disk on idle unload and restore it on reload.
autoUnloadKeepKv: boolean;
};
type ApiOpenAIAutoSwitchSettings = {
@ -21,6 +23,8 @@ type ApiOpenAIAutoSwitchSettings = {
default_enabled: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
idle_unload_active?: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_keep_kv?: boolean;
};
let cachedSettings: OpenAIAutoSwitchSettings | null = null;
@ -34,6 +38,7 @@ function fromApi(
autoUnloadIdleSeconds: settings.auto_unload_idle_seconds,
defaultEnabled: settings.default_enabled,
idleUnloadActive: settings.idle_unload_active ?? false,
autoUnloadKeepKv: settings.auto_unload_keep_kv ?? true,
};
}
@ -66,15 +71,23 @@ export async function loadOpenAIAutoSwitchSettings() {
export async function updateOpenAIAutoSwitchSettings(
enabled: boolean,
autoUnloadIdleSeconds: number,
autoUnloadIdleSeconds?: number,
autoUnloadKeepKv?: boolean,
): Promise<OpenAIAutoSwitchSettings> {
const res = await authFetch("/api/settings/openai-auto-switch", {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
enabled,
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_idle_seconds: autoUnloadIdleSeconds,
// Omitted fields keep their stored value.
...(autoUnloadIdleSeconds === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_unload_idle_seconds: autoUnloadIdleSeconds }),
...(autoUnloadKeepKv === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_unload_keep_kv: autoUnloadKeepKv }),
}),
});
if (!res.ok) {

26
studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
View file

@ -62,13 +62,18 @@ export function ModelAutoSwitchSection() {
const persist = async (
enabled: boolean,
idleSeconds: number,
idleSeconds: number | undefined,
syncDraft = true,
keepKv?: boolean,
) => {
setIsSaving(true);
setError(null);
try {
const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds);
const saved = await updateOpenAIAutoSwitchSettings(
enabled,
idleSeconds,
keepKv,
);
setSettings(saved);
if (syncDraft) {
setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds));
@ -107,6 +112,11 @@ export function ModelAutoSwitchSection() {
void persist(true, idleSeconds);
};
const handleKeepKvToggle = (keepKv: boolean) => {
if (!settings) return;
void persist(settings.enabled, undefined, false, keepKv);
};
return (
<SettingsSection title={t("settings.general.modelAutoSwitch.sectionTitle")}>
<SettingsRow
@ -166,6 +176,18 @@ export function ModelAutoSwitchSection() {
) : null}
</div>
</SettingsRow>
{settings?.idleUnloadActive ? (
<SettingsRow
label={t("settings.general.modelAutoSwitch.keepKv")}
description={t("settings.general.modelAutoSwitch.keepKvDescription")}
>
<Switch
checked={settings.autoUnloadKeepKv}
disabled={isSaving}
onCheckedChange={handleKeepKvToggle}
/>
</SettingsRow>
) : null}
</SettingsSection>
);
}

3
studio/frontend/src/i18n/locales/en.ts
View file

@ -232,6 +232,9 @@ export const en = {
loadError: "Failed to load model auto-switch settings.",
saveError: "Failed to save model auto-switch settings.",
idleError: "Enter 0 to keep the model loaded, or at least 60 seconds.",
keepKv: "Keep chat context across idle unload",
keepKvDescription:
"Save the model's KV cache to disk before an idle unload and restore it on reload, so resumed chats skip re-reading their history. Chat context is written to disk (up to 10 GB) until it is restored or cleaned up.",
},
previewSharing: {
sectionTitle: "Preview sharing",