diff --git a/CHANGELOG.md b/CHANGELOG.md index 51e0319..f7c1fb2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,14 @@ # Changelog +## 1.4.0 (2026-10-25) + +Models can leave the GPU on their own once the server is idle. + +### Added + +- Admin option *Unload models after idle*: once nothing has run for that many minutes, RigShare unloads the models, as the *Unload models* button does, so the VRAM is free for other programs. Off by default (`auto_unload_minutes`: `0`). A run or a waiting prompt restarts the clock, and it unloads once per idle period +- With *Also clear the cache when unloading* (`auto_unload_cache`), it clears ComfyUI's cache too, like *Free models & cache*: RAM is freed as well, but the next run computes every node again + ## 1.3.3 (2026-10-08) Login and request hardening, mainly for servers behind a reverse proxy. diff --git a/docs/settings.md b/docs/settings.md index 0fc5d12..76ac9a4 100644 --- a/docs/settings.md +++ b/docs/settings.md @@ -36,6 +36,8 @@ Admins change these under **Admin → Server settings**. They are stored in `rig | `snapshot_interval_sec` | `300` | Automatic snapshot interval while editing; `0` turns it off | | `snapshot_keep` | `30` | Automatic snapshots kept per workflow, and as many on-save ones; named ones are always kept | | `chat_keep` | `0` | Chat messages kept: `0` keeps the whole history, a number trims it to the latest ones at startup. Admins can also clear it under *Admin → Chat* | +| `auto_unload_minutes` | `0` | Unload models from VRAM once nothing has run for this many minutes, as the *Unload models* button does; `0` turns it off. A run or a waiting prompt restarts the clock | +| `auto_unload_cache` | `false` | When unloading, also clear ComfyUI's cache (*Free models & cache*): frees RAM too, but the next run computes every node again | ## Environment variables diff --git a/pyproject.toml b/pyproject.toml index ba67011..d285455 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-rigshare" description = "Turn one ComfyUI server into a shared team workspace. Sign-in with accounts, permissions (edit, queue, Manager, admin) and personal API keys, enforced by the server. A private folder for each person plus shared folders with access lists, with Save asking where. Live co-editing of every shared workflow, with cursors, follow mode and conflict-free merges, in the graph or App view. Per-workflow history with a snapshot on every save, team chat with its whole history, and a CPU, RAM, GPU and queue monitor." -version = "1.3.3" +version = "1.4.0" license = { file = "LICENSE" } dependencies = [] diff --git a/rigshare_core/hub.py b/rigshare_core/hub.py index 41cad8b..768cc62 100644 --- a/rigshare_core/hub.py +++ b/rigshare_core/hub.py @@ -110,6 +110,8 @@ class Hub: self.workspace = None # set by __init__ (folders, private spaces) self.queue_owners = OrderedDict() # prompt_id -> {user, name, workflow, at} self.run_started = {} # prompt_id -> ms timestamp of its execution_start (set by routes) + self.last_activity = time.time() # last run event or busy queue, for auto unload + self.models_idle = True # nothing to unload until a run loads models # ----- lifecycle ---------------------------------------------------- @@ -135,6 +137,7 @@ class Hub: if expiry > 0 and now - room.updated > expiry and not self.members(room.key): self.rooms.pop(room.key, None) self.store.delete_room(room.key) + self.auto_unload(now) except Exception as e: log.error(f"[RigShare] housekeeping error: {e}") @@ -183,6 +186,8 @@ class Hub: prompt_id = data.get("prompt_id") if isinstance(data, dict) else None if not prompt_id: return + self.last_activity = time.time() + self.models_idle = False stamp = data.get("timestamp") or int(time.time() * 1000) if event == "execution_start": self.run_started[prompt_id] = stamp @@ -221,6 +226,31 @@ class Hub: return {"running": running, "recent": list(reversed(self.store.runs()[-recent:])), "etas": etas, "now": now} + def auto_unload(self, now): + """Unload models once nothing has run for ``auto_unload_minutes`` (0 = off). + + Sets the same queue flags as ComfyUI's /free, so its prompt worker does the + unloading between runs. Happens once per idle period. + """ + try: + minutes = float(self.store.config.get("auto_unload_minutes") or 0) + except (TypeError, ValueError): + minutes = 0 + if minutes <= 0 or self.models_idle or not self.server: + return + queue = self.server.prompt_queue + running, pending = queue.get_current_queue_volatile() + if running or pending: + self.last_activity = now + return + if now - self.last_activity < minutes * 60: + return + queue.set_flag("unload_models", True) + if self.store.config.get("auto_unload_cache"): + queue.set_flag("free_memory", True) + self.models_idle = True + log.info(f"[RigShare] Unloaded models after {minutes:g} min without runs.") + def running_prompt_ids(self): try: running, _ = self.server.prompt_queue.get_current_queue_volatile() diff --git a/rigshare_core/store.py b/rigshare_core/store.py index 6e3a519..b6f4605 100644 --- a/rigshare_core/store.py +++ b/rigshare_core/store.py @@ -58,6 +58,10 @@ DEFAULT_CONFIG = { "snapshot_interval_sec": 300, "snapshot_keep": 30, "chat_keep": 0, # 0 keeps the whole chat; N trims it to the last N messages at startup + # Unload models once the queue has been empty this many minutes (0 = off), + # and with auto_unload_cache also clear ComfyUI's node cache. + "auto_unload_minutes": 0, + "auto_unload_cache": False, } diff --git a/tests/python/helpers.py b/tests/python/helpers.py index 85d4047..09892e2 100644 --- a/tests/python/helpers.py +++ b/tests/python/helpers.py @@ -38,10 +38,14 @@ class FakeQueue: def __init__(self): self.running, self.pending, self.number = [], [], 0 + self.flags = {} def get_current_queue_volatile(self): return list(self.running), list(self.pending) + def set_flag(self, name, data): + self.flags[name] = data + history = {} def get_history(self, prompt_id=None, max_items=None, offset=-1, map_function=None): diff --git a/tests/python/test_server.py b/tests/python/test_server.py index ea51d56..bfda57e 100644 --- a/tests/python/test_server.py +++ b/tests/python/test_server.py @@ -650,6 +650,51 @@ def test_performance_run_log_eta_history_and_clear(run): run(scenario()) +def test_auto_unload_after_idle(run): + async def scenario(): + async with rig_server() as rig: + import time + admin = await rig.setup_admin() + hub, queue = rig.hub, rig.ps.prompt_queue + + def run(pid): + hub.run_event("execution_start", {"prompt_id": pid, "timestamp": 1}) + hub.run_event("execution_success", {"prompt_id": pid, "timestamp": 2}) + + r = await rig.http.patch(rig.url("/rigshare/api/config"), json={"auto_unload_minutes": 15}, headers=admin) + assert (await r.json())["auto_unload_minutes"] == 15 + hub.auto_unload(time.time() + 3600) + assert queue.flags == {}, "nothing ran since startup, nothing to unload" + + run("a") + now = time.time() + hub.auto_unload(now + 14 * 60) + assert queue.flags == {}, "not idle long enough" + queue.pending.append((1, "b", {}, {}, [])) + hub.auto_unload(now + 3600) + assert queue.flags == {}, "a waiting prompt keeps models loaded and restarts the clock" + queue.pending.clear() + hub.auto_unload(now + 3600 + 14 * 60) + assert queue.flags == {} + hub.auto_unload(now + 3600 + 15 * 60) + assert queue.flags == {"unload_models": True} + queue.flags.clear() + hub.auto_unload(now + 4 * 3600) + assert queue.flags == {}, "once per idle period" + + await rig.http.patch(rig.url("/rigshare/api/config"), json={"auto_unload_cache": True}, headers=admin) + run("c") + hub.auto_unload(time.time() + 15 * 60) + assert queue.flags == {"unload_models": True, "free_memory": True} + + queue.flags.clear() + await rig.http.patch(rig.url("/rigshare/api/config"), json={"auto_unload_minutes": 0}, headers=admin) + run("d") + hub.auto_unload(time.time() + 24 * 3600) + assert queue.flags == {}, "0 turns it off" + run(scenario()) + + def test_folder_routes_refuse_paths_outside_the_workspace(run, tmp_path): async def scenario(): async with rig_server() as rig: diff --git a/web/rigshare/panel.js b/web/rigshare/panel.js index 850e3d7..e1190d4 100644 --- a/web/rigshare/panel.js +++ b/web/rigshare/panel.js @@ -861,6 +861,8 @@ export class RigSharePanel { trusted: text((cfg.api_protection?.trusted_ips || []).join(", "), "127.0.0.1, 172.17.0.0/16"), interval: h("input", { class: "rs-input rs-num", type: "number", min: 0, value: Math.round((cfg.snapshot_interval_sec ?? 300) / 60) }), keep: h("input", { class: "rs-input rs-num", type: "number", min: 1, value: cfg.snapshot_keep ?? 30 }), + unloadMinutes: h("input", { class: "rs-input rs-num", type: "number", min: 0, value: cfg.auto_unload_minutes ?? 0 }), + unloadCache: box(cfg.auto_unload_cache), }; const line = (input, label, hint) => h("label", { class: "rs-check rs-block" }, input, h("span", {}, label, hint ? h("div", { class: "rs-muted rs-small" }, hint) : null)); @@ -880,6 +882,8 @@ export class RigSharePanel { api_protection: { enabled: f.protect.checked, trusted_ips: f.trusted.value.split(/[\s,]+/).filter(Boolean) }, snapshot_interval_sec: Math.max(0, +f.interval.value || 0) * 60, snapshot_keep: Math.max(1, +f.keep.value || 30), + auto_unload_minutes: Math.max(0, +f.unloadMinutes.value || 0), + auto_unload_cache: f.unloadCache.checked, }), "Settings saved"); if (ok) { this.client.server.name = ok.server_name; this.renderHeader(); } }, @@ -896,6 +900,8 @@ export class RigSharePanel { h("label", { class: "rs-label" }, "Trusted IPs / networks (skip login and checks)"), f.trusted, h("div", { class: "rs-row" }, h("label", { class: "rs-label rs-grow" }, "Auto snapshot every (min, 0 = off)"), f.interval), h("div", { class: "rs-row" }, h("label", { class: "rs-label rs-grow" }, "Automatic snapshots kept"), f.keep), + h("div", { class: "rs-row" }, h("label", { class: "rs-label rs-grow" }, "Unload models after idle (min, 0 = off)"), f.unloadMinutes), + line(f.unloadCache, "Also clear the cache when unloading", "Frees RAM as well, but the next run computes every node again"), h("button", { class: "rs-btn rs-btn-primary", type: "submit" }, h("i", { class: "pi pi-save" }), "Save settings")); } }