fix(models): do not preload a model on a machine with no local user
The startup preload exists so the first generate feels instant for the person sitting in front of the app. A machine lending its GPU has nobody sitting there, so it was several GB of VRAM held from boot against a request that may never arrive — and the idle sweep could not reclaim it, because the sweep owns the worker executor's engines while this is the default local model. Measured on gpu2: a node that had run nothing still sat at 2.4 GB, and an idle unload after a real job returned it to exactly that floor rather than below it. Worker-mode processes now load on first request and release when idle, which is what a node should do. A machine that is both a desktop app and a worker keeps the warm-up — there is a real user there and the point stands.
This commit is contained in:
@@ -2441,6 +2441,15 @@ def _checkpoint_in_local_cache(checkpoint: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _headless_worker() -> bool:
|
||||
"""True when this process serves remote work and has no local UI."""
|
||||
try:
|
||||
from worker.agent import worker_mode_enabled # noqa: PLC0415
|
||||
except Exception:
|
||||
return False
|
||||
return worker_mode_enabled()
|
||||
|
||||
|
||||
async def preload_model():
|
||||
"""Background model warm-up — call from lifespan startup.
|
||||
|
||||
@@ -2451,6 +2460,23 @@ async def preload_model():
|
||||
global model, _last_used
|
||||
if model is not None:
|
||||
return # already loaded
|
||||
|
||||
# A machine lending its GPU has no local user to warm the model FOR. This
|
||||
# preload exists to make the first /generate feel instant for the person
|
||||
# sitting in front of the app; on a headless node there is nobody sitting
|
||||
# there, so it is several GB of VRAM held from boot against a request that
|
||||
# may never come — and the idle sweep cannot reclaim it, because the sweep
|
||||
# owns the worker executor's engines and this is the default local model.
|
||||
# Observed on hardware: a node that had run nothing still sat at 2.4 GB.
|
||||
#
|
||||
# A machine that is BOTH a desktop app and a worker keeps the warm-up:
|
||||
# there is a real user there, and the whole point stands.
|
||||
if _headless_worker():
|
||||
logger.info(
|
||||
"Preload skipped: this process is running as a remote worker, so the "
|
||||
"model loads on first request and is released when it goes idle."
|
||||
)
|
||||
return
|
||||
try:
|
||||
# Warm-up is gated on LOCAL availability only — never a Hub API
|
||||
# probe. The old `model_info(checkpoint)` probe proved the repo
|
||||
|
||||
@@ -289,3 +289,46 @@ def test_the_sweep_interval_can_be_shortened_with_the_threshold(monkeypatch):
|
||||
|
||||
monkeypatch.setenv("OMNIVOICE_IDLE_SWEEP_SECONDS", "0")
|
||||
assert agent._sweep_seconds_from_env() == 60.0
|
||||
|
||||
|
||||
# ── Preload on a node ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_a_worker_machine_does_not_preload_a_model(monkeypatch, caplog):
|
||||
"""A node has no local user to warm the model for.
|
||||
|
||||
The startup preload exists so the first /generate feels instant for the
|
||||
person in front of the app. On a headless GPU node there is nobody there,
|
||||
so it is several GB held from boot against a request that may never
|
||||
arrive — and the idle sweep cannot reclaim it, because the sweep owns the
|
||||
worker executor's engines while this is the default local model. Observed
|
||||
on hardware: a node that had run nothing still sat at 2.4 GB.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
from services import model_manager
|
||||
|
||||
monkeypatch.setenv("OMNIVOICE_WORKER_MODE", "1")
|
||||
monkeypatch.setattr(model_manager, "model", None)
|
||||
loaded = {"count": 0}
|
||||
monkeypatch.setattr(
|
||||
model_manager,
|
||||
"_checkpoint_in_local_cache",
|
||||
lambda *a, **k: loaded.__setitem__("count", loaded["count"] + 1) or True,
|
||||
)
|
||||
|
||||
with caplog.at_level("INFO"):
|
||||
asyncio.run(model_manager.preload_model())
|
||||
|
||||
assert loaded["count"] == 0, "a worker machine still went looking for a model to preload"
|
||||
assert "loads on first request" in caplog.text
|
||||
|
||||
|
||||
def test_a_desktop_machine_still_preloads(monkeypatch):
|
||||
"""A machine that is both an app and a worker keeps the warm-up — there is
|
||||
a real user in front of it and the whole point of preloading stands."""
|
||||
from services import model_manager
|
||||
|
||||
monkeypatch.delenv("OMNIVOICE_WORKER_MODE", raising=False)
|
||||
|
||||
assert model_manager._headless_worker() is False
|
||||
|
||||
Reference in New Issue
Block a user