50 lines
1.9 KiB
Python
50 lines
1.9 KiB
Python
"""Backend control operations (model unload).
|
|
|
|
llama-server and llama-swap evict a resident model through different routes:
|
|
* llama-server → ``POST {base}/models/unload`` with body ``{"model": id}``
|
|
* llama-swap → ``POST {base}/api/models/unload/{id}`` (path parameter)
|
|
|
|
``unload_model`` dispatches on the configured backend type so callers don't
|
|
have to know which one they are talking to. Both routes live at the endpoint
|
|
root, so any ``/v1`` suffix is stripped first.
|
|
"""
|
|
from typing import Optional
|
|
|
|
import aiohttp
|
|
|
|
from config import get_config
|
|
from state import default_headers
|
|
from backends.sessions import get_probe_session
|
|
from backends.normalize import is_llama_swap
|
|
from backends.health import _format_connection_issue
|
|
|
|
|
|
async def unload_model(endpoint: str, model_id: str) -> bool:
|
|
"""Ask ``endpoint`` to unload ``model_id``. Returns True on a 2xx response.
|
|
|
|
``model_id`` must be the backend's native model identifier (the raw HF id
|
|
for llama-server / llama-swap), not the router-normalized display name.
|
|
"""
|
|
cfg = get_config()
|
|
base_url = endpoint.rstrip("/").removesuffix("/v1")
|
|
headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")}
|
|
api_key: Optional[str] = cfg.api_keys.get(endpoint)
|
|
if api_key is not None:
|
|
headers["Authorization"] = "Bearer " + api_key
|
|
|
|
if is_llama_swap(endpoint):
|
|
url = f"{base_url}/api/models/unload/{model_id}"
|
|
json_body = None
|
|
else:
|
|
url = f"{base_url}/models/unload"
|
|
json_body = {"model": model_id}
|
|
|
|
client: aiohttp.ClientSession = get_probe_session(endpoint)
|
|
try:
|
|
async with client.post(url, json=json_body, headers=headers) as resp:
|
|
ok = resp.status < 400
|
|
print(f"[unload_model] {model_id} on {endpoint}: {resp.status}")
|
|
return ok
|
|
except Exception as e:
|
|
print(f"[unload_model] {_format_connection_issue(url, e)}")
|
|
return False
|