Compare commits
2 commits
main
...
studio-mod
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6590f18638 |
||
|
|
80c15643cf |
6 changed files with 453 additions and 0 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
130
studio/backend/tests/test_model_idle_ttl.py
Normal file
130
studio/backend/tests/test_model_idle_ttl.py
Normal 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
|
||||
71
studio/backend/utils/model_ttl_settings.py
Normal file
71
studio/backend/utils/model_ttl_settings.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue