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.
130 lines
4.7 KiB
Python
130 lines
4.7 KiB
Python
# 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
|