jpmschweitzer/odysseus
1
0
Fork 0
mirror of https://github.com/pewdiepie-archdaemon/odysseus.git synced 2026-11-01 10:42:23 +01:00
Code Issues Packages Projects Releases Wiki Activity

fix(cookbook): preserve endpoints across model launches (#6969)

Browse Source 
* fix(cookbook): match stale endpoints by launched model, not display name

After registering an LLM endpoint, the Cookbook deleted every other local-*
endpoint with the same display name at a different URL. The display name is
the repo short name, which repos from different orgs share (for example
meta-llama/Llama-3.1-8B-Instruct and unsloth/Llama-3.1-8B-Instruct), so
launching one deleted the endpoint of the other while its server was running.

Match stale endpoints by the full model id they were last launched for,
which is the last entry of their pinned_models. Reusing an endpoint now
moves the launched model to the end of its pins so that entry stays current.

* fix(cookbook): guard endpoint cleanup by serve ownership

Preserve independently launched models instead of sweeping endpoints by display name or mutable model pins.

Bind Cookbook endpoint registrations to their serve sessions and prevent stale watchdogs or scheduled stops from deleting newer registrations.

Reuse the existing endpoint deletion contract to preserve settings, preferences, session cleanup, and cache invalidation.

Cover LLM and image endpoint lifecycle regressions.

---------

Co-authored-by: Alexandre Teixeira <general@alexandreteixeira.dev>
This commit is contained in:
Yu-Wei Chang
2026-11-01 03:06:13 +00:00
committed by GitHub
co-authored by Alexandre Teixeira
parent 39303f19cd
commit 6f2b58c896
6 changed files with 805 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/database.py
+22
View File
@@ -568,6 +568,7 @@ class ModelEndpoint(TimestampMixin, Base):
hidden_models = Column(Text, nullable=True) # JSON list of model IDs that failed probing
cached_models = Column(Text, nullable=True) # JSON list of last-known model IDs (avoids probe on list)
pinned_models = Column(Text, nullable=True) # JSON list of admin-pinned model IDs (manual, may not appear in /v1/models)
cookbook_serve_session_id = Column(String, nullable=True) # Current Cookbook serve identity
model_type = Column(String, nullable=True, default="llm") # "llm" or "image"
# auto = classify by URL; local = self-hosted server; api/proxy = external
# OpenAI-compatible API even when reachable through a private/tailnet IP.
@@ -1378,6 +1379,26 @@ def _migrate_add_pinned_models_column():
except Exception:
pass
def _migrate_add_cookbook_serve_session_id_column():
"""Bind Cookbook cleanup to the serve session that owns the endpoint."""
from sqlalchemy import inspect
with engine.begin() as conn:
columns = {
column["name"]
for column in inspect(conn).get_columns("model_endpoints")
}
if "cookbook_serve_session_id" not in columns:
conn.execute(
text(
"ALTER TABLE model_endpoints "
"ADD COLUMN cookbook_serve_session_id VARCHAR"
)
)
def _migrate_add_notes_sort_order():
"""Add sort_order, image_url, repeat columns to notes if they don't exist."""
import sqlite3
@@ -2390,6 +2411,7 @@ def init_db():
_migrate_add_hidden_models_column()
_migrate_add_cached_models_column()
_migrate_add_pinned_models_column()
_migrate_add_cookbook_serve_session_id_column()
_migrate_add_notes_sort_order()
_migrate_add_model_type_column()
_migrate_add_model_endpoint_refresh_columns()
routes/cookbook_routes.py
+86 -66
View File
@@ -1408,7 +1408,11 @@ def setup_cookbook_routes() -> APIRouter:
return {"models": models, "host": host or "local"}
def _auto_register_image_endpoint(req: ServeRequest, remote: str | None) -> str | None:
def _auto_register_image_endpoint(
req: ServeRequest,
remote: str | None,
serve_session_id: str | None = None,
) -> str | None:
"""Register a diffusion model as an image endpoint so it appears in the model selector."""
import re
from core.database import SessionLocal, ModelEndpoint
@@ -1442,6 +1446,7 @@ def setup_cookbook_routes() -> APIRouter:
existing.name = display_name
existing.endpoint_kind = "local"
existing.model_refresh_mode = "manual"
existing.cookbook_serve_session_id = serve_session_id
if pinned_models:
existing.cached_models = json.dumps(pinned_models)
existing.pinned_models = json.dumps(pinned_models)
@@ -1465,6 +1470,7 @@ def setup_cookbook_routes() -> APIRouter:
model_refresh_mode="manual",
cached_models=json.dumps(pinned_models) if pinned_models else None,
pinned_models=json.dumps(pinned_models) if pinned_models else None,
cookbook_serve_session_id=serve_session_id,
)
db.add(ep)
db.commit()
@@ -1610,43 +1616,71 @@ def setup_cookbook_routes() -> APIRouter:
# let the probe layer mark it offline if nothing's listening.
logger.info(f"crash-watchdog: serve {session_id} exited cleanly (0); leaving endpoint {endpoint_id}")
return
# Non-zero exit — drop the endpoint.
# A nonzero exit may delete only the endpoint owned by this serve.
try:
from core.database import SessionLocal as _SL, ModelEndpoint as _ME
from core.constants import internal_api_base as _internal_api_base
from core.middleware import INTERNAL_TOOL_HEADER, INTERNAL_TOOL_TOKEN
from urllib.parse import quote
import httpx
db = _SL()
try:
ep = db.query(_ME).filter(_ME.id == endpoint_id).first()
if ep:
# A scheduled serve can leave old non-zero exit markers
# in tmux scrollback while the current OpenAI endpoint is
# actually alive. Verify reachability before deleting the
# endpoint row; otherwise chats fall back even though the
# served model is ready.
try:
probe_url = ep.base_url.rstrip("/") + "/models"
with urllib.request.urlopen(probe_url, timeout=3) as resp:
if 200 <= getattr(resp, "status", 0) < 300:
logger.info(
f"crash-watchdog: serve {session_id} has exit marker {exit_code} "
f"but endpoint {ep.id} is reachable; leaving it registered"
)
return
except Exception:
pass
if not ep or ep.cookbook_serve_session_id != session_id:
logger.info(
f"crash-watchdog: dropping endpoint {endpoint_id} "
f"({ep.name} @ {ep.base_url}) — serve exited {exit_code}"
"crash-watchdog: endpoint %s no longer owned by %s",
endpoint_id, session_id,
)
db.delete(ep)
db.commit()
return
probe_url = ep.base_url.rstrip("/") + "/models"
finally:
db.close()
# Preserve an endpoint that is already healthy, including
# after an old runner emits a stale failure marker.
try:
with urllib.request.urlopen(probe_url, timeout=3) as resp:
if 200 <= getattr(resp, "status", 0) < 300:
logger.info(
"crash-watchdog: endpoint %s is reachable; keeping it",
endpoint_id,
)
return
except Exception:
pass
async with httpx.AsyncClient(timeout=8) as client:
response = await client.delete(
f"{_internal_api_base()}/api/model-endpoints/{quote(endpoint_id, safe='')}",
params={"expected_cookbook_serve_session_id": session_id},
headers={INTERNAL_TOOL_HEADER: INTERNAL_TOOL_TOKEN},
)
if response.status_code == 200:
logger.info(
"crash-watchdog: deleted endpoint %s after serve %s exited %s",
endpoint_id, session_id, exit_code,
)
elif response.status_code in (404, 409):
logger.info(
"crash-watchdog: endpoint %s absent or reassigned; keeping it",
endpoint_id,
)
else:
logger.warning(
"crash-watchdog: endpoint %s deletion returned HTTP %s",
endpoint_id, response.status_code,
)
except Exception as e:
logger.warning(f"crash-watchdog: endpoint cleanup failed: {e!r}")
logger.warning("crash-watchdog: endpoint cleanup failed: %r", e)
return
logger.debug(f"crash-watchdog: no exit marker for {session_id} within window; leaving endpoint {endpoint_id}")
def _auto_register_llm_endpoint(req: ServeRequest, remote: str | None) -> str | None:
def _auto_register_llm_endpoint(
req: ServeRequest,
remote: str | None,
serve_session_id: str | None = None,
) -> str | None:
"""Register a freshly-served LLM as a model endpoint so it appears in the
model picker without a manual /setup step — the text-model sibling of
_auto_register_image_endpoint.
@@ -1663,6 +1697,7 @@ def setup_cookbook_routes() -> APIRouter:
)
import re
from core.database import SessionLocal, ModelEndpoint
from routes.model_routes import _normalize_model_ids
# Port: ordered fallbacks so we match whatever the user actually
# asked for, not a hardcoded default:
@@ -1724,8 +1759,12 @@ def setup_cookbook_routes() -> APIRouter:
# deployment alias from /v1/models, and a stale server can answer on the
# same port while the new launch failed. Keeping the requested model id
# pinned makes the picker reflect the actual launch intent.
pinned_models = [mlx_shim_model_id] if mlx_shim_model_id else ([req.repo_id] if req.repo_id else [])
launched_model = mlx_shim_model_id or req.repo_id
pinned_models = [launched_model] if launched_model else []
# Launch registration must not delete other endpoints.
# Matching names, model pins, and failed probes do not prove that
# another serve has stopped. Lifecycle cleanup owns termination.
db = SessionLocal()
try:
# Reuse an endpoint already pointed at this URL instead of duplicating.
@@ -1736,16 +1775,16 @@ def setup_cookbook_routes() -> APIRouter:
existing.name = display_name
existing.endpoint_kind = "local"
existing.model_refresh_mode = "auto"
if pinned_models:
try:
existing_pinned = json.loads(existing.pinned_models or "[]")
except Exception:
existing_pinned = []
merged_pinned = []
for mid in [*existing_pinned, *pinned_models]:
if mid and mid not in merged_pinned:
merged_pinned.append(mid)
existing.pinned_models = json.dumps(merged_pinned) if merged_pinned else None
# A reused endpoint belongs to the newest serve session.
existing.cookbook_serve_session_id = serve_session_id
if launched_model:
# Pins represent model selection, not launch history.
existing_pins = _normalize_model_ids(
existing.pinned_models
)
if launched_model not in existing_pins:
existing_pins.append(launched_model)
existing.pinned_models = json.dumps(existing_pins)
if is_ollama_endpoint:
existing.endpoint_kind = "ollama"
if pinned_models:
@@ -1771,22 +1810,6 @@ def setup_cookbook_routes() -> APIRouter:
db.commit()
except Exception as _pe:
logger.warning(f"Re-probe failed for {base_url}: {_pe!r}")
# Sweep stale dupes: other endpoints with the same display name
# at DIFFERENT URLs (likely failed earlier-attempt ports) get
# deleted so the picker doesn't show an offline ghost next to
# the working one. Only sweeps endpoints whose id starts with
# `local-` so we never touch a user's hand-added DeepSeek/OpenAI/
# etc. entry with a coincidentally matching name.
stale = (db.query(ModelEndpoint)
.filter(ModelEndpoint.name == display_name)
.filter(ModelEndpoint.base_url != base_url)
.filter(ModelEndpoint.id.like("local-%"))
.all())
for s in stale:
logger.info(f"Sweeping stale local endpoint {s.id} ({s.base_url})")
db.delete(s)
if stale:
db.commit()
return existing.id
ep_id = f"local-{uuid.uuid4().hex[:8]}"
@@ -1802,22 +1825,11 @@ def setup_cookbook_routes() -> APIRouter:
cached_models=json.dumps(pinned_models) if pinned_models else None,
pinned_models=json.dumps(pinned_models) if pinned_models else None,
supports_tools=supports_tools,
cookbook_serve_session_id=serve_session_id,
)
db.add(ep)
db.commit()
logger.info(f"Auto-registered local model endpoint: {display_name} @ {base_url}")
# Same sweep on first-register path: drop any pre-existing local-*
# endpoints with this display name pointed elsewhere.
stale = (db.query(ModelEndpoint)
.filter(ModelEndpoint.name == display_name)
.filter(ModelEndpoint.id != ep_id)
.filter(ModelEndpoint.id.like("local-%"))
.all())
for s in stale:
logger.info(f"Sweeping stale local endpoint {s.id} ({s.base_url})")
db.delete(s)
if stale:
db.commit()
# Probe /v1/models NOW and write cached_models so the chat
# picker actually shows the model on the next /api/models
# call. Without this immediate probe, the endpoint has empty
@@ -2639,9 +2651,17 @@ def setup_cookbook_routes() -> APIRouter:
endpoint_id = None
is_image_endpoint = "diffusion_server.py" in req.cmd or "mlx_image_server.py" in req.cmd
if is_image_endpoint:
endpoint_id = _auto_register_image_endpoint(req, remote)
endpoint_id = _auto_register_image_endpoint(
req,
remote,
serve_session_id=session_id,
)
elif not is_pip_install:
endpoint_id = _auto_register_llm_endpoint(req, remote)
endpoint_id = _auto_register_llm_endpoint(
req,
remote,
serve_session_id=session_id,
)
# Crash watchdog: the auto-register above writes the endpoint row
# IMMEDIATELY (before the server has even bound its port) so the
routes/model_routes.py
+26 -2
View File
@@ -3044,20 +3044,44 @@ def setup_model_routes(model_discovery):
return {"dependents": _settings_using_endpoint(ep_id)}
@router.delete("/model-endpoints/{ep_id}")
def delete_model_endpoint(ep_id: str, request: Request):
def delete_model_endpoint(
ep_id: str,
request: Request,
expected_cookbook_serve_session_id: Optional[str] = None,
):
require_admin(request)
db = SessionLocal()
try:
ep = db.query(ModelEndpoint).filter(ModelEndpoint.id == ep_id).first()
if not ep or not _chatgpt_endpoint_visible(ep, request):
raise HTTPException(404, "Endpoint not found")
if expected_cookbook_serve_session_id is not None:
# The DELETE itself checks ownership. An old serve must never
# mutate the settings of an endpoint reused by another serve.
if (
not ep_id.startswith(("local-", "img-"))
or not expected_cookbook_serve_session_id
):
raise HTTPException(409, "Cookbook endpoint ownership mismatch")
deleted = (
db.query(ModelEndpoint)
.filter(
ModelEndpoint.id == ep_id,
ModelEndpoint.cookbook_serve_session_id
== expected_cookbook_serve_session_id,
)
.delete(synchronize_session=False)
)
if not deleted:
raise HTTPException(409, "Cookbook endpoint ownership changed")
# Clean up any settings that reference this endpoint
cleared = _clear_settings_for_endpoint(ep_id)
cleared_user_preferences = _clear_user_prefs_for_endpoint(ep_id)
cleared_sessions = _clear_sessions_for_endpoint(db, ep.base_url, ep.id, ep.owner)
cleared_loaded_sessions = _clear_loaded_sessions_for_endpoint(ep.base_url, ep.id, ep.owner)
auth_id = getattr(ep, "provider_auth_id", None)
db.delete(ep)
if expected_cookbook_serve_session_id is None:
db.delete(ep)
cleared_provider_auth = _delete_orphaned_provider_auth(db, auth_id, exclude_ep_id=ep_id)
db.commit()
_invalidate_models_cache()
src/cookbook_serve_lifecycle.py
+31 -51
View File
@@ -31,63 +31,43 @@ def _internal_headers() -> dict:
async def _delete_endpoint_for_task(task: dict) -> None:
"""Drop the auto-registered model endpoint for a scheduled-stop serve.
"""Remove a stopped serve's endpoint using the normal deletion contract."""
import re
from urllib.parse import quote
Without this, killing the tmux session leaves the endpoint sitting in
the picker (probe goes offline; chats still try to route there) and
the user has to delete it by hand in Settings -> Endpoints.
"""
endpoint_id = (task.get("_endpointId") or task.get("endpointId") or "").strip()
if not endpoint_id:
logger.info(
"cookbook_serve_lifecycle: task %s has no endpoint id; skipping endpoint deletion",
task.get("sessionId") or task.get("id") or "",
)
endpoint_id = str(task.get("_endpointId") or task.get("endpointId") or "").strip()
serve_session_id = str(task.get("sessionId") or task.get("id") or "").strip()
if (
not re.fullmatch(r"(?:local|img)-[0-9a-f]{8}", endpoint_id)
or not serve_session_id
):
logger.info("cookbook_serve_lifecycle: missing endpoint/session identity")
return
import re as _re
payload = task.get("payload") or {}
cmd = str(payload.get("_cmd") or "")
remote = task.get("remoteHost") or ""
# Build host the same way _auto_register_llm_endpoint does so URL match wins.
if remote:
host = remote.split("@")[-1] if "@" in remote else remote
else:
host = "host.docker.internal"
port_match = _re.search(r"--port\s+(\d+)", cmd)
ollama_host_match = _re.search(r"OLLAMA_HOST=[^\s]*?:(\d+)", cmd)
if port_match:
port = int(port_match.group(1))
elif ollama_host_match:
port = int(ollama_host_match.group(1))
elif "ollama" in cmd:
port = 11434
else:
port = 8080
base_url = f"http://{host}:{port}/v1"
try:
async with httpx.AsyncClient(timeout=8) as client:
r = await client.get(
f"{internal_api_base()}/api/model-endpoints",
response = await client.delete(
f"{internal_api_base()}/api/model-endpoints/{quote(endpoint_id, safe='')}",
params={"expected_cookbook_serve_session_id": serve_session_id},
headers=_internal_headers(),
)
if r.status_code >= 400:
return
eps = r.json() if r.content else []
# Delete only the endpoint created by this scheduled serve. URL
# matching is unsafe because a later scheduled serve can reuse the
# same host:port after an older task has gone stale.
ep = next((e for e in eps if e.get("id") == endpoint_id), None)
if ep:
await client.delete(
f"{internal_api_base()}/api/model-endpoints/{ep['id']}",
headers=_internal_headers(),
)
logger.info(
f"cookbook_serve_lifecycle: deleted endpoint {ep.get('id')} "
f"({ep.get('base_url')}) after scheduled stop"
)
except Exception as e:
logger.warning(f"cookbook_serve_lifecycle: endpoint delete failed: {e}")
if response.status_code == 200:
logger.info(
"cookbook_serve_lifecycle: deleted endpoint %s for serve %s",
endpoint_id, serve_session_id,
)
elif response.status_code in (404, 409):
logger.info(
"cookbook_serve_lifecycle: endpoint %s absent or reassigned; keeping it",
endpoint_id,
)
else:
logger.warning(
"cookbook_serve_lifecycle: endpoint %s deletion returned HTTP %s",
endpoint_id, response.status_code,
)
except Exception as exc:
logger.warning("cookbook_serve_lifecycle: cleanup failed: %s", exc)
async def _stop_serve(session_id: str, remote_host: str = "", ssh_port: str = "") -> bool:
tests/test_cookbook_endpoint_registration.py
+417
View File
@@ -1,5 +1,9 @@
from pathlib import Path
import pytest
from tests.helpers.database import disposable_database
ROOT = Path(__file__).resolve().parents[1]
COOKBOOK_RUNNING = ROOT / "static" / "js" / "cookbookRunning.js"
@@ -28,3 +32,416 @@ def test_cookbook_advertised_bind_urls_keep_connectable_host():
assert "function _endpointFromAdvertisedUrl" in src
assert "_isAnyBindHost(u.hostname) ? currentHost" in src
assert "host = u.hostname || host;" not in src
@pytest.fixture
def endpoint_db(tmp_path, monkeypatch):
import core.database as database
import routes.model_routes as model_routes
with disposable_database(tmp_path) as factory:
monkeypatch.setattr(database, "SessionLocal", factory)
monkeypatch.setattr(model_routes, "_probe_endpoint", lambda *args, **kwargs: None)
yield factory
def _launch_llm(
repo_id: str,
port: int,
serve_session_id: str | None = None,
) -> str:
import routes.cookbook_routes as cookbook_routes
from routes.cookbook_helpers import ServeRequest
router = cookbook_routes.setup_cookbook_routes()
serve = next(
route.endpoint for route in router.routes
if route.path == "/api/model/serve" and "POST" in route.methods
)
cells = dict(zip(serve.__code__.co_freevars, serve.__closure__ or ()))
register = cells["_auto_register_llm_endpoint"].cell_contents
endpoint_id = register(
ServeRequest(repo_id=repo_id, cmd=f"vllm serve {repo_id} --port {port}"),
None,
serve_session_id=serve_session_id,
)
assert endpoint_id
return endpoint_id
def _endpoint(factory, endpoint_id: str):
import core.database as database
with factory() as db:
return db.get(database.ModelEndpoint, endpoint_id)
@pytest.mark.parametrize("reuse", [False, True])
def test_launch_keeps_endpoint_of_other_model_with_same_short_name(endpoint_db, reuse):
if reuse:
mirror = _launch_llm("unsloth/Llama-3.1-8B-Instruct", 8002)
original = _launch_llm("meta-llama/Llama-3.1-8B-Instruct", 8001)
launched = _launch_llm("unsloth/Llama-3.1-8B-Instruct", 8002)
if reuse:
assert launched == mirror
assert _endpoint(endpoint_db, original) is not None
assert _endpoint(endpoint_db, launched) is not None
@pytest.mark.parametrize("reuse", [False, True])
def test_relaunch_on_another_port_preserves_existing_endpoint(endpoint_db, reuse):
if reuse:
current = _launch_llm("meta-llama/Llama-3.1-8B-Instruct", 8002)
stale = _launch_llm("Qwen/Qwen3-8B", 8001)
launched = _launch_llm("Qwen/Qwen3-8B", 8002)
if reuse:
assert launched == current
assert _endpoint(endpoint_db, stale) is not None
assert _endpoint(endpoint_db, launched) is not None
def test_sweep_ignores_endpoint_now_serving_a_different_model(endpoint_db):
endpoint = _launch_llm("Qwen/Qwen3-8B", 8001)
_launch_llm("meta-llama/Llama-3.1-8B", 8001)
_launch_llm("Qwen/Qwen3-8B", 8001)
_launch_llm("meta-llama/Llama-3.1-8B", 8002)
assert _endpoint(endpoint_db, endpoint) is not None
def test_sweep_ignores_endpoints_not_created_by_cookbook(endpoint_db):
import core.database as database
with endpoint_db() as db:
db.add(database.ModelEndpoint(
id="manual", name="Qwen3-8B", base_url="http://localhost:9000/v1",
pinned_models='["Qwen/Qwen3-8B"]',
))
db.commit()
_launch_llm("Qwen/Qwen3-8B", 8002)
assert _endpoint(endpoint_db, "manual") is not None
def test_manual_pin_order_does_not_change_on_relaunch(endpoint_db):
import json
from core.database import ModelEndpoint
endpoint_id = _launch_llm("organization-a/Model-A", 8001)
with endpoint_db() as db:
endpoint = db.get(ModelEndpoint, endpoint_id)
endpoint.pinned_models = json.dumps([
"organization-a/Model-A",
"organization-b/Model-B",
])
db.commit()
_launch_llm("organization-a/Model-A", 8001)
with endpoint_db() as db:
endpoint = db.get(ModelEndpoint, endpoint_id)
assert json.loads(endpoint.pinned_models) == [
"organization-a/Model-A",
"organization-b/Model-B",
]
def test_manual_pins_cannot_trigger_endpoint_deletion(endpoint_db):
import json
from core.database import ModelEndpoint
original = _launch_llm("organization-a/Model-A", 8001)
with endpoint_db() as db:
endpoint = db.get(ModelEndpoint, original)
endpoint.pinned_models = json.dumps([
"organization-a/Model-A",
"organization-b/Model-B",
])
db.commit()
_launch_llm("organization-b/Model-B", 8002)
assert _endpoint(endpoint_db, original) is not None
def test_unavailable_model_probe_never_deletes_another_endpoint(
endpoint_db, monkeypatch
):
import routes.model_routes as model_routes
original = _launch_llm("organization-a/Model-A", 8001)
monkeypatch.setattr(
model_routes,
"_probe_endpoint",
lambda *args, **kwargs: [],
)
_launch_llm("organization-a/Model-A", 8002)
assert _endpoint(endpoint_db, original) is not None
@pytest.mark.asyncio
async def test_old_watchdog_cannot_delete_reused_endpoint(
endpoint_db,
monkeypatch,
endpoint_http_api,
):
import asyncio
import urllib.request
import routes.cookbook_routes as cookbook_routes
original = _launch_llm(
"organization-a/Model-A",
8001,
serve_session_id="serve-old",
)
reused = _launch_llm(
"organization-b/Model-B",
8001,
serve_session_id="serve-new",
)
assert original == reused
with endpoint_db() as db:
from core.database import ModelEndpoint
endpoint = db.get(ModelEndpoint, reused)
assert endpoint.cookbook_serve_session_id == "serve-new"
router = cookbook_routes.setup_cookbook_routes()
serve = next(
route.endpoint
for route in router.routes
if route.path == "/api/model/serve"
and "POST" in route.methods
)
cells = dict(
zip(
serve.__code__.co_freevars,
serve.__closure__ or (),
)
)
watchdog = cells["_serve_crash_watchdog"].cell_contents
class FakeProcess:
async def communicate(self):
return (
b"=== Process exited with code 1 ===\n",
b"",
)
async def fake_subprocess(*args, **kwargs):
return FakeProcess()
async def fake_sleep(seconds):
return None
def unavailable(*args, **kwargs):
raise OSError("Server still warming up")
monkeypatch.setattr(
asyncio,
"create_subprocess_exec",
fake_subprocess,
)
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
monkeypatch.setattr(
urllib.request,
"urlopen",
unavailable,
)
await watchdog(
endpoint_id=original,
session_id="serve-old",
remote=None,
ssh_port=None,
is_windows=False,
)
assert _endpoint(endpoint_db, reused) is not None
# The current serve's own crash watchdog must still be able
# to remove its endpoint when it actually fails.
await watchdog(
endpoint_id=reused,
session_id="serve-new",
remote=None,
ssh_port=None,
is_windows=False,
)
assert _endpoint(endpoint_db, reused) is None
def test_cookbook_serve_session_migration_is_idempotent(
tmp_path,
monkeypatch,
):
from sqlalchemy import create_engine, inspect, text
import core.database as database
engine = create_engine(
f"sqlite:///{tmp_path / 'legacy-endpoints.db'}"
)
try:
with engine.begin() as conn:
conn.execute(
text(
"CREATE TABLE model_endpoints "
"(id VARCHAR PRIMARY KEY)"
)
)
monkeypatch.setattr(database, "engine", engine)
database._migrate_add_cookbook_serve_session_id_column()
database._migrate_add_cookbook_serve_session_id_column()
columns = {
column["name"]
for column in inspect(engine).get_columns(
"model_endpoints"
)
}
assert "cookbook_serve_session_id" in columns
finally:
engine.dispose()
@pytest.fixture
def endpoint_http_api(endpoint_db, monkeypatch):
"""Route deletion through a disposable in-process FastAPI app."""
import httpx
from fastapi import FastAPI
import routes.model_routes as model_routes
import routes.prefs_routes as prefs_routes
import core.constants as constants
import src.cookbook_serve_lifecycle as lifecycle
monkeypatch.setattr(model_routes, "SessionLocal", endpoint_db)
monkeypatch.setattr(model_routes, "require_admin", lambda request: None)
settings = {}
preferences = {}
monkeypatch.setattr(model_routes, "_load_settings", lambda: settings)
monkeypatch.setattr(model_routes, "_save_settings", lambda value: settings.update(value))
monkeypatch.setattr(prefs_routes, "_load", lambda: preferences)
monkeypatch.setattr(prefs_routes, "_save", lambda value: preferences.update(value))
# Endpoint deletion must not touch the user's running sessions.
import src.ai_interaction as ai_interaction
monkeypatch.setattr(ai_interaction, "get_session_manager", lambda: None)
app = FastAPI()
router = model_routes.setup_model_routes(None)
app.include_router(router)
original_client = httpx.AsyncClient
transport = httpx.ASGITransport(app=app)
def client_factory(*args, **kwargs):
kwargs["transport"] = transport
return original_client(*args, **kwargs)
monkeypatch.setattr(httpx, "AsyncClient", client_factory)
monkeypatch.setattr(constants, "internal_api_base", lambda: "http://testserver")
monkeypatch.setattr(lifecycle, "internal_api_base", lambda: "http://testserver")
return {"settings": settings, "preferences": preferences, "router": router}
@pytest.mark.asyncio
async def test_ownership_condition_preserves_settings_on_mismatch(
endpoint_db, endpoint_http_api,
):
import httpx
original = _launch_llm("organization-a/Model-A", 8001, "serve-old")
reused = _launch_llm("organization-b/Model-B", 8001, "serve-new")
assert reused == original
settings = endpoint_http_api["settings"]
settings.update({"default_endpoint_id": reused, "default_model": "organization-b/Model-B"})
async with httpx.AsyncClient() as client:
response = await client.delete(
f"http://testserver/api/model-endpoints/{reused}",
params={"expected_cookbook_serve_session_id": "serve-old"},
)
assert response.status_code == 409
assert _endpoint(endpoint_db, reused) is not None
assert settings["default_endpoint_id"] == reused
@pytest.mark.asyncio
async def test_owned_deletion_clears_settings_and_model_cache(
endpoint_db, endpoint_http_api,
):
import httpx
from core.database import Session as DbSession
endpoint_id = _launch_llm("organization-a/Model-A", 8001, "serve-current")
settings = endpoint_http_api["settings"]
settings.update({
"default_endpoint_id": endpoint_id,
"default_model": "organization-a/Model-A",
"foreground_model_fallbacks": [{"endpoint_id": endpoint_id, "model": "x"}],
})
with endpoint_db() as db:
db.add(DbSession(
id="owned-delete-session", name="Test", owner=None,
endpoint_url="http://localhost:8001/v1",
endpoint_id=endpoint_id, model="organization-a/Model-A",
headers={"Authorization": "Bearer test-only"},
))
db.commit()
# Seed the actual model listing cache captured by the route closure.
router = endpoint_http_api["router"]
route = next(r for r in router.routes if r.path == "/api/model-endpoints/{ep_id}" and "DELETE" in r.methods)
closure = dict(zip(route.endpoint.__code__.co_freevars, route.endpoint.__closure__ or ()))
invalidate = closure["_invalidate_models_cache"].cell_contents
cache_closure = dict(zip(invalidate.__code__.co_freevars, invalidate.__closure__ or ()))
models_cache = cache_closure["_models_cache"].cell_contents
models_cache["test"] = {"data": [], "time": 0}
async with httpx.AsyncClient() as client:
response = await client.delete(
f"http://testserver/api/model-endpoints/{endpoint_id}",
params={"expected_cookbook_serve_session_id": "serve-current"},
)
assert response.status_code == 200, response.text
assert _endpoint(endpoint_db, endpoint_id) is None
assert settings["default_endpoint_id"] == ""
assert settings["foreground_model_fallbacks"] == []
assert not models_cache
with endpoint_db() as db:
sess = db.get(DbSession, "owned-delete-session")
assert sess.headers == {}
@pytest.mark.asyncio
async def test_ordinary_admin_delete_does_not_require_serve_owner(
endpoint_db, endpoint_http_api,
):
import httpx
endpoint_id = _launch_llm("organization-a/Model-A", 8001, "serve-new")
async with httpx.AsyncClient() as client:
response = await client.delete(f"http://testserver/api/model-endpoints/{endpoint_id}")
assert response.status_code == 200, response.text
assert _endpoint(endpoint_db, endpoint_id) is None
tests/test_cookbook_serve_lifecycle.py
+223
View File
@@ -4,6 +4,13 @@ import pytest
from src import cookbook_serve_lifecycle as lifecycle
from tests.test_cookbook_endpoint_registration import (
endpoint_db,
endpoint_http_api,
_endpoint,
_launch_llm,
)
@pytest.mark.asyncio
async def test_tick_persists_only_successfully_stopped_serves(tmp_path, monkeypatch):
@@ -48,3 +55,219 @@ async def test_tick_persists_only_successfully_stopped_serves(tmp_path, monkeypa
assert tasks["stop-succeeds"]["_scheduledStopAtMs"] is None
assert tasks["stop-fails"]["status"] == "running"
assert tasks["stop-fails"]["_scheduledStopAtMs"] == 0
@pytest.mark.asyncio
async def test_scheduled_stop_preserves_endpoint_reused_by_newer_serve(
endpoint_db,
endpoint_http_api,
):
original = _launch_llm(
"organization-a/Model-A",
8001,
serve_session_id="serve-old",
)
reused = _launch_llm(
"organization-b/Model-B",
8001,
serve_session_id="serve-new",
)
assert original == reused
await lifecycle._delete_endpoint_for_task({
"_endpointId": original,
"sessionId": "serve-old",
"type": "serve",
})
assert _endpoint(endpoint_db, reused) is not None
await lifecycle._delete_endpoint_for_task({
"_endpointId": reused,
"sessionId": "serve-new",
"type": "serve",
})
assert _endpoint(endpoint_db, reused) is None
@pytest.mark.asyncio
async def test_scheduled_stop_preserves_legacy_endpoint_without_owner(
endpoint_db,
endpoint_http_api,
):
endpoint_id = _launch_llm(
"organization-a/Model-A",
8002,
)
await lifecycle._delete_endpoint_for_task({
"_endpointId": endpoint_id,
"sessionId": "serve-old",
"type": "serve",
})
assert _endpoint(endpoint_db, endpoint_id) is not None
@pytest.mark.asyncio
async def test_scheduled_stop_cannot_delete_manual_endpoint(
endpoint_db,
endpoint_http_api,
):
from core.database import ModelEndpoint
with endpoint_db() as db:
db.add(ModelEndpoint(
id="manual-endpoint",
name="Manual",
base_url="http://localhost:9000/v1",
cookbook_serve_session_id="serve-old",
))
db.commit()
await lifecycle._delete_endpoint_for_task({
"_endpointId": "manual-endpoint",
"sessionId": "serve-old",
"type": "serve",
})
assert _endpoint(endpoint_db, "manual-endpoint") is not None
@pytest.mark.asyncio
async def test_scheduled_stop_requires_session_identity(endpoint_db, endpoint_http_api):
endpoint_id = _launch_llm(
"organization-a/Model-A",
8003,
serve_session_id="serve-current",
)
await lifecycle._delete_endpoint_for_task({
"_endpointId": endpoint_id,
"type": "serve",
})
assert _endpoint(endpoint_db, endpoint_id) is not None
def _launch_image(repo_id, port, serve_session_id):
import routes.cookbook_routes as cookbook_routes
from routes.cookbook_helpers import ServeRequest
router = cookbook_routes.setup_cookbook_routes()
serve = next(
route.endpoint
for route in router.routes
if route.path == "/api/model/serve"
and "POST" in route.methods
)
cells = dict(
zip(
serve.__code__.co_freevars,
serve.__closure__ or (),
)
)
register = cells["_auto_register_image_endpoint"].cell_contents
endpoint_id = register(
ServeRequest(
repo_id=repo_id,
cmd=f"python scripts/diffusion_server.py --port {port}",
),
None,
serve_session_id=serve_session_id,
)
assert endpoint_id
return endpoint_id
@pytest.mark.asyncio
async def test_scheduled_stop_deletes_owned_image_endpoint(
endpoint_db,
monkeypatch,
endpoint_http_api,
):
import src.settings as settings
from core.database import ModelEndpoint
monkeypatch.setattr(settings, "load_settings", lambda: {})
monkeypatch.setattr(
settings,
"save_settings",
lambda value: None,
)
endpoint_id = _launch_image(
"organization-a/Image-A",
8101,
"serve-image-a",
)
assert endpoint_id.startswith("img-")
with endpoint_db() as db:
endpoint = db.get(ModelEndpoint, endpoint_id)
assert endpoint.cookbook_serve_session_id == "serve-image-a"
await lifecycle._delete_endpoint_for_task({
"_endpointId": endpoint_id,
"sessionId": "serve-image-a",
"type": "serve",
})
assert _endpoint(endpoint_db, endpoint_id) is None
@pytest.mark.asyncio
async def test_old_image_serve_cannot_delete_reused_endpoint(
endpoint_db,
monkeypatch,
endpoint_http_api,
):
import src.settings as settings
monkeypatch.setattr(settings, "load_settings", lambda: {})
monkeypatch.setattr(
settings,
"save_settings",
lambda value: None,
)
original = _launch_image(
"organization-a/Image-A",
8102,
"serve-old",
)
reused = _launch_image(
"organization-b/Image-B",
8102,
"serve-new",
)
assert original == reused
await lifecycle._delete_endpoint_for_task({
"_endpointId": original,
"sessionId": "serve-old",
"type": "serve",
})
assert _endpoint(endpoint_db, reused) is not None
await lifecycle._delete_endpoint_for_task({
"_endpointId": reused,
"sessionId": "serve-new",
"type": "serve",
})
assert _endpoint(endpoint_db, reused) is None