diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8a7fad9af4..83b56c0091 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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() diff --git a/studio/backend/main.py b/studio/backend/main.py index a56bd46c4b..d9bc74f883 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -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() diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2e2c38933e..a8a75d1cf5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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" diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index d01009fc20..ba2fa21224 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -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) diff --git a/studio/backend/tests/test_model_idle_ttl.py b/studio/backend/tests/test_model_idle_ttl.py new file mode 100644 index 0000000000..0e71f9b2a8 --- /dev/null +++ b/studio/backend/tests/test_model_idle_ttl.py @@ -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 diff --git a/studio/backend/utils/model_ttl_settings.py b/studio/backend/utils/model_ttl_settings.py new file mode 100644 index 0000000000..28d7127568 --- /dev/null +++ b/studio/backend/utils/model_ttl_settings.py @@ -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