refactor(cookbook): move endpoint registration out of the routes closure (#7031)

* refactor(cookbook): move endpoint registration out of the routes closure

_auto_register_llm_endpoint and _auto_register_image_endpoint use no
closure state and model_serve is their only caller, so they can live in a
plain module and be imported and tested directly. Bodies are unchanged
apart from indentation.

Tests now import the helpers instead of reading them from the route's
closure. That lookup also failed if the route stopped referencing a
helper, so add a test that pins which helper model_serve calls.

Part of #4071

* fix(cookbook): bound the OLLAMA_HOST port scan

The pattern rescanned the rest of the token from every OLLAMA_HOST=
occurrence, so a long serve command with repeated assignments took
quadratic time. model_serve runs this on the event loop, and a 192 KB
command stalled the whole app for about 9s.

Cap the host at 253 characters, the longest DNS name, so every valid host
still matches.

Part of #4071
This commit is contained in:
Yu-Wei Chang
2026-11-06 02:00:44 +00:00
committed by GitHub
parent c21980f5a6
commit 1844b084bc
5 changed files with 371 additions and 309 deletions
+1
View File
@@ -0,0 +1 @@
"""Cookbook route domain package (#4071)."""
+284
View File
@@ -0,0 +1,284 @@
"""Auto-register Cookbook-served models as model endpoints."""
import json
import logging
import uuid
from routes.cookbook_helpers import ServeRequest
logger = logging.getLogger(__name__)
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
from routes.model_routes import _normalize_model_ids
from src.settings import load_settings, save_settings
# Parse port from command (--port NNNN), default 8100 for diffusion_server
port_match = re.search(r'--port\s+(\d+)', req.cmd)
port = int(port_match.group(1)) if port_match else 8100
# Determine host
if remote:
# SSH alias — use as hostname (Tailscale resolves it later)
host = remote.split("@")[-1] if "@" in remote else remote
else:
host = "localhost"
base_url = f"http://{host}:{port}/v1"
# Friendly display name from repo_id
short_name = req.repo_id.split("/")[-1] if "/" in req.repo_id else req.repo_id
display_name = f"{short_name} (image)"
pinned_models = [req.repo_id] if req.repo_id else []
db = SessionLocal()
try:
# Check for existing endpoint with same base_url — update it
existing = db.query(ModelEndpoint).filter(ModelEndpoint.base_url == base_url).first()
if existing:
existing.is_enabled = True
existing.model_type = "image"
# Follow the newly served model unless the user renamed the endpoint.
pinned_names = {f"{m.split('/')[-1]} (image)" for m in _normalize_model_ids(existing.pinned_models)}
if existing.name in pinned_names:
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)
db.commit()
settings = load_settings()
if settings.get("image_gen_enabled") is not True:
settings["image_gen_enabled"] = True
save_settings(settings)
logger.info(f"Updated existing image endpoint: {base_url}")
return existing.id
ep_id = f"img-{uuid.uuid4().hex[:8]}"
ep = ModelEndpoint(
id=ep_id,
name=display_name,
base_url=base_url,
api_key=None,
is_enabled=True,
model_type="image",
endpoint_kind="local",
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()
settings = load_settings()
settings["image_gen_enabled"] = True
if not settings.get("image_model"):
settings["image_model"] = req.repo_id
save_settings(settings)
logger.info(f"Auto-registered image endpoint: {display_name} @ {base_url}")
return ep_id
except Exception as e:
logger.error(f"Failed to auto-register image endpoint: {e}")
db.rollback()
return None
finally:
db.close()
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.
Cookbook serve commands launch an OpenAI-compatible server (llama.cpp's
llama-server, vLLM, SGLang, or Ollama) on a known port. We point an
endpoint at that server's /v1; the picker auto-discovers the model id by
probing /v1/models and dims the endpoint until the server is reachable,
so registering immediately (before the server finishes loading) is safe.
"""
logger.info(
f"_auto_register_llm_endpoint: ENTRY repo_id={req.repo_id!r} "
f"remote={remote!r} cmd_prefix={req.cmd[:80]!r}"
)
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:
# 1. explicit `--port N` (vllm / sglang / llama-server)
# 2. `OLLAMA_HOST=host:port` (the way Ollama specifies its bind)
# 3. fallback by backend (11434 ollama / 8080 llama.cpp)
# Previously the OLLAMA_HOST form was silently ignored and we
# registered every Ollama endpoint at 11434 — even if the user
# set OLLAMA_HOST=0.0.0.0:11435 to avoid colliding with an
# existing systemd Ollama, the registered endpoint pointed at
# the OLD port and showed as offline.
port_match = re.search(r'--port\s+(\d+)', req.cmd)
# Host length is capped at the longest DNS name; unbounded, repeated
# assignments make this scan quadratic and stall the event loop.
ollama_host_match = re.search(r'OLLAMA_HOST=[^\s]{0,253}?:(\d+)', req.cmd)
if port_match:
port = int(port_match.group(1))
elif ollama_host_match:
port = int(ollama_host_match.group(1))
elif "ollama" in req.cmd:
port = 11434
else:
port = 8080 # llama.cpp's llama-server default — the Apple Silicon path
# Determine host. The cookbook tmux for `local=true` serves runs INSIDE
# the odysseus container — so the right URL for the in-container
# backend to reach it is `localhost`, NOT `host.docker.internal`
# (the latter points at the docker HOST, which doesn't have a server
# on that port). The previous host.docker.internal fallback only made
# sense for /setup-added external services like systemd Ollama on the
# host — and those go through manual setup, not this auto-register
# code path. For remote serves we still use the SSH host alias.
if remote:
host = remote.split("@")[-1] if "@" in remote else remote
elif re.search(r"\bdocker\s+exec\s+(?:ollama-rocm|ollama-test)\b", req.cmd or ""):
host = "host.docker.internal"
else:
host = "localhost"
base_url = f"http://{host}:{port}/v1"
short_name = req.repo_id.split("/")[-1] if "/" in req.repo_id else req.repo_id
display_name = short_name or "Local model"
is_mlx_deepseek_v4 = (
"mlx_lm.server" in (req.cmd or "")
and "deepseek-v4" in ((req.repo_id or "") + " " + (req.cmd or "")).lower()
)
mlx_shim_model_id = ""
if is_mlx_deepseek_v4 and short_name:
home_match = re.search(r"((?:/Users|/home)/[^/\s'\"]+)", req.cmd or "")
remote_home = home_match.group(1) if home_match else ""
if remote_home:
mlx_shim_model_id = f"{remote_home}/.cache/odysseus/mlx-shims/{short_name}"
# If the serve command opts models into OpenAI tool-calling, record it so
# agent_loop trusts emitted tool_calls instead of the name heuristic.
is_ollama_endpoint = "ollama" in (req.cmd or "").lower()
supports_tools = True if "--enable-auto-tool-choice" in req.cmd else None
# Pin the model the user launched for every Cookbook-created LLM
# endpoint, not just Ollama. Some OpenAI-compatible servers report a
# 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.
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.
existing = db.query(ModelEndpoint).filter(ModelEndpoint.base_url == base_url).first()
if existing:
existing.is_enabled = True
existing.model_type = "llm"
# Follow the newly served model unless the user renamed the endpoint.
pinned_names = {m.split("/")[-1] for m in _normalize_model_ids(existing.pinned_models)}
if existing.name in pinned_names:
existing.name = display_name
existing.endpoint_kind = "local"
existing.model_refresh_mode = "auto"
# 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:
existing.cached_models = json.dumps(pinned_models)
if supports_tools is not None:
existing.supports_tools = supports_tools
db.commit()
logger.info(f"Updated existing local model endpoint: {base_url}")
# Re-probe so cached_models matches what the server actually
# serves right now (the URL may have stayed the same but the
# model behind it changed across launches).
try:
if mlx_shim_model_id:
existing.cached_models = json.dumps([mlx_shim_model_id])
existing.pinned_models = json.dumps([mlx_shim_model_id])
db.commit()
else:
from routes.model_routes import _probe_endpoint
import json as _json2
probed = _probe_endpoint(base_url, existing.api_key, timeout=5)
if probed:
existing.cached_models = _json2.dumps(probed)
db.commit()
except Exception as _pe:
logger.warning(f"Re-probe failed for {base_url}: {_pe!r}")
return existing.id
ep_id = f"local-{uuid.uuid4().hex[:8]}"
ep = ModelEndpoint(
id=ep_id,
name=display_name,
base_url=base_url,
api_key=None,
is_enabled=True,
model_type="llm",
endpoint_kind="ollama" if is_ollama_endpoint else "local",
model_refresh_mode="auto",
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}")
# 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
# cached_models until the next background refresh fires (up
# to a minute later) and the picker shows nothing — even
# though the endpoint is in the DB and the server is up.
try:
if mlx_shim_model_id:
ep.cached_models = json.dumps([mlx_shim_model_id])
ep.pinned_models = json.dumps([mlx_shim_model_id])
db.commit()
logger.info(f"Auto-register: pinned MLX DeepSeek-V4 shim model @ {base_url}")
else:
from routes.model_routes import _probe_endpoint
import json as _json2
probed = _probe_endpoint(base_url, None, timeout=5)
if probed:
ep.cached_models = _json2.dumps(probed)
db.commit()
logger.info(f"Auto-register: probed {len(probed)} models @ {base_url}")
except Exception as _pe:
logger.warning(f"Auto-register: probe-after-create failed for {base_url}: {_pe!r}")
return ep_id
except Exception as e:
logger.error(f"Failed to auto-register local model endpoint: {e}")
db.rollback()
return None
finally:
db.close()
+3 -271
View File
@@ -43,6 +43,9 @@ from routes.cookbook_output import (
error_aware_output_tail, classify_dead_download, error_aware_output_tail, classify_dead_download,
HF_CACHE_COMPLETE_PROBE, HF_CACHE_INCOMPLETE_PROBE, HF_CACHE_COMPLETE_PROBE, HF_CACHE_INCOMPLETE_PROBE,
) )
from routes.cookbook.endpoint_registration import (
_auto_register_image_endpoint, _auto_register_llm_endpoint,
)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -1408,90 +1411,6 @@ def setup_cookbook_routes() -> APIRouter:
return {"models": models, "host": host or "local"} return {"models": models, "host": host or "local"}
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
from routes.model_routes import _normalize_model_ids
from src.settings import load_settings, save_settings
# Parse port from command (--port NNNN), default 8100 for diffusion_server
port_match = re.search(r'--port\s+(\d+)', req.cmd)
port = int(port_match.group(1)) if port_match else 8100
# Determine host
if remote:
# SSH alias — use as hostname (Tailscale resolves it later)
host = remote.split("@")[-1] if "@" in remote else remote
else:
host = "localhost"
base_url = f"http://{host}:{port}/v1"
# Friendly display name from repo_id
short_name = req.repo_id.split("/")[-1] if "/" in req.repo_id else req.repo_id
display_name = f"{short_name} (image)"
pinned_models = [req.repo_id] if req.repo_id else []
db = SessionLocal()
try:
# Check for existing endpoint with same base_url — update it
existing = db.query(ModelEndpoint).filter(ModelEndpoint.base_url == base_url).first()
if existing:
existing.is_enabled = True
existing.model_type = "image"
# Follow the newly served model unless the user renamed the endpoint.
pinned_names = {f"{m.split('/')[-1]} (image)" for m in _normalize_model_ids(existing.pinned_models)}
if existing.name in pinned_names:
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)
db.commit()
settings = load_settings()
if settings.get("image_gen_enabled") is not True:
settings["image_gen_enabled"] = True
save_settings(settings)
logger.info(f"Updated existing image endpoint: {base_url}")
return existing.id
ep_id = f"img-{uuid.uuid4().hex[:8]}"
ep = ModelEndpoint(
id=ep_id,
name=display_name,
base_url=base_url,
api_key=None,
is_enabled=True,
model_type="image",
endpoint_kind="local",
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()
settings = load_settings()
settings["image_gen_enabled"] = True
if not settings.get("image_model"):
settings["image_model"] = req.repo_id
save_settings(settings)
logger.info(f"Auto-registered image endpoint: {display_name} @ {base_url}")
return ep_id
except Exception as e:
logger.error(f"Failed to auto-register image endpoint: {e}")
db.rollback()
return None
finally:
db.close()
def _pick_free_port_for_ollama( def _pick_free_port_for_ollama(
remote: str | None, ssh_port: str | None, start_port: int, max_offset: int remote: str | None, ssh_port: str | None, start_port: int, max_offset: int
) -> int | None: ) -> int | None:
@@ -1680,193 +1599,6 @@ def setup_cookbook_routes() -> APIRouter:
return return
logger.debug(f"crash-watchdog: no exit marker for {session_id} within window; leaving endpoint {endpoint_id}") 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,
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.
Cookbook serve commands launch an OpenAI-compatible server (llama.cpp's
llama-server, vLLM, SGLang, or Ollama) on a known port. We point an
endpoint at that server's /v1; the picker auto-discovers the model id by
probing /v1/models and dims the endpoint until the server is reachable,
so registering immediately (before the server finishes loading) is safe.
"""
logger.info(
f"_auto_register_llm_endpoint: ENTRY repo_id={req.repo_id!r} "
f"remote={remote!r} cmd_prefix={req.cmd[:80]!r}"
)
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:
# 1. explicit `--port N` (vllm / sglang / llama-server)
# 2. `OLLAMA_HOST=host:port` (the way Ollama specifies its bind)
# 3. fallback by backend (11434 ollama / 8080 llama.cpp)
# Previously the OLLAMA_HOST form was silently ignored and we
# registered every Ollama endpoint at 11434 — even if the user
# set OLLAMA_HOST=0.0.0.0:11435 to avoid colliding with an
# existing systemd Ollama, the registered endpoint pointed at
# the OLD port and showed as offline.
port_match = re.search(r'--port\s+(\d+)', req.cmd)
ollama_host_match = re.search(r'OLLAMA_HOST=[^\s]*?:(\d+)', req.cmd)
if port_match:
port = int(port_match.group(1))
elif ollama_host_match:
port = int(ollama_host_match.group(1))
elif "ollama" in req.cmd:
port = 11434
else:
port = 8080 # llama.cpp's llama-server default — the Apple Silicon path
# Determine host. The cookbook tmux for `local=true` serves runs INSIDE
# the odysseus container — so the right URL for the in-container
# backend to reach it is `localhost`, NOT `host.docker.internal`
# (the latter points at the docker HOST, which doesn't have a server
# on that port). The previous host.docker.internal fallback only made
# sense for /setup-added external services like systemd Ollama on the
# host — and those go through manual setup, not this auto-register
# code path. For remote serves we still use the SSH host alias.
if remote:
host = remote.split("@")[-1] if "@" in remote else remote
elif re.search(r"\bdocker\s+exec\s+(?:ollama-rocm|ollama-test)\b", req.cmd or ""):
host = "host.docker.internal"
else:
host = "localhost"
base_url = f"http://{host}:{port}/v1"
short_name = req.repo_id.split("/")[-1] if "/" in req.repo_id else req.repo_id
display_name = short_name or "Local model"
is_mlx_deepseek_v4 = (
"mlx_lm.server" in (req.cmd or "")
and "deepseek-v4" in ((req.repo_id or "") + " " + (req.cmd or "")).lower()
)
mlx_shim_model_id = ""
if is_mlx_deepseek_v4 and short_name:
home_match = re.search(r"((?:/Users|/home)/[^/\s'\"]+)", req.cmd or "")
remote_home = home_match.group(1) if home_match else ""
if remote_home:
mlx_shim_model_id = f"{remote_home}/.cache/odysseus/mlx-shims/{short_name}"
# If the serve command opts models into OpenAI tool-calling, record it so
# agent_loop trusts emitted tool_calls instead of the name heuristic.
is_ollama_endpoint = "ollama" in (req.cmd or "").lower()
supports_tools = True if "--enable-auto-tool-choice" in req.cmd else None
# Pin the model the user launched for every Cookbook-created LLM
# endpoint, not just Ollama. Some OpenAI-compatible servers report a
# 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.
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.
existing = db.query(ModelEndpoint).filter(ModelEndpoint.base_url == base_url).first()
if existing:
existing.is_enabled = True
existing.model_type = "llm"
# Follow the newly served model unless the user renamed the endpoint.
pinned_names = {m.split("/")[-1] for m in _normalize_model_ids(existing.pinned_models)}
if existing.name in pinned_names:
existing.name = display_name
existing.endpoint_kind = "local"
existing.model_refresh_mode = "auto"
# 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:
existing.cached_models = json.dumps(pinned_models)
if supports_tools is not None:
existing.supports_tools = supports_tools
db.commit()
logger.info(f"Updated existing local model endpoint: {base_url}")
# Re-probe so cached_models matches what the server actually
# serves right now (the URL may have stayed the same but the
# model behind it changed across launches).
try:
if mlx_shim_model_id:
existing.cached_models = json.dumps([mlx_shim_model_id])
existing.pinned_models = json.dumps([mlx_shim_model_id])
db.commit()
else:
from routes.model_routes import _probe_endpoint
import json as _json2
probed = _probe_endpoint(base_url, existing.api_key, timeout=5)
if probed:
existing.cached_models = _json2.dumps(probed)
db.commit()
except Exception as _pe:
logger.warning(f"Re-probe failed for {base_url}: {_pe!r}")
return existing.id
ep_id = f"local-{uuid.uuid4().hex[:8]}"
ep = ModelEndpoint(
id=ep_id,
name=display_name,
base_url=base_url,
api_key=None,
is_enabled=True,
model_type="llm",
endpoint_kind="ollama" if is_ollama_endpoint else "local",
model_refresh_mode="auto",
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}")
# 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
# cached_models until the next background refresh fires (up
# to a minute later) and the picker shows nothing — even
# though the endpoint is in the DB and the server is up.
try:
if mlx_shim_model_id:
ep.cached_models = json.dumps([mlx_shim_model_id])
ep.pinned_models = json.dumps([mlx_shim_model_id])
db.commit()
logger.info(f"Auto-register: pinned MLX DeepSeek-V4 shim model @ {base_url}")
else:
from routes.model_routes import _probe_endpoint
import json as _json2
probed = _probe_endpoint(base_url, None, timeout=5)
if probed:
ep.cached_models = _json2.dumps(probed)
db.commit()
logger.info(f"Auto-register: probed {len(probed)} models @ {base_url}")
except Exception as _pe:
logger.warning(f"Auto-register: probe-after-create failed for {base_url}: {_pe!r}")
return ep_id
except Exception as e:
logger.error(f"Failed to auto-register local model endpoint: {e}")
db.rollback()
return None
finally:
db.close()
@router.post("/api/model/serve") @router.post("/api/model/serve")
async def model_serve(request: Request, req: ServeRequest): async def model_serve(request: Request, req: ServeRequest):
"""Launch a model server in a tmux session (or PowerShell background process on Windows). """Launch a model server in a tmux session (or PowerShell background process on Windows).
+81 -18
View File
@@ -50,17 +50,10 @@ def _launch_llm(
port: int, port: int,
serve_session_id: str | None = None, serve_session_id: str | None = None,
) -> str: ) -> str:
import routes.cookbook_routes as cookbook_routes from routes.cookbook.endpoint_registration import _auto_register_llm_endpoint
from routes.cookbook_helpers import ServeRequest from routes.cookbook_helpers import ServeRequest
router = cookbook_routes.setup_cookbook_routes() endpoint_id = _auto_register_llm_endpoint(
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}"), ServeRequest(repo_id=repo_id, cmd=f"vllm serve {repo_id} --port {port}"),
None, None,
serve_session_id=serve_session_id, serve_session_id=serve_session_id,
@@ -76,6 +69,83 @@ def _endpoint(factory, endpoint_id: str):
return db.get(database.ModelEndpoint, endpoint_id) return db.get(database.ModelEndpoint, endpoint_id)
@pytest.mark.parametrize(
("cmd", "registers", "skips"),
[
(
"python scripts/diffusion_server.py --port 8100",
"_auto_register_image_endpoint",
"_auto_register_llm_endpoint",
),
(
"llama-server -m /models/model.gguf --port 8085",
"_auto_register_llm_endpoint",
"_auto_register_image_endpoint",
),
],
)
@pytest.mark.asyncio
async def test_serve_registers_endpoint_for_its_model_type(monkeypatch, tmp_path, cmd, registers, skips):
from unittest.mock import AsyncMock, MagicMock
from starlette.requests import Request
from routes import cookbook_routes
from routes.cookbook_helpers import ServeRequest
launch =MagicMock(returncode=0, wait=AsyncMock(return_value=0))
monkeypatch.setattr(cookbook_routes, "TMUX_LOG_DIR", tmp_path)
monkeypatch.setattr(cookbook_routes, "require_admin", lambda request: None)
monkeypatch.setattr(cookbook_routes, "_binary_available", AsyncMock(return_value=True))
monkeypatch.setattr(cookbook_routes, "load_stored_hf_token", lambda **kwargs: "")
monkeypatch.setattr(cookbook_routes.asyncio, "create_subprocess_shell", AsyncMock(return_value=launch))
# Returning no endpoint id keeps the crash watchdog from being scheduled.
register = MagicMock(return_value=None)
skip = MagicMock(return_value=None)
monkeypatch.setattr(cookbook_routes, registers, register)
monkeypatch.setattr(cookbook_routes, skips, skip)
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
)
req = ServeRequest(repo_id="organization-a/Model-A", cmd=cmd, remote_host="gpu-box")
response = await serve(Request({"type": "http", "headers": []}), req)
register.assert_called_once_with(req, "gpu-box", serve_session_id=response["session_id"])
skip.assert_not_called()
def _register_ollama(cmd: str) -> str:
from routes.cookbook.endpoint_registration import _auto_register_llm_endpoint
from routes.cookbook_helpers import ServeRequest
endpoint_id = _auto_register_llm_endpoint(ServeRequest(repo_id="library/model", cmd=cmd), None)
assert endpoint_id
return endpoint_id
@pytest.mark.parametrize("host", ["0.0.0.0", "h" * 253])
def test_registers_port_from_ollama_host(endpoint_db, host):
endpoint_id = _register_ollama(f"OLLAMA_HOST={host}:11435 ollama serve")
assert _endpoint(endpoint_db, endpoint_id).base_url == "http://localhost:11435/v1"
def test_repeated_ollama_host_assignments_register_quickly(endpoint_db):
import time
start = time.perf_counter()
endpoint_id = _register_ollama("OLLAMA_HOST=" * 32000 + " ollama serve")
elapsed = time.perf_counter() - start
# The unbounded pattern took about 30s on this input and blocked the event loop.
assert elapsed < 4.0, f"registration took {elapsed:.2f}s"
assert _endpoint(endpoint_db, endpoint_id).base_url == "http://localhost:11434/v1"
@pytest.mark.parametrize("reuse", [False, True]) @pytest.mark.parametrize("reuse", [False, True])
def test_launch_keeps_endpoint_of_other_model_with_same_short_name(endpoint_db, reuse): def test_launch_keeps_endpoint_of_other_model_with_same_short_name(endpoint_db, reuse):
if reuse: if reuse:
@@ -448,20 +518,13 @@ async def test_ordinary_admin_delete_does_not_require_serve_owner(
def _launch_image(repo_id: str, port: int, monkeypatch) -> str: def _launch_image(repo_id: str, port: int, monkeypatch) -> str:
import routes.cookbook_routes as cookbook_routes
import src.settings as settings import src.settings as settings
from routes.cookbook.endpoint_registration import _auto_register_image_endpoint
from routes.cookbook_helpers import ServeRequest from routes.cookbook_helpers import ServeRequest
monkeypatch.setattr(settings, "load_settings", lambda: {"image_gen_enabled": True, "image_model": repo_id}) monkeypatch.setattr(settings, "load_settings", lambda: {"image_gen_enabled": True, "image_model": repo_id})
monkeypatch.setattr(settings, "save_settings", lambda _settings: None) monkeypatch.setattr(settings, "save_settings", lambda _settings: None)
router = cookbook_routes.setup_cookbook_routes() endpoint_id = _auto_register_image_endpoint(
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 diffusion_server.py --model {repo_id} --port {port}"), None ServeRequest(repo_id=repo_id, cmd=f"python diffusion_server.py --model {repo_id} --port {port}"), None
) )
assert endpoint_id assert endpoint_id
+2 -20
View File
@@ -156,28 +156,10 @@ async def test_scheduled_stop_requires_session_identity(endpoint_db, endpoint_ht
def _launch_image(repo_id, port, serve_session_id): def _launch_image(repo_id, port, serve_session_id):
import routes.cookbook_routes as cookbook_routes from routes.cookbook.endpoint_registration import _auto_register_image_endpoint
from routes.cookbook_helpers import ServeRequest from routes.cookbook_helpers import ServeRequest
router = cookbook_routes.setup_cookbook_routes() endpoint_id = _auto_register_image_endpoint(
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( ServeRequest(
repo_id=repo_id, repo_id=repo_id,
cmd=f"python scripts/diffusion_server.py --port {port}", cmd=f"python scripts/diffusion_server.py --port {port}",