From 1844b084bc031793c48beb4dc3ec91e0f148b5c1 Mon Sep 17 00:00:00 2001 From: Yu-Wei Chang Date: Thu, 5 Nov 2026 18:00:44 -0800 Subject: [PATCH] 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 --- routes/cookbook/__init__.py | 1 + routes/cookbook/endpoint_registration.py | 284 +++++++++++++++++++ routes/cookbook_routes.py | 274 +----------------- tests/test_cookbook_endpoint_registration.py | 99 +++++-- tests/test_cookbook_serve_lifecycle.py | 22 +- 5 files changed, 371 insertions(+), 309 deletions(-) create mode 100644 routes/cookbook/__init__.py create mode 100644 routes/cookbook/endpoint_registration.py diff --git a/routes/cookbook/__init__.py b/routes/cookbook/__init__.py new file mode 100644 index 000000000..3aba3c152 --- /dev/null +++ b/routes/cookbook/__init__.py @@ -0,0 +1 @@ +"""Cookbook route domain package (#4071).""" diff --git a/routes/cookbook/endpoint_registration.py b/routes/cookbook/endpoint_registration.py new file mode 100644 index 000000000..a4c5602a6 --- /dev/null +++ b/routes/cookbook/endpoint_registration.py @@ -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() diff --git a/routes/cookbook_routes.py b/routes/cookbook_routes.py index 7d9e6e0cb..848b0b0f2 100644 --- a/routes/cookbook_routes.py +++ b/routes/cookbook_routes.py @@ -43,6 +43,9 @@ from routes.cookbook_output import ( error_aware_output_tail, classify_dead_download, 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__) @@ -1408,90 +1411,6 @@ def setup_cookbook_routes() -> APIRouter: 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( remote: str | None, ssh_port: str | None, start_port: int, max_offset: int ) -> int | None: @@ -1680,193 +1599,6 @@ def setup_cookbook_routes() -> APIRouter: 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, - 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") async def model_serve(request: Request, req: ServeRequest): """Launch a model server in a tmux session (or PowerShell background process on Windows). diff --git a/tests/test_cookbook_endpoint_registration.py b/tests/test_cookbook_endpoint_registration.py index 802fbbf0c..614748868 100644 --- a/tests/test_cookbook_endpoint_registration.py +++ b/tests/test_cookbook_endpoint_registration.py @@ -50,17 +50,10 @@ def _launch_llm( port: int, serve_session_id: str | None = None, ) -> str: - import routes.cookbook_routes as cookbook_routes + from routes.cookbook.endpoint_registration import _auto_register_llm_endpoint 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( + endpoint_id = _auto_register_llm_endpoint( ServeRequest(repo_id=repo_id, cmd=f"vllm serve {repo_id} --port {port}"), None, serve_session_id=serve_session_id, @@ -76,6 +69,83 @@ def _endpoint(factory, endpoint_id: str): 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]) def test_launch_keeps_endpoint_of_other_model_with_same_short_name(endpoint_db, 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: - import routes.cookbook_routes as cookbook_routes import src.settings as settings + from routes.cookbook.endpoint_registration import _auto_register_image_endpoint from routes.cookbook_helpers import ServeRequest monkeypatch.setattr(settings, "load_settings", lambda: {"image_gen_enabled": True, "image_model": repo_id}) monkeypatch.setattr(settings, "save_settings", lambda _settings: None) - 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( + endpoint_id = _auto_register_image_endpoint( ServeRequest(repo_id=repo_id, cmd=f"python diffusion_server.py --model {repo_id} --port {port}"), None ) assert endpoint_id diff --git a/tests/test_cookbook_serve_lifecycle.py b/tests/test_cookbook_serve_lifecycle.py index f8fd25720..11aaee39c 100644 --- a/tests/test_cookbook_serve_lifecycle.py +++ b/tests/test_cookbook_serve_lifecycle.py @@ -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): - import routes.cookbook_routes as cookbook_routes + from routes.cookbook.endpoint_registration import _auto_register_image_endpoint 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( + endpoint_id = _auto_register_image_endpoint( ServeRequest( repo_id=repo_id, cmd=f"python scripts/diffusion_server.py --port {port}",