Compare commits

...
Sign in to create a new pull request.

2 commits

Author SHA1 Message Date
Daniel Han
6590f18638
Merge branch 'main' into studio-model-idle-ttl 2026-06-23 08:14:27 -07:00
Daniel Han
80c15643cf studio: add a configurable idle TTL that auto-unloads an unused model
A loaded GGUF model stayed resident until a manual unload or process exit, with
no way to set an idle timeout. Add an idle TTL (seconds): when set, a background
task started by the app lifespan unloads the model after it has been idle that
long. 0 disables eviction, preserving the historical behavior.

- utils/model_ttl_settings.py: the setting (app_settings key), the
  UNSLOTH_MODEL_IDLE_TTL env default, and validation (0 to one week).
- LlamaCppBackend tracks last-activity: refreshed on load, on every generation
  request entry, and on every streamed chunk, so a long generation (e.g. a
  minutes-long reasoning stream) is never evicted mid-flight; idle_seconds is
  None when nothing is loaded.
- main.py runs the eviction loop in the lifespan and cancels it on shutdown.
- GET/PUT /api/settings/model-ttl to read/set the TTL at runtime; the response
  also reports the loaded model's current idle and time-to-eviction.

The request-path activity ping goes through a small _note_idle_activity() guard
so a backend that does not implement the hook degrades gracefully instead of
turning a served request into a 500.

Eviction is made race-safe. unload_model sets _cancel_event, which an in-flight
load of a different model watches, so a naive evict could abort that load. The
evictor now goes through LlamaCppBackend.evict_if_idle, which re-checks idle and
in-flight under _serial_load_lock (the lock load_model holds for its whole
duration) and unloads atomically, so it can never run concurrently with a load
or unload. An in-flight request counter (request_in_flight, applied to the chat
and tool generators) additionally blocks eviction while a request is active,
covering the window before the first streamed token when per-chunk activity has
not started refreshing yet.
2026-06-22 15:39:13 +00:00
6 changed files with 453 additions and 0 deletions

View file

@ -9,6 +9,7 @@ OpenAI-compatible /v1/chat/completions endpoint.
import atexit
import contextlib
import functools
import json
import os
import re
@ -1219,6 +1220,20 @@ def _backfill_usage_from_timings(usage, timings):
return out
def _tracks_inflight(genfunc):
"""Decorate a generator method so it counts as an in-flight request for the
whole time the caller iterates it. Keeps the idle-TTL evictor from unloading
the model mid-generation -- including the window before the first streamed
token, when per-chunk activity has not started refreshing yet."""
@functools.wraps(genfunc)
def wrapper(self, *args, **kwargs):
with self.request_in_flight():
yield from genfunc(self, *args, **kwargs)
return wrapper
class LlamaCppBackend:
"""Manages a llama-server subprocess for GGUF model inference.
@ -1350,6 +1365,18 @@ class LlamaCppBackend:
# to decide whether to wait for the VRAM reclaim to finish.
self._last_kill_monotonic: float = 0.0
# Monotonic timestamp of the last generation activity, refreshed on load
# and on every streamed chunk. Drives the idle-TTL eviction loop so a
# genuinely idle model can be unloaded automatically.
self._last_activity: float = 0.0
# Count of inference requests currently in flight. The idle-TTL evictor
# refuses to unload while this is > 0, so a request that is still waiting
# on its first token (no per-chunk activity yet) is never evicted
# mid-flight. Guarded by its own lock; requests bracket themselves with
# request_in_flight().
self._inflight_requests: int = 0
self._inflight_lock = threading.Lock()
_reaped = self._kill_orphaned_servers()
if _reaped:
# Reaped VRAM frees lazily; arm the settle wait so the first load
@ -1368,6 +1395,76 @@ class LlamaCppBackend:
"""True if a llama-server process exists (loading or loaded)."""
return self._process is not None
def note_activity(self) -> None:
"""Record that the loaded model was just used (drives idle-TTL eviction)."""
self._last_activity = time.monotonic()
@property
def last_activity_monotonic(self) -> float:
return self._last_activity
@property
def idle_seconds(self) -> Optional[float]:
"""Seconds since the model was last used, or None when none is loaded."""
if not self.is_loaded:
return None
if not self._last_activity:
return 0.0
return max(0.0, time.monotonic() - self._last_activity)
def _inflight_lock_obj(self) -> "threading.Lock":
"""The in-flight counter lock, created on demand. __init__ sets it, but
some tests build a backend via __new__ (skipping __init__); creating it
lazily keeps the idle-TTL guard working for those too."""
lock = getattr(self, "_inflight_lock", None)
if lock is None:
lock = threading.Lock()
self._inflight_lock = lock
if not hasattr(self, "_inflight_requests"):
self._inflight_requests = 0
return lock
@contextlib.contextmanager
def request_in_flight(self):
"""Bracket an inference request so the idle-TTL evictor will not unload
the model while it runs. Refreshes activity on entry and exit so the gap
before the first streamed token (when per-chunk activity has not started)
cannot be mistaken for idleness."""
lock = self._inflight_lock_obj()
with lock:
self._inflight_requests += 1
self.note_activity()
try:
yield
finally:
with lock:
self._inflight_requests = max(0, self._inflight_requests - 1)
self.note_activity()
def evict_if_idle(self, ttl: Optional[float]) -> Optional[float]:
"""Unload the model iff it is loaded, has no in-flight request, and has
been idle at least ``ttl`` seconds. Returns the idle seconds at eviction
time, or ``None`` if nothing was evicted.
Serialized via ``_serial_load_lock`` (the same lock ``load_model`` holds
for its whole duration) so eviction can never run concurrently with a
load: ``unload_model`` sets ``_cancel_event``, which would otherwise abort
an in-flight load of a different model. The idle and in-flight checks are
re-evaluated under the lock, so a load or a fresh request that lands first
wins the race.
"""
if not ttl or ttl <= 0:
return None
with self._serial_load_lock:
with self._inflight_lock_obj():
if self._inflight_requests > 0:
return None
idle = self.idle_seconds
if self.is_loaded and idle is not None and idle >= ttl:
self.unload_model()
return idle
return None
@property
def base_url(self) -> str:
return f"http://127.0.0.1:{self._port}"
@ -3580,6 +3677,7 @@ class LlamaCppBackend:
healthy = self._wait_for_health(timeout = 600.0)
if healthy:
self._healthy = True
self.note_activity()
self._gpu_offload_active = True
if extra_args is not None:
self._extra_args = list(extra_args)
@ -6015,6 +6113,7 @@ class LlamaCppBackend:
)
self._healthy = True
self.note_activity()
# Commit caller intent only after _healthy=True so a failed start
# can't poison the next inheritance check. None keeps prior, []
@ -7497,6 +7596,7 @@ class LlamaCppBackend:
logger.error(f"Failed to respawn llama-server: {exc}")
return False
@_tracks_inflight
def generate_chat_completion(
self,
messages: list[dict],
@ -7579,6 +7679,9 @@ class LlamaCppBackend:
first_token_deadline = first_token_deadline,
):
buffer += raw_chunk
# Idle-TTL activity: each streamed chunk keeps the model from
# being auto-evicted mid-generation.
self.note_activity()
while "\n" in buffer:
line, buffer = buffer.split("\n", 1)
line = line.strip()
@ -7704,6 +7807,7 @@ class LlamaCppBackend:
# ── Tool-calling agentic loop ──────────────────────────────
@_tracks_inflight
def generate_chat_completion_with_tools(
self,
messages: list[dict],
@ -8632,6 +8736,9 @@ class LlamaCppBackend:
first_token_deadline = first_token_deadline,
):
buffer += raw_chunk
# Idle-TTL activity: each streamed chunk keeps the model from
# being auto-evicted mid-generation.
self.note_activity()
while "\n" in buffer:
line, buffer = buffer.split("\n", 1)
line = line.strip()

View file

@ -440,6 +440,55 @@ def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None:
).start()
async def _model_idle_eviction_loop() -> None:
"""Unload the loaded GGUF model once it has been idle longer than the
configured TTL. ``0`` disables eviction (the historical behavior). Activity is
recorded on every generation request and streamed chunk, so this only fires on
a genuinely idle model. Started/stopped by the app lifespan."""
import os as _os
import structlog as _structlog
from utils.model_ttl_settings import (
MODEL_IDLE_EVICTION_POLL_SECONDS,
get_model_idle_ttl_seconds,
)
_log = _structlog.get_logger(__name__)
while True:
try:
await asyncio.sleep(MODEL_IDLE_EVICTION_POLL_SECONDS)
ttl = get_model_idle_ttl_seconds()
if ttl <= 0:
continue
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
# Cheap pre-check off the lock to avoid hopping to a thread every
# poll; the authoritative, atomic decision is made in evict_if_idle.
idle = backend.idle_seconds
if not (backend.is_loaded and idle is not None and idle >= ttl):
continue
# Read the label before eviction (model_identifier clears on unload).
_label = _os.path.basename(str(backend.model_identifier or "")) or "model"
# evict_if_idle re-checks idle + in-flight under the load lock and
# unloads atomically, so it cannot race a concurrent load/unload or
# evict a request still waiting on its first token. Blocking, so run
# it off the event loop.
evicted_idle = await asyncio.to_thread(backend.evict_if_idle, ttl)
if evicted_idle is not None:
_log.info(
"model.idle_evict",
model = _label,
idle_seconds = round(evicted_idle, 1),
ttl_seconds = ttl,
)
except asyncio.CancelledError:
break
except Exception as exc:
_log.debug("model idle eviction loop error: %s", exc)
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
@ -521,8 +570,21 @@ async def lifespan(app: FastAPI):
print("=" * 60 + "\n")
else:
app.state.bootstrap_password = storage.get_bootstrap_password()
app.state.model_idle_eviction_task = asyncio.create_task(_model_idle_eviction_loop())
yield
_evict_task = getattr(app.state, "model_idle_eviction_task", None)
if _evict_task is not None:
_evict_task.cancel()
try:
await _evict_task
except asyncio.CancelledError:
pass
except Exception:
pass
from core.inference.llama_http import aclose as _close_llama_http
await _close_llama_http()

View file

@ -2230,6 +2230,19 @@ def get_llama_cpp_backend() -> LlamaCppBackend:
return _llama_cpp_backend
def _note_idle_activity(backend) -> None:
"""Best-effort idle-TTL activity ping at a request boundary.
The real GGUF backend implements ``note_activity`` for the idle-TTL evictor.
Calling it through this guard keeps the request path working for any backend
that does not track activity, so a missing telemetry hook can never turn a
served request into a 500.
"""
note = getattr(backend, "note_activity", None)
if callable(note):
note()
def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool:
"""Effective quantization the loader will use: a LoRA adapter can flip 4-bit to
16-bit via adapter_config.json, so the guard sizes this, not the raw request."""
@ -4778,6 +4791,10 @@ async def openai_chat_completions(
llama_backend = get_llama_cpp_backend()
using_gguf = llama_backend.is_loaded
if using_gguf:
# Refresh idle-TTL activity at request start; streamed chunks refresh it
# again so a long generation is never auto-evicted mid-flight.
_note_idle_activity(llama_backend)
# OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, which
# the SDK spreads into the request body at the top level. Studio's
@ -6390,6 +6407,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
status_code = 503,
detail = "No GGUF model loaded. Load a GGUF model first.",
)
_note_idle_activity(llama_backend)
body = await request.json()
if body.get("max_tokens") is None:
@ -6560,6 +6578,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get
status_code = 503,
detail = "No GGUF model loaded. Load a GGUF model first.",
)
_note_idle_activity(llama_backend)
body = await request.json()
target_url = f"{llama_backend.base_url}/v1/embeddings"

View file

@ -31,6 +31,13 @@ from utils.helper_precache_settings import (
helper_model_disabled_by_env,
set_helper_precache_enabled,
)
from utils.model_ttl_settings import (
MAX_MODEL_IDLE_TTL_SECONDS,
MIN_MODEL_IDLE_TTL_SECONDS,
default_model_idle_ttl_seconds,
get_model_idle_ttl_seconds,
set_model_idle_ttl_seconds,
)
router = APIRouter()
@ -197,3 +204,60 @@ def update_personalization_settings(
log = logger,
) from exc
return payload
class ModelIdleTtlPayload(BaseModel):
idle_ttl_seconds: int = Field(..., ge = MIN_MODEL_IDLE_TTL_SECONDS, le = MAX_MODEL_IDLE_TTL_SECONDS)
class ModelIdleTtlResponse(BaseModel):
idle_ttl_seconds: int
default_idle_ttl_seconds: int
min_idle_ttl_seconds: int = MIN_MODEL_IDLE_TTL_SECONDS
max_idle_ttl_seconds: int = MAX_MODEL_IDLE_TTL_SECONDS
enabled: bool
loaded_model_idle_seconds: float | None = None
evicts_in_seconds: float | None = None
def _model_idle_ttl_response(ttl: int) -> ModelIdleTtlResponse:
idle: float | None = None
evicts_in: float | None = None
try:
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
idle = backend.idle_seconds
if ttl > 0 and idle is not None:
evicts_in = max(0.0, ttl - idle)
except Exception:
pass
return ModelIdleTtlResponse(
idle_ttl_seconds = ttl,
default_idle_ttl_seconds = default_model_idle_ttl_seconds(),
enabled = ttl > 0,
loaded_model_idle_seconds = idle,
evicts_in_seconds = evicts_in,
)
@router.get("/model-ttl", response_model = ModelIdleTtlResponse)
def get_model_ttl(current_subject: str = Depends(get_current_subject)) -> ModelIdleTtlResponse:
return _model_idle_ttl_response(get_model_idle_ttl_seconds())
@router.put("/model-ttl", response_model = ModelIdleTtlResponse)
def update_model_ttl(
payload: ModelIdleTtlPayload, current_subject: str = Depends(get_current_subject)
) -> ModelIdleTtlResponse:
try:
ttl = set_model_idle_ttl_seconds(payload.idle_ttl_seconds)
except ValueError as exc:
raise log_and_http_error(
exc,
400,
safe_error_detail(exc, fallback = "Invalid model idle TTL."),
event = "settings.update_model_ttl_failed",
log = logger,
) from exc
return _model_idle_ttl_response(ttl)

View file

@ -0,0 +1,130 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Idle-TTL settings + backend activity tracking that drive auto-eviction."""
import sys
import time
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import utils.model_ttl_settings as ttl # noqa: E402
def test_coerce_bounds_and_types():
assert ttl._coerce_ttl_seconds(60) == 60
assert ttl._coerce_ttl_seconds("120") == 120
assert ttl._coerce_ttl_seconds(0) == 0
assert ttl._coerce_ttl_seconds(-5) is None
assert ttl._coerce_ttl_seconds(ttl.MAX_MODEL_IDLE_TTL_SECONDS + 1) is None
assert ttl._coerce_ttl_seconds(True) is None
assert ttl._coerce_ttl_seconds("abc") is None
def test_validate_raises_on_invalid():
with pytest.raises(ValueError):
ttl.validate_model_idle_ttl_seconds(-1)
assert ttl.validate_model_idle_ttl_seconds(300) == 300
def test_default_from_env(monkeypatch):
monkeypatch.delenv("UNSLOTH_MODEL_IDLE_TTL", raising = False)
assert ttl.default_model_idle_ttl_seconds() == 0 # disabled by default
monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "900")
assert ttl.default_model_idle_ttl_seconds() == 900
monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "bad")
assert ttl.default_model_idle_ttl_seconds() == 0
def test_get_set_roundtrip(monkeypatch):
store: dict = {}
monkeypatch.setattr("storage.studio_db.get_app_setting", lambda k, d = None: store.get(k, d))
monkeypatch.setattr("storage.studio_db.upsert_app_settings", lambda d: store.update(d))
monkeypatch.delenv("UNSLOTH_MODEL_IDLE_TTL", raising = False)
assert ttl.get_model_idle_ttl_seconds() == 0 # default disabled
assert ttl.set_model_idle_ttl_seconds(600) == 600
assert store[ttl.MODEL_IDLE_TTL_SETTING_KEY] == 600
assert ttl.get_model_idle_ttl_seconds() == 600
def test_backend_idle_seconds_and_activity(monkeypatch):
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
# No model loaded -> idle is undefined (None), so the evictor never fires.
monkeypatch.setattr(type(backend), "is_loaded", property(lambda self: False))
assert backend.idle_seconds is None
# Loaded -> idle measured from the last activity timestamp.
monkeypatch.setattr(type(backend), "is_loaded", property(lambda self: True))
backend._last_activity = time.monotonic() - 5.0
assert backend.idle_seconds >= 4.0
backend.note_activity()
assert backend.idle_seconds < 1.0
def _loaded_backend(monkeypatch):
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
monkeypatch.setattr(type(backend), "is_loaded", property(lambda self: True))
# Spy unload so eviction does not touch a real subprocess.
calls = {"unload": 0}
def _fake_unload():
calls["unload"] += 1
return True
monkeypatch.setattr(backend, "unload_model", _fake_unload)
backend._inflight_requests = 0
return backend, calls
def test_evict_if_idle_disabled_and_not_idle(monkeypatch):
backend, calls = _loaded_backend(monkeypatch)
backend._last_activity = time.monotonic() - 100.0
# ttl 0/None disables eviction entirely.
assert backend.evict_if_idle(0) is None
assert backend.evict_if_idle(None) is None
# Idle below the TTL: nothing happens.
backend.note_activity()
assert backend.evict_if_idle(60) is None
assert calls["unload"] == 0
def test_evict_if_idle_unloads_when_idle(monkeypatch):
backend, calls = _loaded_backend(monkeypatch)
backend._last_activity = time.monotonic() - 120.0
evicted = backend.evict_if_idle(60)
assert evicted is not None and evicted >= 60
assert calls["unload"] == 1
def test_evict_if_idle_skips_when_request_in_flight(monkeypatch):
backend, calls = _loaded_backend(monkeypatch)
backend._last_activity = time.monotonic() - 120.0
with backend.request_in_flight():
# A request is active (even if it has not streamed a token yet): the
# model must not be evicted out from under it.
assert backend.evict_if_idle(60) is None
assert calls["unload"] == 0
# Once it finishes, activity was just refreshed, so still no eviction.
assert backend.idle_seconds < 1.0
def test_request_in_flight_balances_counter(monkeypatch):
backend, _ = _loaded_backend(monkeypatch)
assert backend._inflight_requests == 0
with backend.request_in_flight():
assert backend._inflight_requests == 1
with backend.request_in_flight():
assert backend._inflight_requests == 2
assert backend._inflight_requests == 1
assert backend._inflight_requests == 0

View file

@ -0,0 +1,71 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Idle TTL for a loaded inference model.
A loaded GGUF model otherwise stays resident until a manual unload or process
exit. When the idle TTL is set (> 0), a background task unloads the model after
it has been idle (no generation activity) for that many seconds. 0 disables
eviction, preserving the historical behavior. The value can be set at startup
via the ``UNSLOTH_MODEL_IDLE_TTL`` env var or at runtime via
``PUT /api/settings/model-ttl``.
"""
from __future__ import annotations
import os
from typing import Any
MODEL_IDLE_TTL_SETTING_KEY = "model_idle_ttl_seconds"
DEFAULT_MODEL_IDLE_TTL_SECONDS = 0 # 0 = disabled (never auto-evict)
MIN_MODEL_IDLE_TTL_SECONDS = 0
MAX_MODEL_IDLE_TTL_SECONDS = 7 * 24 * 60 * 60 # one week ceiling
MODEL_IDLE_EVICTION_POLL_SECONDS = 30.0
_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
def _coerce_ttl_seconds(value: Any) -> int | None:
if isinstance(value, bool):
return None
try:
parsed = int(value)
except (TypeError, ValueError):
return None
if parsed < MIN_MODEL_IDLE_TTL_SECONDS or parsed > MAX_MODEL_IDLE_TTL_SECONDS:
return None
return parsed
def default_model_idle_ttl_seconds() -> int:
"""The startup default: the env override when valid, else disabled (0)."""
env_value = _coerce_ttl_seconds(os.environ.get(_ENV_VAR))
return env_value if env_value is not None else DEFAULT_MODEL_IDLE_TTL_SECONDS
def validate_model_idle_ttl_seconds(value: Any) -> int:
parsed = _coerce_ttl_seconds(value)
if parsed is None:
raise ValueError(
"Model idle TTL must be a whole number of seconds from "
f"{MIN_MODEL_IDLE_TTL_SECONDS} to {MAX_MODEL_IDLE_TTL_SECONDS} (0 disables it)."
)
return parsed
def get_model_idle_ttl_seconds() -> int:
"""Effective TTL: the stored setting when valid, else the startup default."""
try:
from storage.studio_db import get_app_setting
stored = get_app_setting(MODEL_IDLE_TTL_SETTING_KEY, None)
except Exception:
stored = None
parsed = _coerce_ttl_seconds(stored)
return parsed if parsed is not None else default_model_idle_ttl_seconds()
def set_model_idle_ttl_seconds(value: Any) -> int:
parsed = validate_model_idle_ttl_seconds(value)
from storage.studio_db import upsert_app_settings
upsert_app_settings({MODEL_IDLE_TTL_SETTING_KEY: parsed})
return parsed