Vito Sansevero
489d449ec8
- py/model_gate.py: one process-wide, non-blocking gate for loading,
swapping, running, downloading and unloading a tagging model. Worker
calls run shielded; a release requested while one still runs takes
effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
gate and answer 409 {code: "model-busy"} while another operation holds
it; their bodies are otherwise unchanged. /autotag/models gains
suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
first, GGUF/HF (name, None) in output order, first `limit` distinct
case-insensitively; generate_tags keeps its return; the image file is
closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
pinned capture, with strict per-engine options, never downloads
(409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
metadata over the inspection budgets no longer blocks it.
149 lines
4.7 KiB
Python
149 lines
4.7 KiB
Python
"""One process-wide, non-blocking gate around the tagging models (#249 B2).
|
|
|
|
Loading, swapping, running, downloading and unloading a tagging model all
|
|
happen under this gate, whether the caller is a legacy Auto-tag route, the
|
|
Tag dialog's Generate or (later) a bulk job. A caller that cannot take the
|
|
gate is refused at once (``409 model-busy``), never queued.
|
|
|
|
Work that runs in a worker thread is admitted with :meth:`ModelGate.run`. A
|
|
release requested while such work still runs (a client disconnect cancels
|
|
the waiting coroutine, not the thread) takes effect only when the last
|
|
admitted call really exits, so nobody can swap the model under a running
|
|
inference.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
import uuid
|
|
from typing import Any, Callable, Optional
|
|
|
|
|
|
class GateNotHeld(RuntimeError):
|
|
"""Raised when work is admitted for an owner that does not hold the gate."""
|
|
|
|
|
|
class ModelGate:
|
|
"""A single-owner gate whose release waits for admitted worker calls."""
|
|
|
|
def __init__(self) -> None:
|
|
self._lock = threading.Lock()
|
|
self._holder: Optional[str] = None
|
|
self._inflight = 0
|
|
self._releasing = False
|
|
|
|
def try_acquire(self, owner: str) -> bool:
|
|
"""Take the gate for *owner*; False at once when anyone holds it."""
|
|
with self._lock:
|
|
if self._holder is not None:
|
|
return False
|
|
self._holder = owner
|
|
self._inflight = 0
|
|
self._releasing = False
|
|
return True
|
|
|
|
def release(self, owner: str) -> None:
|
|
"""Let go; deferred until admitted work of *owner* has exited."""
|
|
with self._lock:
|
|
if self._holder != owner:
|
|
return
|
|
if self._inflight:
|
|
self._releasing = True
|
|
else:
|
|
self._clear_locked()
|
|
|
|
def holder(self) -> Optional[str]:
|
|
with self._lock:
|
|
return self._holder
|
|
|
|
def begin(self, owner: str) -> None:
|
|
"""Admit one worker call; the holder must still hold the gate."""
|
|
with self._lock:
|
|
if self._holder != owner or self._releasing:
|
|
raise GateNotHeld(f"{owner} does not hold the model gate")
|
|
self._inflight += 1
|
|
|
|
def end(self, owner: str) -> None:
|
|
"""One admitted call exited; completes a deferred release."""
|
|
with self._lock:
|
|
self._end_locked(owner)
|
|
|
|
async def run(self, owner: str, func: Callable[..., Any], *args, **kwargs):
|
|
"""Run ``func`` in the default executor while *owner* holds the gate.
|
|
|
|
The executor future is shielded: cancelling the caller never cancels
|
|
a call that has not started yet (which would leave it admitted
|
|
forever) and never ends the admission of a call that is running.
|
|
"""
|
|
self.begin(owner)
|
|
state = {"started": False, "finished": False}
|
|
|
|
def finish() -> None:
|
|
with self._lock:
|
|
if not state["finished"]:
|
|
state["finished"] = True
|
|
self._end_locked(owner)
|
|
|
|
def work():
|
|
with self._lock:
|
|
state["started"] = True
|
|
try:
|
|
return func(*args, **kwargs)
|
|
finally:
|
|
finish()
|
|
|
|
def never_ran(future) -> None:
|
|
# Only a loop shutdown cancels the shielded future; a call that
|
|
# never started must still give its admission back.
|
|
if future.cancelled():
|
|
with self._lock:
|
|
skipped = not state["started"] and not state["finished"]
|
|
if skipped:
|
|
state["finished"] = True
|
|
self._end_locked(owner)
|
|
|
|
try:
|
|
future = asyncio.get_running_loop().run_in_executor(None, work)
|
|
except BaseException:
|
|
finish()
|
|
raise
|
|
future.add_done_callback(never_ran)
|
|
return await asyncio.shield(future)
|
|
|
|
def _end_locked(self, owner: str) -> None:
|
|
if self._holder != owner or self._inflight == 0:
|
|
return
|
|
self._inflight -= 1
|
|
if self._releasing and self._inflight == 0:
|
|
self._clear_locked()
|
|
|
|
def _clear_locked(self) -> None:
|
|
self._holder = None
|
|
self._inflight = 0
|
|
self._releasing = False
|
|
|
|
|
|
_GATE = ModelGate()
|
|
|
|
|
|
def new_owner(kind: str) -> str:
|
|
"""A unique owner name such as ``legacy-single:<uuid>``."""
|
|
return f"{kind}:{uuid.uuid4()}"
|
|
|
|
|
|
def try_acquire(owner: str) -> bool:
|
|
return _GATE.try_acquire(owner)
|
|
|
|
|
|
def release(owner: str) -> None:
|
|
_GATE.release(owner)
|
|
|
|
|
|
def holder() -> Optional[str]:
|
|
return _GATE.holder()
|
|
|
|
|
|
async def run(owner: str, func: Callable[..., Any], *args, **kwargs):
|
|
return await _GATE.run(owner, func, *args, **kwargs)
|